# -*- coding: utf-8 -*- """Runtime-only Albitron-TPH2O V1.0 public interface. This module intentionally contains no model training or model-selection logic. It loads the paired frozen prediction and support assets, reproduces the validated feature construction, predicts with the natural-sample head, and applies the target-specific support layer. No training logic is included. """ from __future__ import annotations import hashlib import json from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple import joblib import numpy as np import pandas as pd OX = ["SiO2", "TiO2", "Al2O3", "FeOt", "MgO", "CaO", "Na2O", "K2O"] UNKNOWN_REF = "__ZERO__" TARGETS = ("T", "P", "H2O") MODEL_FORMAT = "AlbitronPortablePredictionBundleV1" SUPPORT_FORMAT = "AlbitronArchiveSupportPortableV2.1" RAW_COLUMNS = ( [f"{c}_liq" for c in OX] + ["Glassraw_Total"] + [f"{c}_plag" for c in OX] + ["Plag_Total"] ) DISPLAY_TO_RAW = { "SiO2": ("SiO2_liq", "SiO2_plag"), "TiO2": ("TiO2_liq", "TiO2_plag"), "Al2O3": ("Al2O3_liq", "Al2O3_plag"), "FeOt": ("FeOt_liq", "FeOt_plag"), "MgO": ("MgO_liq", "MgO_plag"), "CaO": ("CaO_liq", "CaO_plag"), "Na2O": ("Na2O_liq", "Na2O_plag"), "K2O": ("K2O_liq", "K2O_plag"), "Analytical total": ("Glassraw_Total", "Plag_Total"), } SUPPORT_LABEL = { "Constrained": "High support", "Conditional": "Intermediate support", "Exploratory": "Low support", } def sha256_file(path: Path) -> str: h = hashlib.sha256() with Path(path).open("rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): h.update(chunk) return h.hexdigest() def schema_sha256(columns: Sequence[str]) -> str: payload = json.dumps(list(columns), ensure_ascii=False, separators=(",", ":")) return hashlib.sha256(payload.encode("utf-8")).hexdigest() def closure_100(arr: np.ndarray, eps: float = 1e-6) -> np.ndarray: x = np.asarray(arr, float) x = np.where(np.isfinite(x), x, np.nan) x = np.maximum(x, eps) s = np.nansum(x, axis=1, keepdims=True) return 100.0 * x / s def clr_transform(arr: np.ndarray, eps: float = 1e-6) -> np.ndarray: x = closure_100(arr, eps) lx = np.log(x) return lx - np.nanmean(lx, axis=1, keepdims=True) def stable_ratio(a: Sequence[float], b: Sequence[float], floor: float = 0.02) -> np.ndarray: aa = np.asarray(a, float) bb = np.asarray(b, float) out = np.full_like(aa, np.nan, dtype=float) mask = np.isfinite(aa) & np.isfinite(bb) & (np.abs(bb) >= float(floor)) out[mask] = aa[mask] / bb[mask] return out def an_ab_or_from_oxides( CaO: Sequence[float], Na2O: Sequence[float], K2O: Sequence[float] ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: mw = {"CaO": 56.077, "Na2O": 61.979, "K2O": 94.196} ca = np.asarray(CaO, float) / mw["CaO"] na = 2.0 * np.asarray(Na2O, float) / mw["Na2O"] kk = 2.0 * np.asarray(K2O, float) / mw["K2O"] total = ca + na + kk an = np.where(total > 0, ca / total, np.nan) ab = np.where(total > 0, na / total, np.nan) orr = np.where(total > 0, kk / total, np.nan) return an, ab, orr def canonicalize_natural_columns(df: pd.DataFrame) -> pd.DataFrame: """Accept the canonical Albitron-TPH2O input schema plus safe total aliases.""" out = df.copy() out.columns = [str(c).strip() for c in out.columns] aliases = { "Glassraw.Total": "Glassraw_Total", "Glassraw Total": "Glassraw_Total", "Plag.Total": "Plag_Total", "Plag Total": "Plag_Total", } out = out.rename(columns={c: aliases.get(c, c) for c in out.columns}) missing = [c for c in RAW_COLUMNS if c not in out.columns] if missing: raise ValueError("Missing required input columns: " + ", ".join(missing)) return out def validate_natural_input(df: pd.DataFrame) -> Tuple[pd.DataFrame, list[str]]: out = canonicalize_natural_columns(df) warnings: list[str] = [] for c in RAW_COLUMNS: out[c] = pd.to_numeric(out[c], errors="coerce") missing_mask = out[list(RAW_COLUMNS)].isna() if missing_mask.any().any(): bad = [] for i in range(len(out)): cols = missing_mask.columns[missing_mask.iloc[i].to_numpy()].tolist() if cols: bad.append(f"row {i + 1}: {', '.join(cols)}") raise ValueError("Missing or non-numeric required values: " + "; ".join(bad[:10])) oxide_cols = [c for c in RAW_COLUMNS if c not in {"Glassraw_Total", "Plag_Total"}] if (out[oxide_cols] < 0).any().any(): raise ValueError("Oxide concentrations cannot be negative.") if (out[["Glassraw_Total", "Plag_Total"]] <= 0).any().any(): raise ValueError("Analytical totals must be positive.") for i, row in out.iterrows(): for c, label in [("Glassraw_Total", "liquid"), ("Plag_Total", "plagioclase")]: val = float(row[c]) if val < 85 or val > 105: warnings.append( f"Row {i + 1}: unusual reported {label} analytical total ({val:.2f} wt.%). Verify the EPMA analysis." ) return out, warnings def build_natural_pairs(df_raw: pd.DataFrame) -> Tuple[pd.DataFrame, list[str]]: df, warnings = validate_natural_input(df_raw) pairs = df.copy().reset_index(drop=True) pairs["total_liq"] = pairs["Glassraw_Total"].astype(float) pairs["deficit_liq"] = 100.0 - pairs["total_liq"] pairs["total_plag"] = pairs["Plag_Total"].astype(float) pairs["deficit_plag"] = 100.0 - pairs["total_plag"] an_p, _, _ = an_ab_or_from_oxides( pairs["CaO_plag"], pairs["Na2O_plag"], pairs["K2O_plag"] ) pairs["An_obs_molpct"] = 100.0 * an_p return pairs, warnings def chemistry_feature_df(pairs: pd.DataFrame, clean_femg: bool = True) -> pd.DataFrame: x_liq = pairs[[f"{c}_liq" for c in OX]].to_numpy(float) x_plag = pairs[[f"{c}_plag" for c in OX]].to_numpy(float) clr_l = clr_transform(x_liq) clr_p = clr_transform(x_plag) clr_d = clr_l - clr_p an_l, ab_l, or_l = an_ab_or_from_oxides( pairs["CaO_liq"], pairs["Na2O_liq"], pairs["K2O_liq"] ) ratio_l_alsi = stable_ratio(pairs["Al2O3_liq"], pairs["SiO2_liq"], floor=0.05) ratio_l_femg = stable_ratio(pairs["FeOt_liq"], pairs["MgO_liq"], floor=0.02) ratio_l_alkca = stable_ratio( pairs["Na2O_liq"] + pairs["K2O_liq"], pairs["CaO_liq"], floor=0.02 ) ratio_p_alsi = stable_ratio(pairs["Al2O3_plag"], pairs["SiO2_plag"], floor=0.05) ratio_p_femg = stable_ratio(pairs["FeOt_plag"], pairs["MgO_plag"], floor=0.02) ratio_p_alkca = stable_ratio( pairs["Na2O_plag"] + pairs["K2O_plag"], pairs["CaO_plag"], floor=0.02 ) cols: Dict[str, np.ndarray] = {} for i, ox in enumerate(OX): cols[f"clr_liq_{ox}"] = clr_l[:, i] cols[f"clr_plag_{ox}"] = clr_p[:, i] cols[f"clr_d_{ox}"] = clr_d[:, i] cols.update( { "anL": an_l, "abL": ab_l, "orL": or_l, "ratioL_AlSi": ratio_l_alsi, "ratioL_AlkCa": ratio_l_alkca, "ratioP_AlSi": ratio_p_alsi, "ratioP_AlkCa": ratio_p_alkca, "ratioD_AlSi": ratio_l_alsi - ratio_p_alsi, "ratioD_AlkCa": ratio_l_alkca - ratio_p_alkca, "total_liq": pairs["total_liq"].to_numpy(float), "deficit_liq": pairs["deficit_liq"].to_numpy(float), "total_plag": pairs["total_plag"].to_numpy(float), "deficit_plag": pairs["deficit_plag"].to_numpy(float), "An_obs_molpct": pairs["An_obs_molpct"].to_numpy(float), } ) if not clean_femg: cols.update( { "ratioL_FeMg": ratio_l_femg, "ratioP_FeMg": ratio_p_femg, "ratioD_FeMg": ratio_l_femg - ratio_p_femg, } ) return pd.DataFrame(cols).replace([np.inf, -np.inf], np.nan) def equilibrium_feature_df(pairs: pd.DataFrame) -> pd.DataFrame: x_liq = pairs[[f"{c}_liq" for c in OX]].to_numpy(float) clr_l = clr_transform(x_liq) an_l, ab_l, or_l = an_ab_or_from_oxides( pairs["CaO_liq"], pairs["Na2O_liq"], pairs["K2O_liq"] ) cols: Dict[str, np.ndarray] = {} for i, ox in enumerate(OX): cols[f"clr_liq_{ox}"] = clr_l[:, i] cols.update( { "anL": an_l, "abL": ab_l, "orL": or_l, "ratioL_AlSi": stable_ratio(pairs["Al2O3_liq"], pairs["SiO2_liq"], 0.05), "ratioL_AlkCa": stable_ratio( pairs["Na2O_liq"] + pairs["K2O_liq"], pairs["CaO_liq"], 0.02 ), "total_liq": pairs["total_liq"].to_numpy(float), "deficit_liq": pairs["deficit_liq"].to_numpy(float), } ) return pd.DataFrame(cols).replace([np.inf, -np.inf], np.nan) def append_stage_context( chem: pd.DataFrame, target: str, pred: Mapping[str, np.ndarray] ) -> pd.DataFrame: out = chem.copy() if target in {"P", "H2O"}: out["T_pred_chain"] = np.asarray(pred["T"], float) if target == "H2O": out["P_pred_chain"] = np.asarray(pred["P"], float) return out def _apply_calibrator(values: np.ndarray, packed: Optional[Any]) -> np.ndarray: values = np.asarray(values, float) if packed is None: return values kind = packed.get("kind") if kind == "linear": return float(packed["slope"]) * values + float(packed["intercept"]) if kind == "sklearn": return np.asarray(packed["object"].transform(values), float) raise ValueError(f"Unknown calibrator kind: {kind}") def _transform_numeric( x: pd.DataFrame, refs: pd.Series, encoder: Mapping[str, Any], zero_ref: bool ) -> np.ndarray: columns = list(encoder["columns"]) numeric = encoder["imputer"].transform(x[columns]).astype(np.float32) if not bool(encoder["use_ref"]): return numeric ref_encoder = encoder["ref_encoder"] values = ( np.full(len(x), UNKNOWN_REF, dtype=object) if zero_ref else refs.astype(str).to_numpy(object) ) onehot = ref_encoder.transform(values.reshape(-1, 1)).astype(np.float32) return np.hstack([numeric, onehot]) def _stage_predict( packed: Mapping[str, Any], x: pd.DataFrame, refs: pd.Series, zero_ref: bool ) -> np.ndarray: spec = packed["spec"] family = spec["family"] if family in {"hgb", "et"}: xx = _transform_numeric(x, refs, packed["encoder"], zero_ref) estimator = packed["estimator"] if packed.get("aggregation") == "tree_median": pred = np.median( np.stack([tree.predict(xx) for tree in estimator.estimators_], axis=0), axis=0, ) else: pred = estimator.predict(xx) elif family == "cat": xx = x.copy() if bool(spec["use_ref"]): xx["RefCat"] = UNKNOWN_REF if zero_ref else refs.astype(str).to_numpy() pred = packed["estimator"].predict(xx) else: raise ValueError(f"Unknown model family: {family}") pred = np.asarray(pred, float) if packed["target"] == "H2O" and packed.get("h2o_hi") is not None: pred = np.clip(pred, 0.0, float(packed["h2o_hi"])) return pred def row_iqr(x: np.ndarray) -> np.ndarray: arr = np.asarray(x, float) return np.nanquantile(arr, 0.75, axis=1) - np.nanquantile(arr, 0.25, axis=1) def empirical_percentile(values: np.ndarray, reference: np.ndarray) -> np.ndarray: ref = np.asarray(reference, float) ref = np.sort(ref[np.isfinite(ref)]) x = np.asarray(values, float) out = np.full(len(x), np.nan, dtype=float) good = np.isfinite(x) if len(ref) and np.any(good): out[good] = 100.0 * np.searchsorted(ref, x[good], side="right") / len(ref) return out def distance_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray: a = np.asarray(a, float) b = np.asarray(b, float) aa = np.sum(a * a, axis=1, keepdims=True) bb = np.sum(b * b, axis=1, keepdims=True).T d2 = np.maximum(aa + bb - 2.0 * a @ b.T, 0.0) return np.sqrt(d2) def transform_embedding(x: pd.DataFrame, bundle: Mapping[str, Any]) -> np.ndarray: cols = list(bundle["feature_columns"]) missing = [c for c in cols if c not in x.columns] if missing: raise KeyError("Query feature table is missing columns: " + ", ".join(missing)) return bundle["pca"].transform( bundle["scaler"].transform(bundle["imputer"].transform(x[cols])) ) def eligible_knn_distance( z_query: np.ndarray, z_train: np.ndarray, query_refs: np.ndarray, train_refs: np.ndarray, query_samples: np.ndarray, train_samples: np.ndarray, k: int, exclude_same_sample: bool, exclude_same_ref: bool, ) -> np.ndarray: d = distance_matrix(z_query, z_train) q_refs = np.asarray(query_refs, object).astype(str) t_refs = np.asarray(train_refs, object).astype(str) q_samples = np.asarray(query_samples, object).astype(str) t_samples = np.asarray(train_samples, object).astype(str) out = np.full(len(z_query), np.nan, dtype=float) for i in range(len(z_query)): allowed = np.ones(len(z_train), dtype=bool) if exclude_same_sample: allowed &= t_samples != q_samples[i] if exclude_same_ref: allowed &= t_refs != q_refs[i] vals = d[i, allowed] vals = vals[np.isfinite(vals)] if len(vals) == 0: continue take = min(int(k), len(vals)) out[i] = float(np.median(np.partition(vals, take - 1)[:take])) return out def apply_support_bundle( bundle: Mapping[str, Any], x_query: pd.DataFrame, refs: Iterable[str], samples: Iterable[str], spread_query: Mapping[str, np.ndarray], ) -> pd.DataFrame: z = transform_embedding(x_query, bundle) refs_arr = np.asarray(list(refs), object).astype(str) samples_arr = np.asarray(list(samples), object).astype(str) cfg = dict(bundle["config"]) k = int(cfg.get("knn_k", 10)) constrained_pct = float(cfg.get("constrained_percentile", 75.0)) exploratory_pct = float(cfg.get("exploratory_percentile", 95.0)) cross_d = eligible_knn_distance( z, np.asarray(bundle["train_embedding"], float), refs_arr, np.asarray(bundle["train_refs"], object), samples_arr, np.asarray(bundle["train_samples"], object), k=k, exclude_same_sample=True, exclude_same_ref=True, ) cross_pct = empirical_percentile( cross_d, np.asarray(bundle["cross_ref_distance_reference"], float) ) coverage = np.clip(1.0 - cross_pct / 100.0, 0.0, 1.0) out = pd.DataFrame( { "cross_ref_knn_distance": cross_d, "cross_ref_knn_distance_percentile": cross_pct, "archive_coverage_index": coverage, } ) for target in TARGETS: spread = np.asarray(spread_query[target], float) spread_pct = empirical_percentile( spread, np.asarray(bundle["spread_reference"][target], float) ) stability = np.clip(1.0 - spread_pct / 100.0, 0.0, 1.0) support = np.sqrt(np.clip(coverage * stability, 0.0, 1.0)) cls = np.full(len(out), "Exploratory", dtype=object) conditional = (cross_pct <= exploratory_pct) & (spread_pct <= exploratory_pct) constrained = (cross_pct <= constrained_pct) & (spread_pct <= constrained_pct) cls[conditional] = "Conditional" cls[constrained] = "Constrained" out[f"branch_iqr_{target}"] = spread out[f"branch_iqr_percentile_{target}"] = spread_pct out[f"model_stability_index_{target}"] = stability out[f"support_index_{target}"] = support out[f"evidence_class_{target}"] = cls return out @dataclass class PredictionResult: table: pd.DataFrame warnings: list[str] class AlbitronRuntime: def __init__(self, model_path: Path | str, support_path: Path | str): self.model_path = Path(model_path) self.support_path = Path(support_path) if not self.model_path.exists(): raise FileNotFoundError(f"Prediction bundle not found: {self.model_path}") if not self.support_path.exists(): raise FileNotFoundError(f"Support bundle not found: {self.support_path}") self.model = joblib.load(self.model_path) self.support = joblib.load(self.support_path) self._validate_assets() def _validate_assets(self) -> None: if not isinstance(self.model, dict) or self.model.get("format") != MODEL_FORMAT: raise RuntimeError(f"Unsupported prediction bundle format: {self.model.get('format') if isinstance(self.model, dict) else type(self.model)}") if not isinstance(self.support, dict) or self.support.get("format") != SUPPORT_FORMAT: raise RuntimeError(f"Unsupported support bundle format: {self.support.get('format') if isinstance(self.support, dict) else type(self.support)}") model_cols = list(self.model.get("feature_columns", [])) support_cols = list(self.support.get("feature_columns", [])) if model_cols != support_cols: raise RuntimeError("Prediction and support bundles use different chemistry feature schemas.") if self.model.get("feature_schema_sha256") and ( schema_sha256(model_cols) != self.model.get("feature_schema_sha256") ): raise RuntimeError("Prediction bundle feature schema hash mismatch.") if self.support.get("feature_schema_sha256") and ( schema_sha256(support_cols) != self.support.get("feature_schema_sha256") ): raise RuntimeError("Support bundle feature schema hash mismatch.") m_data = self.model.get("source_data_sha256") s_data = self.support.get("metadata", {}).get("source_data_sha256") if m_data and s_data and m_data != s_data: raise RuntimeError("Prediction and support bundles were built from different source archives.") expected_model_hash = self.support.get("metadata", {}).get("prediction_model_sha256") if expected_model_hash: actual = sha256_file(self.model_path) if actual != expected_model_hash: raise RuntimeError("Support bundle expects a different prediction-model file hash.") @property def metadata(self) -> Dict[str, Any]: return { "model_version": self.model.get("version", "unknown"), "support_version": self.support.get("version", "unknown"), "model_sha256": sha256_file(self.model_path), "support_sha256": sha256_file(self.support_path), "feature_count": len(self.model.get("feature_columns", [])), "source_data_sha256": self.model.get("source_data_sha256"), } def _predict_natural( self, chem: pd.DataFrame, eq: pd.DataFrame ) -> Tuple[pd.DataFrame, Dict[str, np.ndarray]]: expected = list(self.model["feature_columns"]) missing = [c for c in expected if c not in chem.columns] if missing: raise ValueError(f"Missing chemistry features: {missing}") chem = chem[expected].reset_index(drop=True) refs = pd.Series(["__NATURAL__"] * len(chem), dtype=str) branch_predictions: Dict[str, Dict[str, np.ndarray]] = {} for branch_name, branch in self.model["branches"].items(): chain: Dict[str, np.ndarray] = {} for target in TARGETS: xstage = append_stage_context(chem, target, chain) pred = _stage_predict( branch["stage_models"][target], xstage, refs, zero_ref=True ) pred = _apply_calibrator(pred, branch["stage_calibrators"].get(target)) if target == "H2O": hi = branch["stage_models"][target].get("h2o_hi") if hi is not None: pred = np.clip(pred, 0.0, float(hi)) chain[target] = np.asarray(pred, float) branch_predictions[branch_name] = chain out = pd.DataFrame(index=np.arange(len(chem))) for target in TARGETS: blend = self.model["heads"]["natural"][target] z = np.column_stack( [branch_predictions[name][target] for name in blend["branch_names"]] ) pred = z @ np.asarray(blend["weights"], float) + float(blend["intercept"]) pred = _apply_calibrator(pred, blend.get("calibrator")) if target == "H2O": pred = np.clip(pred, 0.0, np.inf) out[target] = pred eq_expected = list(self.model["eq_feature_columns"]) eq_input = eq[eq_expected].reset_index(drop=True) eq_imp = self.model["eq_imputer"].transform(eq_input) out["An_pred"] = self.model["eq_model"].predict(eq_imp) spreads = { target: row_iqr( np.column_stack( [branch_predictions[name][target] for name in self.model["branches"].keys()] ) ) for target in TARGETS } return out, spreads def predict(self, raw: pd.DataFrame) -> PredictionResult: pairs, warnings = build_natural_pairs(raw) expected = list(self.model["feature_columns"]) clean_femg = not any("FeMg" in c for c in expected) chem = chemistry_feature_df(pairs, clean_femg=clean_femg) eq = equilibrium_feature_df(pairs) pred, spreads = self._predict_natural(chem, eq) sample_ids = ( raw["Sample"].astype(str).fillna("").tolist() if "Sample" in raw.columns else [f"row_{i + 1}" for i in range(len(raw))] ) support = apply_support_bundle( self.support, chem[list(self.support["feature_columns"])], refs=["__NATURAL__"] * len(chem), samples=sample_ids, spread_query=spreads, ) result = pd.DataFrame() if "Sample" in raw.columns: result["Sample"] = raw["Sample"].astype(str).to_numpy() else: result["Sample"] = sample_ids result["T_C"] = pred["T"] result["P_MPa"] = pred["P"] result["H2O_wt"] = pred["H2O"] result["An_obs_molpct"] = pairs["An_obs_molpct"].to_numpy(float) result["An_pred_molpct"] = pred["An_pred"].to_numpy(float) result["delta_An"] = result["An_obs_molpct"] - result["An_pred_molpct"] result["archive_distance_percentile"] = support["cross_ref_knn_distance_percentile"] result["archive_coverage_index"] = support["archive_coverage_index"] for target in TARGETS: result[f"{target}_disagreement_percentile"] = support[ f"branch_iqr_percentile_{target}" ] result[f"{target}_support_index"] = support[f"support_index_{target}"] result[f"{target}_evidence_class"] = support[f"evidence_class_{target}"] result[f"{target}_support"] = result[f"{target}_evidence_class"].map(SUPPORT_LABEL) return PredictionResult(result, warnings) def table_from_ui(value: pd.DataFrame | list[list[Any]] | np.ndarray) -> pd.DataFrame: if isinstance(value, pd.DataFrame): ui = value.copy() else: ui = pd.DataFrame(value, columns=["Component", "Liquid wt.%", "Plagioclase wt.%"]) if list(ui.columns) != ["Component", "Liquid wt.%", "Plagioclase wt.%"]: ui.columns = ["Component", "Liquid wt.%", "Plagioclase wt.%"] raw: Dict[str, list[float]] = {c: [np.nan] for c in RAW_COLUMNS} for _, row in ui.iterrows(): comp = str(row["Component"]).strip() if comp not in DISPLAY_TO_RAW: continue liq_col, plag_col = DISPLAY_TO_RAW[comp] raw[liq_col][0] = row["Liquid wt.%"] raw[plag_col][0] = row["Plagioclase wt.%"] return pd.DataFrame(raw)