Repository: DoubleML/doubleml-for-py (language: python)

Issue:
[Feature Request]: Support for multiple discrete treatments in `doubleml.DoubleMLIRM` 
### Describe the feature you want to propose or implement

Would be nice to extend the `DoubleMLIRM` class to accommodate multiple discrete-valued treatments $w \in \mathcal{W}$, where the estimation primatives are the marginal counterfactual mean $Y^{(w)}$ under each treatment, and one can construct treatment effects as contrasts between them. The marginal mean construction is the same as in the binary treatment case

$$
\phi_i^{(w)} = \mu^{(w)}(x) + \frac{1(W_i = w)}{\pi^{w}(x)} \left(Y_i - \mu^{(w)}(x)\right)
$$

which can be constructed at the unit level, and the subsequent analysis can be performed as usual (the binary influence function would be a special case where one constructs $\psi_i = \phi_i^{(1)} - \phi_i^{(0)} - \theta$ and takes the mean and variance for point estimate and standard error. 

---
Nuances:
 - ATT isn't a well defined concept for multiple-treatments (or at the very least one has to stipulate the subpopulation that is allowed to contribute to the estimator of $E[Y^({0}) \mid W = 1]$, typically the 'pure control' (0-valued by convention) subpopulation might be reasonable)


### Propose a possible solution or implementation


This approach is implemented in [Ed Kennedy's package](https://github.com/ehkennedy/npcausal/blob/master/R/ate.R), and I have a minimal implementation in [this](https://github.com/apoorvalal/aipyw) library. Wondering if there's interest in accommodating this use case (since this will involve changing `IIRM` from [`check_data`](https://github.com/DoubleML/doubleml-for-py/blob/main/doubleml/irm/irm.py#L246) onwards.


### Did you consider alternatives to the proposed solution. If yes, please describe

_No response_

### Comments, context or references

_No response_


Discussion / hints:
Thanks @apoorvalal  for adding this feature request. We are working on an extension to the multi-treatment IRM model. Also thanks for the hint for the ATET in multi-treatment, makes sense 😄 

We are currently working on various extensions including a change of the DoubleML interface to make it possible to better combine various DoubleML objects... The development is already pretty far but we are facing a couple of issues regarding parallel changes (new model, also docu updates etc). Maybe @SvenKlaassen can add more info to what extent multi-level treatment IRM is already possible with the dev version. Anyways, I think the refactoring will be completed in the next months.

We'll add an update and an example here
Yes, i think this would be a great addition. 
The with the new `DoubleMLFramework` class it should be possible to substract two potential outcome estimates, such that an implementation for average potential outcomes $\mathbb{E}[Y(a)]$ for action/treatment $a\in A$, would be a first step. Then a simple class for multiple levels (like `DoubleMLQTE`, see https://github.com/DoubleML/doubleml-for-py/blob/main/doubleml/irm/qte.py) could handle the comparisons on different levels. 

Fix patch:
diff --git a/doubleml/did/did.py b/doubleml/did/did.py
index 733b76c9..c62c8e32 100644
--- a/doubleml/did/did.py
+++ b/doubleml/did/did.py
@@ -343,7 +343,9 @@ def _sensitivity_element_est(self, preds):
         element_dict = {'sigma2': sigma2,
                         'nu2': nu2,
                         'psi_sigma2': psi_sigma2,
-                        'psi_nu2': psi_nu2}
+                        'psi_nu2': psi_nu2,
+                        'riesz_rep': rr,
+                        }
         return element_dict
 
     def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv,
diff --git a/doubleml/did/did_cs.py b/doubleml/did/did_cs.py
index 766f0ea3..b1cd3d4f 100644
--- a/doubleml/did/did_cs.py
+++ b/doubleml/did/did_cs.py
@@ -476,7 +476,9 @@ def _sensitivity_element_est(self, preds):
         element_dict = {'sigma2': sigma2,
                         'nu2': nu2,
                         'psi_sigma2': psi_sigma2,
-                        'psi_nu2': psi_nu2}
+                        'psi_nu2': psi_nu2,
+                        'riesz_rep': rr,
+                        }
         return element_dict
 
     def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv,
diff --git a/doubleml/double_ml.py b/doubleml/double_ml.py
index 810686f2..c8bb7604 100644
--- a/doubleml/double_ml.py
+++ b/doubleml/double_ml.py
@@ -8,16 +8,13 @@
 from scipy.stats import norm
 
 from abc import ABC, abstractmethod
-from scipy.optimize import minimize_scalar
 
 from .double_ml_data import DoubleMLBaseData, DoubleMLClusterData
 from .double_ml_framework import DoubleMLFramework
 
 from .utils.resampling import DoubleMLResampling, DoubleMLClusterResampling
 from .utils._estimation import _rmse, _aggregate_coefs_and_ses, _var_est, _set_external_predictions
-from .utils._checks import _check_in_zero_one, _check_integer, _check_float, _check_bool, \
-    _check_benchmarks, _check_external_predictions, _check_sample_splitting
-from .utils._plots import _sensitivity_contour_plot
+from .utils._checks import _check_external_predictions, _check_sample_splitting
 from .utils.gain_statistics import gain_statistics
 
 _implemented_data_backends = ['DoubleMLData', 'DoubleMLClusterData']
@@ -335,8 +332,8 @@ def psi_elements(self):
     def sensitivity_elements(self):
         """
         Values of the sensitivity components after calling :meth:`fit`;
-        If available (e.g., PLR, IRM) a dictionary with entries ``sigma2``, ``nu2``, ``psi_sigma2``
-        and ``psi_nu2``.
+        If available (e.g., PLR, IRM) a dictionary with entries ``sigma2``, ``nu2``, ``psi_sigma2``, ``psi_nu2``
+        and ``riesz_rep``.
         """
         return self._sensitivity_elements
 
@@ -347,7 +344,11 @@ def sensitivity_params(self):
         If available (e.g., PLR, IRM) a dictionary with entries ``theta``, ``se``, ``ci``, ``rv``
         and ``rva``.
         """
-        return self._sensitivity_params
+        if self._framework is None:
+            sensitivity_params = None
+        else:
+            sensitivity_params = self._framework.sensitivity_params
+        return sensitivity_params
 
     @property
     def coef(self):
@@ -534,15 +535,38 @@ def construct_framework(self):
         scaled_psi_reshape = np.transpose(scaled_psi, (0, 2, 1))
 
         doubleml_dict = {
-            "thetas": self.coef,
-            "all_thetas": self.all_coef,
-            "ses": self.se,
-            "all_ses": self.all_se,
-            "var_scaling_factors": self._var_scaling_factors,
-            "scaled_psi": scaled_psi_reshape,
-            "is_cluster_data": self._is_cluster_data
+            'thetas': self.coef,
+            'all_thetas': self.all_coef,
+            'ses': self.se,
+            'all_ses': self.all_se,
+            'var_scaling_factors': self._var_scaling_factors,
+            'scaled_psi': scaled_psi_reshape,
+            'is_cluster_data': self._is_cluster_data
         }
 
+        if self._sensitivity_implemented:
+            # reshape sensitivity elements to (n_obs, n_coefs, n_rep)
+            doubleml_dict.update({
+                'sensitivity_elements': {
+                    'sigma2': np.transpose(self.sensitivity_elements['sigma2'], (0, 2, 1)),
+                    'nu2': np.transpose(self.sensitivity_elements['nu2'], (0, 2, 1)),
+                    'psi_sigma2': np.transpose(self.sensitivity_elements['psi_sigma2'], (0, 2, 1)),
+                    'psi_nu2': np.transpose(self.sensitivity_elements['psi_nu2'], (0, 2, 1)),
+                    'riesz_rep': np.transpose(self.sensitivity_elements['riesz_rep'], (0, 2, 1))
+                }
+            })
+
+        if self._is_cluster_data:
+            doubleml_dict.update({
+                'is_cluster_data': True,
+                'cluster_dict': {
+                    'smpls': self._smpls,
+                    'smpls_cluster': self._smpls_cluster,
+                    'cluster_vars': self._dml_data.cluster_vars,
+                    'n_folds_per_cluster': self._n_folds_per_cluster,
+                }
+            })
+
         doubleml_framework = DoubleMLFramework(doubleml_dict)
         return doubleml_framework
 
@@ -1324,14 +1348,15 @@ def _sensitivity_element_est(self, preds):
 
     @property
     def _sensitivity_element_names(self):
-        return ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2']
+        return ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2', 'riesz_rep']
 
     # the dimensions will usually be (n_obs, n_rep, n_coefs) to be equal to the score dimensions psi
     def _initialize_sensitivity_elements(self, score_dim):
         sensitivity_elements = {'sigma2': np.full((1, score_dim[1], score_dim[2]), np.nan),
                                 'nu2': np.full((1, score_dim[1], score_dim[2]), np.nan),
                                 'psi_sigma2': np.full(score_dim, np.nan),
-                                'psi_nu2': np.full(score_dim, np.nan)}
+                                'psi_nu2': np.full(score_dim, np.nan),
+                                'riesz_rep': np.full(score_dim, np.nan)}
         return sensitivity_elements
 
     def _get_sensitivity_elements(self, i_rep, i_treat):
@@ -1350,122 +1375,6 @@ def _set_sensitivity_elements(self, sensitivity_elements, i_rep, i_treat):
             self.sensitivity_elements[key][:, i_rep, i_treat] = sensitivity_elements[key]
         return
 
-    def _calc_sensitivity_analysis(self, cf_y, cf_d, rho, level):
-        if self._sensitivity_elements is None:
-            raise NotImplementedError(f'Sensitivity analysis not yet implemented for {self.__class__.__name__}.')
-
-        # checks
-        _check_in_zero_one(cf_y, 'cf_y', include_one=False)
-        _check_in_zero_one(cf_d, 'cf_d', include_one=False)
-        if not isinstance(rho, float):
-            raise TypeError(f'rho must be of float type. '
-                            f'{str(rho)} of type {str(type(rho))} was passed.')
-        _check_in_zero_one(abs(rho), 'The absolute value of rho')
-        _check_in_zero_one(level, 'The confidence level', include_zero=False, include_one=False)
-
-        # set elements for readability
-        sigma2 = self.sensitivity_elements['sigma2']
-        nu2 = self.sensitivity_elements['nu2']
-        psi_sigma = self.sensitivity_elements['psi_sigma2']
-        psi_nu = self.sensitivity_elements['psi_nu2']
-        psi_scaled = np.divide(self.psi, np.mean(self.psi_deriv, axis=0))
-
-        if (np.any(sigma2 < 0)) | (np.any(nu2 < 0)):
-            raise ValueError('sensitivity_elements sigma2 and nu2 have to be positive. '
-                             f"Got sigma2 {str(sigma2)} and nu2 {str(nu2)}. "
-                             'Most likely this is due to low quality learners (especially propensity scores).')
-
-        # elementwise operations
-        confounding_strength = np.multiply(np.abs(rho), np.sqrt(np.multiply(cf_y, np.divide(cf_d, 1.0-cf_d))))
-        S = np.sqrt(np.multiply(sigma2, nu2))
-
-        # sigma2 and nu2 are of shape (1, n_rep, n_coefs), whereas the all_coefs is of shape (n_coefs, n_reps)
-        all_theta_lower = self.all_coef - np.multiply(np.transpose(np.squeeze(S, axis=0)), confounding_strength)
-        all_theta_upper = self.all_coef + np.multiply(np.transpose(np.squeeze(S, axis=0)), confounding_strength)
-
-        psi_S2 = np.multiply(sigma2, psi_nu) + np.multiply(nu2, psi_sigma)
-        psi_bias = np.multiply(np.divide(confounding_strength, np.multiply(2.0, S)), psi_S2)
-        psi_lower = psi_scaled - psi_bias
-        psi_upper = psi_scaled + psi_bias
-
-        # transpose to obtain shape (n_coefs, n_reps); includes scaling with n^{-1/2}
-        all_sigma_lower = np.full_like(all_theta_lower, fill_value=np.nan)
-        all_sigma_upper = np.full_like(all_theta_upper, fill_value=np.nan)
-        for i_rep in range(self.n_rep):
-            self._i_rep = i_rep
-            for i_d in range(self._dml_data.n_treat):
-                self._i_treat = i_d
-
-                if not self._is_cluster_data:
-                    cluster_vars = None
-                    smpls_cluster = None
-                    n_folds_per_cluster = None
-                else:
-                    cluster_vars = self._dml_data.cluster_vars
-                    smpls_cluster = self.__smpls_cluster
-                    n_folds_per_cluster = self._n_folds_per_cluster
-
-                sigma2_lower_hat, _ = _var_est(psi=psi_lower[:, i_rep, i_d],
-                                               psi_deriv=np.ones_like(psi_lower[:, i_rep, i_d]),
-                                               smpls=self.__smpls,
-                                               is_cluster_data=self._is_cluster_data,
-                                               cluster_vars=cluster_vars,
-                                               smpls_cluster=smpls_cluster,
-                                               n_folds_per_cluster=n_folds_per_cluster)
-                sigma2_upper_hat, _ = _var_est(psi=psi_upper[:, i_rep, i_d],
-                                               psi_deriv=np.ones_like(psi_upper[:, i_rep, i_d]),
-                                               smpls=self.__smpls,
-                                               is_cluster_data=self._is_cluster_data,
-                                               cluster_vars=cluster_vars,
-                                               smpls_cluster=smpls_cluster,
-                                               n_folds_per_cluster=n_folds_per_cluster)
-
-                all_sigma_lower[self._i_treat, self._i_rep] = np.sqrt(sigma2_lower_hat)
-                all_sigma_upper[self._i_treat, self._i_rep] = np.sqrt(sigma2_upper_hat)
-
-        # aggregate coefs and ses over n_rep
-        theta_lower, sigma_lower = _aggregate_coefs_and_ses(all_theta_lower, all_sigma_lower, self._var_scaling_factors)
-        theta_upper, sigma_upper = _aggregate_coefs_and_ses(all_theta_upper, all_sigma_upper, self._var_scaling_factors)
-
-        quant = norm.ppf(level)
-        ci_lower = theta_lower - np.multiply(quant, sigma_lower)
-        ci_upper = theta_upper + np.multiply(quant, sigma_upper)
-
-        theta_dict = {'lower': theta_lower,
-                      'upper': theta_upper}
-
-        se_dict = {'lower': sigma_lower,
-                   'upper': sigma_upper}
-
-        ci_dict = {'lower': ci_lower,
-                   'upper': ci_upper}
-
-        res_dict = {'theta': theta_dict,
-                    'se': se_dict,
-                    'ci': ci_dict}
-
-        return res_dict
-
-    def _calc_robustness_value(self, null_hypothesis, level, rho, idx_treatment):
-        _check_float(null_hypothesis, "null_hypothesis")
-        _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat-1)
-
-        # check which side is relvant
-        bound = 'upper' if (null_hypothesis > self.coef[idx_treatment]) else 'lower'
-
-        # minimize the square to find boundary solutions
-        def rv_fct(value, param):
-            res = self._calc_sensitivity_analysis(cf_y=value,
-                                                  cf_d=value,
-                                                  rho=rho,
-                                                  level=level)[param][bound][idx_treatment] - null_hypothesis
-            return np.square(res)
-
-        rv = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('theta', )).x
-        rva = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('ci', )).x
-
-        return rv, rva
-
     def sensitivity_analysis(self, cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95, null_hypothesis=0.0):
         """
         Performs a sensitivity analysis to account for unobserved confounders.
@@ -1502,42 +1411,17 @@ def sensitivity_analysis(self, cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95, null_h
         -------
         self : object
         """
-        # compute sensitivity analysis
-        sensitivity_dict = self._calc_sensitivity_analysis(cf_y=cf_y, cf_d=cf_d, rho=rho, level=level)
-
-        if isinstance(null_hypothesis, float):
-            null_hypothesis_vec = np.full(shape=self._dml_data.n_treat, fill_value=null_hypothesis)
-        elif isinstance(null_hypothesis, np.ndarray):
-            if null_hypothesis.shape == (self._dml_data.n_treat,):
-                null_hypothesis_vec = null_hypothesis
-            else:
-                raise ValueError("null_hypothesis is numpy.ndarray but does not have the required "
-                                 f"shape ({self._dml_data.n_treat},). "
-                                 f'Array of shape {str(null_hypothesis.shape)} was passed.')
-        else:
-            raise TypeError("null_hypothesis has to be of type float or np.ndarry. "
-                            f"{str(null_hypothesis)} of type {str(type(null_hypothesis))} was passed.")
 
-        # compute robustess values with respect to null_hypothesis
-        rv = np.full(shape=self._dml_data.n_treat, fill_value=np.nan)
-        rva = np.full(shape=self._dml_data.n_treat, fill_value=np.nan)
-
-        for i_treat in range(self._dml_data.n_treat):
-            rv[i_treat], rva[i_treat] = self._calc_robustness_value(null_hypothesis=null_hypothesis_vec[i_treat],
-                                                                    level=level, rho=rho, idx_treatment=i_treat)
-
-        sensitivity_dict['rv'] = rv
-        sensitivity_dict['rva'] = rva
-
-        # add all input parameters
-        input_params = {'cf_y': cf_y,
-                        'cf_d': cf_d,
-                        'rho': rho,
-                        'level': level,
-                        'null_hypothesis': null_hypothesis_vec}
-        sensitivity_dict['input'] = input_params
+        if self._framework is None:
+            raise ValueError('Apply fit() before sensitivity_analysis().')
+        self._framework.sensitivity_analysis(
+            cf_y=cf_y,
+            cf_d=cf_d,
+            rho=rho,
+            level=level,
+            null_hypothesis=null_hypothesis
+        )
 
-        self._sensitivity_params = sensitivity_dict
         return self
 
     @property
@@ -1560,19 +1444,19 @@ def sensitivity_summary(self):
                               f'rho={self.sensitivity_params["input"]["rho"]}'
 
             theta_and_ci_col_names = ['CI lower', 'theta lower', ' theta', 'theta upper', 'CI upper']
-            theta_and_ci = np.transpose(np.vstack((self._sensitivity_params['ci']['lower'],
-                                                   self._sensitivity_params['theta']['lower'],
+            theta_and_ci = np.transpose(np.vstack((self.sensitivity_params['ci']['lower'],
+                                                   self.sensitivity_params['theta']['lower'],
                                                    self.coef,
-                                                   self._sensitivity_params['theta']['upper'],
-                                                   self._sensitivity_params['ci']['upper'])))
+                                                   self.sensitivity_params['theta']['upper'],
+                                                   self.sensitivity_params['ci']['upper'])))
             df_theta_and_ci = pd.DataFrame(theta_and_ci,
                                            columns=theta_and_ci_col_names,
                                            index=self._dml_data.d_cols)
             theta_and_ci_summary = str(df_theta_and_ci)
 
             rvs_col_names = ['H_0', 'RV (%)', 'RVa (%)']
-            rvs = np.transpose(np.vstack((self._sensitivity_params['rv'],
-                                          self._sensitivity_params['rva']))) * 100
+            rvs = np.transpose(np.vstack((self.sensitivity_params['rv'],
+                                          self.sensitivity_params['rva']))) * 100
 
             df_rvs = pd.DataFrame(np.column_stack((self.sensitivity_params["input"]["null_hypothesis"], rvs)),
                                   columns=rvs_col_names,
@@ -1589,8 +1473,8 @@ def sensitivity_summary(self):
 
         return res
 
-    def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True, benchmarks=None,
-                         fill=True, grid_bounds=(0.15, 0.15), grid_size=100):
+    def sensitivity_plot(self, idx_treatment=0, value='theta', rho=1.0, level=0.95, null_hypothesis=0.0,
+                         include_scenario=True, benchmarks=None, fill=True, grid_bounds=(0.15, 0.15), grid_size=100):
         """
         Contour plot of the sensivity with respect to latent/confounding variables.
 
@@ -1605,6 +1489,19 @@ def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True
             and ``'ci'`` (refers to the bounds including statistical uncertainty).
             Default is ``'theta'``.
 
+        rho: float
+            The correlation between the differences in short and long representations in the main regression and
+            Riesz representer. Has to be in [-1,1]. The absolute value determines the adversarial strength of the
+            confounding (maximizes at 1.0).
+            Default is ``1.0``.
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        null_hypothesis : float
+            Null hypothesis for the effect. Determines the direction of the contour lines.
+
         include_scenario : bool
             Indicates whether to highlight the scenario from the call of :meth:`sensitivity_analysis`.
             Default is ``True``.
@@ -1630,76 +1527,21 @@ def sensitivity_plot(self, idx_treatment=0, value='theta', include_scenario=True
         fig : object
             Plotly figure of the sensitivity contours.
         """
-        if self.sensitivity_params is None:
-            raise ValueError('Apply sensitivity_analysis() to include senario in sensitivity_plot. '
-                             'The values of rho and the level are used for the scenario.')
-        _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._dml_data.n_treat-1)
-        if not isinstance(value, str):
-            raise TypeError('value must be a string. '
-                            f'{str(value)} of type {type(value)} was passed.')
-        valid_values = ['theta', 'ci']
-        if value not in valid_values:
-            raise ValueError('Invalid value ' + value + '. ' +
-                             'Valid values ' + ' or '.join(valid_values) + '.')
-        _check_bool(include_scenario, 'include_scenario')
-        _check_benchmarks(benchmarks)
-        _check_bool(fill, 'fill')
-        _check_in_zero_one(grid_bounds[0], "grid_bounds", include_zero=False, include_one=False)
-        _check_in_zero_one(grid_bounds[1], "grid_bounds", include_zero=False, include_one=False)
-        _check_integer(grid_size, "grid_size", lower_bound=10)
-
-        null_hypothesis = self.sensitivity_params['input']['null_hypothesis'][idx_treatment]
-        unadjusted_theta = self.coef[idx_treatment]
-        # check which side is relvant
-        bound = 'upper' if (null_hypothesis > unadjusted_theta) else 'lower'
-
-        # create evaluation grid
-        cf_d_vec = np.linspace(0, grid_bounds[0], grid_size)
-        cf_y_vec = np.linspace(0, grid_bounds[1], grid_size)
-
-        # compute contour values
-        contour_values = np.full(shape=(grid_size, grid_size), fill_value=np.nan)
-        for i_cf_d_grid, cf_d_grid in enumerate(cf_d_vec):
-            for i_cf_y_grid, cf_y_grid in enumerate(cf_y_vec):
-                sens_dict = self._calc_sensitivity_analysis(cf_y=cf_y_grid,
-                                                            cf_d=cf_d_grid,
-                                                            rho=self.sensitivity_params['input']['rho'],
-                                                            level=self.sensitivity_params['input']['level'])
-                contour_values[i_cf_d_grid, i_cf_y_grid] = sens_dict[value][bound][idx_treatment]
-
-        # get the correct unadjusted value for confidence bands
-        if value == 'theta':
-            unadjusted_value = unadjusted_theta
-        else:
-            assert value == 'ci'
-            ci = self.confint(level=self.sensitivity_params['input']['level'])
-            if bound == 'upper':
-                unadjusted_value = ci.iloc[idx_treatment, 1]
-            else:
-                unadjusted_value = ci.iloc[idx_treatment, 0]
-
-        # compute the values for the benchmarks
-        benchmark_dict = copy.deepcopy(benchmarks)
-        if benchmarks is not None:
-            n_benchmarks = len(benchmarks['name'])
-            benchmark_values = np.full(shape=(n_benchmarks,), fill_value=np.nan)
-            for benchmark_idx in range(len(benchmarks['name'])):
-                sens_dict_bench = self._calc_sensitivity_analysis(cf_y=benchmarks['cf_y'][benchmark_idx],
-                                                                  cf_d=benchmarks['cf_d'][benchmark_idx],
-                                                                  rho=self.sensitivity_params['input']['rho'],
-                                                                  level=self.sensitivity_params['input']['level'])
-                benchmark_values[benchmark_idx] = sens_dict_bench[value][bound][idx_treatment]
-            benchmark_dict['value'] = benchmark_values
-        fig = _sensitivity_contour_plot(x=cf_d_vec,
-                                        y=cf_y_vec,
-                                        contour_values=contour_values,
-                                        unadjusted_value=unadjusted_value,
-                                        scenario_x=self.sensitivity_params['input']['cf_d'],
-                                        scenario_y=self.sensitivity_params['input']['cf_y'],
-                                        scenario_value=self.sensitivity_params[value][bound][idx_treatment],
-                                        include_scenario=include_scenario,
-                                        benchmarks=benchmark_dict,
-                                        fill=fill)
+        if self._framework is None:
+            raise ValueError('Apply fit() before sensitivity_plot().')
+        fig = self._framework.sensitivity_plot(
+            idx_treatment=idx_treatment,
+            value=value,
+            rho=rho,
+            level=level,
+            null_hypothesis=null_hypothesis,
+            include_scenario=include_scenario,
+            benchmarks=benchmarks,
+            fill=fill,
+            grid_bounds=grid_bounds,
+            grid_size=grid_size
+        )
+
         return fig
 
     def sensitivity_benchmark(self, benchmarking_set, fit_args=None):
diff --git a/doubleml/double_ml_framework.py b/doubleml/double_ml_framework.py
index 9ae9fb59..feedee41 100644
--- a/doubleml/double_ml_framework.py
+++ b/doubleml/double_ml_framework.py
@@ -1,11 +1,15 @@
 import numpy as np
 import pandas as pd
+import copy
 
 from scipy.stats import norm
+from scipy.optimize import minimize_scalar
 from statsmodels.stats.multitest import multipletests
 
-from .utils._estimation import _draw_weights, _aggregate_coefs_and_ses
-from .utils._checks import _check_bootstrap, _check_framework_compatibility
+from .utils._estimation import _draw_weights, _aggregate_coefs_and_ses, _var_est
+from .utils._checks import _check_bootstrap, _check_framework_compatibility, _check_in_zero_one, \
+    _check_float, _check_integer, _check_bool, _check_benchmarks
+from .utils._plots import _sensitivity_contour_plot
 
 
 class DoubleMLFramework():
@@ -27,7 +31,8 @@ def __init__(
         self._is_cluster_data = False
 
         # check input
-        assert isinstance(doubleml_dict, dict), "doubleml_dict must be a dictionary."
+        if not isinstance(doubleml_dict, dict):
+            raise TypeError('doubleml_dict must be a dictionary.')
         expected_keys = ['thetas', 'ses', 'all_thetas', 'all_ses', 'var_scaling_factors', 'scaled_psi']
         if not all(key in doubleml_dict.keys() for key in expected_keys):
             raise ValueError('The dict must contain the following keys: ' + ', '.join(expected_keys))
@@ -44,11 +49,14 @@ def __init__(
         self._var_scaling_factors = doubleml_dict['var_scaling_factors']
         self._scaled_psi = doubleml_dict['scaled_psi']
 
-        if "is_cluster_data" in doubleml_dict.keys():
-            self._is_cluster_data = doubleml_dict['is_cluster_data']
+        # initialize cluster data
+        self._check_and_set_cluster_data(doubleml_dict)
+
+        # initialize sensitivity analysis
+        self._check_and_set_sensitivity_elements(doubleml_dict)
 
         # check if all sizes match
-        _check_framework_shapes(self)
+        self._check_framework_shapes()
         # initialize bootstrap distribution
         self._boot_t_stat = None
         self._boot_method = None
@@ -170,12 +178,30 @@ def boot_t_stat(self):
         """
         return self._boot_t_stat
 
+    @property
+    def sensitivity_elements(self):
+        """
+        Values of the sensitivity components.
+        If available (e.g., PLR, IRM) a dictionary with entries ``sigma2``, ``nu2``, ``psi_sigma2``, ``psi_nu2``
+        and ``riesz_rep``.
+        """
+        return self._sensitivity_elements
+
+    @property
+    def sensitivity_params(self):
+        """
+        Values of the sensitivity parameters after calling :meth:`sesitivity_analysis`;
+        If available (e.g., PLR, IRM) a dictionary with entries ``theta``, ``se``, ``ci``, ``rv``
+        and ``rva``.
+        """
+        return self._sensitivity_params
+
     def __add__(self, other):
 
         if isinstance(other, DoubleMLFramework):
             # internal consistency check
-            _check_framework_shapes(self)
-            _check_framework_shapes(other)
+            self._check_framework_shapes()
+            other._check_framework_shapes()
             _check_framework_compatibility(self, other, check_treatments=True)
 
             all_thetas = self._all_thetas + other._all_thetas
@@ -192,7 +218,6 @@ def __add__(self, other):
             all_ses = np.sqrt(sigma2_hat)
             thetas, ses = _aggregate_coefs_and_ses(all_thetas, all_ses, var_scaling_factors)
 
-            is_cluster_data = self._is_cluster_data or other._is_cluster_data
             doubleml_dict = {
                 'thetas': thetas,
                 'ses': ses,
@@ -200,10 +225,28 @@ def __add__(self, other):
                 'all_ses': all_ses,
                 'var_scaling_factors': var_scaling_factors,
                 'scaled_psi': scaled_psi,
-                'is_cluster_data': is_cluster_data,
+                'is_cluster_data': self._is_cluster_data,
+                'cluster_dict': self._cluster_dict,
             }
-            new_obj = DoubleMLFramework(doubleml_dict)
 
+            # sensitivity combination only available for same outcome and cond. expectation (e.g. IRM)
+            if self._sensitivity_implemented and other._sensitivity_implemented:
+                nu2_score_element = self._sensitivity_elements['psi_nu2'] + other._sensitivity_elements['psi_nu2'] - \
+                     np.multiply(2.0, np.multiply(self._sensitivity_elements['riesz_rep'],
+                                                  self._sensitivity_elements['riesz_rep']))
+                nu2 = np.mean(nu2_score_element, axis=0, keepdims=True)
+                psi_nu2 = nu2_score_element - nu2
+
+                sensitivity_elements = {
+                    'sigma2': self._sensitivity_elements['sigma2'],
+                    'nu2': nu2,
+                    'psi_sigma2': self._sensitivity_elements['psi_sigma2'],
+                    'psi_nu2': psi_nu2,
+                    'riesz_rep': self._sensitivity_elements['riesz_rep'] + other._sensitivity_elements['riesz_rep'],
+                }
+                doubleml_dict['sensitivity_elements'] = sensitivity_elements
+
+            new_obj = DoubleMLFramework(doubleml_dict)
         else:
             raise TypeError(f"Unsupported operand type: {type(other)}")
 
@@ -216,8 +259,8 @@ def __sub__(self, other):
 
         if isinstance(other, DoubleMLFramework):
             # internal consistency check
-            _check_framework_shapes(self)
-            _check_framework_shapes(other)
+            self._check_framework_shapes()
+            other._check_framework_shapes()
             _check_framework_compatibility(self, other, check_treatments=True)
 
             all_thetas = self._all_thetas - other._all_thetas
@@ -234,7 +277,6 @@ def __sub__(self, other):
             all_ses = np.sqrt(sigma2_hat)
             thetas, ses = _aggregate_coefs_and_ses(all_thetas, all_ses, var_scaling_factors)
 
-            is_cluster_data = self._is_cluster_data or other._is_cluster_data
             doubleml_dict = {
                 'thetas': thetas,
                 'ses': ses,
@@ -242,10 +284,28 @@ def __sub__(self, other):
                 'all_ses': all_ses,
                 'var_scaling_factors': var_scaling_factors,
                 'scaled_psi': scaled_psi,
-                'is_cluster_data': is_cluster_data,
+                'is_cluster_data': self._is_cluster_data,
+                'cluster_dict': self._cluster_dict,
             }
-            new_obj = DoubleMLFramework(doubleml_dict)
 
+            # sensitivity combination only available for same outcome and cond. expectation (e.g. IRM)
+            if self._sensitivity_implemented and other._sensitivity_implemented:
+                nu2_score_element = self._sensitivity_elements['psi_nu2'] - other._sensitivity_elements['psi_nu2'] + \
+                     np.multiply(2.0, np.multiply(self._sensitivity_elements['riesz_rep'],
+                                                  self._sensitivity_elements['riesz_rep']))
+                nu2 = np.mean(nu2_score_element, axis=0, keepdims=True)
+                psi_nu2 = nu2_score_element - nu2
+
+                sensitivity_elements = {
+                    'sigma2': self._sensitivity_elements['sigma2'],
+                    'nu2': nu2,
+                    'psi_sigma2': self._sensitivity_elements['psi_sigma2'],
+                    'psi_nu2': psi_nu2,
+                    'riesz_rep': self._sensitivity_elements['riesz_rep'] - other._sensitivity_elements['riesz_rep'],
+                }
+                doubleml_dict['sensitivity_elements'] = sensitivity_elements
+
+            new_obj = DoubleMLFramework(doubleml_dict)
         else:
             raise TypeError(f"Unsupported operand type: {type(other)}")
 
@@ -265,7 +325,6 @@ def __mul__(self, other):
             all_ses = np.multiply(other, self._all_ses)
             scaled_psi = np.multiply(other, self._scaled_psi)
 
-            is_cluster_data = self._is_cluster_data
             doubleml_dict = {
                 'thetas': thetas,
                 'ses': ses,
@@ -273,8 +332,25 @@ def __mul__(self, other):
                 'all_ses': all_ses,
                 'var_scaling_factors': var_scaling_factors,
                 'scaled_psi': scaled_psi,
-                'is_cluster_data': is_cluster_data,
+                'is_cluster_data': self._is_cluster_data,
+                'cluster_dict': self._cluster_dict,
             }
+
+            # sensitivity combination only available for linear models
+            if self._sensitivity_implemented:
+                nu2_score_element = np.multiply(np.square(other), self._sensitivity_elements['psi_nu2'])
+                nu2 = np.mean(nu2_score_element, axis=0, keepdims=True)
+                psi_nu2 = nu2_score_element - nu2
+
+                sensitivity_elements = {
+                    'sigma2': self._sensitivity_elements['sigma2'],
+                    'nu2': nu2,
+                    'psi_sigma2': self._sensitivity_elements['psi_sigma2'],
+                    'psi_nu2': psi_nu2,
+                    'riesz_rep': np.multiply(other, self._sensitivity_elements['riesz_rep']),
+                }
+                doubleml_dict['sensitivity_elements'] = sensitivity_elements
+
             new_obj = DoubleMLFramework(doubleml_dict)
         else:
             raise TypeError(f"Unsupported operand type: {type(other)}")
@@ -284,6 +360,206 @@ def __mul__(self, other):
     def __rmul__(self, other):
         return self.__mul__(other)
 
+    def _calc_sensitivity_analysis(self, cf_y, cf_d, rho, level):
+        if not self._sensitivity_implemented:
+            raise NotImplementedError('Sensitivity analysis is not implemented for this model.')
+
+        # input checks
+        _check_in_zero_one(cf_y, 'cf_y', include_one=False)
+        _check_in_zero_one(cf_d, 'cf_d', include_one=False)
+        if not isinstance(rho, float):
+            raise TypeError(f'rho must be of float type. '
+                            f'{str(rho)} of type {str(type(rho))} was passed.')
+        _check_in_zero_one(abs(rho), 'The absolute value of rho')
+        _check_in_zero_one(level, 'The confidence level', include_zero=False, include_one=False)
+
+        # set elements for readability
+        sigma2 = self.sensitivity_elements['sigma2']
+        nu2 = self.sensitivity_elements['nu2']
+        psi_sigma = self.sensitivity_elements['psi_sigma2']
+        psi_nu = self.sensitivity_elements['psi_nu2']
+        psi_scaled = self._scaled_psi
+
+        if (np.any(sigma2 < 0)) | (np.any(nu2 < 0)):
+            raise ValueError('sensitivity_elements sigma2 and nu2 have to be positive. '
+                             f"Got sigma2 {str(sigma2)} and nu2 {str(nu2)}. "
+                             'Most likely this is due to low quality learners (especially propensity scores).')
+
+        # elementwise operations
+        confounding_strength = np.multiply(np.abs(rho), np.sqrt(np.multiply(cf_y, np.divide(cf_d, 1.0-cf_d))))
+        sensitivity_scaling = np.sqrt(np.multiply(sigma2, nu2))
+
+        # sigma2 and nu2 are of shape (1, n_thetas, n_rep), whereas the all_thetas is of shape (n_thetas, n_rep)
+        all_theta_lower = self.all_thetas - np.multiply(np.squeeze(sensitivity_scaling, axis=0), confounding_strength)
+        all_theta_upper = self.all_thetas + np.multiply(np.squeeze(sensitivity_scaling, axis=0), confounding_strength)
+
+        psi_variances = np.multiply(sigma2, psi_nu) + np.multiply(nu2, psi_sigma)
+        psi_bias = np.multiply(np.divide(confounding_strength, np.multiply(2.0, sensitivity_scaling)), psi_variances)
+        psi_lower = psi_scaled - psi_bias
+        psi_upper = psi_scaled + psi_bias
+
+        # shape (n_thetas, n_reps); includes scaling with n^{-1/2}
+        all_sigma_lower = np.full_like(all_theta_lower, fill_value=np.nan)
+        all_sigma_upper = np.full_like(all_theta_upper, fill_value=np.nan)
+
+        for i_rep in range(self.n_rep):
+            for i_theta in range(self.n_thetas):
+
+                if not self._is_cluster_data:
+                    smpls = None
+                    cluster_vars = None
+                    smpls_cluster = None
+                    n_folds_per_cluster = None
+                else:
+                    smpls = self._cluster_dict['smpls'][i_rep]
+                    cluster_vars = self._cluster_dict['cluster_vars']
+                    smpls_cluster = self._cluster_dict['smpls_cluster'][i_rep]
+                    n_folds_per_cluster = self._cluster_dict['n_folds_per_cluster']
+
+                sigma2_lower_hat, _ = _var_est(psi=psi_lower[:, i_theta, i_rep],
+                                               psi_deriv=np.ones_like(psi_lower[:, i_theta, i_rep]),
+                                               smpls=smpls,
+                                               is_cluster_data=self._is_cluster_data,
+                                               cluster_vars=cluster_vars,
+                                               smpls_cluster=smpls_cluster,
+                                               n_folds_per_cluster=n_folds_per_cluster)
+                sigma2_upper_hat, _ = _var_est(psi=psi_upper[:, i_theta, i_rep],
+                                               psi_deriv=np.ones_like(psi_upper[:, i_theta, i_rep]),
+                                               smpls=smpls,
+                                               is_cluster_data=self._is_cluster_data,
+                                               cluster_vars=cluster_vars,
+                                               smpls_cluster=smpls_cluster,
+                                               n_folds_per_cluster=n_folds_per_cluster)
+
+                all_sigma_lower[i_theta, i_rep] = np.sqrt(sigma2_lower_hat)
+                all_sigma_upper[i_theta, i_rep] = np.sqrt(sigma2_upper_hat)
+
+        # aggregate coefs and ses over n_rep
+        theta_lower, sigma_lower = _aggregate_coefs_and_ses(all_theta_lower, all_sigma_lower, self._var_scaling_factors)
+        theta_upper, sigma_upper = _aggregate_coefs_and_ses(all_theta_upper, all_sigma_upper, self._var_scaling_factors)
+
+        # per repetition confidence intervals
+        quant = norm.ppf(level)
+        all_ci_lower = all_theta_lower - np.multiply(quant, all_sigma_lower)
+        all_ci_upper = all_theta_upper + np.multiply(quant, all_sigma_upper)
+
+        ci_lower = np.median(all_ci_lower, axis=1)
+        ci_upper = np.median(all_ci_upper, axis=1)
+
+        theta_dict = {'lower': theta_lower,
+                      'upper': theta_upper}
+
+        se_dict = {'lower': sigma_lower,
+                   'upper': sigma_upper}
+
+        ci_dict = {'lower': ci_lower,
+                   'upper': ci_upper}
+
+        res_dict = {'theta': theta_dict,
+                    'se': se_dict,
+                    'ci': ci_dict}
+
+        return res_dict
+
+    def _calc_robustness_value(self, null_hypothesis, level, rho, idx_treatment):
+        _check_float(null_hypothesis, "null_hypothesis")
+        _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self._n_thetas-1)
+
+        # check which side is relvant
+        bound = 'upper' if (null_hypothesis > self.thetas[idx_treatment]) else 'lower'
+
+        # minimize the square to find boundary solutions
+        def rv_fct(value, param):
+            res = self._calc_sensitivity_analysis(cf_y=value,
+                                                  cf_d=value,
+                                                  rho=rho,
+                                                  level=level)[param][bound][idx_treatment] - null_hypothesis
+            return np.square(res)
+
+        rv = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('theta', )).x
+        rva = minimize_scalar(rv_fct, bounds=(0, 0.9999), method='bounded', args=('ci', )).x
+
+        return rv, rva
+
+    def sensitivity_analysis(self, cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95, null_hypothesis=0.0):
+        """
+        Performs a sensitivity analysis to account for unobserved confounders.
+
+        The evaluated scenario is stored as a dictionary in the property ``sensitivity_params``.
+
+        Parameters
+        ----------
+        cf_y : float
+            Percentage of the residual variation of the outcome explained by latent/confounding variables.
+            Default is ``0.03``.
+
+        cf_d : float
+            Percentage gains in the variation of the Riesz representer generated by latent/confounding variables.
+            Default is ``0.03``.
+
+        rho : float
+            The correlation between the differences in short and long representations in the main regression and
+            Riesz representer. Has to be in [-1,1]. The absolute value determines the adversarial strength of the
+            confounding (maximizes at 1.0).
+            Default is ``1.0``.
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        null_hypothesis : float or numpy.ndarray
+            Null hypothesis for the effect. Determines the robustness values.
+            If it is a single float uses the same null hypothesis for all estimated parameters.
+            Else the array has to be of shape (n_thetas,).
+            Default is ``0.0``.
+
+        Returns
+        -------
+        self : object
+        """
+        # check null_hypothesis
+        if isinstance(null_hypothesis, float):
+            null_hypothesis_vec = np.full(shape=self._n_thetas, fill_value=null_hypothesis)
+        elif isinstance(null_hypothesis, np.ndarray):
+            if null_hypothesis.shape == (self._n_thetas,):
+                null_hypothesis_vec = null_hypothesis
+            else:
+                raise ValueError("null_hypothesis is numpy.ndarray but does not have the required "
+                                 f"shape ({self._n_thetas},). "
+                                 f'Array of shape {str(null_hypothesis.shape)} was passed.')
+        else:
+            raise TypeError("null_hypothesis has to be of type float or np.ndarry. "
+                            f"{str(null_hypothesis)} of type {str(type(null_hypothesis))} was passed.")
+
+        # compute sensitivity analysis
+        sensitivity_dict = self._calc_sensitivity_analysis(cf_y=cf_y, cf_d=cf_d, rho=rho, level=level)
+
+        # compute robustess values with respect to null_hypothesis
+        rv = np.full(shape=self._n_thetas, fill_value=np.nan)
+        rva = np.full(shape=self._n_thetas, fill_value=np.nan)
+
+        for i_theta in range(self._n_thetas):
+            rv[i_theta], rva[i_theta] = self._calc_robustness_value(
+                null_hypothesis=null_hypothesis_vec[i_theta],
+                level=level,
+                rho=rho,
+                idx_treatment=i_theta
+            )
+
+        sensitivity_dict['rv'] = rv
+        sensitivity_dict['rva'] = rva
+
+        # add all input parameters
+        input_params = {'cf_y': cf_y,
+                        'cf_d': cf_d,
+                        'rho': rho,
+                        'level': level,
+                        'null_hypothesis': null_hypothesis_vec}
+        sensitivity_dict['input'] = input_params
+
+        self._sensitivity_params = sensitivity_dict
+        return self
+
     def confint(self, joint=False, level=0.95):
         """
         Confidence intervals for DoubleML models.
@@ -445,6 +721,229 @@ def p_adjust(self, method='romano-wolf'):
 
         return df_p_vals, all_p_vals_corrected
 
+    def sensitivity_plot(self, idx_treatment=0, value='theta', rho=1.0, level=0.95, null_hypothesis=0.0,
+                         include_scenario=True, benchmarks=None, fill=True, grid_bounds=(0.15, 0.15), grid_size=100):
+        """
+        Contour plot of the sensivity with respect to latent/confounding variables.
+
+        Parameters
+        ----------
+        idx_treatment : int
+            Index of the treatment to perform the sensitivity analysis.
+            Default is ``0``.
+
+        value : str
+            Determines which contours to plot. Valid values are ``'theta'`` (refers to the bounds)
+            and ``'ci'`` (refers to the bounds including statistical uncertainty).
+            Default is ``'theta'``.
+
+        rho: float
+            The correlation between the differences in short and long representations in the main regression and
+            Riesz representer. Has to be in [-1,1]. The absolute value determines the adversarial strength of the
+            confounding (maximizes at 1.0).
+            Default is ``1.0``.
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        null_hypothesis : float
+            Null hypothesis for the effect. Determines the direction of the contour lines.
+
+        include_scenario : bool
+            Indicates whether to highlight the scenario from the call of :meth:`sensitivity_analysis`.
+            Default is ``True``.
+
+        benchmarks : dict or None
+            Dictionary of benchmarks to be included in the plot. The keys are ``cf_y``, ``cf_d`` and ``name``.
+            Default is ``None``.
+
+        fill : bool
+            Indicates whether to use a heatmap style or only contour lines.
+            Default is ``True``.
+
+        grid_bounds : tuple
+            Determines the evaluation bounds of the grid for ``cf_d`` and ``cf_y``. Has to contain two floats in [0, 1).
+            Default is ``(0.15, 0.15)``.
+
+        grid_size : int
+            Determines the number of evaluation points of the grid.
+            Default is ``100``.
+
+        Returns
+        -------
+        fig : object
+            Plotly figure of the sensitivity contours.
+        """
+        _check_integer(idx_treatment, "idx_treatment", lower_bound=0, upper_bound=self.n_thetas-1)
+        if not isinstance(value, str):
+            raise TypeError('value must be a string. '
+                            f'{str(value)} of type {type(value)} was passed.')
+        valid_values = ['theta', 'ci']
+        if value not in valid_values:
+            raise ValueError('Invalid value ' + value + '. ' +
+                             'Valid values ' + ' or '.join(valid_values) + '.')
+        _check_float(null_hypothesis, "null_hypothesis")
+        _check_bool(include_scenario, 'include_scenario')
+        if include_scenario and self.sensitivity_params is None:
+            raise ValueError('Apply sensitivity_analysis() to include senario in sensitivity_plot. ')
+        _check_benchmarks(benchmarks)
+        _check_bool(fill, 'fill')
+        _check_in_zero_one(grid_bounds[0], "grid_bounds", include_zero=False, include_one=False)
+        _check_in_zero_one(grid_bounds[1], "grid_bounds", include_zero=False, include_one=False)
+        _check_integer(grid_size, "grid_size", lower_bound=10)
+
+        null_hypothesis = self.sensitivity_params['input']['null_hypothesis'][idx_treatment]
+        unadjusted_theta = self.thetas[idx_treatment]
+        # check which side is relvant
+        bound = 'upper' if (null_hypothesis > unadjusted_theta) else 'lower'
+
+        # create evaluation grid
+        cf_d_vec = np.linspace(0, grid_bounds[0], grid_size)
+        cf_y_vec = np.linspace(0, grid_bounds[1], grid_size)
+
+        # compute contour values
+        contour_values = np.full(shape=(grid_size, grid_size), fill_value=np.nan)
+        for i_cf_d_grid, cf_d_grid in enumerate(cf_d_vec):
+            for i_cf_y_grid, cf_y_grid in enumerate(cf_y_vec):
+
+                sens_dict = self._calc_sensitivity_analysis(
+                    cf_y=cf_y_grid,
+                    cf_d=cf_d_grid,
+                    rho=rho,
+                    level=level,
+                )
+                contour_values[i_cf_d_grid, i_cf_y_grid] = sens_dict[value][bound][idx_treatment]
+
+        # get the correct unadjusted value for confidence bands
+        if value == 'theta':
+            unadjusted_value = unadjusted_theta
+        else:
+            assert value == 'ci'
+            ci = self.confint(level=self.sensitivity_params['input']['level'])
+            if bound == 'upper':
+                unadjusted_value = ci.iloc[idx_treatment, 1]
+            else:
+                unadjusted_value = ci.iloc[idx_treatment, 0]
+
+        # compute the values for the benchmarks
+        benchmark_dict = copy.deepcopy(benchmarks)
+        if benchmarks is not None:
+            n_benchmarks = len(benchmarks['name'])
+            benchmark_values = np.full(shape=(n_benchmarks,), fill_value=np.nan)
+            for benchmark_idx in range(len(benchmarks['name'])):
+                sens_dict_bench = self._calc_sensitivity_analysis(
+                    cf_y=benchmarks['cf_y'][benchmark_idx],
+                    cf_d=benchmarks['cf_d'][benchmark_idx],
+                    rho=self.sensitivity_params['input']['rho'],
+                    level=self.sensitivity_params['input']['level']
+                )
+                benchmark_values[benchmark_idx] = sens_dict_bench[value][bound][idx_treatment]
+            benchmark_dict['value'] = benchmark_values
+        fig = _sensitivity_contour_plot(x=cf_d_vec,
+                                        y=cf_y_vec,
+                                        contour_values=contour_values,
+                                        unadjusted_value=unadjusted_value,
+                                        scenario_x=self.sensitivity_params['input']['cf_d'],
+                                        scenario_y=self.sensitivity_params['input']['cf_y'],
+                                        scenario_value=self.sensitivity_params[value][bound][idx_treatment],
+                                        include_scenario=include_scenario,
+                                        benchmarks=benchmark_dict,
+                                        fill=fill)
+        return fig
+
+    def _check_and_set_cluster_data(self, doubleml_dict):
+        self._cluster_dict = None
+
+        if "is_cluster_data" in doubleml_dict.keys():
+            _check_bool(doubleml_dict['is_cluster_data'], 'is_cluster_data')
+            self._is_cluster_data = doubleml_dict['is_cluster_data']
+
+        if self._is_cluster_data:
+            if not ("cluster_dict" in doubleml_dict.keys()):
+                raise ValueError('If is_cluster_data is True, cluster_dict must be provided.')
+
+            if not isinstance(doubleml_dict['cluster_dict'], dict):
+                raise TypeError('cluster_dict must be a dictionary.')
+
+            expected_keys_cluster = ['smpls', 'smpls_cluster', 'cluster_vars', 'n_folds_per_cluster']
+            if not all(key in doubleml_dict['cluster_dict'].keys() for key in expected_keys_cluster):
+                raise ValueError('The cluster_dict must contain the following keys: ' + ', '.join(expected_keys_cluster)
+                                 + '. Got: ' + ', '.join(doubleml_dict['cluster_dict'].keys()) + '.')
+
+            self._cluster_dict = doubleml_dict['cluster_dict']
+
+        return
+
+    def _check_and_set_sensitivity_elements(self, doubleml_dict):
+        if not ("sensitivity_elements" in doubleml_dict.keys()):
+            sensitivity_implemented = False
+            sensitivity_elements = None
+
+        else:
+            if not isinstance(doubleml_dict['sensitivity_elements'], dict):
+                raise TypeError('sensitivity_elements must be a dictionary.')
+            expected_keys_sensitivity = ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2', 'riesz_rep']
+            if not all(key in doubleml_dict['sensitivity_elements'].keys() for key in expected_keys_sensitivity):
+                raise ValueError('The sensitivity_elements dict must contain the following '
+                                 'keys: ' + ', '.join(expected_keys_sensitivity))
+
+            for key in expected_keys_sensitivity:
+                if not isinstance(doubleml_dict['sensitivity_elements'][key], np.ndarray):
+                    raise TypeError(f'The sensitivity element {key} must be a numpy array.')
+
+            # set sensitivity elements
+            sensitivity_implemented = True
+            sensitivity_elements = {
+                'sigma2': doubleml_dict['sensitivity_elements']['sigma2'],
+                'nu2': doubleml_dict['sensitivity_elements']['nu2'],
+                'psi_sigma2': doubleml_dict['sensitivity_elements']['psi_sigma2'],
+                'psi_nu2': doubleml_dict['sensitivity_elements']['psi_nu2'],
+                'riesz_rep': doubleml_dict['sensitivity_elements']['riesz_rep'],
+            }
+
+        self._sensitivity_implemented = sensitivity_implemented
+        self._sensitivity_elements = sensitivity_elements
+        self._sensitivity_params = None
+
+        return
+
+    def _check_framework_shapes(self):
+        score_dim = (self._n_obs, self._n_thetas, self.n_rep)
+        # check if all sizes match
+        if self._thetas.shape != (self._n_thetas,):
+            raise ValueError(f'The shape of thetas does not match the expected shape ({self._n_thetas},).')
+        if self._ses.shape != (self._n_thetas,):
+            raise ValueError(f'The shape of ses does not match the expected shape ({self._n_thetas},).')
+        if self._all_thetas.shape != (self._n_thetas, self._n_rep):
+            raise ValueError(f'The shape of all_thetas does not match the expected shape ({self._n_thetas}, {self._n_rep}).')
+        if self._all_ses.shape != (self._n_thetas, self._n_rep):
+            raise ValueError(f'The shape of all_ses does not match the expected shape ({self._n_thetas}, {self._n_rep}).')
+        if self._var_scaling_factors.shape != (self._n_thetas,):
+            raise ValueError(f'The shape of var_scaling_factors does not match the expected shape ({self._n_thetas},).')
+        # dimension of scaled_psi is n_obs x n_thetas x n_rep (per default)
+        if self._scaled_psi.shape != score_dim:
+            raise ValueError(('The shape of scaled_psi does not match the expected '
+                              f'shape ({self._n_obs}, {self._n_thetas}, {self._n_rep}).'))
+
+        if self._sensitivity_implemented:
+            if self._sensitivity_elements['sigma2'].shape != (1, self._n_thetas, self.n_rep):
+                raise ValueError('The shape of sigma2 does not match the expected shape '
+                                 f'(1, {self._n_thetas}, {self._n_rep}).')
+            if self._sensitivity_elements['nu2'].shape != (1, self._n_thetas, self.n_rep):
+                raise ValueError(f'The shape of nu2 does not match the expected shape (1, {self._n_thetas}, {self._n_rep}).')
+            if self._sensitivity_elements['psi_sigma2'].shape != score_dim:
+                raise ValueError(('The shape of psi_sigma2 does not match the expected '
+                                 f'shape ({self._n_obs}, {self._n_thetas}, {self._n_rep}).'))
+            if self._sensitivity_elements['psi_nu2'].shape != score_dim:
+                raise ValueError(('The shape of psi_nu2 does not match the expected '
+                                 f'shape ({self._n_obs}, {self._n_thetas}, {self._n_rep}).'))
+            if self._sensitivity_elements['riesz_rep'].shape != score_dim:
+                raise ValueError(('The shape of riesz_rep does not match the expected '
+                                 f'shape ({self._n_obs}, {self._n_thetas}, {self._n_rep}).'))
+
+        return None
+
 
 def concat(objs):
     """
@@ -457,7 +956,7 @@ def concat(objs):
         raise TypeError('All objects must be of type DoubleMLFramework.')
 
     # check on internal consitency of objects
-    _ = [_check_framework_shapes(obj) for obj in objs]
+    _ = [obj._check_framework_shapes() for obj in objs]
     # check if all objects are compatible in n_obs and n_rep
     _ = [_check_framework_compatibility(objs[0], obj, check_treatments=False) for obj in objs[1:]]
 
@@ -469,7 +968,11 @@ def concat(objs):
     thetas = np.concatenate([obj.thetas for obj in objs], axis=0)
     ses = np.concatenate([obj.ses for obj in objs], axis=0)
 
-    is_cluster_data = any(obj._is_cluster_data for obj in objs)
+    if any(obj._is_cluster_data for obj in objs):
+        raise NotImplementedError('concat not yet implemented with clustering.')
+    else:
+        is_cluster_data = False
+
     doubleml_dict = {
         'thetas': thetas,
         'ses': ses,
@@ -479,29 +982,18 @@ def concat(objs):
         'scaled_psi': scaled_psi,
         'is_cluster_data': is_cluster_data,
     }
+
+    if all(obj._sensitivity_implemented for obj in objs):
+        sensitivity_elements = {}
+        for key in ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2', 'riesz_rep']:
+            assert all(key in obj._sensitivity_elements.keys() for obj in objs)
+            sensitivity_elements[key] = np.concatenate([obj._sensitivity_elements[key] for obj in objs], axis=1)
+
+        doubleml_dict['sensitivity_elements'] = sensitivity_elements
+
     new_obj = DoubleMLFramework(doubleml_dict)
 
     # check internal consistency of new object
-    _check_framework_shapes(new_obj)
+    new_obj._check_framework_shapes()
 
     return new_obj
-
-
-def _check_framework_shapes(self):
-    # check if all sizes match
-    if self._thetas.shape != (self._n_thetas,):
-        raise ValueError(f'The shape of thetas does not match the expected shape ({self._n_thetas},).')
-    if self._ses.shape != (self._n_thetas,):
-        raise ValueError(f'The shape of ses does not match the expected shape ({self._n_thetas},).')
-    if self._all_thetas.shape != (self._n_thetas, self._n_rep):
-        raise ValueError(f'The shape of all_thetas does not match the expected shape ({self._n_thetas}, {self._n_rep}).')
-    if self._all_ses.shape != (self._n_thetas, self._n_rep):
-        raise ValueError(f'The shape of all_ses does not match the expected shape ({self._n_thetas}, {self._n_rep}).')
-    if self._var_scaling_factors.shape != (self._n_thetas,):
-        raise ValueError(f'The shape of var_scaling_factors does not match the expected shape ({self._n_thetas},).')
-    # dimension of scaled_psi is n_obs x n_thetas x n_rep (per default)
-    if self._scaled_psi.shape != (self._n_obs, self._n_thetas, self._n_rep):
-        raise ValueError(('The shape of scaled_psi does not match the expected '
-                          f'shape ({self._n_obs}, {self._n_thetas}, {self._n_rep}).'))
-
-    return None
diff --git a/doubleml/irm/irm.py b/doubleml/irm/irm.py
index 5dd2c1b0..82d22d1d 100644
--- a/doubleml/irm/irm.py
+++ b/doubleml/irm/irm.py
@@ -400,7 +400,9 @@ def _sensitivity_element_est(self, preds):
         element_dict = {'sigma2': sigma2,
                         'nu2': nu2,
                         'psi_sigma2': psi_sigma2,
-                        'psi_nu2': psi_nu2}
+                        'psi_nu2': psi_nu2,
+                        'riesz_rep': rr,
+                        }
         return element_dict
 
     def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv,
diff --git a/doubleml/plm/plr.py b/doubleml/plm/plr.py
index ab758c90..3b7d90f5 100644
--- a/doubleml/plm/plr.py
+++ b/doubleml/plm/plr.py
@@ -279,13 +279,17 @@ def _sensitivity_element_est(self, preds):
         sigma2 = np.mean(sigma2_score_element)
         psi_sigma2 = sigma2_score_element - sigma2
 
-        nu2 = np.divide(1.0, np.mean(np.square(d - m_hat)))
-        psi_nu2 = nu2 - np.multiply(np.square(d-m_hat), np.square(nu2))
+        treatment_residual = d - m_hat
+        nu2 = np.divide(1.0, np.mean(np.square(treatment_residual)))
+        psi_nu2 = nu2 - np.multiply(np.square(treatment_residual), np.square(nu2))
+        rr = np.multiply(treatment_residual, nu2)
 
         element_dict = {'sigma2': sigma2,
                         'nu2': nu2,
                         'psi_sigma2': psi_sigma2,
-                        'psi_nu2': psi_nu2}
+                        'psi_nu2': psi_nu2,
+                        'riesz_rep': rr,
+                        }
         return element_dict
 
     def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv,
diff --git a/doubleml/utils/_checks.py b/doubleml/utils/_checks.py
index a5a67110..e54c4041 100644
--- a/doubleml/utils/_checks.py
+++ b/doubleml/utils/_checks.py
@@ -348,6 +348,10 @@ def _check_framework_compatibility(dml_framework_1, dml_framework_2, check_treat
             raise ValueError('The number of parameters theta in DoubleMLFrameworks must be the same. '
                              f'Got {str(dml_framework_1.n_thetas)} and {str(dml_framework_2.n_thetas)}.')
 
+    if dml_framework_1._is_cluster_data != dml_framework_2._is_cluster_data:
+        raise ValueError('The cluster structure in DoubleMLFrameworks must be the same. '
+                         f'Got {str(dml_framework_1._is_cluster_data)} and {str(dml_framework_2._is_cluster_data)}.')
+
 
 def _check_set(x):
     return {x} if x is not None else {}


Test patch:
diff --git a/doubleml/tests/_utils_doubleml_sensitivity_manual.py b/doubleml/tests/_utils_doubleml_sensitivity_manual.py
index a5de4ae6..c20e2e4b 100644
--- a/doubleml/tests/_utils_doubleml_sensitivity_manual.py
+++ b/doubleml/tests/_utils_doubleml_sensitivity_manual.py
@@ -35,8 +35,12 @@ def doubleml_sensitivity_manual(sensitivity_elements, all_coefs, psi, psi_deriv,
     theta_upper, sigma_upper = _aggregate_coefs_and_ses(all_theta_upper, all_sigma_upper, var_scaling_factor)
 
     quant = norm.ppf(level)
-    ci_lower = theta_lower - np.multiply(quant, sigma_lower)
-    ci_upper = theta_upper + np.multiply(quant, sigma_upper)
+
+    all_ci_lower = all_theta_lower - np.multiply(quant, all_sigma_lower)
+    all_ci_upper = all_theta_upper + np.multiply(quant, all_sigma_upper)
+
+    ci_lower = np.median(all_ci_lower, axis=1)
+    ci_upper = np.median(all_ci_upper, axis=1)
 
     theta_dict = {'lower': theta_lower,
                   'upper': theta_upper}
diff --git a/doubleml/tests/conftest.py b/doubleml/tests/conftest.py
index 27666a72..328a1211 100644
--- a/doubleml/tests/conftest.py
+++ b/doubleml/tests/conftest.py
@@ -9,6 +9,8 @@
 from doubleml.datasets import make_plr_turrell2018, make_irm_data, \
     make_pliv_CHS2015
 
+from doubleml import DoubleMLData
+
 
 def _g(x):
     return np.power(np.sin(x), 2)
@@ -22,6 +24,29 @@ def _m2(x):
     return np.power(x, 2)
 
 
+@pytest.fixture(scope='session',
+                params=[(500, 5)])
+def generate_data_simple(request):
+    n_p = request.param
+    np.random.seed(1111)
+    # setting parameters
+    n = n_p[0]
+    p = n_p[1]
+    theta = 1.0
+
+    # generating data
+    D1 = 1.0 * (np.random.uniform(size=n) > 0.5)
+    D2 = 1.0 * (np.random.uniform(size=n) > 0.5)
+    X = np.random.normal(size=(n, p))
+    Y = theta * D1 + np.dot(X, np.ones(p)) + np.random.normal(size=n)
+    df = pd.DataFrame(np.column_stack((X, Y, D1, D2)),
+                      columns=[f'X{i + 1}' for i in np.arange(p)] + ['Y', 'D1', 'D2'])
+    data_d1 = DoubleMLData(df, 'Y', 'D1')
+    data_d2 = DoubleMLData(df, 'Y', 'D2')
+
+    return data_d1, data_d2
+
+
 @pytest.fixture(scope='session',
                 params=[(500, 10),
                         (1000, 20),
diff --git a/doubleml/tests/test_exceptions.py b/doubleml/tests/test_exceptions.py
index 8743b45e..215c0a08 100644
--- a/doubleml/tests/test_exceptions.py
+++ b/doubleml/tests/test_exceptions.py
@@ -1008,10 +1008,11 @@ def test_doubleml_sensitivity_not_yet_implemented():
 
     dml_pliv = DoubleMLPLIV(dml_data_pliv, ml_g, ml_m, ml_r)
     dml_pliv.fit()
-    msg = "Sensitivity analysis not yet implemented for DoubleMLPLIV."
+    msg = 'Sensitivity analysis is not implemented for this model.'
     with pytest.raises(NotImplementedError, match=msg):
         _ = dml_pliv.sensitivity_analysis()
 
+    msg = 'Sensitivity analysis not yet implemented for DoubleMLPLIV.'
     with pytest.raises(NotImplementedError, match=msg):
         _ = dml_pliv.sensitivity_benchmark(benchmarking_set=["X1"])
 
@@ -1025,77 +1026,45 @@ def test_doubleml_sensitivity_inputs():
     msg = "cf_y must be of float type. 1 of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=1)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=1, cf_d=0.03, rho=1.0, level=0.95)
 
     msg = r'cf_y must be in \[0,1\). 1.0 was passed.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=1.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=1.0, cf_d=0.03, rho=1.0, level=0.95)
 
     # test cf_d
     msg = "cf_d must be of float type. 1 of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=1)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=1, rho=1.0, level=0.95)
 
     msg = r'cf_d must be in \[0,1\). 1.0 was passed.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=1.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=1.0, rho=1.0, level=0.95)
 
     # test rho
     msg = "rho must be of float type. 1 of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1, level=0.95)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho=1, null_hypothesis=0.0, level=0.95, idx_treatment=0)
 
     msg = "rho must be of float type. 1 of type <class 'str'> was passed."
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho="1")
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho="1", level=0.95)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho="1", null_hypothesis=0.0, level=0.95, idx_treatment=0)
 
     msg = r'The absolute value of rho must be in \[0,1\]. 1.1 was passed.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.1)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.1, level=0.95)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho=1.1, null_hypothesis=0.0, level=0.95, idx_treatment=0)
 
     # test level
     msg = "The confidence level must be of float type. 1 of type <class 'int'> was passed."
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1)
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho=1.0, level=1, null_hypothesis=0.0, idx_treatment=0)
 
     msg = r'The confidence level must be in \(0,1\). 1.0 was passed.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho=1.0, level=1.0, null_hypothesis=0.0, idx_treatment=0)
 
     msg = r'The confidence level must be in \(0,1\). 0.0 was passed.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=0.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=0.0)
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_robustness_value(rho=1.0, level=0.0, null_hypothesis=0.0, idx_treatment=0)
 
     # test null_hypothesis
     msg = "null_hypothesis has to be of type float or np.ndarry. 1 of type <class 'int'> was passed."
