coordax 0.2.0__tar.gz → 0.2.2__tar.gz

{coordax-0.2.0 → coordax-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.0
+Version: 0.2.2
 Summary: Coordinate axes for scientific computing in JAX
 Author-email: Google LLC <noreply@google.com>
 License-Expression: Apache-2.0

{coordax-0.2.0 → coordax-0.2.2}/coordax/__init__.py

@@ -41,6 +41,7 @@ from coordax.fields import (
     cmap as cmap,
     cpmap as cpmap,
     get_coordinate as get_coordinate,
+    from_xarray as from_xarray,
     wrap_like as wrap_like,  # deprecated
     wrap as wrap,  # deprecated
     tag as tag,
@@ -52,4 +53,4 @@ from coordax.ndarrays import (
 )
 import coordax.testing  # pylint: disable=unused-import
 
-__version__ = '0.2.0'  # keep sync with pyproject.toml
+__version__ = '0.2.2'  # keep sync with pyproject.toml

{coordax-0.2.0 → coordax-0.2.2}/coordax/coordinate_systems.py

@@ -46,7 +46,7 @@ Sequence = collections.abc.Sequence
 @functools.partial(utils.export, module='coordax.coords')
 @dataclasses.dataclass(frozen=True)
 class NoCoordinateMatch:
-  """For use when a coordinate does not match an xarray.Coordinates object."""
+  """For use when no Coordax coordinate matches xarray coordinate."""
 
   reason: str
 
@@ -56,18 +56,15 @@ class Coordinate(abc.ABC):
   """Abstract class for coordinate objects.
 
   Coordinate subclasses are expected to obey several invariants:
-  1. Dimension names may not be repeated: `len(set(dims)) == len(dims)`
-  2. All dimensions must be named: `len(shape) == len(dims)`
+  1. Dimension names may not be repeated: ``len(set(dims)) == len(dims)``
+  2. All dimensions must be named: ``len(shape) == len(dims)``
 
   Every non-abstract Coordinate subclass must be registered as a "static"
   pytree node, e.g., by decorating the class with
-  `@jax.tree_util.register_static`. Static pytrees nodes must implement
-  `__hash__` and `__eq__` according to the requirements of keys in Python
+  ``@jax.tree_util.register_static``. Static pytrees nodes must implement
+  ``__hash__`` and ``__eq__`` according to the requirements of keys in Python
   dictionaries. This is easiest to acheive with frozen dataclasses, but care
   must be taken when working with np.ndarray attributes.
-
-  TODO(shoyer): add documentation examples, including a version using ArrayKey
-  to wrap np.ndarray attributions.
   """
 
   @property
@@ -76,7 +73,7 @@ class Coordinate(abc.ABC):
     """Dimension names of the coordinate.
 
     All subclasses must return a tuple of dimension names as strings, with the
-    exception of `DummyAxis`.
+    exception of ``DummyAxis``.
     """
     raise NotImplementedError()
 
@@ -138,8 +135,12 @@ class Coordinate(abc.ABC):
         values.
 
     Returns:
-      A matching instance of this coordinate or `NoCoordinateMatch` if this
-      coordinate does not match the xarray dimensions and coordinates.
+      A matching instance of this coordinate or ``NoCoordinateMatch`` if this
+      coordinate type does not match the xarray dimensions and coordinates.
+
+    See also:
+      :func:`coordax.from_xarray`
+      :func:`coordax.coords.from_xarray`
     """
     raise NotImplementedError('from_xarray not implemented')
 
@@ -167,7 +168,7 @@ class ArrayKey:
 @jax.tree_util.register_static
 @dataclasses.dataclass(frozen=True)
 class Scalar(Coordinate):
-  """Zero dimensional sentinel coordinate used to label stand alone scalars."""
+  """Zero dimensional sentinel coordinate used to label scalar fields."""
 
   @property
   def dims(self) -> tuple[str, ...]:
@@ -280,7 +281,27 @@ def _consolidate_coordinates(
 
 @functools.partial(utils.export, module='coordax.coords')
 def canonicalize(*coordinates: Coordinate) -> tuple[Coordinate, ...]:
-  """Canonicalize coordinates into a minimum equivalent collection."""
+  """Canonicalize coordinates into a minimum equivalent collection.
+
+  Args:
+    *coordinates: The coordinates to canonicalize.
+
+  Returns:
+    A tuple of canonicalized coordinates, where ``CartesianProduct`` objects are
+    flattened, ``Scalar`` objects are removed and ``SelectedAxis`` objects are
+    merged back into the original coordinate if all axes are selected in order.
+
+  Examples:
+    >>> import coordax as cx
+    >>> x = cx.SizedAxis('x', 2)
+    >>> y = cx.SizedAxis('y', 3)
+    >>> cx.coords.canonicalize(x, y)
+    (coordax.SizedAxis('x', size=2), coordax.SizedAxis('y', size=3))
+
+    >>> xy = cx.coords.compose(x, y)
+    >>> cx.coords.canonicalize(xy)
+    (coordax.SizedAxis('x', size=2), coordax.SizedAxis('y', size=3))
+  """
   coordinates = _expand_coordinates(*coordinates)
   coordinates = _consolidate_coordinates(*coordinates)
   existing_dims = collections.Counter()
@@ -316,7 +337,10 @@ def _merge_dicts(dicts: Iterable[dict[K, V]]) -> dict[K, V]:
 @jax.tree_util.register_static
 @dataclasses.dataclass(frozen=True)
 class CartesianProduct(Coordinate):
-  """Coordinate defined as the outer product of independent coordinates."""
+  """Coordinate defined as the outer product of independent coordinates.
+
+  To construct a ``CartesianProduct``, use :func:`coordax.coords.compose`.
+  """
 
   coordinates: tuple[Coordinate, ...]
 
@@ -492,7 +516,24 @@ class LabeledAxis(Coordinate):
 
 @functools.partial(utils.export, module='coordax.coords')
 def compose(*coordinates: Coordinate) -> Coordinate:
-  """Compose coordinates into a unified coordinate system."""
+  # pylint: disable=line-too-long
+  # fmt: off
+  """Compose coordinates into a unified coordinate system.
+
+  Args:
+    *coordinates: The coordinates to compose.
+
+  Returns:
+    A single coordinate object representing the Cartesian product of the inputs.
+
+  Examples:
+    >>> import coordax as cx
+    >>> x = cx.SizedAxis('x', 2)
+    >>> y = cx.SizedAxis('y', 3)
+    >>> cx.coords.compose(x, y)
+    CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.SizedAxis('y', size=3)))
+  """
+  # fmt: on
   product = CartesianProduct(coordinates)
   match len(product.coordinates):
     case 0:
@@ -508,7 +549,26 @@ def insert_axes(
     coordinate: Coordinate,
     indices_to_axes: dict[int, Coordinate],
 ) -> Coordinate:
-  """Returns `coordinate` with extra axes inserted at specified positions."""
+  # pylint: disable=line-too-long
+  # fmt: off
+  """Returns ``coordinate`` with extra axes inserted at specified positions.
+
+  Args:
+    coordinate: The coordinate system to modify.
+    indices_to_axes: A mapping from insertion index to the new coordinate to
+      insert. Indices are relative to the *output* coordinate system.
+
+  Returns:
+    A new coordinate object with the axes inserted.
+
+  Examples:
+    >>> import coordax as cx
+    >>> x = cx.SizedAxis('x', 2)
+    >>> z = cx.SizedAxis('z', 4)
+    >>> cx.coords.insert_axes(x, {1: z})
+    CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.SizedAxis('z', size=4)))
+  """
+  # fmt: on
   indices_to_axes = indices_to_axes.copy()
   ndim = coordinate.ndim + len(indices_to_axes)
   normalize_idx = lambda i: i + ndim if i < 0 else i
@@ -528,7 +588,9 @@ def replace_axes(
     to_replace: Coordinate,
     replace_with: Coordinate,
 ) -> Coordinate:
-  """Returns `coordinate` with `to_replace` replaced by `replace_with`.
+  # pylint: disable=line-too-long
+  # fmt: off
+  """Returns ``coordinate`` with ``to_replace`` replaced by ``replace_with``.
 
   Args:
     coordinate: The coordinate system to modify.
@@ -536,11 +598,23 @@ def replace_axes(
     replace_with: The new coordinate to insert.
 
   Returns:
-    A new coordinate object with `to_replace` substituted by `replace_with`.
+    A new coordinate object with ``to_replace`` substituted by ``replace_with``.
 
   Raises:
-    ValueError: If `to_replace` is not a contiguous part of `coordinate`.
+    ValueError: If ``to_replace`` is not a contiguous part of ``coordinate``.
+
+  Examples:
+    >>> import coordax as cx
+    >>> x = cx.SizedAxis('x', 2)
+    >>> y = cx.SizedAxis('y', 3)
+    >>> xy = cx.coords.compose(x, y)
+    >>> xy
+    CartesianProduct(coordinates=(coordax.SizedAxis('x', size=2), coordax.SizedAxis('y', size=3)))
+    >>> z = cx.SizedAxis('z', 4)
+    >>> cx.coords.replace_axes(xy, x, z)
+    CartesianProduct(coordinates=(coordax.SizedAxis('z', size=4), coordax.SizedAxis('y', size=3)))
   """
