Merge pull request #501 from The-Strategy-Unit/copilot/fix-59ab13b9-7b4a-45d4-b591-8d9def500081

Update docstrings to Google-style format and fix outdated documentation

Author: tomjemmett

src/nhp/docker/run.py

@@ -23,27 +23,27 @@ class RunWithLocalStorage:
     def __init__(self, filename: str):
         """Initialize the RunWithLocalStorage instance.
 
-        :param filename: Name of the parameter file to load.
-        :type filename: str
+        Args:
+            filename: Name of the parameter file to load.
         """
         self.params = load_params(f"queue/{filename}")
 
     def finish(self, results_file: str, saved_files: list, save_full_model_results: bool) -> None:
         """Post model run steps.
 
-        :param results_file: the path to the results file
-        :type results_file: str
-        :param saved_files: filepaths of results, saved in parquet format and params
-        in json format
-        :type saved_files: list
-        :param save_full_model_results: whether to save the full model results or not
-        :type save_full_model_results: bool
+        Args:
+            results_file: The path to the results file.
+            saved_files: Filepaths of results, saved in parquet format and params in json format.
+            save_full_model_results: Whether to save the full model results or not.
         """
 
     def progress_callback(self) -> Callable[[Any], Callable[[Any], None]]:
         """Progress callback method.
 
-        for local storage do nothing
+        For local storage do nothing.
+
+        Returns:
+            A no-op progress callback function.
         """
         return noop_progress_callback
 