@@ -1104,30 +1073,18 @@ def test_doubleml_sensitivity_inputs():
     msg = r"null_hypothesis is numpy.ndarray but does not have the required shape \(1,\). Array of shape \(2,\) was passed."
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_analysis(null_hypothesis=np.array([1, 2]))
-    msg = "null_hypothesis must be of float type. 1 of type <class 'int'> was passed."
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(null_hypothesis=1, level=0.95, rho=1.0, idx_treatment=0)
-    msg = r"null_hypothesis must be of float type. \[1\] of type <class 'numpy.ndarray'> was passed."
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(null_hypothesis=np.array([1]), level=0.95, rho=1.0, idx_treatment=0)
 
     # test idx_treatment
     dml_irm.sensitivity_analysis()
     msg = "idx_treatment must be an integer. 0.0 of type <class 'float'> was passed."
-    with pytest.raises(TypeError, match=msg):
-        _ = dml_irm._calc_robustness_value(idx_treatment=0.0, null_hypothesis=0.0, level=0.95, rho=1.0)
     with pytest.raises(TypeError, match=msg):
         _ = dml_irm.sensitivity_plot(idx_treatment=0.0)
 
     msg = "idx_treatment must be larger or equal to 0. -1 was passed."
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_robustness_value(idx_treatment=-1, null_hypothesis=0.0, level=0.95, rho=1.0)
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_plot(idx_treatment=-1)
 
     msg = "idx_treatment must be smaller or equal to 0. 1 was passed."