+  # fmt: on
   if not to_replace.dims:
     raise ValueError(f'`to_replace` must have dimensions, got {to_replace}')
 
@@ -567,20 +641,38 @@ def from_xarray(
     data_array: xarray.DataArray,
     coord_types: Sequence[type[Coordinate]] = (LabeledAxis, DummyAxis),
 ) -> Coordinate:
-  """Convert the coordinates of an xarray.DataArray into a coordax.Coordinate.
+  # pylint: disable=line-too-long
+  # fmt: off
+  """Convert coordinates of an ``xarray.DataArray`` into ``coordax.Coordinate``.
 
   Args:
-    data_array: xarray.DataArray whose coordinates should be converted.
-    coord_types: sequence of coordax.Coordinate subclasses with `from_xarray`
-      methods defined. The first coordinate class that returns a coordinate
-      object (indicating a match) will be used. By default, coordinates will use
-      only generic coordax.LabeledAxis objects. CardesianProduct type is omitted
-      from this sequence since it is introduced by the compose() method.
+    data_array: ``xarray.DataArray`` whose coordinates should be converted.
+    coord_types: sequence of ``coordax.Coordinate`` subclasses with
+      ``from_xarray`` methods defined. The first coordinate class that returns a
+      coordinate object (indicating a match) will be used. By default,
+      coordinates will use only generic ``LabeledAxis`` and ``DummyAxis``
+      objects.
 
   Returns:
