diff --git a/docs/conf.py b/docs/conf.py index 442242f8..9cec0db1 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -65,6 +65,7 @@ def __get_newtypes(some_type: type) -> list: "NodeType": "typing.TypeVar", "ParticleWithSpin": ("obj", "qrules.particle.ParticleWithSpin"), "Path": "pathlib.Path", + "QNTransition": ("obj", "qrules.workflow.QNTransition"), "qrules.topology.EdgeType": "typing.TypeVar", "qrules.topology.NodeType": "typing.TypeVar", "Rule": ("obj", "qrules.argument_handling.Rule"), @@ -104,6 +105,7 @@ def __get_newtypes(some_type: type) -> list: "qrules.solving.GraphElementProperties": "obj", "qrules.solving.GraphSettings": "obj", "qrules.transition.StateTransition": "obj", + "qrules.workflow.QNTransition": "obj", } author = "" autodoc_default_options = { @@ -124,6 +126,7 @@ def __get_newtypes(some_type: type) -> list: "GraphElementProperties": "qrules.solving.GraphElementProperties", "GraphSettings": "qrules.solving.GraphSettings", "InitialFacts": "qrules.combinatorics.InitialFacts", + "QNTransition": "qrules.workflow.QNTransition", "StateTransition": "qrules.transition.StateTransition", } autodoc_typehints_format = "short" diff --git a/docs/usage.ipynb b/docs/usage.ipynb index b4cc7853..1a6f9dc3 100644 --- a/docs/usage.ipynb +++ b/docs/usage.ipynb @@ -203,6 +203,7 @@ "maxdepth: 2\n", "---\n", "usage/reaction\n", + "usage/qn-transitions\n", "usage/particle\n", "usage/visualize\n", "usage/conservation\n", diff --git a/docs/usage/qn-transitions.ipynb b/docs/usage/qn-transitions.ipynb new file mode 100644 index 00000000..02b5823c --- /dev/null +++ b/docs/usage/qn-transitions.ipynb @@ -0,0 +1,265 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Transitions without spin projections\n", + "\n", + ":::{autolink-concat}\n", + ":::" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The workflows described in {doc}`/usage/reaction` generate a transition for every allowed combination of **spin projections** of the initial and final state, because a helicity amplitude model requires each of those combinations. If you are only interested in which intermediate states and quantum numbers are allowed — for instance, which $J^{PC}$ resonances can appear in a Dalitz-plot decomposition — the spin projections merely multiply the number of {class}`.QNProblemSet`s that have to be solved. This page shows how to generate transitions directly at the $J^{P(C)}$ level with the {mod}`.workflow` module, which is considerably faster." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-cell" + ] + }, + "outputs": [], + "source": [ + "from IPython.display import Markdown\n", + "\n", + "import qrules.io\n", + "from qrules.quantum_numbers import EdgeQuantumNumbers\n", + "from qrules.workflow import create_qn_problem_sets, find_qn_transitions\n", + "\n", + "PDG = qrules.load_pdg()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Problem sets without spin projections" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As an example, take the reaction $J/\\psi \\to \\gamma\\pi^0\\pi^0$ with two $f_0$ resonances as allowed intermediate states. By default, {func}`.create_qn_problem_sets` expands the initial and final state over all combinations of their allowed spin projections:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expanded = create_qn_problem_sets(\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", + "sum(len(problems) for problems in expanded.problem_sets.values())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The $J/\\psi$ has three allowed spin projections and the photon has two, so every problem set appears in six spin-projection combinations. With :code:`spin_projections=False`, this Cartesian expansion is skipped altogether and the problem sets only differ in decay topology and interaction types:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "qn_problem_sets = create_qn_problem_sets(\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", + " spin_projections=False,\n", + ")\n", + "sum(len(problems) for problems in qn_problem_sets.problem_sets.values())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The resulting problem sets contain no {attr}`~.EdgeQuantumNumbers.spin_projection`, {attr}`~.NodeQuantumNumbers.l_projection`, or {attr}`~.NodeQuantumNumbers.s_projection` quantum numbers at all. They are identical to what {func}`.strip_spin_projections` produces from the expanded collection, but without ever generating the expansion." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Solve at the quantum-number level" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "{func}`.find_qn_transitions` solves these problem sets purely at the quantum-number level: no particle database is consulted for the intermediate states. It returns {obj}`.QNTransition`s, whose states and interactions are property maps of quantum numbers instead of {class}`.State` and {class}`.InteractionProperties` objects." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "qn_transitions = find_qn_transitions(qn_problem_sets)\n", + "len(qn_transitions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The transitions differ only in the quantum numbers of the intermediate state and the $LS$-couplings of the interaction nodes. They can be visualized with {func}`.asdot`, just like ordinary transitions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "source = qrules.io.asmermaid(qn_transitions[0], render_node=True, markdown=True)\n", + "Markdown(source)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Since the property maps can be inspected directly, it is easy to summarize for example the allowed $J^{PC}$ quantum numbers of the intermediate state:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "{\n", + " (\n", + " state[EdgeQuantumNumbers.spin_magnitude],\n", + " state[EdgeQuantumNumbers.parity],\n", + " state[EdgeQuantumNumbers.c_parity],\n", + " )\n", + " for transition in qn_transitions\n", + " for state in transition.intermediate_states.values()\n", + "}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Fewer combinatorics for larger reactions" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The spin-projection expansion multiplies the number of problem sets by $\\prod_i(2s_i+1)$ over all initial and final states $i$ (with the exception of massless states, which have no $0$ projection). Skipping it therefore matters most for many-body final states with spin. Take $J/\\psi \\to p\\bar p\\pi^0\\pi^0$, where the expansion factor is $3 \\times 2 \\times 2 = 12$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "expanded = create_qn_problem_sets(\n", + " initial_state=[\"J/psi(1S)\"],\n", + " final_state=[\"p\", \"p~\", \"pi0\", \"pi0\"],\n", + " particle_db=PDG,\n", + " allowed_intermediate_particles=[\"N(1440)\"],\n", + ")\n", + "sum(len(problems) for problems in expanded.problem_sets.values())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "qn_problem_sets = create_qn_problem_sets(\n", + " initial_state=[\"J/psi(1S)\"],\n", + " final_state=[\"p\", \"p~\", \"pi0\", \"pi0\"],\n", + " particle_db=PDG,\n", + " allowed_intermediate_particles=[\"N(1440)\"],\n", + " spin_projections=False,\n", + ")\n", + "sum(len(problems) for problems in qn_problem_sets.problem_sets.values())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The projection-free problem sets are not only fewer, but each of them is also cheaper to solve, because the spin projections do not appear as variables in the constraint problem:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "qn_transitions = find_qn_transitions(qn_problem_sets)\n", + "len(qn_transitions)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + ":::{seealso}\n", + "To generate particle-level transitions *with* spin projections — as required for a helicity amplitude model — use {func}`.find_solutions` or {func}`.generate_transitions` as described in {doc}`/usage/reaction`. If the problem sets have already been created with spin projections, they can still be reduced afterwards with {func}`.strip_spin_projections`.\n", + ":::" + ] + } + ], + "metadata": { + "colab": { + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/src/qrules/combinatorics.py b/src/qrules/combinatorics.py index 9d79e7d8..2e468e23 100644 --- a/src/qrules/combinatorics.py +++ b/src/qrules/combinatorics.py @@ -210,13 +210,27 @@ def create_initial_facts( initial_state: Sequence[StateDefinitionInput], final_state: Sequence[StateDefinitionInput], particle_db: ParticleCollection, + expand_spin_projections: bool = True, ) -> list[InitialFacts]: + """Attach the initial and final states to the external edges of a `.Topology`. + + By default, one `.InitialFacts` is created for every combination of allowed spin + projections of the initial and final state. With + :code:`expand_spin_projections=False`, this Cartesian expansion is skipped and a + single `.InitialFacts` without spin projections is returned, for solving at the + :math:`J^{P(C)}` level (see `.strip_spin_projections`). + """ states = __create_states_with_spin_projections( list(topology.incoming_edge_ids) + list(topology.outgoing_edge_ids), list(map(as_state_definition, initial_state)) + list(map(as_state_definition, final_state)), particle_db, ) + if not expand_spin_projections: + projection_free_states = { + edge_id: (particle_db[name], None) for edge_id, (name, _) in states.items() + } + return [MutableTransition(topology, projection_free_states)] # type: ignore[arg-type] spin_states = __generate_spin_combinations(states, particle_db) return [MutableTransition(topology, states) for states in spin_states] diff --git a/src/qrules/workflow.py b/src/qrules/workflow.py index 95e0e057..6469426b 100644 --- a/src/qrules/workflow.py +++ b/src/qrules/workflow.py @@ -436,6 +436,7 @@ def create_problem_sets( # ruff: ignore[too-many-positional-arguments] intermediate_particles: AllowedIntermediateParticles, topologies: Iterable[Topology], final_state_groupings: list[list[list[str]]] | None = None, + expand_spin_projections: bool = True, ) -> dict[float, list[ProblemSet]]: """Create a `.ProblemSet` collection over all topologies, grouped by strength.""" initial_state = list(map(as_state_definition, initial_state)) @@ -450,7 +451,11 @@ def create_problem_sets( # ruff: ignore[too-many-positional-arguments] final_state_groupings, ) for initial_facts in create_initial_facts( - permutation, initial_state, final_state, particle_db + permutation, + initial_state, + final_state, + particle_db, + expand_spin_projections, ) for settings in create_graph_settings( permutation, initial_facts, interaction_config, intermediate_particles @@ -720,6 +725,7 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] max_spin_magnitude: float = 2, final_state_groupings: list[list[list[str]]] | None = None, merge_spin_projections: bool = False, + spin_projections: bool = True, ) -> QNProblemSetCollection: """Create a `.QNProblemSet` collection for a reaction, grouped by strength. @@ -733,7 +739,16 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] initial and final state are merged into value ranges on a single `.QNProblemSet` (see `.merge_qn_problem_sets`), which reduces the number of problem sets and speeds up solving. + + With :code:`spin_projections=False`, the problem sets contain no spin projections + at all: the Cartesian expansion over spin-projection combinations is skipped + 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. """ + if not spin_projections and merge_spin_projections: + msg = "merge_spin_projections has no effect when spin_projections=False" + raise ValueError(msg) _validate_formalism(formalism) if particle_db is None: particle_db = load_pdg() @@ -766,6 +781,7 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] intermediate_particles, topologies, final_state_groupings, + expand_spin_projections=spin_projections, ) qn_problem_sets = _to_qn_problem_sets(problem_sets) if merge_spin_projections: @@ -773,12 +789,15 @@ def create_qn_problem_sets( # ruff: ignore[too-many-positional-arguments] strength: merge_qn_problem_sets(problems) for strength, problems in qn_problem_sets.items() } - return QNProblemSetCollection( + collection = QNProblemSetCollection( problem_sets=qn_problem_sets, intermediate_particles=intermediate_particles, final_state=list(map(as_state_definition, final_state)), formalism=formalism, ) + if not spin_projections: + return strip_spin_projections(collection) + return collection def _to_qn_problem_sets( @@ -875,6 +894,11 @@ def strip_spin_projections( identical — such as the expansion over all spin-projection combinations — are deduplicated. Rules that require spin projections are skipped and reported by the solver through the not-executed-rules mechanism. + + .. tip:: `create_qn_problem_sets` with :code:`spin_projections=False` produces the + same problem sets without generating the spin-projection expansion in the + first place, which is considerably faster for reactions with many spin-carrying + states. """ if isinstance(qn_problem_sets, QNProblemSetCollection): stripped_collection = copy(qn_problem_sets) diff --git a/tests/unit/test_combinatorics.py b/tests/unit/test_combinatorics.py index 5c8f6009..bcbdd29c 100644 --- a/tests/unit/test_combinatorics.py +++ b/tests/unit/test_combinatorics.py @@ -37,6 +37,25 @@ def test_create_initial_facts(three_body_decay, particle_database): assert initial_polarization in {-1, +1} +def test_create_initial_facts_without_spin_projections( + three_body_decay, particle_database +): + initial_facts = create_initial_facts( + three_body_decay, + initial_state=[("J/psi(1S)", [-1, +1])], + final_state=["gamma", "pi0", "pi0"], + particle_db=particle_database, + expand_spin_projections=False, + ) + assert len(initial_facts) == 1 + fact = initial_facts[0] + edge_ids = sorted(fact.states) + assert edge_ids == [-1, 0, 1, 2] + particle_names = [fact.states[i][0].name for i in edge_ids] + assert particle_names == ["J/psi(1S)", "gamma", "pi0", "pi0"] + assert all(projection is None for _, projection in fact.states.values()) + + def describe_generate_kinematic_permutations(): def it_groupings(three_body_decay: Topology): topology = three_body_decay diff --git a/tests/unit/test_workflow.py b/tests/unit/test_workflow.py index d68aad78..b2f6f06c 100644 --- a/tests/unit/test_workflow.py +++ b/tests/unit/test_workflow.py @@ -10,6 +10,7 @@ InteractionType, create_interaction_settings, ) +from qrules.solving import _create_merge_key from qrules.transition import ReactionInfo, SolvingMode from qrules.workflow import ( InteractionConfig, @@ -245,3 +246,37 @@ def test_projection_free_qn_transitions(): serialized = json.dumps(asdict(qn_transitions[0])) assert '"spin_projection"' not in serialized assert '"spin_magnitude"' in serialized + + unexpanded = create_qn_problem_sets( + initial_state=[("J/psi(1S)", [-1, 1])], + final_state=["gamma", "pi0", "pi0"], + particle_db=particle_db, + allowed_intermediate_particles=["f(0)(980)", "f(0)(1500)"], + interaction_config=InteractionConfig( + type_settings=create_interaction_settings( + "helicity", particle_db=particle_db, max_angular_momentum=2 + ), + allowed_types=[InteractionType.STRONG], + ), + spin_projections=False, + ) + assert _to_merge_keys(unexpanded) == _to_merge_keys(stripped) + assert find_qn_transitions(unexpanded) == qn_transitions + + +def _to_merge_keys(collection: QNProblemSetCollection) -> set[tuple]: + return { + (strength, _create_merge_key(problem_set, set())) + for strength, problem_sets in collection.problem_sets.items() + for problem_set in problem_sets + } + + +def test_incompatible_spin_projection_flags_raise(): + with pytest.raises(ValueError, match="merge_spin_projections has no effect"): + create_qn_problem_sets( + initial_state=["J/psi(1S)"], + final_state=["gamma", "pi0", "pi0"], + merge_spin_projections=True, + spin_projections=False, + )