@@ -54,10 +54,9 @@ class RunWithAzureStorage:
     def __init__(self, filename: str, config: Config = Config()):
         """Initialise RunWithAzureStorage.
 
-        :param filename:
-        :type filename: str
-        :param config: The configuration for the run
-        :type config: Config
+        Args:
+            filename: Name of the parameter file to load.
+            config: The configuration for the run. Defaults to Config().
         """
         logging.getLogger("azure.storage.common.storageclient").setLevel(logging.WARNING)
         logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(
@@ -86,10 +85,11 @@ def _get_container(self, container_name: str):
     def _get_params(self, filename: str) -> dict:
         """Get the parameters for the model.
 
-        :param filename: the name of the params file
-        :type filename: str
-        :return: the parameters for the model
-        :rtype: dict
+        Args:
+            filename: The name of the params file.
+
+        Returns:
+            The parameters for the model.
         """
         logging.info("downloading params: %s", filename)
 
@@ -102,12 +102,11 @@ def _get_params(self, filename: str) -> dict:
     def _get_data(self, year: str, dataset: str) -> None:
         """Get data to run the model.
 
-        for local storage, the data is already available, so do nothing.
+        Downloads data from Azure storage for the specified year and dataset.
 
-        :param year: the year of data to load
-        :type year: str
-        :param year: the year of data to load
-        :type year: str
+        Args:
+            year: The year of data to load.
+            dataset: The dataset to load.
         """
         logging.info("downloading data (%s / %s)", year, dataset)
         fs_client = DataLakeServiceClient(
@@ -137,12 +136,11 @@ def _get_data(self, year: str, dataset: str) -> None:
     def _upload_results_json(self, results_file: str, metadata: dict) -> None:
         """Upload the results.
 
-        once the model has run, upload the results to blob storage
+        Once the model has run, upload the results to blob storage.
 
-        :param results_file: the saved results file
-        :type results_file: str
-        :param metadata: the metadata to attach to the blob
-        :type metadata: dict
+        Args:
+            results_file: The saved results file.
+            metadata: The metadata to attach to the blob.
         """
         container = self._get_container("results")
 
@@ -157,14 +155,12 @@ def _upload_results_json(self, results_file: str, metadata: dict) -> None:
     def _upload_results_files(self, files: list, metadata: dict) -> None:
         """Upload the results.
 
-        once the model has run, upload the files (parquet for model results and json for
-        model params) to blob storage
-
-        :param files: list of files to be uploaded
-        :type files: list
-        :param metadata: the metadata to attach to the blob
-        :type metadata: dict
+        Once the model has run, upload the files (parquet for model results and json for
+        model params) to blob storage.
 
+        Args:
+            files: List of files to be uploaded.
+            metadata: The metadata to attach to the blob.
         """
         container = self._get_container("results")
         for file in files:
@@ -202,7 +198,7 @@ def _upload_full_model_results(self) -> None:
     def _cleanup(self) -> None:
         """Cleanup.
 
-        once the model has run, remove the file from the queue
+        Once the model has run, remove the file from the queue.
         """
         logging.info("cleaning up queue")
 
@@ -211,13 +207,10 @@ def _cleanup(self) -> None:
     def finish(self, results_file: str, saved_files: list, save_full_model_results: bool) -> None:
         """Post model run steps.
 
-        :param results_file: the path to the results file
-        :type results_file: str
-        :param saved_files: filepaths of results, saved in parquet format and params
-        in json format
-        :type saved_files: list
-        :param save_full_model_results: whether to save the full model results or not
-        :type save_full_model_results: bool
+        Args:
+            results_file: The path to the results file.
+            saved_files: Filepaths of results, saved in parquet format and params in json format.
+            save_full_model_results: Whether to save the full model results or not.
         """
         metadata = {
             k: str(v)
@@ -233,7 +226,10 @@ def finish(self, results_file: str, saved_files: list, save_full_model_results:
     def progress_callback(self) -> Callable[[Any], Callable[[Any], None]]:
         """Progress callback method.
 
-        updates the metadata for the blob in the queue to give progress
+        Updates the metadata for the blob in the queue to give progress.
+
+        Returns:
+            A callback function that updates progress for each model type.
         """
         blob = self._queue_blob
 

src/nhp/model/__main__.py

@@ -7,7 +7,7 @@
 There are existing launch profiles for vscode that use this file, or you can use it directly in the
 console, e.g.
 
-    python run_model.py data/[DATASET]/results/[SCENARIO]/[RUN_TIME] 0 1 --debug -t=ip
+    python -m nhp.model -d data --model-run 1 -t ip
 
 will run a single run of the inpatients model, returning the results to display.
 """

src/nhp/model/aae.py

@@ -18,40 +18,32 @@ class AaEModel(Model):
 
     Implementation of the Model for Accident and Emergency attendances.
 
-    :param params: the parameters to run the model with, or the path to a params file to load
-    :type params: dict or string
-    :param data: a Data class ready to be constructed
-    :type data: Data
-    :param hsa: An instance of the HealthStatusAdjustment class. If left as None an instance is
-    created
-    :type hsa: HealthStatusAdjustment, optional
-    :param run_params: the parameters to use for each model run. generated automatically if left as
-    None
-    :type run_params: dict
-    :param save_full_model_results: whether to save the full model results or not
-    :type save_full_model_results: bool, optional
+    Args:
+        params: The parameters to run the model with, or the path to a params file to load.
+        data: A callable that creates a Data instance.
+        hsa: An instance of the HealthStatusAdjustment class. If left as None an instance is
+            created. Defaults to None.
+        run_params: The parameters to use for each model run. Generated automatically if left as
+            None. Defaults to None.
+        save_full_model_results: Whether to save the full model results or not. Defaults to False.
     """
 
     def __init__(
         self,
-        params: dict,
+        params: dict | str,
         data: Callable[[int, str], Data],
         hsa: Any = None,
         run_params: dict | None = None,
         save_full_model_results: bool = False,
     ) -> None:
         """Initialise the A&E Model.
 
-        :param params: the parameters to use
-        :type params: dict
-        :param data: a method to create a Data instance
-        :type data: Callable[[int, str], Data]
-        :param hsa: _Health Status Adjustment object, defaults to None
-        :type hsa: Any, optional
-        :param run_params: the run parameters to use, defaults to None
-        :type run_params: dict | None, optional
-        :param save_full_model_results: whether to save full model results, defaults to False
-        :type save_full_model_results: bool, optional
+        Args:
+            params: The parameters to use.
+            data: A method to create a Data instance.
+            hsa: Health Status Adjustment object. Defaults to None.
+            run_params: The run parameters to use. Defaults to None.
+            save_full_model_results: Whether to save full model results. Defaults to False.
         """
         # call the parent init function
         super().__init__(
@@ -70,10 +62,11 @@ def _get_data(self, data_loader: Data) -> pd.DataFrame:
     def get_data_counts(self, data: pd.DataFrame) -> np.ndarray:
         """Get row counts of data.
 
-        :param data: the data to get the counts of
-        :type data: pd.DataFrame
-        :return: the counts of the data, required for activity avoidance steps
-        :rtype: np.ndarray
+        Args:
+            data: The data to get the counts of.
+
+        Returns:
+            The counts of the data, required for activity avoidance steps.
         """
         return np.array([data["arrivals"]]).astype(float)
 
@@ -99,15 +92,15 @@ def _load_strategies(self, data_loader: Data) -> None:
     def apply_resampling(self, row_samples: np.ndarray, data: pd.DataFrame) -> pd.DataFrame:
         """Apply row resampling.
 
-        Called from within `model.activity_resampling.ActivityResampling.apply_resampling`
+        Called from within `model.activity_resampling.ActivityResampling.apply_resampling`.
+
+        Args:
+            row_samples: [1xn] array, where n is the number of rows in `data`, containing the new
+                values for `data["arrivals"]`.
+            data: The data that we want to update.
 
-        :param row_samples: [1xn] array, where n is the number of rows in `data`, containing the new
-        values for `data["arrivals"]`
-        :type row_samples: np.ndarray
-        :param data: the data that we want to update
-        :type data: pd.DataFrame
-        :return: the updated data
-        :rtype: pd.DataFrame
+        Returns:
+            The updated data.
         """
         data["arrivals"] = row_samples[0]
         # return the altered data
@@ -118,18 +111,25 @@ def efficiencies(
     ) -> tuple[pd.DataFrame, pd.DataFrame | None]:
         """Run the efficiencies steps of the model.
 
-        :param model_iteration: an instance of the ModelIteration class
-        :type model_iteration: model.model_iteration.ModelIteration
+        Args:
+            data: The data to apply efficiencies to.
+            model_iteration: An instance of the ModelIteration class.
+
+        Returns:
+            Tuple containing the updated data and step counts (None for A&E).
         """
         # A&E doesn't have any efficiencies steps
         return data, None
 
     @staticmethod
     def process_results(data: pd.DataFrame) -> pd.DataFrame:
-        """Processes the data into a format suitable for aggregation in results files.
+        """Process the data into a format suitable for aggregation in results files.
 
-        :param data: Data to be processed. Format should be similar to Model.data
-        :type data: pd.DataFrame
+        Args:
+            data: Data to be processed. Format should be similar to Model.data.
+
+        Returns:
+            Processed results.
         """
         data["measure"] = "walk-in"
         data.loc[data["is_ambulance"], "measure"] = "ambulance"
@@ -161,10 +161,11 @@ def process_results(data: pd.DataFrame) -> pd.DataFrame:
     def specific_aggregations(self, model_results: pd.DataFrame) -> dict[str, pd.Series]:
         """Create other aggregations specific to the model type.
 
-        :param model_results: the results of a model run
-        :type model_results: pd.DataFrame
-        :return: dictionary containing the specific aggregations
-        :rtype: dict[str, pd.Series]
+        Args:
+            model_results: The results of a model run.
+
+        Returns:
+            Dictionary containing the specific aggregations.
         """
         return {
             "acuity": self.get_agg(model_results, "acuity"),
@@ -176,10 +177,12 @@ def calculate_avoided_activity(
     ) -> pd.DataFrame:
         """Calculate the rows that have been avoided.
 
-        :param data: The data before the binomial thinning step
-        :type data: pd.DataFrame
-        :return: The data that was avoided in the binomial thinning step
-        :rtype: pd.DataFrame
+        Args:
+            data: The data before the binomial thinning step.
+            data_resampled: The data after the binomial thinning step.
+
+        Returns:
+            The data that was avoided in the binomial thinning step.
         """
         avoided = data["arrivals"] - data_resampled["arrivals"]
         data["arrivals"] = avoided
@@ -192,10 +195,9 @@ def save_results(self, model_iteration: ModelIteration, path_fn: Callable[[str],
         It saves just the `rn` (row number) column and the `arrivals`, with the intention that
         you rejoin to the original data.
 
-        :param model_iteration: an instance of the `ModelIteration` class
-        :type model_iteration: model.model_iteration.ModelIteration
-        :param path_fn: a function which takes the activity type and returns a path
-        :type path_fn: Callable[[str], str]
+        Args:
+            model_iteration: An instance of the ModelIteration class.
+            path_fn: A function which takes the activity type and returns a path.
         """
         model_iteration.get_model_results().set_index(["rn"])[["arrivals"]].to_parquet(
             f"{path_fn('aae')}/0.parquet"

src/nhp/model/activity_resampling.py

@@ -3,9 +3,14 @@
 Methods for handling row resampling
 """
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import pandas as pd
 
+if TYPE_CHECKING:
+    from nhp.model.model_iteration import ModelIteration
+
 
 class ActivityResampling:
     """Activity Resampling.
@@ -24,16 +29,16 @@ class ActivityResampling:
     method.
     This updates the `model_iteration` which is passed in at initialisation.
 
-    :param model_iteration: the model iteration object, which contains all of the required
-    values to run the model.
-    :type model_iteration: ModelIteration
+    Args:
+        model_iteration: The model iteration object, which contains all of the required
+            values to run the model.
     """
 
-    def __init__(self, model_iteration) -> None:
+    def __init__(self, model_iteration: "ModelIteration") -> None:
         """Initialise ActivityResampling.
 
-        :param model_iteration: the current model iteration we are performing
-        :type model_iteration: ModelIteration
+        Args:
+            model_iteration: The current model iteration we are performing.
         """
         self._model_iteration = model_iteration