def __init__(
self,
potential: T.Any,
*,
frame: TH.OptFrameLikeType = None,
representation_type: TH.OptRepresentationLikeType = None,
):
# if it's a wrapper, have to pop back
if isinstance(potential, PotentialWrapper):
# frame
frame = potential.frame if frame is None else frame
# representation type
representation_type = (potential.representation_type
if representation_type is None else
representation_type)
# potential
potential = potential.__wrapped__
# Initialize wrapper for potential.
self.__wrapped__: object = potential # a la decorator
# the "intrinsic" frame of the potential.
# resolve else-wise (None -> UnFrame)
self._frame = (resolve_framelike(frame)
if frame is not Ellipsis else frame)
self._default_representation = (
resolve_representationlike(representation_type)
if representation_type not in (None,
Ellipsis) else representation_type)
if frame is not Ellipsis and self._default_representation not in (
Ellipsis,
None,
):
self._frame.representation_type = self._default_representation
def test_density(self):
"""Test method ``density``."""
# ---------------
# when there isn't a frame
with pytest.raises(TypeError, match="must have a frame."):
self.subclass.density(self.potential, self.points)
# ---------------
# basic
points, values = self.subclass.density(
self.potential,
self.points.data,
)
# the points are unchanged
assert points is self.points.data
# check data types
assert isinstance(points, coord.BaseRepresentation)
# and on the values
assert u.allclose(values, [0.0, 0.0, 0.0] * u.solMass / u.pc ** 3)
# ---------------
# frame
# test the different inputs
for frame in (
coord.Galactocentric,
coord.Galactocentric(),
"galactocentric",
):
points, values = self.subclass.density(
self.potential,
self.points,
frame=frame,
)
assert isinstance(points, coord.SkyCoord)
assert isinstance(points.frame, resolve_framelike(frame).__class__)
assert u.allclose(values, [0.0, 0.0, 0.0] * u.solMass / u.pc ** 3)
# TODO! test the specific values
# ---------------
# representation_type
points, values = self.subclass.density(
self.potential,
self.points,
frame=self.points.frame.replicate_without_data(),
representation_type=coord.CartesianRepresentation,
)
assert isinstance(points, coord.SkyCoord)
assert isinstance(points.frame, self.frame)
assert isinstance(points.data, coord.CartesianRepresentation)
assert u.allclose(values, [0.0, 0.0, 0.0] * u.solMass / u.pc ** 3)
def _convert_to_frame(
cls,
points: TH.PositionType,
frame: T.Optional[TH.FrameLikeType],
representation_type: TH.OptRepresentationLikeType = None,
) -> T.Tuple[TH.CoordinateType, T.Union[TH.FrameType, str, None]]:
"""Convert points to the given coordinate frame.
Parameters
----------
points : |SkyCoord| or |CoordinateFrame| or |Representation|
The points at which to evaluate the potential.
frame : |CoordinateFrame| or None
The frame in to which `points` are transformed.
If None, then no transformation is applied.
Returns
-------
points : |CoordinateFrame| or |SkyCoord|
Same type as `points`, in the potential's frame.
from_frame : |CoordinateFrame|
The frame of the input points,
"""
resolved_frame = resolve_framelike(frame) # works with None and ...
# -----------
# from_frame
# "from_frame" is the frame of the input points
if isinstance(points, coord.SkyCoord):
from_frame = points.frame.replicate_without_data()
from_rep = points.representation_type
elif isinstance(points, coord.BaseCoordinateFrame):
from_frame = points.replicate_without_data()
from_rep = points.representation_type
elif isinstance(points, coord.BaseRepresentation): # (frame is None)
from_frame = None
from_rep = points.__class__
else:
raise TypeError(
f"points is <{type(points)}> not "
"<SkyCoord, CoordinateFrame, or Representation>.", )
if isinstance(resolved_frame,
UnFrame) and not (from_frame is None
or isinstance(from_frame, UnFrame)):
raise TypeError(
"To pass points as SkyCoord or CoordinateFrame, "
"the potential must have a frame.", )
# -----------
# parse rep
if representation_type is None:
rep_type = from_rep
else:
rep_type = resolve_representationlike(representation_type)
# -----------
# to frame
# potential doesn't have a frame
if frame is None:
p = points
elif from_frame is None: # but frame is not None
p = resolved_frame.realize_frame(
points,
representation_type=rep_type,
)
# don't need to transform
elif (isinstance(frame, coord.BaseCoordinateFrame) # catch comp error
and frame == from_frame # equivalent frames
):
p = points
else:
p = points.transform_to(resolved_frame)
# -----------
# to rep
if isinstance(p, coord.BaseRepresentation):
p = p.represent_as(rep_type)
elif isinstance(p, coord.BaseCoordinateFrame) and rep_type is not None:
p._data = p._data.represent_as(rep_type)
elif isinstance(p, coord.SkyCoord) and rep_type is not None:
p.frame._data = p.frame._data.represent_as(rep_type)
return p, from_frame
def test_specific_force(self):
"""Test method ``specific_force``."""
# ---------------
# when there isn't a frame
with pytest.raises(TypeError, match="must have a frame."):
self.subclass.specific_force(self.potential, self.points)
# ---------------
# basic
vf = self.subclass.specific_force(self.potential, self.points.data)
assert isinstance(vf, vectorfield.BaseVectorField)
assert isinstance(vf.points, coord.CartesianRepresentation)
assert hasattr(vf, "vf_x")
assert hasattr(vf, "vf_y")
assert hasattr(vf, "vf_z")
assert vf.frame is None
# TODO! test the specific values
# ---------------
# frame
# test the different inputs
for frame in (
coord.Galactocentric,
coord.Galactocentric(),
"galactocentric",
):
vf = self.subclass.specific_force(
self.potential,
self.points,
frame=frame,
)
assert isinstance(vf, vectorfield.BaseVectorField)
assert isinstance(vf.points, coord.CartesianRepresentation)
assert hasattr(vf, "vf_x")
assert hasattr(vf, "vf_y")
assert hasattr(vf, "vf_z")
assert isinstance(vf.frame, resolve_framelike(frame).__class__)
# TODO! test the specific values
# ---------------
# representation_type
vf = self.subclass.specific_force(
self.potential,
self.points,
frame=self.points.frame.replicate_without_data(),
representation_type=coord.CylindricalRepresentation,
)
assert isinstance(vf, vectorfield.BaseVectorField)
assert isinstance(vf.points, coord.CylindricalRepresentation)
assert hasattr(vf, "vf_rho")
assert hasattr(vf, "vf_phi")
assert hasattr(vf, "vf_z")
assert isinstance(vf.frame, self.frame)