Add documentation for create_poly_basis

pcntoolkit/util/utils.py

@@ -42,11 +42,31 @@
 
 
 def create_poly_basis(X, dimpoly):
-    """ 
-    Compute a polynomial basis expansion of the specified order
-
     """
-
+    Creates a polynomial basis matrix for the given input matrix.
+
+    This function takes an input matrix `X` and a degree `dimpoly`, and returns a new matrix where each column is `X` raised to the power of a degree. The degrees range from 1 to `dimpoly`. If `X` is a 1D array, it is reshaped into a 2D array with one column.
+
+    Parameters
+    ----------
+    X : numpy.ndarray
+        The input matrix, a 2D array where each row is a sample and each column is a feature. If `X` is a 1D array, it is reshaped into a 2D array with one column.
+    dimpoly : int
+        The degree of the polynomial basis. The output matrix will have `dimpoly` times as many columns as `X`.
+
+    Returns
+    -------
+    Phi : numpy.ndarray
+        The polynomial basis matrix, a 2D array where each row is a sample and each column is a feature raised to a degree. The degrees range from 1 to `dimpoly`.
+
+    Examples
+    --------
+    >>> X = np.array([[1, 2], [3, 4], [5, 6]])
+    >>> create_poly_basis(X, 2)
+    array([[ 1.,  2.,  1.,  4.],
+           [ 3.,  4.,  9., 16.],
+           [ 5.,  6., 25., 36.]])
+    """
     if len(X.shape) == 1:
         X = X[:, np.newaxis]
     D = X.shape[1]
@@ -79,8 +99,8 @@ def create_design_matrix(X, intercept=True, basis='bspline',
                          **kwargs):
     """ 
     Prepare a design matrix from a set of covariates sutiable for
-        running Bayesian linar regression. This design matrix consists of 
-        a set of user defined covariates, optoinal site intercepts 
+        running Bayesian linear regression. This design matrix consists of 
+        a set of user defined covariates, optional site intercepts 
         (fixed effects) and also optionally a nonlinear basis expansion over 
         one of the columns