--- a
+++ b/src/iterpretability/simulators.py
@@ -0,0 +1,1498 @@
+# stdlib
+import random
+from typing import Tuple
+import src.iterpretability.logger as log
+
+# third party
+import numpy as np
+import torch
+from scipy.special import expit
+from scipy.stats import zscore
+from omegaconf import DictConfig, OmegaConf
+from src.iterpretability.utils import enable_reproducible_results
+from abc import ABC, abstractmethod
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# For computing the propensities from scores
+from scipy.special import softmax
+from scipy.stats import zscore
+from sklearn.model_selection import train_test_split
+
+EPS = 0
+class SimulatorBase(ABC):
+    """
+    Base class for simulators.
+    """
+    @abstractmethod
+    def simulate(self, X: np.ndarray, outcomes: np.ndarray = None) -> Tuple:
+        raise NotImplementedError
+    
+    @abstractmethod
+    def get_simulated_data(self, train_ratio: float) -> Tuple:
+        raise NotImplementedError
+    
+    @property
+    @abstractmethod
+    def selective_features(self) -> np.ndarray:
+        raise NotImplementedError
+    
+    @property
+    @abstractmethod
+    def prognostic_features(self) -> np.ndarray:
+        raise NotImplementedError
+    
+    @property
+    @abstractmethod
+    def predictive_features(self) -> np.ndarray:
+        raise NotImplementedError
+    
+class TYSimulator(SimulatorBase):
+    """
+    Data generation process class for simulating treatment selection and outcomes (and effects)
+    """
+    nonlinear_fcts = [
+                #lambda x: np.abs(x),
+                lambda x: np.exp(-(x**2) / 2),
+                #  lambda x: 1 / (1 + x**2),
+                # lambda x: np.sqrt(x)*(1+x),
+                #lambda x: np.cos(5*x),
+                #lambda x: x**2,
+                # lambda x: np.arctan(x),
+                # lambda x: np.tanh(x),
+                # lambda x: np.sin(x),
+                # lambda x: np.log(1 + x**2),
+                #lambda x: np.sqrt(0.02 + x**2),
+                #lambda x: np.cosh(x),
+            ]
+    
+    def __init__(
+        self,
+        # Data dimensionality
+        dim_X: int,
+
+        # Seed
+        seed: int = 42,
+
+        # Simulation type
+        simulation_type: str = "ty",
+
+        # Dimensionality of treatments and outcome
+        num_binary_outcome: int = 0,
+        outcome_unbalancedness_ratio: float = 0,
+        standardize_outcome: bool = False,
+        num_T: int = 3,
+        dim_Y: int = 3,
+
+        # Scale parameters
+        predictive_scale: float = 1,
+        prognostic_scale: float = 1,
+        propensity_scale: float = 1,
+        unbalancedness_exp: float = 0,
+        nonlinearity_scale: float = 1,
+        propensity_type: str = "prog_pred",
+        alpha: float = 0.5,
+        enforce_balancedness: bool = False,
+
+        # Control
+        include_control: bool = False,
+
+        # Important features
+        num_pred_features: int = 5,
+        num_prog_features: int = 5,
+        num_select_features: int = 5,
+        feature_type_overlap: str = "sel_none",
+        treatment_feature_overlap: bool = False,
+
+        # Feature selection
+        random_feature_selection: bool = False,
+        nonlinearity_selection_type: bool = True,
+
+        # Noise
+        noise: bool = True,
+        noise_std: float = 0.1,
+        
+    ) -> None:
+        # Number of features
+        self.dim_X = dim_X
+
+        # Make sure results are reproducible by setting seed for np, torch, random
+        self.seed = seed
+        enable_reproducible_results(seed=self.seed)
+
+        # Simulation type
+        self.simulation_type = simulation_type
+
+        # Store dimensions
+        self.num_binary_outcome = num_binary_outcome
+        self.outcome_unbalancedness_ratio = outcome_unbalancedness_ratio
+        self.standardize_outcome = standardize_outcome
+        self.num_T = num_T
+        self.dim_Y = dim_Y
+
+        # Scale parameters
+        self.predictive_scale = predictive_scale
+        self.prognostic_scale = prognostic_scale
+        self.propensity_scale = propensity_scale
+        self.unbalancedness_exp = unbalancedness_exp
+        self.nonlinearity_scale = nonlinearity_scale
+        self.propensity_type = propensity_type
+        self.alpha = alpha
+        self.enforce_balancedness = enforce_balancedness
+
+        # Control
+        self.include_control = include_control
+
+        # Important features
+        self.num_pred_features = num_pred_features
+        self.num_prog_features = num_prog_features
+        self.num_select_features = num_select_features
+        self.num_important_features = self.num_T*(num_pred_features + num_select_features) + num_prog_features
+        self.feature_type_overlap = feature_type_overlap
+        self.treatment_feature_overlap = treatment_feature_overlap
+
+        # Feature selection
+        self.random_feature_selection = random_feature_selection
+        self.nonlinearity_selection_type = nonlinearity_selection_type
+
+        # Noise
+        self.noise = noise
+        self.noise_std = noise_std
+
+        # Setup variables
+        self.nonlinearities = None
+        self.prog_mask, self.pred_masks, self.select_masks = None, None, None
+        self.prog_weights, self.pred_weights, self.select_weights = None, None, None
+
+        # Setup
+        self.setup()
+
+        # Simulation variables
+        self.X = None
+        self.prog_scores, self.pred_scores, self.select_scores = None, None, None
+        self.select_scores_pred_overlap = None
+        self.select_scores_prog_overlap = None
+        self.propensities, self.outcomes, self.T, self.Y = None, None, None, None
+    
+    def get_simulated_data(self):
+        """
+        Extract results and split into training and test set. Include counterfactual outcomes and propensities.
+        """
+        return self.X, self.T, self.Y, self.outcomes, self.propensities
+
+        ## OLD CODE
+        # Split data
+        # train_size = int(train_ratio * self.X.shape[0])
+
+        # if self.num_binary_outcome > 0:
+        #     (
+        #         X_train, X_test, 
+        #         Y_train, Y_test, 
+        #         T_train, T_test,
+        #         outcomes_train, outcomes_test,
+        #         propensities_train, propensities_test,
+        #     ) = train_test_split(self.X, self.Y, self.T, self.outcomes, self.propensities, train_size=train_size, stratify=self.Y)
+        # else:
+        #     X_train, X_test = self.X[:train_size], self.X[train_size:]
+        #     T_train, T_test = self.T[:train_size], self.T[train_size:]
+        #     Y_train, Y_test = self.Y[:train_size], self.Y[train_size:]
+
+        #     outcomes_train, outcomes_test = self.outcomes[:train_size,:,:], self.outcomes[train_size:,:,:]
+        #     propensities_train, propensities_test = self.propensities[:train_size], self.propensities[train_size:]
+
+        # if train_ratio == 1:
+        #     return self.X, self.T, self.Y, self.outcomes, self.propensities
+        
+        # return X_train, X_test, T_train, T_test, Y_train, Y_test, outcomes_train, outcomes_test, propensities_train, propensities_test
+
+    def simulate(self, X, outcomes=None) -> Tuple:
+        """
+        Simulate treatment and outcome for a dataset based on the configuration.
+        """
+        log.debug(
+            f'Simulating treatment and outcome for a dataset with:'
+            f'\n==================================================================='
+            f'\nDim X: {self.dim_X}'
+            f'\nDim T: {self.num_T}'
+            f'\nDim Y: {self.dim_Y}'
+            f'\nPredictive Scale: {self.predictive_scale}'
+            f'\nPrognostic Scale: {self.prognostic_scale}'
+            f'\nPropensity Scale: {self.propensity_scale}'
+            f'\nUnbalancedness Exponent: {self.unbalancedness_exp}'
+            f'\nNonlinearity Scale: {self.nonlinearity_scale}'
+            f'\nNum Pred Features: {self.num_pred_features}'
+            f'\nNum Prog Features: {self.num_prog_features}'
+            f'\nNum Select Features: {self.num_select_features}'
+            f'\nFeature Overlap: {self.treatment_feature_overlap}'
+            f'\nRandom Feature Selection: {self.random_feature_selection}'
+            f'\nNonlinearity Selection Type: {self.nonlinearity_selection_type}'
+            f'\nNoise: {self.noise}'
+            f'\nNoise Std: {self.noise_std}'
+            f'\n===================================================================\n'
+        )
+
+        # 1. Store data with min max scaling to range [0, 1]
+        self.X = X
+        # self.X = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0) + EPS)
+
+        # 2. Compute scores for prognostic, predictive, and selective features
+        self.compute_scores()
+
+        # 3. Compute factual and counterfactual outcomes based on the data and the predictive and prognostic scores
+        self.compute_all_outcomes()
+
+        # 4. Compute propensities based on the data and the selective scores
+        self.compute_propensities()
+
+        # 5. Sample treatment assignment based on the propensities
+        self.sample_T()
+
+        # 6. Extract the outcome based on the treatment assignment
+        self.extract_Y()
+
+        return None
+    
+    def setup(self) -> None:
+        """
+        Setup the simulator by defining variables which remain the same across simulations with different samples but the same configuration.
+        """
+        # 1. Sample nonlinearities used 
+        num_nonlinearities = 2 + self.dim_Y # Different non-linearities for each outcome (predictive), same for all treatments
+        self.nonlinearities = self.sample_nonlinearities(num_nonlinearities)
+
+        # 2. Set important feature masks - determine which features should be used for treatment selection, outcome prediction
+        self.sample_important_feature_masks()
+
+        # 3. Sample weights for features
+        self.sample_uniform_weights()
+
+    def get_true_cates(self, 
+                       X: np.ndarray, 
+                       T: np.ndarray, 
+                       outcomes: np.ndarray) -> np.ndarray:
+        """
+        Compute true CATEs for each treatment based on the data and the outcomes.
+        Always use the selected treatment as the base treatment.
+        """
+        # Compute CATEs for each treatment
+        cates = np.zeros((X.shape[0], self.num_T, self.dim_Y))
+
+        for i in range(X.shape[0]):
+            for j in range(self.num_T):
+                cates[i,j,:] = outcomes[i,j,:] - outcomes[i,int(T[i]),:]
+
+        log.debug(
+            f'\nCheck if true CATEs are computed correctly:'
+            f'\n==================================================================='
+            f'\nOutcomes: {outcomes.shape}'
+            f'\n{outcomes}'
+            f'\n\nTreatment Assignment: {T.shape}'
+            f'\n{T}'
+            f'\n\nTrue CATEs: {cates.shape}'
+            f'\n{cates}'
+            f'\n===================================================================\n'
+        )
+
+        return cates
+    
+    def extract_Y(self) -> None:
+        """
+        Extract the outcome based on the treatment assignment.
+        """
+        self.Y = self.outcomes[np.arange(self.X.shape[0]), self.T] 
+
+        log.debug(
+            f'\nCheck if outcomes are extracted correctly:'
+            f'\n==================================================================='
+            f'\nOutcomes'
+            f'\n{self.outcomes}'
+            f'\n{self.outcomes.shape}'
+            f'\n\nTreatment Assignment'
+            f'\n{self.T}'
+            f'\n{self.T.shape}'
+            f'\n\nExtracted Outcomes'
+            f'\n{self.Y}'
+            f'\n{self.Y.shape}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+
+    def compute_all_outcomes_toy(self) -> None:
+        # Compute outcomes for each treatment and outcome
+        outcomes = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))
+        X0 = self.X[:,0]
+        X1 = self.X[:,1]
+
+        k=20
+        nonlinearity = lambda x: 1 / (1 + np.exp(-k * (x - 0.5))) #logistic
+        
+        if self.propensity_type.startswith("toy1") or self.propensity_type.startswith("toy3") or self.propensity_type.startswith("toy4"):
+            fun_y0 = lambda X0, X1: X0
+            fun_y1 = lambda X0, X1: 1-X0
+
+        elif self.propensity_type.startswith("toy2"):
+            fun_y0 = lambda X0, X1: X0
+            fun_y1 = lambda X0, X1: 1-X1
+
+        elif self.propensity_type.startswith("toy6"):
+            fun_y0 = lambda X0, X1: X0
+            fun_y1 = lambda X0, X1: X1
+
+        elif self.propensity_type.startswith("toy5"):
+            fun_y0 = lambda X0, X1: np.sin(X0*10*np.pi)
+            fun_y1 = lambda X0, X1: np.sin((1-X0)*10*np.pi)
+
+        elif self.propensity_type.startswith("toy7"):
+            fun_y0 = lambda X0, X1: nonlinearity(X0)-nonlinearity(X1)
+            fun_y1 = lambda X0, X1: nonlinearity(X0)+nonlinearity(X1)
+
+        elif self.propensity_type.startswith("toy8"):
+            fun_y0 = lambda X0, X1: X0
+            fun_y1 = lambda X0, X1: 1-X0
+
+        Y = np.array([fun_y0(X0, X1),fun_y1(X0, X1)]).T
+
+        if self.propensity_type.endswith("nonlinear"):
+            Y = nonlinearity(Y)
+
+        Y = zscore(Y, axis=None)
+
+        outcomes[:,:,0] = Y
+
+
+        return outcomes
+
+    def compute_all_outcomes(self) -> None:
+        """
+        Compute factual and counterfactual outcomes based on the data and the predictive and prognostic scores.
+        """
+        if self.propensity_type.startswith("toy"):
+            outcomes = self.compute_all_outcomes_toy()
+
+        else:
+            # Compute outcomes for each treatment and outcome
+            outcomes = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))
+
+            for i in range(self.num_T):
+                for j in range(self.dim_Y):
+                    if self.include_control and i == 0:
+                        outcomes[:,i,j] = self.prognostic_scale*self.prog_scores[:,j]
+
+                    else:
+                        outcomes[:,i,j] = self.prognostic_scale*self.prog_scores[:,j] + self.predictive_scale*self.pred_scores[:,i,j]
+
+        # Add gaussian noise to outcomes
+        if self.noise:
+            outcomes = outcomes + np.random.normal(0, self.noise_std, size=outcomes.shape)
+
+        # Create binary outcomes and introduce unbalancedness
+        if int(self.num_binary_outcome) > 0:
+            for j in range(self.num_binary_outcome):
+                scores = zscore(outcomes[:,:,j], axis=0)
+                prob = expit(scores)
+                outcomes[:,:,j] = prob > self.outcome_unbalancedness_ratio
+                
+        self.outcomes = outcomes
+
+        # Standardize outcomes
+        if self.standardize_outcome:
+            # normalize outcomes per outcome
+            self.outcomes = zscore(self.outcomes, axis=0)
+
+        log.debug(
+            f'\nCheck if outcomes are computed correctly:'
+            f'\n==================================================================='
+            f'\nProg Scores'
+            f'\n{self.prog_scores}'
+            f'\n{self.prog_scores.shape}'
+            f'\n\nPred Scores'
+            f'\n{self.pred_scores}'
+            f'\n{self.pred_scores.shape}'
+            f'\n\nOutcomes'
+            f'\n{self.outcomes}'
+            f'\n{self.outcomes.shape}'
+            f'\n\nMean Outcomes'
+            f'\n{self.outcomes.mean(axis=0)}'
+            f'\n\nVariance Outcomes'
+            f'\n{self.outcomes.var(axis=0)}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+
+
+    def sample_T(self) -> None:
+        """
+        Sample treatment assignment based on the propensities.
+        """
+        # Sample from the resulting categorical distribution per row
+        self.T = np.array([np.random.choice([tre for tre in range(self.propensities.shape[1])], p=row) for row in self.propensities])
+
+        log.debug(
+            f'\nCheck if treatment assignment is sampled correctly:'
+            f'\n==================================================================='
+            f'\nPropensities'
+            f'\n{self.propensities}'
+            f'\n{self.propensities.shape}'
+            f'\n\nTreatment Assignment'
+            f'\n{self.T}'
+            f'\n{self.T.shape}'
+            f'\n\nUnique Treatment Counts'
+            f'\n{np.unique(self.T, return_counts=True)}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+    
+    def get_unbalancedness_weights(self, size: int) -> np.ndarray:
+        """
+        Create weights for introducing unbalancedness for class probabilities.
+        """
+        # Sample initial distribution of treatment assignment 
+        unb_weights = np.random.uniform(0, 1, size=size) 
+        unb_weights = unb_weights / unb_weights.sum()
+
+        # Standardize the weights and make sure that a treatment doesn't completely disappear for small unbalancedness exponents
+        min_val = unb_weights.min()
+        range_val = unb_weights.max() - min_val
+        unb_weights = (unb_weights - min_val) / range_val
+        unb_weights = 0.01 + unb_weights * 0.98
+
+        return unb_weights
+    
+    def compute_propensity_scores_toy(self) -> np.ndarray:
+        X0 = self.X[:,0]
+        X1 = self.X[:,1]
+
+        if self.propensity_type.startswith("toy1"):
+            fun_t0 = lambda X0, X1: X0
+            fun_t1 = lambda X0, X1: 1-X0
+
+        elif self.propensity_type.startswith("toy2"):
+            fun_t0 = lambda X0, X1: X0
+            fun_t1 = lambda X0, X1: 1-X1
+
+        elif self.propensity_type.startswith("toy3"):
+            fun_t0 = lambda X0, X1: X1
+            fun_t1 = lambda X0, X1: 1-X1
+
+        elif self.propensity_type.startswith("toy4"):
+            fun_t0 = lambda X0, X1: np.sin(X0*10*np.pi)
+            fun_t1 = lambda X0, X1: np.sin((1-X0)*10*np.pi)
+
+        elif self.propensity_type.startswith("toy5"):
+            fun_t0 = lambda X0, X1: 1-X0
+            fun_t1 = lambda X0, X1: X0
+
+        elif self.propensity_type.startswith("toy6"):
+            fun_t0 = lambda X0, X1: 1-X0
+            fun_t1 = lambda X0, X1: X0
+
+        elif self.propensity_type.startswith("toy7"):
+            fun_t0 = lambda X0, X1: 1-X0
+            fun_t1 = lambda X0, X1: X0
+
+        elif self.propensity_type.startswith("toy8"):
+            fun_t0 = lambda X0, X1: 1-X0
+            fun_t1 = lambda X0, X1: X0
+
+        scores = np.array([fun_t0(X0, X1),fun_t1(X0, X1)]).T
+
+        return scores
+
+
+    def compute_propensities(self) -> None:
+        """
+        Compute propensities based on the data and the selective scores.
+        """
+        
+        select_scores_pred_overlap = zscore(self.select_scores_pred_overlap, axis=0) # Comment for Predictive Epertise
+        select_scores_prog_overlap = zscore(self.select_scores_prog_overlap, axis=0) # Comment for Predictive Epertise
+        select_scores_none = zscore(self.select_scores, axis=0) # Comment for Predictive Epertise
+
+        select_scores_pred = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_pred_flipped = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_prog = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_tre = np.zeros((self.X.shape[0], self.num_T))
+
+        select_scores_pred[:,0] = self.outcomes[:,0,0] - self.outcomes[:,1,0]
+        select_scores_pred[:,1] = self.outcomes[:,1,0] - self.outcomes[:,0,0]
+
+        select_scores_pred_flipped[:,0] = self.outcomes[:,1,0] - self.outcomes[:,0,0]
+        select_scores_pred_flipped[:,1] = self.outcomes[:,0,0] - self.outcomes[:,1,0]
+
+        select_scores_prog[:,0] = self.outcomes[:,0,0]
+        select_scores_prog[:,1] = -self.outcomes[:,0,0]
+
+        select_scores_tre[:,0] = -self.outcomes[:,1,0]
+        select_scores_tre[:,1] = self.outcomes[:,1,0]
+
+        if self.propensity_type == "prog_tre":
+            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_prog
+
+        # Standardize all scores
+        select_scores_pred = zscore(select_scores_pred, axis=0)
+        select_scores_pred_flipped = zscore(select_scores_pred_flipped, axis=0)
+        select_scores_prog = zscore(select_scores_prog, axis=0)
+        select_scores_tre = zscore(select_scores_tre, axis=0)
+
+        if self.propensity_type == "prog_pred":
+            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_prog
+
+        elif self.propensity_type == "prog_tre":
+            pass
+
+        elif self.propensity_type == "none_prog":
+            scores = self.alpha * select_scores_prog + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_pred":
+            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_tre":
+            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_pred_flipped":
+            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "pred_pred_flipped":
+            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_pred
+
+        elif self.propensity_type == "none_pred_overlap":
+            scores = self.alpha * select_scores_pred_overlap + (1 - self.alpha) * select_scores_none
+            
+        elif self.propensity_type == "none_prog_overlap":
+            scores = self.alpha * select_scores_prog_overlap + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "pred_overalp_prog_overlap":
+            scores = self.alpha * select_scores_prog_overlap + (1 - self.alpha) * select_scores_pred_overlap
+
+        elif self.propensity_type == "rct_none":
+            scores = select_scores_none
+
+        elif self.propensity_type.startswith("toy"):
+            scores = self.compute_propensity_scores_toy()
+
+        else:
+            raise ValueError(f"Unknown propensity type {self.propensity_type}.")
+
+        if self.enforce_balancedness:
+            scores = zscore(scores, axis=0)
+
+        if self.propensity_type == "rct_none":
+            scores = self.alpha * select_scores_none
+
+        # Introduce unbalancedness and manipulate unbalancedness weights for comparable experiments with different seeds
+        unb_weights = self.get_unbalancedness_weights(size=scores.shape[1])
+
+        # Apply the softmax function to each row to get probabilities
+        p = softmax(self.propensity_scale*scores, axis=1)
+
+        # Scale probabilities to introduce unbalancedness
+        p = p * (1 - unb_weights) ** self.unbalancedness_exp
+
+        # Make sure rows add up to one again
+        row_sums = p.sum(axis=1, keepdims=True)
+        p = p / row_sums
+        self.propensities = p
+
+        log.debug(
+            f'\nCheck if propensities are computed correctly:'
+            f'\n==================================================================='
+            f'\nSelect Scores'
+            f'\n{self.select_scores}'
+            f'\n{self.select_scores.shape}'
+            f'\n\nPropensities'
+            f'\n{self.propensities}'
+            f'\n{self.propensities.shape}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+
+    def compute_scores(self) -> None:
+        """
+        Compute scores for prognostic, predictive, and selective features based on the data and the feature weights.
+        """
+        # Each column of the score matrix corresponds to the score for a specific outcome. Rows correspond to samples.
+        prog_lin = self.X @ self.prog_weights.T
+        select_lin = self.X @ self.select_weights.T
+        select_lin_pred = self.X @ self.select_weights_pred.T
+        select_lin_prog = self.X @ self.select_weights_prog.T
+
+        log.debug(
+            f'\nCheck if linear scores are computed correctly for selective features:'
+            f'\n==================================================================='
+            f'\nself.X'
+            f'\n{self.X}'
+            f'\n{self.X.shape}'
+            f'\n\nSelect Weights'
+            f'\n{self.select_weights}'
+            f'\n{self.select_weights.shape}'
+            f'\n\nSelect Lin'
+            f'\n{select_lin}'
+            f'\n{select_lin.shape}'
+            f'\n===================================================================\n'
+        )
+
+        # Compute scores for predictive and selective features for each treatment and outcome
+        pred_lin = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))
+
+        # This creates a score for each treatment and outcome for each sample
+        for i in range(self.num_T):
+            pred_lin[:,i,:] = self.X @ self.pred_weights[i].T
+        
+        # Introduce non-linearity and get final scores
+        prog_scores = (1 - self.nonlinearity_scale) * prog_lin + self.nonlinearity_scale * self.nonlinearities[0](prog_lin)
+        select_scores = (1 - self.nonlinearity_scale) * select_lin + self.nonlinearity_scale * self.nonlinearities[1](select_lin)
+        select_scores_pred_overlap = (1 - self.nonlinearity_scale) * select_lin_pred + self.nonlinearity_scale * self.nonlinearities[1](select_lin_pred)
+        select_scores_prog_overlap = (1 - self.nonlinearity_scale) * select_lin_prog + self.nonlinearity_scale * self.nonlinearities[1](select_lin_prog)
+
+        pred_scores = np.zeros((self.X.shape[0], self.num_T, self.dim_Y))
+        for i in range(self.dim_Y):
+            pred_scores[:,:,i] = (1 - self.nonlinearity_scale) * pred_lin[:,:,i] + self.nonlinearity_scale * self.nonlinearities[i+2](pred_lin[:,:,i])
+
+        log.debug(
+            f'\nCheck if all scores are computed correctly for predictive features:'
+            f'\n==================================================================='
+            f'\nself.X'
+            f'\n{self.X}'
+            f'\n{self.X.shape}'
+            f'\n\nPred Weights'
+            f'\n{self.pred_weights}'
+            f'\n{self.pred_weights.shape}'
+            f'\n\nPred Lin'
+            f'\n{pred_lin}'
+            f'\n{pred_lin.shape}'
+            f'\n\nPred Scores'
+            f'\n{pred_scores}'
+            f'\n{pred_scores.shape}'
+            f'\n===================================================================\n'
+        )
+
+        self.prog_scores = prog_scores
+        self.select_scores = select_scores
+        self.select_scores_pred_overlap = select_scores_pred_overlap
+        self.select_scores_prog_overlap = select_scores_prog_overlap
+
+        self.pred_scores = pred_scores
+
+        return None
+    
+    @property
+    def weights(self) -> Tuple:
+        """
+        Return weights for prognostic, predictive, and selective features.
+        """
+        return self.prog_weights, self.pred_weights, self.select_weights
+    
+    def sample_uniform_weights(self) -> None:
+        """
+        sample uniform weights for the features.
+        """
+        if self.propensity_type.startswith("toy"):
+            self.prog_weights = np.zeros((self.dim_Y, self.dim_X))
+            self.pred_weights = np.zeros((self.num_T, self.dim_Y, self.dim_X))
+            self.select_weights = np.zeros((self.num_T, self.dim_X))
+            self.select_weights_pred = np.zeros((self.num_T, self.dim_X))
+            self.select_weights_prog = np.zeros((self.num_T, self.dim_X))
+            return None
+
+
+        # Sample weights for prognostic features, a weight for every outcome
+        prog_weights = np.random.uniform(-1, 1, size=(self.dim_Y, self.dim_X)) * self.prog_mask
+
+        # Sample weights for predictive and selective features, a weight for every dimension for every treatment and outcome
+        pred_weights = np.random.uniform(-1, 1, size=(self.num_T, self.dim_Y, self.dim_X))
+        select_weights = np.random.uniform(-1, 1, size=(self.num_T, self.dim_X))
+        select_weights_pred = select_weights.copy()
+        select_weights_prog = select_weights.copy()
+
+        # # Sample weights for prognostic features, a weight for every outcome
+        # prog_weights = np.random.uniform(0, 1, size=(self.dim_Y, self.dim_X)) * self.prog_mask
+
+        # # Sample weights for predictive and selective features, a weight for every dimension for every treatment and outcome
+        # pred_weights = np.random.uniform(0, 1, size=(self.num_T, self.dim_Y, self.dim_X))
+        # select_weights = np.random.uniform(0, 1, size=(self.num_T, self.dim_X))
+
+        # # Make sure treatments are different
+        # pred_weights[0] = -pred_weights[0]
+        # select_weights[0] = -select_weights[0]
+
+        # # Ones as weights
+        # prog_weights = np.ones((self.dim_Y, self.dim_X)) * self.prog_mask#/ self.prog_mask.sum()
+        # pred_weights = np.ones((self.num_T, self.dim_Y, self.dim_X))  #/ self.pred_masks.sum(axis=1, keepdims=True)
+        # select_weights = np.ones((self.num_T, self.dim_X))  #/ self.select_masks.sum(axis=1, keepdims=True)
+
+        # Mask weights for features that are not important
+        for i in range(self.num_T):
+            pred_weights[i] = pred_weights[i] * self.pred_masks[:,i]
+            select_weights[i] = select_weights[i] * self.select_masks[:,i]
+            select_weights_pred[i] = select_weights_pred[i] * self.select_masks_pred[:,i]
+            select_weights_prog[i] = select_weights_prog[i] * self.select_masks_prog[:,i]
+
+        # for i in range(self.num_T):
+        #     row_sums = pred_weights[i].sum(axis=1, keepdims=True)
+        #     pred_weights[i] = pred_weights[i] / row_sums
+
+        #     row_sums = select_weights[i].sum()
+        #     select_weights[i] = select_weights[i] / row_sums
+
+        # # Make sure that prog weights sum to one per outcome
+        # row_sums = prog_weights.sum(axis=1, keepdims=True)
+        # prog_weights = prog_weights / row_sums
+
+        log.debug(
+            f'\nCheck if masks are applied correctly:'
+            f'\n==================================================================='
+            f'\nSelect Weights'
+            f'\n{select_weights}'
+            f'\n{select_weights.shape}'
+            f'\n\nSelect Masks'
+            f'\n{self.select_masks}'
+            f'\n{self.select_masks.shape}'
+            f'\n\nPred Weights'
+            f'\n{pred_weights}'
+            f'\n{pred_weights.shape}'
+            f'\n\nPred Masks'
+            f'\n{self.pred_masks}'
+            f'\n{self.pred_masks.shape}'
+            f'\n===================================================================\n'
+        )
+        
+        self.prog_weights = prog_weights
+        self.pred_weights = pred_weights
+        self.select_weights = select_weights
+        self.select_weights_pred = select_weights_pred
+        self.select_weights_prog = select_weights_prog
+
+        return None
+    
+    @property
+    def all_important_features(self) -> np.ndarray:
+        """
+        Return all important feature indices.
+        """
+        all_important_features = np.union1d(self.predictive_features, self.prognostic_features)
+        all_important_features = np.union1d(all_important_features, self.selective_features)
+
+        log.debug(
+            f'\nCheck if all important features are computed correctly:'
+            f'\n==================================================================='
+            f'\nProg Features'
+            f'\n{self.prognostic_features}'
+            f'\n\nPred Features'
+            f'\n{self.predictive_features}'
+            f'\n\nSelect Features'
+            f'\n{self.selective_features}'
+            f'\n\nAll Important Features'
+            f'\n{all_important_features}'
+            f'\n===================================================================\n'
+        )
+
+        return all_important_features
+
+    @property
+    def prognostic_features(self) -> np.ndarray:
+        """
+        Return prognostic feature indices.
+        """
+        prog_features = np.where((self.prog_mask).astype(np.int32) != 0)
+        return prog_features
+    
+    @property
+    def predictive_features(self) -> np.ndarray:
+        """
+        Return predictive feature indices.
+        """
+        pred_features = np.where((self.pred_masks.sum(axis=1)).astype(np.int32) != 0)
+        return pred_features
+
+    @property
+    def selective_features(self) -> np.ndarray:
+        """
+        Return selective feature indices.
+        """
+        select_features = np.where((self.select_masks.sum(axis=1)).astype(np.int32) != 0)
+        return select_features
+    
+    def sample_important_feature_masks(self) -> None:
+        """
+        Pick features that are important for treatment selection, outcome prediction, and prognostic prediction based on the configuration.
+        """
+        if self.propensity_type.startswith("toy"):
+            self.prog_mask = np.zeros(shape=(self.dim_X))
+            self.pred_masks = np.zeros(shape=(self.dim_X, self.num_T))
+            self.select_masks = np.zeros(shape=(self.dim_X, self.num_T))
+
+            self.prog_mask[0] = 1
+            self.pred_masks[0,0] = 1
+            self.pred_masks[1,1] = 1
+            self.select_masks[0,0] = 1
+            self.select_masks[1,1] = 1
+
+            return None
+
+        # Get indices for features and shuffle if random_feature_selection is True
+        all_indices = np.arange(self.dim_X)
+        n = self.num_pred_features
+
+        if self.random_feature_selection:
+            np.random.shuffle(all_indices)
+
+        # Initialize masks
+        prog_mask = np.zeros(shape=(self.dim_X))
+        pred_masks = np.zeros(shape=(self.dim_X, self.num_T))
+        select_masks = np.zeros(shape=(self.dim_X, self.num_T))
+
+        # Handle case with feature overlap
+        if self.feature_type_overlap == "sel_pred":
+        
+            prog_indices = all_indices[:n]
+            prog_mask[prog_indices] = 1
+
+            if self.treatment_feature_overlap:
+                assert 2*n <= int(self.dim_X)
+                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])
+                select_indices = np.array(self.num_T * [all_indices[n:2*n]])
+
+                prog_mask[prog_indices] = 1
+                pred_masks[pred_indices] = 1
+                select_masks[select_indices] = 1
+
+            else:
+                assert n*(1+self.num_T) <= int(self.dim_X)
+                for i in range(self.num_T):
+                    pred_indices = all_indices[(i+1)*n: (i+2)*n]
+                    select_indices = all_indices[(i+1)*n: (i+2)*n]
+            
+                    pred_masks[pred_indices,i] = 1
+                    select_masks[select_indices,i] = 1
+
+        elif self.feature_type_overlap == "sel_prog":
+
+            if self.treatment_feature_overlap:
+                assert 2*n <= int(self.dim_X)
+                prog_indices = all_indices[:n]
+                prog_mask[prog_indices] = 1
+                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])
+                select_indices = np.array(self.num_T * [all_indices[:n]])
+
+                prog_mask[prog_indices] = 1
+                pred_masks[pred_indices] = 1
+                select_masks[select_indices] = 1
+
+            else:
+                assert 2*n*self.num_T <= int(self.dim_X)
+                prog_indices = all_indices[:n*self.num_T:self.num_T]
+                prog_mask[prog_indices] = 1
+                for i in range(self.num_T):
+                    select_indices = all_indices[i*n: (i+1)*n]
+                    pred_indices = all_indices[(i+self.num_T+1)*n: (i+self.num_T+2)*n]
+            
+                    pred_masks[pred_indices,i] = 1
+                    select_masks[select_indices,i] = 1
+
+        elif self.feature_type_overlap == "sel_none":
+            prog_indices = all_indices[:n]
+            prog_mask[prog_indices] = 1
+
+            if self.treatment_feature_overlap:
+                assert 3*n <= int(self.dim_X)
+                pred_indices = np.array(self.num_T * [all_indices[n:2*n]])
+                select_indices = np.array(self.num_T * [all_indices[2*n:3*n]])
+
+                prog_mask[prog_indices] = 1
+                pred_masks[pred_indices] = 1
+                select_masks[select_indices] = 1
+
+            else:
+                #assert n+2*n*self.num_T <= int(self.dim_X)
+                for i in range(1,self.num_T+1):
+                    select_indices = all_indices[i*n: (i+1)*n]
+                    pred_indices = all_indices[(i+self.num_T)*n: (i+self.num_T+1)*n]
+                    pred_masks[pred_indices,i-1] = 1
+                    select_masks[select_indices,i-1] = 1
+
+        # # Handle case with feature overlap
+        # if self.feature_overlap:
+        #     assert max(self.num_pred_features, self.num_prog_features, self.num_select_features) <= int(self.dim_X)
+        
+        #     prog_indices = all_indices[:self.num_prog_features]
+        #     pred_indices = np.array(self.num_T * [all_indices[:self.num_pred_features]])
+        #     select_indices = np.array(self.num_T * [all_indices[:self.num_select_features]])
+
+        #     prog_mask[prog_indices] = 1
+        #     pred_masks[pred_indices] = 1
+        #     select_masks[select_indices] = 1
+
+        # # Handle case without feature overlap
+        # else:
+        #     assert (self.num_prog_features + self.num_T * (self.num_pred_features + self.num_select_features)) <= int(self.dim_X)
+
+        #     prog_indices = all_indices[:self.num_prog_features]
+        #     prog_mask[prog_indices] = 1
+        #     pred_indices = all_indices[self.num_prog_features : (self.num_prog_features + self.num_T*self.num_pred_features)]
+        #     select_indices = all_indices[(self.num_prog_features + self.num_T*self.num_pred_features):(self.num_prog_features + self.num_T*(self.num_pred_features+self.num_select_features))]
+            
+        #     # Mask features for every treatment
+        #     for i in range(self.num_T):
+        #         pred_masks[pred_indices[i*self.num_pred_features:(i+1)*self.num_pred_features],i] = 1
+        #         select_masks[select_indices[i*self.num_select_features:(i+1)*self.num_select_features],i] = 1
+
+        self.prog_mask = prog_mask
+        self.pred_masks = pred_masks
+        self.select_masks = select_masks
+        self.select_masks_pred = pred_masks.copy()
+        self.select_masks_prog = select_masks.copy()
+
+        log.debug(
+            f'\nCheck if important features are sampled correctly:'
+            f'\n==================================================================='
+            f'\nProg Indices'
+            f'\n{prog_indices}'
+            f'\n\nPred Indices'
+            f'\n{pred_indices}'
+            f'\n\nSelect Indices'
+            f'\n{select_indices}'
+            f'\n\nProg Mask'
+            f'\n{prog_mask}'
+            f'\n\nPred Masks'
+            f'\n{pred_masks}'
+            f'\n\nSelect Masks'
+            f'\n{select_masks}'
+            f'\n===================================================================\n'
+        )
+        return None
+
+    def sample_nonlinearities(self, num_nonlinearities: int):
+        """
+        Sample non-linearities for each outcome.
+        """
+        if self.nonlinearity_selection_type == "random":
+            # pick num_nonlinearities 
+            return random.choices(population=self.nonlinear_fcts, k=num_nonlinearities)
+        
+        else:
+            raise ValueError(f"Unknown nonlinearity selection type {self.selection_type}.")
+        
+
+class TSimulator(SimulatorBase):
+    """
+    Data generation process class for simulating treatment selection only, when counterfactual outcomes are available (as for in-vitro/pharmacoscopy data).
+    """
+    nonlinear_fcts = [
+                lambda x: np.abs(x),
+                lambda x: np.exp(-(x**2) / 2),
+                lambda x: 1 / (1 + x**2),
+                lambda x: np.cos(x),
+                lambda x: np.arctan(x),
+                lambda x: np.tanh(x),
+                lambda x: np.sin(x),
+                lambda x: np.log(1 + x**2),
+                lambda x: np.sqrt(1 + x**2),
+                lambda x: np.cosh(x),
+            ]
+    
+    def __init__(
+        self,
+        # Data dimensionality
+        dim_X: int,
+
+        # Seed
+        seed: int = 42,
+
+        # Simulation type
+        simulation_type: str = "T",
+
+        # Dimensionality of treatments and outcome
+        num_binary_outcome: int = 0,
+        standardize_outcome: bool = False,
+        standardize_per_outcome: bool = False,
+        num_T: int = 3,
+        dim_Y: int = 3,
+
+        # Scale parameters
+        propensity_scale: float = 1,
+        unbalancedness_exp: float = 0,
+        nonlinearity_scale: float = 1,
+        propensity_type: str = "prog_pred",
+        alpha: float = 0.5,
+        enforce_balancedness: bool = False,
+
+        # Important features
+        num_select_features: int = 5,
+        treatment_feature_overlap: bool = False,
+
+        # Feature selection
+        random_feature_selection: bool = True,
+        nonlinearity_selection_type: bool = True,
+
+        
+    ) -> None:
+        # Number of features
+        self.dim_X = dim_X
+
+        # Make sure results are reproducible by setting seed for np, torch, random
+        self.seed = seed
+        enable_reproducible_results(seed=self.seed)
+
+        # Simulation type
+        self.simulation_type = simulation_type
+
+        # Store dimensions
+        self.num_binary_outcome = num_binary_outcome
+        self.standardize_outcome = standardize_outcome
+        self.standardize_per_outcome = standardize_per_outcome
+        self.num_T = num_T
+        self.dim_Y = dim_Y
+
+        # Scale parameters
+        self.propensity_scale = propensity_scale
+        self.unbalancedness_exp = unbalancedness_exp
+        self.nonlinearity_scale = nonlinearity_scale
+        self.propensity_type = propensity_type
+        self.alpha = alpha
+        self.enforce_balancedness = enforce_balancedness
+
+        # Important features
+        self.num_select_features = num_select_features
+        self.treatment_feature_overlap = treatment_feature_overlap
+        self.num_important_features = num_select_features
+
+        # Feature selection
+        self.random_feature_selection = random_feature_selection
+        self.nonlinearity_selection_type = nonlinearity_selection_type
+
+        # Setup variables
+        self.nonlinearities = None
+        self.select_masks = None
+        self.select_weights = None
+
+        # Setup
+        self.setup()
+
+        # Simulation variables
+        self.X = None
+        self.select_scores = None
+        self.propensities, self.outcomes, self.T, self.Y = None, None, None, None
+
+    def get_simulated_data(self, train_ratio: float = 0.8):
+        """
+        Extract results and split into training and test set. Include counterfactual outcomes and propensities.
+        """
+        return self.X, self.T, self.Y, self.outcomes, self.propensities
+        # Split data
+        # train_size = int(train_ratio * self.X.shape[0])
+        # X_train, X_test = self.X[:train_size], self.X[train_size:]
+        # T_train, T_test = self.T[:train_size], self.T[train_size:]
+        # Y_train, Y_test = self.Y[:train_size], self.Y[train_size:]
+        # outcomes_train, outcomes_test = self.outcomes[:train_size,:,:], self.outcomes[train_size:,:,:]
+        # propensities_train, propensities_test = self.propensities[:train_size], self.propensities[train_size:]
+
+        # if train_ratio == 1:
+        #     return self.X, self.T, self.Y, self.outcomes, self.propensities
+        
+        # return X_train, X_test, T_train, T_test, Y_train, Y_test, outcomes_train, outcomes_test, propensities_train, propensities_test
+
+    def simulate(self, X, outcomes=None) -> Tuple:
+        """
+        Simulate treatment and outcome for a dataset based on the configuration.
+        """
+        log.debug(
+            f'Simulating treatment and outcome for a dataset with:'
+            f'\n==================================================================='
+            f'\nDim X: {self.dim_X}'
+            f'\nDim T: {self.num_T}'
+            f'\nDim Y: {self.dim_Y}'
+            f'\nPropensity Scale: {self.propensity_scale}'
+            f'\nUnbalancedness Exponent: {self.unbalancedness_exp}'
+            f'\nNonlinearity Scale: {self.nonlinearity_scale}'
+            f'\nNum Select Features: {self.num_select_features}'
+            f'\nFeature Overlap: {self.treatment_feature_overlap}'
+            f'\nRandom Feature Selection: {self.random_feature_selection}'
+            f'\nNonlinearity Selection Type: {self.nonlinearity_selection_type}'
+            f'\n===================================================================\n'
+        )
+
+        # 1. Store data
+        self.X = X
+
+        # 2. Compute scores for prognostic, predictive, and selective features
+        self.compute_scores()
+
+        # 3. Retrieve factual and counterfactual outcomes based on the data and the predictive and prognostic scores
+        self.outcomes = outcomes
+        assert self.outcomes.shape == (self.X.shape[0], self.num_T, self.dim_Y)
+
+        if self.standardize_outcome:
+            if self.standardize_per_outcome:
+                self.outcomes = zscore(self.outcomes, axis=0) #, axis=None) # add axis=None to make problem easier again
+            else:
+                self.outcomes = zscore(self.outcomes, axis=None) #, axis=None) # add axis=None to make problem easier again
+
+        log.debug(
+            f'\nCheck if outcomes are processed correctly:'
+            f'\n==================================================================='
+            f'\n\nOutcomes'
+            f'\n{self.outcomes}'
+            f'\n{self.outcomes.shape}'
+            f'\n\nMean Outcomes'
+            f'\n{self.outcomes.mean(axis=0)}'
+            f'\n\nVariance Outcomes'
+            f'\n{self.outcomes.var(axis=0)}'
+            f'\n===================================================================\n'
+        )
+
+        # 4. Compute propensities based on the data and the selective scores
+        self.compute_propensities()
+
+        # 5. Sample treatment assignment based on the propensities
+        self.sample_T()
+
+        # 6. Extract the outcome based on the treatment assignment
+        self.extract_Y()
+
+        return None
+    
+    def setup(self) -> None:
+        """
+        Setup the simulator by defining variables which remain the same across simulations with different samples but the same configuration.
+        """
+        # 1. Sample nonlinearities used 
+        num_nonlinearities = 1 # Same non-linearity for all treatment selection mechanisms
+        self.nonlinearities = self.sample_nonlinearities(num_nonlinearities)
+
+        # 2. Set important feature masks - determine which features should be used for treatment selection, outcome prediction
+        self.sample_important_feature_masks()
+
+        # 3. Sample weights for features
+        self.sample_uniform_weights()
+
+    def get_true_cates(self, 
+                       X: np.ndarray, 
+                       T: np.ndarray, 
+                       outcomes: np.ndarray) -> np.ndarray:
+        """
+        Compute true CATEs for each treatment based on the data and the outcomes.
+        Always use the selected treatment as the base treatment.
+        """
+        # Compute CATEs for each treatment
+        cates = np.zeros((X.shape[0], self.num_T, self.dim_Y))
+
+        for i in range(X.shape[0]):
+            for j in range(self.num_T):
+                cates[i,j,:] = outcomes[i,j,:] - outcomes[i,int(T[i]),:]
+
+        log.debug(
+            f'\nCheck if true CATEs are computed correctly:'
+            f'\n==================================================================='
+            f'\nOutcomes: {outcomes.shape}'
+            f'\n{outcomes}'
+            f'\n\nTreatment Assignment: {T.shape}'
+            f'\n{T}'
+            f'\n\nTrue CATEs: {cates.shape}'
+            f'\n{cates}'
+            f'\n===================================================================\n'
+        )
+
+        return cates
+    
+    def extract_Y(self) -> None:
+        """
+        Extract the outcome based on the treatment assignment.
+        """
+        self.Y = self.outcomes[np.arange(self.X.shape[0]), self.T] 
+
+        log.debug(
+            f'\nCheck if outcomes are extracted correctly:'
+            f'\n==================================================================='
+            f'\nOutcomes'
+            f'\n{self.outcomes}'
+            f'\n{self.outcomes.shape}'
+            f'\n\nTreatment Assignment'
+            f'\n{self.T}'
+            f'\n{self.T.shape}'
+            f'\n\nExtracted Outcomes'
+            f'\n{self.Y}'
+            f'\n{self.Y.shape}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+
+    def sample_T(self) -> None:
+        """
+        Sample treatment assignment based on the propensities.
+        """
+        # Sample from the resulting categorical distribution per row
+        self.T = np.array([np.random.choice([tre for tre in range(self.propensities.shape[1])], p=row) for row in self.propensities])
+   
+        log.debug(
+            f'\nCheck if treatment assignment is sampled correctly:'
+            f'\n==================================================================='
+            f'\nPropensities'
+            f'\n{self.propensities}'
+            f'\n{self.propensities.shape}'
+            f'\n\nTreatment Assignment'
+            f'\n{self.T}'
+            f'\n{self.T.shape}'
+            f'\n\nUnique Treatment Counts'
+            f'\n{np.unique(self.T, return_counts=True)}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+    
+    def get_unbalancedness_weights(self, size: int) -> np.ndarray:
+        """
+        Create weights for introducing unbalancedness for class probabilities.
+        """
+        # Sample initial distribution of treatment assignment 
+        unb_weights = np.random.uniform(0, 1, size=size) 
+        unb_weights = unb_weights / unb_weights.sum()
+
+        # Standardize the weights and make sure that a treatment doesn't completely disappear for small unbalancedness exponents
+        min_val = unb_weights.min()
+        range_val = unb_weights.max() - min_val
+        unb_weights = (unb_weights - min_val) / range_val
+        unb_weights = 0.01 + unb_weights * 0.98
+
+        return unb_weights
+    
+    def compute_propensities(self) -> None:
+        """
+        Compute propensities based on the data and the selective scores.
+        """
+        select_scores_none = zscore(self.select_scores, axis=0) # Comment for Predictive Epertise
+
+        select_scores_pred = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_pred_flipped = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_prog = np.zeros((self.X.shape[0], self.num_T))
+        select_scores_tre = np.zeros((self.X.shape[0], self.num_T))
+
+        select_scores_pred[:,0] = self.outcomes[:,0,0] - self.outcomes[:,1,0]
+        select_scores_pred[:,1] = self.outcomes[:,1,0] - self.outcomes[:,0,0]
+
+        select_scores_pred_flipped[:,0] = self.outcomes[:,1,0] - self.outcomes[:,0,0]
+        select_scores_pred_flipped[:,1] = self.outcomes[:,0,0] - self.outcomes[:,1,0]
+
+        select_scores_prog[:,0] = self.outcomes[:,0,0]
+        select_scores_prog[:,1] = -self.outcomes[:,0,0]
+
+        select_scores_tre[:,0] = -self.outcomes[:,1,0]
+        select_scores_tre[:,1] = self.outcomes[:,1,0]
+
+        if self.propensity_type == "prog_tre":
+            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_prog
+
+        # Standardize all scores
+        select_scores_pred = zscore(select_scores_pred, axis=0)
+        select_scores_pred_flipped = zscore(select_scores_pred_flipped, axis=0)
+        select_scores_prog = zscore(select_scores_prog, axis=0)
+        select_scores_tre = zscore(select_scores_tre, axis=0)
+
+        if self.propensity_type == "prog_pred":
+            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_prog
+
+        elif self.propensity_type == "prog_tre":
+            pass
+
+        elif self.propensity_type == "none_prog":
+            scores = self.alpha * select_scores_prog + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_pred":
+            scores = self.alpha * select_scores_pred + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_tre":
+            scores = self.alpha * select_scores_tre + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "none_pred_flipped":
+            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_none
+
+        elif self.propensity_type == "pred_pred_flipped":
+            scores = self.alpha * select_scores_pred_flipped + (1 - self.alpha) * select_scores_pred
+
+        elif self.propensity_type == "rct_none":
+            scores = select_scores_none
+
+        else:
+            raise ValueError(f"Unknown propensity type {self.propensity_type}.")
+
+        if self.enforce_balancedness:
+            scores = zscore(scores, axis=0)
+
+        if self.propensity_type == "rct_none":
+            scores = self.alpha * select_scores_none
+
+        # Introduce unbalancedness and manipulate unbalancedness weights for comparable experiments with different seeds
+        unb_weights = self.get_unbalancedness_weights(size=scores.shape[1])
+
+        # Apply the softmax function to each row to get probabilities
+        p = softmax(self.propensity_scale*scores, axis=1)
+
+        # Scale probabilities to introduce unbalancedness
+        p = p * (1 - unb_weights) ** self.unbalancedness_exp
+
+        # Make sure rows add up to one again
+        row_sums = p.sum(axis=1, keepdims=True)
+        p = p / row_sums
+        self.propensities = p
+
+        log.debug(
+            f'\nCheck if propensities are computed correctly:'
+            f'\n==================================================================='
+            f'\nSelect Scores'
+            f'\n{self.select_scores}'
+            f'\n{self.select_scores.shape}'
+            f'\n\nPropensities'
+            f'\n{self.propensities}'
+            f'\n{self.propensities.shape}'
+            f'\n===================================================================\n'
+        )
+
+        return None
+
+    def compute_scores(self) -> None:
+        """
+        Compute scores for prognostic, predictive, and selective features based on the data and the feature weights.
+        """
+        # Each column of the score matrix corresponds to the score for a specific outcome. Rows correspond to samples.
+        select_lin = self.X @ self.select_weights.T
+
+        log.debug(
+            f'\nCheck if linear scores are computed correctly for selective features:'
+            f'\n==================================================================='
+            f'\nself.X'
+            f'\n{self.X}'
+            f'\n{self.X.shape}'
+            f'\n\nSelect Weights'
+            f'\n{self.select_weights}'
+            f'\n{self.select_weights.shape}'
+            f'\n\nSelect Lin'
+            f'\n{select_lin}'
+            f'\n{select_lin.shape}'
+            f'\n===================================================================\n'
+        )
+
+        # Introduce non-linearity and get final scores
+        select_scores = (1 - self.nonlinearity_scale) * select_lin + self.nonlinearity_scale * self.nonlinearities[0](select_lin)
+        self.select_scores = select_scores
+
+        return None
+    
+    @property
+    def weights(self) -> Tuple:
+        """
+        Return weights for prognostic, predictive, and selective features.
+        """
+        return None, None, self.select_weights
+    
+    def sample_uniform_weights(self) -> None:
+        """
+        sample uniform weights for the features.
+        """
+        # Sample weights for selective features, a weight for every dimension for every treatment and outcome
+        select_weights = np.random.uniform(-1, 1, size=(self.num_T, self.dim_X))
+
+
+        # Mask weights for features that are not important
+        for i in range(self.num_T):
+            select_weights[i] = select_weights[i] * self.select_masks[:,i]
+
+        log.debug(
+            f'\nCheck if masks are applied correctly:'
+            f'\n==================================================================='
+            f'\nSelect Weights'
+            f'\n{select_weights}'
+            f'\n{select_weights.shape}'
+            f'\n\nSelect Masks'
+            f'\n{self.select_masks}'
+            f'\n{self.select_masks.shape}'
+            f'\n===================================================================\n'
+        )
+           
+        self.select_weights = select_weights
+
+        return None
+    @property
+    def all_important_features(self) -> np.ndarray:
+        """
+        Return all important feature indices.
+        """
+        all_important_features = self.selective_features
+        log.debug(
+            f'\nCheck if all important features are computed correctly:'
+            f'\n==================================================================='
+            f'\n\nSelect Features'
+            f'\n{self.selective_features}'
+            f'\n\nAll Important Features'
+            f'\n{all_important_features}'
+            f'\n===================================================================\n'
+        )
+
+        return all_important_features
+    
+    @property
+    def predictive_features(self) -> np.ndarray:
+        """
+        Return predictive feature indices.
+        """
+        return None
+    
+    @property
+    def prognostic_features(self) -> np.ndarray:
+        """
+        Return prognostic feature indices.
+        """
+        return None
+
+    @property
+    def selective_features(self) -> np.ndarray:
+        """
+        Return selective feature indices.
+        """
+        select_features = np.where((self.select_masks.sum(axis=1)).astype(np.int32) != 0)
+        return select_features
+    
+    def sample_important_feature_masks(self) -> None:
+        """
+        Pick features that are important for treatment selection based on the configuration.
+        """
+        # Get indices for features and shuffle if random_feature_selection is True
+        all_indices = np.arange(self.dim_X)
+
+        if self.random_feature_selection:
+            np.random.shuffle(all_indices)
+
+        # Initialize masks
+        select_masks = np.zeros(shape=(self.dim_X, self.num_T))
+
+        # Handle case with feature overlap
+        if self.treatment_feature_overlap:
+            assert self.num_select_features <= int(self.dim_X)
+            select_indices = np.array(self.num_T * [all_indices[:self.num_select_features]])
+            select_masks[select_indices] = 1
+
+        # Handle case without feature overlap
+        else:
+            assert (self.num_T * self.num_select_features) <= int(self.dim_X)
+            select_indices = all_indices[:self.num_select_features*self.num_T]
+            
+            # Mask features for every treatment
+            for i in range(self.num_T):
+                select_masks[select_indices[i*self.num_select_features:(i+1)*self.num_select_features],i] = 1
+
+        self.select_masks = select_masks
+
+        return None
+
+    def sample_nonlinearities(self, num_nonlinearities: int):
+        """
+        Sample non-linearities for each outcome.
+        """
+        if self.nonlinearity_selection_type == "random":
+            # pick num_nonlinearities 
+            return random.choices(population=self.nonlinear_fcts, k=num_nonlinearities)
+        
+        else:
+            raise ValueError(f"Unknown nonlinearity selection type {self.selection_type}.")