coordax 0.2.2__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.2
+Version: 0.2.3
 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.3}/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.3'  # keep sync with pyproject.toml

{coordax-0.2.2 → coordax-0.2.3}/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, TypeVar
 import warnings
 
 from coordax import utils
@@ -636,6 +636,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.3}/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.3}/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.3}/coordax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coordax
-Version: 0.2.2
+Version: 0.2.3
 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.3}/pyproject.toml

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