2. Design a Train/Validation/Test Split for Nonlinear FEA Outputs

Why a Nonlinear FEA Dataset Needs a Purpose-Built Split

A nonlinear finite element dataset is not a bag of independent rows. Each simulation run is a batch: one geometry, one material realization, one mesh, and a load history that produces a full response curve. If you split at the row level — one load step here, another there — the same run leaks across partitions and validation scores become optimistic. The task in this case is to build a split that is leakage-free at the run level and then prove that the split still preserves the distribution of the input parameters that drive the surrogate.

We will work with a representative dataset for a nonlinear cantilever problem: an elastoplastic beam under a tip load, parameterized by yield stress sigma_y, hardening modulus H, beam length L, and load amplitude P_max. Each run returns a tip-deflection response sampled at 40 load steps. The split must keep every run intact, and the resulting train/validation/test partitions must match the parent distribution on all four inputs.

Dataset Contract and Leakage Rules

Before writing any splitting code, fix the contract. The dataset is a single table with one row per (run, load step) and a run_id column that identifies the simulation batch. The split is defined on run_id, never on rows.

• Unit of splitting: run_id. All rows sharing a run_id land in exactly one partition.
• Target proportions: 70% train, 15% validation, 15% test by run count.
• Leakage rule: no run_id may appear in more than one partition; no near-duplicate parameter vector may straddle partitions (handled in the next case, but the split must not create the opportunity).
• Distribution rule: each partition's input-parameter marginals must be statistically consistent with the full dataset.

The figure below shows the intended data flow: raw run-level records are grouped, partitioned by run, and then checked against the parent distribution before being accepted.

Implementing the Split

The split is a short, self-contained script. It groups rows by run_id, assigns each run to a partition using a seeded shuffle, and then runs a Kolmogorov–Smirnov test on each input parameter to confirm the partitions match the parent distribution. The code below is complete and runnable given a pandas DataFrame df with the columns described above.

import numpy as np
import pandas as pd
from scipy.stats import ks_2samp

INPUTS = ["sigma_y", "H", "L", "P_max"]
SEED = 42
FRACTIONS = {"train": 0.70, "val": 0.15, "test": 0.15}

def split_by_run(df, seed=SEED, fractions=FRACTIONS):
    runs = df["run_id"].unique()
    rng = np.random.default_rng(seed)
    rng.shuffle(runs)

    n = len(runs)
    n_train = int(round(fractions["train"] * n))
    n_val = int(round(fractions["val"] * n))

    train_ids = set(runs[:n_train])
    val_ids = set(runs[n_train:n_train + n_val])
    test_ids = set(runs[n_train + n_val:])

    assert train_ids.isdisjoint(val_ids)
    assert train_ids.isdisjoint(test_ids)
    assert val_ids.isdisjoint(test_ids)

    return {
        "train": df[df["run_id"].isin(train_ids)],
        "val":   df[df["run_id"].isin(val_ids)],
        "test":  df[df["run_id"].isin(test_ids)],
    }

def check_distribution(df, parts, alpha=0.05):
    report = []
    for name, part in parts.items():
        for col in INPUTS:
            stat, p = ks_2samp(df[col], part[col])
            report.append({
                "partition": name,
                "input": col,
                "ks_stat": round(stat, 4),
                "p_value": round(p, 4),
                "pass": p > alpha,
            })
    return pd.DataFrame(report)

parts = split_by_run(df)
report = check_distribution(df, parts)
print(report.to_string(index=False))

Two details matter for correctness. First, the shuffle is on run_id, not on rows, so a run's 40 load steps cannot be scattered. Second, the KS test compares each partition's marginal against the full dataset, not against the training set — this is the right reference because the goal is to preserve the parent distribution, not to make partitions identical to each other.

Reading the Verification Report

With a typical run count of 200 simulations, the expected output is a 12-row table (3 partitions × 4 inputs). A healthy split looks like this:

These values are illustrative for a 200-run dataset; the exact numbers depend on the seed and the underlying parameter ranges. What matters is the pattern: KS statistics below roughly 0.15 and p-values comfortably above 0.05 across all partitions. The validation and test partitions show larger KS statistics than train simply because they are smaller — with 30 runs each, sampling noise alone produces statistics in the 0.10–0.15 range even for a perfectly representative split.

Failure Symptoms and Troubleshooting

• KS p-value below 0.05 on one input in one partition: usually sampling noise at small run counts. Re-run with a different seed; if the failure persists across seeds, the parameter space is likely clustered and the split needs stratification on that input.
• All partitions fail on the same input: the parent distribution is bimodal or heavily skewed, and a uniform shuffle cannot preserve it. Stratify by binning that input into quantiles and sampling within bins.
• Overlap assertion fires: the run_id column is not unique per simulation, or rows were duplicated before splitting. Deduplicate on run_id first.
• Validation score much higher than test score: a sign that the split leaked — check that no run_id appears in two partitions and that the test set was not used during tuning.

Deliverable: Split Specification Sheet

Freeze the accepted split as a small specification so downstream cases (leakage detection, high-fidelity validation) can reuse it without re-deriving the partition.

Persist the accepted partition as split_ids.json so that every later experiment uses the same runs. A split that is regenerated with a new seed between experiments silently invalidates comparisons, and that is the most common way a leakage-free split becomes a leakage-prone one in practice.