Docstring compliance

pyproject.toml

@@ -25,6 +25,7 @@ mypy = "^1.10.0"
 requests = "^2.32.3"
 docker = "^7.1.0"
 bayeux-ml = "^0.1.14"
+pydocstyle = "^6.3.0"
 
 
 [tool.poetry.group.dev.dependencies]

src/dynode/__init__.py

@@ -1,4 +1,13 @@
-# needs to exist to define a module
+"""DynODE, a dynamic ordinary differential model framework.
+
+DynODE is a a compartmental mechanistic ODE model that accounts for
+age structure, immunity history, vaccination, immunity waning and
+multiple variants.
+
+DynODE is currently under active development and will be substantially
+refactored in the near future!
+"""
+
 # ruff: noqa: E402
 import jax
 

src/dynode/abstract_initializer.py

@@ -1,5 +1,5 @@
-"""
-A module that creates an abstract class for an initializer object.
+"""A module that creates an abstract class for an initializer object.
+
 An initializer objects primary purpose is initialize the state on which ODEs will be run.
 AbstractInitializers will often be tasked with reading, parsing, and combining data sources
 to produce an initial state representing some analyzed population
@@ -8,18 +8,27 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
+from numpy import ndarray
+
 from . import SEIC_Compartments, utils
 
 
 class AbstractInitializer(ABC):
-    """
-    An Abstract class meant for use by disease-specific initializers.
-    an initializers sole responsibility is to return an INITIAL_STATE
+    """An abstract class meant for use by disease-specific initializers.
+
+    An initializer's sole responsibility is to return an INITIAL_STATE
     parameter via self.get_initial_state().
     """
 
     @abstractmethod
     def __init__(self, initializer_config) -> None:
+        """Load parameters from `initializer_config` and generate self.INITIAL_STATE.
+
+        Parameters
+        ----------
+        initializer_config : str
+            str path to config json holding necessary initializer parameters.
+        """
         # add these for mypy
         self.INITIAL_STATE: SEIC_Compartments | None = None
         self.config: Any = {}
@@ -28,27 +37,32 @@ def __init__(self, initializer_config) -> None:
     def get_initial_state(
         self,
     ) -> SEIC_Compartments:
-        """
-        Returns the initial state of the model as defined by the child class in __init__
+        """Get the initial state of the model as defined by the child class in __init__.
+
+        Returns
+        -------
+        SEIC_Compartments
+            tuple of matricies representing initial state of each compartment
+            in the model.
         """
         assert self.INITIAL_STATE is not None
         return self.INITIAL_STATE
 
-    def load_initial_population_fractions(self) -> None:
-        """
-        a wrapper function which loads age demographics for the US and sets the inital population fraction by age bin.
+    def load_initial_population_fractions(self) -> ndarray:
+        """Load age demographics for the specified region.
 
-        Updates
-        ----------
-        `self.config.INITIAL_POPULATION_FRACTIONS` : numpy.ndarray
-            proportion of the total population that falls into each age group,
-            length of this array is equal the number of age groups and will sum to 1.0.
+        Returns
+        -------
+        numpy.ndarray
+            Proportion of the total population that falls into each age group.
+            `len(self.load_initial_population_fractions()) == self.config.NUM_AGE_GROUPS`
+            `np.sum(self.load_initial_population_fractions()) == 1.0
         """
         populations_path = (
             self.config.DEMOGRAPHIC_DATA_PATH
             + "population_rescaled_age_distributions/"
         )
         # TODO support getting more regions than just 1
-        self.config.INITIAL_POPULATION_FRACTIONS = utils.load_age_demographics(
+        return utils.load_age_demographics(
             populations_path, self.config.REGIONS, self.config.AGE_LIMITS
         )[self.config.REGIONS[0]]