from __future__ import annotations
import functools
from typing import Mapping, TYPE_CHECKING, TypeVar, Generic, Type, ClassVar, Sequence, Callable, Collection, Any, Literal
from typing import Self
import abc
import dataclasses
import numpy as np
import astropy.units as u
import named_arrays as na
import named_arrays._core_array_functions as _core_array_functions
from named_arrays._functions import _fields
import itertools
__all__ = [
"InputValueError",
"AbstractFunctionArray",
"FunctionArray",
"AbstractPolynomialFunctionArray",
"PolynomialFitFunctionArray",
]
InputsT = TypeVar("InputsT", bound=na.AbstractArray, covariant=True)
OutputsT = TypeVar("OutputsT", bound=na.AbstractArray, covariant=True)
class InputValueError(ValueError):
"""
Exception raised when the inputs of two functions do not match
"""
def _is_vertex(
axis: str,
shape_inputs: dict[str, int],
shape_outputs: dict[str, int],
) -> bool:
"""
Whether an axis of a function array is a vertex axis, along which the
inputs are the edges of the cells of the outputs, rather than a center
axis.
Parameters
----------
axis
The name of the axis.
shape_inputs
The shape of the inputs of the function array.
shape_outputs
The shape of the outputs of the function array.
"""
if axis in shape_inputs and axis in shape_outputs:
num_inputs = shape_inputs[axis]
num_outputs = shape_outputs[axis]
if num_inputs == num_outputs + 1:
return True
if num_inputs != num_outputs and 1 not in (num_inputs, num_outputs):
raise ValueError(
f"output dimension, "
f"self.outputs.shape[{axis}]={num_outputs},"
f"must either match input dimension,"
f"self.inputs.shape[{axis}]={num_inputs},"
f"(representing bin centers) "
"or exceed by one (representing bin vertices)."
)
return False
def _item_vertex(
item: int | np.integer | slice,
num: int,
axis: str,
) -> tuple[slice, slice]:
"""
Convert an index of the cells along a vertex axis of a function array
into a slice of its outputs and a slice of its inputs.
Parameters
----------
item
An integer or a slice of the cells along `axis`.
Negative values count from the end, as they do for :mod:`numpy`.
An integer selects one cell and keeps the axis, since the cell
still has two vertices.
num
The number of cells along `axis`.
axis
The name of the vertex axis, for the error messages.
Returns
-------
The slice of the outputs, the cells, and the slice of the inputs,
the vertices which bound those cells.
"""
if isinstance(item, slice):
cells = range(num)[item]
if len(cells) <= 1:
# The step of a selection of at most one cell does not change
# which cells it selects, only their direction.
step = 1 if cells.step > 0 else -1
cells = range(cells.start, cells.start + step * len(cells), step)
elif cells.step not in (1, -1):
raise ValueError(
f"A slice along the vertex axis {axis!r} must have a step of "
f"1 or -1, got {item}, since cells which are not adjacent "
f"do not share vertices."
)
else:
try:
index = range(num)[item]
except IndexError:
raise IndexError(
f"index {item} is out of bounds for the vertex axis {axis!r} "
f"with {num} cells"
) from None
cells = range(index, index + 1)
# A selection of no cells keeps the vertex where it starts,
# so that the axis stays a vertex axis.
if cells.step == 1:
start = cells.start
stop = start + len(cells)
return slice(start, stop), slice(start, stop + 1)
else:
# The cells `start, start - 1, ..., stop + 1` are bounded by the
# vertices `start + 1, start, ..., stop + 1`.
# A stop of -1 is written as `None`, since -1 would count from the end.
start = cells.start
stop = start - len(cells)
stop = stop if stop >= 0 else None
outputs = slice(start, stop, -1) if cells else slice(0, 0)
return outputs, slice(start + 1, stop, -1)
@functools.cache
def _array_function_dispatch() -> tuple[dict[Callable, Callable], dict[Callable, Callable]]:
"""
The handlers of the :mod:`numpy` functions which function arrays support,
built on first use since :mod:`named_arrays._functions.function_array_functions`
imports this module.
Returns
-------
A dictionary mapping each function in a category of
:mod:`~named_arrays._functions.function_array_functions` to the handler of
that category, and the ``HANDLED_FUNCTIONS`` of that module.
"""
from named_arrays._core import _array_function_handlers
from . import function_array_functions as f
handlers = _array_function_handlers([
(f.DEFAULT_FUNCTIONS, f.array_function_default),
(f.CUMULATIVE_REDUCE_FUNCTIONS, f.array_function_cumulative_reduce),
(f.PERCENTILE_LIKE_FUNCTIONS, f.array_function_percentile_like),
(f.ARG_REDUCE_FUNCTIONS, f.array_function_arg_reduce),
(f.STACK_LIKE_FUNCTIONS, f.array_function_stack_like),
])
return handlers, f.HANDLED_FUNCTIONS
@dataclasses.dataclass(eq=False, repr=False)
class AbstractFunctionArray(
na.AbstractArray,
):
@property
@abc.abstractmethod
def inputs(self) -> na.AbstractArray:
"""
The arrays representing the inputs to this function.
"""
@property
@abc.abstractmethod
def outputs(self) -> na.AbstractArray:
"""
The arrays representing the outputs of this function.
"""
__named_array_priority__: ClassVar[float] = (100 * na.AbstractVectorArray.__named_array_priority__)
@property
def axes_center(self) -> tuple(str):
"""
Return keys corresponding to all input axes representing bin centers
"""
shape_inputs = self.inputs.shape
shape_outputs = self.outputs.shape
return tuple(
axis for axis in self.axes
if not _is_vertex(axis, shape_inputs, shape_outputs)
)
@property
def axes_vertex(self) -> tuple(str):
"""
Return keys corresponding to all input axes representing bin vertices
"""
axes_center = self.axes_center
return tuple(axis for axis in self.axes if axis not in axes_center)
@property
def type_abstract(self) -> Type[AbstractFunctionArray]:
return AbstractFunctionArray
@property
def type_explicit(self) -> Type[FunctionArray]:
return type(self)
@property
def value(self) -> FunctionArray:
exp = self.explicit
return exp.replace(
outputs=self.outputs.value,
)
@property
def broadcasted(self) -> FunctionArray:
exp = self.explicit
broadcasted_shape_outputs = exp.shape
broadcasted_shape_inputs = broadcasted_shape_outputs.copy()
axes_vertex = exp.axes_vertex
inputs = exp.inputs
for key in broadcasted_shape_outputs:
axes_vertex = axes_vertex
if key in axes_vertex:
broadcasted_shape_inputs[key] = inputs.shape[key]
return exp.replace(
inputs=exp.inputs.broadcast_to(broadcasted_shape_inputs),
outputs=exp.outputs.broadcast_to(broadcasted_shape_outputs),
)
def astype(
self,
dtype: str | np.dtype | type,
order: str = 'K',
casting='unsafe',
subok: bool = True,
copy: bool = True,
) -> FunctionArray:
exp = self.explicit
return exp.replace(
inputs=exp.inputs,
outputs=exp.outputs.astype(
dtype=dtype,
order=order,
casting=casting,
subok=subok,
copy=copy,
)
)
def to(
self,
unit: u.UnitBase,
equivalencies: None | list[tuple[u.Unit, u.Unit]] = [],
copy: bool = True,
) -> FunctionArray:
exp = self.explicit
return exp.replace(
outputs=exp.outputs.to(
unit=unit,
equivalencies=equivalencies,
copy=copy,
),
)
def to_value(
self,
unit: u.UnitBase,
equivalencies: None | list[tuple[u.Unit, u.Unit]] = [],
) -> FunctionArray:
exp = self.explicit
return exp.replace(
outputs=exp.outputs.to_value(
unit=unit,
equivalencies=equivalencies,
),
)
@property
def length(self) -> FunctionArray:
exp = self.explicit
return exp.replace(
outputs=exp.outputs.length,
)
def add_axes(self, axes: str | Sequence[str]) -> FunctionArray:
exp = self.explicit
return exp.replace(
outputs=exp.outputs.add_axes(axes),
)
def combine_axes(
self,
axes: None | Sequence[str] = None,
axis_new: None | str = None,
) -> FunctionArray:
self = self.explicit
axes = self.axes if axes is None else axes
if not set(axes).issubset(set(self.axes)):
raise ValueError(f"Axes {set(axes) - set(self.axes)} not in {self.axes}")
for axis in axes:
if axis in self.axes_vertex:
raise ValueError(f"Axis {axis} describes input vertices and cannot be used in combine_axes.")
inputs = self.inputs
outputs = self.outputs
inputs_shape_broadcasted = inputs.shape
outputs_shape_broadcasted = outputs.shape
shape = self.shape
for axis in axes:
if axis not in inputs_shape_broadcasted:
inputs_shape_broadcasted[axis] = shape[axis]
if axis not in outputs_shape_broadcasted:
outputs_shape_broadcasted[axis] = shape[axis]
inputs = na.broadcast_to(inputs, inputs_shape_broadcasted)
outputs = na.broadcast_to(outputs, outputs_shape_broadcasted)
def combine(v: na.AbstractArray) -> na.AbstractArray:
v = na.broadcast_to(v, v.shape | {ax: shape[ax] for ax in axes})
return v.combine_axes(axes=axes, axis_new=axis_new)
return self.replace(
inputs=inputs.combine_axes(axes=axes, axis_new=axis_new),
outputs=outputs.combine_axes(axes=axes, axis_new=axis_new),
**_fields.apply(self, combine, axes=axes),
)
def __call__(
self,
inputs: na.AbstractArray,
axis: None | str | tuple[str] = None,
method: Literal['multilinear', 'conservative'] = 'multilinear',
weights: None | tuple[na.AbstractScalar, dict[str, int], dict[str, int]] = None,
) -> Self:
return self.regrid(
inputs=inputs,
axis=axis,
method=method,
weights=weights,
)
def regrid(
self,
inputs: na.AbstractArray,
axis: None | str | tuple[str] = None,
method: Literal['multilinear', 'conservative'] = 'multilinear',
weights: None | tuple[na.AbstractScalar, dict[str, int], dict[str, int]] = None,
) -> Self:
"""
Resample this function array onto a new set of input coordinates
using :func:`named_arrays.regridding.regrid`.
Parameters
----------
inputs
The new input coordinates on which to resample the outputs.
axis
The logical axes of the input over which to resample.
method
The resampling method to use.
weights
Optional weights which can be computed in advance using :meth:`weights`
to greatly speed repeated resampling of the same `inputs`.
See Also
--------
:meth:`weights`: If you need to resample the same coordinates more than once.
"""
_self = self.explicit
inputs_new, _, axis_regrid = _self._normalize__regrid__args(
inputs=inputs,
axis=axis,
)
_fields.check_axes(
_self,
axis_regrid,
operation="regrid resamples",
resamples=True,
)
if weights is not None:
_weights, shape_input, shape_output = weights
else:
_weights, shape_input, shape_output = _self.weights(
inputs=inputs,
axis=axis,
method=method,
)
outputs_new = na.regridding.regrid_from_weights(
weights=_weights,
shape_input=shape_input,
shape_output=shape_output,
values_input=_self.outputs,
)
final_coordinates_dict = {}
if isinstance(inputs, na.AbstractVectorArray) and isinstance(_self.inputs, na.AbstractVectorArray):
for c in _self.inputs.cartesian_nd.components:
if inputs.cartesian_nd.components[c] is None: # pragma: no cover
final_coordinates_dict[c] = _self.inputs.cartesian_nd.components[c]
else:
final_coordinates_dict[c] = inputs.cartesian_nd.components[c]
return _self.replace(
inputs=_self.inputs.type_explicit.from_cartesian_nd(
array=na.CartesianNdVectorArray(final_coordinates_dict),
like=_self.inputs
),
outputs=outputs_new,
)
else:
return _self.replace(
inputs=inputs_new,
outputs=outputs_new,
)
def weights(
self,
inputs: na.AbstractArray,
axis: None | str | tuple[str] = None,
method: Literal['multilinear', 'conservative'] = 'multilinear',
) -> tuple[na.AbstractScalar, dict[str, int], dict[str, int]]:
"""
Compute the resampling weights of this array using
:func:`named_arrays.regridding.weights`.
The output of this method is designed to be used by :meth:`regrid`.
Parameters
----------
inputs
The new input coordinates on which to resample the outputs.
axis
The logical axes of the input over which to resample.
method
The resampling method to use.
See Also
--------
:meth:`regrid`: A method designed to use these weights.
"""
coordinates_new, coordinates_old, axis_input = self._normalize__regrid__args(
inputs=inputs,
axis=axis,
)
return na.regridding.weights(
coordinates_input=coordinates_old,
coordinates_output=coordinates_new,
axis_input=axis_input,
method=method
)
def _normalize__regrid__args(
self,
inputs: na.AbstractArray,
axis: None | str | tuple[str],
):
inputs_old = self.inputs
inputs_new = inputs
if isinstance(inputs_new, na.AbstractVectorArray):
new_input_components = inputs_new.cartesian_nd.components
else:
new_input_components = dict(_dummy=inputs_new.explicit)
if isinstance(inputs_old, na.AbstractVectorArray):
old_input_components = inputs_old.cartesian_nd.components
else:
old_input_components = dict(_dummy=inputs_old.explicit)
# broadcast new inputs against value to be interpolated
if axis is None:
axis = na.shape_broadcasted(
new_input_components[c] for c in new_input_components if new_input_components[c] is not None
)
axis = tuple(axis)
# check physical(vector) dimensions of each input match
if new_input_components.keys() == old_input_components.keys():
coordinates_new = {}
coordinates_old = {}
for c in new_input_components:
component = new_input_components[c]
if component is not None:
# if input components logical axes do not include interp axes, skip
if not set(axis).isdisjoint(component.axes):
coordinates_new[c] = component
coordinates_old[c] = old_input_components[c]
else:
# check if uninterpolated physical axes vary along interpolation axes
if not set(axis).isdisjoint(old_input_components[c].axes): # pragma: no cover
raise ValueError(
f"If a component is marked separable using `None`, its shape, {old_input_components[c].axes},"
f"should be disjoint from the interpolated axes, {axis}.",
)
else:
raise ValueError('Physical axes of new and old inputs must match.') # pragma: no cover
if isinstance(inputs_new, na.AbstractVectorArray):
coordinates_new = na.CartesianNdVectorArray.from_components(coordinates_new)
else:
coordinates_new = coordinates_new['_dummy']
if isinstance(inputs_old, na.AbstractVectorArray):
coordinates_old = na.CartesianNdVectorArray.from_components(coordinates_old)
else:
coordinates_old = coordinates_old['_dummy']
return coordinates_new, coordinates_old, axis
def cell_centers(
self,
axis: None | str | Sequence[str] = None,
random: bool = False,
seed: None | int = None,
) -> na.AbstractExplicitArray:
exp = self.explicit
if axis is None:
axes = tuple(na.shape(exp.outputs))
elif isinstance(axis, str):
axes = (axis,)
else:
axes = tuple(axis)
# the fields are not sampled at the same random points as the outputs,
# and the centers of a mask or a count would be fractions
if random:
_fields.check_axes(
exp,
axes,
operation="cell_centers with random=True samples at random points",
resamples=True,
hint=" Or use random=False.",
)
_fields.check_inexact(exp, axes, operation="cell_centers")
def centers(v: na.AbstractArray) -> na.AbstractArray:
return v.cell_centers(axis=tuple(ax for ax in axes if ax in v.shape))
return exp.replace(
inputs=exp.inputs.cell_centers(axis, random=random, seed=seed),
outputs=exp.outputs.cell_centers(axis, random=random, seed=seed),
**_fields.apply(exp, centers, axes=axes),
)
def integrate(
self,
axis: str | Sequence[str],
component: None | str = None,
) -> FunctionArray:
"""
Integrate the outputs over the given input ``axis`` or axes.
The differential measure is computed from the inputs via
:meth:`AbstractArray.volume_cell`:
* If ``component`` is :obj:`None`, the whole input is used as the
integration variable and ``len(axis)`` must match the dimensionality
of the inputs (one for a scalar, two for a
:class:`Cartesian2dVectorArray`, etc.).
* If ``component`` is a :class:`str`, that named sub-element of the
inputs (a scalar or a sub-vector) supplies the measure, and
``len(axis)`` must match its dimensionality.
A vertex axis (where the inputs represent bin edges) is integrated with
a Riemann sum, while a center axis (where the inputs represent samples)
is integrated with the trapezoidal rule; mixing the two in a single
call is supported. The integrated axes are removed from the outputs and
collapsed in the inputs by averaging.
Parameters
----------
axis
The logical axis or axes to integrate over.
component
The named input sub-element supplying the integration variable.
If :obj:`None`, the whole input is used.
Examples
--------
Integrate a constant function over a 1D domain.
.. jupyter-execute::
import numpy as np
import astropy.units as u
import named_arrays as na
f = na.FunctionArray(
inputs=na.ScalarLinearSpace(0, 2, axis="x", num=101) * u.nm,
outputs=na.ScalarArray(np.full(101, 3.0), axes=("x",)) * u.ph,
)
f.integrate("x")
"""
self = self.explicit
if isinstance(axis, str):
axis = (axis,)
axis = tuple(axis)
if not set(axis).issubset(self.axes):
raise ValueError(f"axes {set(axis) - set(self.axes)} not in {self.axes}")
_fields.check_axes(self, axis, operation="integration removes")
axes_vertex = self.axes_vertex
inputs = self.inputs
outputs = self.outputs
# convert outputs to cell centers along the center axes
# (corner-averaging across all center axes is the multi-D trapezoidal rule);
# vertex axes are already at cell centers relative to their input edges.
axes_center = tuple(ax for ax in axis if ax not in axes_vertex)
if axes_center:
outputs = outputs.cell_centers(axes_center)
# compute the joint cell measure from the selected component.
if component is None:
measure = inputs.volume_cell(axis)
else:
measure = inputs.components[component].volume_cell(axis)
outputs = (outputs * measure).sum(axis)
inputs = inputs.mean(axis)
return self.replace(
inputs=inputs,
outputs=outputs,
)
def gradient(
self,
axis: str,
component: None | str = None,
edge_order: int = 1,
) -> FunctionArray:
"""
Differentiate the outputs with respect to the inputs along ``axis``.
This is the counterpart to :meth:`integrate`: where that one sums the
outputs against a measure computed from the inputs, this one divides
the difference of the outputs by the difference of the inputs.
The result has the same inputs as this array, and its outputs are the
derivative, so the units of the result are those of the outputs
divided by those of the differentiation variable.
* If ``component`` is :obj:`None`, the inputs themselves are the
differentiation variable, and they must be a scalar.
* If ``component`` is a :class:`str`, that named sub-element of the
inputs is the variable, and it must be a scalar.
A center axis, where the inputs are samples, differentiates the
outputs against those samples directly.
A vertex axis, where the inputs are bin edges and so has one more
input than output, differentiates against the cell centers, which is
where the outputs live.
Parameters
----------
axis
The logical axis to differentiate along.
component
The named input sub-element which is the differentiation variable.
If :obj:`None`, the inputs themselves are the variable.
edge_order
The order of the one-sided differences used at the two ends of
``axis``, either 1 or 2.
The interior always uses second-order central differences.
Raises
------
ValueError
If ``axis`` is not an axis of this array, if the differentiation
variable is not a scalar or does not vary along ``axis``, or if
there are too few points along ``axis`` for the given
``edge_order``.
See Also
--------
:meth:`integrate` : The counterpart of this method, which sums the
outputs against the inputs rather than dividing by them.
:func:`numpy.gradient` : The :mod:`numpy` function this method backs.
Notes
-----
Unlike :func:`numpy.gradient`, the differentiation variable may vary
along axes other than ``axis``, which is what a distorted grid gives,
since the differences are taken elementwise.
Examples
--------
Differentiate a parabola, whose derivative is a straight line.
Second-order edges are asked for, since the first-order ones of the
default are inexact wherever the function curves.
.. jupyter-execute::
import numpy as np
import astropy.units as u
import named_arrays as na
x = na.linspace(0, 2, axis="x", num=5) * u.mm
f = na.FunctionArray(
inputs=x,
outputs=np.square(x / u.mm) * u.ph,
)
f.gradient("x", edge_order=2)
"""
self = self.explicit
inputs = self.inputs
outputs = self.outputs
if not isinstance(axis, str):
raise TypeError(
f"`axis` must be the name of a single axis, got {axis!r}. "
f"Differentiating along several axes gives one result per "
f"axis, so use `numpy.gradient` for that."
)
if axis not in self.axes:
raise ValueError(f"{axis=} must be a member of {self.axes}")
if edge_order not in (1, 2):
raise ValueError(f"{edge_order=} must be either 1 or 2")
if component is None:
x = inputs
else:
x = inputs.components[component]
if not isinstance(x, na.AbstractScalar):
raise ValueError(
f"the differentiation variable must be a scalar, got "
f"{x.type_abstract}. Name the component to differentiate "
f"against using the `component` argument."
)
# where the inputs are bin edges there is one more of them than there
# are outputs, so the variable is taken at the cell centers, which is
# where the outputs live
if axis in self.axes_vertex:
x = x.cell_centers(axis)
return self.replace(
# a container of its own, so that the result is not the source's
# inputs under another name, as every neighboring operation gives
inputs=inputs.copy_shallow(),
outputs=_core_array_functions._gradient(
f=outputs,
x=x,
axis=axis,
edge_order=edge_order,
),
)
def to_string_array(
self,
format_value: str = "%.2f",
format_unit: str = "latex_inline",
pad_unit: str = r"$\,$",
):
exp = self.explicit
kwargs = dict(
format_value=format_value,
format_unit=format_unit,
pad_unit=pad_unit,
)
return exp.replace(
inputs=na.as_named_array(exp.inputs).to_string_array(**kwargs),
outputs=na.as_named_array(exp.outputs).to_string_array(**kwargs),
)
def _getitem(
self,
item: Mapping[str, int | slice | na.AbstractArray] | na.AbstractArray | na.AbstractFunctionArray,
) -> FunctionArray:
# Implicit inputs and outputs are made explicit once, as they would be
# by `explicit`, except inputs which are a WCS vector, which can be
# indexed without computing every coordinate and so stay implicit.
array = self
inputs = array.inputs
if not isinstance(inputs, na.AbstractWcsVector):
inputs = na.explicit(inputs)
outputs = na.explicit(array.outputs)
if isinstance(item, na.AbstractArray):
if isinstance(item, na.AbstractFunctionArray):
if not np.all(item.inputs == array.inputs):
raise ValueError("item inputs do not match function inputs.")
axes = set(inputs.axes) - set(array.axes_center)
inputs = inputs.cell_centers(axis=axes)
item_inputs = item.outputs
item_outputs = item.outputs
else:
item_inputs = item_outputs = item
shape_item_inputs = item_inputs.shape
shape_item_outputs = item_outputs.shape
inputs = na.broadcast_to(inputs, na.broadcast_shapes(inputs.shape, shape_item_inputs))
outputs = na.broadcast_to(outputs, na.broadcast_shapes(outputs.shape, shape_item_outputs))
# the fields are selected the same way as the outputs, so they
# are broadcast against the item first
fields = _fields.apply(
array,
lambda v: na.broadcast_to(v, na.broadcast_shapes(v.shape, shape_item_outputs))[item_outputs],
axes=tuple(shape_item_outputs),
)
elif isinstance(item, dict):
# the shapes are found once, since they may need to be computed
shape_inputs = na.shape(inputs)
shape_outputs = na.shape(outputs)
# ignore axes that are not present in this function array, so that
# indexing is consistent with ``ScalarArray`` and composes with
# ``na.getitem`` over heterogeneous structures (e.g. a function
# array that lacks the axis being selected is left untouched)
item = {
ax: item[ax] for ax in item
if ax in shape_inputs or ax in shape_outputs
}
item_inputs = dict()
item_outputs = dict()
for ax in item:
item_ax = item[ax]
if isinstance(item_ax, na.AbstractFunctionArray):
item_inputs[ax] = item_ax.inputs
item_outputs[ax] = item_ax.outputs
elif _is_vertex(ax, shape_inputs, shape_outputs):
if isinstance(item_ax, slice) or np.issubdtype(type(item_ax), np.integer):
item_outputs[ax], item_inputs[ax] = _item_vertex(
item=item_ax,
num=shape_outputs[ax],
axis=ax,
)
else:
return NotImplemented
else:
#can't assume center ax is in both outputs and inputs
if ax in shape_inputs:
item_inputs[ax] = item_ax
if ax in shape_outputs:
item_outputs[ax] = item_ax
# the fields are indexed like the outputs, and ignore the axes
# which they do not have
fields = dict()
if _fields.names(array):
item_fields = {ax: item_outputs.get(ax, item[ax]) for ax in item}
fields = _fields.apply(array, lambda v: v[item_fields], axes=tuple(item_fields))
else:
return NotImplemented
return array.replace(
inputs=inputs[item_inputs],
outputs=outputs[item_outputs],
**fields,
)
def _getitem_reversed(
self,
array: na.AbstractArray,
item: Mapping[str, int | slice | na.AbstractArray] | na.AbstractFunctionArray
):
if isinstance(array, (na.AbstractScalar, na.AbstractVectorArray)):
array = na.FunctionArray(array, array)
else:
return NotImplemented
return array._getitem(item)
def __bool__(self) -> bool:
result = super().__bool__()
return result and bool(self.outputs)
def __mul__(self, other: na.ArrayLike | u.UnitBase) -> FunctionArray:
if isinstance(other, u.UnitBase):
exp = self.explicit
return exp.replace(
outputs=exp.outputs * other,
)
else:
return super().__mul__(other)
def __lshift__(self, other: na.ArrayLike | u.UnitBase) -> FunctionArray:
if isinstance(other, u.UnitBase):
exp = self.explicit
return exp.replace(
outputs=self.outputs << other,
)
else:
return super().__lshift__(other)
def __truediv__(self, other: na.ArrayLike | u.UnitBase) -> FunctionArray:
if isinstance(other, u.UnitBase):
exp = self.explicit
return exp.replace(
outputs=self.outputs / other,
)
else:
return super().__truediv__(other)
def __array_matmul__(
self,
x1: float | u.Quantity | FunctionArray,
x2: float | u.Quantity | FunctionArray,
out: None | na.FunctionArray | tuple[na.FunctionArray] = None,
**kwargs,
) -> FunctionArray:
# numpy passes the `out` argument of a ufunc as a tuple
if isinstance(out, tuple):
(out,) = out
if isinstance(x1, na.AbstractArray):
if isinstance(x1, AbstractFunctionArray):
inputs_1 = x1.inputs
outputs_1 = x1.outputs
else:
return NotImplemented
else:
inputs_1 = None
outputs_1 = x1
if isinstance(x2, na.AbstractArray):
if isinstance(x2, AbstractFunctionArray):
inputs_2 = x2.inputs
outputs_2 = x2.outputs
else:
return NotImplemented
else:
inputs_2 = None
outputs_2 = x2
if inputs_1 is not None:
inputs = inputs_1
if inputs_2 is not None:
if np.any(inputs_1 != inputs_2):
raise InputValueError("`x1.inputs` and `x2.inputs` must be equal")
else:
pass
else:
if inputs_2 is not None:
inputs = inputs_2
else:
return NotImplemented
if out is not None:
if not np.all(inputs == out.inputs):
raise InputValueError("`out.inputs` must be equal to `x1.inputs` and `x2.inputs`")
result = self.explicit.replace(
inputs=inputs,
outputs=np.matmul(
outputs_1,
outputs_2,
out=out.outputs if out is not None else out,
**kwargs,
),
)
if out is not None:
result = _fields.out(out, _fields.values(result))
return result
def __array_ufunc__(
self,
ufunc: np.ufunc,
method: str,
*inputs,
**kwargs,
) -> None | na.AbstractArray | tuple[na.AbstractArray, ...]:
result = super().__array_ufunc__(ufunc, method, *inputs, **kwargs)
if result is not NotImplemented:
return result
nout = ufunc.nout
inputs_inputs = []
inputs_outputs = []
for inp in inputs:
if isinstance(inp, na.AbstractArray):
if isinstance(inp, AbstractFunctionArray):
inputs_inputs.append(inp.inputs)
inputs_outputs.append(inp.outputs)
else:
inputs_outputs.append(inp)
else:
inputs_outputs.append(inp)
for inputs_input in inputs_inputs[1:]:
if np.any(inputs_input != inputs_inputs[0]):
raise InputValueError(f"all inputs to {ufunc} must have the same coordinates")
if "out" in kwargs:
out = kwargs.pop("out")
outputs_out = list()
for o in out:
if o is not None:
if isinstance(o, self.type_explicit):
outputs_o = o.outputs
else:
raise ValueError(
f"each element of `out` must be an instance of {self.type_explicit}, got {type(o)}"
)
else:
outputs_o = None
outputs_out.append(outputs_o)
if nout == 1:
outputs_out = outputs_out[0]
else:
outputs_out = tuple(outputs_out)
kwargs["out"] = outputs_out
else:
out = (None,) * nout
if "where" in kwargs:
where = kwargs.pop("where")
if isinstance(where, na.AbstractArray):
if isinstance(where, AbstractFunctionArray):
if np.any(where.inputs != inputs_inputs[0]): #pragma: nocover
raise InputValueError("`where.inputs` must match the rest of the inputs")
kwargs["where"] = where.outputs
else:
return NotImplemented
else:
kwargs["where"] = where
inputs_result = inputs_inputs[0]
outputs_result = getattr(ufunc, method)(*inputs_outputs, **kwargs)
if nout == 1:
outputs_result = (outputs_result,)
result = list(
self.explicit.replace(
inputs=inputs_result,
outputs=outputs_result[i],
)
for i in range(nout)
)
for i in range(nout):
if out[i] is not None:
out[i].inputs = result[i].inputs
out[i].outputs = result[i].outputs
_fields.out(out[i], _fields.values(result[i]))
result[i] = out[i]
if nout == 1:
result = result[0]
else:
result = tuple(result)
return result
def __array_function__(
self,
func: Callable,
types: Collection,
args: tuple,
kwargs: dict[str, Any],
):
result = super().__array_function__(func=func, types=types, args=args, kwargs=kwargs)
if result is not NotImplemented:
return result
handlers, handled = _array_function_dispatch()
if func in handlers:
return handlers[func](func, *args, **kwargs)
if func in handled:
return handled[func](*args, **kwargs)
return NotImplemented
def __named_array_function__(self, func, *args, **kwargs):
result = super().__named_array_function__(func, *args, **kwargs)
if result is not NotImplemented:
return result
from . import function_named_array_functions
if func in function_named_array_functions.ASARRAY_LIKE_FUNCTIONS:
return function_named_array_functions.asarray_like(func=func, *args, **kwargs)
if func in function_named_array_functions.NDFILTER_FUNCTIONS:
return function_named_array_functions.ndfilter(func, *args, **kwargs)
if func in function_named_array_functions.HANDLED_FUNCTIONS:
return function_named_array_functions.HANDLED_FUNCTIONS[func](*args, **kwargs)
return NotImplemented
def interp_linear(
self,
item: dict[str, na.AbstractArray],
) -> FunctionArray:
raise NotImplementedError
def pcolormesh(
self,
axs: np.ndarray,
input_component_x: str,
input_component_y: str,
input_component_row: None | str = None,
input_component_column: None | str = None,
index: None | dict[str, int] = None,
output_component_color: None | str = None,
**kwargs: Any,
):
"""
Plot a :class:`FunctionArray` via :func:`matplotlib.pyplot.pcolormesh`.
:func:`FunctionArray.pcolormesh` takes in an axes object, or array of axes objects, along with components to
be plotted along the x and y plot axes (:attr:`input_component_x` and :attr:`input_component_y`). Additional
components can be tiled along subplot row/column and are specified in :attr:`input_component_row` and
:attr:`input_component_column`.
.. jupyter-execute::
import named_arrays as na
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
position = na.Cartesian2dVectorLinearSpace(
start=-10,
stop=10,
axis=na.Cartesian2dVectorArray(
x='position_x',
y='position_y',
),
num=21,
) * u.m
x_width = 5 * u.m
y_width = 2 * u.m
velocity = 1 * u.m/u.s
time = na.ScalarLinearSpace(
start=0 * u.s,
stop=3 * u.s,
num=4,
axis='time'
)
intensity = np.exp(-(((position.x + velocity*time)/x_width) ** 2 + ((position.y + 2*velocity*time)/y_width)** 2))
scene = na.FunctionArray(
inputs=position,
outputs=intensity,
)
fig, axs = plt.subplots(
nrows=scene.outputs.shape['time'],
squeeze=False,
sharex=True,
subplot_kw=dict(aspect='equal'),
)
scene.pcolormesh(
axs=axs,
input_component_x='x',
input_component_y='y',
input_component_row='time',
)
"""
import astropy.visualization
if axs.ndim == 1:
if input_component_row is not None:
axs = axs[..., np.newaxis]
if input_component_column is not None:
axs = axs[np.newaxis, ...]
axs = na.ScalarArray(
ndarray=axs,
axes=('row', 'column')
)
if index is None:
index = dict()
with astropy.visualization.quantity_support():
for index_subplot in axs.ndindex():
index_final = index.copy()
if input_component_row is not None:
index_final[input_component_row] = index_subplot['row']
if input_component_column is not None:
index_final[input_component_column] = index_subplot['column']
array = self[index_final]
inp = array.inputs.cartesian_nd
inp_x = inp.components[input_component_x]
inp_y = inp.components[input_component_y]
out = array.outputs
if output_component_color is not None:
out = out.components[output_component_color]
# Broadcast the coordinates and the values to 2D arrays with
# the same axis order so that matplotlib interprets them
# consistently. Along axes where the coordinates are cell
# vertices, the values have one fewer element.
shape_inp = na.shape_broadcasted(inp_x, inp_y)
shape_out = {
axis: shape_inp[axis] - 1 if axis in array.axes_vertex else shape_inp[axis]
for axis in shape_inp
}
inp_x = na.broadcast_to(inp_x, shape_inp).ndarray
inp_y = na.broadcast_to(inp_y, shape_inp).ndarray
out = na.broadcast_to(out, shape_out).ndarray
ax = axs[index_subplot].ndarray
ax.pcolormesh(
inp_x,
inp_y,
out,
shading='auto',
**kwargs,
)
if index_subplot['row'] == axs.shape['row'] - 1:
if isinstance(inp_x, u.Quantity):
ax.set_xlabel(f'{input_component_x} ({inp_x.unit})')
else:
ax.set_xlabel(f'{input_component_x}')
else:
ax.set_xlabel(None)
if index_subplot['column'] == 0:
if isinstance(inp_y, u.Quantity):
ax.set_ylabel(f'{input_component_y} ({inp_y.unit})')
else:
ax.set_ylabel(f'{input_component_y}')
else:
ax.set_ylabel(None)
if input_component_column is not None:
if input_component_column in inp.components:
if index_subplot['row'] == 0:
inp_column = inp.components[input_component_column]
ax.text(
x=0.5,
y=1.01,
s=f'{inp_column.mean().ndarray.value:0.03f} {inp_column.unit:latex_inline}',
transform=ax.transAxes,
ha='center',
va='bottom'
)
if input_component_row is not None:
if input_component_row in inp.components:
if index_subplot['column'] == axs.shape['column'] - 1:
inp_row = inp.components[input_component_row]
ax.text(
x=1.01,
y=0.5,
s=f'{inp_row.mean().ndarray.value:0.03f} {inp_row.unit:latex_inline}',
transform=ax.transAxes,
va='center',
ha='left',
rotation=-90,
)
@dataclasses.dataclass(eq=False, repr=False)
class FunctionArray(
AbstractFunctionArray,
na.AbstractExplicitArray,
Generic[InputsT, OutputsT],
):
"""
A representation of a discrete function.
A composition of two arrays: :attr:`inputs` and :attr:`outputs`.
:attr:`inputs` represents the inputs (or independent variables) of the
function, and :attr:`outputs` represents the outputs (or dependent variables) of
the function.
A subclass can add fields of its own, such as the exposure time of each
image of an observation. Every such field whose value is a named array is
treated as sampled along the axes of the function:
* Indexing (with a :class:`dict`, :meth:`isel`, index arrays, or a boolean
mask), :func:`numpy.stack`, :func:`numpy.concatenate`,
:func:`numpy.moveaxis`, :func:`numpy.repeat`,
:func:`numpy.take_along_axis`, :func:`numpy.reshape`,
:meth:`combine_axes`, and :meth:`cell_centers` do to the field what they
do to the outputs, along the axes which the field has.
:func:`named_arrays.debroadcast` keeps an axis that the field varies
along.
* Assigning another array of the subclass with ``__setitem__``, or
copying one with :func:`numpy.copyto`, writes its fields into those of
this array, in place, like the outputs. A field which does not vary
along an axis being assigned can only be given the value it already
has, and so can a field which is not a named array, such as a string.
Since most operations pass the fields of an array on without copying
them, assigning into the result of, say, a ufunc also changes the
fields of the array it was computed from. A field without a
distribution is instead replaced by an uncertain copy of itself, like
the outputs, when the value is uncertain or the mask selects different
elements in different samples.
* An operation which removes or resamples an axis that the field varies
along raises an error: a reduction with ``keepdims=False``,
:func:`numpy.percentile` (which removes the axes by default),
:meth:`integrate`, :meth:`regrid`, or :meth:`cell_centers` with
``random=True``. So does a reduction with ``keepdims=True`` along an
axis which the inputs do not have, since it keeps only the axes of the
inputs. Whether the field should be summed, averaged, or dropped
depends on what it is, so replace it first, for example with
:func:`dataclasses.replace`.
* :meth:`cell_centers` also raises an error for a field whose elements
are not floating-point numbers, such as a mask, rather than turn them
into fractions.
* The other operations, including the reductions which keep the reduced
axes (the default for a function array), pass the field on unchanged. A
ufunc of several arrays of the subclass takes the fields of the first.
An array given as the ``out`` argument of an operation takes the fields
of the result, if its type has them.
When the arrays combined by :func:`numpy.stack` or
:func:`numpy.concatenate` have fields which are not named arrays, such as
a string or a number, those must be the same in every array. Fields which
cannot be compared, such as models of an instrument, are taken from the
array whose type is a subclass of the types of all the others. Arrays
with a field which that type does not have, such as arrays of two
unrelated subclasses, cannot be combined.
Every function array checks, when it is built, that each such field has
the same number of elements as the outputs along the axes they share, so
that an operation which changes the outputs without the field fails
instead of leaving the field stale. The check has limits:
* It allows outputs with a single element along an axis, as after a
reduction which keeps its axes, and it cannot see an axis which an
operation removed from the outputs but not from the field. Catching
those would take the shape of the inputs, which can be expensive to
compute.
* It skips outputs and fields which are implicit, such as random samples,
since finding their shape would compute them.
* A subclass which defines its own ``__post_init__`` must call
``super().__post_init__()`` for the check to run.
Fields whose values are not named arrays are otherwise passed on unchanged
by every operation.
Examples
--------
The exposure time of each of a sequence of images follows the images when
they are indexed.
.. jupyter-execute::
import dataclasses
import numpy as np
import astropy.units as u
import named_arrays as na
@dataclasses.dataclass(eq=False, repr=False)
class Images(na.FunctionArray):
timedelta: na.AbstractScalar = 0 * u.s
images = Images(
inputs=na.linspace(0, 30, axis="time", num=4) * u.s,
outputs=na.ScalarArray(np.arange(4.0), axes="time") * u.DN,
timedelta=na.ScalarArray(np.array([1, 2, 3, 4]) * u.s, axes="time"),
)
images[dict(time=slice(1, 3))].timedelta
Co-adding the images removes the axis of time, so the exposure times are
added up first.
.. jupyter-execute::
total = np.sum(
dataclasses.replace(images, timedelta=images.timedelta.sum("time")),
axis="time",
keepdims=False,
)
total.outputs / total.timedelta
"""
inputs: InputsT = 0
"""The inputs of the function."""
# The operators declared on `AbstractArray` can only promise the widest
# array type. The result of an operation is the explicit array of the
# highest family involved, so on an explicit array of this family the
# result is `Self` unless a higher family absorbs it. Declarations only;
# the implementation is inherited.
if TYPE_CHECKING: # pragma: nocover
def __add__(self, other: na.ArrayLike) -> Self: ...
def __sub__(self, other: na.ArrayLike) -> Self: ...
def __floordiv__(self, other: na.ArrayLike) -> Self: ...
def __mod__(self, other: na.ArrayLike) -> Self: ...
def __pow__(self, other: na.ArrayLike) -> Self: ...
def __radd__(self, other: na.ArrayLike) -> Self: ...
def __rsub__(self, other: na.ArrayLike) -> Self: ...
def __rmul__(self, other: na.ArrayLike) -> Self: ...
def __rtruediv__(self, other: na.ArrayLike) -> Self: ...
def __rfloordiv__(self, other: na.ArrayLike) -> Self: ...
def __rmod__(self, other: na.ArrayLike) -> Self: ...
def __rpow__(self, other: na.ArrayLike) -> Self: ...
def __lt__(self, other: na.ArrayLike) -> Self: ...
def __le__(self, other: na.ArrayLike) -> Self: ...
def __gt__(self, other: na.ArrayLike) -> Self: ...
def __ge__(self, other: na.ArrayLike) -> Self: ...
def __mul__(self, other: na.ArrayLike | u.UnitBase) -> Self: ...
def __truediv__(self, other: na.ArrayLike | u.UnitBase) -> Self: ...
def __lshift__(self, other: na.ArrayLike | u.UnitBase) -> Self: ...
def __neg__(self) -> Self: ...
def __pos__(self) -> Self: ...
def __abs__(self) -> Self: ...
outputs: OutputsT = 0
"""The outputs of the function."""
def __post_init__(self) -> None:
_fields.check(self)
@classmethod
def from_scalar_array(
cls: type[Self],
a: float | u.Quantity | na.AbstractScalarArray,
like: None | Self = None,
) -> Self:
self = super().from_scalar_array(a=a, like=like)
if like is None:
self.inputs = a
self.outputs = a
else:
if isinstance(like.inputs, na.AbstractArray):
self.inputs = like.inputs.from_scalar_array(a=a, like=like.inputs)
else:
self.inputs = a
if isinstance(like.outputs, na.AbstractArray):
self.outputs = like.outputs.from_scalar_array(a=a, like=like.outputs)
else:
self.outputs = a
return self
@property
def axes(self) -> tuple[str, ...]:
return tuple(self.inputs.shape | self.outputs.shape)
@property
def shape(self) -> dict[str, int]:
axes_center = self.axes_center
axes_vertex = self.axes_vertex
outputs_shape = self.outputs.shape
inputs_shape = self.inputs.shape
outputs_shape_center = {ax: outputs_shape[ax] for ax in outputs_shape if ax in axes_center}
inputs_shape_center = {ax: inputs_shape[ax] for ax in inputs_shape if ax in axes_center}
outputs_shape_vertex = {ax: outputs_shape[ax] for ax in outputs_shape if ax in axes_vertex}
shape = na.broadcast_shapes(
outputs_shape_center,
outputs_shape_vertex,
inputs_shape_center,
)
return shape
@property
def ndim(self) -> int:
return len(self.shape)
@property
def size(self) -> int:
return int(np.array(tuple(self.shape.values())).prod())
@property
def explicit(self) -> FunctionArray:
return self.replace(
inputs=na.explicit(self.inputs),
outputs=na.explicit(self.outputs),
)
def __setitem__(
self,
item: Mapping[str, int | slice | na.AbstractArray] | na.AbstractArray,
value: float | u.Quantity | na.FunctionArray,
):
shape_inputs = na.shape(self.inputs)
shape_outputs = na.shape(self.outputs)
def is_vertex(axis: str) -> bool:
return _is_vertex(axis, shape_inputs, shape_outputs)
# Whether the inputs can be assigned the inputs of `value`.
# Only the cells are assigned along a vertex axis, since `value`
# cannot say where their vertices are.
assign_inputs = True
if isinstance(item, na.AbstractFunctionArray):
if not np.all(item.inputs == self.inputs):
raise ValueError("boolean advanced index does not have the same inputs as the array")
item_inputs = item.outputs
item_outputs = item.outputs
# the mask selects cells, which are not the vertices
if any(is_vertex(ax) for ax in shape_inputs):
assign_inputs = False
elif isinstance(item, dict):
item_inputs = dict()
item_outputs = dict()
for ax in item:
item_ax = item[ax]
if isinstance(item_ax, na.AbstractFunctionArray):
item_inputs[ax] = item_ax.inputs
item_outputs[ax] = item_ax.outputs
elif is_vertex(ax):
# An integer keeps the axis, as it does for indexing,
# so that ``a[item] = a[item]`` leaves `a` unchanged.
if np.issubdtype(type(item_ax), np.integer):
item_ax, _ = _item_vertex(
item=item_ax,
num=shape_outputs[ax],
axis=ax,
)
item_outputs[ax] = item_ax
assign_inputs = False
else:
item_inputs[ax] = item_outputs[ax] = item_ax
else:
raise TypeError(
f"`item` must be an instance of `{dict.__name__}`, or `{na.AbstractFunctionArray.__name__}`, "
f"got `{type(item)}`"
)
if isinstance(value, na.AbstractArray):
if isinstance(value, na.AbstractFunctionArray):
value_inputs = value.inputs
value_outputs = value.outputs
# maybe this should only set to None vertex axes only
for ax in na.shape(value_inputs):
if is_vertex(ax):
value_inputs = None
else:
if value.shape:
raise ValueError(
f"if `value` is an instance of `{na.AbstractArray.__name__}`, "
f"but not an instance of `{na.AbstractFunctionArray.__name__}`, "
f"`value.shape` should be empty, got {value.shape}"
)
else:
value_inputs = None
value_outputs = value
else:
value_inputs = None
value_outputs = value
from named_arrays._scalars.uncertainties import uncertainties
# every field is checked before anything is written, so that an error
# does not leave this array half written
writes = _fields.setitem(self, item_outputs, value)
# Inputs or outputs without a distribution cannot store a selection
# which differs between samples, or an uncertain value, so they are
# replaced by an uncertain copy of themselves once both are assigned.
inputs = self.inputs
if value_inputs is not None and assign_inputs:
inputs = uncertainties._as_uncertain(inputs, item_inputs, value_inputs)
inputs[item_inputs] = value_inputs
outputs = uncertainties._as_uncertain(self.outputs, item_outputs, value_outputs)
outputs[item_outputs] = value_outputs
self.inputs = inputs
self.outputs = outputs
for write in writes:
write()
@dataclasses.dataclass(eq=False, repr=False)
class AbstractPolynomialFunctionArray(
AbstractFunctionArray,
):
@property
@abc.abstractmethod
def coefficients(self) -> na.AbstractVectorArray | na.AbstractMatrixArray:
"""
A vector or matrix representing the coefficients of the polynomial.
If this function is scalar-valued, :attr:`coefficients` should be a vector,
and if this function is vector-valued, :attr`coefficients` should be a matrix.
"""
@property
@abc.abstractmethod
def coefficient_names(self) -> Sequence[str]:
"""The names of the polynomial terms to fit (the design-matrix component keys)."""
@property
@abc.abstractmethod
def axis_polynomial(self) -> None | str | Sequence[str]:
"""The logical axes along which this polynomial is distributed."""
@property
@abc.abstractmethod
def center(self) -> None | InputsT:
"""
The reference point which is subtracted from the inputs before fitting
the polynomial.
"""
@abc.abstractmethod
def design_matrix(
self,
inputs: float | u.Quantity | na.AbstractScalar | na.AbstractVectorArray,
) -> na.AbstractVectorArray:
"""
The `design matrix <https://en.wikipedia.org/wiki/Design_matrix>`_
corresponding to the given inputs.
Note that while this is `called` a matrix, this function returns a vector
since the rows are independent observations, and this is concept is
already captured by the logical axes in the array.
Parameters
----------
inputs
the set of independent variables to convert into the design matrix
"""
def __call__(
self,
inputs: float | u.Quantity | na.ScalarArray | na.AbstractVectorArray,
) -> na.AbstractFunctionArray:
outputs = self.design_matrix(inputs) @ self.coefficients
return na.FunctionArray(inputs, outputs)
@property
def predictions(self) -> OutputsT:
"""
The outputs of the polynomial model given :attr:`inputs`.
Equivalent to ``self(self.inputs).outputs``.
"""
return self(self.inputs).outputs
@dataclasses.dataclass(eq=False, repr=False)
class PolynomialFitFunctionArray(
FunctionArray[InputsT, OutputsT],
AbstractPolynomialFunctionArray,
):
"""
A :class:`named_arrays.PolynomialFitFunctionArray` carries the independent variables, inputs, and dependent variables, outputs,
of a discrete function, and a linear least squares polynomial fit to that function.
The set of polynomial terms to fit is given explicitly by :attr:`coefficient_names`.
Use the :meth:`from_degree` classmethod to generate the standard term set for a
given polynomial degree (a scalar total-degree cap, or a per-component vector of
maximum exponents for a mixed-order fit).
Parameters
----------
inputs
The set of independent variables.
outputs
The set of dependent variables.
coefficient_names
The names of the polynomial terms to fit (the design-matrix component keys).
axis_polynomial
The logical axes along which this polynomial is distributed.
where_polynomial
A boolean mask controlling which elements to use for fitting.
Examples
--------
.. nblinkgallery::
:caption: Relevant Tutorials
:name: rst-link-gallery
../tutorials/PolynomialFunctionArray
"""
center: None | InputsT = None
"""The reference point which is subtracted from the inputs before fitting."""
coefficient_names: None | Sequence[str] = None
"""
The names of the polynomial terms to fit (the design-matrix component keys).
Each name is a ``*``-joined product of input component names (dotted for nested
vectors, e.g. ``"position.x*wavelength"``); the empty string ``""`` is the
constant term. Use :meth:`from_degree` to generate the standard set for a given
degree, optionally pruning terms before constructing the fit.
"""
axis_polynomial: None | str | Sequence[str] = None
"""The logical axes along which this polynomial is distributed."""
where_polynomial: bool | na.ScalarArray = True
"""A boolean mask controlling which elements to use for fitting."""
@functools.cached_property
def coefficients(self) -> na.AbstractVectorArray | na.AbstractMatrixArray:
d = self.design_matrix(self.inputs)
axis = self.axis_polynomial
where = self.where_polynomial
dTd = self._outer(d, d, axis=axis, where=where)
dTo = self._outer(d, self.outputs, axis=axis, where=where)
return dTd.inverse @ dTo
def design_matrix(
self,
inputs: float | u.Quantity | na.AbstractScalar | na.AbstractVectorArray,
) -> na.AbstractVectorArray:
if self.coefficient_names is None:
raise ValueError(
"`coefficient_names` is required; use "
"`PolynomialFitFunctionArray.from_degree` to generate it."
)
if self.center is not None:
inputs = inputs - self.center
inputs = na.as_named_array(inputs)
if isinstance(inputs, na.AbstractScalar):
inputs = na.CartesianNdVectorArray({"dummy": inputs})
inputs = inputs.cartesian_nd.broadcasted.components
design_matrix = {}
for name in self.coefficient_names:
term = 1
if name: # the empty string is the constant term
for k in name.split("*"):
if k not in inputs:
raise ValueError(
f"coefficient name {name!r} references input component "
f"{k!r}, which is not one of the available components "
f"{tuple(inputs)}."
)
term = term * inputs[k]
design_matrix[name] = term
design_matrix = na.CartesianNdVectorArray(design_matrix)
return design_matrix
@classmethod
def from_degree(
cls,
inputs: InputsT,
outputs: OutputsT,
degree: int,
components: None | str | Sequence[str] = None,
center: None | InputsT = None,
axis_polynomial: None | str | Sequence[str] = None,
where_polynomial: bool | na.ScalarArray = True,
) -> Self:
"""
Construct a :class:`PolynomialFitFunctionArray` whose :attr:`coefficient_names`
are the standard monomial terms for a given polynomial degree.
Parameters
----------
inputs
The set of independent variables.
outputs
The set of dependent variables.
degree
The degree of the polynomial. Every monomial with total degree up to
``degree`` (over `components`) is included.
components
The components of the input that the polynomial depends on.
Defaults to all components of `inputs`.
center
The reference point subtracted from the inputs before fitting.
axis_polynomial
The logical axes along which the polynomial is distributed.
where_polynomial
A boolean mask controlling which elements to use for fitting.
"""
if components is None:
inputs_nd = na.as_named_array(inputs)
if isinstance(inputs_nd, na.AbstractScalar):
# mirror `design_matrix`'s wrapping of scalar inputs
inputs_nd = na.CartesianNdVectorArray({"dummy": inputs_nd})
components = tuple(inputs_nd.cartesian_nd.broadcasted.components)
elif isinstance(components, str):
components = (components,)
coefficient_names = []
for i in range(degree + 1):
for combination in itertools.combinations_with_replacement(components, i):
coefficient_names.append("*".join(combination))
return cls(
inputs=inputs,
outputs=outputs,
coefficient_names=coefficient_names,
center=center,
axis_polynomial=axis_polynomial,
where_polynomial=where_polynomial,
)
@classmethod
def _outer(
cls,
v1: na.AbstractVectorArray,
v2: float | u.Quantity | na.AbstractScalar | na.AbstractVectorArray,
axis: None | str | Sequence[str] = None,
where: bool | na.AbstractScalar = True,
):
v1_T_v2_components = {}
if isinstance(v1, na.AbstractVectorArray):
v1_broadcasted = v1.broadcasted.components
if isinstance(v2, na.AbstractVectorArray):
for c1 in v1_broadcasted:
v2_broadcasted = v2.broadcasted.components
row_components = {}
for c2 in v2_broadcasted:
row_components[c2] = (
v1_broadcasted[c1] * v2_broadcasted[c2]
).sum(axis=axis, where=where)
v1_T_v2_components[c1] = v2.type_explicit.from_components(row_components)
v1_T_v2 = v1.type_matrix.from_components(v1_T_v2_components)
else:
for c1 in v1_broadcasted:
row_components = (v1_broadcasted[c1] * v2).sum(axis=axis, where=where)
v1_T_v2_components[c1] = row_components
v1_T_v2 = v1.type_explicit.from_components(v1_T_v2_components)
else:
raise NotImplementedError
return v1_T_v2