diff --git a/src/bayesvalidrox/surrogate_models/polynomial_chaos.py b/src/bayesvalidrox/surrogate_models/polynomial_chaos.py
index 37d34f31de0bc58e2ee56b6b60d7b95895da9d63..e39bc5bbea9d72948db7ac46c4bf7719b95c0f1a 100644
--- a/src/bayesvalidrox/surrogate_models/polynomial_chaos.py
+++ b/src/bayesvalidrox/surrogate_models/polynomial_chaos.py
@@ -865,6 +865,8 @@ class PCE(MetaModel):
         pce_stds_b = {}
 
         for b_i in range(self.n_bootstrap_itrs):
+            pce_means_b[b_i] = {}
+            pce_stds_b[b_i] = {}
             _coeffs_dicts = self._coeffs_dict[f'b_{b_i + 1}'].items()
 
             for output, coef_dict in _coeffs_dicts:
@@ -924,14 +926,16 @@ class PCE(MetaModel):
 
         return pce_means, pce_stds
 
-    def calculate_sobol(self):
+    def calculate_sobol(self, Y_train = None):
         """
         Calculates Sobol' indices of all possible orders and the Total Sobol' indices 
         based on the coefficients in the PCE.
 
         Parameters
         ----------
-        None.
+        Y_train: dict, optional
+            Trainings outputs. They are needed when used in combination with PCA.
+            The default is None
 
         Returns
         -------
@@ -1041,8 +1045,7 @@ class PCE(MetaModel):
 
                 # Compute the sobol indices according to Ref. 2
                 if self.dim_red_method.lower() == "pca":
-                    # TODO: Update so that it does not use the ExpDesign anymore!
-                    n_c_points = self.engine.ExpDesign.Y[output].shape[1]
+                    n_c_points = Y_train[output].shape[1]
                     pca = self.pca[f"b_{b_i+1}"][output]
                     comp_pca = pca.components_
                     n_comp = comp_pca.shape[0]
@@ -1060,7 +1063,7 @@ class PCE(MetaModel):
 
                             for t_idx in range(n_c_points):
                                 var_yt = np.var(
-                                    self.engine.ExpDesign.Y[output][:, t_idx]
+                                    Y_train[output][:, t_idx]
                                 )
                                 term1, term2 = 0.0, 0.0
                                 if var_yt != 0.0:
@@ -1087,7 +1090,7 @@ class PCE(MetaModel):
                         s_zi = total_sobol_array[par_idx]
 
                         for t_idx in range(n_c_points):
-                            var_yt = np.var(self.engine.ExpDesign.Y[output][:, t_idx])
+                            var_yt = np.var(Y_train[output][:, t_idx])
                             term1, term2 = 0.0, 0.0
                             if var_yt != 0.0:
                                 for i in range(n_comp):
diff --git a/tests/test_PolynomialChaosEmulator.py b/tests/test_PolynomialChaosEmulator.py
index b62bd09dab03270783e53fefe2778287bd2b3f4d..64858588d1382caa8df9002fec2ed4c2d7ef9b9a 100644
--- a/tests/test_PolynomialChaosEmulator.py
+++ b/tests/test_PolynomialChaosEmulator.py
@@ -842,3 +842,34 @@ def test_calculate_moments_pca() -> None:
 
 #%% Test PCE.update_metamodel
 # TODO: taken from engine
+
+
+#%% Test PCE.calculate_sobol
+
+def test_calculate_sobol():
+    """
+    Calculate Sobol' indices of a pce-surrogate
+    """
+    inp = Input()
+    inp.add_marginals()
+    inp.Marginals[0].dist_type = 'normal'
+    inp.Marginals[0].parameters = [0, 1]
+    mm = PCE(inp)
+    mm.fit([[0.2], [0.8]], {'Z': [[0.4, 0.4], [0.5, 0.6]]})
+    sobol, totalsobol = mm.calculate_sobol()
+    # TODO are there theory-related checks that could be applied here?
+
+    
+def test_calculate_sobol_pca():
+    """
+    Calculate Sobol' indices of a pce-surrogate with PCA
+    """
+    inp = Input()
+    inp.add_marginals()
+    inp.Marginals[0].dist_type = 'normal'
+    inp.Marginals[0].parameters = [0, 1]
+    mm = PCE(inp)
+    mm.dim_red_method = 'pca'
+    mm.fit([[0.2], [0.8]], {'Z': [[0.4, 0.4], [0.5, 0.6]]})
+    sobol, totalsobol = mm.calculate_sobol({'Z': np.array([[0.4, 0.4], [0.5, 0.6]])})
+    # TODO are there theory-related checks that could be applied here?