WIP: asdf converters for WCS models (#156)

* Start working on tags

* Use asdf-astropy to convert model

* inital work on CoupleCompoundModel tag

* Make ccm converter work

* Update tests

* Add schemas for models

* Depend on the transform schemas package as we reference it from our schemas

* Add changelog

changelog/156.feature.rst

@@ -0,0 +1 @@
+Add ASDF serialization for `VaryingCelestialTransform` and `CoupledCompoundModel`.

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

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

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"])
     ]

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"

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
+    - |
+      !<asdf://dkist.nso.edu/tags/coupled_compound_model-1.0.0>
+        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]
+...

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
+...

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

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)]
 

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