In Drake, how do I convert an NumPy array to different scalar types? (e.g. from float to AutoDiffXd or Expression?)

Question

In Drake, I have NumPy ndarray's (some multidimensional) that are of dtype=float, and I want to convert them to AutoDiffXd, Expression, etc.

In C++, I know I can do things like this:

MatrixXd X(2, 2);
X <<
   1.0, 2.0,
   3.0, 4.0;
MatrixX<AutoDiffXd> X_ad = X.cast<AutoDiffXd>();

However, in Python, I've found myself writing loops like this:

import numpy as np

from pydrake.autodiffutils import AutoDiffXd

X = np.array([
    [1.0, 2.0],
    [3.0, 4.0],
])

X_ad = np.zeros(X.shape, dtype=object)
for i in X.shape[0]:
    for j in X.shape[1]:
        X_ad[i, j] = AutoDiffXd(X[i, j])

Is there a better way to do this?

Answer 1

You can use np.vectorize to shorten and generalize your code:
https://numpy.org/doc/1.13/reference/generated/numpy.vectorize.html

As example:

to_autodiff = np.vectorize(AutoDiffXd)
X_ad = to_autodiff(X)

If you need more complex logic, you can always customize your function and use a decorator:

@np.vectorize
def maybe_to_autodiff(x):
    if isinstance(x, float):
        return AutoDiffXd(x)
    else:
        return x

X_ad = maybe_to_autodiff(X)

---

This is a (more flexible) analog to Eigen's .cast<T>() method (see "Basic matrix manipulation", the row containing "Assignment/copy"):
https://eigen.tuxfamily.org/dox/group__QuickRefPage.html