Make sure neutron, proton and their anti-particles compare equal (#622)

* Make sure neutron, proton and their anti-particles compare equal in different representations

* Add special casing for from_name

* Add more tests for Particle.__hash__

* Return preferred pdgid in from_nucleus_info

* Fix test docstring

* Test commutation

* Apply suggestions from code review

* Fix indentation

* Apply suggestions from code review

* Fix docstring

* Docstring fix in test_particle.py

* style: pre-commit fixes

---------

Co-authored-by: Eduardo Rodrigues <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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)
 

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