type hints added

comma/hypothesis.py

@@ -3,7 +3,7 @@
 import json
 import os
 import pandas as pd
-from typing import Dict, Set
+from typing import Dict, Set, List
 
 PARAMS_INDIVIDUAL = 'params_individual.json'
 PARAMS_IPF_WEIGHTS = "ipf_weights.csv"
@@ -12,8 +12,7 @@
 class Hypothesis:
     """
     The Hypothesis class is responsible for managing and validating
-    the Hypotheses specified by the user.
-    These can be found in the `dir_params` folder.
+    hypotheses specified by the user.
 
     Methods:
         _get_one_hot_encoded_features():
@@ -140,7 +139,7 @@ def read_hypotheses(cls, dir_params: str, policies: Set[str],
         return data_dfs
 
     @staticmethod
-    def _get_one_hot_encoded_features(fpath_params_individual: str):
+    def _get_one_hot_encoded_features(fpath_params_individual: str) -> List:
         """
         One-hot encode categorical features in the
         `params_individual.json` file and return the
@@ -166,14 +165,17 @@ def _get_one_hot_encoded_features(fpath_params_individual: str):
         return features
 
     @classmethod
-    def create_empty_hypotheses(cls, dir_params: str):
+    def create_empty_hypotheses(cls, dir_params: str) -> None:
         """
         Create empty CSV files for storing hypotheses on
         the impact of actions and lockdown policies on different agent statuses
 
         Args:
             dir_params (str): The directory of the folder that contains
             the agent and model parameter files.
+        Returns:
+            None: This function does not return anything
+            as it creates empty csv files int the specified directory
         """
         fpath_params_individual = os.path.join(dir_params, PARAMS_INDIVIDUAL)
 
@@ -198,7 +200,7 @@ def create_empty_hypotheses(cls, dir_params: str):
             df.to_csv(fp, sep=';', index=False)
 
     @classmethod
-    def validate_param_file(cls, dir_params: str):
+    def validate_param_file(cls, dir_params: str) -> None:
         """Validate files in the parameter folder.
 
         Args:

comma/individual.py

@@ -7,6 +7,7 @@
 import numpy as np
 import os
 import pandas as pd
+from typing import List, Tuple
 from tqdm import tqdm
 
 
@@ -20,7 +21,7 @@ def __init__(self, id: int, dir_params: str, features):
         self._features = features
         self.actions = Hypothesis.all_possible_actions
 
-    def get_features(self):
+    def get_features(self) -> pd.Series:
         """
         Get agent's features
 
@@ -30,16 +31,16 @@ def get_features(self):
         """
         return self._features
 
-    def get_status(self):
+    def get_status(self) -> float:
         """
         Get the current agent status (i.e., mental health)
 
         Returns:
-            pd.Series: the current status of the agent
+            float: the current status of the agent
         """
         return self._status
 
-    def get_actions(self):
+    def get_actions(self) -> List[str]:
         """
         Get the current actions chosen by the agent
 
@@ -49,15 +50,17 @@ def get_actions(self):
         return [action_name for action_name, action_was_taken in
                 zip(self.actions, self.chosen_actions) if action_was_taken]
 
-    def choose_actions_on_lockdown(self, lockdown: pd.DataFrame):
-        """Choose the actions to take based on current lockdown policy.
+    def choose_actions_on_lockdown(self, lockdown: pd.DataFrame)\
+            -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Choose the actions to take based on current lockdown policy.
 
         Args:
             lockdown (pd.DataFrame): dataframe of a given lockdown
 
         Returns:
-            actions (pd.Series): list of booleans of taking/not-taking actions
-            actions_probs (pd.Series): probability of taking that action
+            actions (np.ndarray): array of booleans
+            actions_probs (np.ndarray): array of probability
         """
         params_lockdown = lockdown
         n_actions, _ = params_lockdown.shape
@@ -71,19 +74,26 @@ def choose_actions_on_lockdown(self, lockdown: pd.DataFrame):
 
         return actions, action_probs
 
-    def take_actions(self, actions: pd.Series, action_effects: pd.DataFrame):
-        """Update status by taking the given actions.
+    def take_actions(self, actions: pd.Series,
+                     action_effects: pd.DataFrame) -> None:
+        """
+        Update status by taking the given actions.
 
         Args:
             actions (pd.Series): list of booleans of chosen/not-chosen actions.
