coordax 0.2.2__tar.gz → 0.2.4__tar.gz

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

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

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

@@ -53,4 +53,4 @@ from coordax.ndarrays import (
 )
 import coordax.testing  # pylint: disable=unused-import
 
-__version__ = '0.2.2'  # keep sync with pyproject.toml
+__version__ = '0.2.4'  # keep sync with pyproject.toml

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

@@ -25,7 +25,7 @@ import dataclasses
 import functools
 import itertools
 import typing
-from typing import Any, Self, TYPE_CHECKING, TypeAlias, TypeVar
+from typing import Any, Self, TYPE_CHECKING, Type, TypeAlias, TypeGuard, TypeVar
 import warnings
 
 from coordax import utils
@@ -35,7 +35,7 @@ import numpy as np
 if TYPE_CHECKING:
   # import only under TYPE_CHECKING to avoid circular dependency
   # pylint: disable=g-bad-import-order
-  from coordax import fields
+  from coordax import fields  # pylint: disable=unused-import
   import xarray
 
 
@@ -84,7 +84,7 @@ class Coordinate(abc.ABC):
     raise NotImplementedError()
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """Optional dict that maps from field names to their values."""
     return {}
 
@@ -110,7 +110,7 @@ class Coordinate(abc.ABC):
 
   def to_xarray(self) -> dict[str, xarray.Variable]:
     """Convert this coordinate into xarray variables."""
-    import xarray
+    import xarray  # pylint: disable=g-import-not-at-top
 
     variables = {}
     dims_set = {dim for dim in self.dims if dim is not None}
@@ -164,6 +164,12 @@ class ArrayKey:
     return hash((self.value.shape, self.value.tobytes()))
 
 
+@utils.export
+def is_coord(obj: Any) -> TypeGuard[Coordinate]:
+  """Returns True if obj is a Coordinate."""
+  return isinstance(obj, Coordinate)
+
+
 @utils.export
 @jax.tree_util.register_static
 @dataclasses.dataclass(frozen=True)
@@ -211,7 +217,7 @@ class SelectedAxis(Coordinate):
     return (self.coordinate.shape[self.axis],)
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """A maps from field names to their values."""
     return self.coordinate.fields
 
@@ -364,7 +370,7 @@ class CartesianProduct(Coordinate):
     return _concat_tuples(c.shape for c in self.coordinates)
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     """Returns a mapping from field names to their values."""
     return _merge_dicts(c.fields for c in self.coordinates)
 
@@ -479,7 +485,7 @@ class LabeledAxis(Coordinate):
     return self.ticks.shape
 
   @property
-  def fields(self) -> dict[str, fields.Field]:
+  def fields(self) -> dict[str, 'fields.Field']:
     # needs local import to avoid circular dependency
     from coordax import fields  # pylint: disable=g-import-not-at-top
 
@@ -636,6 +642,41 @@ def replace_axes(
   return compose(*axes)
 
 
+@functools.partial(utils.export, module='coordax.coords')
+def extract(
+    coord: Coordinate,
+    component_type: Type[Coordinate] | tuple[Type[Coordinate], ...],
+) -> Coordinate:
+  """Extracts component of type `component_type` from the `coord`.
+
+  Args:
+    coord: The coordinate system to search.
+    component_type: The type(s) of coordinate to extract.
+
+  Returns:
+    The single coordinate component matching the given type(s).
+
+  Raises:
+    ValueError: If there is not exactly one component of the given type(s).
+
+  Examples:
+    >>> import coordax as cx
+    >>> import numpy as np
+    >>> x = cx.LabeledAxis('x', np.linspace(0, np.pi, 4))
+    >>> y = cx.SizedAxis('y', 3)
+    >>> cx.coords.extract(cx.coords.compose(x, y), cx.SizedAxis)
+    coordax.SizedAxis('y', size=3)
+  """
+  components = canonicalize(coord)
+  of_type = [c for c in components if isinstance(c, component_type)]
+  if len(of_type) != 1:
+    raise ValueError(
+        f'Expected exactly one instance of {component_type}, found {of_type}'
+    )
+  [result] = of_type
+  return result
+
+
 @functools.partial(utils.export, module='coordax.coords')
 def from_xarray(
     data_array: xarray.DataArray,

{coordax-0.2.2 → coordax-0.2.4}/coordax/coordinate_systems_test.py

@@ -459,6 +459,27 @@ class CoordinateSystemsTest(parameterized.TestCase):
     ):
       cx.CartesianProduct((x, y))
 
+  def test_extract(self):
+    x = cx.SizedAxis('x', 2)
+    y = cx.SizedAxis('y', 3)
+    z = cx.LabeledAxis('z', np.arange(4))
+    xy = cx.coords.compose(x, y)
+    xz = cx.coords.compose(x, z)
+
+    self.assertEqual(cx.coords.extract(x, cx.SizedAxis), x)
+    self.assertEqual(cx.coords.extract(z, cx.LabeledAxis), z)
+    self.assertEqual(cx.coords.extract(xz, cx.LabeledAxis), z)
+    self.assertEqual(cx.coords.extract(z, (cx.SizedAxis, cx.LabeledAxis)), z)
+
+    with self.assertRaisesRegex(ValueError, 'Expected exactly one instance'):
+      cx.coords.extract(xy, cx.SizedAxis)
+
+    with self.assertRaisesRegex(ValueError, 'Expected exactly one instance'):
+      cx.coords.extract(xz, (cx.SizedAxis, cx.LabeledAxis))
+
+    with self.assertRaisesRegex(ValueError, 'Expected exactly one instance'):
+      cx.coords.extract(x, cx.LabeledAxis)
+
   def test_deprecated_aliases(self):
     with self.assertWarnsRegex(
         DeprecationWarning,

{coordax-0.2.2 → coordax-0.2.4}/coordax/coords.py

@@ -22,6 +22,7 @@ from coordax.coordinate_systems import (
     insert_axes as insert_axes,
     replace_axes as replace_axes,
     from_xarray as from_xarray,
+    extract as extract,
     ArrayKey as ArrayKey,
     NoCoordinateMatch as NoCoordinateMatch,
 )

{coordax-0.2.2 → coordax-0.2.4}/coordax/fields.py

@@ -538,7 +538,7 @@ class Field:
              [0., 0., 0.]], dtype=float32)
       Dimensions without coordinates: x, y
     """
-    import xarray
+    import xarray  # pylint: disable=g-import-not-at-top
 
     if not all(isinstance(dim, str) for dim in self.dims):
       raise ValueError(

{coordax-0.2.2 → coordax-0.2.4}/coordax.egg-info/PKG-INFO

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

{coordax-0.2.2 → coordax-0.2.4}/pyproject.toml

@@ -7,7 +7,7 @@ packages = ["coordax"]
 
 [project]
 name = "coordax"
-version = "0.2.2"  # keep sync with __init__.py
+version = "0.2.4"  # keep sync with __init__.py
 description = "Coordinate axes for scientific computing in JAX"
 authors = [
     {name = "Google LLC", email = "noreply@google.com"},