Handle zero and other weird division correctly when rescaling the per…

…centages to sum to 1 at the unit-level
washingtonpost · Oct 12, 2023 · 3aa0f5b · 3aa0f5b
1 parent 2254a65
commit 3aa0f5b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/elexsolver/TransitionSolver.py b/src/elexsolver/TransitionSolver.py
@@ -52,8 +52,10 @@ def _check_and_rescale(self, A: np.ndarray):
             if isinstance(A, np.ndarray):
                 for j in range(0, A.shape[1]):
                     A[:, j] /= A[:, j].sum()
+                return np.nan_to_num(A, nan=0, posinf=0, neginf=0)
             else:
                 # pandas.DataFrame()
                 for col in A.columns:
                     A[col] /= A[col].sum()
+                return A.fillna(0).replace(np.inf, 0).replace(-np.inf, 0)
         return A