"""Train a leakage-safe nonlinear candidate from the private episode mart. The model is fitted only on non-audit 2016-2022 episodes, selected on 2023, and probability-calibrated on 2024. The locked 2025 partition is evaluated only when ``--evaluate-locked`` is supplied. No row-level predictions are written. """ from __future__ import annotations import argparse import hashlib import json import math import os import statistics import sys from collections import Counter from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Dict, List, Mapping, Optional, Sequence, Tuple try: from scripts import train_baselines as baselines except ModuleNotFoundError: # Direct execution places scripts/ on sys.path. import train_baselines as baselines PROJECT_ROOT = Path(__file__).resolve().parents[1] MODEL_VERSION = "hist_gradient_boosting_v1" DEFAULT_OUTPUT = PROJECT_ROOT / "artifacts/private/tree" / MODEL_VERSION NUMERIC_FEATURES = baselines.NUMERIC_FEATURES CATEGORICAL_FEATURES = baselines.CATEGORICAL_FEATURES MISSING_CATEGORY_CODE = 0 RARE_CATEGORY_CODE = 1 UNKNOWN_CATEGORY_CODE = 2 FIRST_KNOWN_CATEGORY_CODE = 3 MAX_HISTOGRAM_BINS = 255 # These columns are validated by the baseline mart loader but are deliberately # unavailable to the nonlinear feature encoder. EXCLUDED_FROM_PREDICTORS = ( "vehicle_token", "vehicle_bucket", "is_vin_audit", "episode_number", "episode_start", "first_outcome", "target_nonpass", "eligible_returning_target", "temporal_partition", "source_era", "target_outcome_label_source", ) @dataclass(frozen=True) class TreeCandidate: learning_rate: float max_leaf_nodes: int l2_regularization: float def artifact_state(self) -> Dict[str, object]: return { "learning_rate": self.learning_rate, "max_leaf_nodes": self.max_leaf_nodes, "l2_regularization": self.l2_regularization, } @dataclass class TreeFeatureEncoder: """Train-only median imputation and bounded ordinal category encoding.""" min_category_count: int max_categories: int numeric_medians: Dict[str, float] seen_categories: Dict[str, set] known_category_codes: Dict[str, Dict[str, int]] feature_names: Tuple[str, ...] categorical_mask: Tuple[bool, ...] @classmethod def fit( cls, rows: Sequence[baselines.EpisodeRow], min_category_count: int, max_categories: int, ) -> "TreeFeatureEncoder": if not rows: raise baselines.DataValidationError( "Cannot fit tree preprocessing on zero rows" ) if min_category_count < 1: raise baselines.DataValidationError( "min_category_count must be at least 1" ) if not FIRST_KNOWN_CATEGORY_CODE + 1 <= max_categories <= MAX_HISTOGRAM_BINS: raise baselines.DataValidationError( "max_categories must be between {} and {}".format( FIRST_KNOWN_CATEGORY_CODE + 1, MAX_HISTOGRAM_BINS ) ) numeric_medians: Dict[str, float] = {} for name in NUMERIC_FEATURES: observed = [ row.numeric[name] for row in rows if row.numeric[name] is not None ] if not observed: raise baselines.DataValidationError( "Training data has no observed values for numeric feature " + name ) median = float(statistics.median(observed)) if not math.isfinite(median): raise baselines.DataValidationError( "Training median is non-finite for numeric feature " + name ) numeric_medians[name] = median seen_categories: Dict[str, set] = {} known_category_codes: Dict[str, Dict[str, int]] = {} category_capacity = max_categories - FIRST_KNOWN_CATEGORY_CODE for name in CATEGORICAL_FEATURES: observed_values = [ row.categorical[name] for row in rows if row.categorical[name] is not None ] counts = Counter(observed_values) seen_categories[name] = set(counts) candidates = [ (value, count) for value, count in counts.items() if count >= min_category_count ] # Frequency first, lexical second makes capping deterministic. candidates.sort(key=lambda item: (-item[1], item[0])) kept = candidates[:category_capacity] known_category_codes[name] = { value: FIRST_KNOWN_CATEGORY_CODE + index for index, (value, _count) in enumerate(kept) } feature_names = tuple( list(NUMERIC_FEATURES) + ["missing__" + name for name in NUMERIC_FEATURES] + list(CATEGORICAL_FEATURES) ) categorical_mask = tuple( [False] * (2 * len(NUMERIC_FEATURES)) + [True] * len(CATEGORICAL_FEATURES) ) return cls( min_category_count=min_category_count, max_categories=max_categories, numeric_medians=numeric_medians, seen_categories=seen_categories, known_category_codes=known_category_codes, feature_names=feature_names, categorical_mask=categorical_mask, ) def category_code(self, feature: str, value: Optional[str]) -> int: if value is None: return MISSING_CATEGORY_CODE known = self.known_category_codes[feature] if value in known: return known[value] if value in self.seen_categories[feature]: return RARE_CATEGORY_CODE return UNKNOWN_CATEGORY_CODE def transform(self, rows: Sequence[baselines.EpisodeRow], np_module: object) -> object: matrix = np_module.empty((len(rows), len(self.feature_names)), dtype=np_module.float64) for row_index, row in enumerate(rows): column = 0 for name in NUMERIC_FEATURES: value = row.numeric[name] matrix[row_index, column] = ( self.numeric_medians[name] if value is None else float(value) ) column += 1 for name in NUMERIC_FEATURES: matrix[row_index, column] = 1.0 if row.numeric[name] is None else 0.0 column += 1 for name in CATEGORICAL_FEATURES: matrix[row_index, column] = float( self.category_code(name, row.categorical[name]) ) column += 1 if matrix.size and not bool(np_module.isfinite(matrix).all()): raise baselines.DataValidationError( "Tree feature matrix contains non-finite values" ) return matrix def artifact_state(self) -> Dict[str, object]: return { "min_category_count": self.min_category_count, "max_categories": self.max_categories, "reserved_category_codes": { "missing": MISSING_CATEGORY_CODE, "rare_seen_in_training": RARE_CATEGORY_CODE, "unknown_after_training": UNKNOWN_CATEGORY_CODE, "first_known": FIRST_KNOWN_CATEGORY_CODE, }, "numeric_medians": self.numeric_medians, "seen_categories": { name: sorted(values) for name, values in self.seen_categories.items() }, "known_category_codes": self.known_category_codes, "feature_names": self.feature_names, "categorical_mask": self.categorical_mask, } @dataclass class TrainedTreeModel: encoder: TreeFeatureEncoder model: object platt_model: object selected_candidate: TreeCandidate model_n_iter: int platt_n_iter: int max_iter: int calibration_max_iter: int tuning_results: List[Dict[str, object]] def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--mart", required=True, type=Path) parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT) parser.add_argument("--seed", type=int, default=20260715) parser.add_argument("--min-category-count", type=int, default=100) parser.add_argument("--max-categories", type=int, default=128) parser.add_argument("--learning-rates", default="0.05,0.1") parser.add_argument("--max-leaf-nodes", default="15,31") parser.add_argument("--l2-grid", default="1.0") parser.add_argument("--min-samples-leaf", type=int, default=20) parser.add_argument("--max-iter", type=int, default=300) parser.add_argument("--n-iter-no-change", type=int, default=20) parser.add_argument("--calibration-max-iter", type=int, default=5000) parser.add_argument("--calibration-bins", type=int, default=10) parser.add_argument("--evaluate-locked", action="store_true") parser.add_argument("--overwrite", action="store_true") return parser.parse_args(argv) def _parse_float_grid(value: str, label: str, allow_zero: bool) -> List[float]: try: values = [float(part.strip()) for part in value.split(",") if part.strip()] except ValueError as exc: raise baselines.DataValidationError(label + " must contain finite numbers") from exc minimum_ok = (lambda item: item >= 0.0) if allow_zero else (lambda item: item > 0.0) if not values or any(not math.isfinite(item) or not minimum_ok(item) for item in values): raise baselines.DataValidationError(label + " contains an invalid value") return sorted(set(values)) def _parse_int_grid(value: str, label: str, minimum: int) -> List[int]: try: values = [int(part.strip()) for part in value.split(",") if part.strip()] except ValueError as exc: raise baselines.DataValidationError(label + " must contain integers") from exc if not values or any(item < minimum for item in values): raise baselines.DataValidationError(label + " contains an invalid value") return sorted(set(values)) def candidate_grid( learning_rates: Sequence[float], max_leaf_nodes: Sequence[int], l2_values: Sequence[float], ) -> List[TreeCandidate]: return [ TreeCandidate(rate, leaves, l2) for rate in sorted(set(learning_rates)) for leaves in sorted(set(max_leaf_nodes)) for l2 in sorted(set(l2_values)) ] def require_tree_dependencies() -> Tuple[object, object, object, object, object, object]: # joblib's macOS physical-core probe emits a UserWarning on this managed # host. Preserve an explicit user setting; otherwise give it the already # available logical count so warning-as-error validation remains usable. os.environ.setdefault("LOKY_MAX_CPU_COUNT", "1") try: import joblib import numpy as np import sklearn from sklearn.ensemble import HistGradientBoostingClassifier from sklearn.exceptions import ConvergenceWarning from sklearn.linear_model import LogisticRegression except ImportError as exc: raise baselines.DependencyError( "Tree training requires numpy, scikit-learn, and joblib" ) from exc return ( np, sklearn, joblib, HistGradientBoostingClassifier, LogisticRegression, ConvergenceWarning, ) def _non_audit(rows: Sequence[baselines.EpisodeRow]) -> List[baselines.EpisodeRow]: return [row for row in rows if not row.audit_vehicle] def _targets(rows: Sequence[baselines.EpisodeRow]) -> List[int]: return [row.target for row in rows] def _scores_finite(model: object, np_module: object) -> bool: for attribute in ("train_score_", "validation_score_"): values = getattr(model, attribute, None) if values is None: continue array = np_module.asarray(values) if array.size and not bool(np_module.isfinite(array).all()): return False return True def _hist_converged( model: object, convergence_messages: Sequence[str], max_iter: int, np_module: object, ) -> Tuple[bool, Optional[int]]: value = getattr(model, "n_iter_", None) if value is None: return False, None try: n_iter = int(value) except (TypeError, ValueError, OverflowError): return False, None converged = ( not convergence_messages and 0 < n_iter < max_iter and _scores_finite(model, np_module) ) return converged, n_iter def _decision_scores(model: object, features: object, np_module: object) -> object: scores = np_module.asarray(model.decision_function(features), dtype=np_module.float64) scores = scores.reshape(-1, 1) if scores.size == 0 or not bool(np_module.isfinite(scores).all()): raise baselines.DataValidationError( "Tree model produced non-finite decision scores" ) return scores def train_tree_model( mart: baselines.MartData, candidates: Sequence[TreeCandidate], min_category_count: int, max_categories: int, min_samples_leaf: int, max_iter: int, n_iter_no_change: int, calibration_max_iter: int, seed: int, ) -> TrainedTreeModel: ( np, _sklearn, _joblib, hist_class, logistic_class, convergence_warning_class, ) = require_tree_dependencies() if not candidates: raise baselines.DataValidationError("The tree tuning grid is empty") if min_samples_leaf < 1 or max_iter < 2 or n_iter_no_change < 1: raise baselines.DataValidationError("Invalid tree iteration/leaf configuration") if calibration_max_iter < 2: raise baselines.DataValidationError("calibration_max_iter must be at least 2") train_rows = _non_audit(mart.rows_by_partition["train"]) tune_rows = _non_audit(mart.rows_by_partition["tune"]) calibrate_rows = _non_audit(mart.rows_by_partition["calibrate"]) baselines._require_two_classes(train_rows, "Tree training") baselines._require_two_classes(tune_rows, "Tree tuning") baselines._require_two_classes(calibrate_rows, "Tree calibration") encoder = TreeFeatureEncoder.fit( train_rows, min_category_count=min_category_count, max_categories=max_categories, ) x_train = encoder.transform(train_rows, np) x_tune = encoder.transform(tune_rows, np) x_calibrate = encoder.transform(calibrate_rows, np) train_targets = np.asarray(_targets(train_rows), dtype=np.int8) tune_targets = _targets(tune_rows) calibrate_targets = np.asarray(_targets(calibrate_rows), dtype=np.int8) selectable: List[Tuple[float, float, float, int, float, object, int, TreeCandidate]] = [] tuning_results: List[Dict[str, object]] = [] for candidate in candidates: model = hist_class( loss="log_loss", learning_rate=candidate.learning_rate, max_iter=max_iter, max_leaf_nodes=candidate.max_leaf_nodes, min_samples_leaf=min_samples_leaf, l2_regularization=candidate.l2_regularization, max_bins=MAX_HISTOGRAM_BINS, categorical_features=list(encoder.categorical_mask), early_stopping=True, scoring="loss", validation_fraction=0.10, n_iter_no_change=n_iter_no_change, tol=1e-7, random_state=seed, class_weight=None, ) result: Dict[str, object] = { **candidate.artifact_state(), "converged": False, "n_iter": None, "convergence_warning": None, "finite_scores": False, "finite_train_probabilities": False, "finite_tune_probabilities": False, "eligible_for_selection": False, "brier": None, "average_precision": None, } convergence_messages = baselines.fit_with_convergence_capture( model, x_train, train_targets, convergence_warning_class ) converged, n_iter = _hist_converged( model, convergence_messages, max_iter, np ) result["converged"] = converged result["n_iter"] = n_iter result["convergence_warning"] = ( " | ".join(convergence_messages) if convergence_messages else None ) result["finite_scores"] = _scores_finite(model, np) train_probabilities = model.predict_proba(x_train)[:, 1] tune_probabilities = model.predict_proba(x_tune)[:, 1] finite_train = baselines.probability_array_finite(train_probabilities, np) finite_tune = baselines.probability_array_finite(tune_probabilities, np) result["finite_train_probabilities"] = finite_train result["finite_tune_probabilities"] = finite_tune eligible = converged and finite_train and finite_tune result["eligible_for_selection"] = eligible if eligible and n_iter is not None: metrics = baselines.binary_metrics(tune_targets, tune_probabilities.tolist()) brier = float(metrics["brier"]) average_precision = float(metrics["average_precision"]) result["brier"] = brier result["average_precision"] = average_precision selectable.append( ( brier, -average_precision, candidate.learning_rate, candidate.max_leaf_nodes, candidate.l2_regularization, model, n_iter, candidate, ) ) tuning_results.append(result) if not selectable: raise baselines.DataValidationError( "No tree candidate converged with finite train/tune probabilities; " "increase --max-iter or inspect the feature contract" ) ( _brier, _negative_ap, _rate, _leaves, _l2, selected_model, selected_n_iter, selected_candidate, ) = min(selectable, key=lambda item: item[:5]) calibration_scores = _decision_scores(selected_model, x_calibrate, np) platt_model = logistic_class( C=baselines.PLATT_CALIBRATION_CONFIG["c"], penalty=baselines.PLATT_CALIBRATION_CONFIG["penalty"], solver=baselines.PLATT_CALIBRATION_CONFIG["solver"], class_weight=None, max_iter=calibration_max_iter, random_state=seed, ) platt_warnings = baselines.fit_with_convergence_capture( platt_model, calibration_scores, calibrate_targets, convergence_warning_class, ) platt_n_iter = baselines.model_n_iter(platt_model, np) if ( platt_warnings or platt_n_iter is None or platt_n_iter >= calibration_max_iter or not baselines.model_parameters_finite(platt_model, np) ): detail = " | ".join(platt_warnings) if platt_warnings else "finite/convergence check" raise baselines.DataValidationError( "Tree Platt calibration did not converge: " + detail ) calibrated = platt_model.predict_proba(calibration_scores)[:, 1] if not baselines.probability_array_finite(calibrated, np): raise baselines.DataValidationError( "Tree Platt calibration produced non-finite probabilities" ) return TrainedTreeModel( encoder=encoder, model=selected_model, platt_model=platt_model, selected_candidate=selected_candidate, model_n_iter=selected_n_iter, platt_n_iter=int(platt_n_iter), max_iter=max_iter, calibration_max_iter=calibration_max_iter, tuning_results=tuning_results, ) def tree_probabilities( trained: TrainedTreeModel, rows: Sequence[baselines.EpisodeRow], calibrated: bool, ) -> List[float]: np, _sklearn, _joblib, _hist, _logistic, _warning = require_tree_dependencies() matrix = trained.encoder.transform(rows, np) if calibrated: scores = _decision_scores(trained.model, matrix, np) probabilities = trained.platt_model.predict_proba(scores)[:, 1] else: probabilities = trained.model.predict_proba(matrix)[:, 1] if not baselines.probability_array_finite(probabilities, np): raise baselines.DataValidationError( "Tree model produced non-finite probabilities" ) return probabilities.tolist() def evaluate_tree_model( mart: baselines.MartData, trained: TrainedTreeModel, evaluate_locked: bool, bin_count: int, ) -> Tuple[List[Dict[str, object]], List[Dict[str, object]]]: metric_rows: List[Dict[str, object]] = [] calibration_rows: List[Dict[str, object]] = [] partitions = ["train", "tune", "calibrate"] if evaluate_locked: partitions.append("locked_test") for partition in partitions: source_rows = mart.rows_by_partition[partition] cohorts = baselines.evaluation_cohorts( source_rows, include_audit_breakout=(partition == "locked_test") ) for cohort_name, rows in cohorts: if not rows: continue targets = _targets(rows) predictions = [ ( "hist_gradient_boosting_raw", tree_probabilities(trained, rows, calibrated=False), ) ] if partition in {"calibrate", "locked_test"}: predictions.append( ( "hist_gradient_boosting_platt", tree_probabilities(trained, rows, calibrated=True), ) ) for model_name, probabilities in predictions: metrics = baselines.binary_metrics(targets, probabilities) metric_rows.append( { "model": model_name, "partition": partition, "cohort": cohort_name, "episodes": len(rows), "vehicles": len({row.vehicle_token for row in rows}), "nonpass": sum(targets), **metrics, } ) for values in baselines.calibration_bins( targets, probabilities, bin_count ): calibration_rows.append( { "model": model_name, "partition": partition, "cohort": cohort_name, **values, } ) return metric_rows, calibration_rows def _input_sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for block in iter(lambda: handle.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def write_artifacts( output_dir: Path, mart_path: Path, mart: baselines.MartData, trained: TrainedTreeModel, metric_rows: Sequence[Mapping[str, object]], calibration_rows: Sequence[Mapping[str, object]], evaluate_locked: bool, seed: int, overwrite: bool, sklearn_version: str, ) -> None: _np, _sklearn, joblib, _hist, _logistic, _warning = require_tree_dependencies() output_dir.mkdir(parents=True, exist_ok=True) model_path = output_dir / "model.joblib" metrics_path = output_dir / "metrics.json" calibration_path = output_dir / "calibration_bins.csv" manifest_path = output_dir / "manifest.json" expected = (model_path, metrics_path, calibration_path, manifest_path) existing = [path for path in expected if path.exists()] if existing and not overwrite: raise baselines.DataValidationError( "Refusing to overwrite existing tree artifacts: " + ", ".join(path.name for path in existing) ) partial_model = model_path.with_name(model_path.name + ".partial") joblib.dump( { "model_version": MODEL_VERSION, "target_contract": baselines.TARGET_CONTRACT, "numeric_features": NUMERIC_FEATURES, "categorical_features": CATEGORICAL_FEATURES, "excluded_from_predictors": EXCLUDED_FROM_PREDICTORS, "encoder": trained.encoder.artifact_state(), "selected_candidate": trained.selected_candidate.artifact_state(), "hist_gradient_boosting_model": trained.model, "platt_model": trained.platt_model, "platt_calibration_config": baselines.PLATT_CALIBRATION_CONFIG, }, partial_model, ) os.replace(partial_model, model_path) baselines._atomic_text( metrics_path, json.dumps(list(metric_rows), indent=2, sort_keys=True, allow_nan=False) + "\n", ) baselines._atomic_csv(calibration_path, calibration_rows) partition_counts, source_counts, label_source_counts = ( baselines.manifest_audit_counts(mart, evaluate_locked=evaluate_locked) ) manifest = { "model_version": MODEL_VERSION, "generated_at_utc": datetime.now(timezone.utc).isoformat(), "classification": "private_model_artifact_no_row_predictions", "input_file": mart_path.name, "input_sha256": _input_sha256(mart_path), "metrics_sha256": _input_sha256(metrics_path), "input_rows": mart.input_rows, "eligible_rows": mart.eligible_rows, "partition_counts": partition_counts, "source_era_audit_counts": source_counts, "target_label_source_counts": label_source_counts, "target_contract": baselines.TARGET_CONTRACT, "never_fit_audit_rule": "mart is_vin_audit; validated as vehicle_bucket < 10", "locked_test_evaluated": evaluate_locked, "split_bounds": { name: { "start_inclusive": start.isoformat(), "end_exclusive": end.isoformat(), } for name, (start, end) in baselines.SPLIT_BOUNDS.items() if name != "shadow" }, "numeric_features": NUMERIC_FEATURES, "categorical_features": CATEGORICAL_FEATURES, "excluded_from_predictors": EXCLUDED_FROM_PREDICTORS, "preprocessing_fit_partition": "train_2016_2022_non_audit_only", "unknown_category_code": UNKNOWN_CATEGORY_CODE, "selected_candidate": trained.selected_candidate.artifact_state(), "convergence": { "tree_max_iter": trained.max_iter, "tree_selected_n_iter": trained.model_n_iter, "platt_max_iter": trained.calibration_max_iter, "platt_n_iter": trained.platt_n_iter, "requires_early_stop_before_max_iter": True, "requires_finite_scores_and_probabilities": True, }, "platt_calibration_config": baselines.PLATT_CALIBRATION_CONFIG, "tuning_results": trained.tuning_results, "seed": seed, "python_version": sys.version, "scikit_learn_version": sklearn_version, "artifacts": [path.name for path in expected], } baselines._atomic_text( manifest_path, json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n", ) def main(argv: Optional[Sequence[str]] = None) -> int: args = parse_args(argv) os.umask(0o077) try: mart_path = baselines.require_private_path(args.mart, "--mart") output_dir = baselines.require_private_path(args.output_dir, "--output-dir") learning_rates = _parse_float_grid( args.learning_rates, "--learning-rates", allow_zero=False ) leaf_nodes = _parse_int_grid( args.max_leaf_nodes, "--max-leaf-nodes", minimum=2 ) l2_values = _parse_float_grid(args.l2_grid, "--l2-grid", allow_zero=True) if args.calibration_bins < 2: raise baselines.DataValidationError( "--calibration-bins must be at least 2" ) mart = baselines.load_mart(mart_path) trained = train_tree_model( mart=mart, candidates=candidate_grid(learning_rates, leaf_nodes, l2_values), min_category_count=args.min_category_count, max_categories=args.max_categories, min_samples_leaf=args.min_samples_leaf, max_iter=args.max_iter, n_iter_no_change=args.n_iter_no_change, calibration_max_iter=args.calibration_max_iter, seed=args.seed, ) metric_rows, calibration_rows = evaluate_tree_model( mart, trained, evaluate_locked=args.evaluate_locked, bin_count=args.calibration_bins, ) _np, sklearn, _joblib, _hist, _logistic, _warning = ( require_tree_dependencies() ) write_artifacts( output_dir=output_dir, mart_path=mart_path, mart=mart, trained=trained, metric_rows=metric_rows, calibration_rows=calibration_rows, evaluate_locked=args.evaluate_locked, seed=args.seed, overwrite=args.overwrite, sklearn_version=sklearn.__version__, ) except baselines.BaselineError as exc: print("Tree training failed: {}".format(exc), file=sys.stderr) return 2 print("Wrote private nonlinear model artifacts to {}".format(output_dir)) if args.evaluate_locked: print("The explicitly unlocked 2025 test metrics were evaluated.") else: print("The 2025 locked test was not evaluated.") return 0 if __name__ == "__main__": raise SystemExit(main())