Adding surface serialization

src/simweights/_generation_surface.py

@@ -3,14 +3,14 @@
 # SPDX-License-Identifier: BSD-2-Clause
 
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Self
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
 from simweights._pdgcode import PDGCode
-from simweights._powerlaw import PowerLaw
-from simweights._spatial import SpatialDist
+from simweights._powerlaw import PowerLaw, resolve_powerlaw
+from simweights._spatial import SpatialDist, resolve_spatial
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -69,6 +69,61 @@ def get_epdf(self, weight_cols: "Mapping[str, NDArray[np.float64]]") -> NDArray[
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.pdgid.name}, {self.nevents}, {self.power_law}, {self.spatial})"
 
+    def to_dict(self) -> dict[str, Any]:
+        # json safe state
+        return {
+            "pdgid": int(self.pdgid.value),
+            "nevents": float(self.nevents),
+            "power_law": {"cls": type(self.power_law).__name__, "params": self.power_law.to_dict()},
+            "spatial": {"cls": type(self.spatial).__name__, "params": self.spatial.to_dict()},
+        }
+
+    @classmethod
+    def from_dict(cls, state: "Mapping[str, Any]") -> Self:
+        # ensure required params are included
+        # need to check explicitly as we have to rebuild powerlaw and spatial objects before initializing
+        required = ("power_law", "spatial", "pdgid", "nevents")
+        missing = [param for param in required if param not in state]
+        if missing:
+            raise TypeError(f"{cls.__name__}.from_dict: missing required keys {missing}, got {sorted(state)}")
+
+        # ensure nevents is a float or int
+        nevents = state["nevents"]
+        if isinstance(nevents, bool) or not isinstance(nevents, (int, float)):
+            raise TypeError(f"{cls.__name__}.from_dict: 'nevents' must be a number, got {type(nevents).__name__}")
+
+        # ensure pdgid is an int (enumification validates the int is valid later)
+        pdgid = state["pdgid"]
+        if isinstance(pdgid, bool) or not isinstance(pdgid, int):
+            raise TypeError(f"{cls.__name__}.from_dict: 'pdgid' must be an int, got {type(pdgid).__name__}")
+
+        # reconstruct powerlaw and spatial objects
+        rebuilt_state = dict(state)
+        for p, resolve in (("power_law", resolve_powerlaw), ("spatial", resolve_spatial)):
+            # ensure value is a dict
+            sub = state[p]
+            if not isinstance(sub, dict):
+                raise TypeError(f"{cls.__name__}.from_dict: '{p}' must be a dict, got {type(sub).__name__}")
+            if set(sub) != {"cls", "params"}:
+                raise TypeError(f"{cls.__name__}.from_dict: '{p}' must have keys 'cls' and 'params', got {sorted(sub)}")
+
+            # make sure class name is a str
+            name = sub["cls"]
+            if not isinstance(name, str):
+                raise TypeError(f"{cls.__name__}.from_dict: '{p}.cls' must be a str, got {type(name).__name__}")
+
+            # make sure params is a dict
+            params = sub["params"]
+            if not isinstance(params, dict):
+                raise TypeError(f"{cls.__name__}.from_dict: '{p}.params' must be a dict, got {type(params).__name__}")
+
+            # resolver rejects unknown names
+            # class itself validates params
+            rebuilt_state[p] = resolve(name).from_dict(params)
+
+        # rely on init to validate rest
+        return cls(**rebuilt_state)
+
 
 class CompositeSurface:
     """Represents two or more surface on which Monte Carlo simulation was generated on.
@@ -186,3 +241,35 @@ def __str__(self) -> str:
 
     def __repr__(self) -> str:
         return self.__class__.__name__ + "(\n  " + ",\n  ".join(repr(y) for x in self.components.values() for y in x) + ",\n)"
+
+    def to_dict(self) -> dict[str, Any]:
+        # store flattened list of serialized surfaces
+        # init will rebuild
+        return {"components": [s.to_dict() for lst in self.components.values() for s in lst]}
+
+    @classmethod
+    def from_dict(cls, state: dict[str, Any]) -> Self:
+        # ensure all required keys exist
+        required = ("components",)
+        missing = [param for param in required if param not in state]
+        if missing:
+            raise TypeError(f"{cls.__name__}.from_dict: missing required keys {missing}, got {sorted(state)}")
+
+        # ensure components is a list
+        components = state["components"]
+        if not isinstance(components, list):
+            raise TypeError(f"{cls.__name__}.from_dict: 'components' must be a list, got {type(components).__name__}")
+
+        # rebuild each surface
+        surfaces = []
+        for i, surface_dict in enumerate(state["components"]):
+            # ensure surface_dict is a dict
+            if not isinstance(surface_dict, dict):
+                raise TypeError(
+                    f"{cls.__name__}.from_dict: 'components' must be a list of dicts, got {type(surface_dict).__name__} at index {i}"
+                )
+
+            # class itself validates surface_dict
+            surfaces.append(GenerationSurface.from_dict(surface_dict))
+
+        return cls(*surfaces)

src/simweights/_powerlaw.py

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Self
 
 import numpy as np
 
@@ -123,3 +123,30 @@ def __eq__(self: PowerLaw, other: object) -> bool:
             mesg = f"{self} cannot be compared to {other}"
             raise TypeError(mesg)
         return self.g == other.g and self.a == other.a and self.b == other.b
+
+    def to_dict(self: PowerLaw) -> dict[str, float]:
+        # json safe state
+        return {param: float(getattr(self, param)) for param in ("g", "a", "b")}
+
+    @classmethod
+    def from_dict(cls: type[PowerLaw], state: dict[str, float]) -> Self:
+        # ensure correct types
+        for k, v in state.items():
+            if isinstance(v, bool) or not isinstance(v, (float, int)):
+                raise TypeError(f"{cls.__name__}.from_dict: '{k}' must be a number, got {type(v).__name__}")
+
+        # rely on init to validate the rest
+        return cls(**state)
+
+
+# although only one power law class, just adding so if more get added later
+# backwards compatibility wont be a problem
+_POWERLAW_CLASSES = {cls.__name__: cls for cls in (PowerLaw,)}
+
+
+def resolve_powerlaw(name: str) -> type[PowerLaw]:
+    """Resolve a powerlaw class object from its name."""
+    if name not in _POWERLAW_CLASSES:
+        raise ValueError(f"resolve_powerlaw: unknown power law class {name!r}, expected one of {sorted(_POWERLAW_CLASSES)}")
+
+    return _POWERLAW_CLASSES[name]

src/simweights/_spatial.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: BSD-2-Clause
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
 
 import numpy as np
 
@@ -73,6 +73,20 @@ def __eq__(self: CylinderBase, other: object) -> bool:
             and self.cos_zen_max == other.cos_zen_max
         )
 
+    def to_dict(self: CylinderBase) -> dict[str, float]:
+        # json safe state
+        return {param: float(getattr(self, param)) for param in ("length", "radius", "cos_zen_min", "cos_zen_max")}
+
+    @classmethod
+    def from_dict(cls: type[CylinderBase], state: dict[str, int | float]) -> Self:
+        # ensure correct types
+        for k, v in state.items():
+            if isinstance(v, bool) or not isinstance(v, (float, int)):
+                raise TypeError(f"{cls.__name__}.from_dict: '{k}' must be a number, got {type(v).__name__}")
+
+        # rely on init to validate the rest
+        return cls(**state)
+
 
 class UniformSolidAngleCylinder(CylinderBase):
     r"""Events are generated uniformly on the surface of a sphere.
