fix proposed changes in the docstrings across models files

aepsych/models/base.py

@@ -10,7 +10,7 @@
 
 import time
 from collections.abc import Iterable
-from typing import Any, Dict, List, Mapping, Optional, Protocol, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol, Tuple, Union
 
 import gpytorch
 import numpy as np
@@ -132,7 +132,7 @@ def get_max(
             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.
-            max_time (Optional[float], optional): Maximum time to spend optimizing. Defaults to None.
+            max_time (float, optional): Maximum time to spend optimizing. Defaults to None.
 
         Returns:
             Tuple[float, torch.Tensor]: Tuple containing the max and its location (argmax).
@@ -162,7 +162,7 @@ def get_min(
             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.
-            max_time (Optional[float], optional): Maximum time to spend optimizing. Defaults to None.
+            max_time (float, optional): Maximum time to spend optimizing. Defaults to None.
 
         Returns:
             Tuple[float, torch.Tensor]: Tuple containing the min and its location (argmin).
@@ -198,8 +198,8 @@ def inv_query(
             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.
+            max_time (float, optional): Maximum time to spend optimizing. Defaults to None.
+            weights (torch.Tensor, optional): Weights for the optimization. Defaults to None.
 
         Returns:
             Tuple[float, torch.Tensor]: Tuple containing the value of f
@@ -243,7 +243,7 @@ def get_jnd(
         Both definitions are equivalent for linear psychometric functions.
 
         Args:
-            grid (Optional[Union[np.ndarray, torch.Tensor]], optional): Mesh grid over which to find the JND.
+            grid (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.
@@ -254,9 +254,6 @@ def get_jnd(
             method (str, optional): "taylor" or "step" method (see docstring).
                 Defaults to "step".
 
-        Raises:
-            RuntimeError: for passing an unknown method.
-
         Returns:
             Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: either the
                 mean JND, or a median, lower, upper tuple of the JND posterior.
@@ -324,15 +321,16 @@ def dim_grid(
         
         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."""
+            slice_dims (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):
-        """
+        """Set the training data for the model.
         Args:
-            inputs (Optional[torch.Tensor], optional):  The new training inputs.
-            targets (Optional[torch.Tensor], optional): The new training targets.
+            inputs (torch.Tensor, optional):  The new training inputs.
+            targets (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
@@ -351,7 +349,8 @@ def normalize_inputs(self, x: torch.Tensor) -> torch.Tensor:
             x (torch.Tensor): Tensor of points to normalize.
             
         Returns:
-            torch.Tensor: Normalized tensor of points."""
+            torch.Tensor: Normalized tensor of points.
+        """
         scale = self.ub - self.lb
         return (x - self.lb) / scale
 
@@ -375,15 +374,15 @@ def _fit_mll(
         self,
         mll: MarginalLogLikelihood,
         optimizer_kwargs: Optional[Dict[str, Any]] = None,
-        optimizer=fit_gpytorch_mll_scipy,
+        optimizer: Callable = 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.
+            optimizer_kwargs (Dict[str, Any], optional): Keyword arguments for the optimizer.
+            optimizer (Callable): Optimizer to use. Defaults to fit_gpytorch_mll_scipy.
         """
         self.train()
         train_x, train_y = mll.model.train_inputs[0], mll.model.train_targets

aepsych/models/gp_classification.py

@@ -69,7 +69,7 @@ def __init__(
                 gamma prior.
             likelihood (gpytorch.likelihood.Likelihood, optional): The likelihood function to use. If None defaults to
                 Bernouli likelihood.
-            inducing_size (Optional[int], optional): Number of inducing points. Defaults to 99.
+            inducing_size (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, optional): The method to use to select the inducing points. Defaults to "auto".
@@ -344,15 +344,15 @@ def __init__(
         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
+            dim (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
+            mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None.
+            covar_module (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,
+            inducing_size (int, optional): Number of inducing points. Defaults to 100.
+            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. Defaults to None.
             inducing_point_method (string, optional): The method to use to select the inducing points. Defaults to "auto".
             """

aepsych/models/gp_regression.py

@@ -154,8 +154,7 @@ def sample(
 
         Args:
             x (torch.Tensor): Points at which to sample.
-            num_samples (int, optional): Number of samples to return. Defaults to None.
-            kwargs are ignored
+            num_samples (int): Number of samples to return.
 
         Returns:
             torch.Tensor: Posterior samples [num_samples x dim]
@@ -167,7 +166,8 @@ def update(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs):
         
         Args:
             train_x (torch.Tensor): Inputs.
-            train_y (torch.Tensor): Responses."""
+            train_y (torch.Tensor): Responses.
+        """
         return self.fit(train_x, train_y, **kwargs)
 
     def predict(

aepsych/models/monotonic_projection_gp.py

@@ -113,16 +113,16 @@ def __init__(
             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
+            min_f_val (float, optional): If provided, maintains this minimum in the projection in 5. Defaults to None.
+            dim (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
+            mean_module (gpytorch.means.Mean, optional): GP mean class. Defaults to a constant with a normal prior. Defaults to None.
+            covar_module (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
+            likelihood (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,
+            inducing_size (int, optional): The number of inducing points to use. Defaults to None.
+            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. Defaults to None.
             inducing_point_method (string, optional): The method to use to select the inducing points. Defaults to "auto".
             """

aepsych/models/monotonic_rejection_gp.py

@@ -71,12 +71,12 @@ def __init__(
             constraints.
             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.
-            covar_module (Optional[Kernel], optional): Covariance kernel to use. Default is scaled RBF.
-            mean_module (Optional[Mean], optional): Mean module to use. Default is constant mean.
+            dim (int, optional): The number of dimensions in the parameter space. If None, it is inferred from the size.
+            covar_module (Kernel, optional): Covariance kernel to use. Default is scaled RBF.
+            mean_module (Mean, optional): Mean module to use. Default is constant mean.
             likelihood (str): Link function and likelihood. Can be 'probit-bernoulli' or
                 'identity-gaussian'.
-            fixed_prior_mean (Optional[float], optional): Fixed prior mean. If classification, should be the prior
+            fixed_prior_mean (float, optional): Fixed prior mean. If classification, should be the prior
             classification probability (not the latent function value). Defaults to None.
             num_induc (int, optional): Number of inducing points for variational GP.]. Defaults to 25.
             num_samples (int, optional): Number of samples for estimating posterior on preDict or
@@ -178,8 +178,8 @@ def _set_model(
         Args:
             train_x (Tensor): Training x points
             train_y (Tensor): Training y points. Should be (n x 1).
-            model_state_dict (Optional[Dict[str, Tensor]], optional): State dict for the model
-            likelihood_state_dict (Optional[Dict[str, Tensor]], optional): State dict for the likelihood
+            model_state_dict (Dict[str, Tensor], optional): State dict for the model
+            likelihood_state_dict (Dict[str, Tensor], optional): State dict for the likelihood
             """
         train_x_aug = self._augment_with_deriv_index(train_x, 0)
         self.set_train_data(train_x_aug, train_y)
@@ -229,8 +229,8 @@ def sample(
 
         Args:
             x (Tensor): tensor of n points at which to sample
-            num_samples (Optional[int], optional): how many points to sample. Default is self.num_samples.
-            num_rejection_samples (Optional[int], optional): how many samples to use for rejection sampling. Default is self.num_rejection_samples.
+            num_samples (int, optional): how many points to sample. Default is self.num_samples.
+            num_rejection_samples (int, optional): how many samples to use for rejection sampling. Default is self.num_rejection_samples.
 
         Returns: a Tensor of shape [n_samp, n]
         """

aepsych/models/multitask_regression.py

@@ -51,10 +51,10 @@ def __init__(
             num_outputs (int, optional): Number of tasks (outputs). Defaults to 2.
             rank (int, optional): Rank of cross-task covariance. Lower rank is a simpler model.
                 Should be less than or equal to num_outputs. Defaults to 1.
-            mean_module (Optional[gpytorch.means.Mean], optional): GP mean. Defaults to a constant mean.
-            covar_module (Optional[gpytorch.kernels.Kernel], optional): GP kernel module.
+            mean_module (gpytorch.means.Mean, optional): GP mean. Defaults to a constant mean.
+            covar_module (gpytorch.kernels.Kernel, optional): GP kernel module.
                 Defaults to scaled RBF kernel.
-            likelihood (Optional[gpytorch.likelihoods.Likelihood], optional): Likelihood
+            likelihood (gpytorch.likelihoods.Likelihood, optional): Likelihood
                 (should be a multitask-compatible likelihood). Defaults to multitask Gaussian likelihood.
         """
         self._num_outputs = num_outputs
@@ -86,7 +86,8 @@ def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultitaskMultivaria
             
         Returns:
             gpytorch.distributions.MultitaskMultivariateNormal: Distribution object
-                holding the mean and covariance at x."""
+                holding the mean and covariance at x.
+                """
         transformed_x = self.normalize_inputs(x)
         mean_x = self.mean_module(transformed_x)
         covar_x = self.covar_module(transformed_x)
@@ -132,10 +133,10 @@ def __init__(
 
         Args:
             num_outputs (int, optional): Number of tasks (outputs). Defaults to 2.
-            mean_module (Optional[gpytorch.means.Mean], optional): GP mean. Defaults to a constant mean.
-            covar_module (Optional[gpytorch.kernels.Kernel], optional): GP kernel module.
+            mean_module (gpytorch.means.Mean, optional): GP mean. Defaults to a constant mean.
+            covar_module (gpytorch.kernels.Kernel, optional): GP kernel module.
                 Defaults to scaled RBF kernel.
-            likelihood (Optional[gpytorch.likelihoods.Likelihood], optional): Likelihood
+            likelihood (gpytorch.likelihoods.Likelihood, optional): Likelihood
                 (should be a multitask-compatible likelihood). Defaults to multitask Gaussian likelihood.
         """
 
@@ -171,7 +172,8 @@ def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultitaskMultivaria
 
         Returns:
             gpytorch.distributions.MultitaskMultivariateNormal: Distribution object
-                holding the mean and covariance at x."""
+                holding the mean and covariance at x.
+        """
         base_mvn = super().forward(x)  # do transforms
         return gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(
             base_mvn

aepsych/models/ordinal_gp.py

@@ -29,7 +29,7 @@ def __init__(self, likelihood=None, *args, **kwargs):
         """Initialize the OrdinalGPModel
         
         Args:
-            likelihood (Optional[Likelihood], optional): The likelihood function to use. If None defaults to
+            likelihood (Likelihood, optional): The likelihood function to use. If None defaults to
                 Ordinal likelihood.
         """
         covar_module = kwargs.pop("covar_module", None)

aepsych/models/pairwise_probit.py

@@ -32,7 +32,7 @@ def _pairs_to_comparisons(self, x: torch.Tensor, y: torch.Tensor) -> Tuple[torch
         """Convert pairs of points and their judgements to comparisons.
 
         Args:
-            x (torch.Tensor): Tensor of shape (n, 2, d) where n is the number of pairs and d is the dimensionality of the
+            x (torch.Tensor): Tensor of shape (n, d, 2) where n is the number of pairs and d is the dimensionality of the
                 parameter space.
             y (torch.Tensor): Tensor of shape (n,) where n is the number of pairs. Each element is 0 if the first point
                 in the pair is preferred, and 1 if the second point is preferred.
@@ -77,11 +77,11 @@ def __init__(
         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
+            dim (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.
-            covar_module (Optional[gpytorch.kernels.Kernel], optional): GP covariance kernel class. Defaults to scaled RBF with a
+            covar_module (gpytorch.kernels.Kernel, optional): GP covariance kernel class. Defaults to scaled RBF with a
                 gamma prior. Defaults to None.
-            max_fit_time (Optional[float], optional): The maximum amount of time, in seconds, to spend fitting the model. Defaults to None.
+            max_fit_time (float, optional): The maximum amount of time, in seconds, to spend fitting the model. Defaults to None.
             """
         self.lb, self.ub, dim = _process_bounds(lb, ub, dim)
 
@@ -124,7 +124,7 @@ def fit(
         Args:
             train_x (torch.Tensor): Trainin x points.
             train_y (torch.Tensor): Training y points.
-            optimizer_kwargs (Optional[Dict[str, Any]], optional): Keyword arguments to pass to the optimizer. Defaults to None.
+            optimizer_kwargs (Dict[str, Any], optional): Keyword arguments to pass to the optimizer. Defaults to None.
         """
         self.train()
         mll = PairwiseLaplaceMarginalLogLikelihood(self.likelihood, self)