+            action_effects (pd.DataFrame): matrix of actions effects
+
+        Returns:
+            None: This function updates the agent's status
+            but does not return anything.
         """
         params_status = action_effects
         result = params_status.dot(self.get_features()).dot(actions)
 
         self._status = result
 
     @staticmethod
-    def sampling_from_ipf(size: int, dir_params: str):
+    def sampling_from_ipf(size: int, dir_params: str) -> pd.DataFrame:
         """
         Sample from IPF distribution saved
         as `weights.csv` in the parameters folder
@@ -109,14 +119,19 @@ def sampling_from_ipf(size: int, dir_params: str):
         return sample
 
     @staticmethod
-    def populate_ipf(size: int, dir_params: str):
+    def populate_ipf(size: int, dir_params: str) -> List:
         """
         Create a population of individual agents
         with the given weights obtained via IPF
 
         Args:
             size (int): size of data sample.
             dir_params (str): path to parameters folder.
+
+        Returns:
+            List[Individual]: A list containing instances of
+            the individual class, each representing an
+            agent with specific features.
         """
         _features = pd.DataFrame()
 
@@ -137,7 +152,7 @@ def populate_ipf(size: int, dir_params: str):
                 tqdm(range(size), desc="Populating individuals", unit="i")]
 
     @staticmethod
-    def populate(size: int, dir_params: str):
+    def populate(size: int, dir_params: str) -> List:
         """
         Create a population of individual agents
         with the given feature parameters.

comma/model.py

@@ -4,20 +4,22 @@
 from comma.hypothesis import Hypothesis
 import pandas as pd
 import numpy as np
+from typing import List, Tuple
 from tqdm import tqdm
 
 
 class Model:
     cumulative_status = dict()
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.agents: list = None
         self.simulation_id: int = None
         self.current_step: int = 0  # keep track of the current simulation step
         self.lockdown_status: dict = {}
         self.dir_params: str = None
 
-    def setup(self, size: int, dir_params: str, use_ipf: bool = False):
+    def setup(self, size: int, dir_params: str,
+              use_ipf: bool = False) -> None:
         """Setup the model with input parameters.
 
         Args:
@@ -35,7 +37,8 @@ def setup(self, size: int, dir_params: str, use_ipf: bool = False):
         else:
             self.agents = Individual.populate(size, dir_params)
 
-    def step(self, lockdown: pd.DataFrame, action_effects: pd.DataFrame):
+    def step(self, lockdown: pd.DataFrame,
+             action_effects: pd.DataFrame) -> None:
         """Actions to be performed in each step.
 
         Args:
@@ -51,8 +54,9 @@ def step(self, lockdown: pd.DataFrame, action_effects: pd.DataFrame):
             # take those actions, and compute their effect on mental health
             agent.take_actions(actions=actions, action_effects=action_effects)
 
-    def update(self, lockdown: str, step: int):
-        """Update status at every step.
+    def update(self, lockdown: str, step: int) -> None:
+        """
+        It updates statuses at every step.
 
         Args:
             lockdown (str): lockdown type
@@ -81,10 +85,10 @@ def update(self, lockdown: str, step: int):
                 agent_statuses.append(new_status)
             self.cumulative_status[step] = agent_statuses
 
-    def report(self, out_path: str):
+    def report(self, out_path: str) -> None:
         """
-        Collect data recorded at the end of the simulation
-        and export as csv file.
+        It Collects data recorded at the end of the simulation
+        and exports it as csv file.
 
         Args:
             out_path (str): File path of the output file
@@ -106,15 +110,15 @@ def report(self, out_path: str):
         # Export to a csv
         status_df.to_csv(out_path, index=False, sep=";", decimal=",")
 
-    def run(self, steps: int, lockdown_policy: list, out_path: str):
+    def run(self, steps: int, lockdown_policy: list, out_path: str) -> None:
         """Run a simulation
 
         Args:
-            steps (int): Number of steps to run the simulation
+            steps(int): Number of steps to run the simulation
 
-            lockdown (str): Type of lockdown (easy, hard, medium, no_lockdown)
+            lockdown_policy(str): Type of lockdown policy
 
-            out_path (str): File path of the output file
+            out_path(str): File path of the output file
         """
         if len(lockdown_policy) != steps:
             raise ValueError("The length of the lockdown list \