robust_constantan_ml_model.py

"""
Robust ML + Matroid Partitioning model for Constantan-type Cu-Ni alloy design.

Purpose
-------
Train a property surrogate model for candidate alloy-process-test configurations,
estimate prediction uncertainty from an ensemble, evaluate robust feasibility under
composition/process/test perturbations, then assign candidates into interpretable
matroid-constrained design buckets.

Expected training CSV columns
-----------------------------
Features:
    x_Cu_wt, x_Ni_wt, dopant_wt, cold_work_pct, anneal_temp_C,
    anneal_time_min, test_temp_C, cyclic_strain_pct, cooling_rate_C_s,
    grain_size_um, cast_route, cooling_route

Targets:
    resistivity_uohm_cm, tcr_ppm_K, seebeck_uV_K, resistance_drift_ppm,
    hardness_HV, strength_MPa, ductility_pct, stability_score

The script includes a synthetic data generator only for pipeline testing.
Replace it with real experimental/CALPHAD/simulation data before using conclusions.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Optional, Any
import json
import warnings

import numpy as np
import pandas as pd

from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
import joblib


FEATURE_COLUMNS = [
    "x_Cu_wt",
    "x_Ni_wt",
    "dopant_wt",
    "cold_work_pct",
    "anneal_temp_C",
    "anneal_time_min",
    "test_temp_C",
    "cyclic_strain_pct",
    "cooling_rate_C_s",
    "grain_size_um",
    "cast_route",
    "cooling_route",
]

NUMERIC_FEATURES = [
    "x_Cu_wt",
    "x_Ni_wt",
    "dopant_wt",
    "cold_work_pct",
    "anneal_temp_C",
    "anneal_time_min",
    "test_temp_C",
    "cyclic_strain_pct",
    "cooling_rate_C_s",
    "grain_size_um",
]

CATEGORICAL_FEATURES = ["cast_route", "cooling_route"]

TARGET_COLUMNS = [
    "resistivity_uohm_cm",
    "tcr_ppm_K",
    "seebeck_uV_K",
    "resistance_drift_ppm",
    "hardness_HV",
    "strength_MPa",
    "ductility_pct",
    "stability_score",
]


@dataclass
class PropertyWindows:
    """Acceptable target windows for robust feasibility checks."""
    resistivity_uohm_cm: Tuple[float, float] = (45.0, 55.0)
    tcr_ppm_K_abs_max: float = 60.0
    resistance_drift_ppm_abs_max: float = 120.0
    hardness_HV: Tuple[float, float] = (90.0, 230.0)
    strength_MPa: Tuple[float, float] = (250.0, 850.0)
    ductility_pct_min: float = 8.0
    stability_score_min: float = 0.60


@dataclass
class ScenarioConfig:
    """Perturbation ranges used to test robustness."""
    n_scenarios: int = 64
    delta_x_Ni_wt: float = 0.35
    delta_dopant_wt: float = 0.02
    delta_cold_work_pct: float = 2.0
    delta_anneal_temp_C: float = 10.0
    delta_anneal_time_min: float = 5.0
    delta_test_temp_C: float = 3.0
    delta_cyclic_strain_pct: float = 0.015
    delta_cooling_rate_frac: float = 0.10
    measurement_noise_scale: Dict[str, float] = field(default_factory=lambda: {
        "resistivity_uohm_cm": 0.15,
        "tcr_ppm_K": 3.0,
        "seebeck_uV_K": 0.10,
        "resistance_drift_ppm": 5.0,
        "hardness_HV": 2.0,
        "strength_MPa": 5.0,
        "ductility_pct": 0.30,
        "stability_score": 0.02,
    })


@dataclass
class MatroidConfig:
    """Capacity constraints for greedy partition matroids."""
    bucket_size: int = 6
    max_per_nickel_bin: int = 2
    max_per_cold_work_class: int = 2
    max_per_anneal_bin: int = 2
    min_pass_rate: float = 0.80


class RobustConstantanML:
    """
    Multi-output property surrogate + uncertainty estimator + robust feasibility scorer.
    """

    def __init__(
        self,
        n_estimators: int = 400,
        random_state: int = 7,
        property_windows: Optional[PropertyWindows] = None,
        scenario_config: Optional[ScenarioConfig] = None,
    ) -> None:
        self.property_windows = property_windows or PropertyWindows()
        self.scenario_config = scenario_config or ScenarioConfig()
        self.random_state = random_state

        preprocessor = ColumnTransformer(
            transformers=[
                ("num", StandardScaler(), NUMERIC_FEATURES),
                ("cat", OneHotEncoder(handle_unknown="ignore"), CATEGORICAL_FEATURES),
            ]
        )

        model = RandomForestRegressor(
            n_estimators=n_estimators,
            max_depth=None,
            min_samples_leaf=2,
            random_state=random_state,
            n_jobs=-1,
        )

        self.pipeline = Pipeline(
            steps=[
                ("preprocess", preprocessor),
                ("model", model),
            ]
        )

    def _validate_columns(self, df: pd.DataFrame, require_targets: bool = False) -> None:
        missing_features = [c for c in FEATURE_COLUMNS if c not in df.columns]
        if missing_features:
            raise ValueError(f"Missing feature columns: {missing_features}")

        if require_targets:
            missing_targets = [c for c in TARGET_COLUMNS if c not in df.columns]
            if missing_targets:
                raise ValueError(f"Missing target columns: {missing_targets}")

    def fit(self, df: pd.DataFrame) -> Dict[str, Any]:
        """Fit the multi-output surrogate model."""
        self._validate_columns(df, require_targets=True)
        X = df[FEATURE_COLUMNS].copy()
        y = df[TARGET_COLUMNS].copy()
        self.pipeline.fit(X, y)
        return {"status": "fit_complete", "n_rows": len(df)}

    def evaluate(self, df: pd.DataFrame, test_size: float = 0.25) -> pd.DataFrame:
        """Train/test evaluation on a supplied dataset."""
        self._validate_columns(df, require_targets=True)
        train_df, test_df = train_test_split(
            df, test_size=test_size, random_state=self.random_state
        )
        self.fit(train_df)
        pred = self.pipeline.predict(test_df[FEATURE_COLUMNS])

        rows = []
        for j, target in enumerate(TARGET_COLUMNS):
            rows.append({
                "target": target,
                "MAE": mean_absolute_error(test_df[target], pred[:, j]),
                "R2": r2_score(test_df[target], pred[:, j]),
            })
        return pd.DataFrame(rows)

    def predict_mean_std(self, candidates: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Predict mean and uncertainty using the distribution across random-forest trees.
        """
        self._validate_columns(candidates, require_targets=False)

        preprocess = self.pipeline.named_steps["preprocess"]
        rf = self.pipeline.named_steps["model"]
        X_trans = preprocess.transform(candidates[FEATURE_COLUMNS])
        tree_preds = np.stack([tree.predict(X_trans) for tree in rf.estimators_], axis=0)

        mean = tree_preds.mean(axis=0)
        std = tree_preds.std(axis=0)

        mean_df = pd.DataFrame(mean, columns=TARGET_COLUMNS, index=candidates.index)
        std_df = pd.DataFrame(std, columns=[f"{c}_std" for c in TARGET_COLUMNS], index=candidates.index)
        return mean_df, std_df

    def make_scenarios(self, candidates: pd.DataFrame) -> pd.DataFrame:
        """
        Generate perturbed candidate rows for robust feasibility testing.
        """
        self._validate_columns(candidates, require_targets=False)
        cfg = self.scenario_config
        rng = np.random.default_rng(self.random_state)

        repeated = pd.concat([candidates.copy()] * cfg.n_scenarios, ignore_index=False)
        repeated = repeated.reset_index(names="candidate_index")
        repeated["scenario_id"] = np.repeat(np.arange(cfg.n_scenarios), len(candidates))

        def uniform_delta(scale: float, size: int) -> np.ndarray:
            return rng.uniform(-scale, scale, size=size)

        n = len(repeated)

        repeated["x_Ni_wt"] = repeated["x_Ni_wt"] + uniform_delta(cfg.delta_x_Ni_wt, n)
        repeated["dopant_wt"] = (repeated["dopant_wt"] + uniform_delta(cfg.delta_dopant_wt, n)).clip(lower=0)
        repeated["x_Cu_wt"] = 100.0 - repeated["x_Ni_wt"] - repeated["dopant_wt"]

        repeated["cold_work_pct"] = (repeated["cold_work_pct"] + uniform_delta(cfg.delta_cold_work_pct, n)).clip(0, 95)
        repeated["anneal_temp_C"] = repeated["anneal_temp_C"] + uniform_delta(cfg.delta_anneal_temp_C, n)
        repeated["anneal_time_min"] = (repeated["anneal_time_min"] + uniform_delta(cfg.delta_anneal_time_min, n)).clip(lower=1)
        repeated["test_temp_C"] = repeated["test_temp_C"] + uniform_delta(cfg.delta_test_temp_C, n)
        repeated["cyclic_strain_pct"] = (repeated["cyclic_strain_pct"] + uniform_delta(cfg.delta_cyclic_strain_pct, n)).clip(lower=0)
        repeated["cooling_rate_C_s"] = repeated["cooling_rate_C_s"] * rng.uniform(
            1.0 - cfg.delta_cooling_rate_frac,
            1.0 + cfg.delta_cooling_rate_frac,
            size=n,
        )

        return repeated

    def _property_pass_mask(self, y: pd.DataFrame) -> pd.Series:
        """Boolean pass/fail against target windows."""
        w = self.property_windows
        return (
            y["resistivity_uohm_cm"].between(*w.resistivity_uohm_cm)
            & (y["tcr_ppm_K"].abs() <= w.tcr_ppm_K_abs_max)
            & (y["resistance_drift_ppm"].abs() <= w.resistance_drift_ppm_abs_max)
            & y["hardness_HV"].between(*w.hardness_HV)
            & y["strength_MPa"].between(*w.strength_MPa)
            & (y["ductility_pct"] >= w.ductility_pct_min)
            & (y["stability_score"] >= w.stability_score_min)
        )

    def robust_score(self, candidates: pd.DataFrame) -> pd.DataFrame:
        """
        Return nominal predictions, prediction uncertainty, scenario pass-rate,
        worst-case loss, and robust feasibility flag for each candidate.
        """
        mean_df, std_df = self.predict_mean_std(candidates)

        scenarios = self.make_scenarios(candidates)
        scenario_pred, _ = self.predict_mean_std(scenarios)

        # Add measurement-noise safety margin by pessimistically widening predicted response.
        cfg = self.scenario_config
        noisy = scenario_pred.copy()
        for col, scale in cfg.measurement_noise_scale.items():
            noisy[col] += np.random.default_rng(self.random_state + 13).normal(0, scale, len(noisy))

        pass_mask = self._property_pass_mask(noisy)
        scenario_eval = scenarios[["candidate_index", "scenario_id"]].copy()
        scenario_eval["pass"] = pass_mask.to_numpy()

        # A simple interpretable robust loss. Lower is better.
        w = self.property_windows
        rho_mid = 0.5 * (w.resistivity_uohm_cm[0] + w.resistivity_uohm_cm[1])

        loss = (
            (noisy["resistivity_uohm_cm"] - rho_mid).abs() / max(1e-9, rho_mid)
            + noisy["tcr_ppm_K"].abs() / max(1e-9, w.tcr_ppm_K_abs_max)
            + noisy["resistance_drift_ppm"].abs() / max(1e-9, w.resistance_drift_ppm_abs_max)
            + np.maximum(0, w.hardness_HV[0] - noisy["hardness_HV"]) / w.hardness_HV[0]
            + np.maximum(0, noisy["hardness_HV"] - w.hardness_HV[1]) / w.hardness_HV[1]
            + np.maximum(0, w.ductility_pct_min - noisy["ductility_pct"]) / w.ductility_pct_min
            + np.maximum(0, w.stability_score_min - noisy["stability_score"]) / w.stability_score_min
        )
        scenario_eval["loss"] = loss.to_numpy()

        agg = scenario_eval.groupby("candidate_index").agg(
            pass_rate=("pass", "mean"),
            worst_case_loss=("loss", "max"),
            mean_loss=("loss", "mean"),
        )

        out = candidates.copy()
        for c in mean_df.columns:
            out[f"pred_{c}"] = mean_df[c].values
        for c in std_df.columns:
            out[c] = std_df[c].values

        out = out.join(agg, how="left")
        out["robust_feasible"] = out["pass_rate"] >= 0.80

        # Overall score: high pass rate, low worst loss, low uncertainty.
        uncertainty_cols = [f"{c}_std" for c in TARGET_COLUMNS]
        out["mean_prediction_std"] = out[uncertainty_cols].mean(axis=1)
        out["robust_score"] = (
            2.0 * out["pass_rate"]
            - out["worst_case_loss"]
            - 0.01 * out["mean_prediction_std"]
        )
        return out

    def save(self, path: str) -> None:
        payload = {
            "pipeline": self.pipeline,
            "property_windows": self.property_windows,
            "scenario_config": self.scenario_config,
            "random_state": self.random_state,
        }
        joblib.dump(payload, path)

    @classmethod
    def load(cls, path: str) -> "RobustConstantanML":
        payload = joblib.load(path)
        obj = cls(
            property_windows=payload["property_windows"],
            scenario_config=payload["scenario_config"],
            random_state=payload["random_state"],
        )
        obj.pipeline = payload["pipeline"]
        return obj


