Skip to content

Commit

Permalink
Fix lonlat (#84)
Browse files Browse the repository at this point in the history 
* fix: LonLatSphericalVector transformations
* feat: out-of-range lonlat constructor

Signed-off-by: nstarman <[email protected]>
  • Loading branch information
@nstarman
nstarman authored Apr 2, 2024
1 parent c39db99 commit 07fd544
Show file tree
Hide file tree
Showing 3 changed files with 184 additions and 7 deletions.
22 changes: 16 additions & 6 deletions src/coordinax/_checks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,28 +16,38 @@
_2pid = Quantity(360, "deg")


def check_r_non_negative(r: BatchableLength) -> BatchableLength:
def check_r_non_negative(
r: BatchableLength, _lower: Quantity["length"] = _0m
) -> BatchableLength:
"""Check that the radial distance is non-negative."""
return eqx.error_if(
r,
xp.any(r < _0m),
xp.any(r < _lower),
"The radial distance must be non-negative.",
)


def check_azimuth_range(phi: BatchableAngle) -> BatchableAngle:
def check_azimuth_range(
phi: BatchableAngle,
_lower: Quantity["angle"] = _0d,
_upper: Quantity["angle"] = _2pid,
) -> BatchableAngle:
"""Check that the polar angle is in the range [0, 2pi)."""
return eqx.error_if(
phi,
xp.any((phi < _0d) | (phi >= _2pid)),
xp.any((phi < _lower) | (phi >= _upper)),
"The azimuthal (polar) angle must be in the range [0, 2pi).",
)


def check_polar_range(theta: BatchableAngle) -> BatchableAngle:
def check_polar_range(
theta: BatchableAngle,
_lower: Quantity["angle"] = _0d,
_upper: Quantity["angle"] = _pid,
) -> BatchableAngle:
"""Check that the inclination angle is in the range [0, pi]."""
return eqx.error_if(
theta,
xp.any((theta < _0d) | (theta > _pid)),
xp.any((theta < _lower) | (theta > _upper)),
"The inclination angle must be in the range [0, pi].",
)
147 changes: 146 additions & 1 deletion src/coordinax/_d3/sphere.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
import equinox as eqx
import jax

import quaxed.array_api as xp
import quaxed.lax as qlax
from unxt import Distance, Quantity

import coordinax._typing as ct
Expand All @@ -34,6 +36,9 @@
from coordinax._converters import converter_azimuth_to_range
from coordinax._utils import classproperty

_90d = Quantity(90, "deg")
_180d = Quantity(180, "deg")

##############################################################################
# Position

Expand Down Expand Up @@ -175,6 +180,9 @@ def norm(self) -> ct.BatchableDistance:
return self.r


# ============================================================================


@final
class LonLatSphericalVector(AbstractSphericalVector):
"""Spherical vector representation.
Expand All @@ -192,6 +200,47 @@ class LonLatSphericalVector(AbstractSphericalVector):
distance : Distance
Radial distance r (slant distance to origin),
Examples
--------
>>> from unxt import Quantity
>>> import coordinax as cx
>>> cx.LonLatSphericalVector(lon=Quantity(0, "deg"), lat=Quantity(0, "deg"),
... distance=Quantity(3, "kpc"))
LonLatSphericalVector(
lon=Quantity[PhysicalType('angle')](value=f32[], unit=Unit("deg")),
lat=Quantity[PhysicalType('angle')](value=f32[], unit=Unit("deg")),
distance=Distance(value=f32[], unit=Unit("kpc"))
)
The longitude and latitude angles are in the range [0, 360) and [-90, 90] degrees,
and the radial distance is non-negative.
When initializing, the longitude is wrapped to the [0, 360) degrees range.
>>> vec = cx.LonLatSphericalVector(lon=Quantity(365, "deg"),
... lat=Quantity(90, "deg"),
... distance=Quantity(3, "kpc"))
>>> vec.lon
Quantity['angle'](Array(5., dtype=float32), unit='deg')
The latitude is not wrapped, but it is checked to be in the [-90, 90] degrees range.
>>> try:
... cx.LonLatSphericalVector(lon=Quantity(0, "deg"), lat=Quantity(100, "deg"),
... distance=Quantity(3, "kpc"))
... except Exception as e:
... print(e)
The inclination angle must be in the range [0, pi]...
Likewise, the radial distance is checked to be non-negative.
>>> try:
... cx.LonLatSphericalVector(lon=Quantity(0, "deg"), lat=Quantity(0, "deg"),
... distance=Quantity(-3, "kpc"))
... except Exception as e:
... print(e)
The radial distance must be non-negative...
"""

lon: ct.BatchableAngle = eqx.field(
Expand All @@ -214,7 +263,7 @@ class LonLatSphericalVector(AbstractSphericalVector):
def __check_init__(self) -> None:
"""Check the validity of the initialization."""
check_azimuth_range(self.lon)
check_polar_range(self.lat)
check_polar_range(self.lat, -Quantity(90, "deg"), Quantity(90, "deg"))
check_r_non_negative(self.distance)

@classproperty
Expand All @@ -239,6 +288,102 @@ def norm(self) -> ct.BatchableDistance:
return self.distance


@LonLatSphericalVector.constructor._f.register # type: ignore[attr-defined, misc] # noqa: SLF001
def constructor(
cls: type[LonLatSphericalVector],
*,
lon: Quantity["angle"],
lat: Quantity["angle"],
distance: Distance,
) -> LonLatSphericalVector:
"""Construct LonLatSphericalVector, allowing for out-of-range values.
Examples
--------
>>> import coordinax as cx
Let's start with a valid input:
>>> cx.LonLatSphericalVector.constructor(lon=Quantity(0, "deg"),
... lat=Quantity(0, "deg"),
... distance=Quantity(3, "kpc"))
LonLatSphericalVector(
lon=Quantity[PhysicalType('angle')](value=f32[], unit=Unit("deg")),
lat=Quantity[PhysicalType('angle')](value=f32[], unit=Unit("deg")),
distance=Distance(value=f32[], unit=Unit("kpc"))
)
The distance can be negative, which wraps the longitude by 180 degrees and
flips the latitude:
>>> vec = cx.LonLatSphericalVector.constructor(lon=Quantity(0, "deg"),
... lat=Quantity(45, "deg"),
... distance=Quantity(-3, "kpc"))
>>> vec.lon
Quantity['angle'](Array(180., dtype=float32), unit='deg')
>>> vec.lat
Quantity['angle'](Array(-45., dtype=float32), unit='deg')
>>> vec.distance
Distance(Array(3., dtype=float32), unit='kpc')
The latitude can be outside the [-90, 90] deg range, causing the longitude
to be shifted by 180 degrees:
>>> vec = cx.LonLatSphericalVector.constructor(lon=Quantity(0, "deg"),
... lat=Quantity(-100, "deg"),
... distance=Quantity(3, "kpc"))
>>> vec.lon
Quantity['angle'](Array(180., dtype=float32), unit='deg')
>>> vec.lat
Quantity['angle'](Array(-80., dtype=float32), unit='deg')
>>> vec.distance
Distance(Array(3., dtype=float32), unit='kpc')
>>> vec = cx.LonLatSphericalVector.constructor(lon=Quantity(0, "deg"),
... lat=Quantity(100, "deg"),
... distance=Quantity(3, "kpc"))
>>> vec.lon
Quantity['angle'](Array(180., dtype=float32), unit='deg')
>>> vec.lat
Quantity['angle'](Array(80., dtype=float32), unit='deg')
>>> vec.distance
Distance(Array(3., dtype=float32), unit='kpc')
The longitude can be outside the [0, 360) deg range. This is wrapped to the
[0, 360) deg range (actually the base constructor does this):
>>> vec = cx.LonLatSphericalVector.constructor(lon=Quantity(365, "deg"),
... lat=Quantity(0, "deg"),
... distance=Quantity(3, "kpc"))
>>> vec.lon
Quantity['angle'](Array(5., dtype=float32), unit='deg')
"""
# 1) Convert the inputs
fields = LonLatSphericalVector.__dataclass_fields__
lon = fields["lon"].metadata["converter"](lon)
lat = fields["lat"].metadata["converter"](lat)
distance = fields["distance"].metadata["converter"](distance)

# 2) handle negative distances
distance_pred = distance < xp.zeros_like(distance)
distance = qlax.select(distance_pred, -distance, distance)
lon = qlax.select(distance_pred, lon + _180d, lon)
lat = qlax.select(distance_pred, -lat, lat)

# 3) Handle latitude outside of [-90, 90] degrees
lat_pred = lat < -_90d
lat = qlax.select(lat_pred, -_180d - lat, lat)
lon = qlax.select(lat_pred, lon + _180d, lon)

lat_pred = lat > _90d
lat = qlax.select(lat_pred, _180d - lat, lat)
lon = qlax.select(lat_pred, lon + _180d, lon)

# 4) Construct. This also handles the longitude wrapping
return cls(lon=lon, lat=lat, distance=distance)


##############################################################################


Expand Down
22 changes: 22 additions & 0 deletions src/coordinax/_d3/transform.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -654,3 +654,25 @@ def represent_as(
d_lat=current.d_lat,
d_distance=current.d_distance,
)


@dispatch
def represent_as(
current: LonCosLatSphericalDifferential,
target: type[Abstract3DVectorDifferential],
position: AbstractVector | Quantity["length"],
/,
**kwargs: Any,
) -> Abstract3DVectorDifferential:
"""LonCosLatSphericalDifferential -> Abstract3DVectorDifferential."""
# Parse the position to an AbstractVector
if isinstance(position, AbstractVector):
posvec = position
else: # Q -> Cart<X>D
posvec = current.integral_cls._cartesian_cls.constructor( # noqa: SLF001
position
)
# Transform the differential to LonLatSphericalDifferential
current = represent_as(current, LonLatSphericalDifferential, posvec)
# Transform the position to the required type
return represent_as(current, target, posvec)

0 comments on commit 07fd544

Please sign in to comment.