Make EAR and ENAR handle non-empty error_masks

error_generation/api/mid_level.py

@@ -27,6 +27,6 @@ def create_errors(table: pd.DataFrame, config: MidLevelConfig) -> tuple[pd.DataF
             old_error_mask = error_mask.copy()
             error_mask = error_mechanism.sample(table, column, error_rate, error_mask)
 
-            series = error_type.apply(table, old_error_mask != error_mask, column)
+            series = error_type.apply(table_dirty, old_error_mask != error_mask, column)
             set_column(table_dirty, column, series)
     return table_dirty, error_mask

error_generation/error_mechanism/_ear.py

@@ -35,10 +35,20 @@ def _sample(self: EAR, data: pd.DataFrame, column: str | int, error_rate: float,
         se_mask = get_column(error_mask, column)
         n_errors = int(se_data.size * error_rate)
 
-        upper_bound = len(se_data) - n_errors
+        se_mask_error_free = se_mask[~se_mask]
+        data_column_error_free = data.loc[se_mask_error_free.index, :]
+
+        if len(se_mask_error_free) < n_errors:
+            msg = f"The error rate of {error_rate} requires {n_errors} error-free cells. "
+            msg += f"However, only {len(se_mask_error_free)} error-free cells are available."
+            raise ValueError(msg)
+
+        # we offset the upper bound of the lower_error_index by a) the existing number of errors in the row, and b) the number of errors to-be generated.
+        upper_bound = len(se_data) - sum(se_mask) - n_errors
         lower_error_index = np.random.default_rng(self.seed).integers(0, upper_bound) if upper_bound > 0 else 0
         error_index_range = range(lower_error_index, lower_error_index + n_errors)
+        selected_rows = data_column_error_free.sort_values(by=condition_to_column).iloc[error_index_range, :]
 
-        se_mask.loc[data.sort_values(by=condition_to_column).index[error_index_range]] = True
+        se_mask.loc[selected_rows.index] = True
 
         return error_mask

error_generation/error_mechanism/_ecar.py

@@ -26,7 +26,7 @@ def _sample(self: ECAR, data: pd.DataFrame, column: str | int, error_rate: float
         n_errors = int(se_mask.size * error_rate)
 
         if len(se_mask_error_free) < n_errors:
-            msg = f"The error rate of {error_rate} requires {len(se_mask_error_free)} error-free cells. "
+            msg = f"The error rate of {error_rate} requires {n_errors} error-free cells. "
             msg += f"However, only {len(se_mask_error_free)} error-free cells are available."
             raise ValueError(msg)
 

error_generation/error_mechanism/_enar.py

@@ -21,17 +21,24 @@ def _sample(self: ENAR, data: pd.DataFrame, column: str | int, error_rate: float
         if self.condition_to_column is not None:
             warnings.warn("'condition_to_column' is set but will be ignored by ENAR.", stacklevel=1)
 
-        # distribute errors equally over all columns
-        n_errors = int(se_data.size * error_rate)
+        n_errors = int(len(se_data) * error_rate)
 
-        if n_errors < len(se_data):  # noqa: SIM108
-            lower_error_index = np.random.default_rng(seed=self.seed).integers(0, len(se_data) - n_errors)
-        else:  # all cells are errors
+        # if mid-level or high-level API call ENAR, the error_mask already contains errors. Below we make sure that we only sample rows that do not
+        # already contain errors.
+        se_data_error_free = se_data[~se_mask]
+
+        if len(se_data_error_free) < n_errors:
+            msg = f"The error rate of {error_rate} requires {n_errors} error-free cells. "
+            msg += f"However, only {len(se_data_error_free)} error-free cells are available."
+            raise ValueError(msg)
+
+        if len(se_data_error_free) != n_errors:  # noqa: SIM108
+            lower_error_index = np.random.default_rng(seed=self.seed).integers(0, len(se_data_error_free) - n_errors)
+        else:
             lower_error_index = 0
         error_index_range = range(lower_error_index, lower_error_index + n_errors)
+        selected_rows = se_data_error_free.sort_values().iloc[error_index_range]
 
-        se_mask.loc[se_data.sort_values().index[error_index_range]] = True
-
-        # TODO(PJ): Remember to run if isinstance(seed, int): seed += 1 in mid-level API
+        se_mask.loc[selected_rows.index] = True
 
         return error_mask

error_generation/utils/utils.py

@@ -63,7 +63,7 @@ class ErrorTypeConfig:
     keyboard_layout: str = "ansi-qwerty"
     error_period: int = 10
 
-    na_value = None
+    na_value: str | None = None
 
     mislabel_weighing: str = "uniform"
     mislabel_weights: dict[Any, float] | None = None