From da6c639f0cf03cc3278a6981cc68bff741f8d0dc Mon Sep 17 00:00:00 2001 From: ViciousEagle03 Date: Fri, 16 Aug 2024 12:40:18 +0530 Subject: [PATCH 1/4] Add the serialization support for the ResampledLowlevelWCS ReorderedLowLevelWCS CompoundLowLevelWCS --- .../asdf/converters/compoundwcs_converter.py | 18 ++++ ndcube/asdf/converters/ndcube_converter.py | 7 +- .../asdf/converters/reorderedwcs_converter.py | 21 +++++ ndcube/asdf/converters/resampled_converter.py | 21 +++++ .../tests/test_ndcube_wcs_wrappers.py | 92 +++++++++++++++++++ ndcube/asdf/entry_points.py | 6 ++ .../resources/manifests/ndcube-0.1.0.yaml | 9 ++ .../resources/schemas/compoundwcs-0.1.0.yaml | 25 +++++ .../asdf/resources/schemas/ndcube-0.1.0.yaml | 6 +- .../resources/schemas/reorderedwcs-0.1.0.yaml | 23 +++++ .../resources/schemas/resampledwcs-0.1.0.yaml | 23 +++++ 11 files changed, 249 insertions(+), 2 deletions(-) create mode 100644 ndcube/asdf/converters/compoundwcs_converter.py create mode 100644 ndcube/asdf/converters/reorderedwcs_converter.py create mode 100644 ndcube/asdf/converters/resampled_converter.py create mode 100644 ndcube/asdf/converters/tests/test_ndcube_wcs_wrappers.py create mode 100644 ndcube/asdf/resources/schemas/compoundwcs-0.1.0.yaml create mode 100644 ndcube/asdf/resources/schemas/reorderedwcs-0.1.0.yaml create mode 100644 ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml diff --git a/ndcube/asdf/converters/compoundwcs_converter.py b/ndcube/asdf/converters/compoundwcs_converter.py new file mode 100644 index 000000000..b812fd051 --- /dev/null +++ b/ndcube/asdf/converters/compoundwcs_converter.py @@ -0,0 +1,18 @@ +from asdf.extension import Converter + + +class CompoundConverter(Converter): + tags = ["tag:sunpy.org:ndcube/compoundwcs-0.1.0"] + types = ["ndcube.wcs.wrappers.compound_wcs.CompoundLowLevelWCS"] + + def from_yaml_tree(self, node, tag, ctx): + from ndcube.wcs.wrappers import CompoundLowLevelWCS + + return(CompoundLowLevelWCS(*node["wcs"], mapping = node.get("mapping"), pixel_atol = node.get("atol"))) + + def to_yaml_tree(self, compoundwcs, tag, ctx): + node={} + node["wcs"] = compoundwcs._wcs + node["mapping"] = compoundwcs.mapping.mapping + node["atol"] = compoundwcs.atol + return node diff --git a/ndcube/asdf/converters/ndcube_converter.py b/ndcube/asdf/converters/ndcube_converter.py index 34d372019..fe87da261 100644 --- a/ndcube/asdf/converters/ndcube_converter.py +++ b/ndcube/asdf/converters/ndcube_converter.py @@ -35,9 +35,14 @@ def to_yaml_tree(self, ndcube, tag, ctx): This ensures that users are aware of potentially important information that is not included in the serialized output. """ + from astropy.wcs.wcsapi import HighLevelWCSWrapper + node = {} node["data"] = ndcube.data - node["wcs"] = ndcube.wcs + if isinstance(ndcube.wcs, HighLevelWCSWrapper): + node["wcs"] = ndcube.wcs._low_level_wcs + else: + node["wcs"] = ndcube.wcs node["extra_coords"] = ndcube.extra_coords node["global_coords"] = ndcube.global_coords node["meta"] = ndcube.meta diff --git a/ndcube/asdf/converters/reorderedwcs_converter.py b/ndcube/asdf/converters/reorderedwcs_converter.py new file mode 100644 index 000000000..df86a9804 --- /dev/null +++ b/ndcube/asdf/converters/reorderedwcs_converter.py @@ -0,0 +1,21 @@ +from asdf.extension import Converter + + +class ReorderedConverter(Converter): + tags = ["tag:sunpy.org:ndcube/reorderedwcs-0.1.0"] + types = ["ndcube.wcs.wrappers.reordered_wcs.ReorderedLowLevelWCS"] + + def from_yaml_tree(self, node, tag, ctx): + from ndcube.wcs.wrappers import ReorderedLowLevelWCS + + reorderedwcs = ReorderedLowLevelWCS(wcs=node["wcs"], + pixel_order = node.get("pixel_order"), + world_order = node.get("world_order") + ) + return reorderedwcs + def to_yaml_tree(self, reorderedwcs, tag, ctx): + node={} + node["wcs"] = reorderedwcs._wcs + node["pixel_order"] = (reorderedwcs._pixel_order) + node["world_order"] = (reorderedwcs._world_order) + return node diff --git a/ndcube/asdf/converters/resampled_converter.py b/ndcube/asdf/converters/resampled_converter.py new file mode 100644 index 000000000..aea976ceb --- /dev/null +++ b/ndcube/asdf/converters/resampled_converter.py @@ -0,0 +1,21 @@ +from asdf.extension import Converter + + +class ResampledConverter(Converter): + tags = ["tag:sunpy.org:ndcube/resampledwcs-0.1.0"] + types = ["ndcube.wcs.wrappers.resampled_wcs.ResampledLowLevelWCS"] + + def from_yaml_tree(self, node, tag, ctx): + from ndcube.wcs.wrappers import ResampledLowLevelWCS + + resampledwcs = ResampledLowLevelWCS(wcs=node["wcs"], + offset = node.get("offset"), + factor = node.get("factor"), + ) + return resampledwcs + def to_yaml_tree(self, resampledwcs, tag, ctx): + node={} + node["wcs"] = resampledwcs._wcs + node["factor"] = (resampledwcs._factor) + node["offset"] = (resampledwcs._offset) + return node diff --git a/ndcube/asdf/converters/tests/test_ndcube_wcs_wrappers.py b/ndcube/asdf/converters/tests/test_ndcube_wcs_wrappers.py new file mode 100644 index 000000000..fab66d381 --- /dev/null +++ b/ndcube/asdf/converters/tests/test_ndcube_wcs_wrappers.py @@ -0,0 +1,92 @@ +""" +Tests for roundtrip serialization of NDCube with various GWCS types. + +TODO: Add tests for the roundtrip serialization of NDCube with ResampledLowLevelWCS, ReorderedLowLevelWCS, and CompoundLowLevelWCS when using astropy.wcs.WCS. +""" + +import pytest +from gwcs import __version__ as gwcs_version +from packaging.version import Version + +import asdf + +from ndcube import NDCube +from ndcube.conftest import data_nd +from ndcube.tests.helpers import assert_cubes_equal +from ndcube.wcs.wrappers import CompoundLowLevelWCS, ReorderedLowLevelWCS, ResampledLowLevelWCS + + +@pytest.fixture +def create_ndcube_resampledwcs(gwcs_3d_lt_ln_l): + shape = (2, 3, 4) + new_wcs = ResampledLowLevelWCS(wcs = gwcs_3d_lt_ln_l, factor=2 ,offset = 1) + data = data_nd(shape) + return NDCube(data = data, wcs =new_wcs) + + +@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20") +def test_serialization_resampled(create_ndcube_resampledwcs, tmp_path): + ndc = create_ndcube_resampledwcs + file_path = tmp_path / "test.asdf" + with asdf.AsdfFile() as af: + af["ndcube"] = ndc + af.write_to(file_path) + + with asdf.open(file_path) as af: + loaded_ndcube = af["ndcube"] + + loaded_resampledwcs = loaded_ndcube.wcs.low_level_wcs + resampledwcs = ndc.wcs.low_level_wcs + assert (loaded_resampledwcs._factor == resampledwcs._factor).all() + assert (loaded_resampledwcs._offset == resampledwcs._offset).all() + + assert_cubes_equal(loaded_ndcube, ndc) + +@pytest.fixture +def create_ndcube_reorderedwcs(gwcs_3d_lt_ln_l): + shape = (2, 3, 4) + new_wcs = ReorderedLowLevelWCS(wcs = gwcs_3d_lt_ln_l, pixel_order=[1, 2, 0] ,world_order=[2, 0, 1]) + data = data_nd(shape) + return NDCube(data = data, wcs =new_wcs) + + + +@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20") +def test_serialization_reordered(create_ndcube_reorderedwcs, tmp_path): + ndc = create_ndcube_reorderedwcs + file_path = tmp_path / "test.asdf" + with asdf.AsdfFile() as af: + af["ndcube"] = ndc + af.write_to(file_path) + + with asdf.open(file_path) as af: + loaded_ndcube = af["ndcube"] + + loaded_reorderedwcs = loaded_ndcube.wcs.low_level_wcs + reorderedwcs = ndc.wcs.low_level_wcs + assert (loaded_reorderedwcs._pixel_order == reorderedwcs._pixel_order) + assert (loaded_reorderedwcs._world_order == reorderedwcs._world_order) + + assert_cubes_equal(loaded_ndcube, ndc) + +@pytest.fixture +def create_ndcube_compoundwcs(gwcs_2d_lt_ln, time_and_simple_extra_coords_2d): + + shape = (1, 2, 3, 4) + new_wcs = CompoundLowLevelWCS(gwcs_2d_lt_ln, time_and_simple_extra_coords_2d.wcs, mapping = [0, 1, 2, 3]) + data = data_nd(shape) + return NDCube(data = data, wcs = new_wcs) + +@pytest.mark.skipif(Version(gwcs_version) < Version("0.20"), reason="Requires gwcs>=0.20") +def test_serialization_compoundwcs(create_ndcube_compoundwcs, tmp_path): + ndc = create_ndcube_compoundwcs + file_path = tmp_path / "test.asdf" + with asdf.AsdfFile() as af: + af["ndcube"] = ndc + af.write_to(file_path) + + with asdf.open(file_path) as af: + loaded_ndcube = af["ndcube"] + assert_cubes_equal(loaded_ndcube, ndc) + assert (loaded_ndcube.wcs.low_level_wcs.mapping.mapping == ndc.wcs.low_level_wcs.mapping.mapping) + assert (loaded_ndcube.wcs.low_level_wcs.atol == ndc.wcs.low_level_wcs.atol) diff --git a/ndcube/asdf/entry_points.py b/ndcube/asdf/entry_points.py index 55462842e..0919b45f7 100644 --- a/ndcube/asdf/entry_points.py +++ b/ndcube/asdf/entry_points.py @@ -31,9 +31,12 @@ def get_extensions(): """ Get the list of extensions. """ + from ndcube.asdf.converters.compoundwcs_converter import CompoundConverter from ndcube.asdf.converters.extracoords_converter import ExtraCoordsConverter from ndcube.asdf.converters.globalcoords_converter import GlobalCoordsConverter from ndcube.asdf.converters.ndcube_converter import NDCubeConverter + from ndcube.asdf.converters.reorderedwcs_converter import ReorderedConverter + from ndcube.asdf.converters.resampled_converter import ResampledConverter from ndcube.asdf.converters.tablecoord_converter import ( QuantityTableCoordinateConverter, SkyCoordTableCoordinateConverter, @@ -47,6 +50,9 @@ def get_extensions(): QuantityTableCoordinateConverter(), SkyCoordTableCoordinateConverter(), GlobalCoordsConverter(), + ResampledConverter(), + ReorderedConverter(), + CompoundConverter(), ] _manifest_uri = "asdf://sunpy.org/ndcube/manifests/ndcube-0.1.0" diff --git a/ndcube/asdf/resources/manifests/ndcube-0.1.0.yaml b/ndcube/asdf/resources/manifests/ndcube-0.1.0.yaml index 4559a5499..d9e3a02bf 100644 --- a/ndcube/asdf/resources/manifests/ndcube-0.1.0.yaml +++ b/ndcube/asdf/resources/manifests/ndcube-0.1.0.yaml @@ -23,3 +23,12 @@ tags: - tag_uri: "tag:sunpy.org:ndcube/global_coords/globalcoords-0.1.0" schema_uri: "asdf://sunpy.org/ndcube/schemas/global_coords-0.1.0" + + - tag_uri: "tag:sunpy.org:ndcube/resampledwcs-0.1.0" + schema_uri: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0" + + - tag_uri: "tag:sunpy.org:ndcube/reorderedwcs-0.1.0" + schema_uri: "asdf://sunpy.org/ndcube/schemas/reorderedwcs-0.1.0" + + - tag_uri: "tag:sunpy.org:ndcube/compoundwcs-0.1.0" + schema_uri: "asdf://sunpy.org/ndcube/schemas/compoundwcs-0.1.0" diff --git a/ndcube/asdf/resources/schemas/compoundwcs-0.1.0.yaml b/ndcube/asdf/resources/schemas/compoundwcs-0.1.0.yaml new file mode 100644 index 000000000..24028920e --- /dev/null +++ b/ndcube/asdf/resources/schemas/compoundwcs-0.1.0.yaml @@ -0,0 +1,25 @@ +%YAML 1.1 +--- +$schema: "http://stsci.edu/schemas/yaml-schema/draft-01" +id: "asdf://sunpy.org/ndcube/schemas/Compoundwcs-0.1.0" + +title: + Represents the ndcube CompoundLowLevelWCS object + +description: + Represents the ndcube CompoundLowLevelWCS object + +type: object +properties: + wcs: + type: array + items: + tag: "tag:stsci.edu:gwcs/wcs-1.*" + mapping: + type: array + atol: + type: number + +required: [wcs] +additionalProperties: true +... diff --git a/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml b/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml index db10a488e..7491d715c 100644 --- a/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml +++ b/ndcube/asdf/resources/schemas/ndcube-0.1.0.yaml @@ -14,7 +14,11 @@ properties: data: description: "Must be compatible with ASDF serialization/deserialization and supported by NDCube." wcs: - tag: "tag:stsci.edu:gwcs/wcs-1.*" + anyOf: + - tag: "tag:stsci.edu:gwcs/wcs-1.*" + - tag: "tag:sunpy.org:ndcube/resampledwcs-0.1.0" + - tag: "tag:sunpy.org:ndcube/reorderedwcs-0.1.0" + - tag: "tag:sunpy.org:ndcube/compoundwcs-0.1.0" extra_coords: tag: "tag:sunpy.org:ndcube/extra_coords/extra_coords/extracoords-0.*" global_coords: diff --git a/ndcube/asdf/resources/schemas/reorderedwcs-0.1.0.yaml b/ndcube/asdf/resources/schemas/reorderedwcs-0.1.0.yaml new file mode 100644 index 000000000..76e54521b --- /dev/null +++ b/ndcube/asdf/resources/schemas/reorderedwcs-0.1.0.yaml @@ -0,0 +1,23 @@ +%YAML 1.1 +--- +$schema: "http://stsci.edu/schemas/yaml-schema/draft-01" +id: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0" + +title: + Represents the ndcube ReorderedLowLevelWCS object + +description: + Represents the ndcube ReorderedLowLevelWCS object + +type: object +properties: + wcs: + tag: "tag:stsci.edu:gwcs/wcs-1.*" + pixel_order: + type: array + world_order: + type: array + +required: [wcs] +additionalProperties: true +... diff --git a/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml b/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml new file mode 100644 index 000000000..504419de7 --- /dev/null +++ b/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml @@ -0,0 +1,23 @@ +%YAML 1.1 +--- +$schema: "http://stsci.edu/schemas/yaml-schema/draft-01" +id: "asdf://sunpy.org/ndcube/schemas/resampledwcs-0.1.0" + +title: + Represents the ndcube ResampledLowLevelWCS object + +description: + Represents the ndcube ResampledLowLevelWCS object + +type: object +properties: + wcs: + tag: "tag:stsci.edu:gwcs/wcs-1.*" + factor: + type: object + offset: + type: object + +required: [wcs] +additionalProperties: true +... From 1203516408051390a4b9f7e29c22e71b13cdfc0d Mon Sep 17 00:00:00 2001 From: ViciousEagle03 Date: Fri, 23 Aug 2024 22:29:51 +0530 Subject: [PATCH 2/4] Apply suggestions from code review --- ndcube/asdf/converters/reorderedwcs_converter.py | 4 ++-- ndcube/asdf/converters/resampled_converter.py | 5 +++-- ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/ndcube/asdf/converters/reorderedwcs_converter.py b/ndcube/asdf/converters/reorderedwcs_converter.py index df86a9804..4fb2264a3 100644 --- a/ndcube/asdf/converters/reorderedwcs_converter.py +++ b/ndcube/asdf/converters/reorderedwcs_converter.py @@ -16,6 +16,6 @@ def from_yaml_tree(self, node, tag, ctx): def to_yaml_tree(self, reorderedwcs, tag, ctx): node={} node["wcs"] = reorderedwcs._wcs - node["pixel_order"] = (reorderedwcs._pixel_order) - node["world_order"] = (reorderedwcs._world_order) + node["pixel_order"] = reorderedwcs._pixel_order + node["world_order"] = reorderedwcs._world_order return node diff --git a/ndcube/asdf/converters/resampled_converter.py b/ndcube/asdf/converters/resampled_converter.py index aea976ceb..da9545545 100644 --- a/ndcube/asdf/converters/resampled_converter.py +++ b/ndcube/asdf/converters/resampled_converter.py @@ -16,6 +16,7 @@ def from_yaml_tree(self, node, tag, ctx): def to_yaml_tree(self, resampledwcs, tag, ctx): node={} node["wcs"] = resampledwcs._wcs - node["factor"] = (resampledwcs._factor) - node["offset"] = (resampledwcs._offset) + node["factor"] = resampledwcs._factor + node["offset"] = resampledwcs._offset + return node diff --git a/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml b/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml index 504419de7..bfecfa023 100644 --- a/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml +++ b/ndcube/asdf/resources/schemas/resampledwcs-0.1.0.yaml @@ -14,9 +14,9 @@ properties: wcs: tag: "tag:stsci.edu:gwcs/wcs-1.*" factor: - type: object + tag: "tag:stsci.edu:asdf/core/ndarray-1.0.0" offset: - type: object + tag: "tag:stsci.edu:asdf/core/ndarray-1.0.0" required: [wcs] additionalProperties: true From a349d2062e585866248451817fb7316181783c43 Mon Sep 17 00:00:00 2001 From: ViciousEagle03 Date: Sun, 25 Aug 2024 16:06:22 +0530 Subject: [PATCH 3/4] Use Tabular1D, Tabular2D instead of tabular_model --- ndcube/extra_coords/table_coord.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ndcube/extra_coords/table_coord.py b/ndcube/extra_coords/table_coord.py index 503d6f1c7..297ff3ab4 100644 --- a/ndcube/extra_coords/table_coord.py +++ b/ndcube/extra_coords/table_coord.py @@ -10,8 +10,7 @@ import astropy.units as u from astropy.coordinates import SkyCoord from astropy.modeling import models -from astropy.modeling.models import tabular_model -from astropy.modeling.tabular import _Tabular +from astropy.modeling.tabular import Tabular1D, Tabular2D, _Tabular from astropy.time import Time from astropy.wcs.wcsapi.wrappers.sliced_wcs import combine_slices, sanitize_slices @@ -136,7 +135,10 @@ def _generate_tabular(lookup_table, interpolation='linear', points_unit=u.pix, * raise TypeError("lookup_table must be a Quantity.") # pragma: no cover ndim = lookup_table.ndim - TabularND = tabular_model(ndim, name=f"Tabular{ndim}D") + if ndim == 1: + TabularND = Tabular1D + elif ndim == 2: + TabularND = Tabular2D # The integer location is at the centre of the pixel. points = [(np.arange(size) - 0) * points_unit for size in lookup_table.shape] From ddebe6be206a6cc8799875dee8d83dd2ab68c14a Mon Sep 17 00:00:00 2001 From: ViciousEagle03 Date: Sun, 25 Aug 2024 17:29:53 +0530 Subject: [PATCH 4/4] Add comment and update logic --- ndcube/extra_coords/table_coord.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ndcube/extra_coords/table_coord.py b/ndcube/extra_coords/table_coord.py index 297ff3ab4..13bd0f25b 100644 --- a/ndcube/extra_coords/table_coord.py +++ b/ndcube/extra_coords/table_coord.py @@ -10,6 +10,7 @@ import astropy.units as u from astropy.coordinates import SkyCoord from astropy.modeling import models +from astropy.modeling.models import tabular_model from astropy.modeling.tabular import Tabular1D, Tabular2D, _Tabular from astropy.time import Time from astropy.wcs.wcsapi.wrappers.sliced_wcs import combine_slices, sanitize_slices @@ -135,10 +136,17 @@ def _generate_tabular(lookup_table, interpolation='linear', points_unit=u.pix, * raise TypeError("lookup_table must be a Quantity.") # pragma: no cover ndim = lookup_table.ndim + + # Use existing Tabular1D and Tabular2D classes for 1D and 2D models + # to ensure compatibility with asdf-astropy converters and avoid + # dynamically generated classes that are not recognized by asdf. + # See PR #751 for details on the issue this addresses. if ndim == 1: TabularND = Tabular1D elif ndim == 2: TabularND = Tabular2D + else: + TabularND = tabular_model(ndim, name=f"Tabular{ndim}D") # The integer location is at the centre of the pixel. points = [(np.arange(size) - 0) * points_unit for size in lookup_table.shape]