# Copyright (c) 2018-2026, Eduardo Rodrigues and Henry Schreiner.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/decaylanguage for details.
"""
Submodule with classes and utilities to visualize decay chains.
Decay chains are typically provided by the parser of .dec decay files,
see the ``DecFileParser`` class.
"""
from __future__ import annotations
import html
import itertools
from collections.abc import Iterable
from typing import Any
import graphviz
from particle import latex_to_html_name
from particle.converters.bimap import DirectionalMap, DirectionalMaps
from decaylanguage.decay.decay import DecayChain, _has_no_subdecay
_EvtGen2LatexNameMap: DirectionalMap[str, str]
_Latex2EvtGenNameMap: DirectionalMap[str, str]
_EvtGen2LatexNameMap, _Latex2EvtGenNameMap = DirectionalMaps("EvtGenName", "LaTexName")
[docs]
class GraphNotBuiltError(RuntimeError):
pass
[docs]
class DecayChainViewer:
"""
The class to visualize a decay chain.
Examples
--------
>>> dfp = DecFileParser('my-Dst-decay-file.dec') # doctest: +SKIP
>>> dfp.parse() # doctest: +SKIP
>>> chain = dfp.build_decay_chains('D*+') # doctest: +SKIP
>>> dcv = DecayChainViewer(chain) # doctest: +SKIP
>>> # display the SVG figure in a notebook
>>> dcv # doctest: +SKIP
When not in notebooks the graph can easily be visualized with the
``graphviz.Digraph.render`` or ``graphviz.Digraph.view`` functions, e.g.:
>>> dcv.graph.render(filename="test", format="pdf", view=True, cleanup=True) # doctest: +SKIP
"""
__slots__ = ("_chain", "_counter", "_graph", "_show_effective_bf")
def __init__(
self,
decaychain: dict[str, list[dict[str, float | str | list[Any]]]] | DecayChain,
show_effective_bf: bool = False,
**attrs: dict[str, bool | int | float | str],
) -> None:
"""
Default constructor.
Parameters
----------
decaychain: dict
Input decay chain in dict format, typically created from ``decaylanguage.DecFileParser.build_decay_chains``
after parsing a .dec decay file, or from building a decay chain representation with ``decaylanguage.DecayChain``.
show_effective_bf: bool, optional
If True, display the effective branching fraction on the terminal node (node with no
subchains). The effective branching fraction is the product of all branching fractions
along the chain from the mother particle to the terminal node. Default is False.
attrs: optional
User input ``graphviz.Digraph`` class attributes.
See also
--------
decaylanguage.DecFileParser.build_decay_chains for creating a decay chain dict from parsing a .dec file.
decaylanguage.DecayChain: class for building a decay chain programmatically.
"""
# Accept DecayChain objects directly, converting to dict internally
chain: dict[str, list[Any]]
if isinstance(decaychain, DecayChain):
chain = decaychain.to_dict()
else:
chain = decaychain
# Store the input decay chain as dict
self._chain: dict[str, list[Any]] = chain
# Store whether to show effective branching fractions
self._show_effective_bf = show_effective_bf
# Per-instance node counter, so that rendering the same chain
# twice produces identical, reproducible DOT output.
self._counter = itertools.count()
# Instantiate the digraph with defaults possibly overridden by user attributes
self._graph = self._instantiate_graph(**attrs)
# Build the actual graph from the input decay chain structure
self._build_decay_graph()
def _build_decay_graph(self) -> None:
"""
Recursively navigate the decay chain tree and produce a Digraph
in the DOT language.
"""
def safe_html_name(name: str) -> str:
"""
Get a safe HTML name from the EvtGen name.
Note
----
The match is done using a conversion map rather than via
``Particle.from_evtgen_name(name).html_name`` for 2 reasons:
- Some decay-file-specific "particle" names (e.g. cs_0)
are not in the PDG table.
- No need to load all particle information if all that's needed
is a match EvtGen - HTML name.
"""
try:
return latex_to_html_name(_EvtGen2LatexNameMap[name])
except Exception:
# Escape characters such as &, <, > (legal in .dec Alias
# statements) so the fallback yields valid HTML-like DOT labels.
return html.escape(name)
def html_table_label(
names: list[str],
add_tags: bool = False,
bgcolor: str = "#9abad6",
) -> str:
if add_tags:
label = f'<<TABLE BORDER="0" CELLSPACING="0" BGCOLOR="{bgcolor}">'
else:
label = f'<<TABLE BORDER="0" CELLSPACING="0" CELLPADDING="0" BGCOLOR="{bgcolor}"><TR>'
for i, n in enumerate(names):
if add_tags:
label += f'<TR><TD BORDER="1" CELLPADDING="5" PORT="p{i}">{safe_html_name(n)}</TD></TR>'
else:
label += f'<TD BORDER="0" CELLPADDING="2">{safe_html_name(n)}</TD>'
label += "{tr}</TABLE>>".format(tr="" if add_tags else "</TR>")
return label
def new_node_no_subchain(list_parts: list[str], effective_bf: float) -> str:
label = html_table_label(list_parts, bgcolor="#eef3f8")
r = f"dec{next(self._counter)}"
self.graph.node(r, label=label, style="filled", fillcolor="#eef3f8")
if self._show_effective_bf:
bf_node = f"dec{next(self._counter)}"
self.graph.node(
bf_node,
label=f"eff BF: {effective_bf:.4g}",
shape="none",
fontcolor="#4c4c4c",
)
self.graph.edge(r, bf_node, style="invis")
return r
def new_node_with_subchain(list_parts: list[Any]) -> str:
_list_parts = [
next(iter(p.keys())) if isinstance(p, dict) else p for p in list_parts
]
label = html_table_label(_list_parts, add_tags=True)
r = f"dec{next(self._counter)}"
self.graph.node(r, shape="none", label=label)
return r
def iterate_chain(
subchain: list[dict[str, float | str | list[Any]]],
top_node: str | None = None,
link_pos: int | None = None,
effective_bf: float = 1.0,
) -> None:
if not top_node:
top_node = "mother"
self.graph.node("mother", shape="none", label=label)
n_decaymodes = len(subchain)
for idm in range(n_decaymodes):
_list_parts = subchain[idm]["fs"]
_bf = subchain[idm]["bf"]
_effective_bf = effective_bf * float(_bf) # type: ignore[arg-type]
if _has_no_subdecay(_list_parts): # type: ignore[arg-type]
_ref = new_node_no_subchain(_list_parts, _effective_bf) # type: ignore[arg-type]
if link_pos is None:
self.graph.edge(top_node, _ref, label=str(_bf))
else:
self.graph.edge(f"{top_node}:p{link_pos}", _ref, label=str(_bf))
else:
_ref_1 = new_node_with_subchain(_list_parts) # type: ignore[arg-type]
edge_label = str(_bf)
if link_pos is None:
self.graph.edge(top_node, _ref_1, label=edge_label)
else:
self.graph.edge(
f"{top_node}:p{link_pos}",
_ref_1,
label=edge_label,
)
for i, _p in enumerate(_list_parts): # type: ignore[arg-type]
if not isinstance(_p, str):
_k = next(iter(_p.keys()))
iterate_chain(
_p[_k],
top_node=_ref_1,
link_pos=i,
effective_bf=_effective_bf,
)
k = next(iter(self._chain.keys()))
label = html_table_label([k], add_tags=True, bgcolor="#568dba")
sc = self._chain[k]
# Actually build the whole decay chain, iteratively
iterate_chain(sc)
@property
def graph(self) -> graphviz.Digraph:
"""
Get the actual ``graphviz.Digraph`` object.
The user now has full control ...
"""
return self._graph
[docs]
def to_string(self) -> str:
"""
Return a string representation of the built graph in the DOT language.
The function is a trivial shortcut for ``graphviz.Digraph.source``.
"""
return self.graph.source # type: ignore[no-any-return]
def _instantiate_graph(
self, **attrs: dict[str, bool | int | float | str]
) -> graphviz.Digraph:
"""
Return a ``graphviz.Digraph`` class instance using the default attributes
specified in this class:
- Default graph attributes are overridden by input by the user.
- Class and node and edge defaults.
"""
graph_attr = self._get_graph_defaults()
node_attr = self._get_node_defaults()
edge_attr = self._get_edge_defaults()
for key, attr in (
("graph_attr", graph_attr),
("node_attr", node_attr),
("edge_attr", edge_attr),
):
if key in attrs:
attr.update(**attrs.pop(key))
arguments = self._get_default_arguments()
arguments.update(**attrs) # type: ignore[call-overload]
return graphviz.Digraph(
graph_attr=graph_attr, node_attr=node_attr, edge_attr=edge_attr, **arguments
)
def _get_default_arguments(self) -> dict[str, bool | int | float | str]:
"""
``graphviz.Digraph`` default arguments.
"""
return {
"name": "DecayChainGraph",
"comment": "Created by https://github.com/scikit-hep/decaylanguage",
"engine": "dot",
"format": "png",
}
def _get_graph_defaults(self) -> dict[str, bool | int | float | str]:
# Only real DOT graph attributes belong here. Constructor arguments
# such as name/comment/engine/format are handled separately via
# _get_default_arguments, so they do not leak into the DOT body.
return {"rankdir": "LR"}
def _get_node_defaults(self) -> dict[str, bool | int | float | str]:
return {"fontname": "Helvetica", "fontsize": "11", "shape": "oval"}
def _get_edge_defaults(self) -> dict[str, bool | int | float | str]:
return {"fontcolor": "#4c4c4c", "fontsize": "11"}
def _repr_mimebundle_(
self,
include: Iterable[str] | None = None,
exclude: Iterable[str] | None = None,
**kwargs: Any,
) -> Any: # pragma: no cover
"""
IPython display helper.
"""
try:
return self._graph._repr_mimebundle_(
include=include, exclude=exclude, **kwargs
)
except AttributeError:
return {"image/svg+xml": self._graph._repr_svg_()} # for graphviz < 0.19