def add_matroid_classes(df: pd.DataFrame) -> pd.DataFrame:
    """Create discrete blocks used by the partition matroid constraints."""
    out = df.copy()
    out["nickel_bin"] = pd.cut(
        out["x_Ni_wt"],
        bins=[0, 35, 40, 45, 50, 55, 100],
        labels=["<35", "35-40", "40-45", "45-50", "50-55", ">55"],
        include_lowest=True,
    ).astype(str)

    out["cold_work_class"] = pd.cut(
        out["cold_work_pct"],
        bins=[-0.1, 20, 50, 100],
        labels=["low", "moderate", "high"],
        include_lowest=True,
    ).astype(str)

    out["anneal_bin"] = pd.cut(
        out["anneal_temp_C"],
        bins=[0, 350, 500, 650, 2000],
        labels=["low_T", "mid_T", "high_T", "very_high_T"],
        include_lowest=True,
    ).astype(str)

    out["uncertainty_class"] = pd.cut(
        out["mean_prediction_std"],
        bins=[-np.inf, out["mean_prediction_std"].quantile(0.33),
              out["mean_prediction_std"].quantile(0.66), np.inf],
        labels=["low_uq", "mid_uq", "high_uq"],
        include_lowest=True,
    ).astype(str)

    return out


def _bucket_score(df: pd.DataFrame, bucket: str) -> pd.Series:
    """Bucket-specific priority functions."""
    score = df["robust_score"].copy()

    if bucket == "electrical_stability":
        score += (
            -0.015 * df["pred_tcr_ppm_K"].abs()
            -0.005 * df["pred_resistance_drift_ppm"].abs()
        )
    elif bucket == "high_resistivity":
        score += 0.04 * df["pred_resistivity_uohm_cm"]
    elif bucket == "formability":
        score += 0.08 * df["pred_ductility_pct"] - 0.005 * df["pred_hardness_HV"]
    elif bucket == "robust_manufacturing":
        score += 1.5 * df["pass_rate"] - 0.02 * df["mean_prediction_std"]
    elif bucket == "experimental_validation":
        # Prefer good but diverse mid-uncertainty candidates for learning.
        median_uq = df["mean_prediction_std"].median()
        score += -0.02 * (df["mean_prediction_std"] - median_uq).abs()

    return score


