SIENTIAPDE-1579: Updated tests and refactored apply filters method

This commit is contained in:
Kou Kinoshita
2026-02-18 17:44:46 -03:00
parent 43cbc31e10
commit 665870320f
3 changed files with 182 additions and 54 deletions

View File

@@ -47,6 +47,41 @@ def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) ->
return data
def _single_variable_support_mask(
data_view: pd.DataFrame,
var_col: str,
target_variable: str,
config: dict,
) -> np.ndarray | None:
"""Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
if var_col not in data_view.columns:
return None
upper = config.get('upper_line') or config.get('upperLine')
lower = config.get('lower_line') or config.get('lowerLine')
if not upper or not lower:
return None
x_vals = data_view[var_col].astype(float).to_numpy()
y_vals = data_view[target_variable].astype(float).to_numpy()
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
scale_ratio = y_range / x_range
b1 = float(upper.get('intercept', 0))
deg1 = float(upper.get('angle', 0))
b2 = float(lower.get('intercept', 0))
deg2 = float(lower.get('angle', 0))
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
y1 = m1 * x_vals + b1
y2 = m2 * x_vals + b2
lower_bound = np.minimum(y1, y2)
upper_bound = np.maximum(y1, y2)
return (y_vals >= lower_bound) & (y_vals <= upper_bound)
def _apply_support_filters(
data_view: pd.DataFrame,
target_variable: str,
@@ -71,39 +106,10 @@ def _apply_support_filters(
return data_view
combined_keep_mask = np.ones(len(data_view), dtype=bool)
n = len(data_view)
for var_col, config in support_filters.items():
if var_col not in data_view.columns:
continue
upper = config.get('upper_line') or config.get('upperLine')
lower = config.get('lower_line') or config.get('lowerLine')
if not upper or not lower:
continue
x_vals = data_view[var_col].astype(float).to_numpy()
y_vals = data_view[target_variable].astype(float).to_numpy()
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
scale_ratio = y_range / x_range
b1 = float(upper.get('intercept', 0))
deg1 = float(upper.get('angle', 0))
b2 = float(lower.get('intercept', 0))
deg2 = float(lower.get('angle', 0))
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
y1 = m1 * x_vals + b1
y2 = m2 * x_vals + b2
lower_bound = np.minimum(y1, y2)
upper_bound = np.maximum(y1, y2)
keep_mask = (y_vals >= lower_bound) & (y_vals <= upper_bound)
if len(keep_mask) == len(combined_keep_mask):
keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
if keep_mask is not None and len(keep_mask) == n:
combined_keep_mask &= keep_mask
return data_view.loc[combined_keep_mask]