@@ -172,5 +186,32 @@ def __eq__(self: CircleInjector, other: object) -> bool:
             and self.cos_zen_max == other.cos_zen_max
         )
 
+    def to_dict(self: CircleInjector) -> dict[str, float]:
+        # json safe state
+        return {param: float(getattr(self, param)) for param in ("radius", "cos_zen_min", "cos_zen_max")}
+
+    @classmethod
+    def from_dict(cls: type[CircleInjector], state: dict[str, float]) -> Self:
+        # ensure correct types
+        for k, v in state.items():
+            if isinstance(v, bool) or not isinstance(v, (float, int)):
+                raise TypeError(f"{cls.__name__}.from_dict: '{k}' must be a number, got {type(v).__name__}")
+
+        # rely on init to validate the rest
+        return cls(**state)
+
 
 SpatialDist = CylinderBase | CircleInjector
+
+
+_SPATIAL_CLASSES = {cls.__name__: cls for cls in (CylinderBase, UniformSolidAngleCylinder, NaturalRateCylinder, CircleInjector)}
+
+
+def resolve_spatial(name: str) -> type[SpatialDist]:
+    """Resolve a spatial distribution class object from its name."""
+    if name not in _SPATIAL_CLASSES:
+        raise ValueError(
+            f"resolve_spatial: unknown spatial distribution class {name!r}, expected one of {sorted(_SPATIAL_CLASSES)}"
+        )
+
+    return _SPATIAL_CLASSES[name]

tests/test_generation_surface.py

@@ -4,6 +4,7 @@
 #
 # SPDX-License-Identifier: BSD-2-Clause
 
+import json
 import unittest
 from copy import deepcopy
 
@@ -264,6 +265,95 @@ def test_repr_gsc(self):
         self.assertEqual(eval("".join(s[4].split()[-7:-4])[:-1]), self.p1)
         self.assertEqual(s[5], ">")
 
