calpit 0.1__py3-none-any.whl

calpit/__init__.py

@@ -0,0 +1,9 @@
+<<<<<<< before updating
+from .diagnostics_and_calibration import CalPit
+
+__all__ = ["CalPit"]
+=======
+from .example_module import greetings, meaning
+
+__all__ = ["greetings", "meaning"]
+>>>>>>> after updating

calpit/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1'
+__version_tuple__ = version_tuple = (0, 1)

calpit/diagnostics_and_calibration.py

@@ -0,0 +1,305 @@
+from pathlib import Path
+import numpy as np
+import torch
+from torch.utils.data import TensorDataset, DataLoader
+from scipy.interpolate import PchipInterpolator
+from tqdm import trange
+
+
+from calpit.nn.models import MLP
+from calpit.nn.utils import count_parameters, RandomDataset, EarlyStopping
+from calpit.metrics import probability_integral_transform
+from calpit.utils import trapz_grid
+
+
+class CalPit:
+    def __init__(self, model, input_dim=None, hidden_layers=None, **args):
+        """
+        Initializes an instance of the CalPit Class.
+
+        Args:
+            model (str or torch.nn.Module): The model to be used to learn the conditional PIT.
+            Can be any pytorch model that outputs a value between 0 and 1.
+            A string with the name of an inbuilt can also be provided. Currently supports: `mlp`
+
+            input_dim (int, optional): The input dimension of the model. Defaults to None.
+            hidden_layers (list, optional): A list of hidden layer sizes for the MLP models. Defaults to None.
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if model == "mlp":
+            self.model = MLP(input_dim + 1, hidden_layers, 1).to(self.device)
+        else:
+            self.model = model.to(self.device)
+
+        count_parameters(self.model)
+
+        self.training_loss = None
+        self.validation_bce = None
+        self.val_loss_min = None
+
+    def fit(
+        self,
+        x_calib,
+        y_calib=None,
+        cde_calib=None,
+        y_grid=None,
+        pit_calib=None,
+        oversample=1,
+        n_cov_val=201,
+        patience=20,
+        n_epochs=1000,
+        lr=0.001,
+        weight_decay=1e-5,
+        batch_size=2048,
+        frac_train=0.9,
+        lr_decay=0.99,
+        trace_func=print,
+        seed=299792458,
+        num_workers=1,
+        checkpt_path="_results/checkpoint_.pt",
+    ):
+        """
+        Train the model using the calibration data.
+
+        Args:
+            x_calib (numpy.ndarray): The input features for calibration data.
+            y_calib (numpy.ndarray, optional): The target values for calibration data.
+            cde_calib (numpy.ndarray, optional): The conditional density estimates for calibration.
+            y_grid (numpy.ndarray, optional): The grid of target values for calibration.
+            pit_calib (numpy.ndarray, optional): The probability integral transforms for the given CDEs evaluated at y_calib.
+            Either pit_calib or y_calib, cde_calib and y_grid must be provided.
+            oversample (int, optional): The oversampling factor for the training data. Default is 1.
+            This is used to upsample the number of coverage values used for training.
+            n_cov_val (int, optional): The number of coverage values to use for validation. Default is 201.
+            patience (int, optional): The number of epochs to wait for improvement in validation loss before early stopping. Default is 20.
+            n_epochs (int, optional): The maximum number of epochs for training. Default is 1000.
+            lr (float, optional): The initial learning rate for the optimizer (AdamW). Default is 0.001.
+            weight_decay (float, optional): The weight decay for the optimizer. Default is 1e-5.
+            batch_size (int, optional): The batch size for training and validation. Default is 2048.
+            frac_train (float, optional): The fraction of data to use for training.
+            The rest is used for the validation set used to determine when to stop training. Default is 0.9.
+            lr_decay (float, optional): The learning rate decay factor for the rule,
+            learning_rate(epoch) = lr*lr_decay ** epoch. Default is 0.99.
+            trace_func (function, optional): The function used for printing training progress. Default is print.
+            seed (int, optional): The random seed for reproducibility. Default is 299792458.
+            num_workers (int, optional): The number of CPU worker threads for data loading. Default is 1.
+            checkpt_path (str, optional): The path to save the checkpoint of the best model. Default is "_results/checkpoint_.pt".
+
+        Returns:
+            torch.nn.Module: The trained model.
+        """
+        # method implementation
+        if pit_calib is None:
+            if y_calib is None or cde_calib is None or y_grid is None:
+                raise ValueError("Either pit_calib or, y_calib, cde_calib and y_grid must be provided")
+            pit_calib = probability_integral_transform(cde_calib, y_grid, y_calib)
+
+        cov_grid = np.linspace(0.001, 0.999, n_cov_val)
+        # Split into train and valid sets
+        train_size = int(frac_train * len(x_calib))
+        valid_size = len(x_calib) - train_size
+
+        rnd_idx = np.random.default_rng(seed=seed).permutation(len(x_calib))
+        x_train_rnd = x_calib[rnd_idx[:train_size]]
+        x_val_rnd = x_calib[rnd_idx[train_size:]]
+        pit_train_rand = pit_calib[rnd_idx[:train_size]]
+        pit_val_rand = pit_calib[rnd_idx[train_size:]]
+
+        # Creat randomized Data set for training
+        trainset = RandomDataset(x_train_rnd, pit_train_rand, oversample=oversample)
+
+        # Create static dataset for validation
+        feature_val = torch.cat(
+            [
+                torch.Tensor(np.repeat(cov_grid, len(x_val_rnd)))[:, None],
+                torch.Tensor(np.tile(x_val_rnd, (n_cov_val, 1))),
+            ],
+            dim=-1,
+        )
+        target_val = torch.Tensor(
+            np.tile(pit_val_rand, n_cov_val) <= np.repeat(cov_grid, len(x_val_rnd))
+        ).float()[:, None]
+
+        validset = TensorDataset(feature_val, target_val)
+
+        # Create Data loader
+        train_dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+        valid_dataloader = DataLoader(validset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+        # Initialize the Model and optimizer, etc.
+        training_loss = []
+        validation_bce = []
+
+        # Optimizer
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=weight_decay)
+        # Use lr decay
+        schedule_rule = lambda epoch: lr_decay**epoch
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=schedule_rule)
+        checkpt_path = Path(checkpt_path)
+        checkpt_path.parent.mkdir(parents=True, exist_ok=True)
+        early_stopping = EarlyStopping(
+            patience=patience, verbose=True, path=checkpt_path, trace_func=trace_func
+        )
+
+        # Training loop for all epochs
+        for epoch in range(1, n_epochs + 1):
+            training_loss_batch = []
+            validation_mse_batch = []
+            validation_bce_batch = []
+
+            # Training loop per epoch
+            self.model.train()  # prep model for training
+            for batch, (feature, target) in enumerate(train_dataloader, start=1):
+                feature = feature.to(self.device)
+                target = target.to(self.device)
+
+                # Zero your gradients for every batch!
+                optimizer.zero_grad()
+                # Make predictions for this batch
+                output = self.model(feature.float())
+
+                # Compute the loss and its gradients
+
+                loss_fn = torch.nn.BCELoss(reduction="sum")
+                loss = loss_fn(torch.clamp(torch.squeeze(output), min=0.0, max=1.0), torch.squeeze(target))
+
+                loss.backward()
+                # Adjust learning weights
+                optimizer.step()
+
+                # record training loss
+                training_loss_batch.append(loss.item())
+
+            # Validation
+            self.model.eval()  # prep model for evaluation
+
+            for feature, target in valid_dataloader:
+
+                feature = feature.to(self.device)
+                target = target.to(self.device)
+
+                # forward pass: compute predicted outputs by passing inputs to the model
+                output = self.model(feature.float())
+
+                # calculate the loss
+                mse = ((output - target.float()) ** 2).sum()
+                # record validation loss
+                validation_mse_batch.append(mse.item())
+
+                criterion = torch.nn.BCELoss(reduction="sum")
+                bce = criterion(torch.clamp(torch.squeeze(output), min=0, max=1), torch.squeeze(target))
+                validation_bce_batch.append(bce.item())
+
+            # calculate average loss over an epoch
+            train_loss_epoch = np.sum(training_loss_batch) / (train_size * oversample)
+            valid_bce_epoch = np.sum(validation_bce_batch) / (valid_size * n_cov_val)
+            training_loss.append(train_loss_epoch)
+            validation_bce.append(valid_bce_epoch)
+
+            epoch_len = len(str(n_epochs))
+            # print training/validation statistics
+            msg = (
+                f"[{epoch:>{epoch_len}}/{n_epochs:>{epoch_len}}] | "
+                + f"train_loss: {train_loss_epoch:.5f} |"
+                + f"valid_bce: {valid_bce_epoch:.5f} | "
+            )
+
+            trace_func(msg)
+
+            # change the lr
+            scheduler.step()
+
+            # early_stopping needs the validation loss to check if it has decresed,
+            # and if it has, it will make a checkpoint of the current model
+            early_stopping(valid_bce_epoch, self.model)
+
+            if early_stopping.early_stop:
+                print("Early stopping")
+                break
+
+        # # load the last checkpoint with the best model
+        self.model.load_state_dict(torch.load(checkpt_path))
+        self.training_loss = np.array(training_loss)
+        self.validation_bce = np.array(validation_bce)
+        self.val_loss_min = early_stopping.val_loss_min
+        return self.model
+
+    def predict(self, x_test, cov_grid, batch_size=2048):
+        """
+        Predicts the conditional PIT values for the given test data and coverage grid.
+
+        Args:
+            x_test (numpy.ndarray): The input features of the test data.
+            cov_grid (numpy.ndarray): The coverage grid at which the PIT values are to be evaluated.
+            batch_size (int, optional): The batch size for prediction. Defaults to 2048.
+
+        Returns:
+            numpy.ndarray: The predicted conditional PIT values.
+        """
+        self.model.eval()
+        self.model.to(self.device)
+
+        pred_pit = []
+        n_test = len(x_test)
+        n_cov = len(cov_grid)
+        n_batches = (n_test - 1) // batch_size + 1
+
+        for i in trange(n_batches):
+            x = x_test[i * batch_size : (i + 1) * batch_size]
+            if cov_grid.ndim == 1:
+                with torch.no_grad():
+                    pred_pit_batch = (
+                        self.model(
+                            torch.Tensor(
+                                np.hstack([np.repeat(cov_grid, len(x))[:, None], np.tile(x, (n_cov, 1))])
+                            ).to(self.device)
+                        )
+                        .detach()
+                        .cpu()
+                        .numpy()
+                        .reshape(n_cov, -1)
+                        .T
+                    )
+            elif cov_grid.ndim == 2:
+                c = cov_grid[i * batch_size : (i + 1) * batch_size]
+                with torch.no_grad():
+                    pred_pit_batch = (
+                        self.model(
+                            torch.Tensor(
+                                np.hstack([np.ravel(c)[:, None], np.repeat(x, c.shape[1], axis=0)])
+                            ).to(self.device)
+                        )
+                        .detach()
+                        .cpu()
+                        .numpy()
+                        .reshape(len(x), -1)
+                    )
+
+            pred_pit_batch[pred_pit_batch < 0] = 0
+            pred_pit_batch[pred_pit_batch > 1] = 1
+            pred_pit.extend(pred_pit_batch)
+        return np.array(pred_pit)
+
+    def transform(self, x_test, cde_test, y_grid, batch_size=2048):
+        """
+        Transforms the input CDEs for a test data set to calibrated CDEs.
+
+        Args:
+            x_test (array-like): The input features of the test data.
+            cde_test (array-like): The initial CDEs for the test data that is to be transformed.
+            y_grid (array-like): The grid of values for the CDEs.
+            batch_size (int, optional): The batch size for prediction. Defaults to 2048.
+
+        Returns:
+            numpy.ndarray: The transformed CDEs for the given.
+        """
+        cdf_test = trapz_grid(cde_test, y_grid)
+        cdf_test_new = self.predict(x_test, cov_grid=cdf_test, batch_size=batch_size)
+        cdf_funct = PchipInterpolator(y_grid, cdf_test_new, extrapolate=True, axis=1)
+        pdf_func = cdf_funct.derivative(1)
+        cde_test_new = pdf_func(y_grid)
+        return cde_test_new
+
+    def fit_transform(self, **args):
+        """Fit the model and transform the data in one go"""
+        raise NotImplementedError

calpit/example_benchmarks.py

@@ -0,0 +1,14 @@
+"""An example module containing simplistic methods under benchmarking."""
+
+import random
+import time
+
+
+def runtime_computation():
+    """Runtime computation consuming between 0 and 5 seconds."""
+    time.sleep(random.uniform(0, 5))
+
+
+def memory_computation():
+    """Memory computation for a random list up to 512 samples."""
+    return [0] * random.randint(0, 512)

calpit/metrics.py

@@ -0,0 +1,150 @@
+import numpy as np
+
+
+def cde_loss(cde_estimates: np.ndarray, y_grid: np.ndarray, y_test: np.ndarray) -> tuple:
+    """
+    Calculates conditional density estimation loss on holdout data.
+
+    Args:
+        cde_estimates (numpy.array): An array where each row is a density estimate on y_grid.
+        z_grid (numpy.array): An array of the grid points at which cde_estimates is evaluated.
+        z_test (numpy.array): An array of the true y values corresponding to the rows of cde_estimates.
+
+    Returns:
+        tuple: A tuple containing the loss and the standard error of the loss.
+
+    Raises:
+        ValueError: If the dimensions of the input tensors are not compatible.
+    """
+
+    if len(y_test.shape) == 1:
+        y_test = y_test.reshape(-1, 1)
+    if len(y_grid.shape) == 1:
+        y_grid = y_grid.reshape(-1, 1)
+
+    n_obs, n_grid = cde_estimates.shape
+    n_samples, feats_samples = y_test.shape
+    n_grid_points, feats_grid = y_grid.shape
+
+    if n_obs != n_samples:
+        raise ValueError(
+            f"Number of samples in CDEs should be the same as in z_test.Currently {n_obs} and {n_samples}."
+        )
+    if n_grid != n_grid_points:
+        raise ValueError(
+            f"Number of grid points in CDEs should be the same as in z_grid. Currently {n_grid} and {n_grid_points}."
+        )
+
+    if feats_samples != feats_grid:
+        raise ValueError(
+            f"Dimensionality of test points and grid points need to coincise. Currently {feats_samples} and {feats_grid}."
+        )
+
+    integrals = np.trapz(cde_estimates**2, np.squeeze(y_grid), axis=1)
+
+    nn_ids = np.argmin(np.abs(y_grid - y_test.T), axis=0)
+    likeli = cde_estimates[(tuple(np.arange(n_samples)), tuple(nn_ids))]
+
+    losses = integrals - 2 * likeli
+    loss = np.mean(losses)
+    se_error = np.std(losses, axis=0) / (n_obs**0.5)
+
+    return loss, se_error
+
+
+def kolmogorov_smirnov_statistic(cdf_test: np.ndarray, cdf_ref: np.ndarray) -> np.ndarray:
+    """
+    Calculate the Kolmogorov-Smirnov statistic between two cumulative distribution functions (CDFs).
+
+    Parameters:
+    cdf_test (np.ndarray): CDF of the test distribution.
+    cdf_ref (np.ndarray): CDF of the reference distribution on the same grid.
+
+    Returns:
+    np.ndarray: The Kolmogorov-Smirnov statistic.
+
+    """
+    ks = np.max(np.abs(cdf_test - cdf_ref), axis=-1)
+
+    return ks
+
+
+def cramer_von_mises_statistic(cdf_test: np.ndarray, cdf_ref: np.ndarray) -> np.ndarray:
+    """
+    Calculates the Cramer-von Mises statistic between two cumulative distribution functions (CDFs).
+
+    Args:
+        cdf_test (np.ndarray): CDF of the test distribution.
+        cdf_ref (np.ndarray): CDF of the reference distribution on the same grid.
+
+    Returns:
+        np.ndarray: The Cramer-von Mises statistic.
+
+    """
+    diff = (cdf_test - cdf_ref) ** 2
+
+    cvm2 = np.trapz(diff, cdf_ref, axis=-1)
+    return np.sqrt(cvm2)
+
+
+def anderson_darling_statistic(cdf_test: np.ndarray, cdf_ref: np.ndarray, n_tot: int = 1) -> np.ndarray:
+    """
+    Calculates the Anderson-Darling statistic between two cumulative distribution functions (CDFs).
+
+    Args:
+        cdf_test (np.ndarray): CDF of the test distribution (1D array).
+        cdf_ref (np.ndarray): CDF of the reference distribution on the same grid (1D array).
+        n_tot (int): Scaling factor equal to the number of PDFs used to construct ECDF.
+
+    Returns:
+        np.ndarray: The Anderson-Darling statistic.
+
+    """
+    num = (cdf_test - cdf_ref) ** 2
+    den = cdf_ref * (1 - cdf_ref)
+
+    ad2 = n_tot * np.trapz((num / den), cdf_ref, axis=-1)
+    return np.sqrt(ad2)
+
+
+def probability_integral_transform(cde: np.ndarray, y_grid: np.ndarray, y_test: np.ndarray) -> np.ndarray:
+    """
+    Calculates the Probability Integral Transform (PIT) based on Conditional Density Estimates (CDE).
+
+    Args:
+        cde (np.ndarray): A numpy array of conditional density estimates.
+            Each row corresponds to an observation, each column corresponds to a grid point.
+        y_grid (np.ndarray): A numpy array of the grid points at which cde is evaluated.
+        y_test (np.ndarray): A numpy array of the true y values corresponding to the rows of cde.
+
+    Returns:
+        np.ndarray: A numpy array of PIT values.
+
+    Raises:
+        ValueError: If the number of samples in cde is not the same as in y_test,
+            or if the number of grid points in cde is not the same as in y_grid.
+
+    """
+    # flatten the input arrays to 1D
+    y_grid = np.ravel(y_grid)
+    y_test = np.ravel(y_test)
+
+    # Sanity checks
+    nrow_cde, ncol_cde = cde.shape
+    n_samples = y_test.shape[0]
+    n_grid_points = y_grid.shape[0]
+
+    if nrow_cde != n_samples:
+        raise ValueError(
+            f"Number of samples in CDEs should be the same as in z_test. Currently {nrow_cde} and {n_samples}."
+        )
+    if ncol_cde != n_grid_points:
+        raise ValueError(
+            f"Number of grid points in CDEs should be the same as in z_grid. Currently {nrow_cde} and {n_grid_points}."
+        )
+
+    # Vectorized implementation using masked arrays
+    pit = np.ma.masked_array(cde, (y_grid > y_test[:, np.newaxis]))
+    pit = np.trapz(pit, y_grid)
+
+    return np.array(pit)

calpit/nn/__init__.py

@@ -0,0 +1,2 @@
+from .umnn import MonotonicNN  # noqa
+from .ispline_nn import IsplineNN  # noqa

calpit/nn/ispline_nn.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+from splinebasis import ISplineBasis
+
+
+class ISplineLayer(nn.Module):
+    def __init__(self, in_features, num_basis,dropout_p=0):
+        super().__init__()
+        self.in_features = in_features
+        self.num_basis = num_basis
+        self.coefs = nn.Sequential(nn.Linear(in_features, num_basis), nn.Softmax(dim=-1),nn.Dropout(p=dropout_p))
+        self.grid = torch.linspace(0, 1, 1000)
+        self.basis_vectors = ISplineBasis(
+            order=3, num_basis=num_basis, lower=0, upper=1, grid=self.grid
+        ).basis_vectors
+        self.basis_vectors = torch.from_numpy(self.basis_vectors)
+
+#         def init_weights(m):
+#             if isinstance(m, nn.Linear):
+#                 torch.nn.init.kaiming_normal_(m.weight)
+#                 m.bias.data.fill_(0.01)
+
+#         self.coefs.apply(init_weights)
+
+    def interp1d(self, x, y, x_new):
+        # 2. Find where in the original data, the values to interpolate
+        #    would be inserted.
+        #    Note: If x_new[n] == x[m], then m is returned by searchsorted.
+        # y = torch.moveaxis(y,axis,0)
+        # y = y.reshape((y.shape[0],-1))
+
+        x_new_indices = torch.searchsorted(x, x_new)
+
+        # 3. Clip x_new_indices so that they are within the range of
+        #    self.x indices and at least 1. Removes mis-interpolation
+        #    of x_new[n] = x[0]
+        x_new_indices = x_new_indices.clip(1, len(x) - 1)
+
+        # 4. Calculate the slope of regions that each x_new value falls in.
+        lo = x_new_indices - 1
+        hi = x_new_indices
+
+        x_lo = x[lo]
+        x_hi = x[hi]
+        y_lo = y[lo]
+        y_hi = y[hi]
+
+        # Note that the following two expressions rely on the specifics of the
+        # broadcasting semantics.
+        slope = (y_hi - y_lo) / (x_hi - x_lo)[:, None]
+
+        # 5. Calculate the actual value for each entry in x_new.
+        y_new = slope * (x_new - x_lo)[:, None] + y_lo
+
+        return y_new
+
+    def forward(self, x, alpha):
+        grid = self.grid.to(alpha)
+        basis_vectors = self.basis_vectors.to(alpha)
+        basis = self.interp1d(grid, basis_vectors, alpha)
+
+        # print(basis.shape)
+        # print(self.coefs(x).shape)
+        # print(self.coefs(x))
+        weighted_basis = self.coefs(x) * basis
+        # print(weighted_basis.shape)
+        return weighted_basis.sum(axis=-1)
+
+
+class IsplineNN(nn.Module):
+    def __init__(self, input_dim, hidden_layers=[512, 512, 512],dropout_p=0.5, num_basis=10):
+        super().__init__()
+        self.all_layers = [input_dim + 1]
+        self.hidden_layers = hidden_layers
+        self.all_layers.extend(hidden_layers)
+        self.num_basis = num_basis
+        self.dropout_p = dropout_p
+        self.spline_layer = ISplineLayer(in_features=self.hidden_layers[-1], num_basis=self.num_basis,dropout_p=self.dropout_p)
+
+        self.mlp_layer_list = []
+        for i in range(len(self.all_layers) - 1):
+            self.mlp_layer_list.append(nn.Linear(self.all_layers[i], self.all_layers[i + 1]))
+            self.mlp_layer_list.append(nn.PReLU())
+
+        # self.mlp_layer_list.append(nn.Dropout(p=dropout_p))
+        self.mlp_layers = nn.Sequential(*self.mlp_layer_list)
+
+        def init_weights(m):
+            if isinstance(m, nn.Linear):
+                torch.nn.init.kaiming_normal_(m.weight)
+                m.bias.data.fill_(0.01)
+
+        self.mlp_layers.apply(init_weights)
+
+    def forward(self, x):
+        alpha = x[:, 0]
+
+        res = self.mlp_layers(x)
+
+        res = self.spline_layer(res, alpha)
+
+        return res

calpit/nn/models.py

@@ -0,0 +1,52 @@
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    """
+    Multi-Layer Perceptron (MLP) neural network model.
+
+    Args:
+        input_dim (int): The number of input features.
+        hidden_layers (list): A list of integers representing the number of units in each hidden layer.
+        output_dim (int): The number of output units. Defaults to 1.
+        sigmoid (bool): Whether to apply a sigmoid activation function to the output layer. Defaults to True.
+
+    Methods:
+        forward(x): Performs a forward pass through the MLP.
+
+    """
+
+    def __init__(self, input_dim, hidden_layers, output_dim=1, sigmoid=True):
+        super().__init__()
+        self.all_layers = [input_dim]
+        self.all_layers.extend(hidden_layers)
+        self.all_layers.append(output_dim)
+        self.layer_list = []
+
+        for i in range(len(self.all_layers) - 1):
+            self.layer_list.append(nn.Linear(self.all_layers[i], self.all_layers[i + 1]))
+            self.layer_list.append(nn.PReLU())
+        self.layer_list.pop()
+        if sigmoid:
+            self.layer_list.append(nn.Sigmoid())
+        self.layers = nn.Sequential(*self.layer_list)
+
+        def init_weights(m):
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight)
+                m.bias.data.fill_(0.01)
+
+        self.layers.apply(init_weights)
+
+    def forward(self, x):
+        """
+        Performs a forward pass through the MLP.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor.
+
+        """
+        return self.layers(x)