From 3aa0f5bdeb351342b0db3f09fe11869181867644 Mon Sep 17 00:00:00 2001
From: Diane Napolitano <diane.napolitano@washpost.com>
Date: Thu, 12 Oct 2023 16:55:13 -0400
Subject: [PATCH] Handle zero and other weird division correctly when rescaling
 the percentages to sum to 1 at the unit-level

---
 src/elexsolver/TransitionSolver.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/elexsolver/TransitionSolver.py b/src/elexsolver/TransitionSolver.py
index 2dae4c37..945da4ea 100644
--- 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