Propagate field-corrected calibration points

Author: francois-drielsma

src/spine/calib/manager.py

@@ -42,6 +42,7 @@ def __init__(
         # Add the modules to a processor list in configuration order
         parsed = parse_module_config(cfg)
         names = [spec["name"] for spec in parsed.values()]
+        self.update_points = "field" in names
 
         # Make sure the essential calibration modules are present
         if not gain_applied and "recombination" in names and "gain" not in names:
@@ -83,7 +84,7 @@ def __call__(
         track: bool | None = None,
         meta: Meta | None = None,
         module_id: int | None = None,
-    ) -> NDArray[np.floating]:
+    ) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
         """Main calibration driver.
 
         Parameters
@@ -110,10 +111,13 @@ def __call__(
 
         Returns
         -------
+        np.ndarray
+            (N, 3) array of calibrated point coordinates
         np.ndarray
             (N) array of calibrated depositions in ADC, e- or MeV
         """
         # If necessary, convert all points to detector coordinates
+        orig_points = points
         if meta is not None:
             points = meta.to_cm(points, center=True)
         if module_id is not None:
@@ -136,6 +140,7 @@ def __call__(
             tpc_indexes = self.geo.get_closest_tpc_indexes(points)
 
         # Loop over the TPCs, apply the relevant calibration corrections
+        new_points = np.copy(points) if self.update_points else orig_points
         new_values = np.copy(values)
         for t in range(self.geo.tpc.num_chambers):
             # Restrict to the TPC of interest
@@ -164,6 +169,14 @@ def __call__(
                 self.watch.stop(key)
 
             # Append
+            if self.update_points:
+                new_points[tpc_indexes[t]] = tpc_points
             new_values[tpc_indexes[t]] = tpc_values
 
-        return new_values
+        if self.update_points:
+            if module_id is not None:
+                new_points = self.geo.translate(new_points, module_id, 0)
+            if meta is not None:
+                new_points = meta.to_px(new_points)
+
+        return new_points, new_values

src/spine/model/full_chain.py

@@ -950,7 +950,7 @@ def run_calibration(
                 run_id = run_info[b // rep].run if run_info is not None else None
 
                 # Calibrate voxel values
-                values_b = self.calibrator(
+                voxels_b, values_b = self.calibrator(
                     voxels_b,
                     values_b,
                     sources_b,
@@ -959,6 +959,10 @@ def run_calibration(
                     module_id=b % rep,
                 )
 
+                if self.calibrator.update_points:
+                    data.tensor[lower:upper, COORD_COLS] = torch.tensor(
+                        voxels_b, dtype=data.dtype, device=data.device
+                    )
                 data.tensor[lower:upper, VALUE_COL] = torch.tensor(
                     values_b, dtype=data.dtype, device=data.device
                 )

src/spine/post/reco/calo.py

@@ -164,9 +164,11 @@ def process(self, data):
 
                 # Apply calibration
                 if not self.do_tracking or part.shape != TRACK_SHP:
-                    depositions = self.calibrator(points, deps, sources, run_id)
+                    cal_points, depositions = self.calibrator(
+                        points, deps, sources, run_id
+                    )
                 else:
-                    depositions = self.calibrator(
+                    cal_points, depositions = self.calibrator(
                         points, deps, sources, run_id, track=True
                     )
 
@@ -176,22 +178,33 @@ def process(self, data):
                 else:
                     setattr(part, self.truth_dep_mode, depositions)
 
+                if self.calibrator.update_points:
+                    if not part.is_truth:
+                        part.points = cal_points
+                    else:
+                        setattr(part, self.truth_point_mode, cal_points)
+                    data[points_key][part.index] = cal_points
                 data[dep_key][part.index] = depositions
                 unass_mask[part.index] = False
 
             # Apply calibration corrections to unassociated depositions
             unass_index = np.where(unass_mask)[0]
             points = data[points_key][unass_index]
             depositions = data[dep_key][unass_index]
+            sources = None
             if source_key in data:
                 sources = data[source_key][unass_index]
 
-            data[dep_key][unass_index] = self.calibrator(
+            cal_points, depositions = self.calibrator(
                 points, depositions, sources, run_id
             )
+            if self.calibrator.update_points:
+                data[points_key][unass_index] = cal_points
+            data[dep_key][unass_index] = depositions
 
         # If requested, updated the depositions attribute of interactions
         for k in self.interaction_keys:
+            points_key = "points" if not "truth" in k else self.truth_point_key
             dep_key = "depositions" if not "truth" in k else self.truth_dep_key
             for inter in data[k]:
                 # Update depositions for the interaction
@@ -200,3 +213,10 @@ def process(self, data):
                     inter.depositions = depositions
                 else:
                     setattr(inter, self.truth_dep_mode, depositions)
+
+                if self.calibrator.update_points:
+                    points = data[points_key][inter.index]
+                    if not inter.is_truth:
+                        inter.points = points
+                    else:
+                        setattr(inter, self.truth_point_mode, points)

test/test_calib/test_manager.py

@@ -12,6 +12,9 @@ class FakeMeta:
     def to_cm(self, points, center=True):
         return points + 1.0
 
+    def to_px(self, points):
+        return points - 1.0
+
 
 def test_manager_parses_labels_and_explicit_names(monkeypatch, fake_geo):
     monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
@@ -50,7 +53,7 @@ def test_manager_applies_modules_in_config_order(monkeypatch, fake_geo):
     values = np.array([10.0, 10.0])
     sources = np.array([[0, 0], [0, 1]])
 
-    corrected = manager(points, values, sources=sources)
+    _, corrected = manager(points, values, sources=sources)
 
     assert np.allclose(corrected, [20.0, 30.0])
 
@@ -59,7 +62,7 @@ def test_manager_can_infer_tpc_indexes_without_sources(monkeypatch, fake_geo):
     monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
     manager = CalibrationManager(gain={"gain": [2.0, 3.0]})
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[1.0, 0.0, 0.0], [8.0, 0.0, 0.0]]),
         np.array([10.0, 10.0]),
     )
@@ -71,7 +74,7 @@ def test_manager_applies_meta_module_translation_and_empty_tpc(monkeypatch, fake
     monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
     manager = CalibrationManager(gain={"gain": [2.0, 3.0]})
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[1.0, 0.0, 0.0]]),
         np.array([10.0]),
         sources=np.array([[0, 0]]),
@@ -86,7 +89,7 @@ def test_manager_dispatches_lifetime_and_unknown_modules(monkeypatch, fake_geo):
     monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
     manager = CalibrationManager(lifetime={"lifetime": 10.0, "driftv": 2.0})
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[4.0, 0.0, 0.0]]),
         np.array([1.0]),
         sources=np.array([[0, 0]]),
@@ -115,7 +118,7 @@ def test_manager_dispatches_recombination(monkeypatch, fake_geo):
         recombination={"efield": 0.5},
     )
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[1.0, 0.0, 0.0]]),
         np.array([1000.0]),
         sources=np.array([[0, 0]]),
