def _process_deriv_args(f, kwargs):
"""Handle common processing of arguments for derivative functions."""
n = f.ndim
axis = normalize_axis_index(kwargs.get('axis', 0), n)
if f.shape[axis] < 3:
raise ValueError('f must have at least 3 point along the desired axis.')
if 'delta' in kwargs:
if 'x' in kwargs:
raise ValueError('Cannot specify both "x" and "delta".')
delta = atleast_1d(kwargs['delta'])
if delta.size == 1:
diff_size = list(f.shape)
diff_size[axis] -= 1
delta_units = getattr(delta, 'units', None)
delta = np.broadcast_to(delta, diff_size, subok=True)
if delta_units is not None:
delta = delta * delta_units
else:
delta = _broadcast_to_axis(delta, axis, n)
elif 'x' in kwargs:
x = _broadcast_to_axis(kwargs['x'], axis, n)
delta = diff(x, axis=axis)
else:
raise ValueError('Must specify either "x" or "delta" for value positions.')
return n, axis, delta
def _process_deriv_args(f, kwargs):
"""Handle common processing of arguments for derivative functions."""
n = f.ndim
axis = normalize_axis_index(kwargs.get('axis', 0), n)
if f.shape[axis] < 3:
raise ValueError(
'f must have at least 3 point along the desired axis.')
if 'delta' in kwargs:
if 'x' in kwargs:
raise ValueError('Cannot specify both "x" and "delta".')
delta = atleast_1d(kwargs['delta'])
if delta.size == 1:
diff_size = list(f.shape)
diff_size[axis] -= 1
delta_units = getattr(delta, 'units', None)
delta = np.broadcast_to(delta, diff_size, subok=True)
if delta_units is not None:
delta = delta * delta_units
else:
delta = _broadcast_to_axis(delta, axis, n)
elif 'x' in kwargs:
x = _broadcast_to_axis(kwargs['x'], axis, n)
delta = diff(x, axis=axis)
else:
raise ValueError(
'Must specify either "x" or "delta" for value positions.')
return n, axis, delta
def _process_deriv_args(f, kwargs):
"""Handle common processing of arguments for derivative functions."""
n = f.ndim
axis = normalize_axis_index(kwargs.get("axis", 0), n)
if f.shape[axis] < 3:
raise ValueError("f must have at least 3 pts along the desired axis.")
if "delta" in kwargs:
if "x" in kwargs:
raise ValueError('Cannot specify both "x" and "delta".')
delta = np.atleast_1d(kwargs["delta"])
if delta.size == 1:
diff_size = list(f.shape)
diff_size[axis] -= 1
delta_units = getattr(delta, "units", None)
delta = np.broadcast_to(delta, diff_size, subok=True)
if not hasattr(delta, "units") and delta_units is not None:
delta = delta * delta_units
else:
delta = _broadcast_to_axis(delta, axis, n)
elif "x" in kwargs:
x = _broadcast_to_axis(kwargs["x"], axis, n)
delta = np.diff(x, axis=axis)
else:
raise ValueError("Must specify either x or delta for value positions.")
return n, axis, delta
def subdomain(datalist, mask, axes=[-2, -1]):
'''
mask: list of slice objects
'''
if len(mask) != len(axes):
raise ValueError("lengths of mask and axes must match")
dataout = []
for data in datalist:
n = data.ndim
l = len(axes)
if n == l:
expanded_mask = mask
elif n > l:
expanded_mask = n * [slice(None)]
for axis, m in zip(axes, mask):
axis = _nx.normalize_axis_index(axis, n)
expanded_mask[axis] = m
else:
raise ValueError(
"data dimensions must have at least the length of axes")
dataout.append(data[expanded_mask])
return dataout
def _datas_axis_to_first(datas, axis, mom_ndim):
"""move axis to first first position"""
# NOTE: removinvg this. should be handles elsewhere
# datas = np.asarray(datas)
# ndim = datas.ndim - mom_ndim
# if axis < 0:
# axis += ndim
# assert 0 <= axis < ndim
axis = normalize_axis_index(axis, datas.ndim - mom_ndim)
if axis != 0:
datas = np.moveaxis(datas, axis, 0)
return datas, axis
def _wrap_axis(self, axis, default=0, ndim=None):
"""wrap axis to positive value and check"""
if axis is None:
axis = default
if ndim is None:
ndim = self.val_ndim
axis = normalize_axis_index(axis, ndim)
# if axis < 0:
# axis += ndim
# assert 0 <= axis < ndim
return axis
def put_along_axis(arr, indices, values, axis):
# copy from numpy
# normalize inputs
if axis is None:
arr = arr.flat
axis = 0
arr_shape = (len(arr), ) # flatiter has no .shape
else:
axis = normalize_axis_index(axis, arr.ndim)
arr_shape = arr.shape
# use the fancy index
return index_update(arr, _make_along_axis_idx(arr_shape, indices, axis),
values)
def _verify_value(
self,
x,
target=None,
axis=None,
broadcast=False,
expand=False,
shape_flat=None,
other=None,
):
"""
verify input values
Parameters
----------
x : array
target : tuple or array
If tuple, this is the target shape to be used to Make target.
If array, this is the target array
Optinal target that has already been rolled. If this is passed, and
x will be broadcast/expanded, can expand to this shape without the need
to reorder,
"""
x = self._asarray(x)
if isinstance(target, str):
if target == "val":
target = self.val_shape
elif target == "vals":
target = _shape_insert_axis(self.val_shape, axis,
x.shape[axis])
elif target == "data":
target = self.shape
elif target == "datas":
# make sure axis in limits
axis = normalize_axis_index(axis, self.val_ndim + 1)
# if axis < 0:
# axis += self.ndim - self.mom_ndim
target = _shape_insert_axis(self.shape, axis, x.shape[axis])
elif target == "var":
target = self.shape_var
elif target == "vars":
target = _shape_insert_axis(self.shape_var, axis,
other.shape[axis])
if isinstance(target, tuple):
target_shape = target
target_output = x
else:
target_shape = target.shape
target_output = None
x = _axis_expand_broadcast(
x,
target_shape,
axis,
verify=False,
expand=expand,
broadcast=broadcast,
dtype=self.dtype,
roll=False,
)
# check shape:
assert (x.shape == target_shape
), f"x.shape = {x.shape} not equal target_shape={target_shape}"
# move axis
if axis is not None:
if axis != 0:
x = np.moveaxis(x, axis, 0)
nrec = (x.shape[0], )
else:
nrec = ()
if shape_flat is not None:
x = x.reshape(nrec + shape_flat)
if x.ndim == 0:
x = x[()]
if target_output is None:
return x
else:
return x, target_output