From 5dd11fc636c3dcfce3f1f099fef51c715fdceb20 Mon Sep 17 00:00:00 2001 From: Yousif Alsaffar Date: Thu, 31 Oct 2024 23:45:43 +0300 Subject: [PATCH] fix proposed changes in the docstrings across models files --- aepsych/models/base.py | 35 +++++++++++------------ aepsych/models/gp_classification.py | 12 ++++---- aepsych/models/gp_regression.py | 6 ++-- aepsych/models/monotonic_projection_gp.py | 14 ++++----- aepsych/models/monotonic_rejection_gp.py | 16 +++++------ aepsych/models/multitask_regression.py | 18 ++++++------ aepsych/models/ordinal_gp.py | 2 +- aepsych/models/pairwise_probit.py | 10 +++---- aepsych/models/semi_p.py | 26 +++++++++-------- aepsych/models/utils.py | 3 +- 10 files changed, 73 insertions(+), 69 deletions(-) diff --git a/aepsych/models/base.py b/aepsych/models/base.py index f2f36ad05..5b0486aca 100644 --- a/aepsych/models/base.py +++ b/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 diff --git a/aepsych/models/gp_classification.py b/aepsych/models/gp_classification.py index e84c5f8b0..b6b409a13 100644 --- a/aepsych/models/gp_classification.py +++ b/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". """ diff --git a/aepsych/models/gp_regression.py b/aepsych/models/gp_regression.py index cb0c274b4..990dc8209 100644 --- a/aepsych/models/gp_regression.py +++ b/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( diff --git a/aepsych/models/monotonic_projection_gp.py b/aepsych/models/monotonic_projection_gp.py index d8173de82..0911d502f 100644 --- a/aepsych/models/monotonic_projection_gp.py +++ b/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". """ diff --git a/aepsych/models/monotonic_rejection_gp.py b/aepsych/models/monotonic_rejection_gp.py index 970776be0..e350c5992 100644 --- a/aepsych/models/monotonic_rejection_gp.py +++ b/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] """ diff --git a/aepsych/models/multitask_regression.py b/aepsych/models/multitask_regression.py index d2e662c55..322f3a43a 100644 --- a/aepsych/models/multitask_regression.py +++ b/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 diff --git a/aepsych/models/ordinal_gp.py b/aepsych/models/ordinal_gp.py index 1587c1078..bf2ca038d 100644 --- a/aepsych/models/ordinal_gp.py +++ b/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) diff --git a/aepsych/models/pairwise_probit.py b/aepsych/models/pairwise_probit.py index 68cc37baa..b645a4553 100644 --- a/aepsych/models/pairwise_probit.py +++ b/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) diff --git a/aepsych/models/semi_p.py b/aepsych/models/semi_p.py index 52554892a..d6d2c8935 100644 --- a/aepsych/models/semi_p.py +++ b/aepsych/models/semi_p.py @@ -144,8 +144,8 @@ def rsample( """Sample from the posterior distribution using the reparameterization trick Args: - sample_shape (Optional[torch.Size], optional): The desired shape of the samples. Defaults to None. - base_samples (Optional[torch.Tensor], optional): The base samples. Defaults to None. + sample_shape (torch.Size, optional): The desired shape of the samples. Defaults to None. + base_samples (torch.Tensor, optional): The base samples. Defaults to None. Returns: torch.Tensor: The sampled values from the posterior distribution. @@ -171,11 +171,12 @@ def sample_p( """Sample from the likelihood distribution of the modeled function. Args: - sample_shape (Optional[torch.Size], optional): The desired shape of the samples. Defaults to None. - base_samples (Optional[torch.Tensor], optional): The base samples. Defaults to None. + sample_shape (torch.Size, optional): The desired shape of the samples. Defaults to None. + base_samples (torch.Tensor, optional): The base samples. Defaults to None. Returns: - torch.Tensor: The sampled values from the likelihood distribution.""" + torch.Tensor: The sampled values from the likelihood distribution. + """ kcsamps = self.rsample(sample_shape=sample_shape, base_samples=base_samples) return self.likelihood.p(function_samples=kcsamps, Xi=self.Xi).squeeze(-1) @@ -187,11 +188,12 @@ def sample_f( """Sample from the function values of the modeled distribution. Args: - sample_shape (Optional[torch.Size], optional): The desired shape of the samples. Defaults to None. - base_samples (Optional[torch.Tensor], optional): The base samples. Defaults to None. + sample_shape (torch.Size, optional): The desired shape of the samples. Defaults to None. + base_samples (torch.Tensor, optional): The base samples. Defaults to None. Returns: - torch.Tensor: The sampled function values from the likelihood.""" + torch.Tensor: The sampled function values from the likelihood. + """ kcsamps = self.rsample(sample_shape=sample_shape, base_samples=base_samples) return self.likelihood.f(function_samples=kcsamps, Xi=self.Xi).squeeze(-1) @@ -206,8 +208,8 @@ def sample_thresholds( Args: threshold_level (float): The target threshold level for sampling. - sample_shape (Optional[torch.Size], optional): The desired shape of the samples. Defaults to None. - base_samples (Optional[torch.Tensor], optional): The base samples. Defaults to None. + sample_shape (torch.Size, optional): The desired shape of the samples. Defaults to None. + base_samples (torch.Tensor, optional): The base samples. Defaults to None. Returns: SemiPThresholdObjective: The sampled thresholds based on the threshold level. @@ -448,7 +450,7 @@ def posterior(self, X: torch.Tensor, posterior_transform: Optional[PosteriorTran Args: X (torch.Tensor): Points at which to evaluate the posterior. - posterior_transform (Optional[PosteriorTransform], optional): A transform to apply to the posterior. Defaults to None. + posterior_transform (PosteriorTransform, optional): A transform to apply to the posterior. Defaults to None. Returns: SemiPPosterior: The posterior distribution at the given points. @@ -522,7 +524,7 @@ def __init__( offset_mean_module (gpytorch.means.Mean, optional): Mean module to use (default: constant mean) for offset. offset_covar_module (gpytorch.kernels.Kernel, optional): Covariance kernel to use (default: scaled RBF) for offset. likelihood (gpytorch.likelihood.Likelihood, optional)): defaults to bernoulli with logistic input and a floor of .5 - 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". diff --git a/aepsych/models/utils.py b/aepsych/models/utils.py index 145f42526..9cb2503db 100644 --- a/aepsych/models/utils.py +++ b/aepsych/models/utils.py @@ -169,6 +169,7 @@ def get_extremum( inverse is along a slice of the full surface. n_samples (int): number of coarse grid points to sample for optimization estimate. max_time (float): Maximum amount of time in seconds to spend optimizing. + weights (Tensor, optional): Weights to apply to the target value. Defaults to None. Returns: Tuple[float, np.ndarray]: Tuple containing the min and its location (argmin). """ @@ -259,7 +260,7 @@ def __init__( Args: target_value (Union[float, Tensor]): The target value to transform the posterior to. - weights (Optional[Tensor]): Weights to apply to the target value. Defaults to None. + weights (Tensor, optional): Weights to apply to the target value. Defaults to None. """ super().__init__() self.target_value = target_value