Enhance Docstrings in aepsych/acquisition (#424)

Summary:
Improves documentation in `aepsych/acquisition` for better clarity and consistency.

- Adds missing docstrings to functions and methods across all acquisition modules.
- Updates existing docstrings with refined type hints and a unified structure.

Pull Request resolved: #424

Reviewed By: crasanders

Differential Revision: D65924329

Pulled By: JasonKChow

fbshipit-source-id: afad0de60fd644166b293a89a317495232169e0f

aepsych/acquisition/bvn.py

@@ -15,6 +15,16 @@
 
 
 def _gauss_legendre20(dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Computes the abscissae and weights for the Gauss-Legendre quadrature of order 20.
+
+    Args:
+        dtype (torch.dtype): The desired data type of the output tensors.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]:
+            - abscissae: The quadrature points.
+            - weights: The corresponding weights for each quadrature point.
+    """
     _abscissae = torch.tensor(
         [
             0.9931285991850949,
@@ -54,6 +64,12 @@ def _gauss_legendre20(dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]:
 def _ndtr(x: torch.Tensor) -> torch.Tensor:
     """
     Standard normal CDF. Called <phid> in Genz's original code.
+
+    Args:
+        x (torch.Tensor): Input tensor of values.
+
+    Returns:
+        torch.Tensor: CDF values for each element in the input tensor.
     """
     return 0.5 * torch.erfc(_neg_inv_sqrt2 * x)
 
@@ -65,6 +81,14 @@ def _bvnu(
 ) -> torch.Tensor:
     """
     Primary subroutine for bvnu()
+
+    Args:
+        dh (torch.Tensor): Input tensor representing the first variable values.
+        dk (torch.Tensor): Input tensor representing the second variable values.
+        r (torch.Tensor): Input tensor for the correlation coefficient.
+
+    Returns:
+        torch.Tensor: Approximated bivariate normal CDF values.
     """
     # Precompute some terms
     h = dh
@@ -106,11 +130,12 @@ def bvn_cdf(
     Journal of Statist. Comput. Simul. 35, pp. 101-107.
 
     Args:
-        xu: Upper limits for cdf evaluation in x
-        yu: Upper limits for cdf evaluation in y
-        r: BVN correlation
+        xu (torch.Tensor): Upper limits for cdf evaluation in x
+        yu (torch.Tensor): Upper limits for cdf evaluation in y
+        r (torch.Tensor): BVN correlation
 
-    Returns: Tensor of cdf evaluations of same size as xu, yu, and r.
+    Returns: 
+        Tensor of cdf evaluations of same size as xu, yu, and r.
     """
     p = 1 - _ndtr(-xu) - _ndtr(-yu) + _bvnu(xu, yu, r)
     return torch.clip(p, 0, 1)