Coverage for ml_workbench/experiment.py: 14%
163 statements
« prev ^ index » next coverage.py v7.11.0, created at 2026-01-06 16:09 +0200
1from __future__ import annotations
3from collections.abc import Mapping, Sequence
4from dataclasses import dataclass
5from typing import TYPE_CHECKING, Any
6import warnings
8import pandas as pd
10if TYPE_CHECKING: 10 ↛ 11line 10 didn't jump to line 11 because the condition on line 10 was never true
11 from .config import YamlConfig
13# Valid experiment types
14VALID_EXPERIMENT_TYPES = ["regression", "classification"]
17@dataclass(frozen=True)
18class ExperimentSpec:
19 name: str
20 description: str | None
21 models: list[str]
22 dataset: str
23 target: str
24 features: str
25 do_not_split_by: list[str]
26 metrics: list[str]
27 hold_out: dict[str, Any]
28 drop_outliers: float | None
29 type: str | None
32class Experiment:
33 """Encapsulates an experiment definition from configuration.
35 Expected structure under ``experiments``:
37 experiments:
38 exp_name:
39 description: description_text # optional
40 models: model_name or [model1, model2] # required (one or many model names)
41 dataset: dataset_name # required (one dataset name)
42 target: target_name # required (one target name)
43 features: feature_name # required (single feature name)
44 type: regression or classification # optional (experiment type)
45 do_not_split_by: [col] # optional (one or many column names)
46 metrics: metric_name or [metric1, metric2] # optional (one or many metric names)
47 hold_out: { ... } # optional
48 drop_outliers: 3.0 or 0.0 or false # optional (default: 3.0, disable with 0.0 or false)
50 Note: The 'split' component is deprecated and will be ignored if present.
51 Use 'hold_out' instead for creating holdout sets.
53 If ``type`` is not specified, it will remain None until inference in the Runner,
54 where it will be inferred from the target column type:
55 - If target is numeric -> type="regression"
56 - If target is categorical -> type="classification"
57 """
59 def __init__(self, config: YamlConfig, name: str | None = None) -> None:
60 self.config = config
62 experiments = self._get_experiments_section(config)
64 # If name is None, use the first experiment in the config
65 if name is None:
66 if not experiments:
67 raise KeyError("No experiments found in configuration") # noqa: TRY003
68 name = next(iter(experiments))
70 if name not in experiments:
71 raise KeyError(f"Experiment '{name}' not found in configuration") # noqa: TRY003
73 self.name = name
74 raw = experiments[name]
75 if not isinstance(raw, Mapping):
76 raise TypeError(f"Experiment '{name}' specification must be a mapping") # noqa: TRY003
78 description = raw.get("description")
79 models = raw.get("models")
80 dataset = raw.get("dataset")
81 target = raw.get("target")
82 features = raw.get("features")
84 if isinstance(models, str):
85 models_list = [models]
86 elif isinstance(models, Sequence):
87 models_list = [str(m) for m in list(models)]
88 else:
89 raise TypeError( # noqa: TRY003
90 f"Experiment '{name}' missing required 'models' (str or list)"
91 )
93 if not isinstance(dataset, str) or not dataset:
94 raise ValueError(f"Experiment '{name}' missing required 'dataset' string") # noqa: TRY003
96 if not isinstance(target, str) or not target:
97 raise ValueError(f"Experiment '{name}' missing required 'target' string") # noqa: TRY003
99 if not isinstance(features, str) or not features:
100 raise ValueError(f"Experiment '{name}' missing required 'features' string") # noqa: TRY003
102 do_not_split_by = raw.get("do_not_split_by", [])
103 if isinstance(do_not_split_by, str):
104 do_not_split_by_list = [do_not_split_by]
105 elif isinstance(do_not_split_by, Sequence):
106 do_not_split_by_list = [str(c) for c in list(do_not_split_by)]
107 else:
108 do_not_split_by_list = []
110 metrics = raw.get("metrics", [])
111 if isinstance(metrics, str):
112 metrics_list = [metrics]
113 elif isinstance(metrics, Sequence):
114 metrics_list = [str(m) for m in list(metrics)]
115 else:
116 metrics_list = []
118 # Warn if 'split' is present (deprecated)
119 if "split" in raw:
120 warnings.warn(
121 f"Experiment '{name}' contains deprecated 'split' component. "
122 "It will be ignored. Use 'hold_out' instead for creating holdout sets.",
123 DeprecationWarning,
124 stacklevel=2,
125 )
127 hold_out = raw.get("hold_out", {})
128 hold_out_dict = dict(hold_out) if isinstance(hold_out, Mapping) else {}
130 # Parse drop_outliers: default 3.0, disable with 0.0 or false
131 drop_outliers_raw = raw.get("drop_outliers", 3.0)
132 drop_outliers_value: float | None = None
133 if drop_outliers_raw is False or drop_outliers_raw == "false":
134 drop_outliers_value = None
135 elif isinstance(drop_outliers_raw, int | float):
136 drop_outliers_float = float(drop_outliers_raw)
137 if drop_outliers_float == 0.0:
138 drop_outliers_value = None
139 else:
140 drop_outliers_value = drop_outliers_float
141 elif (
142 isinstance(drop_outliers_raw, str) and drop_outliers_raw.lower() == "false"
143 ):
144 drop_outliers_value = None
145 else:
146 # Default to 3.0 if invalid type
147 drop_outliers_value = 3.0
149 experiment_type = raw.get("type")
150 experiment_type_str = None
151 if isinstance(experiment_type, str):
152 # Convert to lowercase and validate
153 experiment_type_lower = experiment_type.lower()
154 if experiment_type_lower not in VALID_EXPERIMENT_TYPES:
155 raise ValueError( # noqa: TRY003
156 f"Experiment '{name}' has invalid type '{experiment_type}'. "
157 f"Valid types are: {', '.join(VALID_EXPERIMENT_TYPES)}"
158 )
159 experiment_type_str = experiment_type_lower
161 self.spec = ExperimentSpec(
162 name=name,
163 description=description if isinstance(description, str) else None,
164 models=models_list,
165 dataset=dataset,
166 target=target,
167 features=features,
168 do_not_split_by=do_not_split_by_list,
169 metrics=metrics_list,
170 hold_out=hold_out_dict,
171 drop_outliers=drop_outliers_value,
172 type=experiment_type_str,
173 )
175 # Store inferred type separately (since spec is frozen)
176 self._inferred_type: str | None = None
178 @staticmethod
179 def _get_experiments_section(config: YamlConfig) -> Mapping[str, Any]:
180 experiments = config.get_data().get("experiments")
181 if not isinstance(experiments, Mapping):
182 raise KeyError("No 'experiments' section found in configuration") # noqa: TRY003
183 return experiments
185 def infer_type_from_dataset(self, dataset: pd.DataFrame) -> str:
186 """Infer experiment type from target column dtype.
188 Parameters
189 ----------
190 dataset : pd.DataFrame
191 Dataset containing the target column
193 Returns
194 -------
195 str
196 Inferred type: "regression" if target is numeric, "classification" if categorical
198 Raises
199 ------
200 ValueError
201 If inferred type is not in VALID_EXPERIMENT_TYPES
202 """
203 target = self.spec.target
204 if target not in dataset.columns:
205 raise ValueError(f"Target column '{target}' not found in dataset") # noqa: TRY003
207 target_dtype = dataset[target].dtype
208 if pd.api.types.is_numeric_dtype(target_dtype):
209 inferred_type = "regression"
210 else:
211 inferred_type = "classification"
213 # Validate inferred type (should always be valid, but check for safety)
214 if inferred_type not in VALID_EXPERIMENT_TYPES:
215 raise ValueError( # noqa: TRY003
216 f"Inferred type '{inferred_type}' is not in valid types: {', '.join(VALID_EXPERIMENT_TYPES)}"
217 )
219 self._inferred_type = inferred_type
220 return inferred_type
222 def get_type(self) -> str | None:
223 """Get experiment type, returning configured type or inferred type.
225 Returns
226 -------
227 Optional[str]
228 Experiment type: "regression", "classification", or None if not yet inferred
229 """
230 if self.spec.type is not None:
231 return self.spec.type
232 return self._inferred_type
234 def to_dict(self) -> dict[str, Any]:
235 return {
236 "name": self.spec.name,
237 "description": self.spec.description,
238 "models": list(self.spec.models),
239 "dataset": self.spec.dataset,
240 "target": self.spec.target,
241 "features": self.spec.features,
242 "type": self.get_type(),
243 "do_not_split_by": list(self.spec.do_not_split_by),
244 "metrics": list(self.spec.metrics),
245 "hold_out": dict(self.spec.hold_out),
246 "drop_outliers": self.spec.drop_outliers,
247 }
249 @classmethod
250 def list_experiment_names(cls, config: YamlConfig) -> list[str]:
251 return list(cls._get_experiments_section(config).keys())
253 @classmethod
254 def verify_config(cls, config: YamlConfig) -> None:
255 """Validate experiments reference existing models, datasets, and features.
257 Only performs validation if an 'experiments' section exists.
258 """
259 experiments = config.get_data().get("experiments")
260 if experiments is None:
261 return
262 if not isinstance(experiments, Mapping):
263 raise TypeError("'experiments' section must be a mapping") # noqa: TRY003
265 datasets = config.get_data().get("datasets")
266 features = config.get_data().get("features")
267 models = config.get_data().get("models")
269 if not isinstance(datasets, Mapping):
270 raise TypeError("No 'datasets' section found while validating experiments") # noqa: TRY003
271 if not isinstance(features, Mapping):
272 raise TypeError("No 'features' section found while validating experiments") # noqa: TRY003
273 if not isinstance(models, Mapping):
274 raise TypeError("No 'models' section found while validating experiments") # noqa: TRY003
276 for exp_name, raw in experiments.items():
277 if not isinstance(raw, Mapping):
278 raise TypeError( # noqa: TRY003
279 f"Experiment '{exp_name}' specification must be a mapping"
280 )
282 # Dataset must exist
283 ds_name = raw.get("dataset")
284 if not isinstance(ds_name, str) or not ds_name:
285 raise ValueError( # noqa: TRY003
286 f"Experiment '{exp_name}' missing required 'dataset' string"
287 )
288 if ds_name not in datasets:
289 raise ValueError( # noqa: TRY003
290 f"Experiment '{exp_name}' references unknown dataset '{ds_name}'"
291 )
293 # Target must be provided as a single string
294 tgt_field = raw.get("target")
295 if not isinstance(tgt_field, str) or not tgt_field:
296 raise ValueError( # noqa: TRY003
297 f"Experiment '{exp_name}' missing required 'target' string"
298 )
300 # Models must exist
301 model_field = raw.get("models")
302 model_names: list[str]
303 if isinstance(model_field, str):
304 model_names = [model_field]
305 elif isinstance(model_field, Sequence):
306 model_names = [str(m) for m in list(model_field)]
307 else:
308 raise TypeError( # noqa: TRY003
309 f"Experiment '{exp_name}' missing required 'models' (str or list)"
310 )
311 missing_models = [m for m in model_names if m not in models]
312 if missing_models:
313 raise ValueError( # noqa: TRY003
314 f"Experiment '{exp_name}' references unknown models: {', '.join(missing_models)}"
315 )
317 # Features must exist
318 feature_field = raw.get("features")
319 if not isinstance(feature_field, str) or not feature_field:
320 raise ValueError( # noqa: TRY003
321 f"Experiment '{exp_name}' missing required 'features' string"
322 )
323 if feature_field not in features:
324 raise ValueError( # noqa: TRY003
325 f"Experiment '{exp_name}' references unknown feature '{feature_field}'"
326 )
328 # Validate type if specified
329 type_field = raw.get("type")
330 if isinstance(type_field, str):
331 type_lower = type_field.lower()
332 if type_lower not in VALID_EXPERIMENT_TYPES:
333 raise ValueError( # noqa: TRY003
334 f"Experiment '{exp_name}' has invalid type '{type_field}'. "
335 f"Valid types are: {', '.join(VALID_EXPERIMENT_TYPES)}"
336 )
339__all__ = ["Experiment", "ExperimentSpec"]