Merge pull request #2 from becksgld/type-hints

Added type hints

Author: bmtgoncalves

src/epidemik/EpiModel.py

@@ -3,15 +3,18 @@
 # @author Bruno Goncalves
 ######################################################
 
+from typing import Dict, List, Set, Union
 import warnings
+import string
+
 import networkx as nx
 import numpy as np
 from numpy import linalg
 from numpy import random
 import scipy.integrate
 import pandas as pd
 import matplotlib.pyplot as plt
-import string
+
 from .utils import *
 
 class EpiModel(object):
@@ -38,7 +41,7 @@ def __init__(self, compartments=None):
         if compartments is not None:
             self.transitions.add_nodes_from([comp for comp in compartments])
 
-    def add_interaction(self, source, target, agent, rate):  
+    def add_interaction(self, source: str, target: str, agent: str, rate: float) -> None:  
         """
         Add an interaction between two compartments
         
@@ -57,7 +60,7 @@ def add_interaction(self, source, target, agent, rate):
         """      
         self.transitions.add_edge(source, target, agent=agent, rate=rate)        
 
-    def add_spontaneous(self, source, target, rate):
+    def add_spontaneous(self, source: str, target: str, rate: float) -> None:
         """
         Add a spontaneous transition between two compartments
         
@@ -74,7 +77,7 @@ def add_spontaneous(self, source, target, rate):
         """
         self.transitions.add_edge(source, target, rate=rate)
 
-    def add_vaccination(self, source, target, rate, start):
+    def add_vaccination(self, source: str, target: str, rate: float, start: int) -> None:
         """
         Add a vaccination transition between two compartments
         
@@ -93,7 +96,14 @@ def add_vaccination(self, source, target, rate, start):
         """
         self.transitions.add_edge(source, target, rate=rate, start=start)
 
-    def add_age_structure(self, matrix, population):
+    def add_age_structure(self, matrix: List, population: List) -> List[List]:
+        """ 
+        Add a vaccination transition between two compartments
+        
+        Parameters:
+        - matrix: List
+        - population: List
+        """
         self.contact = np.asarray(matrix)
         self.population = np.asarray(population).flatten()
 
@@ -131,7 +141,7 @@ def add_age_structure(self, matrix, population):
 
         self.transitions = model.transitions
 
-    def _new_cases(self, population, time, pos):
+    def _new_cases(self, population: np.ndarray, time: float, pos: Dict) -> np.ndarray:
         """
         Internal function used by integration routine
         
@@ -188,7 +198,7 @@ def _new_cases(self, population, time, pos):
 
         return diff
 
-    def plot(self, title=None, normed=True, show=True, ax=None, **kwargs):
+    def plot(self, title: Union[str, None]= None, normed: bool = True, show: bool = True, ax: Union[plt.Axes, None] = None, **kwargs):
         """
         Convenience function for plotting
         
@@ -235,7 +245,7 @@ def plot(self, title=None, normed=True, show=True, ax=None, **kwargs):
             print(e)
             raise NotInitialized('You must call integrate() or simulate() first')
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> pd.Series:
         """
         Dynamic method to return the individual compartment values
         
@@ -252,7 +262,7 @@ def __getattr__(self, name):
         else:
             raise AttributeError("'EpiModel' object has no attribute '%s'" % name)
 
-    def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):
+    def simulate(self, timesteps: int, t_min: int = 1, seasonality: Union[np.ndarray, None] = None, **kwargs) -> None:
         """
         Stochastically simulate the epidemic model
         
@@ -334,7 +344,7 @@ def simulate(self, timesteps, t_min=1, seasonality=None, **kwargs):
         values = np.array(values)
         self.values_ = pd.DataFrame(values[1:], columns=comps, index=time)
 