def greedy_matroid_partition(
    scored_candidates: pd.DataFrame,
    config: Optional[MatroidConfig] = None,
    buckets: Optional[List[str]] = None,
) -> Dict[str, pd.DataFrame]:
    """
    Greedy partitioning under interpretable capacity constraints.
    This is a practical partition-matroid heuristic, not a proof-optimal exact solver.
    """
    config = config or MatroidConfig()
    buckets = buckets or [
        "electrical_stability",
        "high_resistivity",
        "formability",
        "robust_manufacturing",
        "experimental_validation",
    ]

    df = add_matroid_classes(scored_candidates)
    df = df[df["pass_rate"] >= config.min_pass_rate].copy()

    used_indices = set()
    partitions: Dict[str, pd.DataFrame] = {}

    for bucket in buckets:
        pool = df[~df.index.isin(used_indices)].copy()
        if pool.empty:
            partitions[bucket] = pool
            continue

        pool["bucket_priority"] = _bucket_score(pool, bucket)
        pool = pool.sort_values("bucket_priority", ascending=False)

        counts = {
            "nickel_bin": {},
            "cold_work_class": {},
            "anneal_bin": {},
        }
        selected_rows = []

        for idx, row in pool.iterrows():
            if len(selected_rows) >= config.bucket_size:
                break

            checks = [
                counts["nickel_bin"].get(row["nickel_bin"], 0) < config.max_per_nickel_bin,
                counts["cold_work_class"].get(row["cold_work_class"], 0) < config.max_per_cold_work_class,
                counts["anneal_bin"].get(row["anneal_bin"], 0) < config.max_per_anneal_bin,
            ]

            if all(checks):
                selected_rows.append(idx)
                used_indices.add(idx)
                for key in counts:
                    val = row[key]
                    counts[key][val] = counts[key].get(val, 0) + 1

        partitions[bucket] = pool.loc[selected_rows].copy()

    return partitions