-    with pytest.raises(ValueError, match=msg):
-        _ = dml_irm._calc_robustness_value(idx_treatment=1, null_hypothesis=0.0, level=0.95, rho=1.0)
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_plot(idx_treatment=1)
 
@@ -1142,7 +1099,7 @@ def test_doubleml_sensitivity_inputs():
         _ = dml_irm._set_sensitivity_elements(sensitivity_elements=sensitivity_elements, i_rep=0, i_treat=0)
 
     # test variances
-    sensitivity_elements = dict({'sigma2': 1.0, 'nu2': -2.4, 'psi_sigma2': 1.0, 'psi_nu2': 1.0})
+    sensitivity_elements = dict({'sigma2': 1.0, 'nu2': -2.4, 'psi_sigma2': 1.0, 'psi_nu2': 1.0, 'riesz_rep': 1.0})
     _ = dml_irm._set_sensitivity_elements(sensitivity_elements=sensitivity_elements, i_rep=0, i_treat=0)
     msg = ('sensitivity_elements sigma2 and nu2 have to be positive. '
            r'Got sigma2 \[\[\[1.\]\]\] and nu2 \[\[\[-2.4\]\]\]. '
@@ -1176,8 +1133,7 @@ def test_doubleml_sensitivity_plot_input():
     dml_irm = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression(), trimming_threshold=0.1)
     dml_irm.fit()
 