-    def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):
+    def integrate(self, timesteps: int , t_min: int = 1, seasonality: Union[np.ndarray, None] = None, **kwargs) -> None:
         """
         Numerically integrate the epidemic model
         
@@ -375,7 +385,14 @@ def integrate(self, timesteps, t_min=1, seasonality=None, **kwargs):
         time = np.arange(t_min, t_min+timesteps, 1)
 
         self.seasonality = seasonality
-        values = pd.DataFrame(scipy.integrate.odeint(self._new_cases, population, time, args=(pos,)), columns=pos.keys(), index=time)
+        values = pd.DataFrame(
+                scipy.integrate.odeint(
+                    self._new_cases, 
+                    population, 
+                    time, 
+                    args=(pos,)
+                ), columns=pos.keys(), index=time
+        )
 
         if self.population is None:
             self.values_ = values
@@ -398,7 +415,7 @@ def single_step(self, seasonality=None, **kwargs):
             new_values = pd.concat([old_values, self.values_.iloc[[-1]]])
             self.values_ = new_values
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Return a string representation of the EpiModel object
         
@@ -433,7 +450,7 @@ def __repr__(self):
 
         return text
 
-    def _get_active(self):
+    def _get_active(self) -> Set:
         active = set()
 
         for node_i, node_j, data in self.transitions.edges(data=True):
@@ -444,7 +461,7 @@ def _get_active(self):
 
         return active
 
-    def _get_susceptible(self):
+    def _get_susceptible(self) -> Set:
         susceptible = set([node for node, deg in self.transitions.in_degree() if deg==0])
 
         if len(susceptible) == 0:
@@ -454,7 +471,7 @@ def _get_susceptible(self):
 
         return susceptible
 
-    def _get_infections(self):
+    def _get_infections(self) -> Dict:
         inf = {}
 
         for node_i, node_j, data in self.transitions.edges(data=True):
@@ -472,7 +489,7 @@ def _get_infections(self):
 
         return inf
 
-    def draw_model(self, ax=None, show=True):
+    def draw_model(self, ax: Union[plt.Axes, None] = None, show: bool = True) -> None:
         """
         Plot the model structure
 
@@ -516,7 +533,7 @@ def draw_model(self, ax=None, show=True):
         if show:
             plt.show()
 
-    def R0(self):
+    def R0(self) -> Union[float, None]:
         """
         Return the value of the basic reproductive ratio, $R_0$, for the model as defined
 

src/epidemik/MetaEpiModel.py

@@ -22,7 +22,7 @@ class MetaEpiModel:
     
         Provides a way to implement and numerically integrate 
     """
-    def __init__(self, travel_graph, populations, population='Population'):
+    def __init__(self, travel_graph: pd.DataFrame, populations: pd.DataFrame, population: str ='Population'):
         """
         Initialize the EpiModel object
         
@@ -51,7 +51,7 @@ def __init__(self, travel_graph, populations, population='Population'):
 
         self.models = models
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """
         Return a string representation of the EpiModel object
         
@@ -65,7 +65,7 @@ def __repr__(self):
         text = "Metapopulation model with %u populations\n\nThe disease is defined by an %s" % (self.travel_graph.shape[0], model_text)
         return text 
 
-    def add_interaction(self, source, target, agent, rate):  
+    def add_interaction(self, source: str, target: str, agent: str, rate: float) -> None:  
         """
         Add an interaction between two compartments_
         
@@ -85,7 +85,7 @@ def add_interaction(self, source, target, agent, rate):
         for state in self.models:  
             self.models[state].add_interaction(source, target, agent, rate)       
 
-    def add_spontaneous(self, source, target, rate):
+    def add_spontaneous(self, source: str, target: str, rate: float) -> None:
         """
         Add a spontaneous transition between two compartments_
         
@@ -103,7 +103,7 @@ def add_spontaneous(self, source, target, rate):
         for state in self.models:  
             self.models[state].add_spontaneous(source, target, rate)      
 
-    def add_vaccination(self, source, target, rate, start):
+    def add_vaccination(self, source: str, target: str, rate: float, start: int) -> None:
         """
         Add a vaccination transition between two compartments_
         
@@ -123,11 +123,11 @@ def add_vaccination(self, source, target, rate, start):
         for state in self.models:  
             self.models[state].add_vaccination(source, target, rate, start)
 
-    def R0(self):
+    def R0(self) -> Union[float, None]:
         key = list(self.models.keys())[0]
         return self.models[key].R0()
 
-    def get_state(self, state):
+    def get_state(self, state: str) -> EpiModel:
         """
         Return a reference to a state EpiModel object
 
@@ -138,7 +138,7 @@ def get_state(self, state):
 
         return self.models[state]
 
-    def _initialize_populations(self, susceptible, population=None):
+    def _initialize_populations(self, susceptible: str, population: Union[pd.DataFrame, None] =None) -> None:
         columns = list(self.transitions.nodes())
         self.compartments_ = pd.DataFrame(np.zeros((self.travel_graph.shape[0], len(columns)), dtype='int'), columns=columns)
         self.compartments_.index = self.populations.index
@@ -149,8 +149,8 @@ def _initialize_populations(self, susceptible, population=None):
         for state in self.compartments_.index:
             self.compartments_.loc[state, susceptible]  = self.populations.loc[state, population]
 
-    def _run_travel(self, compartments_, travel):
-        def travel_step(x, populations):
+    def _run_travel(self, compartments_: pd.DataFrame, travel: pd.DataFrame) -> pd.DataFrame:
+        def travel_step(x, populations: pd.DataFrame) -> pd.Series:
             n = populations.loc[x.name]
             p = travel.loc[x.name].values.tolist()
             output = np.random.multinomial(n, p)
