use exact astropy method names (#1)

* use exact astropy method names
* fix ruff
* rm quax in pre-commit mypy
* add jaxtyping to mypy
* don’t include pypy

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

.github/workflows/ci.yml

@@ -43,9 +43,9 @@ jobs:
         python-version: ["3.10", "3.12"]
         runs-on: [ubuntu-latest, macos-latest, windows-latest]
 
-        include:
-          - python-version: pypy-3.10
-            runs-on: ubuntu-latest
+        # include:
+        #   - python-version: pypy-3.10
+        #     runs-on: ubuntu-latest
 
     steps:
       - uses: actions/checkout@v4

.pre-commit-config.yaml

@@ -54,7 +54,6 @@ repos:
         args: []
         additional_dependencies:
           - pytest
-          - quax
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.2.6"

pyproject.toml

@@ -115,6 +115,7 @@ warn_return_any = false
       "astropy.units.*",
       "equinox.*",
       "jax.*",
+      "jaxtyping.*",
       "quax.*",
     ]
     ignore_missing_imports = true
@@ -131,6 +132,7 @@ ignore = [
   "ANN101",   # Missing type annotation for `self` in method
   "ANN401",   # Dynamically typed expressions (typing.Any) are disallowed
   "ARG001",   # Unused function argument  # TODO: resolve
+  "COM812",   # <- for ruff.format
   "D103",     # Missing docstring in public function  # TODO: resolve
   "D105",     # Missing docstring in magic method
   "D203",     # 1 blank line required before class docstring

src/jax_quantity/_core.py

@@ -20,7 +20,7 @@ class Quantity(ArrayValue):  # type: ignore[misc]
     """Represents an array, with each axis bound to a name."""
 
     value: jax.Array = eqx.field(converter=jax.numpy.asarray)
-    units: Unit = eqx.field(static=True, converter=Unit)
+    unit: Unit = eqx.field(static=True, converter=Unit)
 
     # ===============================================================
     # Quax
@@ -35,21 +35,21 @@ def materialise(self) -> None:
         raise RuntimeError(msg)
 
     def aval(self) -> jax.core.ShapedArray:
-        return jax.core.get_aval(self.value)  # type: ignore[no-untyped-call]
+        return jax.core.get_aval(self.value)
 
     def enable_materialise(self, _: bool = True) -> Self:  # noqa: FBT001, FBT002
-        return type(self)(self.value, self.units)
+        return type(self)(self.value, self.unit)
 
     # ===============================================================
     # Quantity
 
-    def to_units(self, units: Unit) -> "Quantity":
-        return type(self)(self.value * self.units.to(units), units)
+    def to(self, units: Unit) -> "Quantity":
+        return type(self)(self.value * self.unit.to(units), units)
 
-    def to_units_value(self, units: Unit) -> ArrayLike:
-        if units == self.units:
+    def to_value(self, units: Unit) -> ArrayLike:
+        if units == self.unit:
             return self.value
-        return self.value * self.units.to(units)
+        return self.value * self.unit.to(units)
 
     def __getitem__(self, key: Any) -> "Quantity":
         return replace(self, value=self.value[key])

src/jax_quantity/_register_dispatches.py

@@ -17,7 +17,7 @@ def dispatcher(f: T) -> T:  # TODO: figure out mypy stub issue.
 
 @dispatcher
 def empty_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.empty_like(x.value, dtype=dtype), units=x.units)
+    out = Quantity(jnp.empty_like(x.value, dtype=dtype), units=x.unit)
     return jax.device_put(out, device=device)
 
 
@@ -30,17 +30,17 @@ def full_like(
     dtype: Any = None,
     device: Any = None,
 ) -> Quantity:
-    out = Quantity(jnp.full_like(x.value, fill_value, dtype=dtype), units=x.units)
+    out = Quantity(jnp.full_like(x.value, fill_value, dtype=dtype), units=x.unit)
     return jax.device_put(out, device=device)
 
 
 @dispatcher
 def ones_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.ones_like(x.value, dtype=dtype), units=x.units)
+    out = Quantity(jnp.ones_like(x.value, dtype=dtype), units=x.unit)
     return jax.device_put(out, device=device)
 
 
 @dispatcher
 def zeros_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.zeros_like(x.value, dtype=dtype), units=x.units)
+    out = Quantity(jnp.zeros_like(x.value, dtype=dtype), units=x.unit)
     return jax.device_put(out, device=device)