Ravel model for the illusion of 2D tables

changelog/227.bugfix.rst

@@ -0,0 +1,3 @@
+Add a new model to take a 2D index and return the corresponding correct index for a 1D array, and the inverse model for the reverse operation.
+To be used as a compound with Tabular1D so that it looks like a Tabular2D but the compound model can still be inverted.
+This will be relevant and useful in the dkist-inventory package but users are unlikely to need to use it.

dkist/io/asdf/converters/__init__.py

@@ -1,4 +1,4 @@
 from .dataset import DatasetConverter
 from .file_manager import FileManagerConverter
-from .models import CoupledCompoundConverter, VaryingCelestialConverter
+from .models import CoupledCompoundConverter, RavelConverter, VaryingCelestialConverter
 from .tiled_dataset import TiledDatasetConverter

dkist/io/asdf/converters/models.py

@@ -115,3 +115,23 @@ def from_yaml_tree_transform(self, node, tag, ctx):
                                      shared_inputs=node["shared_inputs"])
 
         return model
+
+
+class RavelConverter(TransformConverterBase):
+    """
+    ASDF serialization support for Ravel
+    """
+
+    tags  = [
+        "asdf://dkist.nso.edu/tags/ravel_model-1.0.0"
+    ]
+
+    types = ["dkist.wcs.models.Ravel"]
+
+    def to_yaml_tree_transform(self, model, tag, ctx):
+        return {"array_shape": model.array_shape, "order": model.order}
+
+    def from_yaml_tree_transform(self, node, tag, ctx):
+        from dkist.wcs.models import Ravel
+
+        return Ravel(node["array_shape"], order=node["order"])

dkist/io/asdf/entry_points.py

@@ -6,14 +6,15 @@
 from asdf.extension import ManifestExtension
 from asdf.resource import DirectoryResourceMapping
 
+from dkist.io.asdf.converters import (CoupledCompoundConverter, DatasetConverter,
+                                      FileManagerConverter, RavelConverter, TiledDatasetConverter,
+                                      VaryingCelestialConverter)
+
 if sys.version_info < (3, 9):
     import importlib_resources
 else:
     import importlib.resources as importlib_resources
 
-from dkist.io.asdf.converters import (CoupledCompoundConverter, DatasetConverter,
-                                      FileManagerConverter, TiledDatasetConverter,
-                                      VaryingCelestialConverter)
 
 
 def get_resource_mappings():
