diff --git a/changelog/156.feature.rst b/changelog/156.feature.rst new file mode 100644 index 00000000..4390777c --- /dev/null +++ b/changelog/156.feature.rst @@ -0,0 +1 @@ +Add ASDF serialization for `VaryingCelestialTransform` and `CoupledCompoundModel`. diff --git a/dkist/io/asdf/converters/__init__.py b/dkist/io/asdf/converters/__init__.py index 21704610..72f606d9 100644 --- a/dkist/io/asdf/converters/__init__.py +++ b/dkist/io/asdf/converters/__init__.py @@ -1,3 +1,4 @@ from .dataset import DatasetConverter from .file_manager import FileManagerConverter +from .models import CoupledCompoundConverter, VaryingCelestialConverter from .tiled_dataset import TiledDatasetConverter diff --git a/dkist/io/asdf/converters/models.py b/dkist/io/asdf/converters/models.py new file mode 100644 index 00000000..21fde2f9 --- /dev/null +++ b/dkist/io/asdf/converters/models.py @@ -0,0 +1,98 @@ +from asdf_astropy.converters.transform.core import TransformConverterBase, parameter_to_value + + +class VaryingCelestialConverter(TransformConverterBase): + tags = [ + "asdf://dkist.nso.edu/tags/varying_celestial_transform-1.0.0", + "asdf://dkist.nso.edu/tags/inverse_varying_celestial_transform-1.0.0", + ] + types = [ + "dkist.wcs.models.VaryingCelestialTransform", + "dkist.wcs.models.InverseVaryingCelestialTransform", + ] + + def select_tag(self, obj, tags, ctx): + from dkist.wcs.models import InverseVaryingCelestialTransform, VaryingCelestialTransform + if isinstance(obj, VaryingCelestialTransform): + return "asdf://dkist.nso.edu/tags/varying_celestial_transform-1.0.0" + elif isinstance(obj, InverseVaryingCelestialTransform): + return "asdf://dkist.nso.edu/tags/inverse_varying_celestial_transform-1.0.0" + else: + raise ValueError(f"Unsupported object: {obj}") # pragma: no cover + + def from_yaml_tree_transform(self, node, tag, ctx): + from dkist.wcs.models import InverseVaryingCelestialTransform, VaryingCelestialTransform + + if tag.endswith("inverse_varying_celestial_transform-1.0.0"): + cls = InverseVaryingCelestialTransform + elif tag.endswith("varying_celestial_transform-1.0.0"): + cls = VaryingCelestialTransform + else: + raise ValueError(f"Unsupported tag: {tag}") # pragma: no cover + + return cls( + crpix=node["crpix"], + cdelt=node["cdelt"], + lon_pole=node["lon_pole"], + crval_table=node["crval_table"], + pc_table=node["pc_table"], + projection=node["projection"], + ) + + def to_yaml_tree_transform(self, model, tag, ctx): + return { + "crpix": parameter_to_value(model.crpix), + "cdelt": parameter_to_value(model.cdelt), + "lon_pole": parameter_to_value(model.lon_pole), + "crval_table": parameter_to_value(model.crval_table), + "pc_table": parameter_to_value(model.pc_table), + "projection": model.projection, + } + + +class CoupledCompoundConverter(TransformConverterBase): + """ + ASDF serialization support for CompoundModel. + """ + tags = [ + "asdf://dkist.nso.edu/tags/coupled_compound_model-1.0.0", + ] + + types = ["dkist.wcs.models.CoupledCompoundModel"] + + def to_yaml_tree_transform(self, model, tag, ctx): + left = model.left + + if isinstance(model.right, dict): + right = { + "keys": list(model.right.keys()), + "values": list(model.right.values()) + } + else: + right = model.right + + return { + "forward": [left, right], + "shared_inputs": model.shared_inputs + } + + def from_yaml_tree_transform(self, node, tag, ctx): + from astropy.modeling.core import Model + + from dkist.wcs.models import CoupledCompoundModel + + oper = "&" + + left = node["forward"][0] + if not isinstance(left, Model): + raise TypeError("Unknown model type '{0}'".format(node["forward"][0]._tag)) # pragma: no cover + + right = node["forward"][1] + if (not isinstance(right, Model) and + not (oper == "fix_inputs" and isinstance(right, dict))): + raise TypeError("Unknown model type '{0}'".format(node["forward"][1]._tag)) # pragma: no cover + + model = CoupledCompoundModel("&", left, right, + shared_inputs=node["shared_inputs"]) + + return model diff --git a/dkist/io/asdf/entry_points.py b/dkist/io/asdf/entry_points.py index c6f0cca4..63c8dac9 100644 --- a/dkist/io/asdf/entry_points.py +++ b/dkist/io/asdf/entry_points.py @@ -11,7 +11,9 @@ else: import importlib.resources as importlib_resources -from dkist.io.asdf.converters import DatasetConverter, FileManagerConverter, TiledDatasetConverter +from dkist.io.asdf.converters import (CoupledCompoundConverter, DatasetConverter, + FileManagerConverter, TiledDatasetConverter, + VaryingCelestialConverter) def get_resource_mappings(): @@ -39,13 +41,16 @@ def get_extensions(): """ Get the list of extensions. """ - converters=[FileManagerConverter(), DatasetConverter(), TiledDatasetConverter()] + dkist_converters = [FileManagerConverter(), DatasetConverter(), TiledDatasetConverter()] + wcs_converters = [VaryingCelestialConverter(), CoupledCompoundConverter()] return [ ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-1.0.0", - converters=converters), + converters=dkist_converters), + ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-wcs-1.0.0", + converters=wcs_converters), # This manifest handles all pre-refactor tags ManifestExtension.from_uri("asdf://dkist.nso.edu/manifests/dkist-0.9.0", - converters=converters, + converters=dkist_converters, # Register that this is a replacement for the old extension legacy_class_names=["dkist.io.asdf.extension.DKISTExtension"]) ] diff --git a/dkist/io/asdf/resources/manifests/dkist-wcs-1.0.0.yaml b/dkist/io/asdf/resources/manifests/dkist-wcs-1.0.0.yaml new file mode 100644 index 00000000..06434bf4 --- /dev/null +++ b/dkist/io/asdf/resources/manifests/dkist-wcs-1.0.0.yaml @@ -0,0 +1,14 @@ +%YAML 1.1 +--- +id: asdf://dkist.nso.edu/dkist/manifests/dkist-wcs-1.0.0 +extension_uri: asdf://dkist.nso.edu/dkist/extensions/dkist-wcs-1.0.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" diff --git a/dkist/io/asdf/resources/schemas/coupled_compound_model-1.0.0.yaml b/dkist/io/asdf/resources/schemas/coupled_compound_model-1.0.0.yaml new file mode 100644 index 00000000..b545cf46 --- /dev/null +++ b/dkist/io/asdf/resources/schemas/coupled_compound_model-1.0.0.yaml @@ -0,0 +1,38 @@ +%YAML 1.1 +--- +$schema: "http://stsci.edu/schemas/yaml-schema/draft-01" +id: "asdf://dkist.nso.edu/schemas/coupled_compound_model-1.0.0" +title: > + Perform a list of subtransforms in series. +description: | + The output of each subtransform is fed into the input of the next + subtransform. + The number of output dimensions of each subtransform must be equal + to the number of input dimensions of the next subtransform in list. + To reorder or add/drop axes, insert `remap_axes` transforms in the + subtransform list. + Invertibility: All ASDF tools are required to be able to compute the + analytic inverse of this transform, by reversing the list of + transforms and applying the inverse of each. +examples: + - + - A series of transforms + - | + ! + shared_inputs: 1 + forward: + - !transform/shift-1.2.0 + offset: 2.0 + - !transform/shift-1.2.0 + offset: 3.0 +allOf: + - $ref: "http://stsci.edu/schemas/asdf/transform/transform-1.2.0" + - properties: + forward: + type: array + items: + $ref: "http://stsci.edu/schemas/asdf/transform/transform-1.2.0" + shared_inputs: + type: number + required: [forward, shared_inputs] +... diff --git a/dkist/io/asdf/resources/schemas/varying_celestial_transform-1.0.0.yaml b/dkist/io/asdf/resources/schemas/varying_celestial_transform-1.0.0.yaml new file mode 100644 index 00000000..443f8b1e --- /dev/null +++ b/dkist/io/asdf/resources/schemas/varying_celestial_transform-1.0.0.yaml @@ -0,0 +1,37 @@ +%YAML 1.1 +--- +$schema: "http://stsci.edu/schemas/yaml-schema/draft-01" +id: "asdf://dkist.nso.edu/schemas/varying_celestial_transform-1.0.0" + +title: A varying FITS-like celestial transform. +description: + A model which represents a FITS-like celestial WCS transform which varies over a third pixel input. + +type: object +properties: + crpix: + anyOf: + - tag: "core/ndarray-1.0.0" + - tag: "tag:stsci.edu:asdf/unit/quantity-1.1.0" + cdelt: + anyOf: + - tag: "core/ndarray-1.0.0" + - tag: "tag:stsci.edu:asdf/unit/quantity-1.1.0" + lon_pole: + anyOf: + - type: number + - tag: "tag:stsci.edu:asdf/unit/quantity-1.1.0" + crval_table: + anyOf: + - tag: "core/ndarray-1.0.0" + - tag: "tag:stsci.edu:asdf/unit/quantity-1.1.0" + pc_table: + anyOf: + - tag: "core/ndarray-1.0.0" + - tag: "tag:stsci.edu:asdf/unit/quantity-1.1.0" + projection: + $ref: "tag:stsci.edu:asdf/transform/transform-1.2.0" + +required: [crpix, cdelt, lon_pole, crval_table, pc_table, projection] +additionalProperties: true +... diff --git a/dkist/io/asdf/tests/test_models.py b/dkist/io/asdf/tests/test_models.py new file mode 100644 index 00000000..89f75a85 --- /dev/null +++ b/dkist/io/asdf/tests/test_models.py @@ -0,0 +1,63 @@ +import numpy as np + +import astropy.modeling.models as m +import astropy.units as u +from asdf.testing.helpers import roundtrip_object +from astropy.coordinates.matrix_utilities import rotation_matrix +from astropy.modeling import CompoundModel + +from dkist.wcs.models import (CoupledCompoundModel, InverseVaryingCelestialTransform, + VaryingCelestialTransform) + + +def test_roundtrip_vct(): + varying_matrix_lt = [rotation_matrix(a)[:2, :2] + for a in np.linspace(0, 90, 10)] * u.arcsec + + vct = VaryingCelestialTransform(crpix=(5, 5) * u.pix, + cdelt=(1, 1) * u.arcsec/u.pix, + crval_table=(0, 0) * u.arcsec, + pc_table=varying_matrix_lt, + lon_pole=180 * u.deg) + new_vct = roundtrip_object(vct) + assert isinstance(new_vct, VaryingCelestialTransform) + new_ivct = roundtrip_object(vct.inverse) + assert isinstance(new_ivct, InverseVaryingCelestialTransform) + + assert u.allclose(u.Quantity(new_vct.crpix), (5, 5) * u.pix) + assert u.allclose(u.Quantity(new_ivct.crpix), (5, 5) * u.pix) + + assert u.allclose(u.Quantity(new_vct.pc_table), varying_matrix_lt) + assert u.allclose(u.Quantity(new_ivct.pc_table), varying_matrix_lt) + + pixel = (0*u.pix, 0*u.pix, 5*u.pix) + world = new_vct(*pixel) + assert u.allclose(world, (359.99804329*u.deg, 0.00017119*u.deg)) + + assert u.allclose(new_ivct(*world, 5*u.pix), pixel[:2], atol=0.01*u.pix) + + +def test_coupled_compound_model(): + ccm = CoupledCompoundModel("&", m.Shift(5), m.Scale(10)) + new = roundtrip_object(ccm) + assert isinstance(new, CoupledCompoundModel) + assert isinstance(new.left, m.Shift) + assert isinstance(new.right, m.Scale) + + assert ccm.n_inputs == new.n_inputs + assert ccm.inputs == new.inputs + + +def test_coupled_compound_model_nested(): + ccm = CoupledCompoundModel("&", m.Shift(5) & m.Scale(2), m.Scale(10) | m.Shift(3)) + new = roundtrip_object(ccm) + assert isinstance(new, CoupledCompoundModel) + assert isinstance(new.left, CompoundModel) + assert isinstance(new.left.left, m.Shift) + assert isinstance(new.left.right, m.Scale) + assert isinstance(new.right, CompoundModel) + assert isinstance(new.right.left, m.Scale) + assert isinstance(new.right.right, m.Shift) + + assert ccm.n_inputs == new.n_inputs + assert ccm.inputs == new.inputs diff --git a/dkist/wcs/tests/test_models.py b/dkist/wcs/tests/test_models.py index 65f20306..c57344f8 100644 --- a/dkist/wcs/tests/test_models.py +++ b/dkist/wcs/tests/test_models.py @@ -67,6 +67,7 @@ def test_varying_transform_pc(): assert u.allclose(vct.inverse(*world, 5*u.pix), pixel[:2], atol=0.01*u.pix) + def test_varying_transform_pc_unitless(): varying_matrix_lt = [rotation_matrix(a)[:2, :2] for a in np.linspace(0, 90, 10)] diff --git a/setup.cfg b/setup.cfg index d5e5fcbb..3e74485c 100644 --- a/setup.cfg +++ b/setup.cfg @@ -22,6 +22,8 @@ include_package_data = True install_requires = appdirs>=1.4 asdf>=2.9.1 + asdf-astropy + asdf-transform-schemas astropy>=4.2 dask[array]>=2 globus-sdk>=1.7,<3.0