@@ -131,7 +134,7 @@ def test_manager_dispatches_transparency(monkeypatch, fake_geo, transparency_db)
         transparency={"transparency_db": str(transparency_db), "run_id": 100}
     )
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[0.0, 1.25, 1.25]]),
         np.array([12.0]),
         sources=np.array([[0, 1]]),
@@ -151,10 +154,47 @@ def test_manager_dispatches_field_before_lifetime(monkeypatch, fake_geo):
         lifetime={"lifetime": 10.0, "driftv": 2.0},
     )
 
-    corrected = manager(
+    _, corrected = manager(
         np.array([[1.0, 0.0, 0.0]]),
         np.array([1.0]),
         sources=np.array([[0, 0]]),
     )
 
     assert np.allclose(corrected, [np.exp(3.0 / 20.0)])
+
+
+def test_manager_can_return_field_corrected_points(monkeypatch, fake_geo):
+    monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
+    field_map = FieldMap(
+        np.full((1, 1, 1, 3), [2.0, 0.0, 0.0], dtype=float),
+        [[0.0, 10.0], [-1.0, 1.0], [-1.0, 1.0]],
+    )
+    manager = CalibrationManager(field={"field_map": field_map})
+
+    points, values = manager(
+        np.array([[1.0, 0.0, 0.0]]),
+        np.array([1.0]),
+        sources=np.array([[0, 0]]),
+    )
+
+    assert np.allclose(points, [[3.0, 0.0, 0.0]])
+    assert np.allclose(values, [1.0])
+
+
+def test_manager_returns_field_corrected_points_in_input_units(monkeypatch, fake_geo):
+    monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
+    field_map = FieldMap(
+        np.full((1, 1, 1, 3), [2.0, 0.0, 0.0], dtype=float),
+        [[0.0, 10.0], [-1.0, 1.0], [-1.0, 1.0]],
+    )
+    manager = CalibrationManager(field={"field_map": field_map})
+
+    points, values = manager(
+        np.array([[1.0, -0.5, -0.5]]),
+        np.array([1.0]),
+        sources=np.array([[0, 0]]),
+        meta=FakeMeta(),
+    )
+
+    assert np.allclose(points, [[3.0, -0.5, -0.5]])
+    assert np.allclose(values, [1.0])

test/test_post/test_calo.py

@@ -0,0 +1,113 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+import spine.calib.manager as manager_mod
+from spine.calib.field import FieldMap
+from spine.data.out import RecoInteraction, RecoParticle
+from spine.post.reco.calo import CalibrationProcessor
+
+
+class FakeTPC:
+    def __init__(self, anode_pos, center):
+        self.anode_pos = anode_pos
+        self.drift_axis = 0
+        self.drift_dir = np.array([1.0, 0.0, 0.0])
+        self.center = np.asarray(center, dtype=float)
+        self.dimensions = np.asarray((10.0, 10.0, 10.0), dtype=float)
+        self.boundaries = np.vstack(
+            (self.center - self.dimensions / 2.0, self.center + self.dimensions / 2.0)
+        ).T
+
+
+class FakeTPCSet:
+    num_modules = 1
+    num_chambers_per_module = 2
+    num_chambers = 2
+
+    def __init__(self):
+        self._tpcs = [
+            [
+                FakeTPC(0.0, center=(0.0, 0.0, 0.0)),
+                FakeTPC(10.0, center=(10.0, 0.0, 0.0)),
+            ]
+        ]
+        self.chambers = [tpc for module in self._tpcs for tpc in module]
+        lower = np.min(
+            np.vstack([tpc.boundaries[:, 0] for tpc in self.chambers]), axis=0
+        )
+        upper = np.max(
+            np.vstack([tpc.boundaries[:, 1] for tpc in self.chambers]), axis=0
+        )
+        self.modules = [
+            SimpleNamespace(
+                center=(lower + upper) / 2.0,
+                dimensions=upper - lower,
+                boundaries=np.vstack((lower, upper)).T,
+            )
+        ]
+
+    def __getitem__(self, index):
+        return self._tpcs[index]
+
+
+class FakeGeo:
+    def __init__(self):
+        self.tpc = FakeTPCSet()
+
+    def get_volume_index(self, sources, module_id, tpc_id):
+        return np.where((sources[:, 0] == module_id) & (sources[:, 1] == tpc_id))[0]
+
+    def get_closest_tpc_indexes(self, points):
+        return [np.where(points[:, 0] <= 5.0)[0], np.where(points[:, 0] > 5.0)[0]]
+
+    def translate(self, points, source_module, target_module):
+        return points + float(target_module - source_module)
+
+
+@pytest.fixture
+def fake_geo():
+    return FakeGeo()
+
+
+def test_calibration_processor_updates_field_corrected_points(monkeypatch, fake_geo):
+    monkeypatch.setattr(manager_mod.GeoManager, "get_instance", lambda: fake_geo)
+    field_map = FieldMap(
+        np.full((1, 1, 1, 3), [2.0, 0.0, 0.0], dtype=float),
+        [[0.0, 10.0], [-1.0, 1.0], [-1.0, 1.0]],
+    )
+    processor = CalibrationProcessor(
+        obj_type=("particle", "interaction"),
+        field={"field_map": field_map},
+    )
+
+    points = np.array([[1.0, 0.0, 0.0], [4.0, 0.0, 0.0]])
+    depositions = np.array([10.0, 20.0])
+    sources = np.array([[0, 0], [0, 0]])
+    particle = RecoParticle(
+        index=np.array([0], dtype=np.int32),
+        points=points[[0]].copy(),
+        depositions=depositions[[0]].copy(),
+        sources=sources[[0]].copy(),
+    )
+    interaction = RecoInteraction(
+        index=np.array([0, 1], dtype=np.int32),
+        points=points.copy(),
+        depositions=depositions.copy(),
+        sources=sources.copy(),
+    )
+    data = {
+        "points": points.copy(),
+        "depositions": depositions.copy(),
+        "sources": sources,
+        "reco_particles": [particle],
+        "reco_interactions": [interaction],
+    }
+
+    processor.process(data)
+
+    expected_points = points + np.array([2.0, 0.0, 0.0])
+    assert np.allclose(data["points"], expected_points)
+    assert np.allclose(particle.points, expected_points[[0]])
+    assert np.allclose(interaction.points, expected_points)