Added type hints #2
Conversation
bmtgoncalves
left a comment
bmtgoncalves
left a comment
There was a problem hiding this comment.
This is great, thank you so much. There's a couple of small things in specifying types with defaults, but otherwise looks good.
| def add_age_structure(self, matrix, population): | ||
| def add_age_structure(self, matrix: List, population: List) -> None: | ||
| """ | ||
| Add a vaccination transition between two compartments |
There was a problem hiding this comment.
Copy-paste issue. The descriptions should probably be something like "Add age structure with a contact matrix and age-structured population"
Also, the matrix is 2D, so maybe something like List[List]? (Not sure about how to specify this)
src/epidemik/EpiModel.py
Outdated
| return diff | ||
|
|
||
| def plot(self, title=None, normed=True, show=True, ax=None, **kwargs): | ||
| def plot(self, title: str = None, normed: bool = True, show: bool = True, ax: plt.Axes = None, **kwargs): |
There was a problem hiding this comment.
Since the detault title is None, the type should be Union[str, None]
src/epidemik/EpiModel.py
Outdated
| 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: np.ndarray = None, **kwargs) -> None: |
There was a problem hiding this comment.
I think we need the Union here as well
src/epidemik/EpiModel.py
Outdated
| 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: np.ndarray = None, **kwargs) -> None: |
src/epidemik/EpiModel.py
Outdated
| return inf | ||
|
|
||
| def draw_model(self, ax=None, show=True): | ||
| def draw_model(self, ax: plt.Axes = None, show: bool = True) -> None: |
src/epidemik/MetaEpiModel.py
Outdated
| self.models[state].add_vaccination(source, target, rate, start) | ||
|
|
||
| def R0(self): | ||
| def R0(self) -> Union[np.ndarray, float, None]: |
src/epidemik/MetaEpiModel.py
Outdated
| return self.models[state] | ||
|
|
||
| def _initialize_populations(self, susceptible, population=None): | ||
| def _initialize_populations(self, susceptible: str, population=None) -> None: |
There was a problem hiding this comment.
population should be a DataFrame or None
| 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): |
There was a problem hiding this comment.
seed_state should be a string. It essentially the name of the population that is being seeded, and corresponds to the index of the population DataFrame
src/epidemik/MetaEpiModel.py
Outdated
| return self.models.iloc[0].draw_model() | ||
|
|
||
| def plot(self, title=None, normed=True, layout=None, **kwargs): | ||
| def plot(self, title: str = None, normed: bool = True, layout=None, **kwargs) -> None: |
src/epidemik/NetworkEpiModel.py
Outdated
| 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[np.ndarray, None]: |
3 participants
Added type hints for classes and their methods to improve readability and traceability of the code