From 22fec61d39e57ec79a97d1f12117b3ae60649b3b Mon Sep 17 00:00:00 2001 From: Taylor St Jean Date: Thu, 10 Sep 2026 17:16:48 -0400 Subject: [PATCH 1/2] added surface serialization --- src/simweights/_generation_surface.py | 100 +++++++++++++++++++++++++- src/simweights/_powerlaw.py | 30 +++++++- src/simweights/_spatial.py | 44 +++++++++++- tests/test_generation_surface.py | 82 +++++++++++++++++++++ tests/test_powerlaw.py | 41 +++++++++++ tests/test_spatial.py | 47 +++++++++++- 6 files changed, 338 insertions(+), 6 deletions(-) diff --git a/src/simweights/_generation_surface.py b/src/simweights/_generation_surface.py index a7df87d..ef962a3 100644 --- a/src/simweights/_generation_surface.py +++ b/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,67 @@ 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 +247,36 @@ 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) + diff --git a/src/simweights/_powerlaw.py b/src/simweights/_powerlaw.py index 285aa76..e7f0859 100644 --- a/src/simweights/_powerlaw.py +++ b/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,31 @@ 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] diff --git a/src/simweights/_spatial.py b/src/simweights/_spatial.py index b56f312..f4e11f9 100644 --- a/src/simweights/_spatial.py +++ b/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,33 @@ 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] diff --git a/tests/test_generation_surface.py b/tests/test_generation_surface.py index 115dc62..8a5061f 100755 --- a/tests/test_generation_surface.py +++ b/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,87 @@ 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() diff --git a/tests/test_powerlaw.py b/tests/test_powerlaw.py index 8f7e9a6..0c7adf9 100755 --- a/tests/test_powerlaw.py +++ b/tests/test_powerlaw.py @@ -4,6 +4,7 @@ # # SPDX-License-Identifier: BSD-2-Clause +import json import unittest import numpy as np @@ -11,6 +12,7 @@ from scipy.integrate import quad from simweights import PowerLaw +from simweights._powerlaw import resolve_powerlaw class TestPowerLaw(unittest.TestCase): @@ -146,6 +148,45 @@ 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() diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 22e04e5..341cbb4 100755 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -4,6 +4,7 @@ # # SPDX-License-Identifier: BSD-2-Clause +import json import unittest import numpy as np @@ -13,10 +14,46 @@ NaturalRateCylinder, UniformSolidAngleCylinder, ) -from simweights._spatial import CylinderBase +from simweights._spatial import CylinderBase, resolve_spatial class TestSpatial(unittest.TestCase): + + def check_round_trip(self, c): + state = c.to_dict() + + # json safe + self.assertEqual(json.loads(json.dumps(state)), state) + + # round trip + self.assertEqual(type(c).from_dict(state), c) + + # class name round trips through resolve + self.assertIs(resolve_spatial(type(c).__name__), type(c)) + self.assertEqual(resolve_spatial(type(c).__name__).from_dict(state), c) + + # extra params raise + with self.assertRaises(TypeError): + type(c).from_dict({**state, "bogus": 1}) + + # missing params raise + random_key = next(iter(state)) + with self.assertRaises(TypeError): + type(c).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(c).from_dict({**state, random_key: bad}) + + def test_resolve_spatial(self): + for cls in (CylinderBase, UniformSolidAngleCylinder, NaturalRateCylinder, CircleInjector): + self.assertIs(resolve_spatial(cls.__name__), cls) + + for bad in ("", "bogus", "np", "resolve_spatial"): + with self.assertRaises(ValueError): + resolve_spatial(bad) + def check_diff_etendue(self, c, le, r): le *= 1e2 r *= 1e2 @@ -91,6 +128,8 @@ def test_cylinder_base(self): c.pdf(0.5) self.assertEqual(c, c) + self.check_round_trip(c) + def test_natural_rate_cylinder(self): last_c1 = None for le in range(100, 1000, 300): @@ -99,6 +138,8 @@ def test_natural_rate_cylinder(self): self.check_diff_etendue(c1, le, r) self.check_pdf_etendue(c1, 2 * np.pi**2 * r * (r + le)) + self.check_round_trip(c1) + c2 = NaturalRateCylinder(le, r, -1, 0) self.check_pdf_etendue(c2, np.pi**2 * r * (r + le)) @@ -145,6 +186,8 @@ def test_uniform_solid_angle(self): self.check_diff_etendue(c1, le, r) self.check_uniform_pdf(c1, 4 * np.pi, np.pi / 2 * r * (r + le)) + self.check_round_trip(c1) + c2 = UniformSolidAngleCylinder(le, r, -1, 0) self.check_uniform_pdf(c2, 2 * np.pi, np.pi / 2 * r * (r + le)) @@ -189,6 +232,8 @@ def test_circle_injector(self): c1 = CircleInjector(r, -1, 1) self.check_circle(c1, 4e4 * np.pi**2 * r**2) + self.check_round_trip(c1) + c2 = CircleInjector(r, -1, 0) self.check_circle(c2, 2e4 * np.pi**2 * r**2) From 72d09b3f42417eaf570380442d8b835cce202fd4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:37:31 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/simweights/_generation_surface.py | 11 ++--------- src/simweights/_powerlaw.py | 5 ++--- src/simweights/_spatial.py | 5 ++--- tests/test_generation_surface.py | 12 ++++++++++-- tests/test_powerlaw.py | 1 - tests/test_spatial.py | 1 - 6 files changed, 16 insertions(+), 19 deletions(-) diff --git a/src/simweights/_generation_surface.py b/src/simweights/_generation_surface.py index ef962a3..7c043af 100644 --- a/src/simweights/_generation_surface.py +++ b/src/simweights/_generation_surface.py @@ -74,14 +74,8 @@ def to_dict(self) -> dict[str, Any]: 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() - } + "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 @@ -279,4 +273,3 @@ def from_dict(cls, state: dict[str, Any]) -> Self: surfaces.append(GenerationSurface.from_dict(surface_dict)) return cls(*surfaces) - diff --git a/src/simweights/_powerlaw.py b/src/simweights/_powerlaw.py index e7f0859..5bf6342 100644 --- a/src/simweights/_powerlaw.py +++ b/src/simweights/_powerlaw.py @@ -143,11 +143,10 @@ def from_dict(cls: type[PowerLaw], state: dict[str, float]) -> Self: # 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)}" - ) + raise ValueError(f"resolve_powerlaw: unknown power law class {name!r}, expected one of {sorted(_POWERLAW_CLASSES)}") return _POWERLAW_CLASSES[name] diff --git a/src/simweights/_spatial.py b/src/simweights/_spatial.py index f4e11f9..733d92f 100644 --- a/src/simweights/_spatial.py +++ b/src/simweights/_spatial.py @@ -204,9 +204,8 @@ def from_dict(cls: type[CircleInjector], state: dict[str, float]) -> Self: SpatialDist = CylinderBase | CircleInjector -_SPATIAL_CLASSES = { - cls.__name__: cls for cls in (CylinderBase, UniformSolidAngleCylinder, NaturalRateCylinder, 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.""" diff --git a/tests/test_generation_surface.py b/tests/test_generation_surface.py index 8a5061f..b120f61 100755 --- a/tests/test_generation_surface.py +++ b/tests/test_generation_surface.py @@ -317,8 +317,16 @@ def test_surface_from_dict_errors(self): 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}): + 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}) diff --git a/tests/test_powerlaw.py b/tests/test_powerlaw.py index 0c7adf9..2b85641 100755 --- a/tests/test_powerlaw.py +++ b/tests/test_powerlaw.py @@ -187,6 +187,5 @@ def test_round_trip(self): self.check_round_trip(PowerLaw(1, 1, 1000)) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 341cbb4..507715d 100755 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -18,7 +18,6 @@ class TestSpatial(unittest.TestCase): - def check_round_trip(self, c): state = c.to_dict()