cosmology <-> ASDF

asdf_astropy/converters/cosmology/cosmology.py

@@ -0,0 +1,33 @@
+from asdf.extension import Converter
+
+from ..utils import import_type
+
+
+class CosmologyConverter(Converter):
+
+    def __init__(self, tag, cosmology_type_name):
+        self._tag = tag
+        self._cosmology_type_name = cosmology_type_name
+        self._cosmology_type = None
+
+    @property
+    def tags(self):
+        return [self._tag]
+
+    @property
+    def types(self):
+        return [self._cosmology_type_name]
+
+    @property
+    def cosmology_type(self):
+        # Delay import until the cosmology class is needed to improve speed
+        # of loading the extension.
+        if self._cosmology_type is None:
+            self._cosmology_type = import_type(self._cosmology_type_name)
+        return self._cosmology_type
+
+    def to_yaml_tree(self, obj, tag, ctx):
+        return obj.to_format("mapping")
+
+    def from_yaml_tree(self, node, tag, ctx):
+        return self.cosmology_type.from_format(node, format="mapping")

asdf_astropy/converters/cosmology/tests/test_cosmology.py

@@ -0,0 +1,29 @@
+
+
+import pytest
+
+import asdf
+from astropy import units as u
+from astropy.cosmology import parameters, realizations
+
+
+# TODO! something that tests every cosmology class
+TEST_COSMOLOGIES = (getattr(realizations, n) for n in parameters.available)
+
+
+def assert_cosmo_equal(a, b):
+    __tracebackhide__ = True
+
+    assert a == b
+    assert a.meta == b.meta
+
+
+@pytest.mark.parametrize("cosmo", TEST_COSMOLOGIES)
+def test_serialization(cosmo, tmp_path):
+    file_path = tmp_path / "test.asdf"
+    with asdf.AsdfFile() as af:
+        af["cosmo"] = cosmo
+        af.write_to(file_path)
+
+    with asdf.open(file_path) as af:
+        assert_cosmo_equal(af["cosmo"], cosmo)

asdf_astropy/extensions.py

@@ -25,6 +25,8 @@
 from .converters.coordinates.sky_coord import SkyCoordConverter
 from .converters.coordinates.spectral_coord import SpectralCoordConverter
 
+from .converters.cosmology.cosmology import CosmologyConverter
+
 from .converters.time.time import TimeConverter
 from .converters.time.time_delta import TimeDeltaConverter
 
@@ -438,6 +440,39 @@
     TimeDeltaConverter(),
     AstropyTableConverter(),
     AstropyFitsConverter(),
+    # Cosmology
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/lambdacdm-*",
+        "astropy.cosmology.flrw.LambdaCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/flatlambdacdm-*",
+        "astropy.cosmology.flrw.FlatLambdaCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/wcdm-*",
+        "astropy.cosmology.flrw.wCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/flatwcdm-*",
+        "astropy.cosmology.flrw.FlatwCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/w0wacdm-*",
+        "astropy.cosmology.flrw.w0waCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/flatw0wacdm-*",
+        "astropy.cosmology.flrw.Flatw0waCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/wpwacdm-*",
+        "astropy.cosmology.flrw.wpwaCDM",
+    ),
+    CosmologyConverter(
+        "tag:astropy.org:astropy/cosmology/w0wzcdm-*",
+        "astropy.cosmology.flrw.w0wzCDM",
+    ),
 ]
 
 

asdf_astropy/resources/schemas/cosmology/basecosmology-1.0.0.yaml

@@ -0,0 +1,31 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/basecosmology-1.0.0"
+
+title: |
+  Represents a cosmology object from astropy
+
+description: |
+  This schema is designed to be extended by other schemas to restrict the
+  allowable parameters.
+
+type: object
+properties:
+  name:
+    description: |
+      The name of the cosmology.
+    type: string
+  parameters:
+    description: |
+      Parameters of the cosmology.
+    type: object
+  meta:
+    description: |
+      Additional free-form metadata about the cosmology.
+    type: object
+    default: {}
+
+additionalProperties: false
+required: [parameters]
+...

asdf_astropy/resources/schemas/cosmology/baseflrw-1.0.0.yaml

@@ -0,0 +1,32 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/baseflrw-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/baseflrw-1.0.0"
+
+title: |
+  Represents a FLRW cosmology object from astropy
+
+description: |
+  This schema is designed to be extended by other schemas to restrict the
+  allowable parameters.
+
+type: object
+properties:
+  name:
+    description: |
+      The name of the cosmology.
+    type: string
+  parameters:
+    description: |
+      Parameters of the cosmology.
+    type: object
+  meta:
+    description: |
+      Additional free-form metadata about the cosmology.
+    type: object
+    default: {}
+
+additionalProperties: false
+required: [parameters]
+...

