# Standard library
import operator
# Third-party
import astropy.coordinates as coord
import astropy.units as u
import numpy as np
from gala._compat_utils import COPY_IF_NEEDED
__all__ = ["NDCartesianRepresentation", "NDCartesianDifferential"]
def _make_getter(component):
"""Make an attribute getter for use in a property.
Removed from Astropy in v6.3 but still used here.
Parameters
----------
component : str
The name of the component that should be accessed. This assumes the
actual value is stored in an attribute of that name prefixed by '_'.
"""
# This has to be done in a function to ensure the reference to component
# is not lost/redirected.
component = "_" + component
def get_component(self):
return getattr(self, component)
return get_component
class NDMixin(object):
def _apply(self, method, *args, **kwargs):
"""Create a new representation with ``method`` applied to the arrays.
In typical usage, the method is any of the shape-changing methods for
`~numpy.ndarray` (``reshape``, ``swapaxes``, etc.), as well as those
picking particular elements (``__getitem__``, ``take``, etc.), which
are all defined in `~astropy.utils.misc.ShapedLikeNDArray`. It will be
applied to the underlying arrays (e.g., ``x``, ``y``, and ``z`` for
`~astropy.coordinates.CartesianRepresentation`), with the results used
to create a new instance.
Internally, it is also used to apply functions to the components
(in particular, `~numpy.broadcast_to`).
Parameters
----------
method : str or callable
If str, it is the name of a method that is applied to the internal
``components``. If callable, the function is applied.
args : tuple
Any positional arguments for ``method``.
kwargs : dict
Any keyword arguments for ``method``.
"""
if callable(method):
apply_method = lambda array: method(array, *args, **kwargs) # noqa
else:
apply_method = operator.methodcaller(method, *args, **kwargs)
return self.__class__(
[apply_method(getattr(self, component)) for component in self.components],
copy=COPY_IF_NEEDED,
)
[docs]
class NDCartesianRepresentation(NDMixin, coord.CartesianRepresentation):
"""
Representation of points in ND cartesian coordinates.
Parameters
----------
x : `~astropy.units.Quantity` or array
The Cartesian coordinates of the point(s). If not quantity,
``unit`` should be set.
differentials : dict, `NDCartesianDifferential` (optional)
Any differential classes that should be associated with this
representation.
unit : `~astropy.units.Unit` or str
If given, the coordinates will be converted to this unit (or taken to
be in this unit if not given.
copy : bool, optional
If `True` (default), arrays will be copied rather than referenced.
"""
attr_classes = dict()
def __init__(self, x, differentials=None, unit=None, copy=True):
if unit is None:
if not hasattr(x[0], "unit"):
unit = u.one
else:
unit = x[0].unit
x = u.Quantity(x, unit, copy=copy, subok=True)
copy = False
self.attr_classes = dict(
[("x" + str(i), u.Quantity) for i in range(1, len(x) + 1)]
)
super(coord.CartesianRepresentation, self).__init__(
*x, differentials=differentials, copy=copy
)
ptype = None
for name, _ in self.attr_classes.items():
if ptype is None:
ptype = getattr(self, "_" + name).unit.physical_type
else:
if getattr(self, "_" + name).unit.physical_type != ptype:
raise u.UnitsError(
"All components should have matching " "physical types"
)
cls = self.__class__
if not hasattr(cls, name):
setattr(
cls,
name,
property(
_make_getter(name),
doc=(f"The '{name}' component of the points(s)."),
),
)
[docs]
def get_xyz(self, xyz_axis=0):
"""Return a vector array of the x, y, and z coordinates.
Parameters
----------
xyz_axis : int, optional
The axis in the final array along which the x, y, z components
should be stored (default: 0).
Returns
-------
xs : `~astropy.units.Quantity`
With dimension 3 along ``xyz_axis``.
"""
# Add new axis in x, y, z so one can concatenate them around it.
# NOTE: just use np.stack once our minimum numpy version is 1.10.
result_ndim = self.ndim + 1
if not -result_ndim <= xyz_axis < result_ndim:
raise IndexError(
"xyz_axis {0} out of bounds [-{1}, {1})".format(xyz_axis, result_ndim)
)
if xyz_axis < 0:
xyz_axis += result_ndim
# Get components to the same units (very fast for identical units)
# since np.concatenate cannot deal with quantity.
unit = self._x1.unit
sh = self.shape
sh = sh[:xyz_axis] + (1,) + sh[xyz_axis:]
components = [
getattr(self, "_" + name).reshape(sh).to(unit).value
for name in self.attr_classes
]
xs_value = np.concatenate(components, axis=xyz_axis)
return u.Quantity(xs_value, unit=unit, copy=COPY_IF_NEEDED)
xyz = property(get_xyz)
[docs]
class NDCartesianDifferential(NDMixin, coord.CartesianDifferential):
"""Differentials in of points in ND cartesian coordinates.
Parameters
----------
*d_x : `~astropy.units.Quantity` or array
The Cartesian coordinates of the differentials. If not quantity,
``unit`` should be set.
unit : `~astropy.units.Unit` or str
If given, the differentials will be converted to this unit (or taken to
be in this unit if not given.
copy : bool, optional
If `True` (default), arrays will be copied rather than referenced.
"""
base_representation = NDCartesianRepresentation
attr_classes = dict()
def __init__(self, d_x, unit=None, copy=True):
if unit is None:
if not hasattr(d_x[0], "unit"):
unit = u.one
else:
unit = d_x[0].unit
d_x = u.Quantity(d_x, unit, copy=copy, subok=True)
copy = False
self.attr_classes = dict(
[("d_x" + str(i), u.Quantity) for i in range(1, len(d_x) + 1)]
)
super(coord.CartesianDifferential, self).__init__(*d_x, copy=copy)
ptype = None
for name, _ in self.attr_classes.items():
if ptype is None:
ptype = getattr(self, "_" + name).unit.physical_type
else:
if getattr(self, "_" + name).unit.physical_type != ptype:
raise u.UnitsError(
"All components should have matching " "physical types"
)
cls = self.__class__
if not hasattr(cls, name):
setattr(
cls,
name,
property(
_make_getter(name),
doc=("The '{0}' component of the points(s).".format(name)),
),
)
[docs]
def get_d_xyz(self, xyz_axis=0):
"""Return a vector array of the x, y, and z coordinates.
Parameters
----------
xyz_axis : int, optional
The axis in the final array along which the x, y, z components
should be stored (default: 0).
Returns
-------
d_xs : `~astropy.units.Quantity`
With dimension 3 along ``xyz_axis``.
"""
# Add new axis in x, y, z so one can concatenate them around it.
# NOTE: just use np.stack once our minimum numpy version is 1.10.
result_ndim = self.ndim + 1
if not -result_ndim <= xyz_axis < result_ndim:
raise IndexError(
"xyz_axis {0} out of bounds [-{1}, {1})".format(xyz_axis, result_ndim)
)
if xyz_axis < 0:
xyz_axis += result_ndim
# Get components to the same units (very fast for identical units)
# since np.concatenate cannot deal with quantity.
unit = self._d_x1.unit
sh = self.shape
sh = sh[:xyz_axis] + (1,) + sh[xyz_axis:]
components = [
getattr(self, "_" + name).reshape(sh).to(unit).value
for name in self.components
]
xs_value = np.concatenate(components, axis=xyz_axis)
return u.Quantity(xs_value, unit=unit, copy=COPY_IF_NEEDED)
d_xyz = property(get_d_xyz)
