from __future__ import annotations
from typing import Mapping, TYPE_CHECKING, TypeVar, Generic, Sequence, Iterator, Union, Type, Callable, Collection, Any
from typing import Self
import abc
import dataclasses
import copy
import functools
import secrets
from types import NotImplementedType
import numpy as np
import numpy.typing as npt
import astropy.units as u
import named_arrays as na
__all__ = [
"QuantityLike",
"threshold_print",
"edgeitems_print",
"StartT",
"StopT",
"named_array_like",
"get_dtype",
"value",
"type_array",
"broadcast_shapes",
"shape_broadcasted",
"ndindex",
"indices",
"flatten_axes",
"axis_normalized",
"explicit",
"getitem",
"pack",
"unpack",
"AbstractArray",
"ArrayLike",
"AbstractExplicitArray",
"AbstractImplicitArray",
"AbstractRandomMixin",
"AbstractRangeMixin",
"AbstractSymmetricRangeMixin",
"AbstractRandomSample",
"AbstractParameterizedArray",
"AbstractLinearParameterizedArrayMixin",
"AbstractArrayRange",
"AbstractSpace",
"AbstractLinearSpace",
"strata",
"AbstractStratifiedRandomSpace",
"AbstractLogarithmicSpace",
"AbstractGeometricSpace",
"AbstractUniformRandomSample",
"AbstractNormalRandomSample",
"AbstractPoissonRandomSample",
]
QuantityLike = Union[int, float, complex, np.ndarray, u.Quantity]
threshold_print = 100
"""
The number of elements an array may have before :meth:`AbstractArray.to_string`
shows only its ends.
This is the ``threshold`` argument of :func:`numpy.array2string`, but it is
smaller than the :mod:`numpy` default of 1000, since an array in this package
usually has more axes than a bare :class:`numpy.ndarray`, and the string
representation of a vector or a function array repeats the representation of
every array it contains.
Assign to :data:`named_arrays.threshold_print` to change it, in the same
spirit as :func:`numpy.set_printoptions`::
import named_arrays as na
na.threshold_print = 1000
"""
edgeitems_print = 2
"""
The number of elements :meth:`AbstractArray.to_string` shows at each end of an
axis of an array which is too large to show in full.
This is the ``edgeitems`` argument of :func:`numpy.array2string`, but it is
smaller than the :mod:`numpy` default of 3, since the number of lines it costs
grows with the number of axes.
Assign to :data:`named_arrays.edgeitems_print` to change it.
"""
def _required() -> Any:
"""
Clear a default value inherited from a base dataclass.
Returns :obj:`dataclasses.MISSING`, which the dataclass machinery reads as
"this field has no default". Writing ``dataclasses.MISSING`` directly is
equivalent at runtime but is rejected by a type checker, since the
sentinel is not a value of the field's declared type.
"""
return dataclasses.MISSING
def _array_function_handlers(
categories: Sequence[tuple[Collection[Callable], Callable]],
) -> dict[Callable, Callable]:
"""
Map each :mod:`numpy` function in `categories` to the handler of its category.
Each family of arrays groups the :mod:`numpy` functions it supports into
categories which share a handler, so that ``__array_function__`` can find
the handler with one dictionary lookup instead of testing the categories
one after another.
A function in more than one category goes to the first of them.
Parameters
----------
categories
A sequence of ``(functions, handler)`` pairs,
where each handler is called as ``handler(func, *args, **kwargs)``.
"""
result = dict()
for functions, handler in categories:
for func in functions:
result.setdefault(func, handler)
return result
@functools.cache
def _handled_functions_core() -> dict[Callable, Callable]:
"""
The ``HANDLED_FUNCTIONS`` of :mod:`named_arrays._core_array_functions`,
imported on first use since that module imports this one.
"""
from . import _core_array_functions
return _core_array_functions.HANDLED_FUNCTIONS
AxisT = TypeVar("AxisT", bound="str | AbstractArray", covariant=True)
NumT = TypeVar("NumT", bound="int | AbstractArray", covariant=True)
StartT = TypeVar("StartT", bound="QuantityLike | AbstractArray", covariant=True)
StopT = TypeVar("StopT", bound="QuantityLike | AbstractArray", covariant=True)
CenterT = TypeVar("CenterT", bound="QuantityLike | AbstractArray", covariant=True)
WidthT = TypeVar("WidthT", bound="QuantityLike | AbstractArray", covariant=True)
StartExponentT = TypeVar("StartExponentT", bound="QuantityLike | AbstractArray", covariant=True)
StopExponentT = TypeVar("StopExponentT", bound="QuantityLike | AbstractArray", covariant=True)
BaseT = TypeVar("BaseT", bound="QuantityLike | AbstractArray", covariant=True)
PrototypeT = TypeVar("PrototypeT")
def named_array_like(a: Any) -> bool:
"""
Check if an object is compatible with the :mod:`named_arrays` API.
If the object has a ``__named_array_function__`` method it is considered compatible.
Parameters
----------
a
Object to be checked for compatibility with the :mod:`named_arrays` API.
Examples
--------
Instances of :class:`named_arrays.ScalarArray` are compatible with the :mod:`named_arrays` API
.. jupyter-execute::
import named_arrays as na
na.named_array_like(na.ScalarArray(2))
But instances of :class:`numpy.ndarray` are not compatible
.. jupyter-execute::
import numpy as np
na.named_array_like(np.empty(3))
"""
if hasattr(a, "__named_array_function__"):
return True
else:
return False
def get_dtype(
value: bool | int | float | complex | str | np.ndarray | AbstractArray,
) -> np.dtype:
"""
Get the equivalent :attr:`numpy.ndarray.dtype` of the argument.
If the argument is an instance of :class:`numpy.ndarray`, this function simply returns :attr:`numpy.ndarray.dtype`.
Otherwise, this function wraps the argument in an :func:`numpy.array()` call and then evaluates the ``dtype``.
Parameters
----------
value
Object to find the ``dtype`` of
Returns
-------
``dtype`` of the argument
"""
if isinstance(value, np.ndarray):
return value.dtype
if isinstance(value, AbstractArray):
if isinstance(value, na.AbstractScalar):
return value.dtype
else:
raise ValueError("non-scalar instances of `na.AbstractArray` may not be represented by a single dtype")
else:
return np.array(value).dtype
def value(a: float | u.Quantity | AbstractArray):
"""
Remove the units (if they exist) from the input.
Parameters
----------
a
A numeric value that may or may not have associated units
Returns
-------
The same numeric value, but guaranteed to be dimensionless.
"""
if isinstance(a, AbstractArray):
return a.value
elif isinstance(a, u.Quantity):
return a.value
else:
return a
def type_array(
*values: bool | int | float | complex | str | np.ndarray | u.Quantity | AbstractArray,
) -> Type[AbstractExplicitArray]:
cls = None
priority_max = 0
for value in values:
if isinstance(value, AbstractArray):
cls_tmp = value.type_explicit
priority_tmp = cls_tmp.__named_array_priority__
if priority_tmp > priority_max:
priority_max = priority_tmp
cls = cls_tmp
return cls
def broadcast_shapes(*shapes: dict[str, int]) -> dict[str, int]:
r"""
An analogue of :func:`numpy.broadcast_shapes` for named axes.
Whereas the :mod:`numpy` version of this function accepts a sequence
of :class:`tuple`\ s, this function accepts a sequence of :class:`dict`\ s
since each axis now has an associated name.
Parameters
----------
shapes
A sequence of shapes to broadcast against one another.
"""
if not shapes:
return dict()
result = shapes[0].copy()
for shape in shapes[1:]:
if shape == result:
continue
for axis in shape:
shape_axis = shape[axis]
if axis in result:
result_axis = result[axis]
if result_axis == shape_axis:
pass
elif shape_axis == 1:
pass
elif result_axis == 1:
result[axis] = shape_axis
else:
raise ValueError(f'shapes {shapes} are not compatible')
else:
result[axis] = shape_axis
return result
def shape_broadcasted(*arrays: Any) -> dict[str, int]:
"""
Find the broadcasted shape of a sequence of arrays.
Parameters
----------
arrays
A sequence of arrays to broadcast against one another.
"""
shapes = [a.shape for a in arrays if hasattr(a, "__named_array_function__")]
return broadcast_shapes(*shapes)
def ndindex(
shape: dict[str, int],
axis_ignored: None | str | Sequence[str] = None,
) -> Iterator[dict[str, int]]:
"""
An analogue of :class:`numpy.ndindex` for named axes.
This version adds a new argument, `axis_ignored`,
to allow some of the axes to be excluded from the returned iterator.
Parameters
----------
shape
The size of each dimension of the array to index over.
axis_ignored
The axes to ignore when iterating.
"""
shape = shape.copy()
if axis_ignored is None:
axis_ignored = []
elif isinstance(axis_ignored, str):
axis_ignored = [axis_ignored]
for axis in axis_ignored:
if axis in shape:
shape.pop(axis)
shape_tuple = tuple(shape.values())
for index in np.ndindex(*shape_tuple):
yield dict(zip(shape.keys(), index))
def indices(shape: dict[str, int]) -> dict[str, na.ScalarArrayRange]:
"""
An analogue of :func:`numpy.indices` for named axes.
Parameters
----------
shape
"""
return {axis: na.ScalarArrayRange(0, shape[axis], axis=axis) for axis in shape}
def flatten_axes(axes: Sequence[str]):
"""
Given a sequence of axes names,
combine them into a single :class:`str` using ``*`` as a delimiter.
This is intended to be used during boolean advanced indexing to
represent the axes selected by the mask.
Parameters
----------
axes
A sequence of axes names.
"""
if not axes:
raise ValueError(f"`axes` must be a non-empty sequence, got {axes}")
return '*'.join(axes)
def axis_normalized(
a: AbstractArray,
axis: None | str | Sequence[str],
) -> tuple[str]:
"""
Convert all the possible values of the `axis` argument to a :class:`tuple` of :class:`str`.
Parameters
----------
a
If `axis` is :obj:`None`, the result is :code:`a.axes`.
axis
The axis to normalize.
"""
if axis is None:
result = a.axes
elif isinstance(axis, str):
result = axis,
else:
result = tuple(axis)
return result
def explicit(value: Any | AbstractArray):
"""
Converts an array to its explicit version if possible.
Parameters
----------
value
any value or an instance of :class:`AbstractArray`
Returns
-------
If an instance of :class:`AbstractArray`, returns :attr:`AbstractArray.explicit` otherwise just returns the
input value.
"""
if isinstance(value, AbstractArray):
return value.explicit
else:
return value
def getitem(
a: Any,
item: dict | AbstractArray,
) -> Any:
"""
Recursively index an arbitrarily-nested structure with a named index.
The index ``item`` is applied only to the leaves of ``a`` that are
compatible with the :mod:`named_arrays` API (see :func:`named_array_like`).
Containers (:class:`dict`, :class:`list`, :class:`tuple`) and
:mod:`dataclasses` are traversed recursively, and any other leaf value is
returned unchanged.
Parameters
----------
a
The structure to index. May be a named array, a (possibly nested)
:class:`dict`, :class:`list`, or :class:`tuple` of named arrays, a
:mod:`dataclasses` instance whose fields are named arrays, or any other
value (which is returned unchanged).
item
The named index to apply to each array-like leaf of ``a``, for example
a :class:`dict` mapping axis names to index arrays.
Returns
-------
A structure with the same shape as ``a`` where every array-like leaf
has been indexed by ``item``.
Examples
--------
.. jupyter-execute::
import numpy as np
import named_arrays as na
a = na.arange(0, 5, axis="x")
index = {"x": na.ScalarArray(np.array([1, 3]), axes="x")}
na.getitem({"foo": a, "bar": 2 * a}, index)
Notes
-----
A :class:`tuple` leaf is rebuilt as a plain :class:`tuple`, so the subtype
of a :class:`typing.NamedTuple` is not preserved.
"""
if named_array_like(a):
return a[item]
elif isinstance(a, dict):
return {
key: getitem(a[key], item)
for key in a
}
elif isinstance(a, list):
return [getitem(a_i, item) for a_i in a]
elif isinstance(a, tuple):
return tuple(getitem(a_i, item) for a_i in a)
elif dataclasses.is_dataclass(a):
changes = {
field.name: getitem(getattr(a, field.name), item)
for field in dataclasses.fields(a)
if field.init
}
return dataclasses.replace(a, **changes)
else:
return a
def pack(a: Any, axis: str = "pack") -> na.ScalarArray:
"""
Flatten the numeric leaves of a nested structure into a 1D array.
The magnitudes of every array-like leaf of ``a`` are stripped of their
units and concatenated into a single, one-dimensional, dimensionless
:class:`named_arrays.ScalarArray` along ``axis``, in a deterministic
traversal order. Containers (:class:`dict`, :class:`list`, :class:`tuple`,
:mod:`dataclasses` instances) and composite named arrays
(e.g. :class:`named_arrays.AbstractVectorArray`) are traversed recursively;
non-numeric values (e.g. :obj:`None`, :class:`str`) are ignored.
This is the inverse of :func:`unpack`. Together they bridge the gap between
the named, united arrays used throughout this package and the flat,
dimensionless float vectors expected by optimizers such as those in
:mod:`scipy.optimize`. The units, axis names, and structure dropped here are
restored by :func:`unpack` from a prototype, so the pair round-trips:
``na.unpack(na.pack(a), a) == a``.
Parameters
----------
a
The structure to flatten. May be a named array, a (possibly nested)
:class:`dict`, :class:`list`, :class:`tuple`, or :mod:`dataclasses`
instance of named arrays, or any other value.
axis
The name of the logical axis of the flattened result.
See Also
--------
unpack : Reconstruct a structure from a flattened array and a prototype.
Examples
--------
.. jupyter-execute::
import astropy.units as u
import named_arrays as na
params = na.CartesianNdVectorArray({
"yaw": 1 * u.deg,
"roll": 2 * u.deg,
"defocus": -3 * u.mm,
})
na.pack(params)
"""
chunks = []
def visit(x: Any) -> None:
if isinstance(x, na.AbstractScalarArray):
chunks.append(np.reshape(value(x).ndarray, -1))
elif isinstance(x, na.AbstractVectorArray):
for component in x.components.values():
visit(component)
elif named_array_like(x):
raise NotImplementedError(
f"pack is not supported for {type(x).__qualname__!r} leaves."
)
elif isinstance(x, dict):
for element in x.values():
visit(element)
elif isinstance(x, (list, tuple)):
for element in x:
visit(element)
elif dataclasses.is_dataclass(x) and not isinstance(x, type):
for field in dataclasses.fields(x):
if field.init:
visit(getattr(x, field.name))
elif isinstance(x, (u.Quantity, int, float, complex, np.ndarray)):
chunks.append(np.reshape(value(x), -1))
visit(a)
flat = np.concatenate(chunks).astype(float) if chunks else np.empty(0)
return na.ScalarArray(flat, axes=axis)
def unpack(
a: na.AbstractScalarArray | npt.ArrayLike,
prototype: PrototypeT,
axis: str = "pack",
) -> PrototypeT:
"""
Reconstruct a structure from a flattened array and a prototype.
This is the inverse of :func:`pack`. The flat values in ``a`` are
distributed, in order, back onto the numeric leaves of a copy of
``prototype``, restoring each leaf's units, axis names, and type from the
prototype. Non-numeric parts of the structure are copied from ``prototype``
unchanged. Thus ``na.unpack(na.pack(a), a) == a``.
Parameters
----------
a
The flattened values, for example the result of :func:`pack` or the
``x`` returned by a :mod:`scipy.optimize` routine.
prototype
A structure with the same layout as the one that produced ``a``,
supplying the units, axes, types, and traversal order used to rebuild it.
axis
The name of the logical axis of ``a`` if it is a named array.
See Also
--------
pack : Flatten the numeric leaves of a structure into a 1D array.
Examples
--------
.. jupyter-execute::
import numpy as np
import astropy.units as u
import named_arrays as na
prototype = na.CartesianNdVectorArray({
"yaw": 0 * u.deg,
"roll": 0 * u.deg,
"defocus": 0 * u.mm,
})
na.unpack(na.ScalarArray(np.array([1.0, 2.0, -3.0]), axes="pack"), prototype)
"""
if isinstance(a, na.AbstractArray):
flat = np.reshape(value(a).ndarray, -1)
else:
flat = np.reshape(np.asarray(a), -1)
position = 0
def visit(p: Any) -> Any:
nonlocal position
if isinstance(p, na.AbstractScalarArray):
shape = na.shape(p)
num = int(np.prod(tuple(shape.values()), dtype=int))
chunk = np.reshape(flat[position : position + num], tuple(shape.values()))
position += num
_unit = na.unit(p)
chunk = chunk if _unit is None else chunk << _unit
return na.ScalarArray(chunk, axes=tuple(shape))
elif isinstance(p, na.AbstractVectorArray):
return p.type_explicit.from_cartesian_nd(
na.CartesianNdVectorArray(
{name: visit(c) for name, c in p.components.items()}
)
)
elif named_array_like(p):
raise NotImplementedError(
f"unpack is not supported for {type(p).__qualname__!r} leaves."
)
elif isinstance(p, dict):
return {key: visit(element) for key, element in p.items()}
elif isinstance(p, (list, tuple)):
return type(p)(visit(element) for element in p)
elif dataclasses.is_dataclass(p) and not isinstance(p, type):
return dataclasses.replace(
p,
**{
field.name: visit(getattr(p, field.name))
for field in dataclasses.fields(p)
if field.init
},
)
elif isinstance(p, (u.Quantity, int, float, complex, np.ndarray)):
shape = np.shape(value(p))
num = int(np.prod(shape, dtype=int))
chunk = flat[position : position + num]
position += num
_unit = na.unit(p)
if shape:
chunk = np.reshape(chunk, shape)
else:
chunk = chunk[0]
if isinstance(p, (int, float, complex)):
chunk = type(p)(chunk)
return chunk if _unit is None else chunk * _unit
else:
return p
result = visit(prototype)
if position != flat.size:
raise ValueError(
f"the size of `a` ({flat.size}) does not match the number of "
f"elements in `prototype` ({position})."
)
return result
@dataclasses.dataclass(eq=False, repr=False)
class AbstractArray(
np.lib.mixins.NDArrayOperatorsMixin,
abc.ABC,
):
"""
The ultimate parent class for all array types defined in this package.
"""
@property
@abc.abstractmethod
def __named_array_priority__(self: Self) -> float:
"""
Attribute used to decide what type of array to return in instances where there is more than one option.
Similar to :attr:`numpy.class.__array_priority__`
:return: :type:`int` describing this class's array priority
"""
@property
@abc.abstractmethod
def type_abstract(self: Self) -> Type[AbstractArray]:
"""
The :class:`named_arrays.AbstractArray` type corresponding to this array.
"""
@property
@abc.abstractmethod
def type_explicit(self: Self) -> Type[AbstractExplicitArray]:
"""
The :class:`named_arrays.AbstractExplicitArray` type corresponding to this array.
"""
@property
@abc.abstractmethod
def axes(self: Self) -> tuple[str, ...]:
"""
A :class:`tuple` of :class:`str` representing the names of each dimension of this array.
Must have the same length as the number of dimensions of this array.
"""
@property
def axes_flattened(self: Self) -> str:
"""
Combine :attr:`axes` into a single :class:`str`.
This is useful for functions like :func:`numpy.flatten` which returns an array with only one dimension.
"""
return flatten_axes(self.axes)
@property
@abc.abstractmethod
def shape(self: Self) -> dict[str, int]:
"""
The number of elements along each axis of the array.
Analogous to :attr:`numpy.ndarray.shape` but represented as a :class:`dict` where the keys
are the axis names and the values are the axis sizes.
"""
@property
@abc.abstractmethod
def ndim(self: Self) -> int:
"""
Number of dimensions of the array. Equivalent to :attr:`numpy.ndarray.ndim`.
"""
@property
@abc.abstractmethod
def size(self: Self) -> int:
"""
Total number of elements in the array. Equivalent to :attr:`numpy.ndarray.size`
"""
@property
@abc.abstractmethod
def value(self: Self) -> AbstractExplicitArray:
"""
Returns a new array with its units removed, if they exist.
The family is preserved, so a vector returns a vector and a function
array returns a function array, but the result is always explicit:
an implicit array materializes into its family's explicit type.
"""
@property
@abc.abstractmethod
def explicit(self: Self) -> AbstractExplicitArray:
"""
Converts this array to an instance of :class:`named_arrays.AbstractExplicitArray`.
"""
@property
def broadcasted(self) -> na.AbstractExplicitArray:
"""
if this array has multiple components, broadcast them against each other.
Equivalent to :code:`a.broadcast_to(a.shape)`.
"""
a = self.explicit
return a.broadcast_to(a.shape)
@abc.abstractmethod
def astype(
self: Self,
dtype: str | np.dtype | Type,
order: str = 'K',
casting='unsafe',
subok: bool = True,
copy: bool = True,
) -> Self:
"""
Copy of the array cast to a specific data type.
Equivalent to :meth:`numpy.ndarray.astype`.
"""
@abc.abstractmethod
def to(
self: Self,
unit: u.UnitBase,
equivalencies: None | list[tuple[u.Unit, u.Unit]] = [],
copy: bool = True,
) -> Self:
"""
Convert this array to a new unit.
Equivalent to :meth:`astropy.units.Quantity.to`.
Parameters
----------
unit
New unit of the returned array
equivalencies
A list of equivalence pairs to try if the units are not directly
convertible.
copy
Boolean flag controlling whether to copy the array.
"""
@abc.abstractmethod
def to_value(
self: Self,
unit: u.UnitBase,
equivalencies: None | list[tuple[u.Unit, u.Unit]] = [],
) -> Self:
"""
The numerical value of this array, possibly in a different unit.
Equivalent to :meth:`astropy.units.Quantity.to_value`.
Parameters
----------
unit
New unit of the returned array
equivalencies
A list of equivalence pairs to try if the units are not directly
convertible.
"""
@property
@abc.abstractmethod
def length(self: Self) -> na.AbstractScalar:
"""
L2-norm of this array.
"""
@property
def indices(self: Self) -> dict[str, na.ScalarArrayRange]:
"""
Compute the index of each element of this array.
See Also
--------
:func:`named_arrays.indices`: A functional version of this method.
"""
return indices(self.shape)
def ndindex(
self: Self,
axis_ignored: None | str | Sequence[str] = None,
) -> Iterator[dict[str, int]]:
"""
An iterator that yields the index of each element of this array.
Parameters
----------
axis_ignored
The of the array to ignore when generating the iterator.
See Also
--------
:func:`named_arrays.ndindex`: A functional version of this method.
"""
return ndindex(
shape=self.shape,
axis_ignored=axis_ignored,
)
@abc.abstractmethod
def add_axes(self: Self, axes: str | Sequence[str]) -> AbstractExplicitArray:
"""
Add new singleton axes to this array.
Parameters
----------
axes
Either a single axis name or a sequence of axis names add to this array.
See Also
--------
:func:`named_arrays.add_axes`: A functional version of this method.
"""
@abc.abstractmethod
def combine_axes(
self: Self,
axes: None | Sequence[str] = None,
axis_new: None | str = None,
) -> AbstractExplicitArray:
"""
Combine some of the axes of the array into a single new axis.
Parameters
----------
axes
The axes to combine into a new axis
axis_new
The name of the new axis
Returns
-------
Array with the specified axes combined
"""
@classmethod
def _lerp(
cls,
i: float | na.AbstractScalar,
a0: float | na.AbstractScalar,
a1: float | na.AbstractScalar,
) -> na.AbstractScalar:
return a0 * (1 - i) + a1 * i
def _nlerp(
self,
i: dict[str, na.AbstractScalar],
) -> na.AbstractExplicitArray:
if not i:
return self.explicit
axis = next(iter(i))
i_new = {ax: i[ax] for ax in i if ax != axis}
a0 = self[{axis: slice(None, ~0)}]
a1 = self[{axis: slice(1, None)}]
if i_new:
a0 = a0._nlerp(i_new)
a1 = a1._nlerp(i_new)
return self._lerp(i[axis], a0, a1)
def cell_centers(
self,
axis: None | str | Sequence[str] = None,
random: bool = False,
seed: None | int = None,
) -> na.AbstractExplicitArray:
"""
Convert an array from cell vertices to cell centers.
Parameters
----------
axis
The axes of the array to average over.
random
If true, select a random point within each cell instead of the
geometric center.
seed
The seed of the random sampling, which has an effect only where
``random`` is true.
If :obj:`None` (the default), the sampling differs from one call
to the next, and anything computed from it cannot be reproduced.
"""
if axis is None:
axis = self.axes
elif isinstance(axis, str):
axis = (axis, )
result = self.explicit
shape = result.shape
axis = tuple(a for a in axis if a in shape)
shape_centers = {
ax: shape[ax] - 1 if ax in axis
else shape[ax]
for ax in shape
}
if not random:
i = {
a: 0.5
for a in axis
}
else:
# Each axis is drawn from its own stream. Giving them all the same
# seed would give them all the same offsets, which would put every
# sample on the diagonal of its cell rather than inside it.
seeds = np.random.SeedSequence(seed).generate_state(len(axis))
i = {
a: na.random.uniform(
0,
1,
shape_random=shape_centers,
seed=int(s),
)
for a, s in zip(axis, seeds)
}
return self._nlerp(i)
def volume_cell(self, axis: None | str | Sequence[str]) -> na.AbstractScalar:
"""
Computes the n-dimensional volume of each cell formed by interpreting
this array as a logically-rectangular grid of vertices.
Note that this method is usually only used for sorted arrays.
If `self` is a scalar, this method computes the length of each edge,
and is equivalent to :func:`numpy.diff`.
If `self` is a 2d vector, this method computes the area of each
quadrilateral, and if `self` is a 3d vector, this method computes
the volume of each cuboid.
Parameters
----------
axis
The axis or axes defining the logically-rectangular grid.
If `self` is a physical scalar, there should only be one axis.
If `self` is a physical vector, there should be one axis for each
component of the vector.
"""
raise NotImplementedError
@property
def _attrs_print(self) -> dict[str, Any]:
"""
The name and value of each attribute shown by :meth:`to_string`.
These are the fields of the dataclass, which is what defines the array,
except that a field named ``axes`` is shown as :attr:`shape`.
The shape names the axes just as ``axes`` does, and gives the length of
each one as well, so it says strictly more in about the same space.
"""
result = dict()
for f in dataclasses.fields(self):
if f.name == "axes":
result["shape"] = self.shape
else:
result[f.name] = getattr(self, f.name)
return result
def to_string(
self,
prefix: None | str = None,
multiline: None | bool = None,
):
"""
Convert this array instance to a string representation.
Parameters
----------
prefix
the length of this string is used to align the output
multiline
flag which controls if the output should be spread over multiple
lines.
Returns
-------
array represented as a :class:`str`
"""
attrs = self._attrs_print
if multiline is None:
multiline_normalized = any(isinstance(v, (np.ndarray, na.AbstractArray)) for v in attrs.values())
else:
multiline_normalized = multiline
if multiline_normalized:
delim_field = "\n"
pre = " " * len(prefix) if prefix is not None else ""
tab = " " * 4
else:
delim_field = " "
pre = tab = ""
result = f"{self.__class__.__qualname__}("
if multiline_normalized:
result += "\n"
for i, name in enumerate(attrs):
field_str = f"{pre}{tab}{name}="
val = attrs[name]
if isinstance(val, AbstractArray):
val_str = val.to_string(prefix=f"{pre}{tab}", multiline=multiline)
elif isinstance(val, np.ndarray):
val_str = np.array2string(
a=val,
max_line_width=120,
separator=", ",
prefix=field_str,
threshold=na.threshold_print,
edgeitems=na.edgeitems_print,
)
if isinstance(val, u.Quantity):
val_str = f"{val_str} {val.unit}"
else:
val_str = repr(val)
field_str += val_str
if multiline_normalized or i < (len(attrs) - 1):
field_str += f",{delim_field}"
result += field_str
result += f"{pre})"
return result
def __repr__(self):
return self.to_string()
@abc.abstractmethod
def to_string_array(
self,
format_value: str = "%.2f",
format_unit: str = "latex_inline",
pad_unit: str = r"$\,$",
) -> AbstractExplicitArray:
"""
Convert to an array of strings where each string has an
appropriately-formatted unit attached to the value.
Parameters
----------
format_value
The string used to format the numeric value of each element.
format_unit
The string used to format the units of each element.
pad_unit
The string used to add horizontal space between the value and unit.
"""
def copy_shallow(self: Self) -> Self:
"""Create a shallow copy of this array."""
return copy.copy(self)
def copy(self: Self) -> Self:
"""Create a deep copy of this array."""
return copy.deepcopy(self)
def __copy__(self: Self) -> Self:
fields = {field.name: getattr(self, field.name) for field in dataclasses.fields(self)}
return type(self)(**fields)
def __deepcopy__(self: Self, memodict={}) -> Self:
fields = {field.name: copy.deepcopy(getattr(self, field.name)) for field in dataclasses.fields(self)}
return type(self)(**fields)
def replace(self, /, **changes) -> Self:
"""
A method version of :func:`dataclasses.replace` for named arrays.
Parameters
----------
changes
The fields of the dataclass to be overwritten
"""
return dataclasses.replace(
self,
**changes,
)
@abc.abstractmethod
def _getitem(
self: Self,
item: Mapping[str, int | slice | AbstractArray] | AbstractArray,
) -> AbstractExplicitArray | NotImplementedType:
pass
@abc.abstractmethod
def _getitem_reversed(
self: Self,
array: AbstractArray,
item: Mapping[str, int | slice | AbstractArray] | AbstractArray
) -> AbstractExplicitArray | NotImplementedType:
pass
def __getitem__(
self: Self,
item: Mapping[str, int | slice | AbstractArray] | AbstractArray,
) -> AbstractExplicitArray:
result = self._getitem(item)
if result is not NotImplemented:
return result
else:
if isinstance(item, dict):
for ax in item:
if isinstance(item[ax], AbstractArray):
result = item[ax]._getitem_reversed(self, item)
if result is not NotImplemented:
return result
elif isinstance(item, AbstractArray):
result = item._getitem_reversed(self, item)
if result is not NotImplemented:
return result
raise ValueError(f"item not supported by array with type {type(self)}")
def isel(
self: Self,
**item: int | slice | AbstractArray,
) -> AbstractExplicitArray:
"""
Index this array along named axes given as keyword arguments.
This is a convenience wrapper around :meth:`__getitem__`:
``a.isel(x=0)`` is equivalent to ``a[dict(x=0)]``.
Since keyword-argument names must be valid Python identifiers, axes
whose names are not valid identifiers can only be indexed using the
:class:`dict` form, ``a[{...}]``.
Parameters
----------
item
The index to apply along each named axis.
"""
return self[item]
@abc.abstractmethod
def __bool__(self: Self) -> bool:
return True
@abc.abstractmethod
def __mul__(self: Self, other: ArrayLike | u.Unit) -> AbstractExplicitArray:
return super().__mul__(other)
@abc.abstractmethod
def __lshift__(self: Self, other: ArrayLike | u.UnitBase) -> AbstractExplicitArray:
return super().__lshift__(other)
@abc.abstractmethod
def __truediv__(self: Self, other: ArrayLike | u.UnitBase) -> AbstractExplicitArray:
return super().__truediv__(other)
# The remaining operators are inherited from
# `numpy.lib.mixins.NDArrayOperatorsMixin`, which declares each of them as
# returning `Any`. That is not merely imprecise: `Any` silences every
# subsequent check, so a single arithmetic operation disables type checking
# for everything downstream of it. Declare them here with an honest return
# type. These are declarations only: at runtime the mixin's implementation
# is inherited unchanged. Subclasses are free to narrow them further, as
# the vectors and function arrays do.
if TYPE_CHECKING: # pragma: nocover
def __add__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __sub__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __matmul__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __floordiv__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __mod__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __pow__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rshift__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __and__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __xor__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __or__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __radd__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rsub__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rmul__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rmatmul__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rtruediv__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rfloordiv__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rmod__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rpow__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rlshift__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rrshift__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rand__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __rxor__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ror__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __iadd__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __isub__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __imul__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __imatmul__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __itruediv__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ifloordiv__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __imod__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ipow__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ilshift__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __irshift__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __iand__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ixor__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ior__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __lt__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __le__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __gt__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __ge__(self, other: ArrayLike) -> AbstractExplicitArray: ...
def __neg__(self) -> AbstractExplicitArray: ...
def __pos__(self) -> AbstractExplicitArray: ...
def __abs__(self) -> AbstractExplicitArray: ...
def __invert__(self) -> AbstractExplicitArray: ...
def __divmod__(self, other: ArrayLike) -> tuple[AbstractExplicitArray, AbstractExplicitArray]: ...
def __rdivmod__(self, other: ArrayLike) -> tuple[AbstractExplicitArray, AbstractExplicitArray]: ...
@abc.abstractmethod
def __array_matmul__(
self: Self,
x1: ArrayLike,
x2: ArrayLike,
out: None | AbstractExplicitArray = None,
**kwargs,
) -> AbstractExplicitArray:
"""
Method to handle the behavior of :func:`numpy.matmul` which has different behavior than the other ufuncs.
"""
return NotImplemented
@abc.abstractmethod
def __array_ufunc__(
self,
ufunc: np.ufunc,
method: str,
*inputs,
**kwargs,
) -> None | AbstractArray | tuple[AbstractArray, ...]:
"""
Method to override the behavior of numpy's ufuncs.
"""
if ufunc is np.matmul:
return self.__array_matmul__(*inputs, **kwargs)
else:
return NotImplemented
@abc.abstractmethod
def __array_function__(
self: Self,
func: Callable,
types: Collection,
args: tuple,
kwargs: dict[str, Any],
):
"""
Method to override the behavior of numpy's array functions.
"""
handled = _handled_functions_core()
if func in handled:
return handled[func](*args, **kwargs)
return NotImplemented
def __named_array_function__(self, func, *args, **kwargs):
"""
Method used to dispatch custom named array functions
"""
return NotImplemented
def broadcast_to(
self: Self,
shape: dict[str, int],
append: bool = False,
) -> Self:
"""
A new view of this array with the specified shape.
Parameters
----------
shape
The shape of the new array.
append
If :obj:`True`, `shape` will be appended to the current shape
of this array before broadcasting.
See Also
--------
:func:`named_arrays.broadcast_to`: A functional version of this method.
"""
return na.broadcast_to(
array=self,
shape=shape,
append=append,
)
def debroadcast(
self: Self,
axes: None | str | Sequence[str] = None,
) -> Self:
"""
Remove redundant axes introduced by broadcasting.
Parameters
----------
axes
The axes to consider removing.
If :obj:`None` (the default), every axis of this array is
considered.
Axes not present in this array are ignored.
See Also
--------
:func:`named_arrays.debroadcast`: A functional version of this method.
"""
return na.debroadcast(
array=self,
axes=axes,
)
def reshape(
self: Self,
shape: dict[str, int],
) -> Self:
"""
Reorganize this array into a new shape.
Parameters
----------
shape
The new shape of the array, must be compatible with this array.
"""
return np.reshape(self, shape)
def min(
self: Self,
axis: None | str | Sequence[str] = None,
initial: npt.ArrayLike = np._NoValue,
where: Self = np._NoValue,
) -> Self:
"""
The minimum value of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
initial
The initial value of the minimum, required if `where` provided.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.min`: A functional version of this method.
"""
return np.min(self, axis=axis, initial=initial, where=where)
def max(
self: Self,
axis: None | str | Sequence[str] = None,
initial: npt.ArrayLike = np._NoValue,
where: Self = np._NoValue,
) -> Self:
"""
The maximum value of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
initial
The initial value of the minimum, required if `where` provided.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.max`: A functional version of this method.
"""
return np.max(self, axis=axis, initial=initial, where=where)
def sum(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
The sum of each element of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.sum`: A functional version of this method.
"""
return np.sum(self, axis=axis, where=where)
def ptp(
self: Self,
axis: None | str | Sequence[str] = None,
) -> Self:
"""
The peak-to-peak value of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
See Also
--------
:func:`numpy.ptp`: A functional version of this method.
"""
return np.ptp(self, axis=axis)
def mean(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
The mean value of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.mean`: A functional version of this method.
"""
return np.mean(self, axis=axis, where=where)
def std(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
The standard deviation of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.std`: A functional version of this method.
"""
return np.std(self, axis=axis, where=where)
def var(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
The variance of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.var`: A functional version of this method.
"""
return np.var(self, axis=axis, where=where)
def median(
self,
axis: None | str | Sequence[str] = None,
):
"""
The median value of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
See Also
--------
:func:`numpy.median`: A functional version of this method.
"""
return np.median(self, axis=axis)
def percentile(
self: Self,
q: int | float | u.Quantity | Self,
axis: None | str | Sequence[str] = None,
out: None | Self = None,
overwrite_input: bool = False,
method: str = 'linear',
keepdims: bool = False,
):
"""
The requested percentile of this array along the given axes.
Parameters
----------
q
The percentile to compute.
axis
The logical axis or axes along which the operation is computed.
out
An optional output array in which to place the result.
overwrite_input
Whether to overwrite the input array.
method
How to interpolate the result.
keepdims
A boolean flag indicating whether to keep the reduced dimensions.
See Also
--------
:func:`numpy.percentile`: A functional version of this method.
"""
return np.percentile(
a=self,
q=q,
axis=axis,
out=out,
overwrite_input=overwrite_input,
method=method,
keepdims=keepdims,
)
def all(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
Return :obj:`True` if `all` of the elements along the given axes are :obj:`True`.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.all`: A functional version of this method.
"""
return np.all(self, axis=axis, where=where)
def any(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
Return :obj:`True` if `any` of the elements along the given axes are :obj:`True`.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
See Also
--------
:func:`numpy.any`: A functional version of this method.
"""
return np.any(self, axis=axis, where=where)
def take_along_axis(
self: Self,
indices: AbstractArray,
axis: str,
) -> Self:
"""
Take values from this array by matching ``indices`` along ``axis``.
Parameters
----------
indices
The integer indices to take along ``axis``.
axis
The axis of this array along which the values are taken.
See Also
--------
:func:`named_arrays.take_along_axis`: A functional version of this method.
"""
return na.take_along_axis(self, indices=indices, axis=axis)
def rms(
self: Self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
) -> Self:
"""
The root-mean-square of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
"""
return np.sqrt(np.mean(np.square(self), axis=axis, where=where))
def vmr(
self,
axis: None | str | Sequence[str] = None,
where: Self = np._NoValue,
):
"""
The variance-to-mean ratio of this array along the given axes.
Parameters
----------
axis
The logical axis or axes along which the operation is computed.
where
An optional mask which selects which elements to be considered.
"""
return na.vmr(self, axis=axis, where=where)
def transpose(
self: Self,
axes: None | Sequence[str] = None,
) -> Self:
"""
Reorder the axes of this array to the given sequence.
Parameters
----------
axes
The new axis ordering of this array.
See Also
--------
:func:`numpy.transpose`: The :mod:`numpy` version of this method.
"""
return np.transpose(self, axes=axes)
def _interp_linear_recursive(
self: Self,
item: dict[str, Self],
item_base: dict[str, Self],
):
item = item.copy()
if not item:
raise ValueError('Item must contain at least one key')
axis = next(iter(item))
x = item.pop(axis)
if x.shape:
where_below = x < 0
where_above = (self.shape[axis] - 1) <= x
x0 = np.floor(x).astype(int)
x0[where_below] = 0
x0[where_above] = self.shape[axis] - 2
else:
if x < 0:
x0 = 0
elif x >= self.shape[axis] - 1:
x0 = self.shape[axis] - 2
else:
x0 = int(x)
x1 = x0 + 1
item_base_0 = {**item_base, axis: x0}
item_base_1 = {**item_base, axis: x1}
if item:
y0 = self._interp_linear_recursive(item=item, item_base=item_base_0, )
y1 = self._interp_linear_recursive(item=item, item_base=item_base_1, )
else:
y0 = self[item_base_0].astype(float)
y1 = self[item_base_1].astype(float)
result = y0 + (x - x0) * (y1 - y0)
return result
def interp_linear(
self: Self,
item: dict[str, Self],
) -> Self:
"""
Linearly interpolate this array to find its value at the given
fractional index.
Parameters
----------
item
A fractional index at which to evaluate the array.
"""
if item:
return self._interp_linear_recursive(
item=item,
item_base=self[{ax: 0 for ax in item}].indices,
)
else:
return self
def __call__(self: Self, item: dict[str, Self]) -> Self:
return self.interp_linear(item=item)
def index_secant(
self: Self,
value: Self,
axis: None | str | Sequence[str] = None,
) -> dict[str, Self]:
import named_arrays.scalars
import named_arrays.vectors
import named_arrays.optimization
if axis is None:
axis = list(self.shape.keys())
elif isinstance(axis, str):
axis = [axis, ]
shape = self.shape
shape_nearest = named_arrays.vectors.CartesianND({ax: shape[ax] for ax in axis})
if isinstance(self, named_arrays.vectors.VectorInterface):
coordinates = self.coordinates
coordinates = {comp: None if value.coordinates[comp] is None else coordinates[comp] for comp in coordinates}
self_subspace = type(self).from_coordinates(coordinates)
else:
self_subspace = self
def indices_factory(index: named_arrays.vectors.CartesianND) -> dict[str, named_arrays.scalars.Scalar]:
return index.coordinates
def get_index(index: named_arrays.vectors.CartesianND) -> named_arrays.vectors.CartesianND:
index = indices_factory(index)
print(self_subspace)
value_new = self_subspace(index)
diff = value_new - value
if isinstance(diff, named_arrays.vectors.AbstractVector):
diff = named_arrays.vectors.CartesianND({c: diff.coordinates[c] for c in diff.coordinates if diff.coordinates[c] is not None})
return diff
result = named_arrays.optimization.root_finding.secant(
func=get_index,
root_guess=shape_nearest // 2,
step_size=named_arrays.vectors.CartesianND({ax: 1e-6 for ax in axis}),
)
return indices_factory(result)
def index(
self: Self,
value: Self,
axis: None | str | Sequence[str] = None,
) -> dict[str, Self]:
return self.index_secant(value=value, axis=axis)
ArrayLike = Union[QuantityLike, AbstractArray]
@dataclasses.dataclass(eq=False, repr=False)
class AbstractExplicitArray(
AbstractArray,
):
"""
An interface describing an explicit array or an actual array of numbers.
This is in contrast to a :class:`named_arrays.AbstractImplicitArray`.
"""
@classmethod
@abc.abstractmethod
def from_scalar_array(
cls: type[Self],
a: float | u.Quantity | na.AbstractScalarArray,
like: None | Self = None,
) -> Self:
"""
Constructs a new version of this array using ``a`` as the underlying data.
Parameters
----------
a
Anything that can be coerced into an instance of :class:`named_arrays.AbstractScalarArray`.
like
Optional reference object.
If provided, the result will be defined by this object.
"""
if like is None:
return cls()
else:
if isinstance(like, cls):
return type(like)()
else:
raise TypeError(
f"If `like` is not `None`, it must be an instance of `{cls.__name__}`, "
f"got `{type(like).__name__}`"
)
@abc.abstractmethod
def __setitem__(
self,
item,
value,
) -> None:
pass
@dataclasses.dataclass(eq=False, repr=False)
class AbstractImplicitArray(
AbstractArray,
):
"""
An interface describing an implicit array or a lazily-evaluated array.
"""
@property
def axes(self: Self) -> tuple[str, ...]:
return self.explicit.axes
@property
def ndim(self: Self) -> int:
return self.explicit.ndim
@property
def size(self: Self) -> int:
return self.explicit.size
@property
def value(self: Self) -> AbstractExplicitArray:
return self.explicit.value
@property
def shape(self: Self) -> dict[str, int]:
return self.explicit.shape
@abc.abstractmethod
def _attr_normalized(self, name: str) -> AbstractExplicitArray:
"""
Similar to :func:`getattr`, but normalizes it to an instance of :class:`AbstractExplicitArray`
"""
@dataclasses.dataclass(eq=False, repr=False)
class AbstractRandomMixin(
abc.ABC,
):
"""A mixin class used for random implicit arrays."""
def __post_init__(self):
if self.seed is None:
self.seed = secrets.randbits(128)
@property
@abc.abstractmethod
def seed(self: Self) -> int:
"""
Seed for the random number generator instance
"""
@seed.setter
@abc.abstractmethod
def seed(self: Self, value: int) -> None:
pass
@property
def _rng(self: Self) -> np.random.Generator:
return np.random.default_rng(seed=self.seed)
@dataclasses.dataclass(eq=False, repr=False)
class AbstractRangeMixin(
abc.ABC,
):
"""A mixin class for implicit arrays which describe a range of values."""
@property
@abc.abstractmethod
def start(self: Self) -> int | AbstractArray:
"""
Starting value of the range.
"""
@property
@abc.abstractmethod
def stop(self: Self) -> int | AbstractArray:
"""
Ending value of the range.
"""
@property
def range(self: Self) -> AbstractArray:
return self.stop - self.start
@dataclasses.dataclass(eq=False, repr=False)
class AbstractSymmetricRangeMixin(
AbstractRangeMixin,
):
@property
@abc.abstractmethod
def center(self: Self) -> ArrayLike:
"""
Center value of the range.
"""
@property
@abc.abstractmethod
def width(self: Self) -> ArrayLike:
"""
Width of the range.
"""
@property
def start(self: Self) -> ArrayLike:
return self.center - self.width
@property
def stop(self: Self) -> ArrayLike:
return self.center + self.width
@dataclasses.dataclass(eq=False, repr=False)
class AbstractRandomSample(
AbstractRandomMixin,
AbstractImplicitArray,
):
@property
@abc.abstractmethod
def shape_random(self: Self) -> None | dict[str, int]:
"""
Dimensions along which the resulting random sample is completely uncorrelated.
"""
@dataclasses.dataclass(eq=False, repr=False)
class AbstractUniformRandomSample(
AbstractRangeMixin,
AbstractRandomSample,
Generic[StartT, StopT],
):
start: StartT = _required()
stop: StopT = _required()
shape_random: None | dict[str, int] = None
seed: None | int = None
@property
def explicit(self) -> AbstractExplicitArray:
start = self._attr_normalized("start")
stop = self._attr_normalized("stop")
return na.random.uniform(
low=start,
high=stop,
shape_random=self.shape_random,
seed=self.seed,
)
@dataclasses.dataclass
class AbstractNormalRandomSample(
AbstractSymmetricRangeMixin,
AbstractRandomSample,
Generic[CenterT, WidthT],
):
center: CenterT = _required()
width: WidthT = _required()
shape_random: None | dict[str, int] = None
seed: None | int = None
@property
def explicit(self) -> AbstractExplicitArray:
center = self._attr_normalized("center")
width = self._attr_normalized("width")
return na.random.normal(
loc=center,
scale=width,
shape_random=self.shape_random,
seed=self.seed,
)
@dataclasses.dataclass
class AbstractPoissonRandomSample(
AbstractRandomSample,
Generic[CenterT],
):
center: CenterT = _required()
shape_random: None | dict[str, int] = None
seed: None | int = None
@property
def explicit(self) -> AbstractExplicitArray:
center = self._attr_normalized("center")
return na.random.poisson(
lam=center,
shape_random=self.shape_random,
seed=self.seed,
)
@dataclasses.dataclass(eq=False, repr=False)
class AbstractParameterizedArray(
AbstractImplicitArray,
):
@property
@abc.abstractmethod
def axis(self: Self) -> str | AbstractArray:
"""
The axis along which the array is parameterized
"""
@property
@abc.abstractmethod
def num(self: Self) -> int | AbstractArray:
"""
Number of elements in the parameterization
"""
@dataclasses.dataclass(eq=False, repr=False)
class AbstractLinearParameterizedArrayMixin(
abc.ABC
):
@property
@abc.abstractmethod
def step(self: Self) -> int | AbstractArray:
"""
Spacing between the values.
"""
@dataclasses.dataclass(eq=False, repr=False)
class AbstractArrayRange(
AbstractLinearParameterizedArrayMixin,
AbstractRangeMixin,
AbstractParameterizedArray,
Generic[StartT, StopT]
):
start: StartT = _required()
stop: StopT = _required()
axis: str | na.AbstractArray= _required()
step: int | float | na.AbstractArray = 1
@property
def explicit(self: Self) -> AbstractExplicitArray:
start = self._attr_normalized("start")
stop = self._attr_normalized("stop")
return na.arange(
start=start,
stop=stop,
axis=self.axis,
step=self.step,
)
@property
def num(self: Self) -> int:
return np.ceil((self.stop - self.start) / self.step).astype(int)
@dataclasses.dataclass(eq=False, repr=False)
class AbstractSpace(
AbstractParameterizedArray,
):
@property
@abc.abstractmethod
def endpoint(self: Self) -> bool:
"""
If ``True``, :attr:`stop` is the last sample, otherwise it is not included.
"""
@dataclasses.dataclass(eq=False, repr=False)
class AbstractLinearSpace(
AbstractLinearParameterizedArrayMixin,
AbstractRangeMixin,
AbstractSpace,
Generic[StartT, StopT, AxisT, NumT],
):
start: StartT = _required()
stop: StopT = _required()
axis: AxisT = _required()
num: NumT = 11
endpoint: bool = True
centers: bool = False
@property
def explicit(self: Self) -> AbstractExplicitArray:
return na.linspace(
start=self._attr_normalized("start"),
stop=self._attr_normalized("stop"),
axis=self.axis,
num=self.num,
endpoint=self.endpoint,
centers=self.centers,
)
@property
def step(self: Self) -> AbstractArray:
return na.step(
start=self.start,
stop=self.stop,
num=self.num,
endpoint=self.endpoint,
centers=self.centers,
)
def strata(a: AbstractArray) -> AbstractArray:
"""
If `a` is an instance of :class:`AbstractStratifiedRandomSpace`,
return :code:`a.strata`, otherwise return `a`.
Parameters
----------
a
An array to isolate the strata of.
Examples
--------
Make a scatterplot of a 2D stratified random array and the centers of
the strata.
.. jupyter-execute::
import matplotlib.pyplot as plt
import named_arrays as na
# Define a 2D stratified random array.
a = na.Cartesian2dVectorStratifiedRandomSpace(
start=-1,
stop=1,
axis=na.Cartesian2dVectorArray("x", "y"),
num=11,
)
# Isolate the strata of the array
strata = na.strata(a)
# Plot the array and the strata
fig, ax = plt.subplots()
na.plt.scatter(
a.x,
a.y,
ax=ax,
)
na.plt.scatter(
strata.x,
strata.y,
ax=ax,
);
"""
if isinstance(a, AbstractStratifiedRandomSpace):
return a.strata
else:
return a
@dataclasses.dataclass(eq=False, repr=False)
class AbstractStratifiedRandomSpace(
AbstractRandomMixin,
AbstractLinearSpace[StartT, StopT, AxisT, NumT],
):
centers: bool = True
seed: None | int = None
@property
def explicit(self: Self) -> AbstractExplicitArray:
result = self.strata
step_size = self.step
delta = na.random.uniform(
low=-step_size / 2,
high=step_size / 2,
shape_random=result.shape,
seed=self.seed,
)
return result + delta
@property
def strata(self: Self) -> AbstractExplicitArray:
return na.linspace(
start=self._attr_normalized("start"),
stop=self._attr_normalized("stop"),
num=self.num,
endpoint=self.endpoint,
axis=self.axis,
centers=self.centers,
)
@dataclasses.dataclass(eq=False, repr=False)
class AbstractLogarithmicSpace(
AbstractRangeMixin,
AbstractSpace,
Generic[StartExponentT, StopExponentT, BaseT, AxisT, NumT]
):
start_exponent: StartExponentT = _required()
stop_exponent: StopExponentT = _required()
base: BaseT = _required()
axis: AxisT = _required()
num: NumT = 11
endpoint: bool = True
@property
def explicit(self: Self) -> AbstractExplicitArray:
return na.logspace(
start=self._attr_normalized("start_exponent"),
stop=self._attr_normalized("stop_exponent"),
axis=self.axis,
num=self.num,
endpoint=self.endpoint,
base=self.base,
)
@property
def start(self: Self) -> ArrayLike:
return self.base ** self.start_exponent
@property
def stop(self: Self) -> ArrayLike:
return self.base ** self.stop_exponent
@dataclasses.dataclass(eq=False, repr=False)
class AbstractGeometricSpace(
AbstractRangeMixin,
AbstractSpace,
Generic[StartT, StopT, AxisT, NumT]
):
start: StartT = _required()
stop: StopT = _required()
axis: AxisT = _required()
num: NumT = 11
endpoint: bool = True
@property
def explicit(self: Self) -> AbstractExplicitArray:
return na.geomspace(
start=self._attr_normalized("start"),
stop=self._attr_normalized("stop"),
axis=self.axis,
num=self.num,
endpoint=self.endpoint,
)