-    msg = (r'Apply sensitivity_analysis\(\) to include senario in sensitivity_plot. '
-           'The values of rho and the level are used for the scenario.')
+    msg = (r'Apply sensitivity_analysis\(\) to include senario in sensitivity_plot. ')
     with pytest.raises(ValueError, match=msg):
         _ = dml_irm.sensitivity_plot()
 
diff --git a/doubleml/tests/test_framework.py b/doubleml/tests/test_framework.py
index 761a65c5..0a447420 100644
--- a/doubleml/tests/test_framework.py
+++ b/doubleml/tests/test_framework.py
@@ -163,7 +163,7 @@ def dml_framework_from_doubleml_fixture(n_rep):
     ml_g = LinearRegression()
     ml_m = LogisticRegression()
 
-    dml_irm_obj = DoubleMLIRM(dml_data, ml_g, ml_m)
+    dml_irm_obj = DoubleMLIRM(dml_data, ml_g, ml_m, n_rep=n_rep)
     dml_irm_obj.fit()
     dml_framework_obj = dml_irm_obj.construct_framework()
 
@@ -179,7 +179,7 @@ def dml_framework_from_doubleml_fixture(n_rep):
 
     # substract objects
     dml_data_2 = make_irm_data()
-    dml_irm_obj_2 = DoubleMLIRM(dml_data_2, ml_g, ml_m)
+    dml_irm_obj_2 = DoubleMLIRM(dml_data_2, ml_g, ml_m, n_rep=n_rep)
     dml_irm_obj_2.fit()
     dml_framework_obj_2 = dml_irm_obj_2.construct_framework()
 
