Added lru caches to speed up validphys multi-replica initalization

conda-recipe/meta.yaml

@@ -57,7 +57,8 @@ requirements:
         - pineappl >=0.6.2
         - eko >=0.14.1
         - fiatlux
-        - curio >=1.0 # reportengine uses it but it's not in its dependencies 
+        - frozendict  # needed for caching of data loading
+        - curio >=1.0 # reportengine uses it but it's not in its dependencies
 
 test:
     requires:

pyproject.toml

@@ -59,7 +59,7 @@ vp-nextfitruncard = "validphys.scripts.vp_nextfitruncard:main"
 vp-hyperoptplot = "validphys.scripts.vp_hyperoptplot:main"
 vp-deltachi2 = "validphys.scripts.vp_deltachi2:main"
 
-[tool.poetry.dependencies] 
+[tool.poetry.dependencies]
 # Generic dependencies (i.e., validphys)
 python = "^3.9"
 matplotlib = ">=3.3.0,<3.8"
@@ -68,6 +68,7 @@ pandas = "<2"
 numpy = "*"
 validobj = "*"
 prompt_toolkit = "*"
+frozendict = "*"  # validphys: needed for caching of data loading
 # Reportengine (and its dependencies) need to be installed in a bit more manual way
 reportengine = { git = "https://github.com/NNPDF/reportengine", rev = "3bb2b1d"}
 ruamel_yaml = {version = "<0.18"}

validphys2/src/validphys/commondata.py

@@ -8,8 +8,9 @@
 """
 from reportengine import collect
 from validphys.commondataparser import load_commondata
+import functools
 
-
+@functools.lru_cache
 def loaded_commondata_with_cuts(commondata, cuts):
     """Load the commondata and apply cuts.
 

validphys2/src/validphys/config.py

@@ -45,10 +45,12 @@
 )
 from validphys.paramfits.config import ParamfitsConfig
 from validphys.plotoptions import get_info
+from validphys.utils import freeze_args
+from frozendict import frozendict
 import validphys.scalevariations
 
-log = logging.getLogger(__name__)
 
+log = logging.getLogger(__name__)
 
 class Environment(Environment):
     """Container for information to be filled at run time"""
@@ -1247,6 +1249,11 @@ def parse_default_filter_rules_recorded_spec_(self, spec):
     def parse_added_filter_rules(self, rules: (list, type(None)) = None):
         return rules
 
+    # Every parallel replica triggers a series of calls to this function,
+    # which should not happen since the rules are identical among replicas.
+    # E.g for NNPDF4.0 with 2 parallel replicas 693 calls, 3 parallel replicas 1001 calls...
+    @freeze_args
+    @functools.lru_cache
     def produce_rules(
         self,
         theoryid,
@@ -1286,7 +1293,7 @@ def produce_rules(
 
         if added_filter_rules:
             for i, rule in enumerate(added_filter_rules):
-                if not isinstance(rule, dict):
+                if not isinstance(rule, (dict, frozendict)):
                     raise ConfigError(f"added rule {i} is not a dict")
                 try:
                     rule_list.append(

validphys2/src/validphys/core.py

@@ -489,8 +489,9 @@ def load_cfactors(self):
 
         return [[parse_cfactor(c.open("rb")) for c in cfacs] for cfacs in self.cfactors]
 
+    @functools.lru_cache()
     def load_with_cuts(self, cuts):
-        """Load the fktable and apply cuts inmediately. Returns a FKTableData"""
+        """Load the fktable and apply cuts immediately. Returns a FKTableData"""
         return load_fktable(self).with_cuts(cuts)
 
 

validphys2/src/validphys/covmats.py

@@ -2,6 +2,7 @@
 matrices on different levels of abstraction
 """
 import logging
+import functools
 
 import numpy as np
 import pandas as pd
@@ -226,6 +227,7 @@ def dataset_inputs_covmat_from_systematics(
 
 
 @check_cuts_considered
+@functools.lru_cache
 def dataset_t0_predictions(dataset, t0set):
     """Returns the t0 predictions for a ``dataset`` which are the predictions
     calculated using the central member of ``pdf``. Note that if ``pdf`` has

validphys2/src/validphys/filters.py

@@ -6,13 +6,15 @@
 from importlib.resources import read_text
 import logging
 import re
+import functools
 
 import numpy as np
 
 from reportengine.checks import check, make_check
 from reportengine.compat import yaml
 from validphys.commondatawriter import write_commondata_to_file, write_systype_to_file
 import validphys.cuts
+from validphys.utils import freeze_args
 
 log = logging.getLogger(__name__)
 
@@ -561,6 +563,8 @@ def _make_point_namespace(self, dataset, idat) -> dict:
         return ns
 
 
+@freeze_args
+@functools.lru_cache
 def get_cuts_for_dataset(commondata, rules) -> list:
     """Function to generate a list containing the index
     of all experimental points that passed kinematic

validphys2/src/validphys/utils.py

@@ -5,12 +5,41 @@
 @author: Zahari Kassabov
 """
 import contextlib
+import functools
 import pathlib
 import shutil
 import tempfile
+from typing import Any, Sequence, Mapping, Hashable
 
 import numpy as np
 from validobj import ValidationError, parse_input
+from frozendict import frozendict
+
+
+def make_hashable(obj: Any):
+    # So that we don't infinitely recurse since frozenset and tuples
+    # are Sequences.
+    if isinstance(obj, Hashable):
+        return obj
+    elif isinstance(obj, Mapping):
+        return frozendict(obj)
+    elif isinstance(obj, Sequence):
+        return tuple([make_hashable(i) for i in obj])
+    else:
+        raise ValueError("Object is not hashable")
+
+
+def freeze_args(func):
+    """Transform mutable dictionary
+    Into immutable
+    Useful to be compatible with cache
+    """
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        args = tuple([make_hashable(arg) for arg in args])
+        kwargs = {k: make_hashable(v) for k, v in kwargs.items()}
+        return func(*args, **kwargs)
+    return wrapped
 
 
 def parse_yaml_inp(inp, spec, path):