-    A coordax.Coordinate object representing the coordinates of the input
-    DataArray.
+    A coordax.Coordinate object, constructed from any of the indicated types
+    (plus ``CartesianProduct`` and ``Scalar``), representing the coordinates of
+    the input DataArray.
+
+  Raises:
+    ValueError: if no matching coordinate is found.
+
+  Examples:
+    >>> import coordax as cx
+    >>> import xarray as xr
+    >>> import numpy as np
+    >>> da = xr.DataArray(np.zeros((2, 3)), dims=('x', 'y'), coords={'x': [1, 2]})
+    >>> cx.coords.from_xarray(da)
+    CartesianProduct(coordinates=(coordax.LabeledAxis('x', ticks=array([1, 2])), coordax.DummyAxis('y', size=3)))
+
+  See also:
+    :func:`coordax.from_xarray`
   """
+  # fmt: on
   dims = data_array.dims
   coords = []
 

coordax-0.2.2/coordax/doctest_test.py

@@ -0,0 +1,31 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import doctest
+from absl.testing import absltest
+import coordax
+import pytest
+
+
+pytest.importorskip('xarray')
+
+
+def load_tests(loader, tests, ignore):
+  tests.addTests(doctest.DocTestSuite(coordax))
+  tests.addTests(doctest.DocTestSuite(coordax.coords))
+  tests.addTests(doctest.DocTestSuite(coordax.testing))
+  return tests
+
+
+if __name__ == '__main__':
+  absltest.main()

coordax-0.2.2/coordax/experimental.py

@@ -0,0 +1,22 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Experimental parts of Coordax."""
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/jax-ml/jax/issues/7570
+# pylint: disable=g-multiple-import,useless-import-alias,g-importing-member,unused-import
+from coordax.ndarrays import (
+    NDArray as NDArray,
+    register_ndarray as register_ndarray,
+)