"""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 :]