msm.py: update covariance_mat checks to Enum

Co-authored-by: Marco Favorito <[email protected]>
bancaditalia · Dec 12, 2022 · 7a74765 · 7a74765
1 parent 80c7805
commit 7a74765
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/black_it/loss_functions/msm.py b/black_it/loss_functions/msm.py
@@ -179,10 +179,10 @@ def compute_loss_1d(
 
         g = real_mom_1d - sim_mom_1d
 
-        if self._covariance_mat == "identity":
+        if self._covariance_mat == _CovarianceMatrixType.IDENTITY:
             loss_1d = g.dot(g)
             return loss_1d
-        if self._covariance_mat == "inverse_variance":
+        if self._covariance_mat == _CovarianceMatrixType.INVERSE_VARIANCE:
             W = np.diag(
                 1.0 / np.mean((real_mom_1d[None, :] - ensemble_sim_mom_1d) ** 2, axis=0)
             )