concrete_gnn/data.py

"""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,
    binder_chem_features,
    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
    # Per-graph concrete tensile/flexural label vector (2,), physical MPa, NaN where
    # unsupervised. Drives the masked tensile/flexural capacity-head losses.
    concrete_aux_target: Optional[Tensor] = None
    mortar_target: Optional[Tensor] = None
    itz_target: Optional[Tensor] = None
    # Per-GRAPH ITZ label vector (n_itz_targets,), NaN where unsupervised. Distinct
    # from the synthetic per-EDGE ``itz_target``; used for masked real-data ITZ
    # supervision against the model's per-graph mean ITZ prediction.
    itz_target_graph: Optional[Tensor] = None
    # Per-graph paste label vector (n_paste_targets,), NaN where unsupervised.
    paste_target: Optional[Tensor] = None
    # Per-graph paste water/binder ratio (scalar), for the Abrams paste-strength
    # prior. NaN where the mix has no usable binder.
    paste_wb: 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]
    itz_target_graph: Optional[Tensor] = None
    paste_target: Optional[Tensor] = None
    paste_wb: Optional[Tensor] = None
    tabular_mix: Optional[Tensor] = None
    concrete_aux_target: 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,
            itz_target_graph=m(self.itz_target_graph) if self.itz_target_graph is not None else None,
            paste_target=m(self.paste_target) if self.paste_target is not None else None,
            paste_wb=m(self.paste_wb) if self.paste_wb is not None else None,
            tabular_mix=m(self.tabular_mix) if self.tabular_mix is not None else None,
            concrete_aux_target=m(self.concrete_aux_target)
            if self.concrete_aux_target is not None
            else None,
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _mix_itz_vector(mix: MixDesign, mortar_mix: MortarMixDesign) -> Tensor:
    """Per-graph static input to the ITZ MLP: curing + binder context PLUS the
    full binder chemistry/phases (the ITZ is a modified cement paste, so it sees
    the same oxide/Bogue chemistry the paste nodes do). Chemistry is read from
    the mortar mix, which carries those fields."""

    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,
            *binder_chem_features(mortar_mix),
            mortar_mix.glass_powder_kg_m3,
            mortar_mix.fly_ash_type_id,
        ],
        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, mortar_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_itz_graph = samples[0].itz_target_graph is not None
    have_paste = samples[0].paste_target is not None
    have_paste_wb = samples[0].paste_wb is not None
    have_tabular = samples[0].tabular_mix is not None
    have_concrete_aux = samples[0].concrete_aux_target is not None
    tabular_mixes: List[Tensor] = []
    itz_graph_targets: List[Tensor] = []
    paste_targets: List[Tensor] = []
    paste_wbs: List[Tensor] = []
    concrete_aux_targets: 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_graph:
            itz_graph_targets.append(sample.itz_target_graph)
        if have_paste:
            paste_targets.append(sample.paste_target)
        if have_paste_wb:
            paste_wbs.append(sample.paste_wb)
        if have_concrete_aux:
            concrete_aux_targets.append(sample.concrete_aux_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,
        itz_target_graph=torch.stack(itz_graph_targets, dim=0) if have_itz_graph else None,
        paste_target=torch.stack(paste_targets, dim=0) if have_paste else None,
        paste_wb=torch.stack(paste_wbs, dim=0) if have_paste_wb else None,
        tabular_mix=torch.stack(tabular_mixes, dim=0) if have_tabular else None,
        concrete_aux_target=torch.stack(concrete_aux_targets, dim=0)
        if have_concrete_aux
        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 :]