GalacticDynamics / unxt

Unitful Quantities in JAX
https://unxt.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
19 stars 3 forks source link

Astropy equivalencies not fully supported #234

Closed adrn closed 2 days ago

adrn commented 1 month ago

We have this doctest example:

x = Quantity([1, 2, 3], "Kelvin")
with u.add_enabled_equivalencies(u.temperature()):
    y = x.to(u.deg_C)

which works, but this one doesn't work:

x = Quantity([1, 2, 3], "radian")
with u.add_enabled_equivalencies(u.dimensionless_angles()):
    y = x.to(u.one)

I think this is because the physical type (dimension) changes in the latter case and not the first. See https://github.com/GalacticDynamics/unxt/blob/main/src/unxt/_src/quantity/base_parametric.py#L43

adrn commented 1 month ago

Actually I think the main issue is with this line - I don't think we want to do a replace here (that line just replaces the unit and value of a quantity), because conversions with equivalencies can change the physical type (dimensions) of a quantity (such as with dimensionless_angles()).

nstarman commented 1 month ago

Yes. The trick is how to do replace while changing the class. replace has the nice property of propagating any unchanged arguments.

adrn commented 1 month ago

Right. Though do we expect other arguments? If we don't mind type coercion, one option would be to swap the replace for:

return Quantity(_apy7_unit_to(x.unit, unit, x.value), unit)
nstarman commented 1 month ago

We can support all fields by using the dataclassish.field_items(). Things like Angle have the wrap_angle field. Preventing unnecessary type coercion is a bit more fiddly. I think we can do this by checking the dimensions of the output unit against the input units and only doing type coercion if they are different.