From d34c02b998966124dfe8e286ab455c695126cae5 Mon Sep 17 00:00:00 2001 From: Ariel Shurygin Date: Tue, 10 Dec 2024 17:43:34 +0000 Subject: [PATCH] adding correct return type hint --- src/dynode/covid_sero_initializer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/dynode/covid_sero_initializer.py b/src/dynode/covid_sero_initializer.py index c200ee8e..9f1c6c18 100644 --- a/src/dynode/covid_sero_initializer.py +++ b/src/dynode/covid_sero_initializer.py @@ -8,7 +8,7 @@ import numpy as np import pandas as pd -from . import utils +from . import SEIC_Compartments, utils from .abstract_initializer import AbstractInitializer from .config import Config @@ -52,7 +52,9 @@ def __init__(self, config_initializer_path, global_variables_path): self.config.INITIAL_INFECTIONS ) - def load_initial_state(self, initial_infections: float) -> None: + def load_initial_state( + self, initial_infections: float + ) -> SEIC_Compartments: """ a function which takes a number of initial infections, disperses them across infectious and exposed compartments according to the INIT_INFECTED_DIST