asdf_astropy/resources/schemas/cosmology/flatlambdacdm-1.0.0.yaml

@@ -0,0 +1,30 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/flatlambdacdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/flatlambdacdm-1.0.0"
+
+title: |
+  Represents a FlatLambdaCDM cosmology object from astropy.
+
+allOf:
+  - $ref: lambdacdm-1.0.0
+  - properties:
+      parameters:
+        type: object
+        properties:
+          H0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Om0:
+            type: float
+          Tcmb0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Neff:
+            type: float
+          m_nu:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Ob0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+
+        required: [H0, Om0, Tcmb0, Neff, m_nu, Ob0]
+...

asdf_astropy/resources/schemas/cosmology/flatw0wacdm-1.0.0.yaml

@@ -0,0 +1,34 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/flatw0wacdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/flatw0wacdm-1.0.0"
+
+title: |
+  Represents a Flatw0waCDM cosmology object from astropy.
+
+allOf:
+  - $ref: baseflrw-1.0.0
+  - properties:
+        parameters:
+          type: object
+          properties:
+            H0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Om0:
+              type: float
+            Tcmb0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Neff:
+              type: float
+            m_nu:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Ob0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            w0:
+              type: float
+            wa:
+              type: float
+
+          required: [H0, Om0, Ode0, Tcmb0, Neff, m_nu, Ob0, w0, wa]
+...

asdf_astropy/resources/schemas/cosmology/flatwcdm-1.0.0.yaml

@@ -0,0 +1,32 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/flatwcdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/flatwcdm-1.0.0"
+
+title: |
+  Represents a wCDM cosmology object from astropy.
+
+allOf:
+  - $ref: baseflrw-1.0.0
+  - properties:
+        parameters:
+          type: object
+          properties:
+            H0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Om0:
+              type: float
+            Tcmb0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Neff:
+              type: float
+            m_nu:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Ob0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            w0:
+              type: float
+
+          required: [H0, Om0, Tcmb0, Neff, m_nu, Ob0, w0]
+...

asdf_astropy/resources/schemas/cosmology/flatwpwacdm-1.0.0.yaml

@@ -0,0 +1,36 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/flatwpwacdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/flatwpwacdm-1.0.0"
+
+title: |
+  Represents a FlatwpwaCDM cosmology object from astropy.
+
+allOf:
+  - $ref: baseflrw-1.0.0
+  - properties:
+      parameters:
+        type: object
+        properties:
+          H0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Om0:
+            type: float
+          Tcmb0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Neff:
+            type: float
+          m_nu:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Ob0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          wp:
+            type: float
+          wa:
+            type: float
+          zp:
+            type: float
+
+        required: [H0, Om0, Ode0, Tcmb0, Neff, m_nu, Ob0, wp, wa, zp]
+...

asdf_astropy/resources/schemas/cosmology/lambdacdm-1.0.0.yaml

@@ -0,0 +1,32 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/lambdacdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/lambdacdm-1.0.0"
+
+title: |
+  Represents a LambdaCDM cosmology object from astropy.
+
+allOf:
+  - $ref: baseflrw-1.0.0
+  - properties:
+      parameters:
+        type: object
+        properties:
+          H0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Om0:
+            type: float
+          Ode0:
+            type: float
+          Tcmb0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Neff:
+            type: float
+          m_nu:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+          Ob0:
+            $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+
+          required: [H0, Om0, Ode0, Tcmb0, Neff, m_nu, Ob0]
+...

asdf_astropy/resources/schemas/cosmology/w0wacdm-1.0.0.yaml

@@ -0,0 +1,36 @@
+%YAML 1.1
+---
+$schema: "http://stsci.edu/schemas/yaml-schema/draft-01"
+id: "http://astropy.org/schemas/astropy/cosmology/w0wacdm-1.0.0"
+tag: "tag:astropy.org:astropy/cosmology/w0wacdm-1.0.0"
+
+title: |
+  Represents a w0waCDM cosmology object from astropy.
+
+allOf:
+  - $ref: baseflrw-1.0.0
+  - properties:
+        parameters:
+          type: object
+          properties:
+            H0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Om0:
+              type: float
+            Ode0:
+              type: float
+            Tcmb0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Neff:
+              type: float
+            m_nu:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            Ob0:
+              $ref: "http://stsci.edu/schemas/asdf/unit/quantity-1.1.0"
+            w0:
+              type: float
+            wa:
+              type: float
+
+          required: [H0, Om0, Ode0, Tcmb0, Neff, m_nu, Ob0, w0, wa]
+...