def jacobianDask(u, v, w, dt, dx, dy, dz):
du = da.gradient(u, dx, dy, dz, dt, axis=(1, 2, 3))
dv = da.gradient(v, dx, dy, dz, dt, axis=(1, 2, 3))
dw = da.gradient(w, dx, dy, dz, dt, axis=(1, 2, 3))
J = da.stack(
(da.stack(du, axis=-1), da.stack(dv, axis=-1), da.stack(dw, axis=-1)),
axis=-2)
return J
def test_gradient(shape, varargs, axis, edge_order):
a = np.random.randint(0, 10, shape)
d_a = da.from_array(a, chunks=(len(shape) * (5,)))
r = np.gradient(a, *varargs, axis=axis, edge_order=edge_order)
r_a = da.gradient(d_a, *varargs, axis=axis, edge_order=edge_order)
if isinstance(axis, Number):
assert_eq(r, r_a)
else:
assert len(r) == len(r_a)
for e_r, e_r_a in zip(r, r_a):
assert_eq(e_r, e_r_a)
def gradient(x, coord, axis, edge_order):
if is_duck_dask_array(x):
return dask_array.gradient(x, coord, axis=axis, edge_order=edge_order)
return np.gradient(x, coord, axis=axis, edge_order=edge_order)
def gradient(x, coord, axis, edge_order):
if isinstance(x, dask_array_type):
return dask_array.gradient(x, coord, axis=axis, edge_order=edge_order)
return np.gradient(x, coord, axis=axis, edge_order=edge_order)