def generate_synthetic_constantan_data(n: int = 600, random_state: int = 7) -> pd.DataFrame:
    """
    Synthetic data generator for debugging the pipeline.
    These formulas are illustrative only and should be replaced by real measurements.
    """
    rng = np.random.default_rng(random_state)

    x_Ni = rng.uniform(35, 55, n)
    dopant = rng.choice([0.0, 0.05, 0.10, 0.15, 0.20], size=n)
    x_Cu = 100 - x_Ni - dopant
    cold = rng.uniform(5, 75, n)
    ann_T = rng.uniform(300, 650, n)
    ann_t = rng.uniform(10, 90, n)
    test_T = rng.uniform(20, 150, n)
    strain = rng.uniform(0.01, 0.25, n)
    cool = np.exp(rng.uniform(np.log(0.1), np.log(20), n))
    grain = np.clip(60 - 0.06 * ann_T + 0.25 * ann_t + rng.normal(0, 5, n), 2, 120)

    cast_route = rng.choice(["slow_cooled", "standard_cast", "controlled_solidification"], size=n)
    cooling_route = rng.choice(["furnace", "air", "rapid"], size=n)

    # Illustrative response surfaces.
    rho = 42 + 0.42 * (x_Ni - 35) - 0.002 * (ann_T - 450) + 2.0 * dopant + rng.normal(0, 0.8, n)
    tcr = 85 - 3.2 * (x_Ni - 40) + 0.07 * (ann_T - 450) + rng.normal(0, 9, n)
    seebeck = -36 + 0.22 * (x_Ni - 45) + rng.normal(0, 0.6, n)
    drift = 40 + 0.85 * cold + 0.18 * (test_T - 25) - 0.10 * ann_t + 40 * strain + rng.normal(0, 16, n)
    hardness = 85 + 2.1 * cold - 0.10 * (ann_T - 300) + 18 * dopant + rng.normal(0, 8, n)
    strength = 230 + 7.8 * cold - 0.35 * (ann_T - 300) + 35 * dopant + rng.normal(0, 25, n)
    ductility = 28 - 0.23 * cold + 0.025 * (ann_T - 300) - 10 * strain + rng.normal(0, 2.0, n)
    stability = 0.78 - 0.002 * np.abs(ann_T - 480) - 0.0015 * cold - 0.25 * strain + rng.normal(0, 0.04, n)
    stability = np.clip(stability, 0, 1)

    return pd.DataFrame({
        "x_Cu_wt": x_Cu,
        "x_Ni_wt": x_Ni,
        "dopant_wt": dopant,
        "cold_work_pct": cold,
        "anneal_temp_C": ann_T,
        "anneal_time_min": ann_t,
        "test_temp_C": test_T,
        "cyclic_strain_pct": strain,
        "cooling_rate_C_s": cool,
        "grain_size_um": grain,
        "cast_route": cast_route,
        "cooling_route": cooling_route,
        "resistivity_uohm_cm": rho,
        "tcr_ppm_K": tcr,
        "seebeck_uV_K": seebeck,
        "resistance_drift_ppm": drift,
        "hardness_HV": hardness,
        "strength_MPa": strength,
        "ductility_pct": ductility,
        "stability_score": stability,
    })


