diff --git a/docs/usage/qn-transitions.ipynb b/docs/usage/qn-transitions.ipynb index 02b5823c..9c2ffb31 100644 --- a/docs/usage/qn-transitions.ipynb +++ b/docs/usage/qn-transitions.ipynb @@ -34,7 +34,11 @@ "\n", "import qrules.io\n", "from qrules.quantum_numbers import EdgeQuantumNumbers\n", - "from qrules.workflow import create_qn_problem_sets, find_qn_transitions\n", + "from qrules.workflow import (\n", + " create_qn_problem_sets,\n", + " find_qn_transitions,\n", + " generate_qn_transitions,\n", + ")\n", "\n", "PDG = qrules.load_pdg()" ] @@ -163,6 +167,52 @@ "}" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Reaction-level interface" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The staged workflow above offers full control over the problem sets, but the default use-case is covered by a single call to {func}`.generate_qn_transitions`, the quantum-number-level counterpart of {func}`.generate_transitions`. It returns a {class}`.QNReactionInfo`, which resolves the initial and final states to {class}`.Particle` instances — they are fully determined by their PID — while the intermediate states remain quantum-number property maps:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "reaction = generate_qn_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", + " particle_db=PDG,\n", + " allowed_intermediate_particles=[\"f(0)(980)\", \"f(0)(1500)\"],\n", + ")\n", + "{i: particle.name for i, particle in reaction.final_state.items()}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The rendered transitions now label the initial and final states by particle name:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "source = qrules.io.asmermaid(reaction.transitions[0], render_node=True, markdown=True)\n", + "Markdown(source)" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -228,6 +278,104 @@ "len(qn_transitions)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Many-body reactions" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Solving at the quantum-number level makes reactions with four- and even five-body final states feasible, for which the workflows *with* spin projections take impractically long (see [ComPWA/qrules#27](https://github.com/ComPWA/qrules/issues/27)). Take $J/\\psi \\to K^+K^-\\pi^+\\pi^-$ with $\\phi(1020)$ and $\\rho(770)$ resonances, where {func}`.create_qn_problem_sets`' :code:`final_state_groupings` argument limits the subsystems to $\\phi\\to K^+K^-$ and $\\rho^0\\to\\pi^+\\pi^-$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "reaction_4body = generate_qn_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"K+\", \"K-\", \"pi+\", \"pi-\"],\n", + " particle_db=PDG,\n", + " allowed_intermediate_particles=[\"phi(1020)\", \"rho(770)\"],\n", + " allowed_interaction_types=\"strong\",\n", + " final_state_groupings=[[[\"K+\", \"K-\"], [\"pi+\", \"pi-\"]]],\n", + ")\n", + "len(reaction_4body.transitions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The transitions can be summarized by collapsing them per decay topology with {func}`.asdot`'s :code:`collapse=\"topology\"` option, where the intermediate edges then list the allowed states in $I^G(J^{PC})$ notation:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "Markdown(qrules.io.asmermaid(reaction_4body, collapse=\"topology\", markdown=True))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "and the intermediate states are limited to the following charges, isospins, spins, and parities:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "{\n", + " (\n", + " state[EdgeQuantumNumbers.charge],\n", + " state[EdgeQuantumNumbers.isospin_magnitude],\n", + " state[EdgeQuantumNumbers.spin_magnitude],\n", + " state[EdgeQuantumNumbers.parity],\n", + " )\n", + " for transition in reaction_4body.transitions\n", + " for state in transition.intermediate_states.values()\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Even a five-body final state remains tractable — here with an additional $\\pi^0$, which lets the number of allowed decay topologies and quantum-number combinations grow substantially:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "reaction_5body = generate_qn_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"K+\", \"K-\", \"pi+\", \"pi-\", \"pi0\"],\n", + " particle_db=PDG,\n", + " allowed_intermediate_particles=[\"phi(1020)\", \"rho(770)\"],\n", + " allowed_interaction_types=\"strong\",\n", + " final_state_groupings=[[[\"K+\", \"K-\"], [\"pi+\", \"pi-\"]]],\n", + ")\n", + "f\"{len(reaction_5body.transitions):,} transitions over {len(reaction_5body.group_by_topology())} topologies\"" + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/src/qrules/__init__.py b/src/qrules/__init__.py index 4972cdab..1f2c7db3 100644 --- a/src/qrules/__init__.py +++ b/src/qrules/__init__.py @@ -28,6 +28,7 @@ InitialFacts, StateDefinitionInput, create_initial_facts, + to_state_definitions, ) from qrules.combinatorics import StateDefinition as StateDefinition from qrules.conservation_rules import ( @@ -72,8 +73,6 @@ if TYPE_CHECKING: from collections.abc import Iterable, Sequence - from typing_extensions import TypeIs - def check_reaction_violations( # ruff: ignore[complex-structure, too-many-positional-arguments] initial_state: StateDefinitionInput | Sequence[StateDefinitionInput], @@ -120,7 +119,7 @@ def check_reaction_violations( # ruff: ignore[complex-structure, too-many-posit .. seealso:: :ref:`usage:Check allowed reactions` """ - initial_state_definitions = _to_state_definitions(initial_state) + initial_state_definitions = to_state_definitions(initial_state) if particle_db is None: particle_db = load_pdg() @@ -356,7 +355,7 @@ def generate_transitions( # ruff: ignore[too-many-positional-arguments] >>> len(reaction.group_by_topology()) 3 """ - initial_state_definitions = _to_state_definitions(initial_state) + initial_state_definitions = to_state_definitions(initial_state) stm = StateTransitionManager( initial_state=initial_state_definitions, final_state=final_state, @@ -382,20 +381,6 @@ def generate_transitions( # ruff: ignore[too-many-positional-arguments] return stm.find_solutions(problem_sets) -def _to_state_definitions( - state: StateDefinitionInput | Sequence[StateDefinitionInput], / -) -> list[StateDefinitionInput]: - if _is_state_definition_input(state): - return [state] - return list(state) - - -def _is_state_definition_input(value: object, /) -> TypeIs[StateDefinitionInput]: - return isinstance(value, str) or ( - isinstance(value, tuple) and len(value) == 2 and isinstance(value[0], str) - ) - - def load_default_particles() -> ParticleCollection: """Load the default particle list that comes with `qrules`. diff --git a/src/qrules/combinatorics.py b/src/qrules/combinatorics.py index 2e468e23..32194c77 100644 --- a/src/qrules/combinatorics.py +++ b/src/qrules/combinatorics.py @@ -46,6 +46,21 @@ def as_state_definition( raise ValueError(msg) +def to_state_definitions( + state: StateDefinitionInput | Sequence[StateDefinitionInput], / +) -> list[StateDefinitionInput]: + """Wrap a single state definition in a `list`, or copy a sequence of them.""" + if _is_state_definition_input(state): + return [state] + return list(state) + + +def _is_state_definition_input(value: object, /) -> TypeIs[StateDefinitionInput]: + return isinstance(value, str) or ( + isinstance(value, tuple) and len(value) == 2 and isinstance(value[0], str) + ) + + class _KinematicRepresentation: # ruff: ignore[eq-without-hash] def __init__( self, diff --git a/src/qrules/io/_dot.py b/src/qrules/io/_dot.py index a49a03ed..412a05e4 100644 --- a/src/qrules/io/_dot.py +++ b/src/qrules/io/_dot.py @@ -71,9 +71,11 @@ def _create_preface(self) -> list[str]: ] def _render(self, obj: Any) -> list[str]: + from qrules.workflow import QNReactionInfo # ruff: ignore[import-outside-top-level] + if isinstance(obj, QNResult): obj = obj.solutions - if isinstance(obj, ReactionInfo): + if isinstance(obj, (QNReactionInfo, ReactionInfo)): obj = obj.transitions if isinstance(obj, abc.Iterable): return self._render_multiple_transitions(obj) diff --git a/src/qrules/io/_labels.py b/src/qrules/io/_labels.py index de721df9..afb2c3aa 100644 --- a/src/qrules/io/_labels.py +++ b/src/qrules/io/_labels.py @@ -2,15 +2,17 @@ import logging import re +from collections import abc from fractions import Fraction from functools import singledispatch from inspect import isfunction from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias import attrs +from attrs import frozen from qrules.particle import Particle, ParticleWithSpin, Spin, _render_fraction -from qrules.quantum_numbers import InteractionProperties +from qrules.quantum_numbers import EdgeQuantumNumbers, InteractionProperties from qrules.solving import ( EdgeSettings, GraphEdgePropertyMap, @@ -21,7 +23,7 @@ from qrules.transition import ProblemSet, State if TYPE_CHECKING: - from collections.abc import Callable, Iterable + from collections.abc import Callable, Iterable, Mapping from typing_extensions import TypeIs @@ -203,6 +205,7 @@ def membership(self, key: str, domain: str) -> str: ... def particle(self, name: str, latex: str | None) -> str: ... def spin(self, magnitude: str, projection: str) -> str: ... def state(self, particle: str, projection: str) -> str: ... + def superscript(self, base: str, exponent: str) -> str: ... class _PlainFormatter: @@ -248,6 +251,10 @@ def spin(magnitude: str, projection: str) -> str: def state(particle: str, projection: str) -> str: return f"{particle}[{projection}]" + @staticmethod + def superscript(base: str, exponent: str) -> str: + return base + exponent.translate(_SUPERSCRIPT_SIGNS) + class _LatexFormatter: @staticmethod @@ -293,7 +300,14 @@ def spin(magnitude: str, projection: str) -> str: def state(particle: str, projection: str) -> str: return Rf"{particle}\left[{projection}\right]" + @staticmethod + def superscript(base: str, exponent: str) -> str: + if not exponent: + return base + return f"{base}^{{{exponent}}}" + +_SUPERSCRIPT_SIGNS = str.maketrans({"+": "⁺", "-": "⁻"}) _PLAIN_FORMATTER = _PlainFormatter() _LATEX_FORMATTER = _LatexFormatter() _PARTICLE_COLUMN_MAX_ROWS = 6 @@ -542,6 +556,95 @@ def __render_state(state: State, formatter: _LabelFormatter) -> str: return formatter.state(particle, spin_projection) +@frozen +class QuantumNumberSignature: + """PDG-style :math:`I^G(J^{PC})` summary of a quantum-number property map. + + `collapse_graphs` converts states that are quantum-number property maps to this + compact form, so that a collapsed edge lists signatures instead of complete maps. + + >>> from qrules.quantum_numbers import EdgeQuantumNumbers as EQN + >>> signature = QuantumNumberSignature.from_property_map({ + ... EQN.spin_magnitude: 1, + ... EQN.parity: -1, + ... EQN.c_parity: -1, + ... EQN.isospin_magnitude: 1, + ... EQN.g_parity: +1, + ... }) + >>> as_string(signature) + '1⁺(1⁻⁻)' + >>> as_latex(signature) + '1^{+}(1^{--})' + >>> as_string(QuantumNumberSignature.from_property_map({EQN.spin_magnitude: 0.5})) + '1/2' + """ + + spin_magnitude: Fraction | None = None + parity: int | None = None + c_parity: int | None = None + isospin_magnitude: Fraction | None = None + g_parity: int | None = None + + @classmethod + def from_property_map(cls, qn_map: Mapping[Any, Any]) -> QuantumNumberSignature: + return cls( + spin_magnitude=_to_optional_fraction( + qn_map.get(EdgeQuantumNumbers.spin_magnitude) + ), + parity=_to_optional_int(qn_map.get(EdgeQuantumNumbers.parity)), + c_parity=_to_optional_int(qn_map.get(EdgeQuantumNumbers.c_parity)), + isospin_magnitude=_to_optional_fraction( + qn_map.get(EdgeQuantumNumbers.isospin_magnitude) + ), + g_parity=_to_optional_int(qn_map.get(EdgeQuantumNumbers.g_parity)), + ) + + +def _to_optional_fraction(value: Any) -> Fraction | None: + if value is None: + return None + return Fraction(value) + + +def _to_optional_int(value: Any) -> int | None: + if value is None: + return None + return int(value) + + +@as_string.register(QuantumNumberSignature) +def _(signature: QuantumNumberSignature) -> str: + return __render_signature(signature, _PLAIN_FORMATTER) + + +@as_latex.register(QuantumNumberSignature) +def _(signature: QuantumNumberSignature) -> str: + return __render_signature(signature, _LATEX_FORMATTER) + + +def __render_signature( + signature: QuantumNumberSignature, formatter: _LabelFormatter +) -> str: + if signature.spin_magnitude is None: + spin = formatter.text("?") + else: + spin = formatter.fraction(signature.spin_magnitude) + jpc = formatter.superscript( + spin, __render_parity_signs(signature.parity, signature.c_parity) + ) + if signature.isospin_magnitude is None: + return jpc + ig = formatter.superscript( + formatter.fraction(signature.isospin_magnitude), + __render_parity_signs(signature.g_parity), + ) + return f"{ig}({jpc})" + + +def __render_parity_signs(*parities: int | None) -> str: + return "".join("+" if p > 0 else "-" for p in parities if p is not None) + + @as_string.register(tuple) def _(obj: tuple) -> str: return __render_tuple(obj, _PLAIN_FORMATTER) @@ -666,12 +769,16 @@ def collapse_graphs( def _strip_properties(state: Any) -> Any: if isinstance(state, State): return state.particle + if isinstance(state, abc.Mapping): + return QuantumNumberSignature.from_property_map(state) return state def _sorting_key(obj: Any) -> Any: if isinstance(obj, State): return obj.particle.name + if isinstance(obj, QuantumNumberSignature): + return as_string(obj) if isinstance(obj, str): return obj.lower() return obj diff --git a/src/qrules/io/_mermaid.py b/src/qrules/io/_mermaid.py index f4d9a7aa..c3a5d6b2 100644 --- a/src/qrules/io/_mermaid.py +++ b/src/qrules/io/_mermaid.py @@ -205,9 +205,11 @@ def __call__(self, obj: Any) -> str: return "\n".join(lines) + "\n" def _render(self, obj: Any) -> list[str]: + from qrules.workflow import QNReactionInfo # ruff: ignore[import-outside-top-level] + if isinstance(obj, QNResult): obj = obj.solutions - if isinstance(obj, ReactionInfo): + if isinstance(obj, (QNReactionInfo, ReactionInfo)): obj = obj.transitions if isinstance(obj, abc.Iterable): return self._render_multiple_transitions(obj) diff --git a/src/qrules/workflow.py b/src/qrules/workflow.py index 6469426b..f9387ffd 100644 --- a/src/qrules/workflow.py +++ b/src/qrules/workflow.py @@ -37,8 +37,9 @@ create_initial_facts, match_external_edges, permutate_topology_kinematically, + to_state_definitions, ) -from qrules.particle import ParticleCollection, load_pdg +from qrules.particle import Particle, ParticleCollection, load_pdg from qrules.quantum_numbers import ( EdgeQuantumNumbers, InteractionProperties, @@ -88,7 +89,7 @@ from collections.abc import Iterable, Sequence from qrules.combinatorics import InitialFacts, StateDefinition, StateDefinitionInput - from qrules.particle import Particle, ParticleWithSpin + from qrules.particle import ParticleWithSpin from qrules.quantum_numbers import EdgeQuantumNumber, NodeQuantumNumber from qrules.solving import ( EdgeSettings, @@ -718,6 +719,7 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] | str | None = None, interaction_config: InteractionConfig | None = None, + allowed_interaction_types: str | Iterable[str] | None = None, formalism: SpinFormalism = "helicity", topology_building: str = "isobar", mass_conservation_factor: float | None = 3.0, @@ -745,6 +747,10 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] entirely, so the problem sets can only be solved at the :math:`J^{P(C)}` level with `find_qn_transitions`. This is equivalent to passing the collection through `strip_spin_projections` afterwards, but much cheaper. + + The :code:`allowed_interaction_types` (e.g. :code:`"strong"` or + :code:`["em", "weak"]`) restrict the interaction types of the default or given + :code:`interaction_config`. """ if not spin_projections and merge_spin_projections: msg = "merge_spin_projections has no effect when spin_projections=False" @@ -770,6 +776,11 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] max_spin_magnitude=max_spin_magnitude, ) ) + if allowed_interaction_types is not None: + interaction_config = copy(interaction_config) # do not mutate the argument + interaction_config.set_allowed_interaction_types( + _parse_interaction_types(allowed_interaction_types) + ) intermediate_particles = _resolve_intermediate_particles( allowed_intermediate_particles, particle_db ) @@ -871,8 +882,117 @@ def find_solutions( # ruff: ignore[too-many-positional-arguments] return collect_reaction_info(results, final_state, formalism) -QNTransition = FrozenTransition[FrozenDict[Any, Any], FrozenDict[Any, Any]] -"""Transition whose states and interactions are quantum-number property maps.""" +QNTransition = FrozenTransition[Particle | FrozenDict[Any, Any], FrozenDict[Any, Any]] +"""Transition whose states and interactions are quantum-number property maps. + +The initial and final states are resolved to `.Particle` instances if a particle +database is provided (see `collect_qn_transitions`); intermediate states always remain +quantum-number property maps, as they represent a *set* of allowed states. +""" + + +def _to_qn_transition_tuple( + transitions: Iterable[QNTransition], +) -> tuple[QNTransition, ...]: + return tuple(transitions) + + +@implement_pretty_repr +@frozen +class QNReactionInfo: + """Ordered collection of `QNTransition` instances. + + The quantum-number-level counterpart of `.ReactionInfo`: the initial and final + states are `.Particle` instances, while each intermediate state is a property map + of the quantum numbers that any particle appearing there must carry. Create an + instance with `generate_qn_transitions`. + """ + + transitions: tuple[QNTransition, ...] = field(converter=_to_qn_transition_tuple) + initial_state: FrozenDict[int, Particle] = field(init=False, repr=False, eq=False) + final_state: FrozenDict[int, Particle] = field(init=False, repr=False, eq=False) + + def __attrs_post_init__(self) -> None: + if not self.transitions: + object.__setattr__(self, "initial_state", FrozenDict({})) + object.__setattr__(self, "final_state", FrozenDict({})) + return + transition = self.transitions[0] + external_states = { + **transition.initial_states, + **transition.final_states, + } + for edge_id, state in external_states.items(): + if not isinstance(state, Particle): + msg = ( + f"External state {edge_id} is of type {type(state).__name__}, not" + f" {Particle.__name__}. Solve with a particle database, e.g." + " through generate_qn_transitions()." + ) + raise TypeError(msg) + object.__setattr__(self, "initial_state", FrozenDict(transition.initial_states)) + object.__setattr__(self, "final_state", FrozenDict(transition.final_states)) + + def get_intermediate_quantum_numbers(self) -> list[FrozenDict[Any, Any]]: + """Collect the distinct quantum-number sets of the intermediate states.""" + unique: dict[FrozenDict[Any, Any], None] = {} + for transition in self.transitions: + for state in transition.intermediate_states.values(): + if not isinstance(state, Particle): + unique.setdefault(state, None) + return list(unique) + + def group_by_topology(self) -> dict[Topology, list[QNTransition]]: + groupings = defaultdict(list) + for transition in self.transitions: + groupings[transition.topology].append(transition) + return dict(groupings) + + +def generate_qn_transitions( # ruff: ignore[too-many-positional-arguments] + initial_state: StateDefinitionInput | Sequence[StateDefinitionInput], + final_state: Sequence[StateDefinitionInput], + particle_db: ParticleCollection | None = None, + allowed_intermediate_particles: Iterable[str] | str | None = None, + allowed_interaction_types: str | Iterable[str] | None = None, + formalism: SpinFormalism = "helicity", + mass_conservation_factor: float | None = 3.0, + max_angular_momentum: int = 1, + max_spin_magnitude: float = 2, + final_state_groupings: list[list[list[str]]] | None = None, + topology_building: str = "isobar", +) -> QNReactionInfo: + """Generate allowed transitions without spin projections. + + The quantum-number-level counterpart of `.generate_transitions`: chains + `create_qn_problem_sets` with :code:`spin_projections=False` and + `find_qn_transitions`, so that the reaction is solved at the :math:`J^{P(C)}` + level. Since the spin-projection combinatorics is avoided both in the problem sets + and in the constraint problem, this is feasible for reactions with many-body final + states, for which `.generate_transitions` would take impractically long. The + arguments mirror those of `.generate_transitions`. + """ + if particle_db is None: + particle_db = load_pdg() + qn_problem_sets = create_qn_problem_sets( + to_state_definitions(initial_state), + final_state, + particle_db=particle_db, + allowed_intermediate_particles=allowed_intermediate_particles, + allowed_interaction_types=allowed_interaction_types, + formalism=formalism, + topology_building=topology_building, + mass_conservation_factor=mass_conservation_factor, + max_angular_momentum=max_angular_momentum, + max_spin_magnitude=max_spin_magnitude, + final_state_groupings=final_state_groupings, + spin_projections=False, + ) + transitions = find_qn_transitions(qn_problem_sets, particle_db) + if not transitions: + msg = "No solutions were found" + raise RuntimeError(msg) + return QNReactionInfo(transitions) @overload @@ -931,6 +1051,7 @@ def _unique_problem_sets(problem_sets: Iterable[QNProblemSet]) -> list[QNProblem def find_qn_transitions( qn_problem_sets: QNProblemSetCollection | dict[float, list[QNProblemSet]], + particle_db: ParticleCollection | None = None, ) -> tuple[QNTransition, ...]: """Find allowed transitions purely at the quantum-number level. @@ -938,7 +1059,9 @@ def find_qn_transitions( states from a particle database: the intermediate states of the returned transitions carry exactly the quantum numbers that the problem sets declare as domains. Combine with `strip_spin_projections` to obtain :math:`J^{P(C)}`-level - transitions for e.g. a Dalitz-plot decomposition. + transitions for e.g. a Dalitz-plot decomposition. A :code:`particle_db` is only + used to resolve the initial and final states to `.Particle` instances (see + `collect_qn_transitions`). """ if isinstance(qn_problem_sets, QNProblemSetCollection): qn_problem_sets = qn_problem_sets.problem_sets @@ -949,27 +1072,36 @@ def find_qn_transitions( (qn_problem_set, solver.find_solutions(qn_problem_set)) for qn_problem_set in qn_problems ) - return collect_qn_transitions(qn_results) + return collect_qn_transitions(qn_results, particle_db) def collect_qn_transitions( qn_results: dict[float, list[tuple[QNProblemSet, QNResult]]], + particle_db: ParticleCollection | None = None, ) -> tuple[QNTransition, ...]: """Summarize solver results as unique quantum-number-level transitions. The counterpart of `collect_reaction_info` for a quantum-number-only workflow: no - particle database is consulted and no `.State` objects are created — each - transition carries exactly the quantum numbers that are known from the initial - facts or were solved for. + `.State` objects are created — each transition carries exactly the quantum numbers + that are known from the initial facts or were solved for. If a :code:`particle_db` + is given, the initial and final states — which are fully determined by their + `~.EdgeQuantumNumbers.pid` — are resolved to `.Particle` instances; + intermediate states always remain quantum-number property maps. """ transitions: dict[QNTransition, None] = {} for qn_result_pairs in qn_results.values(): for qn_problem_set, qn_result in qn_result_pairs: facts = qn_problem_set.initial_facts + topology = facts.topology + external_edge_ids = topology.incoming_edge_ids | topology.outgoing_edge_ids for solution in qn_result.solutions: - states = dict(solution.states) + states: dict[int, Any] = dict(solution.states) for edge_id, edge_facts in facts.states.items(): states[edge_id] = {**edge_facts, **states.get(edge_id, {})} + if particle_db is not None: + for edge_id in external_edge_ids: + pid = states[edge_id][EdgeQuantumNumbers.pid] + states[edge_id] = particle_db.find(int(pid)) interactions = dict(solution.interactions) for node_id, node_facts in facts.interactions.items(): interactions[node_id] = { @@ -978,7 +1110,10 @@ def collect_qn_transitions( } transition: QNTransition = FrozenTransition( solution.topology, - states={i: FrozenDict(m) for i, m in states.items()}, + states={ + i: s if isinstance(s, Particle) else FrozenDict(s) + for i, s in states.items() + }, interactions={i: FrozenDict(m) for i, m in interactions.items()}, ) transitions[transition] = None @@ -1013,6 +1148,14 @@ def _strip_spin(state_definition: Sequence[StateDefinition]) -> list[str]: return particle_names +def _parse_interaction_types( + description: str | Iterable[str], +) -> list[InteractionType]: + if isinstance(description, str): + return [InteractionType.from_str(description)] + return [InteractionType.from_str(i) for i in description] + + def _resolve_intermediate_particles( selection: AllowedIntermediateParticles | Iterable[str] | str | None, particle_db: ParticleCollection, diff --git a/tests/unit/test_workflow.py b/tests/unit/test_workflow.py index b2f6f06c..1a993dc9 100644 --- a/tests/unit/test_workflow.py +++ b/tests/unit/test_workflow.py @@ -2,7 +2,7 @@ import pytest -from qrules.io import asdict +from qrules.io import asdict, asdot, asmermaid from qrules.particle import ParticleCollection, load_pdg from qrules.quantum_numbers import EdgeQuantumNumbers from qrules.settings import ( @@ -15,10 +15,12 @@ from qrules.workflow import ( InteractionConfig, QNProblemSetCollection, + QNReactionInfo, create_qn_problem_sets, filter_intermediate_particles, find_qn_transitions, find_solutions, + generate_qn_transitions, strip_spin_projections, ) @@ -272,6 +274,58 @@ def _to_merge_keys(collection: QNProblemSetCollection) -> set[tuple]: } +def test_generate_qn_transitions(): + particle_db = load_pdg() + reaction = generate_qn_transitions( + initial_state="J/psi(1S)", + final_state=["gamma", "pi0", "pi0"], + particle_db=particle_db, + allowed_intermediate_particles=["f(0)(980)", "f(0)(1500)"], + allowed_interaction_types=["strong", "em"], + ) + assert isinstance(reaction, QNReactionInfo) + assert len(reaction.transitions) > 0 + assert {p.name for p in reaction.initial_state.values()} == {"J/psi(1S)"} + assert [p.name for _, p in sorted(reaction.final_state.items())] == [ + "gamma", + "pi0", + "pi0", + ] + for qn_set in reaction.get_intermediate_quantum_numbers(): + assert qn_set[EdgeQuantumNumbers.spin_magnitude] == 0 + assert qn_set[EdgeQuantumNumbers.parity] == +1 + assert qn_set[EdgeQuantumNumbers.c_parity] == +1 + assert len(reaction.group_by_topology()) == 1 + + dot = asdot(reaction) + assert dot.startswith("digraph {") + assert "J/psi(1S)" in dot + assert "spin_projection" not in dot.replace("isospin_projection", "") + + mermaid = asmermaid(reaction, latex=False) + assert mermaid.startswith("flowchart LR") + assert "J/psi(1S)" in mermaid + assert "spin_projection" not in mermaid.replace("isospin_projection", "") + collapsed_dot = asdot(reaction, collapse="topology") + assert "0⁺(0⁺⁺)" in collapsed_dot + collapsed_mermaid = asmermaid(reaction, collapse="topology") + assert "0^{+}(0^{++})" in collapsed_mermaid + + +def test_qn_reaction_info_requires_particle_states(): + particle_db = load_pdg() + qn_problem_sets = create_qn_problem_sets( + initial_state=["J/psi(1S)"], + final_state=["gamma", "pi0", "pi0"], + particle_db=particle_db, + allowed_intermediate_particles=["f(0)(980)"], + spin_projections=False, + ) + qn_transitions = find_qn_transitions(qn_problem_sets) + with pytest.raises(TypeError, match="is of type FrozenDict, not Particle"): + QNReactionInfo(qn_transitions) + + def test_incompatible_spin_projection_flags_raise(): with pytest.raises(ValueError, match="merge_spin_projections has no effect"): create_qn_problem_sets(