Upload 29 files

.gitattributes

@@ -0,0 +1,2 @@
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.csv filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+ENV PYTHONPATH=/app
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    curl \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 8501
+
+HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health || exit 1
+
+ENTRYPOINT ["streamlit", "run", "streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]

README.md

@@ -0,0 +1,35 @@
+---
+title: Concrete Strength Predict Design
+colorFrom: gray
+colorTo: blue
+sdk: docker
+app_port: 8501
+pinned: false
+license: mit
+---
+
+# Concrete Compressive-Strength Predict / Design
+
+An interactive tool for concrete mix design.
+
+- **Predict strength** - enter a mix design, with optional fields allowed, and get the predicted compressive strength.
+- **Design a mix** - enter a target strength and get several detailed mix designs that reach it.
+
+Trained on about 9,700 mixes spanning about 2 to 198 MPa, from normal-strength concrete through ultra-high-performance concrete.
+
+## Running locally
+
+```bash
+pip install -r requirements.txt
+streamlit run streamlit_app.py
+```
+
+## Files
+
+| file | purpose |
+|---|---|
+| `streamlit_app.py` | the web UI |
+| `inference.py` | model loading plus predict / design logic |
+| `build_inverse_index.py` | offline script to precompute the design index and `model_config.json` |
+| `checkpoints_full_rich/` | trained model files and `model_config.json` |
+| `inverse_index.csv` | precomputed predictions over the training mixes |

build_inverse_index.py

@@ -0,0 +1,114 @@
+"""Precompute the inverse-design retrieval index + the serving config.
+
+Runs the trained GNN over every training row once, so the app's inverse
+("strength -> mixes") can retrieve without any live model calls. Also derives
+the per-feature bounds (for UI defaults, range warnings and refine clipping)
+and copies the held-out metrics out of the checkpoint, writing them next to the
+checkpoints as ``model_config.json``.
+
+Run (after training):
+    python app/build_inverse_index.py \
+        --data Data/combined_rich_model_ready.csv \
+        --checkpoint-dir Hybrid/outputs/checkpoints_full_rich
+
+Writes:
+    app/inverse_index.csv
+    <checkpoint-dir>/model_config.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+
+from inference import (  # local module
+    AGE_COL,
+    AMOUNT_COLS,
+    CURING_ONLY_GLOBALS,
+    CURING_ONLY_MORTAR_GLOBALS,
+    MASKABLE_COLS,
+    Predictor,
+    _HERE,
+    build_input_frame,
+)
+
+BOUND_COLS = AMOUNT_COLS + [AGE_COL] + MASKABLE_COLS
+
+
+def _bounds(df: pd.DataFrame) -> dict:
+    out = {}
+    for c in BOUND_COLS:
+        if c not in df.columns:
+            continue
+        s = pd.to_numeric(df[c], errors="coerce").dropna()
+        if s.empty:
+            continue
+        out[c] = {
+            "min": float(s.min()), "max": float(s.max()),
+            "p01": float(s.quantile(0.01)), "p50": float(s.median()),
+            "p99": float(s.quantile(0.99)),
+        }
+    return out
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="Data/combined_rich_model_ready.csv")
+    ap.add_argument("--checkpoint-dir", default="Hybrid/outputs/checkpoints_full_rich")
+    ap.add_argument("--out", default=str(_HERE / "inverse_index.csv"))
+    ap.add_argument("--batch-size", type=int, default=256)
+    args = ap.parse_args()
+
+    ckpt_dir = Path(args.checkpoint_dir)
+    pred = Predictor(ckpt_dir)
+    df = pd.read_csv(args.data)
+    print(f"scoring {len(df)} rows with the GNN ...")
+
+    # Predict in chunks to bound memory / show progress.
+    preds = np.empty(len(df), dtype=float)
+    bs = args.batch_size
+    for start in range(0, len(df), bs):
+        chunk = df.iloc[start:start + bs].reset_index(drop=True)
+        mixes = chunk.to_dict("records")
+        preds[start:start + len(chunk)] = pred.predict_df(
+            build_input_frame(mixes), which=("gnn",)
+        )["gnn"]
+        print(f"  {min(start + bs, len(df))}/{len(df)}")
+
+    index = df.copy()
+    index["measured"] = pd.to_numeric(index["compressive_strength_mpa"], errors="coerce")
+    index["pred_gnn"] = np.round(preds, 2)
+    index.to_csv(args.out, index=False)
+    print(f"wrote {args.out}  ({len(index)} rows)")
+
+    # ---- model_config.json ----
+    ck = torch.load(ckpt_dir / "hierarchical.pt", map_location="cpu", weights_only=False)
+    config = {
+        "schema": "curing_only",
+        "curing_only_globals": list(CURING_ONLY_GLOBALS),
+        "curing_only_mortar_globals": list(CURING_ONLY_MORTAR_GLOBALS),
+        "strength_head_kind": "mortar_capacity",
+        "pool": "mean",
+        "num_layers": 2,
+        "target_mean": float(ck["target_mean"]),
+        "target_std": float(ck["target_std"]),
+        "test_metrics": ck.get("test_metrics", {}),
+        "train_data": args.data,
+        "n_rows": int(len(df)),
+        "strength_min": float(index["measured"].min()),
+        "strength_max": float(index["measured"].max()),
+        "feature_bounds": _bounds(df),
+    }
+    cfg_path = ckpt_dir / "model_config.json"
+    cfg_path.write_text(json.dumps(config, indent=2))
+    print(f"wrote {cfg_path}")
+    print("test metrics:", config["test_metrics"])
+
+
+if __name__ == "__main__":
+    main()

checkpoints_full_rich/hierarchical.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccaa6a54cf7eb6f31d5081b68900cb762fb0bc9b86e18d3dc1a589f181082039
+size 7744095

checkpoints_full_rich/model_config.json

@@ -0,0 +1,292 @@
+{
+  "schema": "curing_only",
+  "curing_only_globals": [
+    "relative_humidity",
+    "temperature_C",
+    "curing_age_days",
+    "_placeholder_0",
+    "_placeholder_1"
+  ],
+  "curing_only_mortar_globals": [
+    "mortar_curing_relative_humidity",
+    "mortar_curing_temperature_C",
+    "mortar_curing_age_days",
+    "_placeholder_0",
+    "_placeholder_1"
+  ],
+  "strength_head_kind": "mortar_capacity",
+  "pool": "mean",
+  "num_layers": 2,
+  "target_mean": 76.61848449707031,
+  "target_std": 54.902469635009766,
+  "test_metrics": {
+    "rmse": 10.050239562988281,
+    "mae": 6.930269718170166,
+    "r2": 0.9658663272857666
+  },
+  "train_data": "Data\\combined_rich_model_ready.csv",
+  "n_rows": 9701,
+  "strength_min": 2.331807851791382,
+  "strength_max": 298.0,
+  "feature_bounds": {
+    "cement_kg_m3": {
+      "min": 19.0,
+      "max": 1530.0,
+      "p01": 132.0,
+      "p50": 400.0,
+      "p99": 1150.0
+    },
+    "slag_kg_m3": {
+      "min": 0.0,
+      "max": 768.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 564.0
+    },
+    "fly_ash_kg_m3": {
+      "min": 0.0,
+      "max": 700.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 480.0
+    },
+    "silica_fume_kg_m3": {
+      "min": 0.0,
+      "max": 617.64703,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 291.0
+    },
+    "metakaolin_kg_m3": {
+      "min": 0.0,
+      "max": 510.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 78.0
+    },
+    "limestone_powder_kg_m3": {
+      "min": 0.0,
+      "max": 1058.2,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 265.3
+    },
+    "other_scm_kg_m3": {
+      "min": 0.0,
+      "max": 1113.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 722.0
+    },
+    "water_kg_m3": {
+      "min": 67.83000183105469,
+      "max": 476.0,
+      "p01": 120.0,
+      "p50": 177.23077,
+      "p99": 306.0
+    },
+    "superplasticizer_kg_m3": {
+      "min": 0.0,
+      "max": 151.7,
+      "p01": 0.0,
+      "p50": 7.903999805450439,
+      "p99": 62.15
+    },
+    "coarse_aggregate_kg_m3": {
+      "min": 0.0,
+      "max": 1816.0,
+      "p01": 0.0,
+      "p50": 877.0,
+      "p99": 1280.0
+    },
+    "fine_aggregate_kg_m3": {
+      "min": 0.0,
+      "max": 1994.0,
+      "p01": 384.6,
+      "p50": 805.0,
+      "p99": 1600.0
+    },
+    "fibre_content_kg_m3": {
+      "min": 0.0,
+      "max": 780.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 234.0
+    },
+    "fibre_length_mm": {
+      "min": 0.0,
+      "max": 60.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 24.75
+    },
+    "fibre_diameter_mm": {
+      "min": 0.0,
+      "max": 0.55,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 0.25
+    },
+    "fibre_tensile_strength_mpa": {
+      "min": 0.0,
+      "max": 4300.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 3000.0
+    },
+    "fibre_modulus_gpa": {
+      "min": 0.0,
+      "max": 250.0,
+      "p01": 0.0,
+      "p50": 0.0,
+      "p99": 200.0
+    },
+    "age_days": {
+      "min": 1.0,
+      "max": 365.0,
+      "p01": 1.0,
+      "p50": 28.0,
+      "p99": 100.0
+    },
+    "max_coarse_aggregate_size_mm": {
+      "min": 20.0,
+      "max": 40.0,
+      "p01": 20.0,
+      "p50": 20.0,
+      "p99": 40.0
+    },
+    "max_fine_aggregate_size_mm": {
+      "min": 0.045,
+      "max": 4.75,
+      "p01": 0.1,
+      "p50": 0.9,
+      "p99": 4.75
+    },
+    "curing_temperature_c": {
+      "min": 20.0,
+      "max": 210.0,
+      "p01": 20.0,
+      "p50": 22.0,
+      "p99": 210.0
+    },
+    "cement_CaO_pct": {
+      "min": 53.32,
+      "max": 76.2,
+      "p01": 54.63,
+      "p50": 63.59,
+      "p99": 69.02
+    },
+    "cement_SiO2_pct": {
+      "min": 11.8,
+      "max": 29.67,
+      "p01": 15.4,
+      "p50": 21.01,
+      "p99": 24.86
+    },
+    "cement_Al2O3_pct": {
+      "min": 0.21,
+      "max": 8.33,
+      "p01": 2.0,
+      "p50": 5.07,
+      "p99": 8.33
+    },
+    "cement_Fe2O3_pct": {
+      "min": 0.36,
+      "max": 5.4,
+      "p01": 1.8,
+      "p50": 3.318,
+      "p99": 5.1
+    },
+    "cement_MgO_pct": {
+      "min": 0.6,
+      "max": 5.31,
+      "p01": 0.64,
+      "p50": 1.91,
+      "p99": 5.11
+    },
+    "cement_SO3_pct": {
+      "min": 0.35,
+      "max": 4.1,
+      "p01": 0.6,
+      "p50": 2.6,
+      "p99": 3.67
+    },
+    "cement_alkali_pct": {
+      "min": 0.12502,
+      "max": 1.658,
+      "p01": 0.13,
+      "p50": 0.5964,
+      "p99": 1.03956
+    },
+    "cement_LOI_pct": {
+      "min": 0.0,
+      "max": 11.6,
+      "p01": 0.7,
+      "p50": 2.2,
+      "p99": 11.6
+    },
+    "scm_CaO_pct": {
+      "min": 0.01,
+      "max": 47.498333,
+      "p01": 0.03,
+      "p50": 3.3997133,
+      "p99": 43.28320000000008
+    },
+    "scm_SiO2_pct": {
+      "min": 15.435,
+      "max": 99.8,
+      "p01": 27.0,
+      "p50": 79.99951,
+      "p99": 99.8
+    },
+    "scm_Al2O3_pct": {
+      "min": 0.0,
+      "max": 38.139545,
+      "p01": 0.01,
+      "p50": 1.635,
+      "p99": 37.9693
+    },
+    "scm_Fe2O3_pct": {
+      "min": 0.0,
+      "max": 57.0,
+      "p01": 0.043465704,
+      "p50": 1.21,
+      "p99": 57.0
+    },
+    "scm_MgO_pct": {
+      "min": 0.01,
+      "max": 9.5,
+      "p01": 0.05,
+      "p50": 0.8678877,
+      "p99": 8.086817
+    },
+    "scm_LOI_pct": {
+      "min": 0.0,
+      "max": 7.5909767,
+      "p01": 0.02,
+      "p50": 2.1069672,
+      "p99": 6.07
+    },
+    "cement_grade_mpa": {
+      "min": 42.5,
+      "max": 53.0,
+      "p01": 42.5,
+      "p50": 52.5,
+      "p99": 53.0
+    },
+    "curing_humidity_pct": {
+      "min": 50.0,
+      "max": 100.0,
+      "p01": 50.0,
+      "p50": 95.0,
+      "p99": 100.0
+    },
+    "specimen_size_mm": {
+      "min": 40.0,
+      "max": 150.0,
+      "p01": 40.0,
+      "p50": 50.0,
+      "p99": 150.0
+    }
+  }
+}

checkpoints_full_rich/tabular_ann.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d801170b240d582d8f9b579fd6f61b633c8aedec30485dbac592ba457b81ed0a
+size 111795

concrete_gnn/__init__.py

@@ -0,0 +1,183 @@
+"""Hierarchical multiscale concrete GNN — unified package.
+
+This package merges the strongest pieces of the three earlier prototypes:
+
+* Codex-Claude's end-to-end-trainable multiscale architecture, with
+  structure-dependent ground-truth physics and a contact-density diagnostic;
+* Claude's modular package layout, learnable missing-feature encoder,
+  node-/edge-type embeddings, PyG-based ``NNConv`` / ``GINEConv`` operators,
+  and auxiliary mortar / ITZ losses;
+* Codex's bounded output transforms, PyG-optional fallback, and the explicit
+  placeholder slots in every schema vector.
+"""
+
+from .data import (
+    ConcreteSample,
+    MultiscaleBatch,
+    MultiscaleConcreteDataset,
+    collate_multiscale,
+    make_dataloader,
+    split_dataset,
+)
+from .graph_generator import (
+    ConcreteGraph,
+    MixDesign,
+    MortarMicrograph,
+    MortarMixDesign,
+    generate_concrete_graph,
+    generate_mortar_graph,
+)
+from .ground_truth import (
+    GroundTruthProps,
+    homogenize_concrete_targets,
+    itz_target_vector,
+    mortar_target_vector,
+    topology_statistics,
+    true_micro_properties,
+)
+from .missing_features import MissingFeatureEncoder, apply_random_missingness
+from .models import (
+    ConcreteMesoGNN,
+    HAS_PYG,
+    ITZSubModel,
+    IntegratedMultiscaleModel,
+    MortarSubModel,
+    TabularANN,
+    global_add_pool,
+    global_mean_pool,
+    transform_concrete_outputs,
+    transform_itz_outputs,
+    transform_mortar_outputs,
+)
+from .real_data import (
+    COL_STRENGTH,
+    NORM_STRENGTH,
+    TABULAR_DIM,
+    TARGET_NAME,
+    Standardizer,
+    StrengthHead,
+    ConcreteMixDataset,
+    collect_strength_predictions,
+    evaluate_strength,
+    fit_encoder_standardizers,
+    load_strength_checkpoint,
+    metrics,
+    raw_mix_vector,
+    read_table,
+    split_indices,
+    strength_column,
+)
+from .schema import (
+    AGGREGATE_FEATURES,
+    CONCRETE_TARGETS,
+    DEFAULT_SCHEMA,
+    EDGE_TYPE_AGGREGATE_AGGREGATE,
+    EDGE_TYPE_ITZ,
+    EDGE_TYPE_MORTAR_MORTAR,
+    GLOBAL_FEATURES,
+    ITZ_FEATURES,
+    ITZ_TARGETS,
+    MORTAR_EDGE_FEATURES,
+    MORTAR_FEATURES,
+    MORTAR_TARGETS,
+    NODE_TYPE_AGGREGATE,
+    NODE_TYPE_MORTAR,
+    SchemaSpec,
+)
+from .train import (
+    MultiTaskLoss,
+    TrainConfig,
+    compute_sublabel_scales,
+    compute_target_scale,
+    evaluate,
+    evaluate_metrics,
+    load_best_checkpoint,
+    train,
+    train_one_epoch,
+)
+from .visualize import (
+    collect_predictions,
+    plot_parity,
+    plot_r2_bars,
+    plot_rve_graph,
+    plot_training_curves,
+)
+
+__all__ = [
+    "AGGREGATE_FEATURES",
+    "COL_STRENGTH",
+    "CONCRETE_TARGETS",
+    "ConcreteGraph",
+    "ConcreteMesoGNN",
+    "ConcreteSample",
+    "DEFAULT_SCHEMA",
+    "EDGE_TYPE_AGGREGATE_AGGREGATE",
+    "EDGE_TYPE_ITZ",
+    "EDGE_TYPE_MORTAR_MORTAR",
+    "GLOBAL_FEATURES",
+    "GroundTruthProps",
+    "HAS_PYG",
+    "ITZ_FEATURES",
+    "ITZ_TARGETS",
+    "ITZSubModel",
+    "IntegratedMultiscaleModel",
+    "MORTAR_EDGE_FEATURES",
+    "MORTAR_FEATURES",
+    "MORTAR_TARGETS",
+    "collect_predictions",
+    "plot_parity",
+    "plot_r2_bars",
+    "plot_rve_graph",
+    "plot_training_curves",
+    "MissingFeatureEncoder",
+    "MixDesign",
+    "MortarMicrograph",
+    "MortarMixDesign",
+    "MortarSubModel",
+    "MultiTaskLoss",
+    "MultiscaleBatch",
+    "MultiscaleConcreteDataset",
+    "NODE_TYPE_AGGREGATE",
+    "NODE_TYPE_MORTAR",
+    "NORM_STRENGTH",
+    "SchemaSpec",
+    "Standardizer",
+    "StrengthHead",
+    "TABULAR_DIM",
+    "TARGET_NAME",
+    "TabularANN",
+    "TrainConfig",
+    "ConcreteMixDataset",
+    "apply_random_missingness",
+    "collate_multiscale",
+    "compute_sublabel_scales",
+    "compute_target_scale",
+    "collect_strength_predictions",
+    "evaluate",
+    "evaluate_strength",
+    "evaluate_metrics",
+    "fit_encoder_standardizers",
+    "load_best_checkpoint",
+    "load_strength_checkpoint",
+    "generate_concrete_graph",
+    "generate_mortar_graph",
+    "global_add_pool",
+    "global_mean_pool",
+    "homogenize_concrete_targets",
+    "itz_target_vector",
+    "make_dataloader",
+    "metrics",
+    "mortar_target_vector",
+    "raw_mix_vector",
+    "read_table",
+    "split_dataset",
+    "split_indices",
+    "strength_column",
+    "topology_statistics",
+    "train",
+    "train_one_epoch",
+    "transform_concrete_outputs",
+    "transform_itz_outputs",
+    "transform_mortar_outputs",
+    "true_micro_properties",
+]

concrete_gnn/categoricals.py

@@ -0,0 +1,78 @@
+"""Canonical categorical vocabularies for the UHPC rich pipeline.
+
+Single source of truth shared by the data builder (``Data/prepare_uhpc_rich.py``)
+and the loader (``real_data.py``). The builder writes the *canonical* class string
+into the CSV; the loader one-hot-encodes it for the TabularANN and maps the cement
+type to the GNN's ``cement_type_id`` paste slot.
+
+Each ``canonical_*`` mapper collapses the many free-text spellings in the raw
+spreadsheet (e.g. "Straight Steel Fibers", "Steel fibre", "Steel Fiber") into a
+small, stable vocabulary. The first vocabulary entry is the natural default for a
+missing value, except cement/curing which default to "other".
+"""
+
+from __future__ import annotations
+
+from typing import List, Sequence
+
+CEMENT_TYPES: Sequence[str] = ("opc", "type_i_ii", "type_iii", "hs", "other")
+FIBRE_TYPES: Sequence[str] = ("none", "steel", "pe", "pva", "other")
+CURING_REGIMES: Sequence[str] = ("standard", "steam", "heat", "autoclave", "other")
+
+
+def _clean(raw) -> str:
+    s = str(raw).strip().lower()
+    return "" if s in ("", "nan", "none") else s
+
+
+def canonical_cement_type(raw) -> str:
+    t = _clean(raw)
+    if not t:
+        return "other"
+    if "sulfate" in t or "type hs" in t or "hs cement" in t:
+        return "hs"
+    if "type iii" in t or "type-iii" in t or "type 3" in t:
+        return "type_iii"
+    if "i/ii" in t or "i / ii" in t:
+        return "type_i_ii"
+    if any(k in t for k in ("portland", "opc", "cem i", "type i", "type 1", "p.o", "p.i")):
+        return "opc"
+    return "other"
+
+
+def canonical_fibre_type(raw) -> str:
+    t = _clean(raw)
+    if not t:
+        return "none"
+    if "steel" in t or "seel" in t:  # "seel" = common typo in the sheet
+        return "steel"
+    if "pva" in t:
+        return "pva"
+    if t.startswith("pe") or "polyeth" in t:
+        return "pe"
+    return "other"
+
+
+def canonical_curing_regime(raw) -> str:
+    t = _clean(raw)
+    if not t:
+        return "other"
+    if "autoclave" in t:
+        return "autoclave"
+    if "steam" in t:
+        return "steam"
+    if any(k in t for k in ("heat", "hot", "warm", "thermal")):
+        return "heat"
+    if "standard" in t or "normal" in t or "ambient" in t:
+        return "standard"
+    return "other"
+
+
+def one_hot(value: str, vocab: Sequence[str]) -> List[float]:
+    """One-hot encode ``value`` against ``vocab`` (all-zero if not present)."""
+    return [1.0 if value == v else 0.0 for v in vocab]
+
+
+def type_id(value: str, vocab: Sequence[str]) -> int:
+    """Integer index of ``value`` in ``vocab`` (0 if not present)."""
+    return vocab.index(value) if value in vocab else 0

concrete_gnn/data.py

@@ -0,0 +1,372 @@
+"""Dataset, batch container and collate function for the multiscale model.
+
+Every training sample bundles three heterogeneous pieces:
+
+- a mesoscale concrete graph (aggregate + mortar nodes, three edge types);
+- a mortar micrograph (sand + paste nodes, contact edges);
+- a per-sample ITZ mix vector used to build ITZ-edge static inputs.
+
+The custom collate concatenates per-graph tensors, offsets edge indices,
+and keeps the mortar batch aligned with the concrete batch so the
+hierarchical model can wire predictions between them.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import List, Optional, Sized, cast
+
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader, Dataset
+
+from .graph_generator import (
+    ConcreteGraph,
+    MixDesign,
+    MortarMicrograph,
+    MortarMixDesign,
+    generate_concrete_graph,
+    generate_mortar_graph,
+    mortar_volume_fraction,
+)
+from .ground_truth import (
+    homogenize_concrete_targets,
+    itz_target_vector,
+    mortar_target_vector,
+    topology_statistics,
+    true_micro_properties,
+)
+from .schema import DEFAULT_SCHEMA, SchemaSpec
+
+
+# ---------------------------------------------------------------------------
+# Containers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ConcreteSample:
+    graph: ConcreteGraph
+    micrograph: MortarMicrograph
+    mix_itz: Tensor
+    concrete_target: Tensor
+    mortar_target: Optional[Tensor] = None
+    itz_target: Optional[Tensor] = None
+    tabular_mix: Optional[Tensor] = None
+
+
+@dataclass
+class MultiscaleBatch:
+    x: Tensor
+    x_mask: Tensor
+    edge_index: Tensor
+    edge_attr: Tensor
+    edge_attr_mask: Tensor
+    global_features: Tensor
+    global_mask: Tensor
+    batch: Tensor
+    pos: Tensor
+    node_type: Tensor
+    edge_type: Tensor
+    m_x: Tensor
+    m_mask: Tensor
+    m_edge_index: Tensor
+    m_edge_attr: Tensor
+    m_node_type: Tensor
+    m_batch: Tensor
+    mortar_global: Tensor
+    mortar_global_mask: Tensor
+    mix_itz: Tensor
+    concrete_target: Tensor
+    mortar_target: Optional[Tensor]
+    itz_target: Optional[Tensor]
+    tabular_mix: Optional[Tensor] = None
+
+    @property
+    def num_graphs(self) -> int:
+        return self.concrete_target.size(0)
+
+    def to(self, device: torch.device) -> "MultiscaleBatch":
+        def m(t):
+            return t.to(device) if isinstance(t, Tensor) else t
+
+        return MultiscaleBatch(
+            x=m(self.x),
+            x_mask=m(self.x_mask),
+            edge_index=m(self.edge_index),
+            edge_attr=m(self.edge_attr),
+            edge_attr_mask=m(self.edge_attr_mask),
+            global_features=m(self.global_features),
+            global_mask=m(self.global_mask),
+            batch=m(self.batch),
+            pos=m(self.pos),
+            node_type=m(self.node_type),
+            edge_type=m(self.edge_type),
+            m_x=m(self.m_x),
+            m_mask=m(self.m_mask),
+            m_edge_index=m(self.m_edge_index),
+            m_edge_attr=m(self.m_edge_attr),
+            m_node_type=m(self.m_node_type),
+            m_batch=m(self.m_batch),
+            mortar_global=m(self.mortar_global),
+            mortar_global_mask=m(self.mortar_global_mask),
+            mix_itz=m(self.mix_itz),
+            concrete_target=m(self.concrete_target),
+            mortar_target=m(self.mortar_target) if self.mortar_target is not None else None,
+            itz_target=m(self.itz_target) if self.itz_target is not None else None,
+            tabular_mix=m(self.tabular_mix) if self.tabular_mix is not None else None,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _mix_itz_vector(mix: MixDesign) -> Tensor:
+    """Per-graph mix vector used as static input to the ITZ MLP."""
+
+    return torch.tensor(
+        [
+            mix.cement_content_kg_m3,
+            float(mix.scm_type_id),
+            mix.scm_fraction,
+            mix.water_to_binder_ratio,
+            mix.admixture_dosage,
+            mix.relative_humidity,
+            mix.temperature_C,
+            mix.curing_age_days,
+        ],
+        dtype=torch.float32,
+    )
+
+
+def _mortar_mix_for(mix: MixDesign) -> MortarMixDesign:
+    """Derive a mortar-scale mix from the concrete-scale mix.
+
+    Cement content is converted from a concrete basis to a mortar basis by
+    removing the coarse-aggregate volume (see ``mortar_volume_fraction``); w/b,
+    scm_fraction and admixture dosage are intensive ratios and stay invariant.
+    Here the synthetic ``aggregate_volume_fraction`` is the true coarse fraction.
+    """
+
+    mortar_vf = mortar_volume_fraction(mix.aggregate_volume_fraction)
+    return MortarMixDesign(
+        water_to_binder_ratio=mix.water_to_binder_ratio,
+        cement_content_kg_m3=mix.cement_content_kg_m3 / mortar_vf,
+        scm_type_id=mix.scm_type_id,
+        scm_fraction=mix.scm_fraction,
+        admixture_dosage=mix.admixture_dosage,
+        curing_relative_humidity=mix.relative_humidity,
+        curing_temperature_C=mix.temperature_C,
+        curing_age_days=mix.curing_age_days,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+class MultiscaleConcreteDataset(Dataset):
+    """Synthetic dataset whose targets depend on graph structure + micrograph.
+
+    Each sample carries optional ground-truth mortar / ITZ target vectors so
+    the training loop can use auxiliary losses when the ablation calls for it.
+    """
+
+    def __init__(
+        self,
+        num_graphs: int = 200,
+        seed: int = 17,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        missing_rate: float = 0.12,
+        with_sublabels: bool = True,
+    ):
+        self.schema = schema
+        self.missing_rate = missing_rate
+        self.with_sublabels = with_sublabels
+        self.samples: List[ConcreteSample] = []
+
+        mix_rng = random.Random(seed + 1)
+        for k in range(num_graphs):
+            mix = MixDesign(
+                aggregate_volume_fraction=mix_rng.uniform(0.30, 0.48),
+                water_to_binder_ratio=mix_rng.uniform(0.32, 0.55),
+                cement_content_kg_m3=mix_rng.uniform(300.0, 430.0),
+                scm_type_id=int(mix_rng.choice([0, 1, 2])),
+                scm_fraction=mix_rng.uniform(0.0, 0.30),
+                admixture_dosage=mix_rng.uniform(0.0, 1.0),
+                relative_humidity=mix_rng.uniform(0.65, 0.98),
+                temperature_C=mix_rng.uniform(10.0, 35.0),
+                curing_age_days=float(mix_rng.choice([7, 14, 28, 56, 90])),
+            )
+            graph = generate_concrete_graph(
+                mix, schema=schema, seed=seed + 100 + k, missing_rate=missing_rate
+            )
+            mortar_mix = _mortar_mix_for(mix)
+            micro = generate_mortar_graph(
+                mortar_mix, schema=schema, seed=seed + 200 + k
+            )
+
+            agg_node_frac, itz_edge_frac, aa_edge_frac = topology_statistics(graph)
+            props = true_micro_properties(mix, micro.sand_contact_density)
+            y = homogenize_concrete_targets(
+                mix, props, agg_node_frac, itz_edge_frac, aa_edge_frac
+            )
+
+            mortar_y = mortar_target_vector(props) if with_sublabels else None
+
+            itz_count = int((graph.edge_type == 0).sum().item())
+            itz_y = (
+                itz_target_vector(props).unsqueeze(0).expand(itz_count, -1).clone()
+                if with_sublabels and itz_count > 0
+                else None
+            )
+
+            self.samples.append(
+                ConcreteSample(
+                    graph=graph,
+                    micrograph=micro,
+                    mix_itz=_mix_itz_vector(mix),
+                    concrete_target=y,
+                    mortar_target=mortar_y,
+                    itz_target=itz_y,
+                )
+            )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> ConcreteSample:
+        return self.samples[idx]
+
+
+# ---------------------------------------------------------------------------
+# Collate
+# ---------------------------------------------------------------------------
+
+
+def collate_multiscale(samples: List[ConcreteSample]) -> MultiscaleBatch:
+    xs, x_masks, e_idx, e_attr, e_attr_masks = [], [], [], [], []
+    globals_, global_masks = [], []
+    batches, positions = [], []
+    node_types, edge_types = [], []
+    m_xs, m_masks, m_eidx, m_eattr, m_node_types, m_batches = [], [], [], [], [], []
+    mortar_globals, mortar_global_masks = [], []
+    mix_itzs, ys = [], []
+    mortar_targets, itz_targets = [], []
+
+    node_offset = 0
+    micro_offset = 0
+    have_mortar_target = samples[0].mortar_target is not None
+    have_itz_target = samples[0].itz_target is not None
+    have_tabular = samples[0].tabular_mix is not None
+    tabular_mixes: List[Tensor] = []
+
+    for gid, sample in enumerate(samples):
+        graph: ConcreteGraph = sample.graph
+        micro: MortarMicrograph = sample.micrograph
+        n_nodes = graph.x.size(0)
+        n_micro = micro.x.size(0)
+
+        xs.append(graph.x)
+        x_masks.append(graph.x_mask)
+        e_idx.append(graph.edge_index + node_offset)
+        e_attr.append(graph.edge_attr)
+        e_attr_masks.append(graph.edge_attr_mask)
+        globals_.append(graph.global_features)
+        global_masks.append(graph.global_mask)
+        batches.append(torch.full((n_nodes,), gid, dtype=torch.long))
+        positions.append(graph.pos)
+        node_types.append(graph.node_type)
+        edge_types.append(graph.edge_type)
+
+        m_xs.append(micro.x)
+        m_masks.append(micro.x_mask)
+        m_eidx.append(micro.edge_index + micro_offset)
+        m_eattr.append(micro.edge_attr)
+        m_node_types.append(micro.node_type)
+        m_batches.append(torch.full((n_micro,), gid, dtype=torch.long))
+        mortar_globals.append(micro.global_features)
+        mortar_global_masks.append(micro.global_mask)
+
+        mix_itzs.append(sample.mix_itz)
+        ys.append(sample.concrete_target)
+        if have_mortar_target:
+            mortar_targets.append(sample.mortar_target)
+        if have_itz_target and sample.itz_target is not None and sample.itz_target.numel() > 0:
+            itz_targets.append(sample.itz_target)
+        if have_tabular and sample.tabular_mix is not None:
+            tabular_mixes.append(sample.tabular_mix)
+
+        node_offset += n_nodes
+        micro_offset += n_micro
+
+    return MultiscaleBatch(
+        x=torch.cat(xs, dim=0),
+        x_mask=torch.cat(x_masks, dim=0),
+        edge_index=torch.cat(e_idx, dim=1),
+        edge_attr=torch.cat(e_attr, dim=0),
+        edge_attr_mask=torch.cat(e_attr_masks, dim=0),
+        global_features=torch.stack(globals_, dim=0),
+        global_mask=torch.stack(global_masks, dim=0),
+        batch=torch.cat(batches, dim=0),
+        pos=torch.cat(positions, dim=0),
+        node_type=torch.cat(node_types, dim=0),
+        edge_type=torch.cat(edge_types, dim=0),
+        m_x=torch.cat(m_xs, dim=0),
+        m_mask=torch.cat(m_masks, dim=0),
+        m_edge_index=torch.cat(m_eidx, dim=1),
+        m_edge_attr=torch.cat(m_eattr, dim=0),
+        m_node_type=torch.cat(m_node_types, dim=0),
+        m_batch=torch.cat(m_batches, dim=0),
+        mortar_global=torch.stack(mortar_globals, dim=0),
+        mortar_global_mask=torch.stack(mortar_global_masks, dim=0),
+        mix_itz=torch.stack(mix_itzs, dim=0),
+        concrete_target=torch.stack(ys, dim=0),
+        mortar_target=torch.stack(mortar_targets, dim=0) if have_mortar_target else None,
+        itz_target=torch.cat(itz_targets, dim=0)
+        if (have_itz_target and itz_targets)
+        else None,
+        tabular_mix=torch.stack(tabular_mixes, dim=0) if have_tabular else None,
+    )
+
+
+def make_dataloader(
+    dataset: Dataset,
+    batch_size: int = 8,
+    shuffle: bool = True,
+    num_workers: int = 0,
+    pin_memory: bool = False,
+    drop_last: bool = False,
+) -> DataLoader:
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=drop_last,
+        collate_fn=collate_multiscale,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Splits
+# ---------------------------------------------------------------------------
+
+
+def split_dataset(
+    dataset: Dataset,
+    seed: int = 0,
+    train_fraction: float = 0.7,
+    val_fraction: float = 0.15,
+) -> tuple:
+    n = len(cast(Sized, dataset))
+    idx = list(range(n))
+    random.Random(seed).shuffle(idx)
+    n_train = int(train_fraction * n)
+    n_val = int(val_fraction * n)
+    return idx[:n_train], idx[n_train : n_train + n_val], idx[n_train + n_val :]

concrete_gnn/graph_generator.py

@@ -0,0 +1,800 @@
+"""Synthetic graph generators for mesoscale concrete and microscale mortar.
+
+Geometry pipeline (concrete RVE):
+    1. Random sequential adsorption of circular aggregate particles sampled
+       from a discrete grain-size distribution until the requested area
+       fraction is approximately reached.
+    2. Sample a regular grid of mortar seed points and discard any that fall
+       inside an aggregate, giving the mortar segmentation centroids.
+    3. Approximate the Voronoi adjacency between every (aggregate, mortar)
+       point pair by sampling points in the RVE and counting how often each
+       pair appears as the two nearest neighbours.  This produces both
+       adjacency information and a proxy for shared-boundary length.
+    4. Mortar-mortar adjacencies in the Voronoi partition are often
+       interrupted by aggregate cells; we therefore *augment* mortar-mortar
+       edges using k-nearest-neighbour links between mortar seeds.
+    5. Aggregate-aggregate proximity edges are added when two aggregates are
+       within a small surface gap.
+
+The mortar microscale graph reuses the same Voronoi adjacency primitive but
+with its own grain-size distribution and node features.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from .schema import (
+    DEFAULT_SCHEMA,
+    EDGE_TYPE_AGGREGATE_AGGREGATE,
+    EDGE_TYPE_ITZ,
+    EDGE_TYPE_MORTAR_MORTAR,
+    MORTAR_NODE_TYPE_PASTE,
+    MORTAR_NODE_TYPE_SAND,
+    NODE_TYPE_AGGREGATE,
+    NODE_TYPE_MORTAR,
+    SchemaSpec,
+)
+
+
+# ---------------------------------------------------------------------------
+# Mix-design containers (concrete + mortar scales kept separate)
+# ---------------------------------------------------------------------------
+
+
+def mortar_volume_fraction(coarse_aggregate_volume_fraction: float) -> float:
+    """Fraction of the concrete volume occupied by the mortar phase.
+
+    The mortar phase is everything except coarse aggregate, so it occupies
+    ``1 - phi_coarse`` of the concrete volume. Per-volume *densities* reported
+    on a concrete basis (cement, SCM, fibre content in kg/m^3) are converted to
+    a mortar basis by dividing by this fraction.
+
+    ``phi_coarse`` must be the *raw* coarse-aggregate volume fraction (0 for
+    UHPC, which has no coarse aggregate) -- not the clamped
+    ``aggregate_volume_fraction`` global feature. ``phi_coarse`` is capped at
+    0.70 (mortar floor 0.30) -- the exact complement of the
+    ``aggregate_volume_fraction`` clamp -- so the two stay consistent and the
+    conversion factor is bounded for unusually coarse mixes.
+    """
+
+    phi_coarse = min(0.70, max(0.0, coarse_aggregate_volume_fraction))
+    return max(0.30, 1.0 - phi_coarse)
+
+
+
+@dataclass
+class MixDesign:
+    """Concrete-scale mix design and curing metadata."""
+
+    rve_size_mm: float = 150.0
+    aggregate_size_bins_mm: Sequence[float] = (8.0, 12.0, 16.0, 20.0, 24.0, 32.0)
+    aggregate_size_fractions: Sequence[float] = (0.10, 0.20, 0.25, 0.20, 0.15, 0.10)
+    aggregate_volume_fraction: float = 0.40
+    water_to_binder_ratio: float = 0.42
+    cement_content_kg_m3: float = 360.0
+    scm_type_id: int = 1
+    scm_fraction: float = 0.12
+    admixture_dosage: float = 0.5
+    relative_humidity: float = 0.95
+    temperature_C: float = 20.0
+    curing_age_days: float = 28.0
+    fibre_content_kg_m3: float = 0.0
+    fibre_length_mm: float = 0.0
+    fibre_diameter_mm: float = 0.0
+    fibre_tensile_strength_MPa: float = 0.0
+    fibre_modulus_GPa: float = 0.0
+    # Fibre material-class id (index into categoricals.FIBRE_TYPES; 0 = none).
+    fibre_type_id: float = 0.0
+    # Measured nominal max grain sizes (mm); ``None`` means not reported, in
+    # which case the default bin distribution is used and the corresponding
+    # global feature is masked out. ``temperature_observed`` drives the same
+    # masking for curing temperature (which always carries a numeric fallback).
+    max_coarse_aggregate_size_mm: Optional[float] = None
+    max_fine_aggregate_size_mm: Optional[float] = None
+    temperature_observed: bool = True
+
+    def effective_aggregate_bins_mm(self) -> np.ndarray:
+        """Default bins rescaled so their top equals the measured max size.
+
+        Rescaling (rather than clipping) preserves the relative gradation shape
+        while respecting the reported nominal maximum coarse-aggregate size.
+        """
+
+        bins = np.asarray(self.aggregate_size_bins_mm, dtype=float)
+        if self.max_coarse_aggregate_size_mm and self.max_coarse_aggregate_size_mm > 0:
+            top = bins.max()
+            if top > 0:
+                bins = bins * (float(self.max_coarse_aggregate_size_mm) / top)
+        return bins
+
+    def sample_radius(self, rng: np.random.Generator) -> float:
+        bins = self.effective_aggregate_bins_mm()
+        weights = np.asarray(self.aggregate_size_fractions, dtype=float)
+        weights = weights / weights.sum()
+        idx = int(rng.choice(len(bins), p=weights))
+        return float(bins[idx] * 0.5 * rng.uniform(0.85, 1.0))
+
+
+@dataclass
+class MortarMixDesign:
+    """Mortar-scale mix description, distinct from :class:`MixDesign`.
+
+    The mortar scale operates on a small RVE (default 6 mm) so each micrograph
+    holds on the order of 10-30 nodes - enough for relational message passing
+    but cheap enough to train end-to-end alongside the mesoscale GNN.
+    """
+
+    rve_size_mm: float = 6.0
+    sand_size_bins_mm: Sequence[float] = (0.6, 1.0, 1.5, 2.0)
+    sand_size_fractions: Sequence[float] = (0.2, 0.3, 0.3, 0.2)
+    sand_volume_fraction: float = 0.35
+    water_to_binder_ratio: float = 0.45
+    cement_content_kg_m3: float = 360.0
+    scm_type_id: int = 1
+    scm_fraction: float = 0.20
+    admixture_dosage: float = 0.5
+    cement_chem: Tuple[float, float, float, float] = (64.0, 21.0, 5.0, 3.0)
+    # Extra cement oxides (MgO, SO3, Na2O-equivalent alkali) for the formerly
+    # placeholder paste slots; OPC-typical defaults when not measured.
+    cement_chem_ext: Tuple[float, float, float] = (2.0, 2.5, 0.6)
+    scm_chem: Tuple[float, float] = (4.0, 55.0)
+    # Extended SCM oxides (Al2O3, Fe2O3, MgO, LOI); generic-pozzolan defaults.
+    scm_chem_ext: Tuple[float, float, float, float] = (15.0, 5.0, 2.0, 3.0)
+    # Cement-type categorical id (index into categoricals.CEMENT_TYPES). Default
+    # 1.0 preserves the previous constant paste ``cement_type_id`` slot value.
+    cement_type_id: float = 1.0
+    curing_relative_humidity: float = 0.95
+    curing_temperature_C: float = 23.0
+    curing_age_days: float = 28.0
+    fibre_content_kg_m3: float = 0.0
+    fibre_length_mm: float = 0.0
+    fibre_diameter_mm: float = 0.0
+    fibre_tensile_strength_MPa: float = 0.0
+    fibre_modulus_GPa: float = 0.0
+    # Measured nominal max fine/sand grain size (mm); ``None`` means not
+    # reported. ``curing_temperature_observed`` masks the mortar curing-temp
+    # global when the source did not report it.
+    max_sand_size_mm: Optional[float] = None
+    curing_temperature_observed: bool = True
+
+    def effective_sand_bins_mm(self) -> np.ndarray:
+        """Default sand bins rescaled to the measured max fine grain size."""
+
+        bins = np.asarray(self.sand_size_bins_mm, dtype=float)
+        if self.max_sand_size_mm and self.max_sand_size_mm > 0:
+            top = bins.max()
+            if top > 0:
+                bins = bins * (float(self.max_sand_size_mm) / top)
+        return bins
+
+
+# ---------------------------------------------------------------------------
+# Output containers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class ConcreteGraph:
+    x: Tensor
+    x_mask: Tensor
+    edge_index: Tensor
+    edge_attr: Tensor
+    edge_attr_mask: Tensor
+    global_features: Tensor
+    global_mask: Tensor
+    pos: Tensor
+    node_type: Tensor
+    edge_type: Tensor
+
+
+@dataclass
+class MortarMicrograph:
+    x: Tensor
+    x_mask: Tensor
+    edge_index: Tensor
+    edge_attr: Tensor
+    node_type: Tensor
+    global_features: Tensor
+    global_mask: Tensor
+    sand_contact_density: float
+
+
+# ---------------------------------------------------------------------------
+# Geometry helpers
+# ---------------------------------------------------------------------------
+
+
+def _sample_aggregates(
+    mix: MixDesign,
+    rng: np.random.Generator,
+    max_attempts: int = 6000,
+) -> Tuple[np.ndarray, np.ndarray]:
+    rve_area = mix.rve_size_mm ** 2
+    target_area = mix.aggregate_volume_fraction * rve_area
+
+    centroids: List[np.ndarray] = []
+    radii: List[float] = []
+    placed_area = 0.0
+    attempts = 0
+    while placed_area < target_area and attempts < max_attempts:
+        attempts += 1
+        r = mix.sample_radius(rng)
+        c = rng.uniform(r, mix.rve_size_mm - r, size=2)
+        ok = True
+        for cj, rj in zip(centroids, radii):
+            if np.linalg.norm(c - cj) < (r + rj) * 1.05:
+                ok = False
+                break
+        if not ok:
+            continue
+        centroids.append(c)
+        radii.append(r)
+        placed_area += math.pi * r * r
+
+    if not centroids:
+        centroids = [np.array([mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0])]
+        radii = [mix.rve_size_mm * 0.05]
+    return np.asarray(centroids), np.asarray(radii)
+
+
+def _mortar_seed_grid(
+    mix: MixDesign,
+    agg_centroids: np.ndarray,
+    agg_radii: np.ndarray,
+    n_side: int = 8,
+) -> np.ndarray:
+    xs = np.linspace(0, mix.rve_size_mm, n_side + 2)[1:-1]
+    grid = np.array([(x, y) for x in xs for y in xs])
+    if len(agg_centroids) == 0:
+        return grid
+    diff = grid[:, None, :] - agg_centroids[None, :, :]
+    d2 = np.einsum("nki,nki->nk", diff, diff)
+    inside = np.any(d2 < (agg_radii * 1.05) ** 2, axis=1)
+    kept = grid[~inside]
+    if len(kept) == 0:
+        kept = np.array([[mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0]])
+    return kept
+
+
+def _voronoi_adjacency(
+    points: np.ndarray,
+    rve_size: float,
+    rng: np.random.Generator,
+    n_samples: int = 4000,
+) -> Dict[Tuple[int, int], int]:
+    """Approximate Voronoi adjacency by sampling.
+
+    For each random sample we look up its two nearest seeds; that pair shares
+    a Voronoi boundary.  The number of samples falling on each pair is a
+    monotone proxy for the shared boundary length.
+    """
+
+    if points.shape[0] < 2:
+        return {}
+    samples = rng.uniform(0.0, rve_size, size=(n_samples, 2))
+    diff = samples[:, None, :] - points[None, :, :]
+    d2 = np.einsum("nki,nki->nk", diff, diff)
+    nearest_two = np.argpartition(d2, kth=1, axis=1)[:, :2]
+    adjacency: Dict[Tuple[int, int], int] = {}
+    a = nearest_two[:, 0]
+    b = nearest_two[:, 1]
+    lo = np.minimum(a, b)
+    hi = np.maximum(a, b)
+    for i, j in zip(lo, hi):
+        if i == j:
+            continue
+        key = (int(i), int(j))
+        adjacency[key] = adjacency.get(key, 0) + 1
+    return adjacency
+
+
+def _voronoi_cell_fraction(
+    seeds: np.ndarray,
+    rve_size: float,
+    rng: np.random.Generator,
+    n_samples: int = 4000,
+) -> np.ndarray:
+    """Fraction of the RVE area nearest to each seed (Monte-Carlo Voronoi cells).
+
+    Returns an array summing to 1.0; used to apportion the non-particle area
+    across mortar / paste seed nodes so each carries a heterogeneous area share.
+    """
+
+    n = len(seeds)
+    if n == 0:
+        return np.zeros(0, dtype=np.float64)
+    if n == 1:
+        return np.ones(1, dtype=np.float64)
+    samples = rng.uniform(0.0, rve_size, size=(n_samples, 2))
+    diff = samples[:, None, :] - seeds[None, :, :]
+    d2 = np.einsum("nki,nki->nk", diff, diff)
+    nearest = np.argmin(d2, axis=1)
+    counts = np.bincount(nearest, minlength=n).astype(np.float64)
+    total = counts.sum()
+    if total <= 0:
+        return np.full(n, 1.0 / n, dtype=np.float64)
+    return counts / total
+
+
+# ---------------------------------------------------------------------------
+# Feature synthesis (concrete scale)
+# ---------------------------------------------------------------------------
+
+
+def _aggregate_feature_vector(
+    radius: float,
+    rng: np.random.Generator,
+    schema: SchemaSpec,
+    area_fraction: float = 0.0,
+) -> np.ndarray:
+    feats = np.zeros(len(schema.aggregate), dtype=np.float32)
+    feats[0] = 2.0 * radius
+    feats[1] = float(rng.uniform(0.6, 1.0))
+    feats[2] = float(rng.normal(120.0, 15.0))
+    feats[3] = float(rng.normal(20.0, 4.0))
+    feats[4] = float(rng.uniform(0.7, 1.0))
+    feats[5] = float(rng.uniform(0.0, 0.05))
+    feats[6] = float(rng.uniform(0.2, 2.0))
+    feats[7] = float(rng.uniform(0.4, 1.0))
+    feats[8] = float(rng.normal(2.65, 0.05))
+    feats[9] = float(rng.normal(1600.0, 50.0))
+    feats[schema.aggregate.index("area_fraction")] = float(area_fraction)
+    return feats
+
+
+def _mortar_placeholder_vector(
+    mix: MixDesign,
+    rng: np.random.Generator,
+    schema: SchemaSpec,
+    area_fraction: float = 0.0,
+) -> np.ndarray:
+    """A coarse prior for mortar features; overwritten by the sub-model."""
+
+    feats = np.zeros(len(schema.mortar), dtype=np.float32)
+    w = mix.water_to_binder_ratio
+    base = 85.0 - 65.0 * w
+    feats[0] = float(rng.normal(base, 5.0))
+    feats[1] = float(rng.normal(0.1 * base, 0.5))
+    feats[2] = float(rng.normal(0.16 * base, 0.5))
+    feats[3] = float(rng.normal(15.0 + 0.35 * base, 2.0))
+    feats[4] = float(rng.uniform(1e-12, 1e-10))
+    feats[5] = float(rng.uniform(1e-12, 1e-10))
+    feats[6] = float(rng.uniform(0.08, 0.20))
+    feats[7] = float(rng.uniform(1.0, 3.0))
+    # Fibre conditioning (slots 8-11): broadcast mix-level fibre attributes
+    # to every mortar node. Aspect ratio = length / diameter when both > 0.
+    feats[8] = float(mix.fibre_content_kg_m3)
+    if mix.fibre_diameter_mm > 0.0:
+        feats[9] = float(mix.fibre_length_mm / mix.fibre_diameter_mm)
+    feats[10] = float(mix.fibre_tensile_strength_MPa)
+    feats[11] = float(mix.fibre_modulus_GPa)
+    feats[schema.mortar.index("mortar_area_fraction")] = float(area_fraction)
+    feats[schema.mortar.index("fibre_type_id")] = float(mix.fibre_type_id)
+    return feats
+
+
+def _itz_placeholder_vector(
+    mix: MixDesign, rng: np.random.Generator, schema: SchemaSpec
+) -> np.ndarray:
+    feats = np.zeros(len(schema.itz), dtype=np.float32)
+    w = mix.water_to_binder_ratio
+    feats[0] = float(rng.uniform(5.0, 80.0))
+    # Porosity rises with w/b (kept physical) but spans the wider [0.05, 0.5].
+    feats[1] = float(np.clip(0.05 + 0.45 * w + rng.uniform(-0.03, 0.03), 0.05, 0.5))
+    feats[2] = float(rng.normal(20.0, 4.0))
+    feats[3] = float(rng.normal(15.0, 3.0))
+    feats[4] = float(rng.uniform(1e-11, 1e-9))
+    feats[5] = float(rng.uniform(1e-11, 1e-9))
+    return feats
+
+
+def _global_feature_vector(
+    mix: MixDesign, schema: SchemaSpec
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Return ``(values, observed_mask)`` for the concrete global vector.
+
+    Maskable descriptors (curing temperature, max coarse / fine grain size)
+    are flagged 0 in the mask when the mix did not report them, so the
+    MissingFeatureEncoder treats them as missing rather than as a real value.
+    """
+
+    values = {
+        "relative_humidity": mix.relative_humidity,
+        "temperature_C": mix.temperature_C,
+        "curing_age_days": mix.curing_age_days,
+        "water_to_binder_ratio": mix.water_to_binder_ratio,
+        "cement_content_kg_m3": mix.cement_content_kg_m3,
+        "scm_fraction": mix.scm_fraction,
+        "aggregate_volume_fraction": mix.aggregate_volume_fraction,
+        "fibre_content_kg_m3": mix.fibre_content_kg_m3,
+        "max_coarse_aggregate_size_mm": mix.max_coarse_aggregate_size_mm or 0.0,
+        "max_fine_aggregate_size_mm": mix.max_fine_aggregate_size_mm or 0.0,
+    }
+    observed = {
+        "temperature_C": mix.temperature_observed,
+        "max_coarse_aggregate_size_mm": mix.max_coarse_aggregate_size_mm is not None,
+        "max_fine_aggregate_size_mm": mix.max_fine_aggregate_size_mm is not None,
+    }
+    g = np.zeros(len(schema.glob), dtype=np.float32)
+    mask = np.zeros(len(schema.glob), dtype=np.float32)
+    for i, name in enumerate(schema.glob):
+        if name in values:
+            g[i] = values[name]
+            mask[i] = 1.0 if observed.get(name, True) else 0.0
+    return g, mask
+
+
+# ---------------------------------------------------------------------------
+# Concrete graph generator
+# ---------------------------------------------------------------------------
+
+
+def generate_concrete_graph(
+    mix: MixDesign,
+    schema: SchemaSpec = DEFAULT_SCHEMA,
+    seed: Optional[int] = None,
+    missing_rate: float = 0.10,
+) -> ConcreteGraph:
+    rng = np.random.default_rng(seed)
+    agg_xy, agg_r = _sample_aggregates(mix, rng)
+    mortar_xy = _mortar_seed_grid(mix, agg_xy, agg_r)
+
+    n_agg = len(agg_xy)
+    n_mortar = len(mortar_xy)
+    all_points = np.concatenate([agg_xy, mortar_xy], axis=0)
+    adjacency = _voronoi_adjacency(all_points, rve_size=mix.rve_size_mm, rng=rng)
+
+    # Augment mortar-mortar edges via k-nearest neighbours (Voronoi cells of
+    # mortar seeds are frequently interrupted by aggregate cells).
+    if n_mortar > 1:
+        k = min(4, n_mortar - 1)
+        d2 = np.sum((mortar_xy[:, None, :] - mortar_xy[None, :, :]) ** 2, axis=-1)
+        np.fill_diagonal(d2, np.inf)
+        nbr = np.argpartition(d2, kth=k - 1, axis=1)[:, :k]
+        for i in range(n_mortar):
+            for j in nbr[i]:
+                a = n_agg + int(i)
+                b = n_agg + int(j)
+                if a == b:
+                    continue
+                key = (min(a, b), max(a, b))
+                adjacency.setdefault(key, 1)
+
+    pad_node = schema.node_pad_dim
+    pad_edge = schema.edge_pad_dim
+    n_total = n_agg + n_mortar
+
+    x = np.zeros((n_total, pad_node), dtype=np.float32)
+    x_mask = np.zeros((n_total, pad_node), dtype=np.float32)
+
+    # Per-node area fractions. Aggregates use exact disk area; mortar seeds split
+    # the remaining (non-aggregate) area by their Monte-Carlo Voronoi cell share.
+    rve_area = float(mix.rve_size_mm ** 2)
+    agg_area_frac = (np.pi * np.asarray(agg_r, dtype=np.float64) ** 2) / rve_area
+    mortar_total_frac = max(0.0, 1.0 - float(agg_area_frac.sum()))
+    mortar_cell = _voronoi_cell_fraction(mortar_xy, mix.rve_size_mm, rng)
+    mortar_area_frac = mortar_cell * mortar_total_frac
+
+    for i in range(n_agg):
+        feats = _aggregate_feature_vector(
+            agg_r[i], rng, schema, area_fraction=float(agg_area_frac[i])
+        )
+        agg_dim = len(schema.aggregate)
+        x[i, :agg_dim] = feats
+        x_mask[i, :agg_dim] = 1.0
+
+    for j in range(n_mortar):
+        feats = _mortar_placeholder_vector(
+            mix, rng, schema, area_fraction=float(mortar_area_frac[j])
+        )
+        mortar_dim = len(schema.mortar)
+        x[n_agg + j, :mortar_dim] = feats
+        x_mask[n_agg + j, :mortar_dim] = 1.0
+
+    node_type = np.concatenate(
+        [
+            np.full(n_agg, NODE_TYPE_AGGREGATE, dtype=np.int64),
+            np.full(n_mortar, NODE_TYPE_MORTAR, dtype=np.int64),
+        ]
+    )
+
+    # Edges
+    edge_index: List[List[int]] = []
+    edge_attr: List[np.ndarray] = []
+    edge_mask: List[np.ndarray] = []
+    edge_type: List[int] = []
+
+    for (a, b), shared in adjacency.items():
+        type_a = NODE_TYPE_AGGREGATE if a < n_agg else NODE_TYPE_MORTAR
+        type_b = NODE_TYPE_AGGREGATE if b < n_agg else NODE_TYPE_MORTAR
+        pa = all_points[a]
+        pb = all_points[b]
+        dist = float(np.linalg.norm(pa - pb))
+        boundary = float(shared) * (mix.rve_size_mm / 60.0)
+
+        feats = np.zeros(pad_edge, dtype=np.float32)
+        mask = np.zeros(pad_edge, dtype=np.float32)
+
+        if type_a != type_b:
+            etype = EDGE_TYPE_ITZ
+            itz_vals = _itz_placeholder_vector(mix, rng, schema)
+            feats[: len(itz_vals)] = itz_vals
+            mask[: len(itz_vals)] = 1.0
+        elif type_a == NODE_TYPE_MORTAR:
+            etype = EDGE_TYPE_MORTAR_MORTAR
+            feats[0] = dist
+            feats[1] = boundary
+            feats[2] = float(pb[0] - pa[0])
+            feats[3] = float(pb[1] - pa[1])
+            mask[:4] = 1.0
+        else:
+            ai = a
+            bi = b
+            surface_gap = max(0.0, dist - float(agg_r[ai]) - float(agg_r[bi]))
+            if surface_gap > 1.5 * (mix.rve_size_mm / 60.0):
+                continue
+            etype = EDGE_TYPE_AGGREGATE_AGGREGATE
+            feats[0] = surface_gap
+            feats[1] = dist
+            mask[:2] = 1.0
+
+        for src, dst in ((a, b), (b, a)):
+            edge_index.append([src, dst])
+            edge_attr.append(feats.copy())
+            edge_mask.append(mask.copy())
+            edge_type.append(etype)
+
+    if not edge_index:
+        edge_index = [[0, 0]]
+        edge_attr = [np.zeros(pad_edge, dtype=np.float32)]
+        edge_mask = [np.zeros(pad_edge, dtype=np.float32)]
+        edge_type = [EDGE_TYPE_ITZ]
+
+    edge_index_t = torch.tensor(np.asarray(edge_index).T, dtype=torch.long)
+    edge_attr_t = torch.tensor(np.stack(edge_attr), dtype=torch.float32)
+    edge_mask_t = torch.tensor(np.stack(edge_mask), dtype=torch.float32)
+    edge_type_t = torch.tensor(edge_type, dtype=torch.long)
+    pos_t = torch.tensor(all_points, dtype=torch.float32)
+    x_t = torch.tensor(x, dtype=torch.float32)
+    x_mask_t = torch.tensor(x_mask, dtype=torch.float32)
+    node_type_t = torch.tensor(node_type, dtype=torch.long)
+
+    if missing_rate > 0.0:
+        protected = _protected_node_columns(schema)
+        drop = (torch.rand_like(x_mask_t) < missing_rate) & (x_mask_t > 0)
+        drop[:, protected] = False
+        x_mask_t = x_mask_t * (~drop).float()
+        x_t = x_t * x_mask_t
+
+    g, g_observed = _global_feature_vector(mix, schema)
+    g_t = torch.tensor(g, dtype=torch.float32)
+    g_mask = torch.tensor(g_observed, dtype=torch.float32)
+    g_t = g_t * g_mask  # zero out unreported descriptors
+    if missing_rate > 0.0:
+        drop = torch.rand_like(g_t) < missing_rate
+        g_mask = g_mask * (~drop).float()
+        g_t = g_t * g_mask
+
+    return ConcreteGraph(
+        x=x_t,
+        x_mask=x_mask_t,
+        edge_index=edge_index_t,
+        edge_attr=edge_attr_t,
+        edge_attr_mask=edge_mask_t,
+        global_features=g_t,
+        global_mask=g_mask,
+        pos=pos_t,
+        node_type=node_type_t,
+        edge_type=edge_type_t,
+    )
+
+
+def _protected_node_columns(schema: SchemaSpec) -> List[int]:
+    """Indices of features that should never be hidden by random missingness."""
+
+    return [0]  # protect the first slot (e.g. aggregate size / strength prior)
+
+
+# ---------------------------------------------------------------------------
+# Mortar micrograph generator
+# ---------------------------------------------------------------------------
+
+
+def _sample_sand(
+    mix: MortarMixDesign,
+    rng: np.random.Generator,
+    max_attempts: int = 800,
+) -> Tuple[np.ndarray, np.ndarray]:
+    bins = mix.effective_sand_bins_mm()
+    weights = np.asarray(mix.sand_size_fractions, dtype=float)
+    weights = weights / weights.sum()
+
+    rve_area = mix.rve_size_mm ** 2
+    target_area = mix.sand_volume_fraction * rve_area
+
+    centroids: List[np.ndarray] = []
+    radii: List[float] = []
+    placed = 0.0
+    attempts = 0
+    while placed < target_area and attempts < max_attempts:
+        attempts += 1
+        idx = int(rng.choice(len(bins), p=weights))
+        r = float(bins[idx] * 0.5 * rng.uniform(0.85, 1.0))
+        c = rng.uniform(r, mix.rve_size_mm - r, size=2)
+        if any(np.linalg.norm(c - cj) < (r + rj) * 1.05 for cj, rj in zip(centroids, radii)):
+            continue
+        centroids.append(c)
+        radii.append(r)
+        placed += math.pi * r * r
+
+    if not centroids:
+        centroids = [np.array([mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0])]
+        radii = [0.5]
+    return np.asarray(centroids), np.asarray(radii)
+
+
+def _paste_grid(mix: MortarMixDesign, sand_xy: np.ndarray, sand_r: np.ndarray) -> np.ndarray:
+    n_side = 5
+    xs = np.linspace(0, mix.rve_size_mm, n_side + 2)[1:-1]
+    grid = np.array([(x, y) for x in xs for y in xs])
+    if len(sand_xy) == 0:
+        return grid
+    diff = grid[:, None, :] - sand_xy[None, :, :]
+    d2 = np.einsum("nki,nki->nk", diff, diff)
+    inside = np.any(d2 < (sand_r * 1.05) ** 2, axis=1)
+    kept = grid[~inside]
+    if len(kept) == 0:
+        kept = np.array([[mix.rve_size_mm / 2.0, mix.rve_size_mm / 2.0]])
+    return kept
+
+
+def _sand_feature_vector(
+    radius: float,
+    rng: np.random.Generator,
+    schema: SchemaSpec,
+    area_fraction: float = 0.0,
+) -> np.ndarray:
+    v = np.zeros(len(schema.sand), dtype=np.float32)
+    v[0] = 2.0 * radius
+    v[1] = float(rng.uniform(2.0, 3.5))
+    v[2] = float(rng.uniform(2.4, 3.2))
+    v[3] = float(rng.uniform(0.5, 2.0))
+    v[4] = float(rng.normal(80.0, 10.0))
+    v[5] = float(rng.normal(70.0, 8.0))
+    v[schema.sand.index("sand_area_fraction")] = float(area_fraction)
+    return v
+
+
+def _paste_feature_vector(
+    mix: MortarMixDesign, schema: SchemaSpec, area_fraction: float = 0.0
+) -> np.ndarray:
+    v = np.zeros(len(schema.paste), dtype=np.float32)
+    v[0] = float(mix.cement_type_id)
+    v[1:5] = np.asarray(mix.cement_chem, dtype=np.float32)
+    v[5] = mix.cement_content_kg_m3
+    v[6] = float(mix.scm_type_id)
+    v[7:9] = np.asarray(mix.scm_chem, dtype=np.float32)
+    v[9] = mix.scm_fraction * mix.cement_content_kg_m3
+    v[10] = mix.water_to_binder_ratio
+    v[11] = mix.admixture_dosage
+    v[schema.paste.index("paste_area_fraction")] = float(area_fraction)
+    ext = np.asarray(mix.cement_chem_ext, dtype=np.float32)
+    v[schema.paste.index("cement_MgO_pct")] = ext[0]
+    v[schema.paste.index("cement_SO3_pct")] = ext[1]
+    v[schema.paste.index("cement_alkali_pct")] = ext[2]
+    sext = np.asarray(mix.scm_chem_ext, dtype=np.float32)
+    v[schema.paste.index("scm_Al2O3_pct")] = sext[0]
+    v[schema.paste.index("scm_Fe2O3_pct")] = sext[1]
+    v[schema.paste.index("scm_MgO_pct")] = sext[2]
+    v[schema.paste.index("scm_LOI_pct")] = sext[3]
+    return v
+
+
+def _mortar_global_vector(
+    mix: MortarMixDesign, schema: SchemaSpec
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Return ``(values, observed_mask)`` for the mortar global vector."""
+
+    values = {
+        "mortar_water_to_binder_ratio": mix.water_to_binder_ratio,
+        "mortar_cement_content_kg_m3": mix.cement_content_kg_m3,
+        "mortar_scm_fraction": mix.scm_fraction,
+        "mortar_admixture_dosage": mix.admixture_dosage,
+        "mortar_curing_relative_humidity": mix.curing_relative_humidity,
+        "mortar_curing_temperature_C": mix.curing_temperature_C,
+        "mortar_curing_age_days": mix.curing_age_days,
+        "mortar_fibre_content_kg_m3": mix.fibre_content_kg_m3,
+    }
+    observed = {"mortar_curing_temperature_C": mix.curing_temperature_observed}
+    g = np.zeros(len(schema.mortar_global), dtype=np.float32)
+    mask = np.zeros(len(schema.mortar_global), dtype=np.float32)
+    for i, name in enumerate(schema.mortar_global):
+        if name in values:
+            g[i] = values[name]
+            mask[i] = 1.0 if observed.get(name, True) else 0.0
+    return g, mask
+
+
+def generate_mortar_graph(
+    mix: MortarMixDesign,
+    schema: SchemaSpec = DEFAULT_SCHEMA,
+    seed: Optional[int] = None,
+) -> MortarMicrograph:
+    rng = np.random.default_rng(seed)
+    sand_xy, sand_r = _sample_sand(mix, rng)
+    paste_xy = _paste_grid(mix, sand_xy, sand_r)
+
+    n_sand = len(sand_xy)
+    n_paste = len(paste_xy)
+    pts = np.concatenate([sand_xy, paste_xy], axis=0)
+    adjacency = _voronoi_adjacency(pts, rve_size=mix.rve_size_mm, rng=rng, n_samples=3000)
+
+    pad = schema.mortar_node_pad_dim
+    n_total = n_sand + n_paste
+    x = np.zeros((n_total, pad), dtype=np.float32)
+    mask = np.zeros((n_total, pad), dtype=np.float32)
+    node_type = np.zeros(n_total, dtype=np.int64)
+
+    # Sand uses exact disk area; paste seeds split the remaining area by their
+    # Monte-Carlo Voronoi cell share.
+    mortar_rve_area = float(mix.rve_size_mm ** 2)
+    sand_area_frac = (np.pi * np.asarray(sand_r, dtype=np.float64) ** 2) / mortar_rve_area
+    paste_total_frac = max(0.0, 1.0 - float(sand_area_frac.sum()))
+    paste_cell = _voronoi_cell_fraction(paste_xy, mix.rve_size_mm, rng, n_samples=3000)
+    paste_area_frac = paste_cell * paste_total_frac
+
+    for i in range(n_sand):
+        v = _sand_feature_vector(
+            sand_r[i], rng, schema, area_fraction=float(sand_area_frac[i])
+        )
+        x[i, : len(v)] = v
+        mask[i, : len(v)] = 1.0
+        node_type[i] = MORTAR_NODE_TYPE_SAND
+
+    for j in range(n_paste):
+        v = _paste_feature_vector(mix, schema, area_fraction=float(paste_area_frac[j]))
+        x[n_sand + j, : len(v)] = v
+        mask[n_sand + j, : len(v)] = 1.0
+        node_type[n_sand + j] = MORTAR_NODE_TYPE_PASTE
+
+    edge_index: List[List[int]] = []
+    edge_attr: List[List[float]] = []
+    sand_sand_pairs = 0
+    for (a, b), shared in adjacency.items():
+        d = float(np.linalg.norm(pts[a] - pts[b]))
+        same_phase = 1.0 if node_type[a] == node_type[b] else 0.0
+        if node_type[a] == MORTAR_NODE_TYPE_SAND and node_type[b] == MORTAR_NODE_TYPE_SAND:
+            sand_sand_pairs += 1
+        for src, dst in ((a, b), (b, a)):
+            edge_index.append([src, dst])
+            edge_attr.append([d, same_phase])
+
+    if not edge_index:
+        edge_index = [[0, 0]]
+        edge_attr = [[0.0, 0.0]]
+
+    contact_density = min(1.0, (sand_sand_pairs / max(1, n_sand)) / 4.0)
+
+    g, g_observed = _mortar_global_vector(mix, schema)
+    g_t = torch.tensor(g, dtype=torch.float32)
+    g_mask = torch.tensor(g_observed, dtype=torch.float32)
+    g_t = g_t * g_mask  # zero out unreported descriptors
+
+    return MortarMicrograph(
+        x=torch.tensor(x, dtype=torch.float32),
+        x_mask=torch.tensor(mask, dtype=torch.float32),
+        edge_index=torch.tensor(np.asarray(edge_index).T, dtype=torch.long),
+        edge_attr=torch.tensor(edge_attr, dtype=torch.float32),
+        node_type=torch.tensor(node_type, dtype=torch.long),
+        global_features=g_t,
+        global_mask=g_mask,
+        sand_contact_density=float(contact_density),
+    )

concrete_gnn/ground_truth.py

@@ -0,0 +1,181 @@
+"""Synthetic physics that produces concrete-level targets dependent on graph topology.
+
+The targets generated here intentionally depend on (a) microscale properties
+that are *relational* with respect to the mortar micrograph (sand-sand
+contact density), and (b) the *realised* mesoscale topology (aggregate node
+fraction, ITZ edge fraction, aggregate-aggregate edge fraction).  Because
+both quantities vary per RVE even at a fixed mix design, a model that sees
+only global mix features cannot reproduce that per-RVE variance, while a
+hierarchical GNN can.  This makes the question "does the architecture
+matter?" empirically testable.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+from .graph_generator import ConcreteGraph, MixDesign, MortarMicrograph
+from .schema import (
+    EDGE_TYPE_AGGREGATE_AGGREGATE,
+    EDGE_TYPE_ITZ,
+    NODE_TYPE_AGGREGATE,
+)
+
+
+@dataclass
+class GroundTruthProps:
+    mortar_elastic: float
+    mortar_strength: float
+    mortar_permeability: float
+    mortar_porosity: float
+    mortar_creep: float
+    mortar_tensile: float
+    mortar_flexural: float
+    mortar_diffusivity: float
+    itz_strength: float
+    itz_permeability: float
+    itz_thickness: float
+    itz_porosity: float
+    itz_elastic: float
+    itz_diffusivity: float
+
+
+def true_micro_properties(mix: MixDesign, sand_contact_density: float) -> GroundTruthProps:
+    """Closed-form true mortar and ITZ properties.
+
+    ``sand_contact_density`` is a relational quantity (sand-sand contacts per
+    sand particle) recoverable only from micrograph connectivity (sum-based
+    message passing), not from pooled node features.
+    """
+
+    w = mix.water_to_binder_ratio
+    cem = mix.cement_content_kg_m3 / 360.0
+    age = min(1.25, math.log1p(mix.curing_age_days) / math.log1p(28.0))
+    base = 85.0 - 65.0 * w
+
+    mortar_elastic = (15.0 + 0.35 * base) * (0.8 + 0.4 * sand_contact_density)
+    mortar_strength = base * age * (0.95 + 0.1 * cem)
+    mortar_permeability = 1.0e-12 * (1.0 + 3.0 * w)
+    mortar_tensile = 0.1 * mortar_strength
+    mortar_flexural = 1.8 * mortar_tensile
+    mortar_diffusivity = 1.0e-11 * (1.0 + 4.0 * w)
+    mortar_porosity = 0.08 + 0.25 * w
+    mortar_creep = 1.0 + 2.0 * w
+
+    itz_strength = 0.70 * base * (0.9 + 0.1 * age)
+    itz_permeability = 2.0e-12 * (1.0 + 4.0 * w)
+    itz_thickness = 20.0 + 30.0 * w
+    itz_porosity = 0.12 + 0.28 * w
+    itz_elastic = 12.0 + 10.0 * (1.0 - w)
+    itz_diffusivity = 2.0e-11 * (1.0 + 5.0 * w)
+
+    return GroundTruthProps(
+        mortar_elastic=mortar_elastic,
+        mortar_strength=mortar_strength,
+        mortar_permeability=mortar_permeability,
+        mortar_porosity=mortar_porosity,
+        mortar_creep=mortar_creep,
+        mortar_tensile=mortar_tensile,
+        mortar_flexural=mortar_flexural,
+        mortar_diffusivity=mortar_diffusivity,
+        itz_strength=itz_strength,
+        itz_permeability=itz_permeability,
+        itz_thickness=itz_thickness,
+        itz_porosity=itz_porosity,
+        itz_elastic=itz_elastic,
+        itz_diffusivity=itz_diffusivity,
+    )
+
+
+def homogenize_concrete_targets(
+    mix: MixDesign,
+    props: GroundTruthProps,
+    agg_node_fraction: float,
+    itz_edge_fraction: float,
+    aa_edge_fraction: float,
+) -> Tensor:
+    """Combine microscale properties with realised topology into concrete targets.
+
+    Output order matches ``schema.CONCRETE_TARGETS``.
+    """
+
+    w = mix.water_to_binder_ratio
+    V = mix.aggregate_volume_fraction
+    e_agg = 60.0
+
+    e_voigt = V * e_agg + (1.0 - V) * props.mortar_elastic
+    e_reuss = 1.0 / (V / e_agg + (1.0 - V) / props.mortar_elastic)
+    elastic = 0.5 * (e_voigt + e_reuss) * (1.0 - 0.12 * itz_edge_fraction)
+
+    itz_quality = props.itz_strength / 60.0
+    compressive = (
+        props.mortar_strength
+        * (1.0 - 0.2 * V)
+        * (0.8 + 0.3 * itz_quality)
+        * (1.0 - 0.15 * aa_edge_fraction)
+        * (1.0 + 0.1 * (agg_node_fraction - 0.3))
+    )
+    compressive = max(5.0, compressive)
+    tensile = 0.1 * compressive * (0.85 + 0.3 * itz_quality)
+    flexural = 0.15 * compressive + 0.6
+
+    permeability = props.mortar_permeability * (1.0 - itz_edge_fraction) + (
+        props.itz_permeability * itz_edge_fraction * (1.0 + 2.0 * aa_edge_fraction)
+    )
+    diffusivity = (0.5 + 0.5 * w) * 1.0e-11 + 5.0e-11 * itz_edge_fraction * (
+        1.0 + aa_edge_fraction
+    )
+
+    return torch.tensor(
+        [compressive, tensile, flexural, elastic, permeability, diffusivity],
+        dtype=torch.float32,
+    )
+
+
+def mortar_target_vector(props: GroundTruthProps) -> Tensor:
+    """Mortar sub-model target vector (order matches ``schema.MORTAR_TARGETS``)."""
+
+    return torch.tensor(
+        [
+            props.mortar_strength,
+            props.mortar_tensile,
+            props.mortar_flexural,
+            props.mortar_elastic,
+            props.mortar_permeability,
+            props.mortar_diffusivity,
+            props.mortar_porosity,
+            props.mortar_creep,
+        ],
+        dtype=torch.float32,
+    )
+
+
+def itz_target_vector(props: GroundTruthProps) -> Tensor:
+    """ITZ sub-model target vector (order matches ``schema.ITZ_TARGETS``)."""
+
+    return torch.tensor(
+        [
+            props.itz_thickness,
+            props.itz_porosity,
+            props.itz_strength,
+            props.itz_elastic,
+            props.itz_permeability,
+            props.itz_diffusivity,
+        ],
+        dtype=torch.float32,
+    )
+
+
+def topology_statistics(graph: ConcreteGraph) -> tuple:
+    """Return (agg_node_fraction, itz_edge_fraction, aa_edge_fraction)."""
+
+    agg_node_fraction = (graph.node_type == NODE_TYPE_AGGREGATE).float().mean().item()
+    itz_edge_fraction = (graph.edge_type == EDGE_TYPE_ITZ).float().mean().item()
+    aa_edge_fraction = (
+        (graph.edge_type == EDGE_TYPE_AGGREGATE_AGGREGATE).float().mean().item()
+    )
+    return agg_node_fraction, itz_edge_fraction, aa_edge_fraction

concrete_gnn/missing_features.py

@@ -0,0 +1,159 @@
+"""Missing-feature handling.
+
+For every input feature we store a binary mask (1 = observed, 0 = missing).
+Observed values are first standardized per channel (z-score using frozen stats
+fit on the training set), so a small-magnitude but informative feature (e.g.
+water/binder ratio ~0.4, admixture dosage ~0.01) is not swamped by a large one
+(cement content ~500) in the input projection.  Missing entries are then
+replaced by a learnable per-channel embedding, observed entries pick up a
+learnable per-channel bias, and the value-plus-mask pair is projected through a
+small Linear + LayerNorm + SiLU stack.  The mask is fed in both as a
+multiplicative gate and as a side channel so the network can still learn whether
+a value was measured, imputed, or structurally absent.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+
+
+def masked_feature_stats(
+    x: torch.Tensor, mask: torch.Tensor, eps: float = 1e-5
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-feature (mean, std) over observed entries only (``mask > 0``).
+
+    Columns with no observed entries (or with near-constant values) return
+    ``mean=0, std=1`` so standardization is a no-op there.
+    """
+
+    mask = (mask > 0).to(x.dtype)
+    x = torch.nan_to_num(x, nan=0.0)
+    count = mask.sum(dim=0)
+    safe = count.clamp(min=1.0)
+    mean = (x * mask).sum(dim=0) / safe
+    var = (((x - mean) ** 2) * mask).sum(dim=0) / safe
+    std = var.clamp(min=0.0).sqrt()
+    has = count > 0
+    mean = torch.where(has, mean, torch.zeros_like(mean))
+    std = torch.where(has & (std > eps), std, torch.ones_like(std))
+    return mean, std
+
+
+def masked_feature_stats_by_type(
+    x: torch.Tensor,
+    mask: torch.Tensor,
+    type_index: torch.Tensor,
+    num_types: int,
+    eps: float = 1e-5,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-(type, feature) masked stats, shape ``(num_types, D)`` each.
+
+    Used where one encoder multiplexes several feature schemas into the same
+    columns (e.g. aggregate vs mortar nodes), so each row type is standardized
+    against its own statistics. Types with no rows fall back to ``mean=0, std=1``.
+    """
+
+    in_dim = x.size(1)
+    means = torch.zeros(num_types, in_dim, dtype=x.dtype)
+    stds = torch.ones(num_types, in_dim, dtype=x.dtype)
+    for t in range(num_types):
+        sel = type_index == t
+        if bool(sel.any()):
+            means[t], stds[t] = masked_feature_stats(x[sel], mask[sel], eps)
+    return means, stds
+
+
+class MissingFeatureEncoder(nn.Module):
+    """Encode a possibly-missing feature vector into a dense hidden representation.
+
+    Parameters
+    ----------
+    in_dim: int
+        Number of raw input channels.
+    out_dim: int
+        Output hidden dimension.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int, num_stat_groups: int = 1):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.num_stat_groups = num_stat_groups
+
+        self.missing_embedding = nn.Parameter(torch.zeros(in_dim))
+        nn.init.normal_(self.missing_embedding, std=0.02)
+        self.observed_bias = nn.Parameter(torch.zeros(in_dim))
+
+        # Frozen per-feature input standardization (z-score). Default identity;
+        # call ``set_feature_stats`` with training-set stats to activate. Stored
+        # as buffers (shape ``(num_stat_groups, in_dim)``) so they travel with
+        # .to(device) and persist in checkpoints. ``num_stat_groups`` > 1 lets one
+        # encoder hold per-node-type / per-edge-type stats, selected per row via a
+        # ``type_index`` in ``forward`` (the columns mean different features for
+        # different types, so a single shared z-score would blend them).
+        self.register_buffer("feat_mean", torch.zeros(num_stat_groups, in_dim))
+        self.register_buffer("feat_std", torch.ones(num_stat_groups, in_dim))
+
+        self.proj = nn.Sequential(
+            nn.Linear(in_dim * 2, out_dim),
+            nn.LayerNorm(out_dim),
+            nn.SiLU(),
+        )
+
+    @torch.no_grad()
+    def set_feature_stats(
+        self, mean: torch.Tensor, std: torch.Tensor, eps: float = 1e-5
+    ) -> None:
+        """Freeze standardization stats (from the train split).
+
+        ``mean`` / ``std`` are ``(in_dim,)`` for a single group or
+        ``(num_stat_groups, in_dim)`` for per-type stats.
+        """
+
+        mean = torch.as_tensor(mean, dtype=self.feat_mean.dtype, device=self.feat_mean.device)
+        std = torch.as_tensor(std, dtype=self.feat_std.dtype, device=self.feat_std.device)
+        if mean.dim() == 1:
+            mean = mean.unsqueeze(0)
+        if std.dim() == 1:
+            std = std.unsqueeze(0)
+        std = torch.where(std > eps, std, torch.ones_like(std))
+        self.feat_mean.copy_(mean.expand_as(self.feat_mean))
+        self.feat_std.copy_(std.expand_as(self.feat_std))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        type_index: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if mask is None:
+            mask = torch.ones_like(x)
+        x = torch.nan_to_num(x, nan=0.0)
+        if type_index is None:
+            mean, std = self.feat_mean[0], self.feat_std[0]
+        else:
+            mean = self.feat_mean.index_select(0, type_index)
+            std = self.feat_std.index_select(0, type_index)
+        x = (x - mean) / std
+        imputed = x * mask + self.missing_embedding * (1.0 - mask)
+        imputed = imputed + self.observed_bias * mask
+        joined = torch.cat([imputed, mask], dim=-1)
+        return self.proj(joined)
+
+
+def apply_random_missingness(
+    x: torch.Tensor,
+    rate: float = 0.15,
+    generator: Optional[torch.Generator] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Randomly drop entries of ``x`` and return ``(x_masked, mask)``."""
+
+    if generator is None:
+        mask = (torch.rand_like(x) > rate).float()
+    else:
+        rnd = torch.rand(x.shape, generator=generator, device=x.device)
+        mask = (rnd > rate).float()
+    return x * mask, mask

concrete_gnn/models/__init__.py

@@ -0,0 +1,28 @@
+"""Model package exports."""
+
+from .concrete_gnn import ConcreteMesoGNN
+from .hierarchical import IntegratedMultiscaleModel, TabularANN
+from .itz_model import ITZSubModel
+from .layers import HAS_PYG, global_add_pool, global_mean_pool, make_gine_conv, make_nn_conv
+from .mortar_gnn import MortarSubModel
+from .outputs import (
+    transform_concrete_outputs,
+    transform_itz_outputs,
+    transform_mortar_outputs,
+)
+
+__all__ = [
+    "ConcreteMesoGNN",
+    "HAS_PYG",
+    "ITZSubModel",
+    "IntegratedMultiscaleModel",
+    "MortarSubModel",
+    "TabularANN",
+    "global_add_pool",
+    "global_mean_pool",
+    "make_gine_conv",
+    "make_nn_conv",
+    "transform_concrete_outputs",
+    "transform_itz_outputs",
+    "transform_mortar_outputs",
+]

concrete_gnn/models/concrete_gnn.py

@@ -0,0 +1,100 @@
+"""Mesoscale concrete GNN.
+
+Edge-conditioned message passing on the concrete RVE graph.  Each layer uses
+``NNConv`` (PyG when available, otherwise a native fallback) to let the edge
+attributes generate per-edge weight matrices - the right inductive bias when
+ITZ vs. mortar-mortar vs. aggregate-aggregate interactions are physically
+distinct.  Node-type and edge-type embeddings carry the discrete semantics
+that the raw feature one-hots cannot fully express.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+from ..missing_features import MissingFeatureEncoder
+from ..schema import DEFAULT_SCHEMA, SchemaSpec
+from .layers import global_mean_pool, global_min_pool, make_nn_conv
+from .outputs import transform_concrete_outputs
+
+
+class ConcreteMesoGNN(nn.Module):
+    def __init__(
+        self,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        hidden_dim: int = 96,
+        num_layers: int = 2,  # see receptive-field ablation; depth >2 gives no gain
+        pool: str = "mean",  # graph readout: "mean" or "min" (weakest-link)
+    ):
+        super().__init__()
+        self.schema = schema
+        self.hidden_dim = hidden_dim
+        if pool not in ("mean", "min"):
+            raise ValueError(f"pool must be 'mean' or 'min', got {pool!r}")
+        self.pool = pool
+
+        self.node_encoder = MissingFeatureEncoder(
+            schema.node_pad_dim, hidden_dim, num_stat_groups=schema.num_node_types
+        )
+        self.node_type_embed = nn.Embedding(schema.num_node_types, hidden_dim)
+
+        self.edge_encoder = MissingFeatureEncoder(
+            schema.edge_pad_dim, hidden_dim, num_stat_groups=schema.num_edge_types
+        )
+        self.edge_type_embed = nn.Embedding(schema.num_edge_types, hidden_dim)
+
+        self.global_encoder = MissingFeatureEncoder(schema.global_dim, hidden_dim)
+
+        self.layers = nn.ModuleList()
+        self.norms = nn.ModuleList()
+        for _ in range(num_layers):
+            edge_nn = nn.Sequential(
+                nn.Linear(hidden_dim, hidden_dim),
+                nn.SiLU(),
+                nn.Linear(hidden_dim, hidden_dim * hidden_dim),
+            )
+            self.layers.append(
+                make_nn_conv(hidden_dim, hidden_dim, edge_nn, aggr="mean")
+            )
+            self.norms.append(nn.LayerNorm(hidden_dim))
+
+        self.feature_dim = hidden_dim * 2
+        self.head = nn.Sequential(
+            nn.Linear(self.feature_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, len(schema.concrete_targets)),
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        x_mask: Tensor,
+        edge_index: Tensor,
+        edge_attr: Tensor,
+        edge_attr_mask: Tensor,
+        node_type: Tensor,
+        edge_type: Tensor,
+        global_features: Tensor,
+        global_mask: Tensor,
+        batch: Tensor,
+        num_graphs: int,
+        return_features: bool = False,
+    ):
+        h = self.node_encoder(x, x_mask, node_type) + self.node_type_embed(node_type)
+        e = self.edge_encoder(edge_attr, edge_attr_mask, edge_type) + self.edge_type_embed(edge_type)
+
+        for conv, norm in zip(self.layers, self.norms):
+            h_new = conv(h, edge_index, e)
+            h = norm(h + h_new)
+            h = torch.nn.functional.silu(h)
+
+        pool_fn = global_min_pool if self.pool == "min" else global_mean_pool
+        pooled = pool_fn(h, batch, num_graphs=num_graphs)
+        g = self.global_encoder(global_features, global_mask)
+        features = torch.cat([pooled, g], dim=-1)
+        raw = self.head(features)
+        concrete_pred = transform_concrete_outputs(raw)
+        if return_features:
+            return concrete_pred, features
+        return concrete_pred

concrete_gnn/models/hierarchical.py

@@ -0,0 +1,521 @@
+"""Integrated multiscale model + tabular ANN baseline.
+
+The hierarchical model wires the mortar GNN and the ITZ MLP into the
+mesoscale GNN via *differentiable* feature injection: mortar predictions
+overwrite the mortar-node feature slots, ITZ predictions overwrite the ITZ
+edge feature slots, and gradients flow through both sub-models from the
+concrete-level loss.  Sub-model targets, when available, are exposed via
+``mortar_pred`` and ``itz_pred`` so the trainer can attach auxiliary losses.
+
+``TabularANN`` is a structure-blind reference baseline: a 4-layer MLP over
+the raw mix-design columns from the CSV. Any test-set gap between the
+hierarchical model and this baseline is attributable to structure.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+
+from ..missing_features import MissingFeatureEncoder
+from ..schema import (
+    DEFAULT_SCHEMA,
+    EDGE_TYPE_ITZ,
+    NODE_TYPE_AGGREGATE,
+    NODE_TYPE_MORTAR,
+    SchemaSpec,
+)
+from .concrete_gnn import ConcreteMesoGNN
+from .itz_model import ITZSubModel
+from .mortar_gnn import MortarSubModel
+from .outputs import transform_concrete_outputs
+
+
+def _slot_indices(target_names, feature_names):
+    return [feature_names.index(n) for n in target_names if n in feature_names]
+
+
+def _mortar_slot_indices(schema: SchemaSpec):
+    return _slot_indices(schema.mortar_targets, schema.mortar)
+
+
+def _itz_slot_indices(schema: SchemaSpec):
+    return _slot_indices(schema.itz_targets, schema.itz)
+
+
+def inject_mortar(
+    x: Tensor,
+    x_mask: Tensor,
+    mortar_pred: Tensor,
+    node_type: Tensor,
+    batch: Tensor,
+    slot_indices: list,
+) -> Tuple[Tensor, Tensor]:
+    """Write mortar predictions into mortar-node feature slots, keeping gradients."""
+
+    mortar_nodes = (node_type == NODE_TYPE_MORTAR).to(x.dtype)
+    pred_per_node = mortar_pred[batch]
+    injected = torch.zeros_like(x)
+    inj_mask = torch.zeros_like(x)
+    for k, idx in enumerate(slot_indices):
+        injected[:, idx] = pred_per_node[:, k] * mortar_nodes
+        inj_mask[:, idx] = mortar_nodes
+    x_new = x * (1.0 - inj_mask) + injected
+    x_mask_new = torch.maximum(x_mask, inj_mask)
+    return x_new, x_mask_new
+
+
+def inject_itz(
+    edge_attr: Tensor,
+    edge_attr_mask: Tensor,
+    itz_pred: Tensor,
+    itz_edge_mask: Tensor,
+    slot_indices: list,
+) -> Tuple[Tensor, Tensor]:
+    """Write ITZ predictions into ITZ-edge feature slots, keeping gradients."""
+
+    itz_idx = itz_edge_mask.nonzero(as_tuple=True)[0]
+    injected = torch.zeros_like(edge_attr)
+    inj_mask = torch.zeros_like(edge_attr)
+    for k, idx in enumerate(slot_indices):
+        injected[itz_idx, idx] = itz_pred[:, k]
+        inj_mask[itz_idx, idx] = 1.0
+    edge_new = edge_attr * (1.0 - inj_mask) + injected
+    edge_mask_new = torch.maximum(edge_attr_mask, inj_mask)
+    return edge_new, edge_mask_new
+
+
+def build_itz_inputs(
+    edge_index: Tensor,
+    edge_type: Tensor,
+    node_type: Tensor,
+    x: Tensor,
+    pos: Tensor,
+    batch_per_node: Tensor,
+    mix_itz_per_graph: Tensor,
+    schema: SchemaSpec,
+    input_dim: int,
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Assemble ITZ sub-model inputs for every ITZ edge."""
+
+    itz_edge_mask = edge_type == EDGE_TYPE_ITZ
+    src_all = edge_index[0]
+    edge_graph = batch_per_node[src_all]
+    itz_src = edge_index[0][itz_edge_mask]
+    itz_dst = edge_index[1][itz_edge_mask]
+    itz_graph = edge_graph[itz_edge_mask]
+    n_itz = int(itz_edge_mask.sum().item())
+
+    features = torch.zeros(n_itz, input_dim, device=x.device, dtype=x.dtype)
+    masks = torch.zeros_like(features)
+    if n_itz == 0:
+        return features, masks, itz_edge_mask, itz_graph
+
+    mix_dim = mix_itz_per_graph.size(1)
+    features[:, :mix_dim] = mix_itz_per_graph[itz_graph]
+
+    # Geometric distance between the two endpoints (positions are in mm).
+    edge_dist = torch.linalg.norm(pos[itz_src] - pos[itz_dst], dim=-1)
+    features[:, mix_dim] = edge_dist
+
+    # Aggregate-side surface texture (read from the aggregate endpoint).
+    src_is_agg = node_type[itz_src] == NODE_TYPE_AGGREGATE
+    agg_nodes = torch.where(src_is_agg, itz_src, itz_dst)
+    surface_texture_col = schema.aggregate.index("surface_texture")
+    features[:, mix_dim + 1] = x[agg_nodes, surface_texture_col]
+
+    masks[:, : mix_dim + 2] = 1.0
+    return features, masks, itz_edge_mask, itz_graph
+
+
+def aggregate_itz_per_graph(
+    itz_pred: Tensor, itz_graph: Tensor, num_graphs: int
+) -> Tensor:
+    """Mean ITZ prediction per concrete graph, shape ``(num_graphs, n_itz_targets)``.
+
+    Falls back to zero for graphs that contain no ITZ edges (empty RVE).
+    """
+
+    if itz_pred.numel() == 0:
+        return torch.zeros(
+            num_graphs, itz_pred.size(-1) if itz_pred.dim() == 2 else 0,
+            device=itz_pred.device, dtype=itz_pred.dtype,
+        )
+    n_targets = itz_pred.size(-1)
+    out = torch.zeros(num_graphs, n_targets, device=itz_pred.device, dtype=itz_pred.dtype)
+    counts = torch.zeros(num_graphs, device=itz_pred.device, dtype=itz_pred.dtype)
+    out.index_add_(0, itz_graph, itz_pred)
+    counts.index_add_(0, itz_graph, torch.ones_like(itz_graph, dtype=itz_pred.dtype))
+    return out / counts.clamp(min=1.0).unsqueeze(-1)
+
+
+def _scatter_mean_1d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
+    out = torch.zeros(num_graphs, device=src.device, dtype=src.dtype)
+    cnt = torch.zeros(num_graphs, device=src.device, dtype=src.dtype)
+    out.index_add_(0, index, src)
+    cnt.index_add_(0, index, torch.ones_like(src))
+    return out / cnt.clamp(min=1.0)
+
+
+def _scatter_amin_1d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
+    base = src.new_full((num_graphs,), float("inf"))
+    out = base.scatter_reduce(0, index, src, reduce="amin", include_self=True)
+    return torch.where(torch.isinf(out), torch.zeros_like(out), out)
+
+
+def _scatter_amin_2d(src: Tensor, index: Tensor, num_graphs: int) -> Tensor:
+    base = src.new_full((num_graphs, src.size(-1)), float("inf"))
+    out = base.scatter_reduce(
+        0, index.unsqueeze(-1).expand_as(src), src, reduce="amin", include_self=True
+    )
+    return torch.where(torch.isinf(out), torch.zeros_like(out), out)
+
+
+def build_itz_summary(
+    itz_pred: Tensor,
+    itz_graph: Tensor,
+    edge_index: Tensor,
+    edge_type: Tensor,
+    pos: Tensor,
+    node_type: Tensor,
+    x: Tensor,
+    batch_per_node: Tensor,
+    num_graphs: int,
+    schema: SchemaSpec,
+) -> Tensor:
+    """Per-graph ITZ + interface descriptor summary for the eta_ITZ head.
+
+    Columns: [mean ITZ targets (6), min ITZ targets (6),
+    aggregate-mortar contact density, mean ITZ edge distance,
+    min ITZ edge distance, mean aggregate surface texture].
+    """
+
+    n_itz_targets = len(schema.itz_targets)
+    itz_mean = aggregate_itz_per_graph(itz_pred, itz_graph, num_graphs)
+    if itz_pred.numel() > 0:
+        itz_min = _scatter_amin_2d(itz_pred, itz_graph, num_graphs)
+    else:
+        itz_min = torch.zeros(num_graphs, n_itz_targets, device=x.device, dtype=x.dtype)
+
+    # ITZ edge geometry (distance between the two endpoints of each ITZ edge).
+    itz_edge_mask = edge_type == EDGE_TYPE_ITZ
+    isrc = edge_index[0][itz_edge_mask]
+    idst = edge_index[1][itz_edge_mask]
+    if isrc.numel() > 0:
+        dist = torch.linalg.norm(pos[isrc] - pos[idst], dim=-1)
+        dist_mean = _scatter_mean_1d(dist, itz_graph, num_graphs)
+        dist_min = _scatter_amin_1d(dist, itz_graph, num_graphs)
+        n_itz_per_graph = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
+        n_itz_per_graph.index_add_(0, itz_graph, torch.ones_like(dist))
+    else:
+        dist_mean = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
+        dist_min = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
+        n_itz_per_graph = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
+
+    # Aggregate-mortar contact density = ITZ edges per aggregate node.
+    agg_mask = node_type == NODE_TYPE_AGGREGATE
+    agg_graph = batch_per_node[agg_mask]
+    n_agg = torch.zeros(num_graphs, device=x.device, dtype=x.dtype)
+    n_agg.index_add_(0, agg_graph, torch.ones(int(agg_mask.sum()), device=x.device, dtype=x.dtype))
+    contact_density = n_itz_per_graph / n_agg.clamp(min=1.0)
+
+    # Mean aggregate surface texture per graph (an explicit interface descriptor).
+    st_col = schema.aggregate.index("surface_texture")
+    st_mean = _scatter_mean_1d(x[agg_mask, st_col], agg_graph, num_graphs)
+
+    return torch.cat(
+        [
+            itz_mean,
+            itz_min,
+            contact_density.unsqueeze(-1),
+            dist_mean.unsqueeze(-1),
+            dist_min.unsqueeze(-1),
+            st_mean.unsqueeze(-1),
+        ],
+        dim=-1,
+    )
+
+
+class PhysicsStrengthHead(nn.Module):
+    """Concrete compressive strength as a physics-informed combination.
+
+    ``f_c = mortar_compressive * η_ITZ * η_matrix``
+
+    Both ``η_ITZ`` and ``η_matrix`` live in ``(0, 1]`` (sigmoids), so the
+    predicted concrete strength is bounded above by the predicted mortar
+    strength. This forces the mortar sub-model to carry signal — collapse to
+    a constant directly bounds concrete predictions and is no longer free.
+
+    * ``η_ITZ`` is derived from per-graph mean ITZ predictions + mesoscale
+      features (so weak ITZ pulls concrete strength down).
+    * ``η_matrix`` is derived from mesoscale graph features alone (captures
+      aggregate volume fraction, paste reactivity, geometry).
+    """
+
+    def __init__(self, feature_dim: int, itz_dim: int, hidden_dim: int = 64):
+        super().__init__()
+        self.itz_eff = nn.Sequential(
+            nn.Linear(feature_dim + itz_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, 1),
+        )
+        self.matrix_eff = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(
+        self,
+        mortar_compressive: Tensor,
+        itz_mean: Tensor,
+        graph_features: Tensor,
+    ) -> Tensor:
+        """Return per-graph concrete compressive strength, shape ``(G,)``."""
+
+        eta_itz = torch.sigmoid(
+            self.itz_eff(torch.cat([graph_features, itz_mean], dim=-1))
+        ).squeeze(-1)
+        eta_matrix = torch.sigmoid(self.matrix_eff(graph_features)).squeeze(-1)
+        return mortar_compressive * eta_itz * eta_matrix
+
+
+class MortarCapacityStrengthHead(nn.Module):
+    """Concrete strength with mortar capacity central, ITZ/matrix as reducers.
+
+    ``f_c = (mortar_compressive * eta_mortar) * eta_itz * eta_matrix``
+
+    * ``eta_mortar`` in (0, 1] is learned from the *other* mortar properties
+      (tensile / flexural / modulus / permeability / diffusivity / porosity /
+      creep), so compressive strength stays physically central while the rest
+      of the mortar state adjusts the usable capacity.
+    * ``eta_itz`` in (0, 1] is learned only from explicit ITZ + interface
+      descriptors (mean/min ITZ properties, contact density, edge-distance
+      stats, aggregate surface texture) - not the shared graph vector.
+    * ``eta_matrix`` in (0, 1] is learned from the mesoscale graph
+      representation (packing, topology, global curing context).
+
+    The point is to avoid feeding every factor the same latent graph vector.
+    """
+
+    def __init__(
+        self,
+        feature_dim: int,
+        mortar_other_dim: int,
+        itz_summary_dim: int,
+        hidden_dim: int = 64,
+    ):
+        super().__init__()
+        self.mortar_eff = nn.Sequential(
+            nn.Linear(mortar_other_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
+        )
+        self.itz_eff = nn.Sequential(
+            nn.Linear(itz_summary_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
+        )
+        self.matrix_eff = nn.Sequential(
+            nn.Linear(feature_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1)
+        )
+
+    def forward(
+        self, mortar_pred: Tensor, itz_summary: Tensor, graph_features: Tensor
+    ) -> Tensor:
+        eta_mortar = torch.sigmoid(self.mortar_eff(mortar_pred[:, 1:])).squeeze(-1)
+        capacity = mortar_pred[:, 0] * eta_mortar
+        eta_itz = torch.sigmoid(self.itz_eff(itz_summary)).squeeze(-1)
+        eta_matrix = torch.sigmoid(self.matrix_eff(graph_features)).squeeze(-1)
+        return capacity * eta_itz * eta_matrix
+
+
+class IntegratedMultiscaleModel(nn.Module):
+    """Mortar GNN + ITZ MLP + mesoscale GNN, trained jointly end-to-end."""
+
+    def __init__(
+        self,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        hidden_dim: int = 96,
+        num_layers: int = 2,  # ablation (1/3/5) showed depth >2 doesn't help this target
+        itz_input_dim: int = 24,
+        physics_aware_strength: bool = True,  # legacy flag; default head is "mortar_capacity"
+        pool: str = "mean",  # mesoscale graph readout: "mean" or "min"
+        strength_head_kind: Optional[str] = None,  # "free" | "physics" | "mortar_capacity"
+    ):
+        super().__init__()
+        self.schema = schema
+        self.itz_input_dim = itz_input_dim
+        self.mortar_slot_indices = _mortar_slot_indices(schema)
+        self.itz_slot_indices = _itz_slot_indices(schema)
+        # Back-compat: derive head kind from the legacy boolean when unset.
+        # Default head is the mortar-capacity head (keeps ITZ non-degenerate at
+        # no accuracy cost vs the plain physics head).
+        if strength_head_kind is None:
+            strength_head_kind = "mortar_capacity" if physics_aware_strength else "free"
+        self.strength_head_kind = strength_head_kind
+        self.physics_aware_strength = strength_head_kind != "free"
+
+        self.mortar = MortarSubModel(
+            schema=schema,
+            hidden_dim=64,
+            num_layers=2,
+        )
+        self.itz = ITZSubModel(
+            schema=schema,
+            hidden_dim=64,
+            input_dim=itz_input_dim,
+        )
+        self.mesoscale = ConcreteMesoGNN(
+            schema=schema,
+            hidden_dim=hidden_dim,
+            num_layers=num_layers,
+            pool=pool,
+        )
+        # ITZ summary dim: mean(n) + min(n) + [contact density, dist mean, dist min, surface texture]
+        itz_summary_dim = 2 * len(schema.itz_targets) + 4
+        if strength_head_kind == "physics":
+            self.strength_head = PhysicsStrengthHead(
+                feature_dim=self.mesoscale.feature_dim,
+                itz_dim=len(schema.itz_targets),
+            )
+        elif strength_head_kind == "mortar_capacity":
+            self.strength_head = MortarCapacityStrengthHead(
+                feature_dim=self.mesoscale.feature_dim,
+                mortar_other_dim=len(schema.mortar_targets) - 1,
+                itz_summary_dim=itz_summary_dim,
+            )
+        elif strength_head_kind == "free":
+            self.strength_head = None
+        else:
+            raise ValueError(f"unknown strength_head_kind {strength_head_kind!r}")
+
+    def forward(self, batch) -> Dict[str, Tensor]:
+        mortar_pred = self.mortar(
+            x=batch.m_x,
+            x_mask=batch.m_mask,
+            edge_index=batch.m_edge_index,
+            edge_attr=batch.m_edge_attr,
+            node_type=batch.m_node_type,
+            batch=batch.m_batch,
+            num_graphs=batch.num_graphs,
+            global_features=batch.mortar_global,
+            global_mask=batch.mortar_global_mask,
+        )
+
+        x, x_mask = inject_mortar(
+            batch.x,
+            batch.x_mask,
+            mortar_pred,
+            batch.node_type,
+            batch.batch,
+            slot_indices=self.mortar_slot_indices,
+        )
+
+        itz_feat, itz_mask, itz_edge_mask, itz_graph = build_itz_inputs(
+            edge_index=batch.edge_index,
+            edge_type=batch.edge_type,
+            node_type=batch.node_type,
+            x=batch.x,
+            pos=batch.pos,
+            batch_per_node=batch.batch,
+            mix_itz_per_graph=batch.mix_itz,
+            schema=self.schema,
+            input_dim=self.itz_input_dim,
+        )
+        if itz_feat.size(0) > 0:
+            itz_pred = self.itz(itz_feat, itz_mask)
+            edge_attr, edge_attr_mask = inject_itz(
+                batch.edge_attr,
+                batch.edge_attr_mask,
+                itz_pred,
+                itz_edge_mask,
+                slot_indices=self.itz_slot_indices,
+            )
+        else:
+            itz_pred = torch.zeros(
+                0, len(self.schema.itz_targets), device=batch.edge_attr.device
+            )
+            edge_attr, edge_attr_mask = batch.edge_attr, batch.edge_attr_mask
+
+        mesoscale_out = self.mesoscale(
+            x=x,
+            x_mask=x_mask,
+            edge_index=batch.edge_index,
+            edge_attr=edge_attr,
+            edge_attr_mask=edge_attr_mask,
+            node_type=batch.node_type,
+            edge_type=batch.edge_type,
+            global_features=batch.global_features,
+            global_mask=batch.global_mask,
+            batch=batch.batch,
+            num_graphs=batch.num_graphs,
+            return_features=self.strength_head is not None,
+        )
+        if self.strength_head_kind == "physics":
+            concrete_pred, graph_features = mesoscale_out
+            itz_mean = aggregate_itz_per_graph(itz_pred, itz_graph, batch.num_graphs)
+            strength = self.strength_head(
+                mortar_compressive=mortar_pred[:, 0],
+                itz_mean=itz_mean,
+                graph_features=graph_features,
+            )
+            concrete_pred = concrete_pred.clone()
+            concrete_pred[:, 0] = strength
+        elif self.strength_head_kind == "mortar_capacity":
+            concrete_pred, graph_features = mesoscale_out
+            itz_summary = build_itz_summary(
+                itz_pred,
+                itz_graph,
+                batch.edge_index,
+                batch.edge_type,
+                batch.pos,
+                batch.node_type,
+                batch.x,
+                batch.batch,
+                batch.num_graphs,
+                self.schema,
+            )
+            strength = self.strength_head(mortar_pred, itz_summary, graph_features)
+            concrete_pred = concrete_pred.clone()
+            concrete_pred[:, 0] = strength
+        else:
+            concrete_pred = mesoscale_out
+        return {
+            "concrete_pred": concrete_pred,
+            "mortar_pred": mortar_pred,
+            "itz_pred": itz_pred,
+        }
+
+
+class TabularANN(nn.Module):
+    """Tabular ANN baseline over the raw mix-design columns from the CSV.
+
+    Consumes the full ``batch.tabular_mix`` vector directly (cement, slag, fly
+    ash, SCMs, water, superplasticizer, coarse/fine aggregate, fibre fields,
+    and age). A BatchNorm at the input handles the wildly different scales of
+    these raw columns (kg/m^3 vs. days vs. dosage).
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        hidden_dim: int = 96,
+    ):
+        super().__init__()
+        self.schema = schema
+        self.in_dim = in_dim
+        self.input_bn = nn.BatchNorm1d(in_dim)
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, len(schema.concrete_targets)),
+        )
+
+    def forward(self, batch) -> Dict[str, Tensor]:
+        x = self.input_bn(batch.tabular_mix)
+        raw = self.net(x)
+        return {"concrete_pred": transform_concrete_outputs(raw)}

concrete_gnn/models/itz_model.py

@@ -0,0 +1,44 @@
+"""ITZ sub-model.
+
+The ITZ is treated as a modified cement paste whose properties depend on
+cement chemistry, SCM chemistry, water-to-binder ratio, admixture dosage,
+aggregate surface descriptors and curing condition.  Inputs are tabular per
+aggregate-mortar pair, so a small MLP with a missing-feature encoder and a
+bounded output transform is the natural fit.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+
+from ..missing_features import MissingFeatureEncoder
+from ..schema import DEFAULT_SCHEMA, SchemaSpec
+from .outputs import transform_itz_outputs
+
+
+class ITZSubModel(nn.Module):
+    def __init__(
+        self,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        hidden_dim: int = 64,
+        input_dim: int = 24,
+    ):
+        super().__init__()
+        self.schema = schema
+        self.input_dim = input_dim
+        self.encoder = MissingFeatureEncoder(input_dim, hidden_dim)
+        self.body = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+        )
+        self.head = nn.Linear(hidden_dim, len(schema.itz_targets))
+
+    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        h = self.encoder(x, mask)
+        h = self.body(h)
+        return transform_itz_outputs(self.head(h))

concrete_gnn/models/layers.py

@@ -0,0 +1,161 @@
+"""PyG-optional message-passing layers and pooling.
+
+When ``torch_geometric`` is importable we expose its battle-tested
+``GINEConv`` and ``NNConv`` operators (plus ``global_mean_pool`` and
+``global_add_pool``).  When PyG is not available we fall back to
+mathematically equivalent native PyTorch implementations so the prototype
+still runs end-to-end - useful when wheels are unavailable on the target
+machine or for unit-testing in lightweight environments.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+
+try:  # pragma: no cover - depends on local environment.
+    from torch_geometric.nn import (
+        GINEConv as _PygGINEConv,
+        NNConv as _PygNNConv,
+        global_add_pool as _pyg_add_pool,
+        global_mean_pool as _pyg_mean_pool,
+    )
+
+    HAS_PYG = True
+except Exception:  # pragma: no cover - PyG missing
+    _PygGINEConv = None
+    _PygNNConv = None
+    _pyg_add_pool = None
+    _pyg_mean_pool = None
+    HAS_PYG = False
+
+
+# ---------------------------------------------------------------------------
+# Pooling
+# ---------------------------------------------------------------------------
+
+
+def global_mean_pool(x: Tensor, batch: Tensor, num_graphs: Optional[int] = None) -> Tensor:
+    """Mean-aggregate node embeddings into one row per graph."""
+
+    if HAS_PYG and num_graphs is None:
+        return _pyg_mean_pool(x, batch)
+    if num_graphs is None:
+        num_graphs = int(batch.max().item()) + 1
+    pooled = torch.zeros(num_graphs, x.size(-1), device=x.device, dtype=x.dtype)
+    pooled.index_add_(0, batch, x)
+    counts = torch.zeros(num_graphs, 1, device=x.device, dtype=x.dtype)
+    counts.index_add_(0, batch, torch.ones(x.size(0), 1, device=x.device, dtype=x.dtype))
+    return pooled / counts.clamp_min(1.0)
+
+
+def global_add_pool(x: Tensor, batch: Tensor, num_graphs: Optional[int] = None) -> Tensor:
+    """Sum-aggregate node embeddings into one row per graph."""
+
+    if HAS_PYG and num_graphs is None:
+        return _pyg_add_pool(x, batch)
+    if num_graphs is None:
+        num_graphs = int(batch.max().item()) + 1
+    pooled = torch.zeros(num_graphs, x.size(-1), device=x.device, dtype=x.dtype)
+    pooled.index_add_(0, batch, x)
+    return pooled
+
+
+def global_min_pool(x: Tensor, batch: Tensor, num_graphs: Optional[int] = None) -> Tensor:
+    """Min-aggregate node embeddings into one row per graph (weakest-link readout).
+
+    Per-channel minimum over the nodes of each graph. Empty graphs (no nodes)
+    yield a zero row rather than +inf.
+    """
+
+    if num_graphs is None:
+        num_graphs = int(batch.max().item()) + 1
+    base = x.new_full((num_graphs, x.size(-1)), float("inf"))
+    idx = batch.unsqueeze(-1).expand_as(x)
+    out = base.scatter_reduce(0, idx, x, reduce="amin", include_self=True)
+    return torch.where(torch.isinf(out), torch.zeros_like(out), out)
+
+
+# ---------------------------------------------------------------------------
+# Native fallbacks
+# ---------------------------------------------------------------------------
+
+
+class _NativeGINE(nn.Module):
+    """Sum-aggregating edge-aware conv used when PyG is unavailable.
+
+    Approximates ``GINEConv`` semantics: messages are ReLU(x_j + e_ij), summed
+    over neighbours, then passed through the user-supplied MLP.
+    """
+
+    def __init__(self, mlp: nn.Module, edge_dim: int, in_dim: int, train_eps: bool = True):
+        super().__init__()
+        self.mlp = mlp
+        self.eps = nn.Parameter(torch.zeros(1)) if train_eps else None
+        self.edge_align = nn.Linear(edge_dim, in_dim) if edge_dim != in_dim else nn.Identity()
+
+    def forward(self, x: Tensor, edge_index: Tensor, edge_attr: Tensor) -> Tensor:
+        src, dst = edge_index
+        e = self.edge_align(edge_attr)
+        messages = torch.relu(x[src] + e)
+        aggregated = torch.zeros_like(x)
+        aggregated.index_add_(0, dst, messages)
+        eps = self.eps if self.eps is not None else torch.zeros(1, device=x.device)
+        return self.mlp((1.0 + eps) * x + aggregated)
+
+
+class _NativeNNConv(nn.Module):
+    """Edge-conditioned conv used when PyG is unavailable.
+
+    Mirrors ``NNConv`` with mean aggregation: an edge MLP produces a
+    (out_dim*in_dim) weight matrix per edge, applied to the source node.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int, edge_nn: nn.Module, aggr: str = "mean"):
+        super().__init__()
+        assert aggr in ("mean", "sum"), "aggr must be 'mean' or 'sum'"
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.edge_nn = edge_nn
+        self.aggr = aggr
+        self.root = nn.Linear(in_dim, out_dim, bias=True)
+
+    def forward(self, x: Tensor, edge_index: Tensor, edge_attr: Tensor) -> Tensor:
+        src, dst = edge_index
+        weights = self.edge_nn(edge_attr).view(-1, self.out_dim, self.in_dim)
+        msg = torch.bmm(weights, x[src].unsqueeze(-1)).squeeze(-1)
+        out = torch.zeros(x.size(0), self.out_dim, device=x.device, dtype=x.dtype)
+        out.index_add_(0, dst, msg)
+        if self.aggr == "mean":
+            deg = torch.zeros(x.size(0), 1, device=x.device, dtype=x.dtype)
+            deg.index_add_(0, dst, torch.ones(msg.size(0), 1, device=x.device, dtype=x.dtype))
+            out = out / deg.clamp_min(1.0)
+        return out + self.root(x)
+
+
+# ---------------------------------------------------------------------------
+# Public factories
+# ---------------------------------------------------------------------------
+
+
+def make_gine_conv(mlp: nn.Module, edge_dim: int, in_dim: int) -> nn.Module:
+    """``GINEConv`` from PyG when available, otherwise a sum-aggregating fallback."""
+
+    if HAS_PYG:
+        return _PygGINEConv(mlp, train_eps=True, edge_dim=edge_dim)
+    return _NativeGINE(mlp, edge_dim=edge_dim, in_dim=in_dim, train_eps=True)
+
+
+def make_nn_conv(
+    in_dim: int,
+    out_dim: int,
+    edge_nn: nn.Module,
+    aggr: str = "mean",
+) -> nn.Module:
+    """``NNConv`` from PyG when available, otherwise an edge-conditioned fallback."""
+
+    if HAS_PYG:
+        return _PygNNConv(in_dim, out_dim, edge_nn, aggr=aggr)
+    return _NativeNNConv(in_dim, out_dim, edge_nn, aggr=aggr)

concrete_gnn/models/mortar_gnn.py

@@ -0,0 +1,89 @@
+"""Microscale mortar GNN.
+
+The mortar sub-model consumes a small sand/paste micrograph and produces a
+mortar property vector that is later injected into mortar nodes of the
+mesoscale concrete graph. Message passing uses ``GINEConv`` (sum
+aggregation), the minimum needed to count sand-sand contacts - a relational
+quantity that drives mortar stiffness in the synthetic ground truth.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+from ..missing_features import MissingFeatureEncoder
+from ..schema import DEFAULT_SCHEMA, SchemaSpec
+from .layers import global_mean_pool, make_gine_conv
+from .outputs import transform_mortar_outputs
+
+
+class MortarSubModel(nn.Module):
+    """Sand/paste GNN that produces a mortar property vector per micrograph."""
+
+    def __init__(
+        self,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+        hidden_dim: int = 64,
+        num_layers: int = 2,
+    ):
+        super().__init__()
+        self.schema = schema
+
+        self.node_encoder = MissingFeatureEncoder(
+            schema.mortar_node_pad_dim, hidden_dim, num_stat_groups=schema.num_mortar_node_types
+        )
+        self.type_embed = nn.Embedding(schema.num_mortar_node_types, hidden_dim)
+
+        self.edge_encoder = nn.Sequential(
+            nn.Linear(2, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.SiLU(),
+        )
+        self.layers = nn.ModuleList(
+            [
+                make_gine_conv(
+                    nn.Sequential(
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.SiLU(),
+                        nn.Linear(hidden_dim, hidden_dim),
+                    ),
+                    edge_dim=hidden_dim,
+                    in_dim=hidden_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norms = nn.ModuleList(
+            [nn.LayerNorm(hidden_dim) for _ in range(num_layers)]
+        )
+        self.global_encoder = MissingFeatureEncoder(schema.mortar_global_dim, hidden_dim)
+
+        self.head = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, len(schema.mortar_targets)),
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        x_mask: Tensor,
+        edge_index: Tensor,
+        edge_attr: Tensor,
+        node_type: Tensor,
+        batch: Tensor,
+        num_graphs: int,
+        global_features: Tensor,
+        global_mask: Tensor,
+    ) -> Tensor:
+        h = self.node_encoder(x, x_mask, node_type) + self.type_embed(node_type)
+        e = self.edge_encoder(edge_attr)
+        for conv, norm in zip(self.layers, self.norms):
+            h_new = conv(h, edge_index, e)
+            h = norm(h + h_new)
+            h = torch.nn.functional.silu(h)
+        pooled = global_mean_pool(h, batch, num_graphs=num_graphs)
+        g = self.global_encoder(global_features, global_mask)
+        raw = self.head(torch.cat([pooled, g], dim=-1))
+        return transform_mortar_outputs(raw)

concrete_gnn/models/outputs.py

@@ -0,0 +1,75 @@
+"""Output transforms that constrain predictions to plausible engineering ranges.
+
+Free linear heads are easy to optimise but tend to drift into nonphysical
+territory (negative strengths, impossibly high moduli, etc.).  Each transform
+maps an unconstrained model output through a sigmoid scaled to a sensible
+engineering interval, so predictions stay positive and bounded without
+flattening the gradient where the data actually lives.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+
+def transform_mortar_outputs(raw: Tensor) -> Tensor:
+    """Map mortar sub-model outputs to plausible engineering ranges.
+
+    Output order matches ``schema.MORTAR_TARGETS``:
+    (compressive, tensile, flexural, elastic_modulus, permeability,
+    diffusivity, porosity, creep_coefficient).
+    """
+
+    compressive = 15.0 + 285.0 * torch.sigmoid(raw[..., 0:1])
+    tensile = 1.0 + 8.0 * torch.sigmoid(raw[..., 1:2])
+    flexural = 2.0 + 48.0 * torch.sigmoid(raw[..., 2:3])
+    elastic = 10.0 + 40.0 * torch.sigmoid(raw[..., 3:4])
+    permeability = 1.0e-13 + 1.0e-11 * torch.sigmoid(raw[..., 4:5])
+    diffusivity = 1.0e-12 + 1.0e-10 * torch.sigmoid(raw[..., 5:6])
+    porosity = 0.03 + 0.25 * torch.sigmoid(raw[..., 6:7])
+    creep = 0.5 + 3.0 * torch.sigmoid(raw[..., 7:8])
+    return torch.cat(
+        [compressive, tensile, flexural, elastic, permeability, diffusivity, porosity, creep],
+        dim=-1,
+    )
+
+
+def transform_itz_outputs(raw: Tensor) -> Tensor:
+    """Map ITZ sub-model outputs to plausible engineering ranges.
+
+    Output order matches ``schema.ITZ_TARGETS``:
+    (thickness_um, porosity, strength, elastic_modulus, permeability,
+    diffusivity).
+    """
+
+    thickness = 5.0 + 75.0 * torch.sigmoid(raw[..., 0:1])
+    porosity = 0.05 + 0.45 * torch.sigmoid(raw[..., 1:2])
+    strength = 5.0 + 65.0 * torch.sigmoid(raw[..., 2:3])
+    elastic = 5.0 + 35.0 * torch.sigmoid(raw[..., 3:4])
+    permeability = 1.0e-13 + 2.0e-11 * torch.sigmoid(raw[..., 4:5])
+    diffusivity = 1.0e-12 + 2.0e-10 * torch.sigmoid(raw[..., 5:6])
+    return torch.cat(
+        [thickness, porosity, strength, elastic, permeability, diffusivity],
+        dim=-1,
+    )
+
+
+def transform_concrete_outputs(raw: Tensor) -> Tensor:
+    """Map mesoscale concrete predictions to plausible engineering ranges.
+
+    Output order matches ``schema.CONCRETE_TARGETS``:
+    (compressive, tensile, flexural, elastic_modulus, permeability,
+    diffusivity).
+    """
+
+    compressive = 350.0 * torch.sigmoid(raw[..., 0:1])
+    tensile = 0.5 + 29.5 * torch.sigmoid(raw[..., 1:2])
+    flexural = 1.0 + 49.0 * torch.sigmoid(raw[..., 2:3])
+    elastic = 8.0 + 92.0 * torch.sigmoid(raw[..., 3:4])
+    permeability = 1.0e-13 + 3.0e-11 * torch.sigmoid(raw[..., 4:5])
+    diffusivity = 1.0e-12 + 3.0e-10 * torch.sigmoid(raw[..., 5:6])
+    return torch.cat(
+        [compressive, tensile, flexural, elastic, permeability, diffusivity],
+        dim=-1,
+    )

concrete_gnn/real_data.py

@@ -0,0 +1,648 @@
+"""Utilities for training Hybrid models on real concrete mix datasets."""
+
+from __future__ import annotations
+
+import os
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Union
+
+import pandas as pd
+import torch
+from torch import Tensor, nn
+from torch.utils.data import DataLoader, Dataset
+
+from .categoricals import (
+    CEMENT_TYPES,
+    CURING_REGIMES,
+    FIBRE_TYPES,
+    canonical_cement_type,
+    canonical_curing_regime,
+    canonical_fibre_type,
+    one_hot,
+    type_id,
+)
+from .data import ConcreteSample
+from .graph_generator import (
+    MixDesign,
+    MortarMixDesign,
+    generate_concrete_graph,
+    generate_mortar_graph,
+    mortar_volume_fraction,
+)
+from .schema import DEFAULT_SCHEMA, SchemaSpec
+
+
+COL_CEMENT = "Cement (component 1)(kg in a m^3 mixture)"
+COL_SLAG = "Blast Furnace Slag (component 2)(kg in a m^3 mixture)"
+COL_FLY_ASH = "Fly Ash (component 3)(kg in a m^3 mixture)"
+COL_WATER = "Water  (component 4)(kg in a m^3 mixture)"
+COL_SUPERPLASTICIZER = "Superplasticizer (component 5)(kg in a m^3 mixture)"
+COL_COARSE = "Coarse Aggregate  (component 6)(kg in a m^3 mixture)"
+COL_FINE = "Fine Aggregate (component 7)(kg in a m^3 mixture)"
+COL_AGE = "Age (day)"
+COL_STRENGTH = "Concrete compressive strength(MPa, megapascals) "
+TARGET_NAME = "compressive_strength_MPa"
+
+NORM_CEMENT = "cement_kg_m3"
+NORM_SLAG = "slag_kg_m3"
+NORM_FLY_ASH = "fly_ash_kg_m3"
+NORM_SILICA_FUME = "silica_fume_kg_m3"
+NORM_METAKAOLIN = "metakaolin_kg_m3"
+NORM_LIMESTONE = "limestone_powder_kg_m3"
+NORM_OTHER_SCM = "other_scm_kg_m3"
+NORM_WATER = "water_kg_m3"
+NORM_SUPERPLASTICIZER = "superplasticizer_kg_m3"
+NORM_COARSE = "coarse_aggregate_kg_m3"
+NORM_FINE = "fine_aggregate_kg_m3"
+NORM_FIBRE_CONTENT = "fibre_content_kg_m3"
+NORM_FIBRE_LENGTH = "fibre_length_mm"
+NORM_FIBRE_DIAMETER = "fibre_diameter_mm"
+NORM_FIBRE_TENSILE = "fibre_tensile_strength_mpa"
+NORM_FIBRE_MODULUS = "fibre_modulus_gpa"
+NORM_MAX_COARSE = "max_coarse_aggregate_size_mm"
+NORM_MAX_FINE = "max_fine_aggregate_size_mm"
+NORM_CURING_TEMP = "curing_temperature_c"
+NORM_AGE = "age_days"
+NORM_STRENGTH = "compressive_strength_mpa"
+
+# Rich UHPC inputs (Data/prepare_uhpc_rich.py): per-mix cement & SCM oxide
+# chemistry plus a few extra numerics. All maskable -> NaN/absent for sources
+# that do not report them, in which case the synthetic OPC prior is used for the
+# GNN paste slots and the tabular mask channel is set to 0.
+NORM_CEMENT_CAO = "cement_CaO_pct"
+NORM_CEMENT_SIO2 = "cement_SiO2_pct"
+NORM_CEMENT_AL2O3 = "cement_Al2O3_pct"
+NORM_CEMENT_FE2O3 = "cement_Fe2O3_pct"
+NORM_CEMENT_MGO = "cement_MgO_pct"
+NORM_CEMENT_SO3 = "cement_SO3_pct"
+NORM_CEMENT_ALKALI = "cement_alkali_pct"
+NORM_CEMENT_LOI = "cement_LOI_pct"
+NORM_SCM_CAO = "scm_CaO_pct"
+NORM_SCM_SIO2 = "scm_SiO2_pct"
+NORM_SCM_AL2O3 = "scm_Al2O3_pct"
+NORM_SCM_FE2O3 = "scm_Fe2O3_pct"
+NORM_SCM_MGO = "scm_MgO_pct"
+NORM_SCM_LOI = "scm_LOI_pct"
+NORM_CEMENT_GRADE = "cement_grade_mpa"
+NORM_CURING_HUMIDITY = "curing_humidity_pct"
+NORM_SPECIMEN_SIZE = "specimen_size_mm"
+
+# Canonical free-text categoricals (string columns; one-hot for the TabularANN).
+NORM_CEMENT_TYPE = "cement_type_norm"
+NORM_FIBRE_TYPE = "fibre_type_norm"
+NORM_CURING_REGIME = "curing_regime_norm"
+
+# Maskable rich columns fed (value + observed-mask channel) to the TabularANN.
+RICH_TABULAR_COLUMNS = (
+    NORM_CEMENT_CAO,
+    NORM_CEMENT_SIO2,
+    NORM_CEMENT_AL2O3,
+    NORM_CEMENT_FE2O3,
+    NORM_CEMENT_MGO,
+    NORM_CEMENT_SO3,
+    NORM_CEMENT_ALKALI,
+    NORM_CEMENT_LOI,
+    NORM_SCM_CAO,
+    NORM_SCM_SIO2,
+    NORM_SCM_AL2O3,
+    NORM_SCM_FE2O3,
+    NORM_SCM_MGO,
+    NORM_SCM_LOI,
+    NORM_CEMENT_GRADE,
+    NORM_CURING_HUMIDITY,
+    NORM_SPECIMEN_SIZE,
+)
+
+# One-hot categorical blocks appended to the tabular vector (after the rich block).
+CATEGORICAL_TABULAR = (
+    (NORM_CEMENT_TYPE, canonical_cement_type, CEMENT_TYPES, "other"),
+    (NORM_FIBRE_TYPE, canonical_fibre_type, FIBRE_TYPES, "none"),
+    (NORM_CURING_REGIME, canonical_curing_regime, CURING_REGIMES, "other"),
+)
+
+# OPC-typical fallbacks used for the GNN paste chemistry when a row does not
+# carry measured oxides (keeps non-UHPC behavior identical to the old defaults).
+_DEFAULT_CEMENT_CHEM = (64.0, 21.0, 5.0, 3.0)       # CaO, SiO2, Al2O3, Fe2O3
+_DEFAULT_CEMENT_CHEM_EXT = (2.0, 2.5, 0.6)          # MgO, SO3, Na2O-eq alkali
+_DEFAULT_SCM_CHEM = (4.0, 55.0)                     # CaO, SiO2
+_DEFAULT_SCM_CHEM_EXT = (15.0, 5.0, 2.0, 3.0)       # Al2O3, Fe2O3, MgO, LOI
+
+TableInput = Union[str, os.PathLike, pd.DataFrame]
+
+
+def read_table(data: TableInput) -> pd.DataFrame:
+    if isinstance(data, pd.DataFrame):
+        return data.copy()
+    path = Path(data)
+    if path.suffix.lower() == ".csv":
+        return pd.read_csv(path)
+    return pd.read_excel(path)
+
+
+def row_value(row, name: str, default: float = 0.0) -> float:
+    value = row.get(name, default)
+    if pd.isna(value):
+        return default
+    return float(value)
+
+
+def row_value_optional(row, name: str) -> Optional[float]:
+    """Return the float value, or ``None`` when the column is absent / NaN.
+
+    Used for maskable descriptors (max grain sizes, curing temperature) where
+    "not reported" must stay distinct from a genuine zero.
+    """
+
+    if name not in row.index:
+        return None
+    value = row[name]
+    if pd.isna(value):
+        return None
+    return float(value)
+
+
+def is_normalized_row(row) -> bool:
+    return NORM_CEMENT in row.index
+
+
+def _chem_or_default(row, name: str, default: float) -> float:
+    """Measured oxide value when the column is present & non-NaN, else ``default``."""
+    value = row_value_optional(row, name)
+    return default if value is None else value
+
+
+def _category(row, name: str, canon_fn, vocab, default: str) -> str:
+    """Read a categorical column as a canonical class string.
+
+    The rich CSV already stores canonical labels (in ``vocab``); raw labels from
+    other sources are canonicalized on the fly; an absent/blank column -> ``default``.
+    """
+    if name not in row.index:
+        return default
+    value = row[name]
+    if value is None or (isinstance(value, float) and pd.isna(value)):
+        return default
+    text = str(value).strip()
+    if text == "" or text.lower() == "nan":
+        return default
+    return text if text in vocab else canon_fn(text)
+
+
+@dataclass
+class Standardizer:
+    mean: Tensor
+    std: Tensor
+
+    @classmethod
+    def from_values(cls, values: Tensor) -> "Standardizer":
+        return cls(values.mean(), values.std().clamp_min(1.0e-6))
+
+    def encode(self, values: Tensor) -> Tensor:
+        mean = self.mean.to(device=values.device, dtype=values.dtype)
+        std = self.std.to(device=values.device, dtype=values.dtype)
+        return (values - mean) / std
+
+    def decode(self, values: Tensor) -> Tensor:
+        mean = self.mean.to(device=values.device, dtype=values.dtype)
+        std = self.std.to(device=values.device, dtype=values.dtype)
+        return values * std + mean
+
+
+def strength_column(df: pd.DataFrame) -> str:
+    return NORM_STRENGTH if NORM_STRENGTH in df.columns else COL_STRENGTH
+
+
+def row_to_mix(row) -> MixDesign:
+    if is_normalized_row(row):
+        cement = row_value(row, NORM_CEMENT)
+        slag = row_value(row, NORM_SLAG)
+        fly_ash = row_value(row, NORM_FLY_ASH)
+        silica_fume = row_value(row, NORM_SILICA_FUME)
+        metakaolin = row_value(row, NORM_METAKAOLIN)
+        limestone = row_value(row, NORM_LIMESTONE)
+        other_scm = row_value(row, NORM_OTHER_SCM)
+        water = row_value(row, NORM_WATER)
+        coarse = row_value(row, NORM_COARSE)
+        superplasticizer = row_value(row, NORM_SUPERPLASTICIZER)
+        age = row_value(row, NORM_AGE)
+    else:
+        cement = float(row[COL_CEMENT])
+        slag = float(row[COL_SLAG])
+        fly_ash = float(row[COL_FLY_ASH])
+        silica_fume = 0.0
+        metakaolin = 0.0
+        limestone = 0.0
+        other_scm = 0.0
+        water = float(row[COL_WATER])
+        coarse = float(row[COL_COARSE])
+        superplasticizer = float(row[COL_SUPERPLASTICIZER])
+        age = float(row[COL_AGE])
+
+    extra_scm = metakaolin + limestone + other_scm
+    binder = max(cement + slag + fly_ash + silica_fume + extra_scm, 1.0)
+    scm = slag + fly_ash + silica_fume + extra_scm
+    primary_count = sum(v > 0.0 for v in (slag, fly_ash, silica_fume))
+    if primary_count > 1 or (primary_count >= 1 and extra_scm > 0.0):
+        scm_type = 3
+    elif slag > 0.0:
+        scm_type = 1
+    elif fly_ash > 0.0:
+        scm_type = 2
+    elif silica_fume > 0.0 or extra_scm > 0.0:
+        scm_type = 4
+    else:
+        scm_type = 0
+
+    aggregate_volume_fraction = max(0.0, min(0.70, coarse / 2650.0))
+    water_to_binder = water / binder
+
+    fibre_content = row_value(row, NORM_FIBRE_CONTENT) if is_normalized_row(row) else 0.0
+    fibre_length = row_value(row, NORM_FIBRE_LENGTH) if is_normalized_row(row) else 0.0
+    fibre_diameter = row_value(row, NORM_FIBRE_DIAMETER) if is_normalized_row(row) else 0.0
+    fibre_tensile = row_value(row, NORM_FIBRE_TENSILE) if is_normalized_row(row) else 0.0
+    fibre_modulus = row_value(row, NORM_FIBRE_MODULUS) if is_normalized_row(row) else 0.0
+    # Fibre material class -> id for the mesoscale mortar-node fibre_type_id slot.
+    fibre_type = _category(row, NORM_FIBRE_TYPE, canonical_fibre_type, FIBRE_TYPES, "none")
+
+    # Maskable descriptors: None when the source did not report them.
+    max_coarse = row_value_optional(row, NORM_MAX_COARSE)
+    max_fine = row_value_optional(row, NORM_MAX_FINE)
+    curing_temp = row_value_optional(row, NORM_CURING_TEMP)
+    # Measured curing humidity (%) when reported (UHPC), else the 0.95 default.
+    humidity_pct = row_value_optional(row, NORM_CURING_HUMIDITY)
+    relative_humidity = humidity_pct / 100.0 if humidity_pct is not None else 0.95
+
+    return MixDesign(
+        aggregate_volume_fraction=aggregate_volume_fraction,
+        water_to_binder_ratio=max(0.10, min(1.20, water_to_binder)),
+        cement_content_kg_m3=cement,
+        scm_type_id=scm_type,
+        scm_fraction=max(0.0, min(0.70, scm / binder)),
+        admixture_dosage=superplasticizer / binder,
+        relative_humidity=relative_humidity,
+        temperature_C=curing_temp if curing_temp is not None else 20.0,
+        temperature_observed=curing_temp is not None,
+        curing_age_days=age,
+        max_coarse_aggregate_size_mm=max_coarse,
+        max_fine_aggregate_size_mm=max_fine,
+        fibre_content_kg_m3=fibre_content,
+        fibre_length_mm=fibre_length,
+        fibre_diameter_mm=fibre_diameter,
+        fibre_tensile_strength_MPa=fibre_tensile,
+        fibre_modulus_GPa=fibre_modulus,
+        fibre_type_id=float(type_id(fibre_type, FIBRE_TYPES)),
+    )
+
+
+def mortar_mix_from_row(row, mix: MixDesign) -> MortarMixDesign:
+    fine = row_value(row, NORM_FINE) if is_normalized_row(row) else float(row[COL_FINE])
+    coarse = row_value(row, NORM_COARSE) if is_normalized_row(row) else float(row[COL_COARSE])
+    # Convert concrete-basis densities (kg per m^3 concrete) to a mortar basis
+    # (kg per m^3 mortar) by removing the coarse-aggregate volume. The mortar
+    # phase is denser in cement and fibre because it excludes coarse aggregate.
+    # w/b, scm_fraction and admixture_dosage are intensive ratios -> invariant.
+    mortar_vf = mortar_volume_fraction(coarse / 2650.0)
+    # Real cement / SCM oxide chemistry where the row carries it (UHPC rich CSV);
+    # OPC-typical synthetic priors otherwise (keeps other sources unchanged).
+    cement_chem = (
+        _chem_or_default(row, NORM_CEMENT_CAO, _DEFAULT_CEMENT_CHEM[0]),
+        _chem_or_default(row, NORM_CEMENT_SIO2, _DEFAULT_CEMENT_CHEM[1]),
+        _chem_or_default(row, NORM_CEMENT_AL2O3, _DEFAULT_CEMENT_CHEM[2]),
+        _chem_or_default(row, NORM_CEMENT_FE2O3, _DEFAULT_CEMENT_CHEM[3]),
+    )
+    cement_chem_ext = (
+        _chem_or_default(row, NORM_CEMENT_MGO, _DEFAULT_CEMENT_CHEM_EXT[0]),
+        _chem_or_default(row, NORM_CEMENT_SO3, _DEFAULT_CEMENT_CHEM_EXT[1]),
+        _chem_or_default(row, NORM_CEMENT_ALKALI, _DEFAULT_CEMENT_CHEM_EXT[2]),
+    )
+    scm_chem = (
+        _chem_or_default(row, NORM_SCM_CAO, _DEFAULT_SCM_CHEM[0]),
+        _chem_or_default(row, NORM_SCM_SIO2, _DEFAULT_SCM_CHEM[1]),
+    )
+    scm_chem_ext = (
+        _chem_or_default(row, NORM_SCM_AL2O3, _DEFAULT_SCM_CHEM_EXT[0]),
+        _chem_or_default(row, NORM_SCM_FE2O3, _DEFAULT_SCM_CHEM_EXT[1]),
+        _chem_or_default(row, NORM_SCM_MGO, _DEFAULT_SCM_CHEM_EXT[2]),
+        _chem_or_default(row, NORM_SCM_LOI, _DEFAULT_SCM_CHEM_EXT[3]),
+    )
+    cement_type = _category(row, NORM_CEMENT_TYPE, canonical_cement_type, CEMENT_TYPES, "other")
+    return MortarMixDesign(
+        sand_volume_fraction=max(0.20, min(0.60, fine / 2650.0)),
+        water_to_binder_ratio=mix.water_to_binder_ratio,
+        cement_content_kg_m3=mix.cement_content_kg_m3 / mortar_vf,
+        scm_type_id=mix.scm_type_id,
+        scm_fraction=mix.scm_fraction,
+        admixture_dosage=mix.admixture_dosage,
+        cement_chem=cement_chem,
+        cement_chem_ext=cement_chem_ext,
+        scm_chem=scm_chem,
+        scm_chem_ext=scm_chem_ext,
+        cement_type_id=float(type_id(cement_type, CEMENT_TYPES)),
+        curing_relative_humidity=mix.relative_humidity,
+        curing_temperature_C=mix.temperature_C,
+        curing_temperature_observed=mix.temperature_observed,
+        max_sand_size_mm=mix.max_fine_aggregate_size_mm,
+        curing_age_days=mix.curing_age_days,
+        fibre_content_kg_m3=mix.fibre_content_kg_m3 / mortar_vf,
+        fibre_length_mm=mix.fibre_length_mm,
+        fibre_diameter_mm=mix.fibre_diameter_mm,
+        fibre_tensile_strength_MPa=mix.fibre_tensile_strength_MPa,
+        fibre_modulus_GPa=mix.fibre_modulus_GPa,
+    )
+
+
+def mix_itz_vector(mix: MixDesign) -> Tensor:
+    return torch.tensor(
+        [
+            mix.cement_content_kg_m3,
+            float(mix.scm_type_id),
+            mix.scm_fraction,
+            mix.water_to_binder_ratio,
+            mix.admixture_dosage,
+            mix.relative_humidity,
+            mix.temperature_C,
+            mix.curing_age_days,
+        ],
+        dtype=torch.float32,
+    )
+
+
+RAW_MIX_COLUMNS = (
+    NORM_CEMENT,
+    NORM_SLAG,
+    NORM_FLY_ASH,
+    NORM_SILICA_FUME,
+    NORM_METAKAOLIN,
+    NORM_LIMESTONE,
+    NORM_OTHER_SCM,
+    NORM_WATER,
+    NORM_SUPERPLASTICIZER,
+    NORM_COARSE,
+    NORM_FINE,
+    NORM_FIBRE_CONTENT,
+    NORM_FIBRE_LENGTH,
+    NORM_FIBRE_DIAMETER,
+    NORM_FIBRE_TENSILE,
+    NORM_FIBRE_MODULUS,
+    NORM_AGE,
+)
+# Tabular input = raw mix columns, then the rich maskable columns, then one
+# observed-mask channel per rich column (1 = measured, 0 = missing), then one-hot
+# blocks for the categorical columns. The mask channels let the ANN distinguish
+# "not reported" from a genuine zero.
+_CATEGORICAL_WIDTH = sum(len(vocab) for _, _, vocab, _ in CATEGORICAL_TABULAR)
+TABULAR_DIM = len(RAW_MIX_COLUMNS) + 2 * len(RICH_TABULAR_COLUMNS) + _CATEGORICAL_WIDTH
+
+
+def raw_mix_vector(row) -> Tensor:
+    """Build the per-sample raw mix-design vector used by TabularANN.
+
+    Layout: raw mix columns, then the rich maskable block (value + observed-mask
+    channels), then one-hot categorical blocks. Uses the normalized columns when
+    available, otherwise maps the original Yeh XLS columns into the leading slots;
+    sources lacking the rich/categorical columns contribute zeros and the default
+    categorical bucket.
+    """
+
+    if is_normalized_row(row):
+        base = [row_value(row, name) for name in RAW_MIX_COLUMNS]
+    else:
+        base = [
+            float(row[COL_CEMENT]),
+            float(row[COL_SLAG]),
+            float(row[COL_FLY_ASH]),
+            0.0,  # silica_fume
+            0.0,  # metakaolin
+            0.0,  # limestone
+            0.0,  # other_scm
+            float(row[COL_WATER]),
+            float(row[COL_SUPERPLASTICIZER]),
+            float(row[COL_COARSE]),
+            float(row[COL_FINE]),
+            0.0,  # fibre_content
+            0.0,  # fibre_length
+            0.0,  # fibre_diameter
+            0.0,  # fibre_tensile
+            0.0,  # fibre_modulus
+            float(row[COL_AGE]),
+        ]
+
+    rich_values: List[float] = []
+    rich_mask: List[float] = []
+    for name in RICH_TABULAR_COLUMNS:
+        v = row_value_optional(row, name)
+        rich_values.append(0.0 if v is None else v)
+        rich_mask.append(0.0 if v is None else 1.0)
+
+    categorical: List[float] = []
+    for name, canon_fn, vocab, default in CATEGORICAL_TABULAR:
+        categorical.extend(one_hot(_category(row, name, canon_fn, vocab, default), vocab))
+
+    return torch.tensor(base + rich_values + rich_mask + categorical, dtype=torch.float32)
+
+
+class ConcreteMixDataset(Dataset):
+    def __init__(
+        self,
+        data: TableInput,
+        indices: Optional[Sequence[int]] = None,
+        missing_rate: float = 0.0,
+        seed: int = 17,
+        target_standardizer: Optional[Standardizer] = None,
+        schema: SchemaSpec = DEFAULT_SCHEMA,
+    ):
+        df = read_table(data)
+        if indices is not None:
+            df = df.iloc[list(indices)].reset_index(drop=True)
+        self.df = df.reset_index(drop=True)
+        self.schema = schema
+        y = torch.tensor(
+            self.df[strength_column(self.df)].to_numpy(dtype="float32"),
+            dtype=torch.float32,
+        )
+        self.y_scaled = (
+            target_standardizer.encode(y) if target_standardizer is not None else y
+        )
+        self.samples: List[ConcreteSample] = []
+        self._build_samples(seed, missing_rate)
+
+    def _build_samples(self, seed: int, missing_rate: float) -> None:
+        for i, row in self.df.iterrows():
+            mix = row_to_mix(row)
+            mortar_mix = mortar_mix_from_row(row, mix)
+            graph = generate_concrete_graph(
+                mix,
+                schema=self.schema,
+                seed=seed + 1000 + i,
+                missing_rate=missing_rate,
+            )
+            micro = generate_mortar_graph(
+                mortar_mix, schema=self.schema, seed=seed + 2000 + i
+            )
+            y = torch.zeros(6, dtype=torch.float32)
+            y[0] = self.y_scaled[i]
+            self.samples.append(
+                ConcreteSample(
+                    graph=graph,
+                    micrograph=micro,
+                    mix_itz=mix_itz_vector(mix),
+                    concrete_target=y,
+                    mortar_target=None,
+                    itz_target=None,
+                    tabular_mix=raw_mix_vector(row),
+                )
+            )
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> ConcreteSample:
+        return self.samples[idx]
+
+
+class StrengthHead(nn.Module):
+    """Wrap a Hybrid model and train only compressive strength."""
+
+    def __init__(self, base: nn.Module, standardizer: Standardizer):
+        super().__init__()
+        self.base = base
+        self.standardizer = standardizer
+
+    def forward(self, batch) -> Tensor:
+        out = self.base(batch)["concrete_pred"][:, 0]
+        return self.standardizer.encode(out)
+
+
+def split_indices(n: int, seed: int) -> tuple:
+    idx = list(range(n))
+    random.Random(seed).shuffle(idx)
+    n_train = int(0.70 * n)
+    n_val = int(0.15 * n)
+    return idx[:n_train], idx[n_train : n_train + n_val], idx[n_train + n_val :]
+
+
+def metrics(pred_mpa: Tensor, true_mpa: Tensor) -> Dict[str, float]:
+    resid = true_mpa - pred_mpa
+    ss_res = resid.pow(2).sum()
+    ss_tot = (true_mpa - true_mpa.mean()).pow(2).sum().clamp_min(1.0e-12)
+    return {
+        "rmse": float(resid.pow(2).mean().sqrt()),
+        "mae": float(resid.abs().mean()),
+        "r2": float(1.0 - ss_res / ss_tot),
+    }
+
+
+@torch.no_grad()
+def collect_strength_predictions(
+    model: StrengthHead, loader: DataLoader, device: torch.device
+) -> tuple:
+    model.eval()
+    pred_scaled, true_scaled = [], []
+    for batch in loader:
+        batch = batch.to(device)
+        pred_scaled.append(model(batch).cpu())
+        true_scaled.append(batch.concrete_target[:, 0].cpu())
+    pred = model.standardizer.decode(torch.cat(pred_scaled))
+    true = model.standardizer.decode(torch.cat(true_scaled))
+    return pred, true
+
+
+def evaluate_strength(
+    model: StrengthHead, loader: DataLoader, device: torch.device
+) -> Dict[str, float]:
+    pred, true = collect_strength_predictions(model, loader, device)
+    return metrics(pred, true)
+
+
+def fit_encoder_standardizers(model: nn.Module, samples: Sequence) -> None:
+    """Fit per-feature input standardization for the GNN encoders that carry
+    real mix-design signal: the mortar sub-model nodes (paste w/b, cement,
+    admixture) and globals, the mesoscale globals, and the ITZ mix vector.
+
+    This mirrors the input ``BatchNorm`` the ``TabularANN`` baseline already
+    uses, so small-magnitude-but-informative features (w/b ratio ~0.4, admixture
+    dosage ~0.01) are not swamped by large ones (cement ~500). Stats are masked
+    so unobserved / placeholder slots are ignored.
+
+    Encoders whose columns multiplex several feature schemas (mesoscale nodes:
+    aggregate vs mortar; mesoscale edges: ITZ vs mortar-mortar vs agg-agg; mortar
+    nodes: sand vs paste) are standardized per row TYPE, so each schema is scaled
+    against its own statistics rather than a blended one. Stats are fit from the
+    raw graph priors, which ``_mortar_placeholder_vector`` / ``_itz_feature_vector``
+    deliberately scale-calibrate to the values the sub-models later inject.
+    """
+
+    from .missing_features import masked_feature_stats, masked_feature_stats_by_type
+
+    if not samples or not hasattr(model, "mortar"):
+        return
+
+    schema = model.schema
+
+    # Mesoscale globals (single schema; curing context under curing-only).
+    gf = torch.stack([s.graph.global_features for s in samples])
+    gm = torch.stack([s.graph.global_mask for s in samples])
+    model.mesoscale.global_encoder.set_feature_stats(*masked_feature_stats(gf, gm))
+
+    # Mesoscale nodes (aggregate vs mortar) -- per node type.
+    nx = torch.cat([s.graph.x for s in samples], dim=0)
+    nxm = torch.cat([s.graph.x_mask for s in samples], dim=0)
+    nnt = torch.cat([s.graph.node_type for s in samples], dim=0)
+    model.mesoscale.node_encoder.set_feature_stats(
+        *masked_feature_stats_by_type(nx, nxm, nnt, schema.num_node_types)
+    )
+
+    # Mesoscale edges (ITZ vs mortar-mortar vs agg-agg) -- per edge type.
+    ex = torch.cat([s.graph.edge_attr for s in samples], dim=0)
+    exm = torch.cat([s.graph.edge_attr_mask for s in samples], dim=0)
+    eet = torch.cat([s.graph.edge_type for s in samples], dim=0)
+    model.mesoscale.edge_encoder.set_feature_stats(
+        *masked_feature_stats_by_type(ex, exm, eet, schema.num_edge_types)
+    )
+
+    # Mortar sub-model nodes (sand vs paste) -- per type; paste carries w/b,
+    # cement content and admixture dosage, the features the diagnosis flagged.
+    mx = torch.cat([s.micrograph.x for s in samples], dim=0)
+    mxm = torch.cat([s.micrograph.x_mask for s in samples], dim=0)
+    mnt = torch.cat([s.micrograph.node_type for s in samples], dim=0)
+    model.mortar.node_encoder.set_feature_stats(
+        *masked_feature_stats_by_type(mx, mxm, mnt, schema.num_mortar_node_types)
+    )
+
+    # Mortar sub-model globals (single schema).
+    mg = torch.stack([s.micrograph.global_features for s in samples])
+    mgm = torch.stack([s.micrograph.global_mask for s in samples])
+    model.mortar.global_encoder.set_feature_stats(*masked_feature_stats(mg, mgm))
+
+    # ITZ encoder: standardize only the leading mix-vector columns; the trailing
+    # geometry / texture / padding columns keep their identity stats.
+    mix_itz = torch.stack([s.mix_itz for s in samples])
+    mean, std = masked_feature_stats(mix_itz, torch.ones_like(mix_itz))
+    enc = model.itz.encoder
+    full_mean = enc.feat_mean[0].clone()
+    full_std = enc.feat_std[0].clone()
+    k = mix_itz.size(1)
+    full_mean[:k] = mean
+    full_std[:k] = std
+    enc.set_feature_stats(full_mean, full_std)
+
+
+def load_strength_checkpoint(
+    base_model: nn.Module,
+    checkpoint_path: Union[str, os.PathLike],
+    map_location: Optional[Union[str, torch.device]] = None,
+    strict: bool = True,
+) -> dict:
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    state = checkpoint.get("base_model_state_dict")
+    if state is None:
+        wrapper_state = checkpoint["model_state_dict"]
+        state = {
+            key.removeprefix("base."): value
+            for key, value in wrapper_state.items()
+            if key.startswith("base.")
+        }
+    base_model.load_state_dict(state, strict=strict)
+    return checkpoint

concrete_gnn/schema.py

@@ -0,0 +1,249 @@
+"""Feature schema for the hierarchical concrete GNN.
+
+A single schema fixes the names, order, and sizes of every feature vector used
+by the graph generator, the missing-feature encoder, the sub-models and the
+mesoscale GNN.  Every vector reserves a small number of explicit placeholder
+slots so additional physical descriptors can be added later without changing
+the encoder dimensions or retraining the input projections.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Tuple
+
+
+NODE_TYPE_AGGREGATE = 0
+NODE_TYPE_MORTAR = 1
+
+EDGE_TYPE_ITZ = 0
+EDGE_TYPE_MORTAR_MORTAR = 1
+EDGE_TYPE_AGGREGATE_AGGREGATE = 2
+
+MORTAR_NODE_TYPE_SAND = 0
+MORTAR_NODE_TYPE_PASTE = 1
+
+
+def _with_placeholders(names: Tuple[str, ...], n: int = 4) -> Tuple[str, ...]:
+    return tuple(names) + tuple(f"_placeholder_{i}" for i in range(n))
+
+
+AGGREGATE_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "aggregate_size_mm",
+        "shape_descriptor",
+        "crushing_strength_MPa",
+        "toughness_J",
+        "abrasion_resistance",
+        "moisture_content",
+        "water_absorption_pct",
+        "surface_texture",
+        "specific_gravity",
+        "bulk_density_kg_m3",
+        "area_fraction",  # node area / RVE area (exact pi r^2 for aggregate disks)
+    ),
+    n=3,
+)
+
+MORTAR_FEATURES: Tuple[str, ...] = (
+    "mortar_compressive_strength_MPa",
+    "mortar_tensile_strength_MPa",
+    "mortar_flexural_strength_MPa",
+    "mortar_elastic_modulus_GPa",
+    "mortar_permeability",
+    "mortar_diffusivity",
+    "mortar_porosity",
+    "mortar_creep_coefficient",
+    # Fibre conditioning (zero = no fibre). Aspect ratio = length / diameter.
+    "fibre_content_kg_m3",
+    "fibre_aspect_ratio",
+    "fibre_tensile_strength_MPa",
+    "fibre_modulus_GPa",
+    # Voronoi-cell area fraction of this mortar region (survives sub-model injection).
+    "mortar_area_fraction",
+    # Fibre material class id (index into categoricals.FIBRE_TYPES). Appended last
+    # so it survives mortar sub-model injection (only mortar_targets slots 0-7 are
+    # overwritten). Steel/PE/PVA differ in stiffness & bond beyond the raw numbers.
+    "fibre_type_id",
+)
+
+ITZ_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "itz_thickness_um",
+        "itz_porosity",
+        "itz_strength_MPa",
+        "itz_elastic_modulus_GPa",
+        "itz_permeability",
+        "itz_diffusivity",
+    )
+)
+
+MORTAR_EDGE_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "edge_distance_mm",
+        "shared_boundary_length_mm",
+        "centroid_dx_mm",
+        "centroid_dy_mm",
+    ),
+    n=2,
+)
+
+AGGREGATE_EDGE_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "agg_agg_surface_gap_mm",
+        "agg_agg_centroid_distance_mm",
+    ),
+    n=2,
+)
+
+GLOBAL_FEATURES: Tuple[str, ...] = (
+    "relative_humidity",
+    "temperature_C",  # curing temperature; observed for UHPC, masked elsewhere
+    "curing_age_days",
+    "water_to_binder_ratio",
+    "cement_content_kg_m3",
+    "scm_fraction",
+    "aggregate_volume_fraction",
+    "fibre_content_kg_m3",  # total fibre dosage; mirror of paste composition pattern
+    # Nominal max grain sizes (mm). Masked when the source does not report them.
+    "max_coarse_aggregate_size_mm",
+    "max_fine_aggregate_size_mm",
+)
+
+CONCRETE_TARGETS: Tuple[str, ...] = (
+    "compressive_strength_MPa",
+    "tensile_strength_MPa",
+    "flexural_strength_MPa",
+    "elastic_modulus_GPa",
+    "permeability",
+    "diffusivity",
+)
+
+MORTAR_TARGETS: Tuple[str, ...] = (
+    "mortar_compressive_strength_MPa",
+    "mortar_tensile_strength_MPa",
+    "mortar_flexural_strength_MPa",
+    "mortar_elastic_modulus_GPa",
+    "mortar_permeability",
+    "mortar_diffusivity",
+    "mortar_porosity",
+    "mortar_creep_coefficient",
+)
+
+ITZ_TARGETS: Tuple[str, ...] = (
+    "itz_thickness_um",
+    "itz_porosity",
+    "itz_strength_MPa",
+    "itz_elastic_modulus_GPa",
+    "itz_permeability",
+    "itz_diffusivity",
+)
+
+SAND_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "sand_size_mm",
+        "sand_gradation",
+        "sand_fineness_modulus",
+        "sand_water_absorption_pct",
+        "sand_compressive_strength_MPa",
+        "sand_elastic_modulus_GPa",
+        "sand_area_fraction",  # exact pi r^2 / mortar-RVE area
+    ),
+    n=3,
+)
+
+CEMENT_PASTE_FEATURES: Tuple[str, ...] = (
+    "cement_type_id",
+    "cement_CaO_pct",
+    "cement_SiO2_pct",
+    "cement_Al2O3_pct",
+    "cement_Fe2O3_pct",
+    "cement_content_kg_m3",
+    "scm_type_id",
+    "scm_CaO_pct",
+    "scm_SiO2_pct",
+    "scm_content_kg_m3",
+    "water_to_binder_ratio",
+    "admixture_dosage",
+    "paste_area_fraction",  # Voronoi-cell area fraction of this paste region
+    # Extra cement oxides (formerly the 3 reserved placeholder slots). Populated
+    # from the UHPC ``Sheet1`` composition table when available, otherwise a
+    # synthetic OPC prior. Alkali is the Na2O equivalent (Na2O + 0.658*K2O).
+    "cement_MgO_pct",
+    "cement_SO3_pct",
+    "cement_alkali_pct",
+    # Extended SCM (binder-blend) oxides, content-weighted across the mix's SCM
+    # ingredients in ``Sheet1``. CaO/SiO2 above carry the dominant signal; these
+    # add the rest of the reactive-phase chemistry.
+    "scm_Al2O3_pct",
+    "scm_Fe2O3_pct",
+    "scm_MgO_pct",
+    "scm_LOI_pct",
+)
+
+MORTAR_GLOBAL_FEATURES: Tuple[str, ...] = _with_placeholders(
+    (
+        "mortar_water_to_binder_ratio",
+        "mortar_cement_content_kg_m3",
+        "mortar_scm_fraction",
+        "mortar_admixture_dosage",
+        "mortar_curing_relative_humidity",
+        "mortar_curing_temperature_C",
+        "mortar_curing_age_days",
+        "mortar_fibre_content_kg_m3",  # same dosage signal at the mortar level
+    ),
+    n=2,
+)
+
+
+@dataclass(frozen=True)
+class SchemaSpec:
+    """Bundle of schema vectors with derived size helpers."""
+
+    aggregate: Tuple[str, ...] = field(default=AGGREGATE_FEATURES)
+    mortar: Tuple[str, ...] = field(default=MORTAR_FEATURES)
+    itz: Tuple[str, ...] = field(default=ITZ_FEATURES)
+    mortar_edge: Tuple[str, ...] = field(default=MORTAR_EDGE_FEATURES)
+    aggregate_edge: Tuple[str, ...] = field(default=AGGREGATE_EDGE_FEATURES)
+    glob: Tuple[str, ...] = field(default=GLOBAL_FEATURES)
+    concrete_targets: Tuple[str, ...] = field(default=CONCRETE_TARGETS)
+    mortar_targets: Tuple[str, ...] = field(default=MORTAR_TARGETS)
+    itz_targets: Tuple[str, ...] = field(default=ITZ_TARGETS)
+    sand: Tuple[str, ...] = field(default=SAND_FEATURES)
+    paste: Tuple[str, ...] = field(default=CEMENT_PASTE_FEATURES)
+    mortar_global: Tuple[str, ...] = field(default=MORTAR_GLOBAL_FEATURES)
+
+    @property
+    def node_pad_dim(self) -> int:
+        return max(len(self.aggregate), len(self.mortar))
+
+    @property
+    def edge_pad_dim(self) -> int:
+        return max(len(self.itz), len(self.mortar_edge), len(self.aggregate_edge))
+
+    @property
+    def mortar_node_pad_dim(self) -> int:
+        return max(len(self.sand), len(self.paste))
+
+    @property
+    def global_dim(self) -> int:
+        return len(self.glob)
+
+    @property
+    def mortar_global_dim(self) -> int:
+        return len(self.mortar_global)
+
+    @property
+    def num_node_types(self) -> int:
+        return 2
+
+    @property
+    def num_edge_types(self) -> int:
+        return 3
+
+    @property
+    def num_mortar_node_types(self) -> int:
+        return 2
+
+
+DEFAULT_SCHEMA = SchemaSpec()

concrete_gnn/train.py

@@ -0,0 +1,316 @@
+"""Training utilities for the hierarchical concrete GNN.
+
+The training loop supports:
+
+* MSE on standardised targets via :class:`MultiTaskLoss` (per-target scaling),
+* optional auxiliary losses on mortar / ITZ sub-model outputs when ground-truth
+  sub-labels are available,
+* gradient clipping (off by default; set ``TrainConfig.grad_clip > 0``),
+* per-target R\\ :sup:`2` reporting on a held-out loader.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional
+
+import torch
+from torch import Tensor, nn
+
+from .data import MultiscaleBatch
+
+
+# ---------------------------------------------------------------------------
+# Loss / config
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class MultiTaskLoss:
+    """Per-target standardised MSE.
+
+    ``target_scales`` should be precomputed on the training set so val / test
+    losses are reported in the same units as training (i.e. relative
+    deviations).  If ``None``, the per-batch absolute mean is used as a
+    fallback - useful for ad-hoc experiments but not for cross-run comparison.
+    """
+
+    target_scales: Optional[Tensor] = None
+    weights: Optional[Tensor] = None
+
+    def __call__(
+        self,
+        pred: Tensor,
+        target: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        scales = self.target_scales
+        if scales is None:
+            scales = target.detach().abs().mean(dim=0).clamp_min(1e-6)
+        diff = (pred - target) / scales
+        sq = diff.pow(2)
+        if mask is not None:
+            sq = sq * mask
+            denom = mask.sum().clamp_min(1.0)
+            sq = sq.sum() / denom
+        elif self.weights is not None:
+            sq = (sq * self.weights).mean()
+        else:
+            sq = sq.mean()
+        return sq
+
+
+@dataclass
+class TrainConfig:
+    epochs: int = 60
+    lr: float = 1.0e-3
+    weight_decay: float = 1.0e-4
+    grad_clip: float = 0.0
+    aux_mortar_weight: float = 0.0
+    aux_itz_weight: float = 0.0
+    log_every: int = 5
+    name: str = "model"
+    checkpoint_path: Optional[str] = None  # save best-val-loss weights here
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_target_scale(loader: Iterable[MultiscaleBatch], device: torch.device) -> Tensor:
+    """Per-target std on the loader; used as the loss normaliser."""
+
+    ys = [batch.concrete_target for batch in loader]
+    y = torch.cat(ys, dim=0)
+    return y.std(dim=0).clamp_min(1.0e-12).to(device)
+
+
+def compute_sublabel_scales(
+    loader: Iterable[MultiscaleBatch], device: torch.device
+) -> Dict[str, Optional[Tensor]]:
+    """Per-target std for the auxiliary mortar / ITZ labels (None when absent)."""
+
+    mortar_targets, itz_targets = [], []
+    for batch in loader:
+        if batch.mortar_target is not None:
+            mortar_targets.append(batch.mortar_target)
+        if batch.itz_target is not None:
+            itz_targets.append(batch.itz_target)
+    mortar_scale = (
+        torch.cat(mortar_targets, dim=0).std(dim=0).clamp_min(1.0e-12).to(device)
+        if mortar_targets
+        else None
+    )
+    itz_scale = (
+        torch.cat(itz_targets, dim=0).std(dim=0).clamp_min(1.0e-12).to(device)
+        if itz_targets
+        else None
+    )
+    return {"mortar": mortar_scale, "itz": itz_scale}
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+
+
+def train_one_epoch(
+    model: nn.Module,
+    loader: Iterable[MultiscaleBatch],
+    optimiser: torch.optim.Optimizer,
+    loss_fn: MultiTaskLoss,
+    config: TrainConfig,
+    device: torch.device,
+    mortar_loss_fn: Optional[MultiTaskLoss] = None,
+    itz_loss_fn: Optional[MultiTaskLoss] = None,
+) -> Dict[str, float]:
+    if mortar_loss_fn is None:
+        mortar_loss_fn = MultiTaskLoss()
+    if itz_loss_fn is None:
+        itz_loss_fn = MultiTaskLoss()
+
+    model.train()
+    running = {"loss": 0.0, "concrete": 0.0, "mortar": 0.0, "itz": 0.0, "n": 0}
+
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch)
+        concrete_pred = out["concrete_pred"]
+        concrete_loss = loss_fn(concrete_pred, batch.concrete_target)
+
+        total = concrete_loss
+        mortar_loss_val = 0.0
+        itz_loss_val = 0.0
+
+        if (
+            config.aux_mortar_weight > 0
+            and batch.mortar_target is not None
+            and "mortar_pred" in out
+        ):
+            mortar_loss = mortar_loss_fn(out["mortar_pred"], batch.mortar_target)
+            total = total + config.aux_mortar_weight * mortar_loss
+            mortar_loss_val = float(mortar_loss.detach().cpu())
+
+        if (
+            config.aux_itz_weight > 0
+            and batch.itz_target is not None
+            and "itz_pred" in out
+            and out["itz_pred"].size(0) == batch.itz_target.size(0)
+        ):
+            itz_loss = itz_loss_fn(out["itz_pred"], batch.itz_target)
+            total = total + config.aux_itz_weight * itz_loss
+            itz_loss_val = float(itz_loss.detach().cpu())
+
+        optimiser.zero_grad()
+        total.backward()
+        if config.grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip)
+        optimiser.step()
+
+        running["loss"] += float(total.detach().cpu())
+        running["concrete"] += float(concrete_loss.detach().cpu())
+        running["mortar"] += mortar_loss_val
+        running["itz"] += itz_loss_val
+        running["n"] += 1
+
+    n = max(running["n"], 1)
+    return {
+        "loss": running["loss"] / n,
+        "concrete": running["concrete"] / n,
+        "mortar": running["mortar"] / n,
+        "itz": running["itz"] / n,
+    }
+
+
+@torch.no_grad()
+def evaluate(
+    model: nn.Module,
+    loader: Iterable[MultiscaleBatch],
+    loss_fn: MultiTaskLoss,
+    device: torch.device,
+) -> Dict[str, float]:
+    model.eval()
+    running = {"concrete": 0.0, "n": 0}
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch)
+        running["concrete"] += float(loss_fn(out["concrete_pred"], batch.concrete_target).cpu())
+        running["n"] += 1
+    return {"concrete": running["concrete"] / max(running["n"], 1)}
+
+
+@torch.no_grad()
+def evaluate_metrics(
+    model: nn.Module,
+    loader: Iterable[MultiscaleBatch],
+    device: torch.device,
+) -> Dict[str, Tensor]:
+    """Return per-target R^2 and MAE on the loader."""
+
+    model.eval()
+    preds: List[Tensor] = []
+    trues: List[Tensor] = []
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch)
+        preds.append(out["concrete_pred"].cpu())
+        trues.append(batch.concrete_target.cpu())
+    pred = torch.cat(preds, dim=0)
+    true = torch.cat(trues, dim=0)
+    ss_res = ((true - pred) ** 2).sum(dim=0)
+    ss_tot = ((true - true.mean(dim=0)) ** 2).sum(dim=0).clamp_min(1.0e-12)
+    r2 = 1.0 - ss_res / ss_tot
+    mae = (true - pred).abs().mean(dim=0)
+    return {"r2": r2, "mae": mae}
+
+
+def train(
+    model: nn.Module,
+    train_loader: Iterable[MultiscaleBatch],
+    config: TrainConfig = TrainConfig(),
+    val_loader: Optional[Iterable[MultiscaleBatch]] = None,
+    device: Optional[torch.device] = None,
+    loss_fn: Optional[MultiTaskLoss] = None,
+    mortar_loss_fn: Optional[MultiTaskLoss] = None,
+    itz_loss_fn: Optional[MultiTaskLoss] = None,
+) -> List[Dict[str, float]]:
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    if loss_fn is None:
+        loss_fn = MultiTaskLoss()
+
+    optimiser = torch.optim.AdamW(
+        model.parameters(), lr=config.lr, weight_decay=config.weight_decay
+    )
+
+    history: List[Dict[str, float]] = []
+    best_val = math.inf
+    best_epoch = -1
+    if config.checkpoint_path is not None:
+        os.makedirs(os.path.dirname(os.path.abspath(config.checkpoint_path)) or ".", exist_ok=True)
+
+    for epoch in range(1, config.epochs + 1):
+        stats = train_one_epoch(
+            model,
+            train_loader,
+            optimiser,
+            loss_fn,
+            config,
+            device,
+            mortar_loss_fn=mortar_loss_fn,
+            itz_loss_fn=itz_loss_fn,
+        )
+        if val_loader is not None:
+            val_stats = evaluate(model, val_loader, loss_fn, device)
+            stats["val_concrete"] = val_stats["concrete"]
+            if config.checkpoint_path is not None and val_stats["concrete"] < best_val:
+                best_val = val_stats["concrete"]
+                best_epoch = epoch
+                torch.save(
+                    {
+                        "epoch": epoch,
+                        "val_concrete": best_val,
+                        "model_state_dict": model.state_dict(),
+                        "config_name": config.name,
+                    },
+                    config.checkpoint_path,
+                )
+        history.append(stats)
+        if epoch % max(1, config.log_every) == 0 or epoch == 1:
+            msg = (
+                f"[{config.name}] epoch={epoch:03d} "
+                f"loss={stats['loss']:.4f} "
+                f"concrete={stats['concrete']:.4f} "
+                f"mortar={stats['mortar']:.4f} "
+                f"itz={stats['itz']:.4f}"
+            )
+            if "val_concrete" in stats:
+                msg += f"  val={stats['val_concrete']:.4f}"
+            print(msg)
+
+    if config.checkpoint_path is not None and best_epoch > 0:
+        print(
+            f"[{config.name}] best val_concrete={best_val:.4f} at epoch {best_epoch} "
+            f"-> saved to {config.checkpoint_path}"
+        )
+    return history
+
+
+def load_best_checkpoint(
+    model: nn.Module, path: str, device: Optional[torch.device] = None
+) -> Dict[str, float]:
+    """Load best-val-loss weights into ``model`` in place and return metadata."""
+
+    if device is None:
+        device = next(model.parameters()).device
+    state = torch.load(path, map_location=device, weights_only=False)
+    model.load_state_dict(state["model_state_dict"])
+    return {
+        "epoch": state.get("epoch", -1),
+        "val_concrete": state.get("val_concrete", float("nan")),
+        "config_name": state.get("config_name", ""),
+    }

concrete_gnn/visualize.py

@@ -0,0 +1,329 @@
+"""Plotting helpers for model diagnostics.
+
+Four figures:
+    * :func:`plot_parity`           - test-set pred vs true per target.
+    * :func:`plot_r2_bars`          - per-target R^2 across models.
+    * :func:`plot_training_curves`  - train + val concrete loss vs epoch.
+    * :func:`plot_rve_graph`        - one generated 2D concrete RVE.
+
+All functions accept a ``save_path``; if provided the figure is written and
+the matplotlib ``Figure`` is also returned for further customisation.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, Iterable, List, Optional, Sequence
+
+import matplotlib
+
+matplotlib.use("Agg")  # safe headless backend
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .graph_generator import ConcreteGraph
+from .schema import (
+    CONCRETE_TARGETS,
+    EDGE_TYPE_AGGREGATE_AGGREGATE,
+    EDGE_TYPE_ITZ,
+    EDGE_TYPE_MORTAR_MORTAR,
+    NODE_TYPE_AGGREGATE,
+    NODE_TYPE_MORTAR,
+)
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _ensure_dir(path: str) -> None:
+    parent = os.path.dirname(os.path.abspath(path))
+    os.makedirs(parent, exist_ok=True)
+
+
+def _r2(true: np.ndarray, pred: np.ndarray) -> float:
+    ss_res = float(np.sum((true - pred) ** 2))
+    ss_tot = float(np.sum((true - true.mean()) ** 2))
+    if ss_tot < 1e-20:
+        return float("nan")
+    return 1.0 - ss_res / ss_tot
+
+
+# ---------------------------------------------------------------------------
+# 1. Parity plots
+# ---------------------------------------------------------------------------
+
+
+def plot_parity(
+    predictions: Dict[str, np.ndarray],
+    truths: np.ndarray,
+    model_order: Sequence[str],
+    save_path: Optional[str] = None,
+    target_names: Sequence[str] = CONCRETE_TARGETS,
+):
+    """One subplot per target, one row per model.
+
+    Parameters
+    ----------
+    predictions: dict
+        Mapping ``model_name -> (N, T)`` test predictions (numpy).
+    truths: ndarray
+        Test ground-truth array of shape ``(N, T)``.
+    model_order: sequence of str
+        Which models (and in which order) to draw.
+    """
+
+    n_models = len(model_order)
+    n_targets = len(target_names)
+    fig, axes = plt.subplots(
+        n_models,
+        n_targets,
+        figsize=(2.6 * n_targets, 2.6 * n_models),
+        squeeze=False,
+    )
+    for i, name in enumerate(model_order):
+        pred = predictions[name]
+        for t in range(n_targets):
+            ax = axes[i, t]
+            y_true = truths[:, t]
+            y_pred = pred[:, t]
+            r2 = _r2(y_true, y_pred)
+            lo = float(min(y_true.min(), y_pred.min()))
+            hi = float(max(y_true.max(), y_pred.max()))
+            span = hi - lo if hi > lo else max(1.0, abs(hi))
+            lo -= 0.05 * span
+            hi += 0.05 * span
+
+            ax.plot([lo, hi], [lo, hi], "k--", lw=0.8, alpha=0.6)
+            ax.scatter(y_true, y_pred, s=22, alpha=0.7, edgecolor="none")
+            ax.set_xlim(lo, hi)
+            ax.set_ylim(lo, hi)
+            ax.set_aspect("equal", adjustable="box")
+            ax.tick_params(labelsize=7)
+            ax.text(
+                0.04,
+                0.93,
+                f"$R^2$={r2:.2f}",
+                transform=ax.transAxes,
+                fontsize=8,
+                va="top",
+            )
+            if i == n_models - 1:
+                ax.set_xlabel(target_names[t], fontsize=8)
+            if t == 0:
+                ax.set_ylabel(f"{name}\npred", fontsize=8)
+    fig.suptitle("Test-set parity (true on x, predicted on y)", fontsize=11)
+    fig.tight_layout(rect=(0, 0, 1, 0.97))
+    if save_path is not None:
+        _ensure_dir(save_path)
+        fig.savefig(save_path, dpi=160)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# 2. Per-target R^2 bars
+# ---------------------------------------------------------------------------
+
+
+def plot_r2_bars(
+    r2_by_model: Dict[str, np.ndarray],
+    model_order: Sequence[str],
+    save_path: Optional[str] = None,
+    target_names: Sequence[str] = CONCRETE_TARGETS,
+):
+    n_targets = len(target_names)
+    n_models = len(model_order)
+    x = np.arange(n_targets)
+    width = 0.8 / max(1, n_models)
+
+    fig, ax = plt.subplots(figsize=(1.2 * n_targets + 2.0, 4.0))
+    colours = plt.cm.tab10(np.linspace(0, 1, max(3, n_models)))
+    for i, name in enumerate(model_order):
+        r2 = np.asarray(r2_by_model[name])
+        offset = (i - (n_models - 1) / 2.0) * width
+        ax.bar(x + offset, r2, width=width * 0.95, label=name, color=colours[i])
+    ax.axhline(0.0, color="k", lw=0.6, alpha=0.5)
+    ax.set_xticks(x)
+    ax.set_xticklabels(target_names, rotation=30, ha="right", fontsize=8)
+    ax.set_ylabel("Test-set $R^2$")
+    ax.set_title("Per-target $R^2$ by model")
+    ax.legend(fontsize=8, frameon=False)
+    ax.set_ylim(min(-0.2, min((r2_by_model[n].min() for n in model_order))) - 0.05, 1.05)
+    fig.tight_layout()
+    if save_path is not None:
+        _ensure_dir(save_path)
+        fig.savefig(save_path, dpi=160)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# 3. Training curves
+# ---------------------------------------------------------------------------
+
+
+def plot_training_curves(
+    histories: Dict[str, List[Dict[str, float]]],
+    model_order: Sequence[str],
+    save_path: Optional[str] = None,
+):
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
+    train_ax, val_ax = axes
+    colours = plt.cm.tab10(np.linspace(0, 1, max(3, len(model_order))))
+
+    for i, name in enumerate(model_order):
+        history = histories[name]
+        epochs = np.arange(1, len(history) + 1)
+        train_loss = np.asarray([h["concrete"] for h in history])
+        train_ax.plot(epochs, train_loss, label=name, color=colours[i], lw=1.4)
+        if "val_concrete" in history[0]:
+            val_loss = np.asarray([h.get("val_concrete", np.nan) for h in history])
+            val_ax.plot(epochs, val_loss, label=name, color=colours[i], lw=1.4)
+            if np.isfinite(val_loss).any():
+                best_idx = int(np.nanargmin(val_loss))
+                val_ax.scatter(
+                    [epochs[best_idx]],
+                    [val_loss[best_idx]],
+                    color=colours[i],
+                    edgecolor="black",
+                    zorder=5,
+                    s=42,
+                )
+
+    for ax, title in zip(axes, ("train concrete loss", "val concrete loss")):
+        ax.set_yscale("log")
+        ax.set_xlabel("epoch")
+        ax.set_ylabel("normalised MSE")
+        ax.set_title(title)
+        ax.grid(True, alpha=0.3)
+        ax.legend(fontsize=8, frameon=False)
+    fig.suptitle("Training curves (best-val checkpoint marked)", fontsize=11)
+    fig.tight_layout(rect=(0, 0, 1, 0.95))
+    if save_path is not None:
+        _ensure_dir(save_path)
+        fig.savefig(save_path, dpi=160)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# 4. Concrete RVE graph visualization
+# ---------------------------------------------------------------------------
+
+
+def plot_rve_graph(
+    graph: ConcreteGraph,
+    rve_size_mm: float = 150.0,
+    save_path: Optional[str] = None,
+):
+    """Draw aggregates as circles, mortar nodes as dots, edges coloured by type."""
+
+    pos = graph.pos.detach().cpu().numpy()
+    node_type = graph.node_type.detach().cpu().numpy()
+    edge_index = graph.edge_index.detach().cpu().numpy()
+    edge_type = graph.edge_type.detach().cpu().numpy()
+    x = graph.x.detach().cpu().numpy()
+
+    fig, ax = plt.subplots(figsize=(7, 7))
+
+    # Edges first (so nodes draw on top).
+    seen = set()
+    edge_styles = {
+        EDGE_TYPE_ITZ: dict(color="#d62728", alpha=0.6, lw=0.8, ls="-", label="ITZ"),
+        EDGE_TYPE_MORTAR_MORTAR: dict(
+            color="#7f7f7f", alpha=0.45, lw=0.6, ls="-", label="mortar-mortar"
+        ),
+        EDGE_TYPE_AGGREGATE_AGGREGATE: dict(
+            color="#1f77b4", alpha=0.6, lw=0.7, ls=":", label="agg-agg"
+        ),
+    }
+    legend_seen = set()
+    for k in range(edge_index.shape[1]):
+        a = int(edge_index[0, k])
+        b = int(edge_index[1, k])
+        if a == b:
+            continue
+        key = (min(a, b), max(a, b))
+        if key in seen:
+            continue
+        seen.add(key)
+        style = edge_styles[int(edge_type[k])]
+        label = style["label"] if int(edge_type[k]) not in legend_seen else None
+        legend_seen.add(int(edge_type[k]))
+        ax.plot(
+            [pos[a, 0], pos[b, 0]],
+            [pos[a, 1], pos[b, 1]],
+            color=style["color"],
+            alpha=style["alpha"],
+            lw=style["lw"],
+            ls=style["ls"],
+            label=label,
+        )
+
+    # Aggregate circles (diameter is the first feature channel).
+    for i in range(pos.shape[0]):
+        if int(node_type[i]) == NODE_TYPE_AGGREGATE:
+            diameter = float(x[i, 0])
+            radius = max(1.0, 0.5 * diameter)
+            ax.add_patch(
+                plt.Circle(
+                    (pos[i, 0], pos[i, 1]),
+                    radius=radius,
+                    facecolor="#bdbdbd",
+                    edgecolor="black",
+                    lw=0.6,
+                    alpha=0.85,
+                    zorder=2,
+                )
+            )
+
+    mortar_pts = pos[node_type == NODE_TYPE_MORTAR]
+    ax.scatter(
+        mortar_pts[:, 0],
+        mortar_pts[:, 1],
+        s=18,
+        c="#2ca02c",
+        edgecolor="black",
+        lw=0.3,
+        zorder=3,
+        label="mortar patch",
+    )
+
+    ax.set_xlim(0, rve_size_mm)
+    ax.set_ylim(0, rve_size_mm)
+    ax.set_aspect("equal")
+    ax.set_xlabel("x [mm]")
+    ax.set_ylabel("y [mm]")
+    ax.set_title("Generated 2D concrete RVE (aggregates + mortar + edge types)")
+    ax.legend(fontsize=8, loc="upper right", frameon=True)
+    fig.tight_layout()
+    if save_path is not None:
+        _ensure_dir(save_path)
+        fig.savefig(save_path, dpi=160)
+    return fig
+
+
+# ---------------------------------------------------------------------------
+# Prediction collection helper
+# ---------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def collect_predictions(
+    model: torch.nn.Module,
+    loader: Iterable,
+    device: torch.device,
+) -> Dict[str, np.ndarray]:
+    """Return ``{"pred": (N, T), "true": (N, T)}`` arrays on CPU."""
+
+    model.eval()
+    preds, trues = [], []
+    for batch in loader:
+        batch = batch.to(device)
+        out = model(batch)
+        preds.append(out["concrete_pred"].cpu())
+        trues.append(batch.concrete_target.cpu())
+    return {
+        "pred": torch.cat(preds, dim=0).numpy(),
+        "true": torch.cat(trues, dim=0).numpy(),
+    }

inference.py

@@ -0,0 +1,441 @@
+"""Inference layer for the deployed concrete-strength app.
+
+Single source of truth for:
+  * loading the trained hierarchical GNN + tabular-ANN checkpoints,
+  * Function 1 (forward): a (possibly incomplete) mix design -> compressive strength,
+  * Function 2 (inverse): a target strength -> several detailed mix designs.
+
+Imported by ``streamlit_app.py`` and runnable as a CLI for verification:
+
+    # forward (fields left out are treated as "not measured")
+    python app/inference.py predict --cement 500 --water 175 --age 28 \
+        --coarse 950 --fine 750
+
+    # inverse
+    python app/inference.py suggest --target 120 --k 5
+    python app/inference.py suggest --target 40 --k 5 --no-fibre
+
+The checkpoints were trained with the *curing-only* schema and the
+``mortar_capacity`` strength head; that is NOT stored in the checkpoint, so we
+rebuild the model with exactly those settings here (see CURING_ONLY_* below).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+
+# --- make the concrete_gnn package importable whether we run from the repo
+#     (package lives in ../Hybrid) or from a self-contained HF Space bundle
+#     (package copied next to this file). -------------------------------------
+_HERE = Path(__file__).resolve().parent
+_PKG_ROOT = None
+for _cand in (_HERE, _HERE.parent / "Hybrid", _HERE.parent):
+    if (_cand / "concrete_gnn" / "__init__.py").exists():
+        sys.path.insert(0, str(_cand))
+        _PKG_ROOT = _cand
+        break
+if _PKG_ROOT is None:
+    raise RuntimeError(
+        "Missing local concrete_gnn package. Deploy the contents of "
+        "app/space_bundle, not app/ alone; the Space root must contain "
+        "concrete_gnn/__init__.py next to streamlit_app.py."
+    )
+
+from concrete_gnn import (  # noqa: E402
+    ConcreteMixDataset,
+    IntegratedMultiscaleModel,
+    SchemaSpec,
+    Standardizer,
+    StrengthHead,
+    TABULAR_DIM,
+    TabularANN,
+    collate_multiscale,
+    collect_strength_predictions,
+    fit_encoder_standardizers,
+    strength_column,
+)
+
+# ---------------------------------------------------------------------------
+# Schema used at training time (curing-only globals). MUST match train_real_data.py.
+# ---------------------------------------------------------------------------
+CURING_ONLY_GLOBALS = (
+    "relative_humidity", "temperature_C", "curing_age_days",
+    "_placeholder_0", "_placeholder_1",
+)
+CURING_ONLY_MORTAR_GLOBALS = (
+    "mortar_curing_relative_humidity", "mortar_curing_temperature_C",
+    "mortar_curing_age_days", "_placeholder_0", "_placeholder_1",
+)
+
+
+def make_schema() -> SchemaSpec:
+    return SchemaSpec(glob=CURING_ONLY_GLOBALS, mortar_global=CURING_ONLY_MORTAR_GLOBALS)
+
+
+# ---------------------------------------------------------------------------
+# Column layout (normalized names the ConcreteMixDataset loader expects).
+# ---------------------------------------------------------------------------
+TARGET_COL = "compressive_strength_mpa"
+
+# Binder + water + aggregate + fibre amounts: a missing value means "none of
+# this component", i.e. a genuine 0.0 (not a masked/unknown channel).
+AMOUNT_COLS = [
+    "cement_kg_m3", "slag_kg_m3", "fly_ash_kg_m3", "silica_fume_kg_m3",
+    "metakaolin_kg_m3", "limestone_powder_kg_m3", "other_scm_kg_m3",
+    "water_kg_m3", "superplasticizer_kg_m3",
+    "coarse_aggregate_kg_m3", "fine_aggregate_kg_m3",
+    "fibre_content_kg_m3", "fibre_length_mm", "fibre_diameter_mm",
+    "fibre_tensile_strength_mpa", "fibre_modulus_gpa",
+]
+# Maskable descriptors: a missing value means "not reported" -> NaN -> masked.
+MASKABLE_COLS = [
+    "max_coarse_aggregate_size_mm", "max_fine_aggregate_size_mm",
+    "curing_temperature_c",
+    "cement_CaO_pct", "cement_SiO2_pct", "cement_Al2O3_pct", "cement_Fe2O3_pct",
+    "cement_MgO_pct", "cement_SO3_pct", "cement_alkali_pct", "cement_LOI_pct",
+    "scm_CaO_pct", "scm_SiO2_pct", "scm_Al2O3_pct", "scm_Fe2O3_pct",
+    "scm_MgO_pct", "scm_LOI_pct",
+    "cement_grade_mpa", "curing_humidity_pct", "specimen_size_mm",
+]
+CATEGORICAL_COLS = ["cement_type_norm", "fibre_type_norm", "curing_regime_norm"]
+
+# age has its own default (0 d would imply no strength); everything else 0.0.
+AGE_COL = "age_days"
+DEFAULT_AGE = 28.0
+
+# Design knobs the inverse "refine" step is allowed to perturb.
+REFINE_KNOBS = [
+    "cement_kg_m3", "water_kg_m3", "superplasticizer_kg_m3",
+    "slag_kg_m3", "fly_ash_kg_m3", "silica_fume_kg_m3",
+]
+# Columns shown / clustered for inverse diversity (the design vector).
+DESIGN_COLS = AMOUNT_COLS + [AGE_COL]
+
+
+def build_input_frame(mixes: Sequence[Dict[str, float]]) -> pd.DataFrame:
+    """Turn a list of partial mix dicts into a full normalized DataFrame.
+
+    Unspecified amount columns -> 0.0 (component absent); unspecified maskable
+    descriptors / chemistry -> NaN (not reported -> masked); unspecified
+    categoricals -> NaN (default bucket); age defaults to 28 d.
+    """
+    rows = []
+    for mix in mixes:
+        row: Dict[str, object] = {}
+        for c in AMOUNT_COLS:
+            row[c] = float(mix.get(c, 0.0) or 0.0)
+        row[AGE_COL] = float(mix.get(AGE_COL, DEFAULT_AGE) or DEFAULT_AGE)
+        for c in MASKABLE_COLS:
+            v = mix.get(c, None)
+            row[c] = np.nan if v is None or v == "" else float(v)
+        for c in CATEGORICAL_COLS:
+            v = mix.get(c, None)
+            row[c] = np.nan if v is None or v == "" else str(v)
+        row[TARGET_COL] = float(mix.get(TARGET_COL, 0.0) or 0.0)  # placeholder
+        rows.append(row)
+    return pd.DataFrame(rows)
+
+
+def add_derived(df: pd.DataFrame) -> pd.DataFrame:
+    """Append water/binder ratio and SCM fraction for display."""
+    df = df.copy()
+    binder = (
+        df["cement_kg_m3"] + df["slag_kg_m3"] + df["fly_ash_kg_m3"]
+        + df["silica_fume_kg_m3"] + df["metakaolin_kg_m3"]
+        + df["limestone_powder_kg_m3"] + df["other_scm_kg_m3"]
+    ).clip(lower=1.0)
+    df["water_binder_ratio"] = (df["water_kg_m3"] / binder).round(3)
+    df["scm_fraction"] = ((binder - df["cement_kg_m3"]) / binder).round(3)
+    return df
+
+
+# ---------------------------------------------------------------------------
+# Predictor
+# ---------------------------------------------------------------------------
+def _detect_head_kind(state: dict) -> str:
+    keys = list(state.keys())
+    if any(k.startswith("strength_head.mortar_eff") for k in keys):
+        return "mortar_capacity"
+    if any(k.startswith("strength_head.") for k in keys):
+        return "physics"
+    return "free"
+
+
+@dataclass
+class Predictor:
+    checkpoint_dir: Path
+    device: torch.device = field(default_factory=lambda: torch.device("cpu"))
+    seed: int = 17
+    config: dict = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        self.schema = make_schema()
+        cfg_path = Path(self.checkpoint_dir) / "model_config.json"
+        if cfg_path.exists():
+            self.config = json.loads(cfg_path.read_text())
+        self.gnn, self.std = self._load("hierarchical.pt", kind="gnn")
+        self.tabular, _ = self._load("tabular_ann.pt", kind="tabular")
+        self.bounds: dict = self.config.get("feature_bounds", {})
+
+    def _load(self, fname: str, kind: str):
+        path = Path(self.checkpoint_dir) / fname
+        ck = torch.load(path, map_location=self.device, weights_only=False)
+        std = Standardizer(
+            mean=torch.tensor(float(ck["target_mean"])),
+            std=torch.tensor(float(ck["target_std"])),
+        )
+        state = ck["base_model_state_dict"]
+        if kind == "gnn":
+            head_kind = self.config.get("strength_head_kind") or _detect_head_kind(state)
+            base = IntegratedMultiscaleModel(schema=self.schema, strength_head_kind=head_kind)
+        else:
+            base = TabularANN(in_dim=TABULAR_DIM, schema=self.schema)
+        base.load_state_dict(state)
+        model = StrengthHead(base, std).to(self.device).eval()
+        return model, std
+
+    # ---- forward ---- #
+    @torch.no_grad()
+    def predict_df(self, df: pd.DataFrame, which=("gnn", "tabular")) -> Dict[str, np.ndarray]:
+        df = df.copy()
+        if TARGET_COL not in df.columns:
+            df[TARGET_COL] = 0.0
+        ds = ConcreteMixDataset(df, None, 0.0, self.seed, self.std, schema=self.schema)
+        dl = DataLoader(ds, batch_size=128, shuffle=False, collate_fn=collate_multiscale)
+        models = {"gnn": self.gnn, "tabular": self.tabular}
+        out: Dict[str, np.ndarray] = {}
+        for name in which:
+            pred, _ = collect_strength_predictions(models[name], dl, self.device)
+            out[name] = pred.numpy()
+        return out
+
+    def predict_strength(self, mix: Dict[str, float]) -> Dict[str, float]:
+        """Single partial mix -> {'gnn': MPa, 'tabular': MPa}."""
+        df = build_input_frame([mix])
+        res = self.predict_df(df)
+        return {k: float(v[0]) for k, v in res.items()}
+
+    def age_curve(self, mix: Dict[str, float], ages=None) -> pd.DataFrame:
+        """Predict strength across curing ages holding the rest of the mix fixed."""
+        if ages is None:
+            ages = [1, 3, 7, 14, 28, 56, 90, 180, 365]
+        mixes = [{**mix, AGE_COL: float(a)} for a in ages]
+        res = self.predict_df(build_input_frame(mixes))
+        return pd.DataFrame({"age_days": ages, "gnn_mpa": res["gnn"], "tabular_mpa": res["tabular"]})
+
+    def out_of_range(self, mix: Dict[str, float]) -> List[str]:
+        """Names of supplied fields that fall outside the training p1-p99 range."""
+        flags = []
+        for c, v in mix.items():
+            b = self.bounds.get(c)
+            if b and v not in (None, "") and isinstance(v, (int, float)):
+                if v < b["p01"] or v > b["p99"]:
+                    flags.append(c)
+        return flags
+
+    # ---- inverse ---- #
+    def suggest_mixes(
+        self,
+        target: float,
+        index: pd.DataFrame,
+        k: int = 5,
+        tol: Optional[float] = None,
+        allow_fibre: bool = True,
+        allow_scm: bool = True,
+        age: Optional[float] = None,
+        domain: str = "any",  # "any" | "uhpc" | "normal"
+        refine: bool = True,
+    ) -> pd.DataFrame:
+        df = index.copy()
+        if not allow_fibre:
+            df = df[df["fibre_content_kg_m3"].fillna(0) <= 0]
+        if not allow_scm:
+            scm = (df[["slag_kg_m3", "fly_ash_kg_m3", "silica_fume_kg_m3",
+                       "metakaolin_kg_m3", "limestone_powder_kg_m3", "other_scm_kg_m3"]]
+                   .fillna(0).sum(axis=1))
+            df = df[scm <= 0]
+        if domain == "uhpc":
+            df = df[df["source"] == "UHPC"]
+        elif domain == "normal":
+            df = df[df["source"] != "UHPC"]
+        if age is not None:
+            df = df[np.isclose(pd.to_numeric(df["age_days"], errors="coerce"), age)]
+        if df.empty:
+            return df
+
+        df = df.copy()
+        df["err"] = (df["pred_gnn"] - target).abs()
+        # Keep the candidate pool within a tolerance band so even the most
+        # "diverse" pick stays near the target; widen if too few rows qualify.
+        base_tol = tol if tol is not None else max(5.0, 0.10 * target)
+        need = max(k * 4, 12)
+        pool, mult = df[df["err"] <= base_tol], 1.0
+        while len(pool) < need and mult < 8:
+            mult *= 2
+            pool = df[df["err"] <= base_tol * mult]
+        if pool.empty:
+            pool = df.nsmallest(need, "err")
+        pool = pool.sort_values("err").head(200)
+
+        seeds = self._diversify(pool, k)
+        if refine and len(seeds) > 0:
+            seeds = self._refine(seeds, target)
+        out = self._assemble(seeds, target)
+        # Return the k closest-to-target after refinement.
+        order = out["pred_gnn"].sub(target).abs().sort_values().index
+        return out.loc[order].head(k).reset_index(drop=True)
+
+    def _diversify(self, df: pd.DataFrame, k: int) -> pd.DataFrame:
+        """Greedy max-min selection over the normalized design vector."""
+        if len(df) <= k:
+            return df
+        X = df[DESIGN_COLS].fillna(0.0).to_numpy(dtype=float)
+        mu, sigma = X.mean(0), X.std(0) + 1e-6
+        Z = (X - mu) / sigma
+        chosen = [0]  # closest-to-target row is first (df is sorted by err)
+        while len(chosen) < k:
+            d = np.min(
+                np.linalg.norm(Z[:, None, :] - Z[chosen][None, :, :], axis=2), axis=1
+            )
+            d[chosen] = -1.0
+            chosen.append(int(np.argmax(d)))
+        return df.iloc[chosen]
+
+    def _refine(self, seeds: pd.DataFrame, target: float, n_per_seed: int = 32) -> pd.DataFrame:
+        """Perturb a few knobs per seed, predict, keep the closest-to-target variant."""
+        rng = np.random.default_rng(self.seed)
+        candidates: List[Dict[str, float]] = []
+        owners: List[int] = []
+        seed_rows = seeds.reset_index(drop=True)
+        for i, row in seed_rows.iterrows():
+            base = {c: (None if pd.isna(row[c]) else row[c]) for c in
+                    AMOUNT_COLS + [AGE_COL] + MASKABLE_COLS + CATEGORICAL_COLS}
+            candidates.append(dict(base)); owners.append(i)  # keep the seed itself
+            for _ in range(n_per_seed):
+                cand = dict(base)
+                for knob in REFINE_KNOBS:
+                    cur = float(base.get(knob) or 0.0)
+                    if cur <= 0 and knob in ("slag_kg_m3", "fly_ash_kg_m3", "silica_fume_kg_m3"):
+                        continue  # don't invent an SCM that wasn't there
+                    factor = float(rng.uniform(0.85, 1.15))
+                    cand[knob] = self._clip(knob, cur * factor)
+                candidates.append(cand); owners.append(i)
+        preds = self.predict_df(build_input_frame(candidates), which=("gnn",))["gnn"]
+        owners = np.asarray(owners)
+        best_rows = []
+        for i in range(len(seed_rows)):
+            sel = np.where(owners == i)[0]
+            err = np.abs(preds[sel] - target)
+            best = sel[int(np.argmin(err))]
+            best_rows.append({**candidates[best], "pred_gnn": float(preds[best]),
+                              "source": seed_rows.loc[i].get("source", "refined"),
+                              "measured": seed_rows.loc[i].get("measured", np.nan)})
+        return pd.DataFrame(best_rows)
+
+    def _clip(self, col: str, value: float) -> float:
+        b = self.bounds.get(col)
+        if b:
+            return float(min(max(value, b["min"]), b["max"]))
+        return float(max(value, 0.0))
+
+    def _assemble(self, seeds: pd.DataFrame, target: float) -> pd.DataFrame:
+        df = seeds.copy().reset_index(drop=True)
+        if "pred_gnn" not in df.columns:
+            df["pred_gnn"] = self.predict_df(build_input_frame(df.to_dict("records")),
+                                             which=("gnn",))["gnn"]
+        df = add_derived(df)
+        for c in DESIGN_COLS:  # tidy kg/mm/day values for display
+            if c in df.columns:
+                df[c] = pd.to_numeric(df[c], errors="coerce").round(1)
+        df["pred_gnn"] = df["pred_gnn"].round(1)
+        df["target_mpa"] = float(target)
+        front = (["pred_gnn", "target_mpa", "measured", "source",
+                  "water_binder_ratio", "scm_fraction"] + DESIGN_COLS)
+        cols = [c for c in front if c in df.columns] + \
+               [c for c in df.columns if c not in front]
+        return df[cols]
+
+
+# ---------------------------------------------------------------------------
+# CLI (verification)
+# ---------------------------------------------------------------------------
+def _default_ckpt_dir() -> Path:
+    for c in (_HERE / "checkpoints_full_rich",
+              _HERE.parent / "Hybrid" / "outputs" / "checkpoints_full_rich"):
+        if (c / "hierarchical.pt").exists():
+            return c
+    return _HERE / "checkpoints_full_rich"
+
+
+def _default_index() -> Optional[Path]:
+    for c in (_HERE / "inverse_index.csv",):
+        if c.exists():
+            return c
+    return None
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    p = sub.add_parser("predict", help="forward: mix -> strength")
+    p.add_argument("--checkpoint-dir", default=str(_default_ckpt_dir()))
+    p.add_argument("--cement", type=float); p.add_argument("--slag", type=float)
+    p.add_argument("--fly-ash", type=float); p.add_argument("--silica-fume", type=float)
+    p.add_argument("--water", type=float); p.add_argument("--sp", type=float)
+    p.add_argument("--coarse", type=float); p.add_argument("--fine", type=float)
+    p.add_argument("--age", type=float, default=28.0)
+
+    s = sub.add_parser("suggest", help="inverse: strength -> mixes")
+    s.add_argument("--checkpoint-dir", default=str(_default_ckpt_dir()))
+    s.add_argument("--index", default=str(_default_index() or ""))
+    s.add_argument("--target", type=float, required=True)
+    s.add_argument("--k", type=int, default=5)
+    s.add_argument("--tol", type=float, default=None)
+    s.add_argument("--no-fibre", action="store_true")
+    s.add_argument("--no-scm", action="store_true")
+    s.add_argument("--no-refine", action="store_true")
+    args = ap.parse_args()
+
+    pred = Predictor(Path(args.checkpoint_dir))
+    if args.cmd == "predict":
+        mix = {
+            "cement_kg_m3": args.cement, "slag_kg_m3": args.slag,
+            "fly_ash_kg_m3": args.fly_ash, "silica_fume_kg_m3": args.silica_fume,
+            "water_kg_m3": args.water, "superplasticizer_kg_m3": args.sp,
+            "coarse_aggregate_kg_m3": args.coarse, "fine_aggregate_kg_m3": args.fine,
+            "age_days": args.age,
+        }
+        mix = {k: v for k, v in mix.items() if v is not None}
+        res = pred.predict_strength(mix)
+        print(f"input: {mix}")
+        print(f"GNN     : {res['gnn']:.1f} MPa")
+        print(f"tabular : {res['tabular']:.1f} MPa")
+        flags = pred.out_of_range(mix)
+        if flags:
+            print(f"[warning] outside training range: {flags}")
+    else:
+        if not args.index:
+            ap.error("suggest needs --index (run app/build_inverse_index.py first)")
+        index = pd.read_csv(args.index)
+        out = pred.suggest_mixes(
+            args.target, index, k=args.k, tol=args.tol,
+            allow_fibre=not args.no_fibre, allow_scm=not args.no_scm,
+            refine=not args.no_refine,
+        )
+        pd.set_option("display.width", 220, "display.max_columns", 40)
+        print(out.to_string(index=False))
+
+
+if __name__ == "__main__":
+    main()

inverse_index.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2f2031959a2949acfc1e4ea8cc3081f4c666c89d5c4af3ed019ec0f787b7009
+size 1848726

requirements.txt

@@ -0,0 +1,13 @@
+# CPU PyTorch wheels (keeps the Hugging Face Space image small).
+--extra-index-url https://download.pytorch.org/whl/cpu
+
+# Pinned to the training environment so the saved (PyG-style) weights load
+# identically. PyG 2.6 on torch 2.x uses torch.scatter_reduce for GINEConv /
+# NNConv, so the compiled torch-scatter / torch-sparse extensions are NOT needed.
+torch==2.5.1
+torch_geometric==2.6.1
+
+numpy
+pandas
+matplotlib
+streamlit>=1.49

streamlit_app.py

@@ -0,0 +1,459 @@
+"""Concrete Compressive Strength — Predict / Design (Streamlit).
+
+Two tools:
+  * Predict strength — enter a mix design (any field may be left blank), get the
+                       predicted compressive strength.
+  * Design a mix     — enter a target strength, get detailed mix designs that
+                       reach it (drawn from real data, optionally refined). All
+                       suggestions are reported at 28-day strength.
+
+Run locally:   streamlit run app/streamlit_app.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pandas as pd
+import streamlit as st
+
+# Make sibling modules importable no matter how the script is launched
+# (`streamlit run`, AppTest, or a Hugging Face Space at repo root).
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from inference import (
+    AGE_COL,
+    DESIGN_COLS,
+    Predictor,
+    _default_ckpt_dir,
+    _default_index,
+)
+
+st.set_page_config(
+    page_title="Concrete Compressive Strength",
+    page_icon="◧",
+    layout="wide",
+    initial_sidebar_state="collapsed",
+)
+
+CKPT_DIR = _default_ckpt_dir()
+INDEX_PATH = _default_index()
+DESIGN_AGE = 28.0  # all design suggestions are reported at 28-day strength
+
+# ---------------------------------------------------------------------------
+# Styling — editorial: cream canvas, black frame + top bar, big grotesque type.
+# ---------------------------------------------------------------------------
+INK = "#1a1a1a"          # near-black text / accent (monochrome)
+SUBTLE = "#6b6b63"       # warm gray secondary text
+CREAM = "#ffffff"        # page background / light text on dark elements
+CARD = "#f6f6f7"         # faint card surface (separates groups on white)
+HAIRLINE = "rgba(0,0,0,0.10)"
+FONT = ('"Inter", -apple-system, BlinkMacSystemFont, "Segoe UI", '
+        'Roboto, Helvetica, Arial, sans-serif')
+
+CSS = f"""
+<style>
+  @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800;900&display=swap');
+
+  /* hide default Streamlit chrome */
+  #MainMenu, footer, header[data-testid="stHeader"] {{ display: none; }}
+  [data-testid="stSidebar"] {{ display: none; }}
+
+  html, body, .stApp, [class*="css"] {{ font-family: {FONT}; }}
+  .stApp {{ background-color: {CREAM}; color: {INK}; }}
+  .block-container {{ padding-top: 66px; padding-bottom: 3rem; max-width: 1300px; }}
+
+  /* black top bar with header labels (no side/bottom frame) */
+  .topbar {{ position: fixed; top: 0; left: 0; right: 0; height: 46px; z-index: 10000;
+             background: #000; display: flex; align-items: center; justify-content: space-between;
+             padding: 0 24px; pointer-events: none; }}
+  .topbar span {{ color: {CREAM}; font-size: 0.72rem; font-weight: 700;
+                  letter-spacing: 0.14em; text-transform: uppercase; white-space: nowrap; }}
+  .topbar .center {{ position: absolute; left: 50%; transform: translateX(-50%); }}
+
+  /* hero: cube on the left, big editorial title on the right */
+  .hero {{ display: flex; align-items: center; gap: 52px; padding: 6px 0 4px; }}
+  .htext {{ text-align: left; }}
+  .htext h1 {{ font-size: 3rem; font-weight: 800; color: {INK};
+               letter-spacing: -0.03em; line-height: 1.0; margin: 0; }}
+  .htext h2 {{ font-size: 1.9rem; font-weight: 500; color: {INK}; opacity: 0.82;
+               letter-spacing: -0.02em; line-height: 1.05; margin: 4px 0 0; }}
+  .htext p  {{ font-size: 1.06rem; font-weight: 400; color: {SUBTLE};
+               margin: 12px 0 0; max-width: 680px; }}
+
+  /* 3D rotating concrete block (transparent background, CSS only) */
+  .scene {{ width: 118px; height: 118px; flex-shrink: 0; perspective: 700px; }}
+  .cube {{ width: 90px; height: 90px; position: relative; transform-style: preserve-3d;
+           margin: 14px auto; animation: spin 16s linear infinite; }}
+  @keyframes spin {{
+    from {{ transform: rotateX(-24deg) rotateY(0deg); }}
+    to   {{ transform: rotateX(-24deg) rotateY(360deg); }}
+  }}
+  .face {{ position: absolute; width: 90px; height: 90px;
+           border: 1px solid rgba(0,0,0,0.10);
+           box-shadow: inset 0 0 22px rgba(0,0,0,0.12); }}
+  .front  {{ transform: translateZ(45px);            background: linear-gradient(145deg,#cfccc4,#a8a59c); }}
+  .back   {{ transform: rotateY(180deg) translateZ(45px); background: linear-gradient(145deg,#cfccc4,#a8a59c); }}
+  .right  {{ transform: rotateY(90deg)  translateZ(45px); background: linear-gradient(145deg,#c2bfb6,#9a978e); }}
+  .left   {{ transform: rotateY(-90deg) translateZ(45px); background: linear-gradient(145deg,#c2bfb6,#9a978e); }}
+  .top    {{ transform: rotateX(90deg)  translateZ(45px); background: linear-gradient(145deg,#e0ddd3,#c4c1b8); }}
+  .bottom {{ transform: rotateX(-90deg) translateZ(45px); background: linear-gradient(145deg,#9a978e,#82807a); }}
+
+  /* tabs as a segmented control (left-aligned, black selected) */
+  .stTabs [data-baseweb="tab-list"] {{
+    background: #e7e3d7; border-radius: 10px; padding: 4px; gap: 4px;
+    width: fit-content; margin: 18px 0 12px; border: none;
+  }}
+  .stTabs [data-baseweb="tab"] {{ height: 40px; padding: 0 28px; background: transparent; border-radius: 7px; }}
+  .stTabs [data-baseweb="tab"] p {{ font-size: 1.0rem; font-weight: 700; color: {INK}; }}
+  .stTabs [aria-selected="true"] {{ background: {INK}; }}
+  .stTabs [aria-selected="true"] p {{ color: {CREAM}; }}
+  .stTabs [data-baseweb="tab-highlight"] {{ display: none; }}
+
+  /* step + section headings */
+  .step {{ font-size: 1.35rem; font-weight: 700; color: {INK};
+           letter-spacing: -0.01em; margin: 8px 0 0; }}
+  .sec  {{ font-size: 0.76rem; font-weight: 700; color: {SUBTLE};
+           text-transform: uppercase; letter-spacing: 0.08em; margin: 2px 0 6px; }}
+
+  /* group cards */
+  [data-testid="stVerticalBlockBorderWrapper"] {{
+    background: {CARD}; border: 1px solid {HAIRLINE}; border-radius: 14px;
+    padding: 10px 20px 16px; margin-bottom: 16px; box-shadow: 0 1px 3px rgba(0,0,0,0.04);
+  }}
+
+  /* result cards */
+  [data-testid="stMetric"] {{
+    background: {CARD}; border: 1px solid {HAIRLINE}; border-radius: 16px; padding: 22px 24px;
+  }}
+  [data-testid="stMetricValue"] {{ font-size: 2.5rem; font-weight: 800; color: {INK};
+                                   letter-spacing: -0.02em; }}
+  [data-testid="stMetricLabel"] p {{ color: {SUBTLE}; font-weight: 600; font-size: 0.95rem; }}
+
+  /* inputs */
+  [data-testid="stNumberInput"] input, [data-baseweb="select"] > div {{ border-radius: 9px; }}
+
+  /* buttons — solid black */
+  .stButton > button {{
+    background: {INK}; color: {CREAM}; border: 0; border-radius: 8px;
+    padding: 0.6rem 1.8rem; font-weight: 700; font-size: 1.0rem;
+  }}
+  .stButton > button:hover {{ background: #000; color: {CREAM}; }}
+  .stDownloadButton > button {{
+    background: transparent; color: {INK}; border: 1px solid {INK}; border-radius: 8px;
+    padding: 0.5rem 1.4rem; font-weight: 700;
+  }}
+</style>
+"""
+
+FRAME = """
+<div class="frame"></div>
+<div class="topbar">
+  <span class="left">Concrete Mix Toolkit</span>
+  <span class="center">Compressive Strength</span>
+  <span class="right">Predict / Design</span>
+</div>
+"""
+
+HERO = """
+<div class="hero">
+  <div class="scene"><div class="cube">
+    <div class="face front"></div><div class="face back"></div>
+    <div class="face right"></div><div class="face left"></div>
+    <div class="face top"></div><div class="face bottom"></div>
+  </div></div>
+  <div class="htext">
+    <h1>Concrete Compressive Strength</h1>
+    <h2>Predict &amp; Design</h2>
+    <p>Estimate the strength of any concrete mix — or design a mix to hit a target strength.</p>
+  </div>
+</div>
+"""
+
+
+@st.cache_resource(show_spinner="Loading…")
+def get_predictor() -> Predictor:
+    return Predictor(CKPT_DIR)
+
+
+@st.cache_data(show_spinner="Loading mix database…")
+def get_index() -> pd.DataFrame:
+    return pd.read_csv(INDEX_PATH) if INDEX_PATH else pd.DataFrame()
+
+
+# ---- friendly field definitions: (column, label, unit) ----
+BINDER = [
+    ("cement_kg_m3", "Cement", "kg/m³"),
+    ("slag_kg_m3", "GGBS / slag", "kg/m³"),
+    ("fly_ash_kg_m3", "Fly ash", "kg/m³"),
+    ("silica_fume_kg_m3", "Silica fume", "kg/m³"),
+    ("metakaolin_kg_m3", "Metakaolin", "kg/m³"),
+    ("limestone_powder_kg_m3", "Limestone powder", "kg/m³"),
+    ("other_scm_kg_m3", "Other SCM / filler", "kg/m³"),
+]
+AGG = [
+    ("coarse_aggregate_kg_m3", "Coarse aggregate", "kg/m³"),
+    ("fine_aggregate_kg_m3", "Fine aggregate / sand", "kg/m³"),
+]
+FIBRE = [
+    ("fibre_content_kg_m3", "Fibre content", "kg/m³"),
+    ("fibre_length_mm", "Fibre length", "mm"),
+    ("fibre_diameter_mm", "Fibre diameter", "mm"),
+    ("fibre_tensile_strength_mpa", "Fibre tensile strength", "MPa"),
+    ("fibre_modulus_gpa", "Fibre modulus", "GPa"),
+]
+CHEM = [
+    ("cement_CaO_pct", "Cement CaO", "%"), ("cement_SiO2_pct", "Cement SiO₂", "%"),
+    ("cement_Al2O3_pct", "Cement Al₂O₃", "%"), ("cement_Fe2O3_pct", "Cement Fe₂O₃", "%"),
+    ("cement_MgO_pct", "Cement MgO", "%"), ("cement_SO3_pct", "Cement SO₃", "%"),
+    ("cement_alkali_pct", "Cement alkali (Na₂O-eq)", "%"), ("cement_LOI_pct", "Cement LOI", "%"),
+    ("scm_CaO_pct", "SCM CaO", "%"), ("scm_SiO2_pct", "SCM SiO₂", "%"),
+    ("scm_Al2O3_pct", "SCM Al₂O₃", "%"), ("scm_Fe2O3_pct", "SCM Fe₂O₃", "%"),
+    ("scm_MgO_pct", "SCM MgO", "%"), ("scm_LOI_pct", "SCM LOI", "%"),
+    ("cement_grade_mpa", "Cement grade", "MPa"),
+    ("specimen_size_mm", "Specimen size", "mm"),
+]
+
+# Official-style labels -> the model's canonical categorical values.
+CEMENT_TYPE_OPTIONS = {
+    "(unspecified)": None,
+    "CEM I — Ordinary Portland (OPC)": "opc",
+    "Portland Type I/II": "type_i_ii",
+    "Type III — Rapid-hardening": "type_iii",
+    "Sulfate-resisting (HS)": "hs",
+    "Other": "other",
+}
+FIBRE_TYPE_OPTIONS = {
+    "(unspecified)": None, "None": "none", "Steel": "steel",
+    "PVA": "pva", "PE / polyethylene": "pe", "Other": "other",
+}
+CURING_OPTIONS = {
+    "(unspecified)": None, "Standard": "standard", "Steam": "steam",
+    "Heat-cured": "heat", "Autoclave": "autoclave", "Other": "other",
+}
+
+# Map every raw column to a clear, units-bearing label for the results table.
+COLUMN_LABELS = {c: f"{lab} ({u})" for c, lab, u in (BINDER + AGG + FIBRE + CHEM)}
+COLUMN_LABELS.update({
+    "max_coarse_aggregate_size_mm": "Max coarse agg. size (mm)",
+    "max_fine_aggregate_size_mm": "Max fine agg. size (mm)",
+    "curing_temperature_c": "Curing temperature (°C)",
+    "curing_humidity_pct": "Curing humidity (%)",
+    "pred_gnn": "Predicted strength (MPa)",
+    "target_mpa": "Target (MPa)",
+    "water_binder_ratio": "Water / binder",
+    "scm_fraction": "SCM fraction",
+    "water_kg_m3": "Water (kg/m³)",
+    "superplasticizer_kg_m3": "Superplasticizer (kg/m³)",
+    "age_days": "Curing age (days)",
+})
+
+
+def _num(pred, col, label, unit, default=0.0, optional=False,
+         lo=None, hi=None, min_value=0.0, max_value=None):
+    """Render a number_input. ``lo``/``hi`` override the 'typical' hint range;
+    ``min_value``/``max_value`` set hard entry limits."""
+    b = pred.bounds.get(col, {})
+    p_lo = lo if lo is not None else b.get("p01")
+    p_hi = hi if hi is not None else b.get("p99")
+    help_txt = (f"typical {p_lo:.0f}–{p_hi:.0f} {unit}"
+                if (p_lo is not None and p_hi is not None) else None)
+    return st.number_input(
+        f"{label} ({unit})",
+        value=(None if optional else float(default)),
+        min_value=float(min_value),
+        max_value=(float(max_value) if max_value is not None else None),
+        step=1.0, help=help_txt, key=f"in_{col}",
+    )
+
+
+def _step(title: str) -> None:
+    st.markdown(f"<div class='step'>{title}</div>", unsafe_allow_html=True)
+
+
+def _sec(title: str) -> None:
+    st.markdown(f"<div class='sec'>{title}</div>", unsafe_allow_html=True)
+
+
+def _category_select(label, options, mix, key):
+    choice = st.selectbox(label, list(options.keys()), index=0)
+    if options[choice] is not None:
+        mix[key] = options[choice]
+
+
+def format_suggestions(out: pd.DataFrame) -> pd.DataFrame:
+    """Clean, clearly-named subset of the suggested mixes."""
+    always = {"cement_kg_m3", "water_kg_m3", "age_days"}
+    ing = [c for c in DESIGN_COLS if c in out.columns]
+    keep_ing = [c for c in ing if c in always
+                or pd.to_numeric(out[c], errors="coerce").fillna(0).abs().sum() > 0]
+    head = [c for c in ["pred_gnn", "water_binder_ratio", "scm_fraction"]
+            if c in out.columns]
+    disp = out[head + keep_ing].copy()
+    return disp.rename(columns=COLUMN_LABELS)
+
+
+# ---------------------------------------------------------------------------
+# Tab 1 — forward
+# ---------------------------------------------------------------------------
+def forward_tab(pred: Predictor) -> None:
+    _step("Enter your mix design")
+    st.caption("Leave any field blank if you don't have it. Amounts are per m³ of concrete.")
+
+    mix: dict = {}
+    c1, c2, c3 = st.columns(3, gap="large")
+    with c1:
+        with st.container(border=True):
+            _sec("Binder")
+            med = pred.bounds.get("cement_kg_m3", {}).get("p50", 350.0)
+            for col, label, unit in BINDER:
+                default = med if col == "cement_kg_m3" else 0.0
+                mix[col] = _num(pred, col, label, unit, default=default)
+    with c2:
+        with st.container(border=True):
+            _sec("Water & admixture")
+            mix["water_kg_m3"] = _num(
+                pred, "water_kg_m3", "Water", "kg/m³",
+                default=pred.bounds.get("water_kg_m3", {}).get("p50", 175.0), lo=120, hi=600,
+            )
+            mix["superplasticizer_kg_m3"] = _num(
+                pred, "superplasticizer_kg_m3", "Superplasticizer", "kg/m³",
+                default=pred.bounds.get("superplasticizer_kg_m3", {}).get("p50", 0.0),
+            )
+        with st.container(border=True):
+            _sec("Aggregate")
+            for col, label, unit in AGG:
+                d = pred.bounds.get(col, {}).get("p50", 0.0)
+                mix[col] = _num(pred, col, label, unit, default=d)
+            mix["max_coarse_aggregate_size_mm"] = _num(
+                pred, "max_coarse_aggregate_size_mm", "Max coarse agg. size", "mm",
+                optional=True, lo=10, hi=40, min_value=10, max_value=40,
+            )
+            mix["max_fine_aggregate_size_mm"] = _num(
+                pred, "max_fine_aggregate_size_mm", "Max fine agg. size", "mm", optional=True,
+            )
+    with c3:
+        with st.container(border=True):
+            _sec("Curing & age")
+            mix[AGE_COL] = st.number_input("Curing age (days)", value=28.0, min_value=1.0, step=1.0)
+            mix["curing_temperature_c"] = _num(
+                pred, "curing_temperature_c", "Curing temperature", "°C", optional=True)
+            mix["curing_humidity_pct"] = _num(
+                pred, "curing_humidity_pct", "Curing humidity", "%", optional=True)
+
+    with st.expander("Fibres (optional)"):
+        fc = st.columns(len(FIBRE))
+        for (col, label, unit), c in zip(FIBRE, fc):
+            with c:
+                mix[col] = _num(pred, col, label, unit, optional=True)
+        _category_select("Fibre type", FIBRE_TYPE_OPTIONS, mix, "fibre_type_norm")
+
+    with st.expander("Cement / SCM chemistry & type (advanced, optional)"):
+        gcols = st.columns(4)
+        for i, (col, label, unit) in enumerate(CHEM):
+            with gcols[i % 4]:
+                mix[col] = _num(pred, col, label, unit, optional=True)
+        s1, s2 = st.columns(2)
+        with s1:
+            _category_select("Cement type", CEMENT_TYPE_OPTIONS, mix, "cement_type_norm")
+        with s2:
+            _category_select("Curing regime", CURING_OPTIONS, mix, "curing_regime_norm")
+
+    show_curve = st.checkbox("Show strength gain with curing age", value=True)
+    st.write("")
+    go = st.button("Predict strength", type="primary")
+
+    if go:
+        clean = {k: v for k, v in mix.items() if v not in (None, "")}
+        with st.spinner("Predicting…"):
+            res = pred.predict_strength(clean)
+        st.write("")
+        m1, m2 = st.columns(2)
+        m1.metric("Predicted compressive strength", f"{res['gnn']:.1f} MPa")
+        m2.metric("Secondary estimate", f"{res['tabular']:.1f} MPa",
+                  help="An independent model used as a cross-check.")
+        rmse = pred.config.get("test_metrics", {}).get("rmse")
+        if rmse:
+            st.caption(f"Typical accuracy ≈ ±{rmse:.0f} MPa on held-out test data.")
+
+        flags = pred.out_of_range(clean)
+        if flags:
+            pretty = ", ".join(COLUMN_LABELS.get(f, f) for f in flags)
+            st.warning(f"Some inputs are outside the usual data range ({pretty}); "
+                       "the prediction there is an extrapolation.")
+
+        if show_curve:
+            curve = pred.age_curve(clean).rename(
+                columns={"gnn_mpa": "Predicted (MPa)", "tabular_mpa": "Secondary (MPa)"}
+            )
+            st.markdown("##### Predicted strength gain with curing age")
+            st.line_chart(curve.set_index("age_days"))
+
+
+# ---------------------------------------------------------------------------
+# Tab 2 — inverse
+# ---------------------------------------------------------------------------
+def inverse_tab(pred: Predictor, index: pd.DataFrame) -> None:
+    _step("Choose a target strength")
+    if index.empty:
+        st.info("The mix database (`inverse_index.csv`) isn't built yet. "
+                "Run `python app/build_inverse_index.py` first.")
+        return
+    smin = float(pred.config.get("strength_min", 10))
+    smax = float(pred.config.get("strength_max", 200))
+
+    c1, c2 = st.columns([2, 1])
+    with c1:
+        target = st.slider("Target compressive strength (MPa)", smin, smax,
+                           min(80.0, smax), step=1.0)
+    with c2:
+        k = st.slider("Number of mix options", 1, 4, 1)
+
+    o1, o2, o3 = st.columns(3)
+    allow_fibre = o1.checkbox("Allow fibres", value=True)
+    allow_scm = o2.checkbox("Allow SCMs", value=True)
+    refine = o3.checkbox("Refine to target", value=False,
+                         help="Fine-tune each candidate so its predicted strength matches the target.")
+    st.caption("All suggestions are reported at 28-day strength.")
+
+    st.write("")
+    if st.button("Design mixes", type="primary"):
+        with st.spinner("Searching mix designs…"):
+            out = pred.suggest_mixes(
+                target, index, k=k, allow_fibre=allow_fibre, allow_scm=allow_scm,
+                domain="any", age=DESIGN_AGE, refine=refine,
+            )
+        if out.empty:
+            st.warning("No mixes match those constraints — try widening them.")
+            return
+        st.success(f"{len(out)} mix design(s) predicted to reach ≈ {target:.0f} MPa at 28 days")
+        disp = format_suggestions(out)
+        st.dataframe(disp, width="stretch", hide_index=True)
+        st.download_button("Download these mixes (CSV)", disp.to_csv(index=False),
+                           file_name=f"mix_designs_{int(target)}MPa.csv", mime="text/csv")
+
+
+def main() -> None:
+    st.markdown(CSS, unsafe_allow_html=True)
+    st.markdown(FRAME, unsafe_allow_html=True)
+    st.markdown(HERO, unsafe_allow_html=True)
+
+    if not (Path(CKPT_DIR) / "hierarchical.pt").exists():
+        st.error(f"No trained model found in `{CKPT_DIR}`.")
+        st.stop()
+
+    pred = get_predictor()
+    index = get_index()
+
+    tab1, tab2 = st.tabs(["Predict strength", "Design a mix"])
+    with tab1:
+        forward_tab(pred)
+    with tab2:
+        inverse_tab(pred, index)
+
+
+if __name__ == "__main__":
+    main()
+else:
+    main()