update z and x cols setter

Author: SvenKlaassen

doubleml/data/base_data.py

@@ -463,7 +463,9 @@ def x_cols(self):
     @x_cols.setter
     def x_cols(self, value):
         reset_value = hasattr(self, "_x_cols")
+
         if value is not None:
+            # Basic checks
             if isinstance(value, str):
                 value = [value]
             if not isinstance(value, list):
@@ -476,7 +478,17 @@ def x_cols(self, value):
             if not set(value).issubset(set(self.all_variables)):
                 raise ValueError("Invalid covariates x_cols. At least one covariate is no data column.")
             assert set(value).issubset(set(self.all_variables))
-            self._x_cols = value
+
+            if reset_value:
+                previous_value = self._x_cols
+                self._x_cols = value
+                try:
+                    self._check_disjoint_sets()
+                except ValueError:
+                    self._x_cols = previous_value
+                    raise
+            else:
+                self._x_cols = value
 
         else:
             excluded_cols = {self.y_col} | set(self.d_cols)
@@ -486,8 +498,6 @@ def x_cols(self, value):
             self._x_cols = [col for col in self.data.columns if col not in excluded_cols]
 
         if reset_value:
-            self._check_disjoint_sets()
-            # by default, we initialize to the first treatment variable
             self.set_x_d(self.d_cols[0])
 
     @property
@@ -500,6 +510,8 @@ def d_cols(self):
     @d_cols.setter
     def d_cols(self, value):
         reset_value = hasattr(self, "_d_cols")
+
+        # Basic checks
         if isinstance(value, str):
             value = [value]
         if not isinstance(value, list):
@@ -511,10 +523,19 @@ def d_cols(self, value):
             raise ValueError("Invalid treatment variable(s) d_cols: Contains duplicate values.")
         if not set(value).issubset(set(self.all_variables)):
             raise ValueError("Invalid treatment variable(s) d_cols. At least one treatment variable is no data column.")
-        self._d_cols = value
+
+        if reset_value:
+            previous_value = self._d_cols
+            self._d_cols = value
+            try:
+                self._check_disjoint_sets()
+            except ValueError:
+                self._d_cols = previous_value
+                raise
+        else:
+            self._d_cols = value
+
         if reset_value:
-            self._check_disjoint_sets()
-            # by default, we initialize to the first treatment variable
             self.set_x_d(self.d_cols[0])
 
     @property
@@ -541,9 +562,9 @@ def y_col(self, value):
             self._y_col = value
             try:
                 self._check_disjoint_sets()
-            except ValueError as e:
+            except ValueError:
                 self._y_col = previous_value
-                raise e
+                raise
         else:
             self._y_col = value
 
@@ -560,7 +581,9 @@ def z_cols(self):
     @z_cols.setter
     def z_cols(self, value):
         reset_value = hasattr(self, "_z_cols")
+
         if value is not None:
+            # Basic validation
             if isinstance(value, str):
                 value = [value]
             if not isinstance(value, list):
@@ -574,12 +597,22 @@ def z_cols(self, value):
                 raise ValueError(
                     "Invalid instrumental variable(s) z_cols. At least one instrumental variable is no data column."
                 )
-            self._z_cols = value
+
+            if reset_value:
+                previous_value = self._z_cols
+                self._z_cols = value
+                try:
+                    self._check_disjoint_sets()
+                except ValueError:
+                    self._z_cols = previous_value
+                    raise
+            else:
+                self._z_cols = value
+
         else:
             self._z_cols = None
 
         if reset_value:
-            self._check_disjoint_sets()
             self._set_y_z()
 
     @property