@@ -42,12 +43,14 @@ def get_extensions():
     Get the list of extensions.
     """
     dkist_converters = [FileManagerConverter(), DatasetConverter(), TiledDatasetConverter()]
-    wcs_converters = [VaryingCelestialConverter(), CoupledCompoundConverter()]
+    wcs_converters = [VaryingCelestialConverter(), CoupledCompoundConverter(), RavelConverter()]
     return [
         ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-1.1.0",
                                    converters=dkist_converters),
         ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-1.0.0",
                                    converters=dkist_converters),
+        ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-wcs-1.1.0",
+                                   converters=wcs_converters),
         ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-wcs-1.0.0",
                                    converters=wcs_converters),
         # This manifest handles all pre-refactor tags

dkist/io/asdf/resources/manifests/dkist-wcs-1.1.0.yaml

@@ -0,0 +1,20 @@
+%YAML 1.1
+---
+id: asdf://dkist.nso.edu/dkist/manifests/dkist-wcs-1.1.0
+extension_uri: asdf://dkist.nso.edu/dkist/extensions/dkist-wcs-1.1.0
+title: DKIST WCS extension
+description: ASDF schemas and tags for models and WCS related classes.
+
+tags:
+  - schema_uri: "asdf://dkist.nso.edu/schemas/varying_celestial_transform-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/varying_celestial_transform-1.0.0"
+  - schema_uri: "asdf://dkist.nso.edu/schemas/varying_celestial_transform-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/inverse_varying_celestial_transform-1.0.0"
+  - schema_uri: "asdf://dkist.nso.edu/schemas/coupled_compound_model-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/coupled_compound_model-1.0.0"
+  - schema_uri: "asdf://dkist.nso.edu/schemas/varying_celestial_transform-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/varying_celestial_transform_slit-1.0.0"
+  - schema_uri: "asdf://dkist.nso.edu/schemas/varying_celestial_transform-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/inverse_varying_celestial_transform_slit-1.0.0"
+  - schema_uri: "asdf://dkist.nso.edu/schemas/ravel_model-1.0.0"
+    tag_uri: "asdf://dkist.nso.edu/tags/ravel_model-1.0.0"

dkist/io/asdf/resources/schemas/ravel_model-1.0.0.yaml

@@ -0,0 +1,16 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "asdf://dkist.nso.edu/schemas/ravel_model-1.0.0"
+
+title: A model to flatten 2D indices into 1D
+description:
+  A model which takes a pair of indices and flattens them into the corresponding index for a 1D array. This can be used as a compound with Tabular1D to enable it to be indexed as if it were Tabular2D.
+
+type: object
+properties:
+  array_shape:
+    type: array
+  order:
+    type: string
+required: [array_shape, order]

dkist/io/asdf/tests/test_models.py

@@ -9,9 +9,9 @@
 from dkist.wcs.models import (CoupledCompoundModel, InverseVaryingCelestialTransform,
                               InverseVaryingCelestialTransform2D,
                               InverseVaryingCelestialTransformSlit,
-                              InverseVaryingCelestialTransformSlit2D, VaryingCelestialTransform,
-                              VaryingCelestialTransform2D, VaryingCelestialTransformSlit,
-                              VaryingCelestialTransformSlit2D)
+                              InverseVaryingCelestialTransformSlit2D, Ravel, Unravel,
+                              VaryingCelestialTransform, VaryingCelestialTransform2D,
+                              VaryingCelestialTransformSlit, VaryingCelestialTransformSlit2D)
 
 
 def test_roundtrip_vct():
@@ -122,3 +122,13 @@ def test_coupled_compound_model_nested():
 
     assert ccm.n_inputs == new.n_inputs
     assert ccm.inputs == new.inputs
+
+
+def test_ravel_model():
+    ravel = Ravel((10, 10))
+    new = roundtrip_object(ravel)
+
+    assert isinstance(new, Ravel)
+    assert isinstance(new.inverse, Unravel)
+
+    assert new.array_shape == ravel.array_shape

dkist/wcs/models.py

@@ -14,7 +14,8 @@
 __all__ = ['CoupledCompoundModel',
            'InverseVaryingCelestialTransform',
            'VaryingCelestialTransform',
-           'BaseVaryingCelestialTransform']
+           'BaseVaryingCelestialTransform',
+           'Ravel']
 
 
 def generate_celestial_transform(crpix: Union[Iterable[float], u.Quantity],
@@ -733,3 +734,54 @@ def varying_celestial_transform_from_tables(crpix: Union[Iterable[float], u.Quan
         lon_pole=lon_pole,
         projection=projection
     )
+
+
+class Ravel(Model):
+    array_shape = (0, 0)
+    n_inputs = 2
+    n_outputs = 1
+    _separable = False
+
+    def __init__(self, array_shape, order='C', **kwargs):
+        super().__init__(**kwargs)
+
+        self.array_shape = tuple(array_shape)
+        self.order = order  # TODO: Validate order is either C or F
+
+    def evaluate(self, *inputs):
+        ravel_shape = self.array_shape[1]
+        if self.order == 'F':
+            inputs = inputs[::-1]
+            ravel_shape = self.array_shape[0]
+        return (np.round(inputs[0]) * ravel_shape) + inputs[1]
+
+    @property
+    def inverse(self):
+        return Unravel(self.array_shape, order=self.order)
+
+    def __repr__(self):
+        return f"<Ravel(array_shape={self.array_shape}, order={self.order})>"
+
+
+class Unravel(Model):
+    array_shape = (0, 0)
+    n_inputs = 1
+    n_outputs = 2
+    _separable = False
+
+    def __init__(self, array_shape, order='C', **kwargs):
+        super().__init__(**kwargs)
+
+        self.array_shape = array_shape
+        self.order = order  # TODO: Validate order is either C or F
+
+    def evaluate(self, input_):
+        i = 1 if self.order == 'C' else 0
+        return (input_ // self.array_shape[i], input_ % self.array_shape[i])
+
+    @property
+    def inverse(self):
+        return Ravel(self.array_shape)
+
+    def __repr__(self):
+        return f"<Unravel(array_shape={self.array_shape}, order={self.order})>"

dkist/wcs/tests/test_models.py

@@ -5,8 +5,9 @@
 import astropy.units as u
 from astropy.coordinates.matrix_utilities import rotation_matrix
 from astropy.modeling import CompoundModel
+from astropy.modeling.models import Tabular1D
 
-from dkist.wcs.models import (VaryingCelestialTransform, VaryingCelestialTransform2D,
+from dkist.wcs.models import (Ravel, VaryingCelestialTransform, VaryingCelestialTransform2D,
                               VaryingCelestialTransformSlit, VaryingCelestialTransformSlit2D,
                               generate_celestial_transform,
                               varying_celestial_transform_from_tables)
@@ -371,3 +372,67 @@ def test_vct_slit2d_unitless():
     world = vct_slit(*pixel)
     ipixel = vct_slit.inverse(*world, 0, 0)
     assert u.allclose(ipixel, pixel[0], atol=1e-5)
+
+
+@pytest.mark.parametrize("array_shape",
+                         [(i, 100 // i) for i in range(2, 21)])
+def test_ravel_model(array_shape):
+    ravel = Ravel(array_shape)
+
+    # Make 10 attempts with random numbers
+    for _ in range(10):
+        x, y = np.random.random() * (array_shape[0]-1), np.random.random() * array_shape[1]
+        expected_val = ((round(x) * array_shape[1]) + y)
+        assert ravel(x, y) == expected_val
+        assert np.allclose(ravel.inverse(expected_val), (round(x), y))
+        assert ravel.inverse.inverse(x, y) == expected_val
+
+
+@pytest.mark.parametrize("array_shape",
+                         [(i, 100 // i) for i in range(2, 21)])
+def test_raveled_tabular1d(array_shape):
+    values = np.arange(100)
+
+    ravel = Ravel(array_shape)
+    tabular = Tabular1D(values,
+                        values*u.nm,
+                        bounds_error=False,
+                        fill_value=np.nan,
+                        method="linear")
+
+    raveled_tab = ravel | tabular
+
+    # Make 10 attempts with random numbers
+    for _ in range(10):
+        x, y = (np.random.random() * (array_shape[0]-1),
+                np.random.random() * (array_shape[1]-1))
+        expected_val = ((round(x) * array_shape[1]) + y) * u.nm
+        assert raveled_tab(x, y) == expected_val
+        assert np.allclose(raveled_tab.inverse(expected_val), (round(x), y))
+        assert raveled_tab.inverse.inverse(x, y) == expected_val
+
+
+@pytest.mark.parametrize("array_shape",
+                         [(i, 100 // i) for i in range(2, 21)])
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_ravel_ordering(array_shape, order):
+    ravel = Ravel(array_shape, order=order)
+
+    values = np.arange(array_shape[0]*array_shape[1]).reshape(array_shape, order=order)
+
+    # Make 10 attempts with random numbers
+    for _ in range(10):
+        x, y = (np.random.randint((array_shape[0]-1)),
+                np.random.randint((array_shape[1]-1)))
+        assert ravel(x, y) == values[x, y]
+
+
+@pytest.mark.parametrize("array_shape",
+                         [(i, 100 // i) for i in range(2, 21)])
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_ravel_repr(array_shape, order):
+    ravel = Ravel(array_shape, order=order)
+    unravel = ravel.inverse
+
+    assert str(array_shape) in repr(ravel) and order in repr(ravel)
+    assert str(array_shape) in repr(unravel) and order in repr(unravel)