FEAT: add serialization workspace loader

Author: redeboer

src/ampform_dpd/io/serialization/__init__.py

@@ -7,3 +7,7 @@
     This module is in preview, see https://github.com/ComPWA/ampform-dpd/issues/133 for
     updates.
 """
+
+from ampform_dpd.io.serialization.workspace import Workspace, load_workspace
+
+__all__ = ["Workspace", "load_workspace"]

src/ampform_dpd/io/serialization/workspace.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import json
+from collections import Counter
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, cast
+
+from ampform_dpd.io.serialization.amplitude import formulate
+from ampform_dpd.io.serialization.dynamics import (
+    PropagatorDynamicsBuilder,
+    formulate_dynamics,
+    formulate_form_factor,
+)
+
+if TYPE_CHECKING:
+    from os import PathLike
+
+    from ampform_dpd import AmplitudeModel, DefinedExpression
+    from ampform_dpd.io.serialization.format import ModelDefinition
+
+
+@dataclass(frozen=True)
+class Workspace:
+    """Backend-independent formulation of a serialized amplitude model."""
+
+    definition: Mapping[str, Any]
+    distributions: Mapping[str, AmplitudeModel]
+    functions: Mapping[str, DefinedExpression]
+    kinematics: Mapping[str, Any]
+    reference_points: tuple[Mapping[str, Any], ...]
+    checksums: tuple[Mapping[str, Any], ...]
+
+
+def load_workspace(
+    source: str | PathLike[str] | Mapping[str, Any],
+    *,
+    builders: Mapping[str, PropagatorDynamicsBuilder] | None = None,
+) -> Workspace:
+    """Load and formulate every distribution in a serialized model."""
+    definition = _load_definition(source)
+    _raise_on_duplicate_names(definition)
+    distributions = {
+        distribution["name"]: formulate(
+            _select_distribution(definition, distribution),
+            additional_builders=dict(builders) if builders is not None else None,
+        )
+        for distribution in definition["distributions"]
+    }
+    functions = _formulate_functions(definition, builders)
+    kinematics = {
+        distribution["name"]: distribution["decay_description"]["kinematics"]
+        for distribution in definition["distributions"]
+    }
+    checksums = definition.get("misc", {}).get("amplitude_model_checksums", [])
+    return Workspace(
+        definition=_freeze(definition),
+        distributions=MappingProxyType(distributions),
+        functions=MappingProxyType(functions),
+        kinematics=_freeze(kinematics),
+        reference_points=tuple(
+            _freeze(point) for point in definition.get("parameter_points", [])
+        ),
+        checksums=tuple(_freeze(checksum) for checksum in checksums),
+    )
+
+
+def _load_definition(
+    source: str | PathLike[str] | Mapping[str, Any],
+) -> ModelDefinition:
+    if isinstance(source, Mapping):
+        return cast("ModelDefinition", dict(source))
+    with Path(source).open() as stream:
+        return cast("ModelDefinition", json.load(stream))
+
+
+def _raise_on_duplicate_names(definition: ModelDefinition) -> None:
+    for collection_name in ("distributions", "functions"):
+        names = [item["name"] for item in definition[collection_name]]
+        duplicates = sorted(name for name, count in Counter(names).items() if count > 1)
+        if duplicates:
+            msg = f"Duplicate {collection_name} names: {', '.join(duplicates)}"
+            raise ValueError(msg)
+
+
+def _select_distribution(
+    definition: ModelDefinition, distribution: Mapping[str, Any]
+) -> ModelDefinition:
+    selected = dict(definition)
+    selected["distributions"] = [dict(distribution)]
+    return cast("ModelDefinition", selected)
+
+
+def _formulate_functions(
+    definition: ModelDefinition,
+    builders: Mapping[str, PropagatorDynamicsBuilder] | None,
+) -> dict[str, DefinedExpression]:
+    formulated = {}
+    for function in definition["functions"]:
+        name = function["name"]
+        formulated[name] = _formulate_function(name, definition, builders)
+    return formulated
+
+
+def _formulate_function(
+    name: str,
+    definition: ModelDefinition,
+    builders: Mapping[str, PropagatorDynamicsBuilder] | None,
+) -> DefinedExpression:
+    for distribution in definition["distributions"]:
+        model = _select_distribution(definition, distribution)
+        for chain in distribution["decay_description"]["chains"]:
+            for propagator in chain["propagators"]:
+                if propagator.get("parametrization") == name:
+                    single_propagator_chain = dict(chain)
+                    single_propagator_chain["propagators"] = [propagator]
+                    return formulate_dynamics(
+                        cast("Any", single_propagator_chain),
+                        model,
+                        additional_definitions=(
+                            dict(builders) if builders is not None else None
+                        ),
+                    )
+            for vertex in chain["vertices"]:
+                if vertex.get("formfactor") == name:
+                    return formulate_form_factor(vertex, model)
+    function_type = next(
+        function["type"]
+        for function in definition["functions"]
+        if function["name"] == name
+    )
+    msg = (
+        f"Cannot formulate function {name!r} of type {function_type!r}: "
+        "it has no propagator or form-factor context"
+    )
+    raise NotImplementedError(msg)
+
+
+def _freeze(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return MappingProxyType({key: _freeze(item) for key, item in value.items()})
+    if isinstance(value, list):
+        return tuple(_freeze(item) for item in value)
+    return value

tests/io_serialization/test_workspace.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import json
+from copy import deepcopy
+from types import MappingProxyType
+from typing import TYPE_CHECKING
+
+import pytest
+
+from ampform_dpd.io.serialization import load_workspace
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from ampform_dpd.io.serialization.format import ModelDefinition
+
+
+@pytest.mark.parametrize("from_path", [False, True], ids=["mapping", "path"])
+def test_load_workspace(
+    model_definition: ModelDefinition, tmp_path: Path, from_path: bool
+):
+    source = model_definition
+    if from_path:
+        source = tmp_path / "model.json"
+        source.write_text(json.dumps(model_definition))
+    workspace = load_workspace(source)
+    assert tuple(workspace.distributions) == ("default_model",)
+    assert set(workspace.functions) == {
+        item["name"] for item in model_definition["functions"]
+    }
+    assert isinstance(workspace.definition, MappingProxyType)
+    with pytest.raises(TypeError):
+        workspace.distributions["new"] = workspace.distributions["default_model"]  # ty: ignore[invalid-assignment]
+
+
+def test_loads_multiple_distributions(model_definition: ModelDefinition):
+    definition = deepcopy(model_definition)
+    second = deepcopy(definition["distributions"][0])
+    second["name"] = "second"
+    definition["distributions"].append(second)
+    workspace = load_workspace(definition)
+    assert tuple(workspace.distributions) == ("default_model", "second")
+
+
+@pytest.mark.parametrize("collection", ["distributions", "functions"])
+def test_rejects_duplicate_names(model_definition: ModelDefinition, collection: str):
+    definition = deepcopy(model_definition)
+    if collection == "distributions":
+        definition["distributions"].append(deepcopy(definition["distributions"][0]))
+    else:
+        definition["functions"].append(deepcopy(definition["functions"][0]))
+    with pytest.raises(ValueError, match=rf"Duplicate {collection} names"):
+        load_workspace(definition)
+
+
+def test_reports_unsupported_unreferenced_function(model_definition: ModelDefinition):
+    definition = deepcopy(model_definition)
+    definition["functions"].append({"name": "orphan", "type": "Unknown"})
+    with pytest.raises(NotImplementedError, match=r"orphan.*Unknown"):
+        load_workspace(definition)