diff --git a/docs/usage.ipynb b/docs/usage.ipynb index bebf67de..b4cc7853 100644 --- a/docs/usage.ipynb +++ b/docs/usage.ipynb @@ -64,7 +64,7 @@ "source": [ "from IPython.display import Markdown\n", "\n", - "source = qrules.io.asmermaid(reaction, collapse_graphs=True, markdown=True)\n", + "source = qrules.io.asmermaid(reaction, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] }, diff --git a/docs/usage/custom-topology.ipynb b/docs/usage/custom-topology.ipynb index ba27e2c6..1d27912b 100644 --- a/docs/usage/custom-topology.ipynb +++ b/docs/usage/custom-topology.ipynb @@ -153,7 +153,7 @@ }, "outputs": [], "source": [ - "source = qrules.io.asmermaid(reaction_kk, collapse_graphs=True, markdown=True)\n", + "source = qrules.io.asmermaid(reaction_kk, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] }, @@ -192,7 +192,7 @@ }, "outputs": [], "source": [ - "source = qrules.io.asmermaid(reaction_ep, collapse_graphs=True, markdown=True)\n", + "source = qrules.io.asmermaid(reaction_ep, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] }, @@ -231,7 +231,7 @@ }, "outputs": [], "source": [ - "source = qrules.io.asmermaid(reaction_ep_no_mass, collapse_graphs=True, markdown=True)\n", + "source = qrules.io.asmermaid(reaction_ep_no_mass, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] } diff --git a/docs/usage/ls-coupling.ipynb b/docs/usage/ls-coupling.ipynb index fe1cb927..e1b4e7e8 100644 --- a/docs/usage/ls-coupling.ipynb +++ b/docs/usage/ls-coupling.ipynb @@ -277,7 +277,7 @@ " allowed_interaction_types=\"strong\",\n", " max_angular_momentum=3,\n", ")\n", - "source = qrules.io.asmermaid(reaction, markdown=True, render_node=True, strip_spin=True)\n", + "source = qrules.io.asmermaid(reaction, collapse=\"spin\", markdown=True, render_node=True)\n", "Markdown(source)" ] } diff --git a/docs/usage/reaction.ipynb b/docs/usage/reaction.ipynb index ebf2a254..d61e9365 100644 --- a/docs/usage/reaction.ipynb +++ b/docs/usage/reaction.ipynb @@ -379,7 +379,7 @@ }, "outputs": [], "source": [ - "source = io.asmermaid(lc2pkpi_reaction, collapse_graphs=True, markdown=True)\n", + "source = io.asmermaid(lc2pkpi_reaction, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] }, @@ -480,7 +480,7 @@ }, "outputs": [], "source": [ - "source = io.asmermaid(reaction, collapse_graphs=True, markdown=True, render_node=False)\n", + "source = io.asmermaid(reaction, collapse=\"topology\", markdown=True, render_node=False)\n", "Markdown(source)" ] }, diff --git a/docs/usage/visualize.ipynb b/docs/usage/visualize.ipynb index dbf1d86a..2c0c3283 100644 --- a/docs/usage/visualize.ipynb +++ b/docs/usage/visualize.ipynb @@ -268,7 +268,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Since this list of all possible spin projections {attr}`~.ReactionInfo.transitions` is rather long, it is often useful to use `strip_spin=True` or `collapse_graphs=True` to bundle comparable graphs. First, {code}`strip_spin=True` allows one collapse (ignore) the spin projections (we again show a selection only):" + "Since this list of all possible spin projections {attr}`~.ReactionInfo.transitions` is rather long, it is often useful to collapse comparable graphs. First, {code}`collapse=\"spin\"` allows one to collapse (ignore) the spin projections (we again show a selection only):" ] }, { @@ -278,7 +278,7 @@ "outputs": [], "source": [ "first_transitions = reaction.transitions[:3]\n", - "source = qrules.io.asmermaid(first_transitions, markdown=True, strip_spin=True)\n", + "source = qrules.io.asmermaid(first_transitions, collapse=\"spin\", markdown=True)\n", "Markdown(source)" ] }, @@ -299,7 +299,7 @@ " first_transitions,\n", " markdown=True,\n", " render_node=True,\n", - " strip_spin=True,\n", + " collapse=\"spin\",\n", ")\n", "Markdown(source)" ] @@ -317,7 +317,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "If that list is still too much, there is {code}`collapse_graphs=True`, which bundles all graphs with the same final state groupings:" + "If that list is still too much, {code}`collapse=\"topology\"` bundles all graphs with the same final state groupings:" ] }, { @@ -327,7 +327,7 @@ "outputs": [], "source": [ "source = qrules.io.asmermaid(\n", - " reaction, collapse_graphs=True, markdown=True, render_node=False\n", + " reaction, collapse=\"topology\", markdown=True, render_node=False\n", ")\n", "Markdown(source)" ] @@ -345,7 +345,7 @@ "metadata": {}, "outputs": [], "source": [ - "dot = qrules.io.asdot(reaction, collapse_graphs=True, render_node=False)\n", + "dot = qrules.io.asdot(reaction, collapse=\"topology\", render_node=False)\n", "graphviz.Source(dot)" ] }, @@ -491,7 +491,7 @@ " )\n", " for t in reaction.transitions\n", "})\n", - "source = qrules.io.asmermaid(jpc_ig_transitions, collapse_graphs=True, markdown=True)\n", + "source = qrules.io.asmermaid(jpc_ig_transitions, collapse=\"topology\", markdown=True)\n", "Markdown(source)" ] }, diff --git a/src/qrules/io/__init__.py b/src/qrules/io/__init__.py index df57950d..6e176d33 100644 --- a/src/qrules/io/__init__.py +++ b/src/qrules/io/__init__.py @@ -14,10 +14,12 @@ import attrs import yaml -from qrules.io import _dict, _dot, _mermaid +from qrules.io import _dict, _dot, _labels, _mermaid from qrules.particle import Particle, ParticleCollection from qrules.topology import Topology +CollapseMode = _labels.CollapseMode + if TYPE_CHECKING: from _typeshed import StrPath @@ -63,8 +65,7 @@ def asdot( render_final_state_id: bool = True, render_resonance_id: bool = False, render_initial_state_id: bool = False, - strip_spin: bool = False, - collapse_graphs: bool = False, + collapse: CollapseMode | None = None, edge_style: dict[str, Any] | None = None, node_style: dict[str, Any] | None = None, **figure_style: Any, @@ -77,12 +78,11 @@ def asdot( Args: instance: the input `object` that is to be rendered as DOT (graphviz) language. - strip_spin: Normally, each `.MutableTransition` has a `.Particle` with a spin - projection on its edges. This option hides the projections, leaving only - `.Particle` names on edges. - - collapse_graphs: Group all transitions by equivalent kinematic topology - and combine all allowed particles on each edge. + collapse: Optionally combine transitions. With ``"spin"``, transitions that + differ only in their spin projections are combined. With ``"topology"``, + transitions with equivalent kinematic topologies are grouped and all allowed + particles on each edge are collected. Node properties cannot be rendered + when collapsing by topology. render_node: Whether or not to render node ID (in the case of a `.Topology`) and/or node properties (in the case of a `.MutableTransition`). Meaning of @@ -114,8 +114,7 @@ def asdot( render_final_state_id=render_final_state_id, render_resonance_id=render_resonance_id, render_initial_state_id=render_initial_state_id, - strip_spin=strip_spin, - collapse_graphs=collapse_graphs, + collapse=collapse, figure_style=figure_style, edge_style=edge_style, node_style=node_style, @@ -131,8 +130,7 @@ def asmermaid( render_final_state_id: bool = True, render_resonance_id: bool = False, render_initial_state_id: bool = False, - strip_spin: bool = False, - collapse_graphs: bool = False, + collapse: CollapseMode | None = None, figure_style: dict[str, Any] | None = None, edge_style: dict[str, Any] | None = None, node_style: dict[str, Any] | None = None, @@ -148,14 +146,12 @@ def asmermaid( instance: the input `object` that is to be rendered as Mermaid flowchart source. - strip_spin: Normally, each `.MutableTransition` has a `.Particle` with a spin - projection on its edges. This option hides the projections, leaving only - `.Particle` names on edges. - - collapse_graphs: Group all transitions by equivalent kinematic topology - and combine all allowed particles on each edge. With LaTeX rendering, - particle lists longer than six entries are arranged in columns of at most - six rows. + collapse: Optionally combine transitions. With ``"spin"``, transitions that + differ only in their spin projections are combined. With ``"topology"``, + transitions with equivalent kinematic topologies are grouped and all allowed + particles on each edge are collected. Node properties cannot be rendered + when collapsing by topology. With LaTeX rendering, particle lists longer + than six entries are arranged in columns of at most six rows. render_node: Whether or not to render node ID (in the case of a `.Topology`) and/or node properties (in the case of a `.MutableTransition`). Meaning of @@ -191,8 +187,7 @@ def asmermaid( render_final_state_id=render_final_state_id, render_resonance_id=render_resonance_id, render_initial_state_id=render_initial_state_id, - strip_spin=strip_spin, - collapse_graphs=collapse_graphs, + collapse=collapse, figure_style=figure_style, edge_style=edge_style, node_style=node_style, diff --git a/src/qrules/io/_dot.py b/src/qrules/io/_dot.py index cb0728b3..a49a03ed 100644 --- a/src/qrules/io/_dot.py +++ b/src/qrules/io/_dot.py @@ -10,7 +10,7 @@ from collections import abc from typing import TYPE_CHECKING, Any -from attrs import Attribute, define, field +from attrs import define, field from attrs.converters import default_if_none from qrules.io import _labels @@ -24,19 +24,6 @@ _LOGGER = logging.getLogger(__name__) -def _check_booleans( - instance: GraphvizPrinter, - _attribute: Attribute, - _value: bool, -) -> None: - if instance.strip_spin and instance.collapse_graphs: - msg = "Cannot both strip spin and collapse graphs" - raise ValueError(msg) - if instance.collapse_graphs and instance.render_node: - msg = "Collapsed graphs cannot be rendered with node properties" - raise ValueError(msg) - - def _create_default_figure_style(style: dict[str, Any] | None) -> dict[str, Any]: figure_style = {"bgcolor": None} if style is None: @@ -45,14 +32,13 @@ def _create_default_figure_style(style: dict[str, Any] | None) -> dict[str, Any] return figure_style -@define(on_setattr=_check_booleans) +@define(kw_only=True) class GraphvizPrinter: render_node: bool | None = None render_final_state_id: bool = True render_resonance_id: bool = False render_initial_state_id: bool = False - strip_spin: bool = False - collapse_graphs: bool = False + collapse: _labels.CollapseMode | None = None figure_style: dict[str, Any] = field( converter=_create_default_figure_style, default=None @@ -99,9 +85,8 @@ def _render(self, obj: Any) -> list[str]: def _render_multiple_transitions(self, obj: Iterable) -> list[str]: transitions = _labels.select_transitions( obj, - collapse=self.collapse_graphs, + collapse=self.collapse, render_node=self.render_node, - strip_spin=self.strip_spin, ) lines = [] for i, graph in enumerate(reversed(transitions)): diff --git a/src/qrules/io/_labels.py b/src/qrules/io/_labels.py index 72957ebc..4fa1d55a 100644 --- a/src/qrules/io/_labels.py +++ b/src/qrules/io/_labels.py @@ -5,7 +5,7 @@ from fractions import Fraction from functools import singledispatch from inspect import isfunction -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias import attrs @@ -29,6 +29,7 @@ _LOGGER = logging.getLogger(__name__) +CollapseMode: TypeAlias = Literal["spin", "topology"] RenderedGraph = ProblemSet | QNProblemSet | Topology | Transition RenderPair = tuple[Topology, RenderedGraph] RenderInput = RenderedGraph | RenderPair @@ -85,23 +86,28 @@ def select_transitions( graphs: Iterable[Any], /, *, - collapse: bool, + collapse: CollapseMode | None, render_node: bool | None, - strip_spin: bool, ) -> list[Any]: """Reduce a collection of transitions to the graphs that are worth rendering. - The ``collapse`` and ``strip_spin`` flags are the printer attributes - :code:`collapse_graphs` and :code:`strip_spin`. Spin projections can only be - stripped from the interaction nodes if those nodes are not rendered. + The ``collapse`` mode is the printer attribute :code:`collapse`. Spin projections + are only stripped from the interaction nodes if those nodes are not rendered, and + topologies cannot be collapsed at all while node properties are rendered. """ - if collapse: - return collapse_graphs(graphs) - if strip_spin: + if collapse is None: + return list(graphs) + if collapse == "spin": if render_node: return sorted({strip_projections(g) for g in graphs}) return get_particle_graphs(graphs) - return list(graphs) + if collapse == "topology": + if render_node: + msg = "Transitions collapsed by topology cannot render node properties" + raise ValueError(msg) + return collapse_graphs(graphs) + msg = f"Unknown collapse mode {collapse!r}; expected None, 'spin', or 'topology'" + raise ValueError(msg) def create_edge_label( diff --git a/src/qrules/io/_mermaid.py b/src/qrules/io/_mermaid.py index 16239713..f4d9a7aa 100644 --- a/src/qrules/io/_mermaid.py +++ b/src/qrules/io/_mermaid.py @@ -192,8 +192,7 @@ class MermaidPrinter: render_final_state_id: bool = True render_resonance_id: bool = False render_initial_state_id: bool = False - strip_spin: bool = False - collapse_graphs: bool = False + collapse: _labels.CollapseMode | None = None figure_style: dict[str, Any] = field(converter=_to_style_dict, default=None) edge_style: dict[str, Any] = field(converter=_to_style_dict, default=None) node_style: dict[str, Any] = field(converter=_to_style_dict, default=None) @@ -220,9 +219,8 @@ def _render(self, obj: Any) -> list[str]: def _render_multiple_transitions(self, obj: Iterable) -> list[str]: transitions = _labels.select_transitions( obj, - collapse=self.collapse_graphs, + collapse=self.collapse, render_node=self.render_node, - strip_spin=self.strip_spin, ) lines: list[str] = [] for i, graph in enumerate(reversed(transitions)): diff --git a/tests/unit/io/test_dot.py b/tests/unit/io/test_dot.py index 0531b1f3..5b012536 100644 --- a/tests/unit/io/test_dot.py +++ b/tests/unit/io/test_dot.py @@ -20,9 +20,9 @@ def it_serializes_a_reaction(reaction: ReactionInfo): assert is_valid_dot(src) src = io.asdot(reaction) assert is_valid_dot(src) - src = io.asdot(reaction, strip_spin=True) + src = io.asdot(reaction, collapse="spin") assert is_valid_dot(src) - src = io.asdot(reaction, collapse_graphs=True) + src = io.asdot(reaction, collapse="topology") assert is_valid_dot(src) def it_exact_format(reaction: ReactionInfo): @@ -186,10 +186,10 @@ def it_write_graph_list(output_dir: str, reaction: ReactionInfo): src = stream.read() assert is_valid_dot(src) - def it_write_strip_spin(output_dir: str, reaction: ReactionInfo): + def it_write_collapsed_spin(output_dir: str, reaction: ReactionInfo): output_file = output_dir + "test_particle_graphs.gv" io.write( - instance=io.asdot(reaction, strip_spin=True), + instance=io.asdot(reaction, collapse="spin"), filename=output_file, ) with open(output_file) as stream: diff --git a/tests/unit/io/test_labels.py b/tests/unit/io/test_labels.py index 97408197..ef58a6b2 100644 --- a/tests/unit/io/test_labels.py +++ b/tests/unit/io/test_labels.py @@ -3,6 +3,7 @@ from textwrap import dedent import attrs +import pytest import qrules from qrules.io._labels import ( @@ -11,6 +12,7 @@ collapse_graphs, create_edge_label, get_particle_graphs, + select_transitions, strip_projections, ) from qrules.particle import Particle, ParticleCollection @@ -355,6 +357,44 @@ def test_get_particle_graphs( assert graphs[0].states[i] is graphs[1].states[i] +def describe_select_transitions(): + def it_keeps_all_transitions(reaction: ReactionInfo): + transitions = reaction.transitions + selection = select_transitions(transitions, collapse=None, render_node=False) + assert selection == list(transitions) + + def it_collapses_spin_into_particle_graphs(reaction: ReactionInfo): + transitions = reaction.transitions + selection = select_transitions(transitions, collapse="spin", render_node=False) + assert selection == get_particle_graphs(transitions) + + def it_strips_projections_when_rendering_nodes(reaction: ReactionInfo): + transitions = reaction.transitions + selection = select_transitions(transitions, collapse="spin", render_node=True) + assert selection == sorted({strip_projections(t) for t in transitions}) + + def it_collapses_topologies(reaction: ReactionInfo): + transitions = reaction.transitions + selection = select_transitions( + transitions, collapse="topology", render_node=False + ) + assert selection == collapse_graphs(transitions) + + def it_rejects_node_properties_on_collapsed_topologies(reaction: ReactionInfo): + with pytest.raises(ValueError, match="cannot render node properties"): + select_transitions( + reaction.transitions, collapse="topology", render_node=True + ) + + def it_rejects_an_unknown_mode(reaction: ReactionInfo): + with pytest.raises(ValueError, match="Unknown collapse mode 'invalid'"): + select_transitions( + reaction.transitions, + collapse="invalid", # ty: ignore[invalid-argument-type] + render_node=False, + ) + + def test_strip_projections(skh_particle_version: str): assert skh_particle_version is not None # skips test if particle version too low resonance = "Sigma(1670)~-" diff --git a/tests/unit/io/test_mermaid.py b/tests/unit/io/test_mermaid.py index 8af31221..bee8182c 100644 --- a/tests/unit/io/test_mermaid.py +++ b/tests/unit/io/test_mermaid.py @@ -174,13 +174,13 @@ def it_latex_reaction(reaction: ReactionInfo): assert labeled_lines assert all(line.count("$$") == 2 for line in labeled_lines) - def it_latex_collapsed_graph(reaction: ReactionInfo): - src = io.asmermaid(reaction, collapse_graphs=True, latex=True) + def it_latex_collapsed_topology(reaction: ReactionInfo): + src = io.asmermaid(reaction, collapse="topology", latex=True) assert R"$$\begin{gathered} f_{0}(980)" in src assert R"\\\ f_{0}(1500) \end{gathered}$$" in src - def it_latex_strip_spin(reaction: ReactionInfo): - src = io.asmermaid(reaction, strip_spin=True, latex=True) + def it_latex_collapsed_spin(reaction: ReactionInfo): + src = io.asmermaid(reaction, collapse="spin", latex=True) assert R"J/\psi(1S)$$" in src assert R"\gamma$$" in src assert R"\left[" not in src @@ -211,10 +211,10 @@ def it_reaction(reaction: ReactionInfo): src = io.asmermaid(reaction, latex=False) assert src.startswith("flowchart LR\n") assert " --- " in src - src = io.asmermaid(reaction, strip_spin=True, latex=False) + src = io.asmermaid(reaction, collapse="spin", latex=False) assert src.startswith("flowchart LR\n") assert " --- " in src - src = io.asmermaid(reaction, collapse_graphs=True, latex=False) + src = io.asmermaid(reaction, collapse="topology", latex=False) assert src.startswith("flowchart LR\n") assert " --- " in src