Repository: scikit-hep/particle (language: python)

Issue:
Vectorised lookup
Hi there,

Love the project. Just wondering whether there is a way to use something like `from_pdgid` with an iterable of integers. Instantiating a new particle object for each is very expensive, and if I'm just trying to quickly look up a restricted selection of properties en masse it isn't very practical for my use case.

Kind regards,
Jacan.


Discussion / hints:
You don't instantiate a new object for each; they are already instantiated in a table and you are just pulling them out. The only cost is the search for the PDGID, which is a little expensive - we don't assume anything about PDGID and use it as a hash, for example - it's a linear search. If you have a limited subset you are interested in, you could cache these in a dictionary and that would be a bit faster; probably we could optimize this a bit.
Can you see if #349 is faster for you?
Hi. Also, can you explain a bit better what you are trying to do? If you want a retricted selection then calls such as `Particle.finditer(lambda p: p.pdgid.is_meson and p.pdgid.has_charm and p.spin_type==SpinType.PseudoScalar)` might be your friend. (BTW, I realise that we need to document this new iterator method better since the README rather talks about `Particle.findall`.)

Fix patch:
diff --git a/noxfile.py b/noxfile.py
index 33647111..a430d74f 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -17,9 +17,6 @@ def tests(session: nox.Session) -> None:
     session.install(".[test]")
     session.run(
         "pytest",
-        "--doctest-modules",
-        "--cov=particle",
-        "--cov-report=xml",
         *session.posargs,
     )
 
diff --git a/setup.cfg b/setup.cfg
index 86944458..c5ac8b86 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -58,18 +58,21 @@ where = src
 all =
     check-manifest>=0.42
     pytest
+    pytest-benchmark
     tabulate
     pandas;python_version>"3.4"
 dev =
     check-manifest>=0.42
     pytest
+    pytest-benchmark
     tabulate
     pandas;python_version>"3.4"
 test =
     pytest
+    pytest-benchmark
     pytest-cov
     tabulate
-    pandas;python_version>"3.4" and python_version<"3.10"
+    pandas;python_version>"3.4"
 
 [options.package_data]
 * = *.csv, *.fwf, *.mcd, *.py.typed
@@ -81,7 +84,9 @@ universal = 1
 testpaths =
     tests
 junit_family = xunit2
-addopts = '-Wd'
+addopts = --benchmark-disable -ra --showlocals
+filterwarnings =
+    error
 
 [flake8]
 max-complexity = 24
diff --git a/src/particle/converters/bimap.py b/src/particle/converters/bimap.py
index 31de13b7..fe7d19ab 100644
--- a/src/particle/converters/bimap.py
+++ b/src/particle/converters/bimap.py
@@ -247,7 +247,7 @@ def __getitem__(self, value):
             msg = "Matching {a}->{b} for input {v} not found !".format(
                 a=self.name_A, b=self.name_B, v=value
             )
-            raise MatchingIDNotFound(msg)
+            raise MatchingIDNotFound(msg)  # noqa: B904 Remove when dropping Python 2
 
     def __iter__(self):
         # type: () -> Iterator[str]
diff --git a/src/particle/particle/particle.py b/src/particle/particle/particle.py
index 9a4de3ef..aa243ea7 100644
--- a/src/particle/particle/particle.py
+++ b/src/particle/particle/particle.py
@@ -233,9 +233,12 @@ def __repr__(self):
             self=self, pdgid=int(self.pdgid), mass=self._str_mass()
         )
 
-    # Loaded table of entries
+    # Ordered loaded table of entries
     _table = None  # type: Optional[List[Particle]]
 
+    # Hash optimized table of particles
+    _hash_table = None  # type: Optional[Dict[int, Particle]]
+
     # Names of loaded tables
     _table_names = None  # type: Optional[List[str]]
 
@@ -557,10 +560,12 @@ def load_table(cls, filename=None, append=False, _name=None):
             cls.load_table(append=False)  # default load
         elif not append:
             cls._table = []
+            cls._hash_table = {}
             cls._table_names = []
 
         # Tell MyPy that this is true
         assert cls._table is not None
+        assert cls._hash_table is not None
         assert cls._table_names is not None
 
         if filename is None:
@@ -585,18 +590,10 @@ def load_table(cls, filename=None, append=False, _name=None):
         with open_file as f:
             r = csv.DictReader(line for line in f if not line.startswith("#"))
 
-            # Use a set to avoid making this O(n^2)
-            known_particles = set(cls._table)
-
             for v in r:
                 try:
                     value = int(v["ID"])
 
-                    # Replace the previous value if it exists
-                    # We can remove an int; ignore typing thinking we need a particle
-                    if value in known_particles:
-                        known_particles.remove(value)  # type: ignore
-
                     p = cls(
                         pdgid=value,
                         mass=float(v["Mass"]),
@@ -617,11 +614,14 @@ def load_table(cls, filename=None, append=False, _name=None):
                         quarks=v["Quarks"],
                         latex_name=v["Latex"],
                     )
-                    known_particles.add(p)
+
+                    # Replace the previous value if it exists
+                    cls._hash_table[value] = p
+
                 except ValueError:
                     pass
 
-        cls._table = sorted(known_particles)
+        cls._table = sorted(cls._hash_table.values())
 
     # The following __le__ and __eq__ needed for total ordering (sort, etc)
 
@@ -1030,11 +1030,17 @@ def from_pdgid(cls, value):
         if not is_valid(value):
             raise InvalidParticle("Input PDGID {} is invalid!".format(value))
 
-        for item in cls.all():
-            if item.pdgid == value:
-                return item
-        else:
-            raise ParticleNotFound("Could not find PDGID {}".format(value))
+        if not cls.table_loaded():
+            cls.load_table()
+
+        assert cls._hash_table is not None
+
+        try:
+            return cls._hash_table[int(value)]
+        except KeyError:
+            raise ParticleNotFound(  # noqa: B904  <- use from None when Python 2 is dropped
+                "Could not find PDGID {}".format(value)
+            )
 
     @classmethod
     def finditer(


Test patch:
diff --git a/tests/particle/test_performance.py b/tests/particle/test_performance.py
new file mode 100644
index 00000000..68f3cc26
--- /dev/null
+++ b/tests/particle/test_performance.py
@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+
+from particle import Particle, data
+
+
+def test_load_particle_table(benchmark):
+    benchmark(Particle.load_table, data.basepath / "particle2021.csv")
+
+
+def test_load_nuclei_append(benchmark):
+    def load_two():
+        Particle.load_table(data.basepath / "particle2021.csv")
+        Particle.load_table(data.basepath / "nuclei2020.csv", append=True)
+
+    benchmark(load_two)
+
+
+def test_from_pdgid(benchmark):
+    Particle.load_table(data.basepath / "particle2021.csv")
+    table = [int(s.pdgid) for s in Particle.all()]
+
+    def get_all(listing):
+        for pdgid in listing:
+            Particle.from_pdgid(pdgid)
+
+    benchmark(get_all, table)