+    def check_surface_round_trip(self, s):
+        state = s.to_dict()
+
+        # json safe
+        self.assertEqual(json.loads(json.dumps(state)), state)
+
+        # round trip
+        rebuilt = GenerationSurface.from_dict(state)
+        self.assertEqual(rebuilt, s)
+
+        # from_dict shouldnt mutate the state it was handed
+        self.assertEqual(state, s.to_dict())
+
+    def test_surface_to_from_dict(self):
+        for s in (self.s0, self.s1, self.s2, self.s3, self.s4):
+            self.check_surface_round_trip(s)
+
+    def test_composite_to_from_dict(self):
+        for c in (self.gsc1, self.gsc2, self.gsc3, self.gsc4, CompositeSurface()):
+            state = c.to_dict()
+
+            # json safe
+            self.assertEqual(json.loads(json.dumps(state)), state)
+
+            # round trip
+            self.assertEqual(CompositeSurface.from_dict(state), c)
+
+        # merged nevents survive the flatten and rebuild
+        rebuilt = CompositeSurface.from_dict(self.gsc1.to_dict())
+        self.assertEqual(len(rebuilt.components[2212]), 1)
+        self.assertEqual(rebuilt.components[2212][0].nevents, 30000)
+
+    def test_surface_from_dict_errors(self):
+        state = self.s0.to_dict()
+
+        for key in ("pdgid", "nevents", "power_law", "spatial"):
+            with self.assertRaises(TypeError):
+                GenerationSurface.from_dict({k: v for k, v in state.items() if k != key})
+
+        for bad in ("2212", None, True, 2212.0, [2212]):
+            with self.assertRaises(TypeError):
+                GenerationSurface.from_dict({**state, "pdgid": bad})
+
+        for bad in ("10000", None, True, [10000]):
+            with self.assertRaises(TypeError):
+                GenerationSurface.from_dict({**state, "nevents": bad})
+
+        # correct type but invalid particle
+        with self.assertRaises(ValueError):
+            GenerationSurface.from_dict({**state, "pdgid": 999999})
+
+        for p in ("power_law", "spatial"):
+            for bad in (
+                None,
+                [],
+                "PowerLaw",
+                {},
+                {"cls": "PowerLaw"},
+                {"cls": 1, "params": {}},
+                {"cls": "PowerLaw", "params": None},
+                {**state[p], "extra": 1},
+            ):
+                with self.assertRaises(TypeError):
+                    GenerationSurface.from_dict({**state, p: bad})
+
+            # resolve reject unknown names
+            with self.assertRaises(ValueError):
+                GenerationSurface.from_dict({**state, p: {**state[p], "cls": "Bogus"}})
+
+        # init still validates own params
+        with self.assertRaises(ValueError):
+            GenerationSurface.from_dict(
+                {**state, "spatial": {**state["spatial"], "params": {**state["spatial"]["params"], "cos_zen_min": 2.0}}},
+            )
+
+    def test_composite_from_dict_errors(self):
+        state = self.gsc1.to_dict()
+
+        with self.assertRaises(TypeError):
+            CompositeSurface.from_dict({})
+
+        for bad in (None, 47, {}, "components"):
+            with self.assertRaises(TypeError):
+                CompositeSurface.from_dict({"components": bad})
+
+        for bad in (None, 47, "surface", []):
+            with self.assertRaises(TypeError):
+                CompositeSurface.from_dict({"components": [*state["components"], bad]})
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_powerlaw.py

@@ -4,13 +4,15 @@
 #
 # SPDX-License-Identifier: BSD-2-Clause
 
+import json
 import unittest
 
 import numpy as np
 from scipy import stats
 from scipy.integrate import quad
 
 from simweights import PowerLaw
+from simweights._powerlaw import resolve_powerlaw
 
 
 class TestPowerLaw(unittest.TestCase):
@@ -146,6 +148,44 @@ def test_raises(self):
         with self.assertRaises(TypeError):
             p == np.array([])  # noqa: B015
 
+    def check_round_trip(self, p):
+        state = p.to_dict()
+
+        # json safe
+        self.assertEqual(json.loads(json.dumps(state)), state)
+
+        # round trip
+        self.assertEqual(type(p).from_dict(state), p)
+
+        # class name round trips through resolve
+        self.assertIs(resolve_powerlaw(type(p).__name__), type(p))
+        self.assertEqual(resolve_powerlaw(type(p).__name__).from_dict(state), p)
+
+        # extra params raise
+        with self.assertRaises(TypeError):
+            type(p).from_dict({**state, "bogus": 1})
+
+        # missing params raise
+        random_key = next(iter(state))
+        with self.assertRaises(TypeError):
+            type(p).from_dict({k: v for k, v in state.items() if k != random_key})
+
+        # non numeric values rejected
+        for bad in ("1.0", None, True, [1.0]):
+            with self.assertRaises(TypeError):
+                type(p).from_dict({**state, random_key: bad})
+
+    def test_resolve_powerlaw(self):
+        for cls in (PowerLaw,):
+            self.assertIs(resolve_powerlaw(cls.__name__), cls)
+
+        for bad in ("", "bogus", "np", "resolve_powerlaw"):
+            with self.assertRaises(ValueError):
+                resolve_powerlaw(bad)
+
+    def test_round_trip(self):
+        self.check_round_trip(PowerLaw(1, 1, 1000))
+
 
 if __name__ == "__main__":
     unittest.main()