@@ -218,6 +218,7 @@ def dml_framework_from_doubleml_fixture(n_rep):
         'ci_joint_sub_obj': ci_joint_sub_obj,
         'ci_joint_mul_obj': ci_joint_mul_obj,
         'ci_joint_concat': ci_joint_concat,
+        'n_rep': n_rep,
     }
     return result_dict
 
@@ -257,14 +258,19 @@ def test_dml_framework_from_doubleml_se(dml_framework_from_doubleml_fixture):
         dml_framework_from_doubleml_fixture['dml_framework_obj_add_obj'].all_ses,
         2*dml_framework_from_doubleml_fixture['dml_obj'].all_se
     )
-    scaling = np.array([dml_framework_from_doubleml_fixture['dml_obj']._var_scaling_factors]).reshape(-1, 1)
-    sub_var = np.mean(
-        np.square(dml_framework_from_doubleml_fixture['dml_obj'].psi - dml_framework_from_doubleml_fixture['dml_obj_2'].psi),
-        axis=0)
-    assert np.allclose(
-        dml_framework_from_doubleml_fixture['dml_framework_obj_sub_obj'].all_ses,
-        np.sqrt(sub_var / scaling)
-    )
+
+    if dml_framework_from_doubleml_fixture['n_rep'] == 1:
+        # formula only valid for n_rep = 1
+        scaling = np.array([dml_framework_from_doubleml_fixture['dml_obj']._var_scaling_factors]).reshape(-1, 1)
+        sub_var = np.mean(
+            np.square(dml_framework_from_doubleml_fixture['dml_obj'].psi
+                      - dml_framework_from_doubleml_fixture['dml_obj_2'].psi),
+            axis=0)
+        assert np.allclose(
+            dml_framework_from_doubleml_fixture['dml_framework_obj_sub_obj'].all_ses,
+            np.sqrt(sub_var / scaling)
+        )
+
     assert np.allclose(
         dml_framework_from_doubleml_fixture['dml_framework_obj_mul_obj'].all_ses,
         2*dml_framework_from_doubleml_fixture['dml_obj'].all_se
diff --git a/doubleml/tests/test_framework_exceptions.py b/doubleml/tests/test_framework_exceptions.py
index c061e1a9..7dc8849b 100644
--- a/doubleml/tests/test_framework_exceptions.py
+++ b/doubleml/tests/test_framework_exceptions.py
@@ -1,5 +1,6 @@
 import pytest
 import numpy as np
+import copy
 
 from doubleml.double_ml_framework import DoubleMLFramework, concat
 from ._utils import generate_dml_dict
@@ -12,6 +13,14 @@
 psi_a = np.ones(shape=(n_obs, n_thetas, n_rep))
 psi_b = np.random.normal(size=(n_obs, n_thetas, n_rep))
 doubleml_dict = generate_dml_dict(psi_a, psi_b)
+# add sensitivity elements
+doubleml_dict['sensitivity_elements'] = {
+    'sigma2': np.ones(shape=(1, n_thetas, n_rep)),
+    'nu2': np.ones(shape=(1, n_thetas, n_rep)),
+    'psi_sigma2': np.ones(shape=(n_obs, n_thetas, n_rep)),
+    'psi_nu2': np.ones(shape=(n_obs, n_thetas, n_rep)),
+    'riesz_rep': np.ones(shape=(n_obs, n_thetas, n_rep))
+}
 
 # combine objects and estimate parameters
 dml_framework_obj_1 = DoubleMLFramework(doubleml_dict)
@@ -61,9 +70,78 @@ def test_input_exceptions():
         DoubleMLFramework(test_dict)
 
     msg = "doubleml_dict must be a dictionary."
-    with pytest.raises(AssertionError, match=msg):
+    with pytest.raises(TypeError, match=msg):
         DoubleMLFramework(1.0)
 
+    msg = "sensitivity_elements must be a dictionary."
+    with pytest.raises(TypeError, match=msg):
+        test_dict = doubleml_dict.copy()
+        test_dict['sensitivity_elements'] = 1
+        DoubleMLFramework(test_dict)
+
+    msg = 'The sensitivity_elements dict must contain the following keys: sigma2, nu2, psi_sigma2, psi_nu2'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = doubleml_dict.copy()
+        test_dict['sensitivity_elements'] = {'sensitivities': np.ones(shape=(n_obs, n_thetas, n_rep))}
+        DoubleMLFramework(test_dict)
+
+    msg = r'The shape of sigma2 does not match the expected shape \(1, 2, 5\)\.'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['sensitivity_elements']['sigma2'] = np.ones(shape=(n_obs, n_rep))
+        DoubleMLFramework(test_dict)
+
+    msg = r'The shape of nu2 does not match the expected shape \(1, 2, 5\)\.'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['sensitivity_elements']['nu2'] = np.ones(shape=(n_obs, n_rep))
+        DoubleMLFramework(test_dict)
+
+    msg = r'The shape of psi_sigma2 does not match the expected shape \(10, 2, 5\)\.'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['sensitivity_elements']['psi_sigma2'] = np.ones(shape=(n_obs, n_thetas, n_rep, 3))
+        DoubleMLFramework(test_dict)
+
+    msg = r'The shape of psi_nu2 does not match the expected shape \(10, 2, 5\)\.'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['sensitivity_elements']['psi_nu2'] = np.ones(shape=(n_obs, n_thetas, n_rep, 3))
+        DoubleMLFramework(test_dict)
+
+    msg = r'The shape of riesz_rep does not match the expected shape \(10, 2, 5\)\.'
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['sensitivity_elements']['riesz_rep'] = np.ones(shape=(n_obs, n_thetas, n_rep, 3))
+        DoubleMLFramework(test_dict)
+
+    msg = "is_cluster_data has to be boolean. 1.0 of type <class 'float'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['is_cluster_data'] = 1.0
+        DoubleMLFramework(test_dict)
+
+    msg = "If is_cluster_data is True, cluster_dict must be provided."
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['is_cluster_data'] = True
+        DoubleMLFramework(test_dict)
+
+    msg = "cluster_dict must be a dictionary."
+    with pytest.raises(TypeError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['is_cluster_data'] = True
+        test_dict['cluster_dict'] = 1.0
+        DoubleMLFramework(test_dict)
+
+    msg = ('The cluster_dict must contain the following keys: smpls, smpls_cluster,'
+           ' cluster_vars, n_folds_per_cluster. Got: cluster_ids.')
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['is_cluster_data'] = True
+        test_dict['cluster_dict'] = {'cluster_ids': np.ones(shape=(n_obs, n_rep))}
+        DoubleMLFramework(test_dict)
+
 
 def test_operation_exceptions():
     # addition
@@ -151,6 +229,24 @@ def test_operation_exceptions():
         dml_framework_obj_2 = DoubleMLFramework(doubleml_dict_2)
         _ = concat([dml_framework_obj_1, dml_framework_obj_2])
 
+    msg = 'concat not yet implemented with clustering.'
+    with pytest.raises(NotImplementedError, match=msg):
+        doubleml_dict_cluster = generate_dml_dict(psi_a_2, psi_b_2)
+        doubleml_dict_cluster['is_cluster_data'] = True
+        doubleml_dict_cluster['cluster_dict'] = {
+            'smpls': np.ones(shape=(n_obs, n_rep)),
+            'smpls_cluster': np.ones(shape=(n_obs, n_rep)),
+            'cluster_vars': np.ones(shape=(n_obs, n_rep)),
+            'n_folds_per_cluster': 2
+        }
+        dml_framework_obj_cluster = DoubleMLFramework(doubleml_dict_cluster)
+        _ = concat([dml_framework_obj_cluster, dml_framework_obj_cluster])
+
+    # cluster compatibility
+    msg = "The cluster structure in DoubleMLFrameworks must be the same. Got False and True."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_2 + dml_framework_obj_cluster
+
 
 @pytest.mark.ci
 def test_p_adjust_exceptions():
@@ -161,3 +257,212 @@ def test_p_adjust_exceptions():
     msg = r'Apply bootstrap\(\) before p_adjust\("rw"\)\.'
     with pytest.raises(ValueError, match=msg):
         _ = dml_framework_obj_1.p_adjust(method='rw')
+
+
+@pytest.mark.ci
+def test_sensitivity_exceptions():
+    dml_framework_no_sensitivity = DoubleMLFramework(generate_dml_dict(psi_a, psi_b))
+    msg = 'Sensitivity analysis is not implemented for this model.'
+    with pytest.raises(NotImplementedError, match=msg):
+        _ = dml_framework_no_sensitivity._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.1, rho=1.0, level=0.95)
+
+    # test cf_y
+    msg = "cf_y must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=1)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=1, cf_d=0.03, rho=1.0, level=0.95)
+
+    msg = r'cf_y must be in \[0,1\). 1.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=1.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=1.0, cf_d=0.03, rho=1.0, level=0.95)
+
+    # test cf_d
+    msg = "cf_d must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=1)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=1, rho=1.0, level=0.95)
+
+    msg = r'cf_d must be in \[0,1\). 1.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=1.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=1.0, rho=1.0, level=0.95)
+
+    # test rho
+    msg = "rho must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1, level=0.95)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho=1, null_hypothesis=0.0, level=0.95, idx_treatment=0)
+
+    msg = "rho must be of float type. 1 of type <class 'str'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho="1")
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho="1", level=0.95)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho="1", null_hypothesis=0.0, level=0.95, idx_treatment=0)
+
+    msg = r'The absolute value of rho must be in \[0,1\]. 1.1 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.1)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.1, level=0.95)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho=1.1, null_hypothesis=0.0, level=0.95, idx_treatment=0)
+
+    # test level
+    msg = "The confidence level must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1)
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho=1.0, level=1, null_hypothesis=0.0, idx_treatment=0)
+
+    msg = r'The confidence level must be in \(0,1\). 1.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=1.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho=1.0, level=1.0, null_hypothesis=0.0, idx_treatment=0)
+
+    msg = r'The confidence level must be in \(0,1\). 0.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=0.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=0.0)
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(rho=1.0, level=0.0, null_hypothesis=0.0, idx_treatment=0)
+
+    # test null_hypothesis
+    msg = "null_hypothesis has to be of type float or np.ndarry. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(null_hypothesis=1)
+    msg = r"null_hypothesis is numpy.ndarray but does not have the required shape \(2,\). Array of shape \(3,\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_analysis(null_hypothesis=np.array([1, 2, 3]))
+    msg = "null_hypothesis must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(null_hypothesis=1, level=0.95, rho=1.0, idx_treatment=0)
+    msg = r"null_hypothesis must be of float type. \[1\] of type <class 'numpy.ndarray'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1._calc_robustness_value(null_hypothesis=np.array([1]), level=0.95, rho=1.0, idx_treatment=0)
+
+    sensitivity_dict = generate_dml_dict(psi_a, psi_b)
+    sensitivity_dict['sensitivity_elements'] = {
+        'sigma2': np.ones(shape=(1, n_thetas, n_rep)),
+        'nu2': -1.0 * np.ones(shape=(1, n_thetas, n_rep)),
+        'psi_sigma2': np.ones(shape=(n_obs, n_thetas, n_rep)),
+        'psi_nu2': np.ones(shape=(n_obs, n_thetas, n_rep)),
+        'riesz_rep': np.ones(shape=(n_obs, n_thetas, n_rep))
+    }
+    dml_framework_sensitivity = DoubleMLFramework(sensitivity_dict)
+
+    # test idx_treatment
+    dml_framework_obj_1.sensitivity_analysis()
+    msg = "idx_treatment must be an integer. 0.0 of type <class 'float'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_plot(idx_treatment=0.0)
+
+    msg = "idx_treatment must be larger or equal to 0. -1 was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_plot(idx_treatment=-1)
+
+    msg = "idx_treatment must be smaller or equal to 1. 2 was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_1.sensitivity_plot(idx_treatment=2)
+
+    # test variances
+    msg = (
+        r'sensitivity_elements sigma2 and nu2 have to be positive\. '
+        r'Got sigma2 \[\[\[1\. 1\. 1\. 1\. 1\.\]\n\s+\[1\. 1\. 1\. 1\. 1\.\]\]\] '
+        r'and nu2 \[\[\[-1\. -1\. -1\. -1\. -1\.\]\n\s+\[-1\. -1\. -1\. -1\. -1\.\]\]\]\. '
+        r'Most likely this is due to low quality learners \(especially propensity scores\)\.'
+    )
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_sensitivity._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95)
+    with pytest.raises(ValueError, match=msg):
+        dml_framework_sensitivity.sensitivity_analysis(cf_y=0.1, cf_d=0.15, rho=1.0, level=0.95)
+
+
+@pytest.mark.ci
+def test_framework_sensitivity_plot_input():
+    dml_framework_obj_plot = DoubleMLFramework(doubleml_dict)
+
+    msg = (r'Apply sensitivity_analysis\(\) to include senario in sensitivity_plot. ')
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot()
+
+    dml_framework_obj_plot.sensitivity_analysis()
+    msg = "null_hypothesis must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(null_hypothesis=1)
+
+    msg = "include_scenario has to be boolean. True of type <class 'str'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(include_scenario="True")
+
+    msg = "benchmarks has to be either None or a dictionary. True of type <class 'str'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks="True")
+    msg = r"benchmarks has to be a dictionary with keys cf_y, cf_d and name. Got dict_keys\(\['cf_y', 'cf_d'\]\)."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks={'cf_y': 0.1, 'cf_d': 0.15})
+    msg = r"benchmarks has to be a dictionary with values of same length. Got \[1, 2, 2\]."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks={'cf_y': [0.1], 'cf_d': [0.15, 0.2],
+                                                                'name': ['test', 'test2']})
+    msg = "benchmarks cf_y must be of float type. 2 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks={'cf_y': [0.1, 2], 'cf_d': [0.15, 0.2],
+                                                                'name': ['test', 'test2']})
+    msg = r'benchmarks cf_y must be in \[0,1\). 1.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks={'cf_y': [0.1, 1.0], 'cf_d': [0.15, 0.2],
+                                                                'name': ['test', 'test2']})
+    msg = "benchmarks name must be of string type. 2 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(benchmarks={'cf_y': [0.1, 0.2], 'cf_d': [0.15, 0.2],
+                                                                'name': [2, 2]})
+
+    msg = "value must be a string. 2 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(value=2)
+    msg = "Invalid value test. Valid values theta or ci."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(value='test')
+
+    msg = "fill has to be boolean. True of type <class 'str'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(fill="True")
+
+    msg = "grid_size must be an integer. 0.0 of type <class 'float'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_size=0.0)
+    msg = "grid_size must be larger or equal to 10. 9 was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_size=9)
+
+    msg = "grid_bounds must be of float type. 1 of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(0.15, 1))
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(1, 0.15))
+    msg = r'grid_bounds must be in \(0,1\). 1.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(1.0, 0.15))
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(0.15, 1.0))
+    msg = r'grid_bounds must be in \(0,1\). 0.0 was passed.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(0.0, 0.15))
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_framework_obj_plot.sensitivity_plot(grid_bounds=(0.15, 0.0))
diff --git a/doubleml/tests/test_framework_sensitivity.py b/doubleml/tests/test_framework_sensitivity.py
new file mode 100644
index 00000000..5b1be8bd
--- /dev/null
+++ b/doubleml/tests/test_framework_sensitivity.py
@@ -0,0 +1,83 @@
+import pytest
+
+from doubleml.irm.irm import DoubleMLIRM
+from doubleml.double_ml_framework import concat
+
+from sklearn.linear_model import LinearRegression, LogisticRegression
+
+
+@pytest.fixture(scope='module',
+                params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def dml_framework_sensitivity_fixture(n_rep, generate_data_simple):
+    dml_data, dml_data_2 = generate_data_simple
+
+    ml_g = LinearRegression()
+    ml_m = LogisticRegression()
+
+    dml_irm_obj = DoubleMLIRM(dml_data, ml_g, ml_m, n_rep=n_rep)
+    dml_irm_obj.fit()
+    dml_framework_obj = dml_irm_obj.construct_framework()
+    dml_framework_obj.sensitivity_analysis()
+
+    # add objects
+    dml_framework_obj_add_obj = dml_framework_obj + dml_framework_obj
+
+    # substract objects
+    dml_irm_obj_2 = DoubleMLIRM(dml_data_2, ml_g, ml_m, n_rep=n_rep)
+    dml_irm_obj_2.fit()
+    dml_framework_obj_2 = dml_irm_obj_2.construct_framework()
+    dml_framework_obj_sub_obj = dml_framework_obj - dml_framework_obj_2
+
+    # multiply objects
+    dml_framework_obj_mul_obj = dml_framework_obj * 2
+
+    # concat objects
+    dml_framework_obj_concat = concat([dml_framework_obj, dml_framework_obj])
+
+    result_dict = {
+        'dml_obj': dml_irm_obj,
+        'dml_obj_2': dml_irm_obj_2,
+        'dml_framework_obj': dml_framework_obj,
+        'dml_framework_obj_add_obj': dml_framework_obj_add_obj,
+        'dml_framework_obj_sub_obj': dml_framework_obj_sub_obj,
+        'dml_framework_obj_mul_obj': dml_framework_obj_mul_obj,
+        'dml_framework_obj_concat': dml_framework_obj_concat,
+        'n_rep': n_rep,
+    }
+    return result_dict
+
+
+@pytest.mark.ci
+def test_dml_framework_sensitivity_shapes(dml_framework_sensitivity_fixture):
+    n_rep = dml_framework_sensitivity_fixture['dml_framework_obj'].n_rep
+    n_thetas = dml_framework_sensitivity_fixture['dml_framework_obj'].n_thetas
+    n_obs = dml_framework_sensitivity_fixture['dml_framework_obj'].n_obs
+
+    object_list = ['dml_framework_obj',
+                   'dml_framework_obj_add_obj',
+                   'dml_framework_obj_sub_obj',
+                   'dml_framework_obj_mul_obj']
+    var_keys = ['sigma2', 'nu2']
+    score_keys = ['psi_sigma2', 'psi_nu2', 'riesz_rep']
+
+    for obj in object_list:
+        assert dml_framework_sensitivity_fixture[obj]._sensitivity_implemented
+        for key in var_keys:
+            assert dml_framework_sensitivity_fixture[obj]._sensitivity_elements[key].shape == \
+                (1, n_thetas, n_rep)
+        for key in score_keys:
+            assert dml_framework_sensitivity_fixture[obj]._sensitivity_elements[key].shape == \
+                (n_obs, n_thetas, n_rep)
+
+    # separate test for concat
+    for key in var_keys:
+        assert dml_framework_sensitivity_fixture['dml_framework_obj_concat']._sensitivity_elements[key].shape == \
+            (1, 2, n_rep)
+    for key in score_keys:
+        assert dml_framework_sensitivity_fixture['dml_framework_obj_concat']._sensitivity_elements[key].shape == \
+            (n_obs, 2, n_rep)
diff --git a/doubleml/tests/test_model_defaults.py b/doubleml/tests/test_model_defaults.py
index 4c3a2f61..5da2fbca 100644
--- a/doubleml/tests/test_model_defaults.py
+++ b/doubleml/tests/test_model_defaults.py
@@ -195,7 +195,7 @@ def test_sensitivity_defaults():
                   'null_hypothesis': np.array([0.])}
 
     dml_plr.sensitivity_analysis()