def demo() -> None:
    data = generate_synthetic_constantan_data(n=750)
    model = RobustConstantanML(n_estimators=250)
    metrics = model.evaluate(data)
    print("Holdout metrics:")
    print(metrics.to_string(index=False))

    # Fit on all synthetic data after evaluation, then score candidate library.
    model.fit(data)
    candidate_library = data[FEATURE_COLUMNS].sample(80, random_state=11).reset_index(drop=True)
    scored = model.robust_score(candidate_library)

    partitions = greedy_matroid_partition(scored)
    for bucket, part in partitions.items():
        print(f"\nBUCKET: {bucket}")
        cols = [
            "x_Ni_wt", "cold_work_pct", "anneal_temp_C",
            "pred_resistivity_uohm_cm", "pred_tcr_ppm_K",
            "pred_hardness_HV", "pred_ductility_pct",
            "pass_rate", "worst_case_loss", "robust_score",
            "nickel_bin", "cold_work_class", "anneal_bin",
        ]
        if len(part) == 0:
            print("No candidates selected.")
        else:
            print(part[cols].round(3).to_string())

    model.save("robust_constantan_surrogate.joblib")
    scored.to_csv("scored_candidates.csv", index=False)
    print("\nSaved: robust_constantan_surrogate.joblib and scored_candidates.csv")


if __name__ == "__main__":
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        demo()