adding docstrings to models files

aepsych/models/base.py

@@ -125,14 +125,17 @@ def get_max(
         max_time: Optional[float] = None,
     ) -> Tuple[float, torch.Tensor]:
         """Return the maximum of the modeled function, subject to constraints
+
         Args:
-            locked_dims (Mapping[int, List[float]]): Dimensions to fix, so that the
-                inverse is along a slice of the full surface.
-            probability_space (bool): Is y (and therefore the returned nearest_y) in
+            locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the
+                inverse is along a slice of the full surface. Defaults to None.
+            probability_space (bool, optional): Is y (and therefore the returned nearest_y) in
                 probability space instead of latent function space? Defaults to False.
-            n_samples int: number of coarse grid points to sample for optimization estimate.
+            n_samples (int, optional): number of coarse grid points to sample for optimization estimate.
+            max_time (Optional[float], optional): Maximum time to spend optimizing. Defaults to None.
+
         Returns:
-            Tuple[float, np.ndarray]: Tuple containing the max and its location (argmax).
+            Tuple[float, torch.Tensor]: Tuple containing the max and its location (argmax).
         """
         locked_dims = locked_dims or {}
         _, _arg = get_extremum(
@@ -154,11 +157,13 @@ def get_min(
     ) -> Tuple[float, torch.Tensor]:
         """Return the minimum of the modeled function, subject to constraints
         Args:
-            locked_dims (Mapping[int, List[float]]): Dimensions to fix, so that the
+            locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the
                 inverse is along a slice of the full surface.
-            probability_space (bool): Is y (and therefore the returned nearest_y) in
+            probability_space (bool, optional): Is y (and therefore the returned nearest_y) in
                 probability space instead of latent function space? Defaults to False.
-            n_samples int: number of coarse grid points to sample for optimization estimate.
+            n_samples (int, optional): number of coarse grid points to sample for optimization estimate.
+            max_time (Optional[float], optional): Maximum time to spend optimizing. Defaults to None.
+
         Returns:
             Tuple[float, torch.Tensor]: Tuple containing the min and its location (argmin).
         """
@@ -185,12 +190,17 @@ def inv_query(
         """Query the model inverse.
         Return nearest x such that f(x) = queried y, and also return the
             value of f at that point.
+
         Args:
             y (float): Points at which to find the inverse.
-            locked_dims (Mapping[int, List[float]]): Dimensions to fix, so that the
+            locked_dims (Mapping[int, List[float]], optional): Dimensions to fix, so that the
                 inverse is along a slice of the full surface.
-            probability_space (bool): Is y (and therefore the returned nearest_y) in
+            probability_space (bool, optional): Is y (and therefore the returned nearest_y) in
                 probability space instead of latent function space? Defaults to False.
+            n_samples (int, optional): number of coarse grid points to sample for optimization estimate. Defaults to 1000.
+            max_time (Optional[float], optional): Maximum time to spend optimizing. Defaults to None.
+            weights (Optional[torch.Tensor], optional): Weights for the optimization. Defaults to None.
+
         Returns:
             Tuple[float, torch.Tensor]: Tuple containing the value of f
                 nearest to queried y and the x position of this value.
@@ -233,8 +243,8 @@ def get_jnd(
         Both definitions are equivalent for linear psychometric functions.
 
         Args:
-            grid (Optional[np.ndarray], optional): Mesh grid over which to find the JND.
-                Defaults to a square grid of size as determined by aepsych.utils.dim_grid
+            grid (Optional[Union[np.ndarray, torch.Tensor]], optional): Mesh grid over which to find the JND.
+                Defaults to a square grid of size as determined by aepsych.utils.dim_grid. 
             cred_level (float, optional): Credible level for computing an interval.
                 Defaults to None, computing no interval.
             intensity_dim (int, optional): Dimension over which to compute the JND.
@@ -310,13 +320,21 @@ def dim_grid(
         gridsize: int = 30,
         slice_dims: Optional[Mapping[int, float]] = None,
     ) -> torch.Tensor:
+        """Generate a grid based on lower, upper, and dim.
+        
+        Args:
+            gridsize (int, optional): Number of points in each dimension. Defaults to 30.
+            slice_dims (Optional[Mapping[int, float]], optional): Dimensions to fix at a 
+            certain value. Defaults to None."""
         return dim_grid(self.lb, self.ub, gridsize, slice_dims)
 
     def set_train_data(self, inputs: Optional[torch.Tensor] = None, targets: Optional[torch.Tensor] = None, strict: bool = False):
         """
-        :param torch.Tensor inputs: The new training inputs.
-        :param torch.Tensor targets: The new training targets.
-        :param bool strict: (default False, ignored). Here for compatibility with
+        Args:
+            inputs (Optional[torch.Tensor], optional):  The new training inputs.
+            targets (Optional[torch.Tensor], optional): The new training targets.
+            strict (bool, optional):  Default is False. Ignored, just for compatibility.
+
         input transformers. TODO: actually use this arg or change input transforms
         to not require it.
         """
@@ -327,6 +345,13 @@ def set_train_data(self, inputs: Optional[torch.Tensor] = None, targets: Optiona
             self.train_targets = targets
 
     def normalize_inputs(self, x: torch.Tensor) -> torch.Tensor:
+        """Normalize the input based on the defined bounds.
+        
+        Args:
+            x (torch.Tensor): Tensor of points to normalize.
+            
+        Returns:
+            torch.Tensor: Normalized tensor of points."""
         scale = self.ub - self.lb
         return (x - self.lb) / scale
 
@@ -353,6 +378,13 @@ def _fit_mll(
         optimizer=fit_gpytorch_mll_scipy,
         **kwargs,
     ) -> None:
+        """Fits the model by maximizing the marginal log likelihood.
+        
+        Args:
+            mll (MarginalLogLikelihood): Marginal log likelihood object.
+            optimizer_kwargs (Optional[Dict[str, Any]], optional): Keyword arguments for the optimizer.
+            optimizer: Optimizer to use. Defaults to fit_gpytorch_mll_scipy.
+        """
         self.train()
         train_x, train_y = mll.model.train_inputs[0], mll.model.train_targets
         optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs.copy()
@@ -375,6 +407,15 @@ def _fit_mll(
         return res
 
     def p_below_threshold(self, x: torch.Tensor, f_thresh: torch.Tensor) -> torch.Tensor: 
+        """Compute the probability that the latent function is below a threshold.
+        
+        Args:
+            x (torch.Tensor): Points at which to evaluate the probability.
+            f_thresh (torch.Tensor): Threshold value.
+            
+        Returns:
+            torch.Tensor: Probability that the latent function is below the threshold.
+        """
         f, var = self.predict(x)
         f_thresh = f_thresh.reshape(-1, 1)
         f = f.reshape(1, -1)

aepsych/models/gp_classification.py

@@ -69,10 +69,10 @@ def __init__(
                 gamma prior.
             likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to
                 Bernouli likelihood.
-            inducing_size (int, optional): Number of inducing points. Defaults to 99.
+            inducing_size (Optional[int], optional): Number of inducing points. Defaults to 99.
             max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. If None,
                 there is no limit to the fitting time.
-            inducing_point_method (string): The method to use to select the inducing points. Defaults to "auto".
+            inducing_point_method (string, optional): The method to use to select the inducing points. Defaults to "auto".
                 If "sobol", a number of Sobol points equal to inducing_size will be selected.
                 If "pivoted_chol", selects points based on the pivoted Cholesky heuristic.
                 If "kmeans++", selects points by performing kmeans++ clustering on the training data.
@@ -179,6 +179,7 @@ def from_config(cls, config: Config) -> GPClassificationModel:
         )
 
     def _reset_hyperparameters(self) -> None:
+        """Reset hyperparameters to their initial values."""
         # warmstart_hyperparams affects hyperparams but not the variational strat,
         # so we keep the old variational strat (which is only refreshed
         # if warmstart_induc=False).
@@ -190,6 +191,7 @@ def _reset_hyperparameters(self) -> None:
         self.likelihood.load_state_dict(self._fresh_likelihood_dict)
 
     def _reset_variational_strategy(self) -> None:
+        """Reset the variational strategy."""
         inducing_points = select_inducing_points(
             inducing_size=self.inducing_size,
             covar_module=self.covar_module,
@@ -221,9 +223,9 @@ def fit(
         Args:
             train_x (torch.Tensor): Inputs.
             train_y (torch.LongTensor): Responses.
-            warmstart_hyperparams (bool): Whether to reuse the previous hyperparameters (True) or fit from scratch
+            warmstart_hyperparams (bool, optional): Whether to reuse the previous hyperparameters (True) or fit from scratch
                 (False). Defaults to False.
-            warmstart_induc (bool): Whether to reuse the previous inducing points or fit from scratch (False).
+            warmstart_induc (bool, optional): Whether to reuse the previous inducing points or fit from scratch (False).
                 Defaults to False.
         """
         self.set_train_data(train_x, train_y)
@@ -300,10 +302,23 @@ def predict(
     def predict_probability(
         self, x: torch.Tensor
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Query the model for posterior mean and variance in probability space.
+        
+        Args:
+            x (torch.Tensor): Points at which to predict from the model.
+            
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Posterior mean and variance at queries points.
+        """
         return self.predict(x, probability_space=True)
 
     def update(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs):
-        """Perform a warm-start update of the model from previous fit."""
+        """Perform a warm-start update of the model from previous fit.
+        
+        Args:
+            train_x (torch.Tensor): Inputs.
+            train_y (torch.Tensor): Responses.
+           """
         return self.fit(
             train_x, train_y, warmstart_hyperparams=True, warmstart_induc=True, **kwargs
         )
@@ -324,6 +339,23 @@ def __init__(
         max_fit_time: Optional[float] = None,
         inducing_point_method: str = "auto",
     ) -> None:
+        """Initialize the GP Beta Regression model
+        
+        Args:
+            lb (torch.Tensor): Lower bounds of the parameters.
+            ub (torch.Tensor): Upper bounds of the parameters.
+            dim (Optional[int], optional): The number of dimensions in the parameter space. If None, it is inferred from the size
+                of lb and ub. Defaults to None.
+            mean_module (Optional[gpytorch.means.Mean], optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None.
+            covar_module (Optional[gpytorch.kernels.Kernel], optional): GP covariance kernel class. Defaults to scaled RBF with a
+                gamma prior.
+            likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to
+                Beta likelihood.
+            inducing_size (Optional[int], optional): Number of inducing points. Defaults to 100.
+            max_fit_time (Optional[float], optional): The maximum amount of time, in seconds, to spend fitting the model. If None,
+                there is no limit to the fitting time. Defaults to None.
+            inducing_point_method (string, optional): The method to use to select the inducing points. Defaults to "auto".
+            """
         if likelihood is None:
             likelihood = BetaLikelihood()
         super().__init__(
@@ -336,4 +368,4 @@ def __init__(
             inducing_size=inducing_size,
             max_fit_time=max_fit_time,
             inducing_point_method=inducing_point_method,
-        )
+        )

aepsych/models/gp_regression.py

@@ -33,8 +33,8 @@ class GPRegressionModel(AEPsychMixin, ExactGP):
 
     def __init__(
         self,
-        lb: Union[np.ndarray, torch.Tensor],
-        ub: Union[np.ndarray, torch.Tensor],
+        lb: torch.Tensor,
+        ub: torch.Tensor,
         dim: Optional[int] = None,
         mean_module: Optional[gpytorch.means.Mean] = None,
         covar_module: Optional[gpytorch.kernels.Kernel] = None,
@@ -44,8 +44,8 @@ def __init__(
         """Initialize the GP regression model
 
         Args:
-            lb (Union[numpy.ndarray, torch.Tensor]): Lower bounds of the parameters.
-            ub (Union[numpy.ndarray, torch.Tensor]): Upper bounds of the parameters.
+            lb (torch.Tensor): Lower bounds of the parameters.
+            ub (torch.Tensor): Upper bounds of the parameters.
             dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size
                 of lb and ub.
             mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior.
@@ -77,6 +77,14 @@ def __init__(
 
     @classmethod
     def construct_inputs(cls, config: Config) -> Dict:
+        """Construct inputs for the GP regression model from configuration.
+        
+        Args:
+            config (Config): A configuration containing keys/values matching this class.
+            
+        Returns:
+            Dict: Dictionary of inputs for the GP regression model.
+        """
         classname = cls.__name__
 
         lb = config.gettensor(classname, "lb")
@@ -118,7 +126,7 @@ def from_config(cls, config: Config) -> GPRegressionModel:
         from a configuration. TODO: document how this works in some tutorial.
 
         Args:
-            config (Config): A configuration containing keys/values matching this class
+            config (Config): A configuration containing keys/values matching this class.
 
         Returns:
             GPRegressionModel: Configured class instance.
@@ -140,7 +148,7 @@ def fit(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs) -> None:
         return self._fit_mll(mll, **kwargs)
 
     def sample(
-        self, x: Union[torch.Tensor, np.ndarray], num_samples: int
+        self, x: torch.Tensor, num_samples: int
     ) -> torch.Tensor:
         """Sample from underlying model.
 
@@ -155,11 +163,15 @@ def sample(
         return self.posterior(x).rsample(torch.Size([num_samples])).detach().squeeze()
 
     def update(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs):
-        """Perform a warm-start update of the model from previous fit."""
+        """Perform a warm-start update of the model from previous fit.
+        
+        Args:
+            train_x (torch.Tensor): Inputs.
+            train_y (torch.Tensor): Responses."""
         return self.fit(train_x, train_y, **kwargs)
 
     def predict(
-        self, x: Union[torch.Tensor, np.ndarray], **kwargs
+        self, x: torch.Tensor, **kwargs
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Query the model for posterior mean and variance.
 
@@ -169,7 +181,7 @@ def predict(
                 response probability instead of latent function value. Defaults to False.
 
         Returns:
-            Tuple[np.ndarray, np.ndarray]: Posterior mean and variance at queries points.
+            Tuple[torch.Tensor, torch.Tensor]: Posterior mean and variance at queries points.
         """
         with torch.no_grad():
             post = self.posterior(x)

aepsych/models/monotonic_projection_gp.py

@@ -105,6 +105,27 @@ def __init__(
         max_fit_time: Optional[float] = None,
         inducing_point_method: str = "auto",
     ) -> None:
+        """Initialize the MonotonicProjectionGP model.
+        
+        Args:
+            lb (torch.Tensor): Lower bounds of the parameters.
+            ub (torch.Tensor): Upper bounds of the parameters.
+            monotonic_dims (List[int]): A list of the dimensions on which monotonicity should
+                be enforced.
+            monotonic_grid_size (int, optional): The size of the grid, s, in 1. above. Defaults to 20.
+            min_f_val (Optional[float], optional): If provided, maintains this minimum in the projection in 5. Defaults to None.
+            dim (Optional[int], optional): The number of dimensions in the parameter space. If None, it is inferred from the size
+                of lb and ub. Defaults to None.
+            mean_module (Optional[gpytorch.means.Mean], optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None.
+            covar_module (Optional[gpytorch.kernels.Kernel], optional): GP covariance kernel class. Defaults to scaled RBF with a
+                gamma prior. Defaults to None.
+            likelihood (Optional[Likelihood], optional): The likelihood function to use. If None defaults to
+                Gaussian likelihood. Defaults to None.
+            inducing_size (Optional[int], optional): The number of inducing points to use. Defaults to None.
+            max_fit_time (Optional[float], optional): The maximum amount of time, in seconds, to spend fitting the model. If None,
+                there is no limit to the fitting time. Defaults to None.
+            inducing_point_method (string, optional): The method to use to select the inducing points. Defaults to "auto".
+            """
         assert len(monotonic_dims) > 0
         self.monotonic_dims = [int(d) for d in monotonic_dims]
         self.mon_grid_size = monotonic_grid_size
@@ -127,6 +148,16 @@ def posterior(
         observation_noise: Union[bool, torch.Tensor] = False,
         **kwargs: Any,
     ) -> GPyTorchPosterior:
+        """Compute the posterior at X, projecting to enforce monotonicity.
+
+        Args:
+            X (torch.Tensor): The input points at which to compute the posterior.
+            observation_noise (Union[bool, torch.Tensor], optional): Whether or not to include the observation noise in the
+                posterior. Defaults to False.
+
+        Returns:
+            GPyTorchPosterior: The posterior at X.
+        """
         # Augment X with monotonicity grid points, for each monotonic dim
         n, d = X.shape  # Require no batch dimensions
         m = len(self.monotonic_dims)
@@ -169,6 +200,15 @@ def posterior(
     def sample(
         self, x: torch.Tensor, num_samples: int
     ) -> torch.Tensor:
+        """Sample from the model.
+
+        Args:
+            x (torch.Tensor): The input points at which to sample.
+            num_samples (int): The number of samples to draw.
+
+        Returns:
+            torch.Tensor: The samples at x.
+        """
         samps = super().sample(x=x, num_samples=num_samples)
         if self.min_f_val is not None:
             samps = samps.clamp(min=self.min_f_val)