@@ -163,17 +163,24 @@ def travel_step(x, populations):
         # Travel occurs independently for each compartment
         # since we don't allow in-flight transitions
         for comp in compartments_.columns:
-            new_compartments[comp] = travel.apply(travel_step, populations=compartments_[comp]).sum(axis=1)
+            new_compartments[comp] = travel.apply(
+                    travel_step, 
+                    populations=compartments_[comp]
+            ).sum(axis=1)
 
         return new_compartments
 
-    def _run_spread(self):
+    def _run_spread(self) -> None:
         for state in self.compartments_.index:
             pop = self.compartments_.loc[state].to_dict()
             self.models[state].single_step(**pop)
             self.compartments_.loc[state] = self.models[state].values_.iloc[[-1]].values[0]
 
-    def simulate(self, timestamp, t_min=1, seasonality=None, seed_state=None, susceptible='S', **kwargs):
+    def simulate(
+            self, timestamp: int, t_min: int = 1, 
+            seasonality=None, seed_state: [str, None] = None, 
+            susceptible: str ='S', **kwargs
+    ) -> None:
         if seed_state is None:
             raise NotInitialized("You have to specify the seed_state")
 
@@ -193,10 +200,10 @@ def simulate(self, timestamp, t_min=1, seasonality=None, seed_state=None, suscep
     def integrate(self, **kwargs):
         raise NotImplementedError("MetaEpiModel doesn't support direct integration of the ODE")
 
-    def draw_model(self):
+    def draw_model(self) -> None:
         return self.models.iloc[0].draw_model()
 
-    def plot(self, title=None, normed=True, layout=None, **kwargs):
+    def plot(self, title: Union[str, None] = None, normed: bool = True, layout=None, **kwargs) -> None:
         if layout is None:
             n_pop = self.travel_graph.shape[0]
             N = int(np.round(np.sqrt(n_pop), 0)+1)
@@ -270,7 +277,7 @@ def plot(self, title=None, normed=True, layout=None, **kwargs):
         fig.patch.set_facecolor('#FFFFFF')
         fig.tight_layout()
 
-    def plot_peaks(self):
+    def plot_peaks(self) -> None:
         peaks = None
 
         for state in self.models.values():
@@ -301,4 +308,4 @@ def plot_peaks(self):
         ax.set_xticks(np.arange(0, peaks.shape[1], 3))
         ax.set_xticklabels(np.arange(0, peaks.shape[1], 3), fontsize=10)
         # ax.set_aspect(1)
-        fig.patch.set_facecolor('#FFFFFF')
+        fig.patch.set_facecolor('#FFFFFF')

src/epidemik/NetworkEpiModel.py

@@ -3,6 +3,7 @@
 # @author Bruno Goncalves
 ######################################################
 
+from typing import Union
 import networkx as nx
 import numpy as np
 from numpy import linalg
@@ -24,7 +25,7 @@ def __init__(self, network, compartments=None):
     def integrate(self, timesteps, **kwargs):
         raise NotImplementedError("Network Models don't support numerical integration")
 
-    def add_interaction(self, source, target, agent, rate, rescale=False):
+    def add_interaction(self, source: str, target: str, agent: str, rate: float, rescale: bool = False) -> None:
         if rescale:
             rate /= self.kavg_
 
@@ -38,7 +39,7 @@ def add_interaction(self, source, target, agent, rate, rescale=False):
 
         self.interactions[source][agent] = {'target': target, 'rate': rate}
 
-    def add_spontaneous(self, source, target, rate):
+    def add_spontaneous(self, source: str, target: str, rate: float) -> None:
         super(NetworkEpiModel, self).add_spontaneous(source, target, rate=rate)
         if source not in self.spontaneous:
             self.spontaneous[source] = {}
@@ -49,7 +50,7 @@ def add_spontaneous(self, source, target, rate):
         self.spontaneous[source][target] = rate
 
 
-    def simulate(self, timesteps, seeds, **kwargs):
+    def simulate(self, timesteps: int, seeds, **kwargs) -> None:
         """Stochastically simulate the epidemic model"""
         pos = {comp: i for i, comp in enumerate(self.transitions.nodes())}
         N = self.network.number_of_nodes()
@@ -130,7 +131,7 @@ def simulate(self, timesteps, seeds, **kwargs):
         self.population_ = pd.DataFrame(population)
         self.values_ = pd.DataFrame.from_records(self.population_.apply(lambda x: Counter(x), axis=1)).fillna(0).astype('int')
 
-    def R0(self):
+    def R0(self) -> Union[float, None]:
         if 'R' not in set(self.transitions.nodes):
             return None
-        return np.round(super(NetworkEpiModel, self).R0()*self.kavg_, 2)
+        return np.round(super(NetworkEpiModel, self).R0()*self.kavg_, 2)