From e4b3cb33c0ac76e65ef7121b8e3f44d39e22d362 Mon Sep 17 00:00:00 2001
From: afmagee42 <rzg0@cdc.gov>
Date: Mon, 25 Nov 2024 14:41:47 -0800
Subject: [PATCH] just do it all at once

---
 ngm/__init__.py | 63 +++++++++++++++++--------------------------------
 1 file changed, 21 insertions(+), 42 deletions(-)

diff --git a/ngm/__init__.py b/ngm/__init__.py
index 7e175dc..2d96460 100644
--- a/ngm/__init__.py
+++ b/ngm/__init__.py
@@ -1,59 +1,38 @@
+from collections import namedtuple
+
 import numpy as np
 
+NgmSummary = namedtuple(
+    "NgmSummary", ["r_eff", "infectious_dist", "outcome_dist"]
+)
+
 
 class CoreGroupNgm:
     def __init__(
         self,
-        num_groups: int,
         r0_matrix: np.ndarray,
         # Other parameters go here
     ):
-        self.k = self.make_k(num_groups, r0_matrix)
+        assert (
+            len(r0_matrix) == 2 and r0_matrix.shape[0] == r0_matrix.shape[1]
+        ), "r0_matrix must be square 2-D matrix"
+        self.k = self.make_k(r0_matrix)
         # Other parameters get stored here
         raise NotImplementedError()
 
-    def make_k(self, num_groups: int, r0_matrix: np.ndarray):
+    def make_k(self, r0_matrix: np.ndarray):
         raise NotImplementedError()
 
-    def calculate_r_eff(self) -> float:
-        r"""
-        Compute $R_e$
-
-        Returns
-        -------
-        float
-            $R_e$ via the spectral radius of the next generation matrix
-        """
-        # Could consider caching eigenvalues/vectors, but always re-computing feels safer
-        eval = np.linalg.eig(self.k).eigenvalues
-        # eval at index (via argmax) of maximum absolute value
-        # @TODO: do we need to check for imaginary components?
-        return eval[np.argmax(np.abs(eval))]
-
-    def calculate_infectious_distribution(self) -> np.ndarray:
-        r"""
-        Compute distribution of infections
-
-        Returns
-        -------
-        np.ndarray
-            The PMF on infections, I think?
-        """
-        # evec = np.linalg.eig(self.k).eigenvectors
-        raise NotImplementedError()
+    def summarize(self) -> NgmSummary:
+        raise NotImplementedError
 
-    def calculate_severe_outcomes(self, inf_dist: np.ndarray):
-        r"""
-        Compute distribution of severe outcomes from the distribution of infections
+        # eigendecomp = np.linalg.eig(self.k)
 
-        Parameters
-        ----------
-        inf_dist : np.ndarray
-            Output of self.calculate_infectious_distribution()
+        # # Index of dominant eigenvalue determines r_eff and infectious distribution
+        # dom = np.argmax(np.abs(eigendecomp.eigenvalues))
 
-        Returns
-        -------
-        np.ndarray
-            The PMF on severe outcomes, I think?
-        """
-        raise NotImplementedError()
+        # return NgmSummary(
+        #     r_eff=eigendecomp.eigenvalues[dom],
+        #     infectious_dist=None, #TBD: is this just eigendecomp.eigenvectors[dom] / eigendecomp.eigenvectors[dom].sum()?
+        #     outcome_dist=None, #TBD
+        # )