unified algorithms params

Author: MaximeLagrange

muograph/reconstruction/asr.py

@@ -42,7 +42,7 @@ class ASR(AbsSave, AbsVoxelInferer):
     _n_mu_per_vox: Optional[Tensor] = None  # (Nx, Ny, Nz)
     _recompute_preds = True
 
-    _asr_params: ASRParams = ASRParams()
+    _params: ASRParams = ASRParams()
 
     _vars_to_save = ["triggered_voxels"]
 
@@ -345,7 +345,7 @@ def _find_triggered_voxels(
 
     @staticmethod
     def get_name_from_params(
-        asr_params: ASRParams,
+        params: ASRParams,
     ) -> str:
         r"""
         Returns a string representing the ASR configuration based on its parameters.
@@ -361,12 +361,12 @@ def get_partial_name_args(func: partial) -> str:
                 func_name += "_{}={}".format(arg, values[i])
             return func_name
 
-        method = "method_{}_".format(get_partial_name_args(asr_params.score_method))  # type: ignore
-        dtheta = "{:.2f}_{:.2f}_rad_".format(asr_params.dtheta_range[0], asr_params.dtheta_range[1])  # type: ignore
-        dp = "{:.0f}_{:.0f}_MeV_".format(asr_params.p_range[0], asr_params.p_range[1])  # type: ignore
-        use_p = "use_p_{}".format(asr_params.use_p)
-        p_clamp = "_pclamp_{:.3f}".format(asr_params.p_clamp) if asr_params.use_p else ""
-        dtheta_clamp = "_dthetaclamp_{:.3f}".format(asr_params.dtheta_clamp)
+        method = "method_{}_".format(get_partial_name_args(params.score_method))  # type: ignore
+        dtheta = "{:.2f}_{:.2f}_rad_".format(params.dtheta_range[0], params.dtheta_range[1])  # type: ignore
+        dp = "{:.0f}_{:.0f}_MeV_".format(params.p_range[0], params.p_range[1])  # type: ignore
+        use_p = "use_p_{}".format(params.use_p)
+        p_clamp = "_pclamp_{:.3f}".format(params.p_clamp) if params.use_p else ""
+        dtheta_clamp = "_dthetaclamp_{:.3f}".format(params.dtheta_clamp)
 
         asr_name = method + dtheta + dp + use_p + p_clamp + dtheta_clamp
         asr_name = asr_name.replace(".", "p")
@@ -432,11 +432,11 @@ def get_xyz_voxel_pred(self) -> Tensor:
             [[[] for _ in range(self.voi.n_vox_xyz[2])] for _ in range(self.voi.n_vox_xyz[1])] for _ in range(self.voi.n_vox_xyz[0])
         ]
 
-        dtheta_max = torch.quantile(self.tracks.dtheta, q=self.asr_params.dtheta_clamp)
+        dtheta_max = torch.quantile(self.tracks.dtheta, q=self.params.dtheta_clamp)
         dtheta = torch.clamp(self.tracks.dtheta, max=dtheta_max)
 
-        if self._asr_params.use_p:
-            p_max = torch.quantile(self.tracks.p, q=self.asr_params.p_clamp)
+        if self._params.use_p:
+            p_max = torch.quantile(self.tracks.p, q=self.params.p_clamp)
             p = torch.clamp(self.tracks.p, max=p_max)
             score = np.log(0.0000001 + dtheta.detach().cpu().numpy() * p.detach().cpu().numpy())
         else:
@@ -445,14 +445,14 @@ def get_xyz_voxel_pred(self) -> Tensor:
             # score = self.theta_xy_in[0].detach().cpu().numpy()
             # score = self.tracks.theta_in.detach().cpu().numpy()
 
-        mask_E = (self.tracks.E >= self.asr_params.p_range[0]) & (  # type: ignore
-            self.tracks.E <= self.asr_params.p_range[1]  # type: ignore
+        mask_E = (self.tracks.E >= self.params.p_range[0]) & (  # type: ignore
+            self.tracks.E <= self.params.p_range[1]  # type: ignore
         )
-        mask_theta = (self.tracks.dtheta >= self.asr_params.dtheta_range[0]) & (  # type: ignore
-            self.tracks.dtheta <= self.asr_params.dtheta_range[1]  # type: ignore
+        mask_theta = (self.tracks.dtheta >= self.params.dtheta_range[0]) & (  # type: ignore
+            self.tracks.dtheta <= self.params.dtheta_range[1]  # type: ignore
         )
 
-        if self.asr_params.use_p:  # type: ignore
+        if self.params.use_p:  # type: ignore
             mask = mask_E & mask_theta
         else:
             mask = mask_E
@@ -472,14 +472,14 @@ def get_xyz_voxel_pred(self) -> Tensor:
             for j in range(self.voi.n_vox_xyz[1]):
                 for k in range(self.voi.n_vox_xyz[2]):
                     if score_list[i][j][k] != []:
-                        vox_density_preds[i, j, k] = self.asr_params.score_method(score_list[i][j][k])  # type: ignore
+                        vox_density_preds[i, j, k] = self.params.score_method(score_list[i][j][k])  # type: ignore
                         self.n_mu_per_vox_test[i, j, k] = len(score_list[i][j][k])
         if vox_density_preds.isnan().any():
             raise ValueError("Prediction contains NaN values")
         self.score_list = score_list
         self._recompute_preds = False
 
-        if self.asr_params.use_p:
+        if self.params.use_p:
             return torch.exp(vox_density_preds)
         else:
             return vox_density_preds
@@ -599,29 +599,29 @@ def theta_xy_out(self) -> Tuple[Tensor, Tensor]:
         return (self.tracks.theta_xy_out[0], self.tracks.theta_xy_out[1])
 
     @property
-    def asr_params(self) -> ASRParams:
+    def params(self) -> ASRParams:
         r"""
         The parameters of the ASR algorithm.
         """
-        return self._asr_params
+        return self._params
 
-    @asr_params.setter
-    def asr_params(self, value: ASRParams) -> None:
+    @params.setter
+    def params(self, value: ASRParams) -> None:
         r"""
         Sets the parameters of the ASR algorithm.
         Args:
             - Dict containing the parameters name and value. Only parameters with
             valid name and non `None` values will be updated.
         """
         if not isinstance(value, ASRParams):
-            raise TypeError("asr_params must be an instance of ASRParams")
+            raise TypeError("params must be an instance of ASRParams")
 
-        if not hasattr(self, "_asr_params") or self._asr_params is None:
-            self._asr_params = ASRParams()
+        if not hasattr(self, "_params") or self._params is None:
+            self._params = ASRParams()
 
         for key, val in value.__dict__.items():
             if val is not None:
-                setattr(self._asr_params, key, val)
+                setattr(self._params, key, val)
 
         self._recompute_preds = True
 
@@ -659,4 +659,4 @@ def name(self) -> str:
         r"""
         The name of the ASR configuration based on its parameters.
         """
-        return self.get_name_from_params(self.asr_params)
+        return self.get_name_from_params(self.params)

muograph/reconstruction/binned_clustered.py

@@ -9,7 +9,7 @@
 
 from muograph.utils.device import DEVICE
 from muograph.tracking.tracking import TrackingMST
-from muograph.reconstruction.poca import POCA
+from muograph.reconstruction.poca import POCA, POCAParams
 from muograph.volume.volume import Volume
 from muograph.reconstruction.voxel_inferer import AbsVoxelInferer
 
@@ -24,20 +24,17 @@
 
 
 @dataclass
-class BCAParams:
+class BCAParams(POCAParams):
     n_max_per_vox: int = 50
     n_min_per_vox: int = 3
     score_method: partial = partial(torch.quantile, q=0.5)
     metric_method: partial = partial(torch.log)
     p_range: Tuple[float, float] = (0.0, 10000000)
     dtheta_range: Tuple[float, float] = (0.0, math.pi / 3)
-    use_p: bool = False
-    p_clamp: float = 0.999
-    dtheta_clamp: float = 0.999
 
 
-class BCA(POCA, AbsVoxelInferer):
-    _bca_params: BCAParams = BCAParams()
+class BCA(POCA[BCAParams], AbsVoxelInferer):
+    _params: BCAParams = BCAParams()
 
     _vars_to_save = ["xyz_voxel_pred", "n_poca_per_vox"]
 
@@ -386,7 +383,7 @@ def get_xyz_voxel_pred(self) -> Tensor:
 
         Compute voxel-wise scattering density predictions.
 
-        Uses parameters stored in `_bca_params`. The algorithm calculates:
+        Uses parameters stored in `_params`. The algorithm calculates:
             - Scattering density predictions (`pred`).
             - Number of POCA points used for each voxel's prediction (`hit_per_voxel`).
 
@@ -405,41 +402,41 @@ def get_xyz_voxel_pred(self) -> Tensor:
             self.nhit,
             self.nhit_rejected,
         ) = self.compute_low_theta_events_voxel_wise_mask(
-            n_max_per_voxel=int(self.bca_params.n_max_per_vox),  # type: ignore
+            n_max_per_voxel=int(self.params.n_max_per_vox),  # type: ignore
             voi=self.voi,
             bca_indices=self.bca_indices,
             dtheta=self.bca_tracks.dtheta,
         )
         self._filter_events(self.mask)
 
         # momentum cut
-        if self.bca_params.use_p:
-            p_mask = (self.bca_tracks.E > self.bca_params.p_range[0]) & (  # type: ignore
-                self.bca_tracks.E < self.bca_params.p_range[1]  # type: ignore
+        if self.params.use_p:
+            p_mask = (self.bca_tracks.E > self.params.p_range[0]) & (  # type: ignore
+                self.bca_tracks.E < self.params.p_range[1]  # type: ignore
             )
         else:
             p_mask = torch.ones_like(self.bca_tracks.dtheta, dtype=torch.bool, device=DEVICE)
 
         # scattering angle cut
-        dtheta_mask = (self.bca_tracks.dtheta > self.bca_params.dtheta_range[0]) & (  # type: ignore
-            self.bca_tracks.dtheta < self.bca_params.dtheta_range[1]  # type: ignore
+        dtheta_mask = (self.bca_tracks.dtheta > self.params.dtheta_range[0]) & (  # type: ignore
+            self.bca_tracks.dtheta < self.params.dtheta_range[1]  # type: ignore
         )
 
         # apply dtheta, p cuts
         self._filter_events(mask=p_mask & dtheta_mask)
 
         # prepare variables
-        p_max = torch.quantile(self.bca_tracks.p, q=self.bca_params.p_clamp)
+        p_max = torch.quantile(self.bca_tracks.p, q=self.params.p_clamp)
         p = torch.clamp(self.bca_tracks.p, max=p_max)
 
-        dtheta_max = torch.quantile(self.bca_tracks.dtheta, q=self.bca_params.dtheta_clamp)
+        dtheta_max = torch.quantile(self.bca_tracks.dtheta, q=self.params.dtheta_clamp)
         dtheta = torch.clamp(self.bca_tracks.dtheta, max=dtheta_max)
 
         # compute voxels distribution
         self.score_list = self.compute_voxels_distribution(
-            metric_method=self.bca_params.metric_method,  # type: ignore
-            use_p=self.bca_params.use_p,  # type: ignore
-            n_min_per_vox=self.bca_params.n_min_per_vox,  # type: ignore
+            metric_method=self.params.metric_method,  # type: ignore
+            use_p=self.params.use_p,  # type: ignore
+            n_min_per_vox=self.params.n_min_per_vox,  # type: ignore
             voi=self.voi,
             momentum=p,
             bca_indices=self.bca_indices,
@@ -448,7 +445,10 @@ def get_xyz_voxel_pred(self) -> Tensor:
         )
 
         # compute fina scores
-        pred, self._hit_per_voxel = self.compute_final_scores(score_list=self.score_list, score_method=self.bca_params.score_method)  # type: ignore
+        pred, self._hit_per_voxel = self.compute_final_scores(score_list=self.score_list, score_method=self.params.score_method)  # type: ignore
+
+        pred_max = torch.quantile(pred, q=self.params.preds_clamp)
+        pred = torch.clamp(pred, max=pred_max)
 
         self._recompute_preds = False
 
@@ -471,16 +471,17 @@ def get_partial_name_args(func: partial) -> str:
                 func_name += "_{}={}".format(arg, values[i])
             return func_name
 
-        method = "method_{}_".format(get_partial_name_args(self.bca_params.score_method))  # type: ignore
-        metric = "metric_{}_".format(get_partial_name_args(self.bca_params.metric_method))  # type: ignore
-        dtheta = "{:.1f}_{:.1f}_mrad_".format(self.bca_params.dtheta_range[0] * 1000, self.bca_params.dtheta_range[1] * 1000)  # type: ignore
-        dp = "{:.0f}_{:.0f}_MeV_".format(self.bca_params.p_range[0], self.bca_params.p_range[1])  # type: ignore
-        n_min_max = "n_min_max_{}_{}_".format(self.bca_params.n_min_per_vox, self.bca_params.n_max_per_vox)
-        use_p = "use_p_{}".format(self.bca_params.use_p)
-        p_clamp = "_pclamp_{:.3f}".format(self.bca_params.p_clamp) if self.bca_params.use_p else ""
-        dtheta_clamp = "_dthetaclamp_{:.3f}".format(self.bca_params.dtheta_clamp)
-
-        bca_name = method + metric + dtheta + dp + n_min_max + use_p + p_clamp + dtheta_clamp
+        method = "method_{}_".format(get_partial_name_args(self.params.score_method))  # type: ignore
+        metric = "metric_{}_".format(get_partial_name_args(self.params.metric_method))  # type: ignore
+        dtheta = "{:.1f}_{:.1f}_mrad_".format(self.params.dtheta_range[0] * 1000, self.params.dtheta_range[1] * 1000)  # type: ignore
+        dp = "{:.0f}_{:.0f}_MeV_".format(self.params.p_range[0], self.params.p_range[1])  # type: ignore
+        n_min_max = "n_min_max_{}_{}_".format(self.params.n_min_per_vox, self.params.n_max_per_vox)
+        use_p = "use_p_{}".format(self.params.use_p)
+        p_clamp = "_pclamp_{:.3f}".format(self.params.p_clamp) if self.params.use_p else ""
+        dtheta_clamp = "_dthetaclamp_{:.3f}".format(self.params.dtheta_clamp)
+        preds_clamp = "_preds_clamp_{:.3f}".format(self.params.preds_clamp)
+
+        bca_name = method + metric + dtheta + dp + n_min_max + use_p + p_clamp + dtheta_clamp + preds_clamp
         bca_name = bca_name.replace(".", "p")
         return bca_name
 
@@ -492,29 +493,29 @@ def name(self) -> str:
         return self.get_bca_name()
 
     @property
-    def bca_params(self) -> BCAParams:
+    def params(self) -> BCAParams:
         r"""
         The parameters of the bca algorithm.
         """
-        return self._bca_params
+        return self._params
 
-    @bca_params.setter
-    def bca_params(self, value: BCAParams) -> None:
+    @params.setter
+    def params(self, value: BCAParams) -> None:
         r"""
         Sets the parameters of the bca algorithm.
         Args:
             - Dict containing the parameters name and value. Only parameters with
             valid name and non `None` values wil be updated.
         """
         if not isinstance(value, BCAParams):
-            raise TypeError("bca_params must be an instance of BCAParams")
+            raise TypeError("params must be an instance of BCAParams")
 
-        if not hasattr(self, "_bca_params") or self._bca_params is None:
-            self._bca_params = BCAParams()
+        if not hasattr(self, "_params") or self._params is None:
+            self._params = BCAParams()
 
         for key, val in value.__dict__.items():
             if val is not None:
-                setattr(self._bca_params, key, val)
+                setattr(self._params, key, val)
 
         self._recompute_preds = True