1pub mod angular_momentum {
2 use laddu_core::{
3 allowed_projections, AngularMomentum, LadduError, LadduResult, OrbitalAngularMomentum,
4 Projection,
5 };
6 use num::rational::Ratio;
7 use pyo3::{
8 prelude::*,
9 types::{PyAny, PyBool, PyModule},
10 IntoPyObjectExt,
11 };
12 type PyQuantumNumber = Py<PyAny>;
13
14 pub fn parse_angular_momentum(input: &Bound<'_, PyAny>) -> PyResult<AngularMomentum> {
15 Ok(parse_ratio_like(input).and_then(AngularMomentum::try_from)?)
16 }
17
18 fn parse_ratio_like(input: &Bound<'_, PyAny>) -> LadduResult<Ratio<i32>> {
19 if input.is_instance_of::<PyBool>() {
20 return Err(LadduError::Custom(
21 "quantum number cannot be a bool".to_string(),
22 ));
23 }
24 if let Ok(value) = input.extract::<i32>() {
25 return Ok(Ratio::from_integer(value));
26 }
27 if let Ok(value) = input.extract::<f64>() {
28 let twice = Projection::try_from(value)?.value();
29 return Ok(Ratio::new(twice, 2));
30 }
31 let numerator = input
32 .getattr("numerator")
33 .and_then(|value| value.extract::<i32>());
34 let denominator = input
35 .getattr("denominator")
36 .and_then(|value| value.extract::<i32>());
37 if let (Ok(numerator), Ok(denominator)) = (numerator, denominator) {
38 if denominator == 0 {
39 return Err(LadduError::Custom(
40 "quantum number denominator cannot be zero".to_string(),
41 ));
42 }
43 return Ok(Ratio::new(numerator, denominator));
44 }
45 Err(LadduError::Custom(
46 "quantum number must be an int, float, or fractions.Fraction".to_string(),
47 ))
48 }
49
50 pub fn parse_projection(input: &Bound<'_, PyAny>) -> PyResult<Projection> {
51 Ok(parse_ratio_like(input).and_then(Projection::try_from)?)
52 }
53
54 pub fn parse_orbital_angular_momentum(
55 input: &Bound<'_, PyAny>,
56 ) -> PyResult<OrbitalAngularMomentum> {
57 Ok(parse_ratio_like(input).and_then(OrbitalAngularMomentum::try_from)?)
58 }
59
60 pub fn angular_momentum_to_python(
61 py: Python<'_>,
62 angular_momentum: laddu_core::AngularMomentum,
63 ) -> PyResult<PyQuantumNumber> {
64 let twice = angular_momentum.value() as i32;
65 if twice % 2 == 0 {
66 Ok((twice / 2).into_bound_py_any(py)?.unbind())
67 } else {
68 let fractions = PyModule::import(py, "fractions")?;
69 let fraction = fractions.getattr("Fraction")?;
70 Ok(fraction.call1((twice, 2))?.unbind())
71 }
72 }
73
74 pub fn projection_to_python(
75 py: Python<'_>,
76 projection: Projection,
77 ) -> PyResult<PyQuantumNumber> {
78 let twice = projection.value();
79 if twice % 2 == 0 {
80 Ok((twice / 2).into_bound_py_any(py)?.unbind())
81 } else {
82 let fractions = PyModule::import(py, "fractions")?;
83 let fraction = fractions.getattr("Fraction")?;
84 Ok(fraction.call1((twice, 2))?.unbind())
85 }
86 }
87
88 #[pyfunction(name = "allowed_projections")]
90 pub fn py_allowed_projections(
91 py: Python<'_>,
92 spin: &Bound<'_, PyAny>,
93 ) -> PyResult<Vec<PyQuantumNumber>> {
94 allowed_projections(parse_angular_momentum(spin)?)
95 .into_iter()
96 .map(|projection| projection_to_python(py, projection))
97 .collect()
98 }
99}
100
101use laddu_core::{
102 AllowedPartialWave, Charge, Isospin, LadduError, OrbitalAngularMomentum, Parity, PartialWave,
103 ParticleProperties, RuleSet, SelectionRules, Statistics,
104};
105use pyo3::{
106 exceptions::PyTypeError,
107 prelude::*,
108 types::{PyAny, PyBool},
109 IntoPyObjectExt,
110};
111
112use self::angular_momentum::{
113 angular_momentum_to_python, parse_angular_momentum, parse_orbital_angular_momentum,
114 parse_projection, projection_to_python,
115};
116
117type PyQuantumNumber = Py<PyAny>;
118
119fn parse_parity(input: &Bound<'_, PyAny>) -> PyResult<Parity> {
120 if let Ok(value) = input.extract::<PyParity>() {
121 return Ok(value.0);
122 }
123 if let Ok(value) = input.extract::<String>() {
124 return Ok(value.parse()?);
125 }
126 Err(PyTypeError::new_err(
127 "parity must be a Parity or sign string",
128 ))
129}
130
131fn parse_statistics(input: &Bound<'_, PyAny>) -> PyResult<Statistics> {
132 if let Ok(value) = input.extract::<PyStatistics>() {
133 return Ok(value.into());
134 }
135 if let Ok(value) = input.extract::<String>() {
136 return match value.to_ascii_lowercase().as_str() {
137 "boson" | "bosonic" => Ok(Statistics::Boson),
138 "fermion" | "fermionic" => Ok(Statistics::Fermion),
139 _ => Err(LadduError::ParseError {
140 name: value,
141 object: "Statistics".to_string(),
142 }
143 .into()),
144 };
145 }
146 Err(PyTypeError::new_err(
147 "statistics must be a Statistics value or string",
148 ))
149}
150
151fn parse_charge_input(input: &Bound<'_, PyAny>) -> PyResult<Charge> {
152 if let Ok(value) = input.extract::<PyCharge>() {
153 return Ok(value.0);
154 }
155 if input.is_instance_of::<PyBool>() {
156 return Err(LadduError::Custom("electric charge cannot be a bool".to_string()).into());
157 }
158 if let Ok(value) = input.extract::<i32>() {
159 return Ok(Charge::try_from(num::rational::Ratio::from_integer(value))?);
160 }
161 let numerator = input
162 .getattr("numerator")
163 .and_then(|value| value.extract::<i32>());
164 let denominator = input
165 .getattr("denominator")
166 .and_then(|value| value.extract::<i32>());
167 if let (Ok(numerator), Ok(denominator)) = (numerator, denominator) {
168 if denominator == 0 {
169 return Err(LadduError::Custom(
170 "electric charge denominator cannot be zero".to_string(),
171 )
172 .into());
173 }
174 return Ok(Charge::try_from(num::rational::Ratio::new(
175 numerator,
176 denominator,
177 ))?);
178 }
179 if let Ok(value) = input.extract::<f64>() {
180 return Ok(Charge::try_from(value)?);
181 }
182 Err(PyTypeError::new_err(
183 "electric charge must be an int, float, fractions.Fraction, or Charge",
184 ))
185}
186
187fn charge_to_python(py: Python<'_>, charge: Charge) -> PyResult<PyQuantumNumber> {
188 let thirds = charge.value();
189 if thirds % 3 == 0 {
190 Ok((thirds / 3).into_bound_py_any(py)?.unbind())
191 } else {
192 let fractions = pyo3::types::PyModule::import(py, "fractions")?;
193 let fraction = fractions.getattr("Fraction")?;
194 Ok(fraction.call1((thirds, 3))?.unbind())
195 }
196}
197
198#[pyclass(eq, name = "Parity", module = "laddu", from_py_object)]
199#[derive(Clone, Copy, PartialEq)]
200pub struct PyParity(pub Parity);
201
202#[pymethods]
203impl PyParity {
204 #[new]
205 fn new(value: &str) -> PyResult<Self> {
206 Ok(Self(value.parse()?))
207 }
208
209 #[staticmethod]
210 fn positive() -> Self {
211 Self(Parity::Positive)
212 }
213
214 #[staticmethod]
215 fn negative() -> Self {
216 Self(Parity::Negative)
217 }
218
219 #[getter]
220 fn value(&self) -> i32 {
221 self.0.value()
222 }
223
224 fn __repr__(&self) -> String {
225 format!("Parity('{}')", self.0)
226 }
227
228 fn __str__(&self) -> String {
229 self.0.to_string()
230 }
231}
232
233#[pyclass(eq, eq_int, name = "Statistics", module = "laddu", from_py_object)]
234#[derive(Clone, Copy, PartialEq)]
235pub enum PyStatistics {
236 Boson,
237 Fermion,
238}
239
240impl From<PyStatistics> for Statistics {
241 fn from(value: PyStatistics) -> Self {
242 match value {
243 PyStatistics::Boson => Self::Boson,
244 PyStatistics::Fermion => Self::Fermion,
245 }
246 }
247}
248
249impl From<Statistics> for PyStatistics {
250 fn from(value: Statistics) -> Self {
251 match value {
252 Statistics::Boson => Self::Boson,
253 Statistics::Fermion => Self::Fermion,
254 }
255 }
256}
257
258#[pyclass(eq, name = "Charge", module = "laddu", from_py_object)]
259#[derive(Clone, Copy, PartialEq)]
260pub struct PyCharge(pub Charge);
261
262#[pymethods]
263impl PyCharge {
264 #[new]
265 fn new(value: &Bound<'_, PyAny>) -> PyResult<Self> {
266 Ok(Self(parse_charge_input(value)?))
267 }
268
269 #[getter]
270 fn value(&self, py: Python<'_>) -> PyResult<PyQuantumNumber> {
271 charge_to_python(py, self.0)
272 }
273
274 fn __repr__(&self) -> String {
275 format!("Charge({})", self.0)
276 }
277
278 fn __str__(&self) -> String {
279 self.0.to_string()
280 }
281}
282
283#[pyclass(eq, name = "Isospin", module = "laddu", from_py_object)]
284#[derive(Clone, Copy, PartialEq)]
285pub struct PyIsospin(pub Isospin);
286
287#[pymethods]
288impl PyIsospin {
289 #[new]
290 #[pyo3(signature = (isospin, *, projection=None))]
291 fn new(isospin: &Bound<'_, PyAny>, projection: Option<&Bound<'_, PyAny>>) -> PyResult<Self> {
292 Ok(Self(Isospin::new(
293 parse_angular_momentum(isospin)?,
294 projection.map(parse_projection).transpose()?,
295 )?))
296 }
297
298 #[getter]
299 fn isospin(&self, py: Python<'_>) -> PyResult<PyQuantumNumber> {
300 angular_momentum_to_python(py, self.0.isospin())
301 }
302
303 #[getter]
304 fn projection(&self, py: Python<'_>) -> PyResult<Option<PyQuantumNumber>> {
305 self.0
306 .projection()
307 .map(|projection| projection_to_python(py, projection))
308 .transpose()
309 }
310
311 fn __repr__(&self) -> String {
312 match self.0.projection() {
313 Some(projection) => format!("Isospin({}, projection={})", self.0.isospin(), projection),
314 None => format!("Isospin({})", self.0.isospin()),
315 }
316 }
317}
318
319#[pyclass(name = "ParticleProperties", module = "laddu", from_py_object)]
320#[derive(Clone)]
321pub struct PyParticleProperties(pub ParticleProperties);
322
323#[pymethods]
324impl PyParticleProperties {
325 #[new]
326 #[pyo3(signature = (name=None, *, species=None, antiparticle_species=None, self_conjugate=None, spin=None, parity=None, c_parity=None, g_parity=None, charge=None, isospin=None, strangeness=None, charm=None, bottomness=None, topness=None, baryon_number=None, electron_lepton_number=None, muon_lepton_number=None, tau_lepton_number=None, statistics=None))]
327 #[allow(clippy::too_many_arguments)]
328 fn new(
329 name: Option<String>,
330 species: Option<String>,
331 antiparticle_species: Option<String>,
332 self_conjugate: Option<bool>,
333 spin: Option<&Bound<'_, PyAny>>,
334 parity: Option<&Bound<'_, PyAny>>,
335 c_parity: Option<&Bound<'_, PyAny>>,
336 g_parity: Option<&Bound<'_, PyAny>>,
337 charge: Option<&Bound<'_, PyAny>>,
338 isospin: Option<PyIsospin>,
339 strangeness: Option<i32>,
340 charm: Option<i32>,
341 bottomness: Option<i32>,
342 topness: Option<i32>,
343 baryon_number: Option<i32>,
344 electron_lepton_number: Option<i32>,
345 muon_lepton_number: Option<i32>,
346 tau_lepton_number: Option<i32>,
347 statistics: Option<&Bound<'_, PyAny>>,
348 ) -> PyResult<Self> {
349 let mut properties = ParticleProperties::unknown();
350 if let Some(name) = name {
351 properties = properties.with_name(name);
352 }
353 if let Some(species) = species {
354 properties = properties.with_species(species);
355 }
356 if let Some(antiparticle_species) = antiparticle_species {
357 properties = properties.with_antiparticle_species(antiparticle_species);
358 }
359 if let Some(self_conjugate) = self_conjugate {
360 properties = properties.with_self_conjugate(self_conjugate);
361 }
362 if let Some(spin) = spin {
363 properties = properties.with_spin(parse_angular_momentum(spin)?);
364 }
365 if let Some(parity) = parity {
366 properties = properties.with_parity(parse_parity(parity)?);
367 }
368 if let Some(c_parity) = c_parity {
369 properties = properties.with_c_parity(parse_parity(c_parity)?);
370 }
371 if let Some(g_parity) = g_parity {
372 properties = properties.with_g_parity(parse_parity(g_parity)?);
373 }
374 if let Some(charge) = charge {
375 properties = properties.with_charge(parse_charge_input(charge)?);
376 }
377 if let Some(isospin) = isospin {
378 properties = properties.with_isospin(isospin.0);
379 }
380 if let Some(strangeness) = strangeness {
381 properties = properties.with_strangeness(strangeness);
382 }
383 if let Some(charm) = charm {
384 properties = properties.with_charm(charm);
385 }
386 if let Some(bottomness) = bottomness {
387 properties = properties.with_bottomness(bottomness);
388 }
389 if let Some(topness) = topness {
390 properties = properties.with_topness(topness);
391 }
392 if let Some(baryon_number) = baryon_number {
393 properties = properties.with_baryon_number(baryon_number);
394 }
395 if let Some(electron_lepton_number) = electron_lepton_number {
396 properties = properties.with_electron_lepton_number(electron_lepton_number);
397 }
398 if let Some(muon_lepton_number) = muon_lepton_number {
399 properties = properties.with_muon_lepton_number(muon_lepton_number);
400 }
401 if let Some(tau_lepton_number) = tau_lepton_number {
402 properties = properties.with_tau_lepton_number(tau_lepton_number);
403 }
404 if let Some(statistics) = statistics {
405 properties = properties.with_statistics(parse_statistics(statistics)?)?;
406 }
407 Ok(Self(properties))
408 }
409
410 #[getter]
411 fn name(&self) -> Option<String> {
412 self.0.name.clone()
413 }
414
415 #[getter]
416 fn spin(&self, py: Python<'_>) -> PyResult<Option<PyQuantumNumber>> {
417 self.0
418 .spin
419 .map(|spin| angular_momentum_to_python(py, spin))
420 .transpose()
421 }
422
423 #[getter]
424 fn parity(&self) -> Option<PyParity> {
425 self.0.parity.map(PyParity)
426 }
427
428 #[getter]
429 fn c_parity(&self) -> Option<PyParity> {
430 self.0.c_parity.map(PyParity)
431 }
432
433 #[getter]
434 fn g_parity(&self) -> Option<PyParity> {
435 self.0.g_parity.map(PyParity)
436 }
437
438 #[getter]
439 fn charge(&self) -> Option<PyCharge> {
440 self.0.charge.map(PyCharge)
441 }
442
443 #[getter]
444 fn isospin(&self) -> Option<PyIsospin> {
445 self.0.isospin.map(PyIsospin)
446 }
447
448 fn __repr__(&self) -> String {
449 format!("{:?}", self.0)
450 }
451}
452
453#[pyclass(eq, name = "PartialWave", module = "laddu", from_py_object)]
454#[derive(Clone, PartialEq)]
455pub struct PyPartialWave(pub PartialWave);
456
457#[pymethods]
458impl PyPartialWave {
459 #[new]
460 #[pyo3(signature = (*, j, l, s, label=None))]
461 fn new(
462 j: &Bound<'_, PyAny>,
463 l: &Bound<'_, PyAny>,
464 s: &Bound<'_, PyAny>,
465 label: Option<String>,
466 ) -> PyResult<Self> {
467 let wave = PartialWave::new(
468 parse_angular_momentum(j)?,
469 parse_orbital_angular_momentum(l)?,
470 parse_angular_momentum(s)?,
471 )?;
472 Ok(Self(match label {
473 Some(label) => wave.with_label(label),
474 None => wave,
475 }))
476 }
477
478 #[getter]
479 fn j(&self, py: Python<'_>) -> PyResult<PyQuantumNumber> {
480 angular_momentum_to_python(py, self.0.j)
481 }
482
483 #[getter]
484 fn l(&self) -> u32 {
485 self.0.l.value()
486 }
487
488 #[getter]
489 fn s(&self, py: Python<'_>) -> PyResult<PyQuantumNumber> {
490 angular_momentum_to_python(py, self.0.s)
491 }
492
493 #[getter]
494 fn label(&self) -> String {
495 self.0.label.clone()
496 }
497
498 fn __repr__(&self) -> String {
499 format!("PartialWave('{}')", self.0.label)
500 }
501
502 fn __str__(&self) -> String {
503 self.0.to_string()
504 }
505}
506
507#[pyclass(eq, name = "AllowedPartialWave", module = "laddu", from_py_object)]
508#[derive(Clone, PartialEq)]
509pub struct PyAllowedPartialWave(pub AllowedPartialWave);
510
511#[pymethods]
512impl PyAllowedPartialWave {
513 #[getter]
514 fn wave(&self) -> PyPartialWave {
515 PyPartialWave(self.0.wave.clone())
516 }
517
518 #[getter]
519 fn parity(&self) -> Option<PyParity> {
520 self.0.parity.map(PyParity)
521 }
522
523 #[getter]
524 fn c_parity(&self) -> Option<PyParity> {
525 self.0.c_parity.map(PyParity)
526 }
527
528 fn __repr__(&self) -> String {
529 format!("{:?}", self.0)
530 }
531}
532
533#[pyclass(eq, name = "RuleSet", module = "laddu", from_py_object)]
534#[derive(Clone, PartialEq)]
535pub struct PyRuleSet(pub RuleSet);
536
537#[pymethods]
538impl PyRuleSet {
539 #[new]
540 fn new() -> Self {
541 Self(RuleSet::default())
542 }
543
544 #[staticmethod]
545 fn angular() -> Self {
546 Self(RuleSet::angular())
547 }
548
549 #[staticmethod]
550 fn strong() -> Self {
551 Self(RuleSet::strong())
552 }
553
554 #[staticmethod]
555 fn electromagnetic() -> Self {
556 Self(RuleSet::electromagnetic())
557 }
558
559 #[staticmethod]
560 fn weak() -> Self {
561 Self(RuleSet::weak())
562 }
563
564 fn __repr__(&self) -> String {
565 format!("{:?}", self.0)
566 }
567}
568
569fn parse_rules(rules: Option<&Bound<'_, PyAny>>) -> PyResult<RuleSet> {
570 let Some(rules) = rules else {
571 return Ok(RuleSet::strong());
572 };
573 if let Ok(rules) = rules.extract::<PyRuleSet>() {
574 return Ok(rules.0);
575 }
576 if let Ok(name) = rules.extract::<String>() {
577 return match name.to_ascii_lowercase().as_str() {
578 "angular" => Ok(RuleSet::angular()),
579 "strong" => Ok(RuleSet::strong()),
580 "electromagnetic" | "em" => Ok(RuleSet::electromagnetic()),
581 "weak" => Ok(RuleSet::weak()),
582 _ => Err(LadduError::ParseError {
583 name,
584 object: "RuleSet".to_string(),
585 }
586 .into()),
587 };
588 }
589 Err(PyTypeError::new_err(
590 "rules must be a RuleSet or preset string",
591 ))
592}
593
594#[pyclass(name = "SelectionRules", module = "laddu", from_py_object)]
595#[derive(Clone)]
596pub struct PySelectionRules(pub SelectionRules);
597
598#[pymethods]
599impl PySelectionRules {
600 #[new]
601 #[pyo3(signature = (*, max_l=6, rules=None))]
602 fn new(max_l: u32, rules: Option<&Bound<'_, PyAny>>) -> PyResult<Self> {
603 Ok(Self(SelectionRules {
604 max_l: OrbitalAngularMomentum::integer(max_l),
605 rules: parse_rules(rules)?,
606 }))
607 }
608
609 #[staticmethod]
610 fn coupled_spins(
611 py: Python<'_>,
612 spin_1: &Bound<'_, PyAny>,
613 spin_2: &Bound<'_, PyAny>,
614 ) -> PyResult<Vec<PyQuantumNumber>> {
615 SelectionRules::coupled_spins(
616 parse_angular_momentum(spin_1)?,
617 parse_angular_momentum(spin_2)?,
618 )
619 .into_iter()
620 .map(|spin| angular_momentum_to_python(py, spin))
621 .collect()
622 }
623
624 fn allowed_partial_waves(
625 &self,
626 parent: &PyParticleProperties,
627 daughter_1: &PyParticleProperties,
628 daughter_2: &PyParticleProperties,
629 ) -> Vec<PyAllowedPartialWave> {
630 self.0
631 .allowed_partial_waves(&parent.0, (&daughter_1.0, &daughter_2.0))
632 .into_iter()
633 .map(PyAllowedPartialWave)
634 .collect()
635 }
636
637 fn __repr__(&self) -> String {
638 format!("{:?}", self.0)
639 }
640}
641
642#[pyfunction(name = "coupled_spins")]
644pub fn py_coupled_spins(
645 py: Python<'_>,
646 spin_1: &Bound<'_, PyAny>,
647 spin_2: &Bound<'_, PyAny>,
648) -> PyResult<Vec<PyQuantumNumber>> {
649 PySelectionRules::coupled_spins(py, spin_1, spin_2)
650}
651
652#[pyfunction(name = "allowed_partial_waves", signature = (parent, daughter_1, daughter_2, *, max_l=6, rules=None))]
654pub fn py_allowed_partial_waves(
655 parent: &PyParticleProperties,
656 daughter_1: &PyParticleProperties,
657 daughter_2: &PyParticleProperties,
658 max_l: u32,
659 rules: Option<&Bound<'_, PyAny>>,
660) -> PyResult<Vec<PyAllowedPartialWave>> {
661 Ok(PySelectionRules::new(max_l, rules)?
662 .0
663 .allowed_partial_waves(&parent.0, (&daughter_1.0, &daughter_2.0))
664 .into_iter()
665 .map(PyAllowedPartialWave)
666 .collect())
667}