add docs for utils

docs/modules/root.rst

@@ -50,6 +50,13 @@ ABC
     :undoc-members:
 
 
+STS
+--------------
+.. automodule:: tsgm.models.sts
+    :members:
+    :undoc-members:
+
+
 Visualization
 --------------
 .. automodule:: tsgm.utils.visualization
@@ -62,3 +69,17 @@ Monitors
 .. automodule:: tsgm.models.monitors
     :members:
     :undoc-members:
+
+
+Datasets
+--------------
+.. automodule:: tsgm.utils.datasets
+    :members:
+    :undoc-members:
+
+
+Data Processing Utils
+--------------
+.. automodule:: tsgm.utils.data_processing
+    :members:
+    :undoc-members:

tsgm/models/monitors.py

@@ -19,6 +19,7 @@
 class GANMonitor(keras.callbacks.Callback):
     """
     GANMonitor is a Keras callback for monitoring and visualizing generated samples during training.
+
     :param num_samples: The number of samples to generate and visualize.
     :type num_samples: int
 

tsgm/models/sts.py

@@ -13,13 +13,33 @@
 
 
 class STS:
-    def __init__(self, model=None):
+    """
+    Class for training and generating from a structural time series model.
+    """
+
+    def __init__(self, model: tfp.sts.StructuralTimeSeries = None) -> None:
+        """
+        Initializes a new instance of the STS class.
+
+        :param model: Structural time series model to use. If None, default model is used.
+        :type model: tfp.sts.StructuralTimeSeriesModel or None
+        """
         self._model = model or DEFAULT_MODEL
         self._dist = None
         self._elbo_loss = None
 
     def train(self, ds: tsgm.dataset.Dataset, num_variational_steps: int = 200,
               steps_forw: int = 10) -> None:
+        """
+        Trains the structural time series model.
+
+        :param ds: Dataset containing time series data.
+        :type ds: tsgm.dataset.Dataset
+        :param num_variational_steps: Number of variational optimization steps, defaults to 200.
+        :type num_variational_steps: int
+        :param steps_forw: Number of steps to forecast, defaults to 10.
+        :type steps_forw: int
+        """
         assert ds.shape[0] == 1  # now works only with 1 TS
         X = ds.X.astype(np.float32)
         variational_posteriors = tfp.sts.build_factored_surrogate_posterior(
@@ -39,9 +59,24 @@ def train(self, ds: tsgm.dataset.Dataset, num_variational_steps: int = 200,
             parameter_samples=q_samples, num_steps_forecast=steps_forw)
 
     def elbo_loss(self) -> float:
+        """
+        Returns the evidence lower bound (ELBO) loss from training.
+
+        :returns: The value of the ELBO loss.
+        :rtype: float
+        """
         return self._elbo_loss
 
     def generate(self, num_samples: int) -> tsgm.types.Tensor:
+        """
+        Generates samples from the trained model.
+
+        :param num_samples: Number of samples to generate.
+        :type num_samples: int
+
+        :returns: Generated samples.
+        :rtype: tsgm.types.Tensor
+        """
         assert self._dist is not None
 
         return self._dist.sample(num_samples).numpy()[..., 0]

tsgm/utils/data_processing.py

@@ -7,33 +7,109 @@
 
 
 class TSGlobalScaler():
+    """
+    Scales time series data globally.
+
+    Attributes:
+    -----------
+    min : float
+        Minimum value encountered in the data.
+    max : float
+        Maximum value encountered in the data.
+    """
     def fit(self, X: TensorLike) -> "TSGlobalScaler":
+        """
+        Fits the scaler to the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: The fitted scaler object.
+        :rtype: TSGlobalScaler
+        """
         self.min = np.min(X)
         self.max = np.max(X)
         return self
 
     def transform(self, X: TensorLike) -> TensorLike:
+        """
+        Transforms the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: Scaled X.
+        :rtype: TensorLike
+        """
         return (X - self.min) / (self.max - self.min + EPS)
 
     def inverse_transform(self, X: TensorLike) -> TensorLike:
+        """
+        Inverse-transforms the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: Original data.
+        :rtype: TensorLike
+        """
         X *= (self.max - self.min + EPS)
         X += self.min
         return X
 
     def fit_transform(self, X: TensorLike) -> TensorLike:
+        """
+        Fits the scaler to the data and transforms it.
+
+        :parameter X: Input data
+        :type X: TensorLike
+
+        :returns: Scaled input data X
+        :rtype: TensorLike
+        """
         self.fit(X)
         scaled_X = self.transform(X)
         return scaled_X
 
 
 class TSFeatureWiseScaler():
+    """
+    Scales time series data feature-wise.
+
+    Parameters:
+    -----------
+    feature_range : tuple(float, float), optional
+        Tuple representing the minimum and maximum feature values (default is (0, 1)).
+
+    Attributes:
+    -----------
+    _min_v : float
+        Minimum feature value.
+    _max_v : float
+        Maximum feature value.
+    """
     def __init__(self, feature_range: T.Tuple[float, float] = (0, 1)) -> None:
+        """
+        Initializes a new instance of the TSFeatureWiseScaler class.
+
+        :parameter feature_range: Tuple representing the minimum and maximum feature values, defaults to (0, 1)
+        :type tuple(float, float), optional:
+        """
         assert len(feature_range) == 2
 
         self._min_v, self._max_v = feature_range
 
     # X: N x T x D
     def fit(self, X: TensorLike) -> "TSFeatureWiseScaler":
+        """
+        Fits the scaler to the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: The fitted scaler object.
+        :rtype: TSGlobalScaler
+        """
         D = X.shape[2]
         self.mins = np.zeros(D)
         self.maxs = np.zeros(D)
@@ -45,15 +121,42 @@ def fit(self, X: TensorLike) -> "TSFeatureWiseScaler":
         return self
 
     def transform(self, X: TensorLike) -> TensorLike:
+        """
+        Transforms the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: Scaled X.
+        :rtype: TensorLike
+        """
         return ((X - self.mins) / (self.maxs - self.mins + EPS)) * (self._max_v - self._min_v) + self._min_v
 
     def inverse_transform(self, X: TensorLike) -> TensorLike:
+        """
+        Inverse-transforms the data.
+
+        :parameter X: Input data.
+        :type X: TensorLike
+
+        :returns: Original data.
+        :rtype: TensorLike
+        """
         X -= self._min_v
         X /= self._max_v - self._min_v
         X *= (self.maxs - self.mins + EPS)
         X += self.mins
         return X
 
     def fit_transform(self, X: TensorLike) -> TensorLike:
+        """
+        Fits the scaler to the data and transforms it.
+
+        :parameter X: Input data
+        :type X: TensorLike
+
+        :returns: Scaled input data X
+        :rtype: TensorLike
+        """
         self.fit(X)
         return self.transform(X)