binning generation tests passed; factory interface operational; pytes…

…ts integration in setup

ddmisid/utils/binning.py

@@ -11,27 +11,31 @@ class PIDCalibAliasFactory:
     """Factory to process variable names based on year and species."""
 
     @staticmethod
-    def process_variable(var: str, year: str, spc: str) -> str:
+    def process_variable(var: str, year: str, spc: str | None = None) -> str:
         """
         Apply year-based processing to variable names.
         Source: https://twiki.cern.ch/twiki/bin/view/LHCb/PIDCalibPackage
         """
         # Run 1 vs Run 2 aliases
-        if year in ["2016", "2017", "2018"]:
-            if var in ["P", "PT"]:
-                var = f"Brunel_{var}"
-            elif var == "nTracks":
-                var = f"{var}_Brunel"
-        elif year not in ["2011", "2012", "2015"]:
-            raise ValueError(f"Year {year} not recognized.")
-
-        # bespoke handling of electrons
-        if spc == "e_B_Jpsi" and year == "2016":
-            if var in ["Brunel_P", "Brunel_PT"]:
-                var = var.replace("Brunel_", "")
-            elif var == "nTracks_Brunel":
-                var = var.replace("_Brunel", "")
-
+        match var:
+            case var if var in ["P", "PT", "ETA"]:
+                if year in ["2016", "2017", "2018"]:
+                    var = f"Brunel_{var}"
+                elif year in ["2011", "2012", "2015"]:
+                    pass
+                else:
+                    raise ValueError(f"Year {year} not supported for {var} assignment.")
+
+            case var if var == "nTracks":
+                if year in ["2016", "2017", "2018"]:
+                    var = f"{var}_Brunel"
+                elif year in ["2011", "2012", "2015"]:
+                    pass
+                else:
+                    raise ValueError(f"Year {year} not supported for {var} assignment.")
+
+            case _:
+                raise ValueError(f"Variable {var} not supported (yet). If appropriate, please add to the `PIDCalibAliasFactory` in ddmisid/utils/binning.py.")
         return var
 
 
@@ -63,7 +67,7 @@ def fetch_info(self) -> Tuple[Dict[str, str], Dict[str, Dict[str, List[float]]]]
 
     def build(self, year: str, outdir: str = ".data", verbose: bool = True) -> None:
         """Build the JSON spec file for the supplied years and species."""
-        bins, species = self.fetch_info()
+        species, bins = self.fetch_info()
 
         for spc in species.values():
             binning2json = {spc: {}}
@@ -82,6 +86,7 @@ def build(self, year: str, outdir: str = ".data", verbose: bool = True) -> None:
                 )
             else:
                 outfile_path = Path(f"{outdir}/binning_{year}/{spc}.json")
+
             outfile_path.parent.mkdir(parents=True, exist_ok=True)
             with open(outfile_path, "w") as f:
                 json.dump(binning2json, f, indent=4)

setup.py

@@ -27,7 +27,8 @@
         "cabinetry",
         "polars",
         "pydantic",
-        "snakemake",
+        "pytest",
+        # "snakemake", # recommended install via mamba and running `pip install -e .` after loading mamba environment
     ],
     entry_points="""
         [console_scripts]

tests/test_binning_generator.py

@@ -0,0 +1,71 @@
+import pytest
+from pathlib import Path
+import json
+from ddmisid.utils import DefaultBinningGenerator, PIDCalibAliasFactory
+
+
+@pytest.fixture
+def setup_binning():
+    """Test binning and species dictionaries."""
+    example_species = {
+        "kaon": "K",
+    }
+
+    example_binning = {
+        "P": [10000.0, 20000.0, 30000.0, 40000.0, 100000.0],
+        "PT": [500.0, 1000.0, 1500.0, 2000.0, 4000.0],
+        "ETA": [2.0, 2.5, 3.0, 3.5, 4.0],
+        "nTracks": [0, 100, 200, 300, 400],
+    }
+
+    binning_generator = DefaultBinningGenerator(
+        species=example_species,
+        binning=example_binning,
+        binning_alias="TestBinning",
+    )
+
+    yield binning_generator
+
+    # Clean up after test
+    output_path = Path("tests/data/binning_2016/TestBinning.json")
+    if output_path.exists():
+        output_path.unlink()
+
+
+def test_process_variable():
+    """Verify PIDCalib2 alias compliance."""
+    assert PIDCalibAliasFactory.process_variable("P", "2011", "K") == "P"
+    assert PIDCalibAliasFactory.process_variable("PT", "2018", "K") == "Brunel_PT"
+    assert (
+        PIDCalibAliasFactory.process_variable("nTracks", "2016", "K")
+        == "nTracks_Brunel"
+    )
+    assert (
+        PIDCalibAliasFactory.process_variable("ETA", "2016", "e_B_Jpsi") == "Brunel_ETA"
+    )
+    assert PIDCalibAliasFactory.process_variable("ETA", "2012", "e_B_Jpsi") == "ETA"
+    assert (
+        PIDCalibAliasFactory.process_variable("nTracks", "2011", "e_B_Jpsi")
+        == "nTracks"
+    )
+
+
+def test_binning_build(setup_binning):
+    """Test binning build method."""
+    expected_binning = {
+        "K": {
+            "Brunel_P": [10000.0, 20000.0, 30000.0, 40000.0, 100000.0],
+            "Brunel_PT": [500.0, 1000.0, 1500.0, 2000.0, 4000.0],
+            "Brunel_ETA": [2.0, 2.5, 3.0, 3.5, 4.0],
+            "nTracks_Brunel": [0, 100, 200, 300, 400],
+        }
+    }
+
+    # Generate binning and write to JSON file
+    setup_binning.build(year="2016", outdir="tests/data", verbose=False)
+
+    # Read JSON file back in, and compare to expected binning
+    with open(Path("tests/data/binning_2016/TestBinning.json")) as f:
+        binning = json.load(f)
+
+    assert binning == expected_binning