FEAT: select allowed resonances with regex (#212)

* DOC: find f(2) resonances in `reaction.ipynb` notebook
* DX: add pattern test for `allowed_intermediate_particles`
* ENH: avoid duplicate intermediate state definitions
* FEAT: allow simplified syntax without list
* MAINT: extend mass leak factor in Lc to pKpi
* MAINT: increase max angular momentum in Lc to pKpi
* MAINT: remove Delta(19**) resonances in Lc to pKpi
* MAINT: remove redundant cast in doctest
* MAINT: write `particle_filter` instead of `particle`
* MAINT: write `state` instead of `particle for property maps

docs/usage/reaction.ipynb

@@ -132,6 +132,7 @@
  " initial_state=[\"J/psi(1S)\"],\n",
  " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n",
  " formalism=\"helicity\",\n",
+ " max_angular_momentum=2,\n",
  ")"
  ]
  },
@@ -416,30 +417,49 @@
  "cell_type": "markdown",
  "metadata": {},
  "source": [
- "Unfortunately, the $\\rho^0$ mainly decays into $\\pi^0+\\pi^0$, not $\\gamma+\\pi^0$ and is therefore suppressed. This information is currently not known to {mod}`qrules`, but it is possible to hand {mod}`qrules` a list of allowed intermediate states."
+ "Unfortunately, the $\\rho^0$ mainly decays into $\\pi^0+\\pi^0$, not $\\gamma+\\pi^0$ and is therefore suppressed. This information is currently not known to {mod}`qrules`, but it is possible to hand {mod}`qrules` a list of allowed intermediate states,"
  ]
  },
  {
  "cell_type": "code",
  "execution_count": null,
- "metadata": {},
+ "metadata": {
+ "tags": [
+ "hide-output"
+ ]
+ },
  "outputs": [],
  "source": [
- "# particles are found by name comparison,\n",
- "# i.e. f2 will find all f2's and f all f's independent of their spin\n",
  "stm.set_allowed_intermediate_particles([\"f(0)\", \"f(2)\"])\n",
- "\n",
  "reaction = stm.find_solutions(problem_sets)\n",
- "\n",
- "print(\"found\", len(reaction.transitions), \"solutions!\")\n",
  "reaction.get_intermediate_particles().names"
  ]
  },
  {
  "cell_type": "markdown",
  "metadata": {},
  "source": [
- "Now we have selected all amplitudes that involve **f** states. The warnings appear only to notify the user that the list of solutions is not exhaustive: for certain edges in the graph, no suitable particle was found (since only f states were allowed)."
+ "or, using regular expressions,"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "tags": []
+ },
+ "outputs": [],
+ "source": [
+ "stm.set_allowed_intermediate_particles(r\"f\\([02]\\)\", regex=True)\n",
+ "reaction = stm.find_solutions(problem_sets)\n",
+ "assert len(reaction.get_intermediate_particles().names) == 11"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Now we have selected all amplitudes that involve **f** states:"
  ]
  },
  {
@@ -521,7 +541,9 @@
  },
  {
  "cell_type": "markdown",
- "metadata": {},
+ "metadata": {
+ "tags": []
+ },
  "source": [
  "```{tip}\n",
  "The {mod}`ampform` package can formulate amplitude models based on the state transitions created by {mod}`qrules`. See {doc}`ampform:usage/amplitude`.\n",

src/qrules/particle.py

@@ -388,7 +388,7 @@ def filter( # noqa: A003
  ... and p.spin == 2
  ... and p.strangeness == 1
  ... )
- >>> sorted(list(subset.names))
+ >>> sorted(subset.names)
  ['K(2)(1820)+', 'K(2)(1820)0', 'K(2)*(1980)+', 'K(2)*(1980)0']
  """
  return ParticleCollection({particle for particle in self if function(particle)})

src/qrules/transition.py

@@ -1,6 +1,7 @@
 """Find allowed transitions between an initial and final state."""
 
 import logging
+import re
 import sys
 from collections import defaultdict
 from copy import copy, deepcopy
@@ -321,30 +322,33 @@ def __init__( # pylint: disable=too-many-arguments, too-many-branches, too-many
  max_spin_magnitude=max_spin_magnitude,
  )
 
- self.__user_allowed_intermediate_particles = allowed_intermediate_particles
- self.__allowed_intermediate_particles: List[GraphEdgePropertyMap] = []
- if allowed_intermediate_particles is not None:
- self.set_allowed_intermediate_particles(allowed_intermediate_particles)
- else:
- self.__allowed_intermediate_particles = [
+ self.__intermediate_particle_filters = allowed_intermediate_particles
+ if allowed_intermediate_particles is None:
+ self.__allowed_intermediate_states: List[GraphEdgePropertyMap] = [
  create_edge_properties(x) for x in self.__particles
  ]
+ else:
+ self.set_allowed_intermediate_particles(allowed_intermediate_particles)
 
- def set_allowed_intermediate_particles(self, particle_names: List[str]) -> None:
- self.__allowed_intermediate_particles = []
- for particle_name in particle_names:
+ def set_allowed_intermediate_particles(
+ self, name_patterns: Union[Iterable[str], str], regex: bool = False
+ ) -> None:
+ if isinstance(name_patterns, str):
+ name_patterns = [name_patterns]
+ selected_particles = ParticleCollection()
+ for pattern in name_patterns:
  # pylint: disable=cell-var-from-loop
- matches = self.__particles.filter(
- lambda p: particle_name in p.name # noqa: B023
- )
+ matches = _filter_by_name_pattern(self.__particles, pattern, regex)
  if len(matches) == 0:
  raise LookupError(
  "Could not find any matches for allowed intermediate"
- f' particle "{particle_name}"'
+ f' particle pattern "{pattern}"'
  )
- self.__allowed_intermediate_particles += [
- create_edge_properties(x) for x in matches
- ]
+ selected_particles.update(matches)
+ self.__allowed_intermediate_states = [
+ create_edge_properties(x)
+ for x in sorted(selected_particles, key=lambda p: p.name)
+ ]
 
  @property
  def formalism(self) -> str:
@@ -432,7 +436,7 @@ def __determine_graph_settings(
  def create_intermediate_edge_qn_domains() -> Dict:
  # if a list of intermediate states is given by user,
  # built a domain based on these states
- if self.__user_allowed_intermediate_particles:
+ if self.__intermediate_particle_filters is not None:
  intermediate_edge_domains: Dict[Type[EdgeQuantumNumber], Set] = (
  defaultdict(set)
  )
@@ -441,7 +445,7 @@ def create_intermediate_edge_qn_domains() -> Dict:
  EdgeQuantumNumbers.spin_projection
  ]
  )
- for particle_props in self.__allowed_intermediate_particles:
+ for particle_props in self.__allowed_intermediate_states:
  for edge_qn, qn_value in particle_props.items():
  intermediate_edge_domains[edge_qn].add(qn_value)
 
@@ -653,7 +657,7 @@ def find_quantum_number_transitions(
  return qn_results
 
  def _solve(self, qn_problem_set: QNProblemSet) -> Tuple[QNProblemSet, QNResult]:
- solver = CSPSolver(self.__allowed_intermediate_particles)
+ solver = CSPSolver(self.__allowed_intermediate_states)
  solutions = solver.find_solutions(qn_problem_set)
  return qn_problem_set, solutions
 
@@ -701,6 +705,20 @@ def _safe_wrap_list(nested_list: Union[List[str], List[List[str]]]) -> List[List
  )
 
 
+def _filter_by_name_pattern(
+ particles: ParticleCollection, pattern: str, regex: bool
+) -> ParticleCollection:
+ def match_regex(particle: Particle) -> bool:
+ return re.match(pattern, particle.name) is not None
+
+ def match_substring(particle: Particle) -> bool:
+ return pattern in particle.name
+
+ if regex:
+ return particles.filter(match_regex)
+ return particles.filter(match_substring)
+
+
 def _match_final_state_ids(
  graph: "MutableTransition[ParticleWithSpin, InteractionProperties]",
  state_definition: Sequence[StateDefinition],

tests/channels/test_lc_to_p_km_pip.py

@@ -6,21 +6,29 @@ def test_resonances():
  stm = StateTransitionManager(
  initial_state=["Lambda(c)+"],
  final_state=["p", "K-", "pi+"],
- allowed_intermediate_particles=["Delta", "K", "Lambda"],
- mass_conservation_factor=0,
+ mass_conservation_factor=0.6,
+ max_angular_momentum=2,
  )
  stm.set_allowed_interaction_types([InteractionType.STRONG], node_id=1)
+ stm.set_allowed_intermediate_particles([r"Delta..(?!9)", r"^K", r"^L"], regex=True)
  problem_sets = stm.create_problem_sets()
  reaction = stm.find_solutions(problem_sets)
  resonances = reaction.get_intermediate_particles().names
+ # https://lc2pkpi-polarimetry.docs.cern.ch/amplitude-model.html#resonances-and-ls-scheme
  assert resonances == [
  "Delta(1232)++",
  "Delta(1600)++",
  "Delta(1620)++",
+ "Delta(1700)++",
  "K(0)*(700)~0",
  "K*(892)~0",
+ "K*(1410)~0",
+ "K(0)*(1430)~0",
+ "Lambda(1405)",
+ "Lambda(1520)",
  "Lambda(1600)",
  "Lambda(1670)",
+ "Lambda(1690)",
  "Lambda(1810)",
  "Lambda(1800)",
  "Lambda(1890)",

tests/unit/test_transition.py

@@ -7,6 +7,7 @@
 
 from qrules.particle import Parity, Particle, ParticleCollection, Spin # noqa: F401
 from qrules.quantum_numbers import InteractionProperties # noqa: F401
+from qrules.settings import InteractionType
 from qrules.topology import ( # noqa: F401
  Edge,
  FrozenDict,
@@ -74,4 +75,31 @@ def test_allowed_intermediate_particles(self):
  LookupError,
  match=r"Could not find any matches for allowed intermediate particle",
  ):
- stm.set_allowed_intermediate_particles([particle_name])
+ stm.set_allowed_intermediate_particles(particle_name)
+
+ def test_regex_pattern(self):
+ stm = StateTransitionManager(
+ initial_state=["Lambda(c)+"],
+ final_state=["p", "K-", "pi+"],
+ allowed_intermediate_particles=["Delta"],
+ )
+ stm.set_allowed_interaction_types([InteractionType.STRONG], node_id=1)
+ problem_sets = stm.create_problem_sets()
+ reaction = stm.find_solutions(problem_sets)
+ assert reaction.get_intermediate_particles().names == [
+ "Delta(1232)++",
+ "Delta(1600)++",
+ "Delta(1620)++",
+ "Delta(1900)++",
+ "Delta(1910)++",
+ "Delta(1920)++",
+ ]
+
+ stm.set_allowed_intermediate_particles(r"^Delta\(\d(60|9[02])0\)", regex=True)
+ problem_sets = stm.create_problem_sets()
+ reaction = stm.find_solutions(problem_sets)
+ assert reaction.get_intermediate_particles().names == [
+ "Delta(1600)++",
+ "Delta(1900)++",
+ "Delta(1920)++",
+ ]