diff --git a/src/particle/particle/particle.py b/src/particle/particle/particle.py index 103d7fb6..e8ccd8fe 100644 --- a/src/particle/particle/particle.py +++ b/src/particle/particle/particle.py @@ -50,6 +50,28 @@ class InvalidParticle(RuntimeError): pass +# The proton and the neutron (and their anti-particles) have two possible PDG ID representations, +# a) a particle ("bag of quarks") or b) a nucleus. +_NON_UNIQUE_PDGIDS = { + 2112: 1000000010, + 2212: 1000010010, +} +for pdgid1, pdgid2 in list(_NON_UNIQUE_PDGIDS.items()): + # add reverse lookup + _NON_UNIQUE_PDGIDS[pdgid2] = pdgid1 + # add anti-particles + _NON_UNIQUE_PDGIDS[-pdgid1] = -pdgid2 + _NON_UNIQUE_PDGIDS[-pdgid2] = -pdgid1 + +# lookup for hash and from_name which representation is the preferred one +# this results in the "bag-of-quarks" representation +_PREFERRED_PDGID = {} +for pdgid1, pdgid2 in _NON_UNIQUE_PDGIDS.items(): + sign = -1 if pdgid1 < 0 else 1 + _PREFERRED_PDGID[pdgid1] = sign * min(abs(pdgid1), abs(pdgid2)) + _PREFERRED_PDGID[pdgid2] = _PREFERRED_PDGID[pdgid1] + + def _isospin_converter(isospin: str) -> float | None: vals: dict[str | None, float | None] = { "0": 0.0, @@ -604,12 +626,25 @@ def __lt__(self, other: Particle | int) -> bool: return int(self.pdgid) < other def __eq__(self, other: object) -> bool: + """ + Compare with another Particle instance based on PDG IDs. + + Note + ---- + Ensure the comparison also works for the special cases of the proton and the neutron, + which have two PDG ID representations as particles or nuclei. + """ if isinstance(other, Particle): - return self.pdgid == other.pdgid + other = other.pdgid + + if self.pdgid in _NON_UNIQUE_PDGIDS: + return other in {self.pdgid, _NON_UNIQUE_PDGIDS[self.pdgid]} + return self.pdgid == other - # Only one particle can exist per PDGID number def __hash__(self) -> int: + if self.pdgid in _PREFERRED_PDGID: + return hash(_PREFERRED_PDGID[self.pdgid]) return hash(self.pdgid) # Shared with PDGID @@ -992,6 +1027,13 @@ def from_name(cls: type[Self], name: str) -> Self: ParticleNotFound If no particle matches the input name uniquely and exactly. """ + # special handling for the particles with two possible pdgids + if name in {"p", "n", "p~", "n~"}: + + def find_preferred_id(p: Self) -> bool: + return int(p.pdgid) in _PREFERRED_PDGID.values() + + return next(filter(find_preferred_id, cls.finditer(name=name))) try: (particle,) = cls.finditer( name=name @@ -1083,7 +1125,10 @@ def from_nucleus_info( pdgid = int(1e9 + l_strange * 1e5 + z * 1e4 + a * 10 + i) if anti: - return cls.from_pdgid(-pdgid) + pdgid = -pdgid + + # replace nucleon ids of single hadrons with quark version + pdgid = _PREFERRED_PDGID.get(pdgid, pdgid) return cls.from_pdgid(pdgid) diff --git a/tests/particle/test_particle.py b/tests/particle/test_particle.py index 9c28e4bb..cafb3f43 100644 --- a/tests/particle/test_particle.py +++ b/tests/particle/test_particle.py @@ -172,6 +172,15 @@ def test_from_nucleus_info(): assert p.pdgid == -1000010020 +def test_from_nucleus_info_special_cases(): + """ + The proton and the neutron should return the preferred quark representation + rather than the representation as a nucleus. + """ + assert Particle.from_nucleus_info(a=1, z=1).pdgid == 2212 + assert Particle.from_nucleus_info(a=1, z=0).pdgid == 2112 + + def test_from_nucleus_info_ParticleNotFound(): with pytest.raises(ParticleNotFound): _ = Particle.from_nucleus_info(z=999, a=999) @@ -706,3 +715,65 @@ def test_decfile_style_names(name, pid): @pytest.mark.parametrize(("name", "pid"), decfile_style_names) def test_evtgen_name(name, pid): # noqa: ARG001 assert Particle.from_evtgen_name(name).evtgen_name == name + + +@pytest.mark.parametrize( + ("pdgid1", "pdgid2"), + [ + pytest.param(2212, 1000010010, id="p"), + pytest.param(2112, 1000000010, id="n"), + pytest.param(-2212, -1000010010, id="p~"), + pytest.param(-2112, -1000000010, id="n~"), + ], +) +def test_eq_non_unique_pdgids(pdgid1, pdgid2): + """The proton and the neutron have two PDG ID representations. Make sure they still compare equal.""" + + p1 = Particle.from_pdgid(pdgid1) + p2 = Particle.from_pdgid(pdgid2) + assert p1.pdgid != p2.pdgid + assert p1 == p2 + assert p2 == p1 + assert hash(p1) == hash(p2) + + +@pytest.mark.parametrize( + ("name", "pdgid"), + [ + pytest.param("p", 2212, id="p"), + pytest.param("n", 2112, id="n"), + pytest.param("p~", -2212, id="p~"), + pytest.param("n~", -2112, id="n~"), + ], +) +def test_from_name_non_unique_pdgids(name, pdgid): + """ + Test that Particle.from_name works for p and n, returning the preferred quark representation + rather than the representation as a nucleus. + """ + + p = Particle.from_name(name) + assert p.name == name + assert p.pdgid == pdgid + + +@pytest.mark.parametrize("sign", [1, -1]) +def test_particle_hash(sign): + proton1 = Particle.from_pdgid(sign * 2212) + proton2 = Particle.from_pdgid(sign * 1000010010) + neutron1 = Particle.from_pdgid(sign * 2112) + neutron2 = Particle.from_pdgid(sign * 1000000010) + + s1 = {proton1, neutron1} + + assert proton1 in s1 + assert proton2 in s1 + assert neutron1 in s1 + assert neutron2 in s1 + + assert len({proton1, proton2}) == 1 + assert len({neutron1, neutron2}) == 1 + + d = {proton1: 5, neutron2: 3} + assert d[proton2] == 5 + assert d[neutron1] == 3