-    assert dml_plr._sensitivity_params['input'] == input_dict
+    assert dml_plr.sensitivity_params['input'] == input_dict
 
 
 @pytest.mark.ci
diff --git a/doubleml/tests/test_return_types.py b/doubleml/tests/test_return_types.py
index 45aae6d9..d76f2d14 100644
--- a/doubleml/tests/test_return_types.py
+++ b/doubleml/tests/test_return_types.py
@@ -397,44 +397,95 @@ def test_nuisance_loss():
 
 @pytest.mark.ci
 def test_sensitivity():
+
+    var_keys = ['sigma2', 'nu2']
+    score_keys = ['psi_sigma2', 'psi_nu2', 'riesz_rep']
     benchmarks = {'cf_y': [0.1, 0.2], 'cf_d': [0.15, 0.2], 'name': ["test1", "test2"]}
+
+    # PLR
+    assert isinstance(plr_obj.sensitivity_elements, dict)
+    for key in var_keys:
+        assert isinstance(plr_obj.sensitivity_elements[key], np.ndarray)
+        assert plr_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
+    for key in score_keys:
+        assert isinstance(plr_obj.sensitivity_elements[key], np.ndarray)
+        assert plr_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
+
     assert isinstance(plr_obj.sensitivity_summary, str)
     plr_obj.sensitivity_analysis()
     assert isinstance(plr_obj.sensitivity_summary, str)
     assert isinstance(plr_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
     assert isinstance(plr_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(plr_obj._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(plr_obj._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0), tuple)
+    assert isinstance(plr_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
+    assert isinstance(
+        plr_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
+        tuple)
     plr_benchmark = plr_obj.sensitivity_benchmark(benchmarking_set=["X1"])
     assert isinstance(plr_benchmark, pd.DataFrame)
 
+    # DID
+    assert isinstance(irm_obj.sensitivity_elements, dict)
+    for key in var_keys:
+        assert isinstance(irm_obj.sensitivity_elements[key], np.ndarray)
+        assert irm_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
+    for key in score_keys:
+        assert isinstance(irm_obj.sensitivity_elements[key], np.ndarray)
+        assert irm_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
+
     assert isinstance(irm_obj.sensitivity_summary, str)
     irm_obj.sensitivity_analysis()
     assert isinstance(irm_obj.sensitivity_summary, str)
     assert isinstance(irm_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
     assert isinstance(irm_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(irm_obj._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(irm_obj._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0), tuple)
+    assert isinstance(irm_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
+    assert isinstance(
+        irm_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
+        tuple
+    )
     irm_benchmark = irm_obj.sensitivity_benchmark(benchmarking_set=["X1"])
     assert isinstance(irm_benchmark, pd.DataFrame)
 
+    # DID
+    assert isinstance(did_obj.sensitivity_elements, dict)
+    for key in var_keys:
+        assert isinstance(did_obj.sensitivity_elements[key], np.ndarray)
+        assert did_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
+    for key in score_keys:
+        assert isinstance(did_obj.sensitivity_elements[key], np.ndarray)
+        assert did_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
+
     assert isinstance(did_obj.sensitivity_summary, str)
     did_obj.sensitivity_analysis()
     assert isinstance(did_obj.sensitivity_summary, str)
     assert isinstance(did_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
     assert isinstance(did_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_obj._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(did_obj._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0), tuple)
+    assert isinstance(did_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
+    assert isinstance(
+        did_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
+        tuple
+    )
     did_benchmark = did_obj.sensitivity_benchmark(benchmarking_set=['Z1'])
     assert isinstance(did_benchmark, pd.DataFrame)
 
+    # DIDCS
+    assert isinstance(did_cs_obj.sensitivity_elements, dict)
+    for key in var_keys:
+        assert isinstance(did_cs_obj.sensitivity_elements[key], np.ndarray)
+        assert did_cs_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
+    for key in score_keys:
+        assert isinstance(did_cs_obj.sensitivity_elements[key], np.ndarray)
+        assert did_cs_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
+
     assert isinstance(did_cs_obj.sensitivity_summary, str)
     did_cs_obj.sensitivity_analysis()
     assert isinstance(did_cs_obj.sensitivity_summary, str)
     assert isinstance(did_cs_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
     assert isinstance(did_cs_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_cs_obj._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(did_cs_obj._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0), tuple)
+    assert isinstance(did_cs_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
+    assert isinstance(
+        did_cs_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
+        tuple
+    )
     did_cs_benchmark = did_cs_obj.sensitivity_benchmark(benchmarking_set=['Z1'])
     assert isinstance(did_cs_benchmark, pd.DataFrame)