diff --git a/changelog/227.bugfix.rst b/changelog/227.bugfix.rst new file mode 100644 index 00000000..4796ce50 --- /dev/null +++ b/changelog/227.bugfix.rst @@ -0,0 +1,2 @@ +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. diff --git a/dkist/io/asdf/converters/__init__.py b/dkist/io/asdf/converters/__init__.py index 72f606d9..3c159bb8 100644 --- a/dkist/io/asdf/converters/__init__.py +++ b/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 diff --git a/dkist/io/asdf/converters/models.py b/dkist/io/asdf/converters/models.py index 07a1761b..30a549aa 100644 --- a/dkist/io/asdf/converters/models.py +++ b/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"]) diff --git a/dkist/io/asdf/entry_points.py b/dkist/io/asdf/entry_points.py index 1e52731c..b90f6eb0 100644 --- a/dkist/io/asdf/entry_points.py +++ b/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 diff --git a/dkist/io/asdf/resources/manifests/dkist-wcs-1.1.0.yaml b/dkist/io/asdf/resources/manifests/dkist-wcs-1.1.0.yaml new file mode 100644 index 00000000..ead05544 --- /dev/null +++ b/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" diff --git a/dkist/io/asdf/resources/schemas/ravel_model-1.0.0.yaml b/dkist/io/asdf/resources/schemas/ravel_model-1.0.0.yaml new file mode 100644 index 00000000..09ebb423 --- /dev/null +++ b/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] diff --git a/dkist/io/asdf/tests/test_models.py b/dkist/io/asdf/tests/test_models.py index 000b990e..07f8d355 100644 --- a/dkist/io/asdf/tests/test_models.py +++ b/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 diff --git a/dkist/wcs/models.py b/dkist/wcs/models.py index 7ddfe272..f9182ed4 100644 --- a/dkist/wcs/models.py +++ b/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,58 @@ 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) + if order not in ("C", "F"): + raise ValueError("order kwarg must be one of 'C' or 'F'") + self.order = order + + 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"" + + +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 + if order not in ("C", "F"): + raise ValueError("order kwarg must be one of 'C' or 'F'") + self.order = order + + 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"" diff --git a/dkist/wcs/tests/test_models.py b/dkist/wcs/tests/test_models.py index 7de181b7..64c368c4 100644 --- a/dkist/wcs/tests/test_models.py +++ b/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)