Enhance Docstrings in aepsych/factory

aepsych/factory/default.py

@@ -35,6 +35,7 @@ def default_mean_covar_factory(
             config details).
         dim (int, optional): Dimensionality of the parameter space. Must be provided
             if config is None.
+        stimuli_per_trial (int): Number of stimuli per trial. Defaults to 1.
 
     Returns:
         Tuple[gpytorch.means.Mean, gpytorch.kernels.Kernel]: Instantiated
@@ -68,6 +69,14 @@ def default_mean_covar_factory(
 def _get_default_mean_function(
     config: Optional[Config] = None,
 ) -> gpytorch.means.ConstantMean:
+    """Creates a default mean function for Gaussian Processes.
+
+    Args:
+        config (Config, optional): Configuration object.
+
+    Returns:
+        gpytorch.means.ConstantMean: An instantiated mean function with appropriate priors based on the configuration.
+    """
     # default priors
     fixed_mean = False
     mean = gpytorch.means.ConstantMean()
@@ -93,6 +102,18 @@ def _get_default_cov_function(
     stimuli_per_trial: int,
     active_dims: Optional[List[int]] = None,
 ) -> gpytorch.kernels.Kernel:
+    """Creates a default covariance function for Gaussian Processes.
+
+    Args:
+        config (Config, optional): Configuration object.
+        dim (int): Dimensionality of the parameter space.
+        stimuli_per_trial (int): Number of stimuli per trial.
+        active_dims (List[int], optional): List of dimensions to use in the covariance function. Defaults to None.
+
+    Returns:
+        gpytorch.kernels.Kernel: An instantiated kernel with appropriate priors based on the configuration.
+    """
+
     # default priors
     lengthscale_prior = "lognormal" if stimuli_per_trial == 1 else "gamma"
     ls_loc = torch.tensor(math.sqrt(2.0), dtype=torch.float64)
@@ -178,4 +199,4 @@ def _get_default_cov_function(
             outputscale_prior=os_prior,
             outputscale_constraint=gpytorch.constraints.GreaterThan(1e-4),
         )
-    return covar
+    return covar

aepsych/factory/ordinal.py

@@ -17,6 +17,16 @@
 def ordinal_mean_covar_factory(
     config: Config,
 ) -> Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]:
+    """ Create a mean and covariance function for ordinal GPs.
+
+    Args:
+        config (Config): Config object containing bounds.
+
+    Returns:
+        Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]: A tuple containing
+        the mean function (ConstantMean) and the covariance function (ScaleKernel).
+    """
+
     try:
         base_factory = config.getobj("ordinal_mean_covar_factory", "base_factory")
     except NoOptionError:
@@ -31,4 +41,4 @@ def ordinal_mean_covar_factory(
     else:
         covar = base_covar
 
-    return mean, covar
+    return mean, covar

aepsych/factory/pairwise.py

@@ -15,6 +15,14 @@
 def pairwise_mean_covar_factory(
     config: Config,
 ) -> Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]:
+    """ Creates a mean and covariance function for pairwise GPs.
+
+    Args:
+        config (Config): Config object containing bounds.
+
+    Returns:
+        Tuple[gpytorch.means.ConstantMean, gpytorch.kernels.ScaleKernel]: A tuple containing
+        the mean function (ConstantMean) and the covariance function (ScaleKernel)."""
     lb = config.gettensor("common", "lb")
     ub = config.gettensor("common", "ub")
     assert lb.shape[0] == ub.shape[0], "bounds shape mismatch!"
@@ -65,4 +73,4 @@ def pairwise_mean_covar_factory(
         cov = _get_default_cov_function(config, config_dim // 2, stimuli_per_trial=1)
         covar = PairwiseKernel(cov)
 
-    return mean, covar
+    return mean, covar