ENH: support more types in normalize_state_ids() (#166)

* DX: add test for `normalize_state_ids()` with `ProblemSet`
* ENH: support `normalize_state_ids()` with mutable transitions and `ProblemSet`

Author: redeboer

src/ampform_dpd/adapter/qrules.py

@@ -10,8 +10,14 @@
 import attrs
 import qrules
 from qrules.quantum_numbers import InteractionProperties
-from qrules.topology import EdgeType, FrozenTransition, NodeType
-from qrules.transition import ReactionInfo, StateTransition, Topology
+from qrules.topology import (
+    EdgeType,
+    FrozenTransition,
+    MutableTransition,
+    NodeType,
+    Transition,
+)
+from qrules.transition import ProblemSet, ReactionInfo, Topology
 
 from ampform_dpd.decay import (
     FinalStateID,
@@ -240,15 +246,28 @@ def _(obj: ReactionInfo) -> ReactionInfo:
     )
 
 
-@_impl_normalize_state_ids.register(FrozenTransition)  # type:ignore[attr-defined]
-def _(obj: StateTransition) -> StateTransition:
+_Transition = TypeVar("_Transition", FrozenTransition, MutableTransition)
+
+
+@_impl_normalize_state_ids.register(FrozenTransition)
+@_impl_normalize_state_ids.register(MutableTransition)
+def _(obj: _Transition) -> _Transition:
     return attrs.evolve(
         obj,
         topology=_impl_normalize_state_ids(obj.topology),
         states={new: obj.states[old] for new, old in enumerate(sorted(obj.states))},
     )
 
 
+@_impl_normalize_state_ids.register(ProblemSet)
+def _(obj: ProblemSet) -> ProblemSet:
+    return ProblemSet(
+        initial_facts=_impl_normalize_state_ids(obj.initial_facts),
+        solving_settings=_impl_normalize_state_ids(obj.solving_settings),
+        topology=_impl_normalize_state_ids(obj.topology),
+    )
+
+
 @_impl_normalize_state_ids.register(Topology)  # type:ignore[attr-defined]
 def _(obj: Topology) -> Topology:
     mapping = {old: new for new, old in enumerate(sorted(obj.edges))}
@@ -260,7 +279,15 @@ def _(obj: abc.Iterable[T]) -> list[T]:
     return [_impl_normalize_state_ids(x) for x in obj]
 
 
-T = TypeVar("T", ReactionInfo, StateTransition, Topology)
+T = TypeVar(
+    "T",
+    FrozenTransition,
+    MutableTransition,
+    ProblemSet,
+    ReactionInfo,
+    Topology,
+    Transition,
+)
 """Type variable for the input and output of :func:`normalize_state_ids`."""
 
 

tests/adapter/test_qrules.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, Callable, SupportsFloat
 
 import pytest
+from qrules import InteractionType, StateTransitionManager
 
 from ampform_dpd.adapter.qrules import (
     _convert_transition,
@@ -117,6 +118,21 @@ def test_normalize_state_ids_reaction(jpsi2pksigma_reaction: ReactionInfo):
             assert transition012.states[i] == transition123.states[i + 1]
 
 
+def test_normalize_state_ids_problem_set():
+    stm = StateTransitionManager(
+        initial_state=[("J/psi(1S)", [-1, +1])],
+        final_state=["K0", "Sigma+", "p~"],
+        allowed_intermediate_particles=["N(1700)", "Sigma(1750)"],
+        formalism="helicity",
+        mass_conservation_factor=0,
+    )
+    stm.set_allowed_interaction_types([InteractionType.STRONG, InteractionType.EM])
+    problem_sets = stm.create_problem_sets()
+    some_problem_set = normalize_state_ids(problem_sets[3600.0][0])
+    assert set(some_problem_set.initial_facts.initial_states) == {0}
+    assert set(some_problem_set.initial_facts.final_states) == {1, 2, 3}
+
+
 def test_permute_equal_final_states(
     a2pipipi_reaction: ReactionInfo,
     jpsi2pksigma_reaction: ReactionInfo,