afem/outlook/check_correction_data.py

458 lines
18 KiB
Python

#!/usr/bin/env python3
"""
check_correction_data.py — Validate step-wise correction dataset.
Checks:
1. Global statistics (elements, marks, overlap, IoU, correction ratios)
2. Per-step trend analysis (correction signal vs AMR step)
3. Spatial visualization of 5 random samples
4. Field & shape consistency
5. Final verdict: PASS / WARNING / FAIL
Usage:
python outlook/check_correction_data.py \
--data-dir outlook/data_correction \
--output-dir outlook/data_correction_check
"""
import argparse
import sys
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
# ──────────────────────────────────────────────────────────────────────────────
# Required fields and their expected properties
# ──────────────────────────────────────────────────────────────────────────────
REQUIRED_FIELDS = {
"features": {"ndim": 2, "dtype_kind": "f"},
"edge_index": {"ndim": 2, "dtype_kind": "i"},
"physics_score": {"ndim": 1, "dtype_kind": "f"},
"teacher_eta": {"ndim": 1, "dtype_kind": "f"},
"teacher_mark": {"ndim": 1, "dtype_kind": "i"},
"physics_mark": {"ndim": 1, "dtype_kind": "i"},
"correction_label": {"ndim": 1, "dtype_kind": "i"},
"elements": {"ndim": 0},
"step": {"ndim": 0},
"aw_rel_before": {"ndim": 0},
"max_err_before": {"ndim": 0},
"k": {"ndim": 0},
"eps_r": {"ndim": 0},
"cx": {"ndim": 0},
"cy": {"ndim": 0},
"radius": {"ndim": 0},
}
# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────
def iou_binary(a: np.ndarray, b: np.ndarray) -> float:
"""IoU between two binary (0/1) arrays."""
a = a.astype(bool)
b = b.astype(bool)
inter = np.sum(a & b)
union = np.sum(a | b)
return float(inter) / float(union) if union > 0 else 1.0
def load_all_samples(data_dir: Path):
"""Load all sampleXXXX_stepYYY.npz files, return list of dicts."""
files = sorted(data_dir.glob("sample*_step*.npz"))
samples = []
for f in files:
try:
d = dict(np.load(f, allow_pickle=True))
d["_path"] = str(f)
d["_name"] = f.name
samples.append(d)
except Exception as e:
print(f" [WARN] Cannot load {f.name}: {e}")
return samples
# ──────────────────────────────────────────────────────────────────────────────
# 1. Field & shape consistency
# ──────────────────────────────────────────────────────────────────────────────
def check_fields(samples, issues):
"""Check every file has the required fields with correct ndim."""
print("\n" + "=" * 70)
print(" [1/5] Field & Shape Consistency")
print("=" * 70)
n_files = len(samples)
n_issues_before = len(issues)
for s in samples:
name = s["_name"]
for field, spec in REQUIRED_FIELDS.items():
if field not in s:
issues.append(f"MISSING field '{field}' in {name}")
continue
arr = s[field]
if arr.ndim != spec["ndim"]:
issues.append(
f"BAD ndim for '{field}' in {name}: "
f"expected {spec['ndim']}, got {arr.ndim}"
)
if "dtype_kind" in spec and arr.dtype.kind != spec["dtype_kind"]:
issues.append(
f"BAD dtype for '{field}' in {name}: "
f"expected kind={spec['dtype_kind']}, got {arr.dtype}"
)
# Cross-check: features rows == len(physics_score) == elements
if "features" in s and "physics_score" in s:
nf = s["features"].shape[0]
ns = len(s["physics_score"])
ne = int(s["elements"])
if not (nf == ns == ne):
issues.append(
f"SIZE MISMATCH in {name}: features={nf}, "
f"physics_score={ns}, elements={ne}"
)
n_new = len(issues) - n_issues_before
if n_new == 0:
print(f" All {n_files} files have consistent fields & shapes.")
else:
print(f" Found {n_new} field/shape issues.")
# ──────────────────────────────────────────────────────────────────────────────
# 2. Global statistics
# ──────────────────────────────────────────────────────────────────────────────
def compute_global_stats(samples):
"""Compute and print global statistics across all files."""
print("\n" + "=" * 70)
print(" [2/5] Global Statistics")
print("=" * 70)
elems = []
teacher_ratios = []
physics_ratios = []
overlaps = []
overlap_teacher = []
overlap_physics = []
ious = []
pos_ratios = []
neg_ratios = []
for s in samples:
n = int(s["elements"])
tm = s["teacher_mark"]
pm = s["physics_mark"]
cl = s["correction_label"]
elems.append(n)
n_tm = int(np.sum(tm))
n_pm = int(np.sum(pm))
ov = int(np.sum((tm == 1) & (pm == 1)))
teacher_ratios.append(n_tm / max(n, 1))
physics_ratios.append(n_pm / max(n, 1))
overlaps.append(ov)
overlap_teacher.append(ov / max(n_tm, 1))
overlap_physics.append(ov / max(n_pm, 1))
ious.append(iou_binary(tm, pm))
pos_ratios.append(int(np.sum(cl == 1)) / max(n, 1))
neg_ratios.append(int(np.sum(cl == -1)) / max(n, 1))
stats = {
"total_files": len(samples),
"elements": elems,
"teacher_mark_ratio": teacher_ratios,
"physics_mark_ratio": physics_ratios,
"overlap_count": overlaps,
"overlap/teacher": overlap_teacher,
"overlap/physics": overlap_physics,
"IoU": ious,
"correction_pos_ratio": pos_ratios,
"correction_neg_ratio": neg_ratios,
}
print(f" {'Metric':<30s} {'mean':>10s} {'min':>10s} {'max':>10s}")
print(f" {'-'*30} {'-'*10} {'-'*10} {'-'*10}")
for key in [
"elements", "teacher_mark_ratio", "physics_mark_ratio",
"overlap_count", "overlap/teacher", "overlap/physics",
"IoU", "correction_pos_ratio", "correction_neg_ratio",
]:
arr = np.array(stats[key])
print(f" {key:<30s} {arr.mean():10.4f} {arr.min():10.4f} {arr.max():10.4f}")
return stats
# ──────────────────────────────────────────────────────────────────────────────
# 3. Per-step trend analysis
# ──────────────────────────────────────────────────────────────────────────────
def check_per_step_trends(samples, issues):
"""Group by step, print trend table, flag anomalies."""
print("\n" + "=" * 70)
print(" [3/5] Per-Step Trend Analysis")
print("=" * 70)
from collections import defaultdict
step_data = defaultdict(lambda: {
"elems": [], "tm_ratio": [], "pm_ratio": [],
"overlap": [], "iou": [], "pos_ratio": [], "neg_ratio": [],
})
for s in samples:
step = int(s["step"])
n = int(s["elements"])
tm = s["teacher_mark"]
pm = s["physics_mark"]
cl = s["correction_label"]
n_tm = int(np.sum(tm))
n_pm = int(np.sum(pm))
ov = int(np.sum((tm == 1) & (pm == 1)))
d = step_data[step]
d["elems"].append(n)
d["tm_ratio"].append(n_tm / max(n, 1))
d["pm_ratio"].append(n_pm / max(n, 1))
d["overlap"].append(ov)
d["iou"].append(iou_binary(tm, pm))
d["pos_ratio"].append(int(np.sum(cl == 1)) / max(n, 1))
d["neg_ratio"].append(int(np.sum(cl == -1)) / max(n, 1))
steps = sorted(step_data.keys())
header = f" {'step':>4s} {'n_files':>7s} {'mean_elem':>9s} {'tm_ratio':>9s} {'pm_ratio':>9s} {'IoU':>7s} {'pos%':>7s} {'neg%':>7s}"
print(header)
print(f" {'-'*4} {'-'*7} {'-'*9} {'-'*9} {'-'*9} {'-'*7} {'-'*7} {'-'*7}")
for st in steps:
d = step_data[st]
nf = len(d["elems"])
me = np.mean(d["elems"])
mt = np.mean(d["tm_ratio"])
mp = np.mean(d["pm_ratio"])
mi = np.mean(d["iou"])
mpos = np.mean(d["pos_ratio"])
mneg = np.mean(d["neg_ratio"])
print(f" {st:4d} {nf:7d} {me:9.1f} {mt:9.4f} {mp:9.4f} {mi:7.4f} {mpos:7.4f} {mneg:7.4f}")
# Trend checks
if len(steps) >= 3:
# Check: correction signal should diminish at late steps (elements grow → marks shrink)
late = steps[-1]
early = steps[0]
late_pos = np.mean(step_data[late]["pos_ratio"])
early_pos = np.mean(step_data[early]["pos_ratio"])
late_neg = np.mean(step_data[late]["neg_ratio"])
early_neg = np.mean(step_data[early]["neg_ratio"])
late_iou = np.mean(step_data[late]["iou"])
early_iou = np.mean(step_data[early]["iou"])
print(f"\n Trend summary (step {early}{late}):")
print(f" pos_ratio: {early_pos:.4f}{late_pos:.4f}")
print(f" neg_ratio: {early_neg:.4f}{late_neg:.4f}")
print(f" IoU: {early_iou:.4f}{late_iou:.4f}")
# Warning: if late-step IoU is much lower than early, marks diverge
if late_iou < early_iou * 0.5:
issues.append(
f"TREND: IoU drops significantly from step {early} ({early_iou:.3f}) "
f"to step {late} ({late_iou:.3f}) — marks diverge at late steps."
)
# Warning: if correction signal vanishes (pos+neg → 0) at late steps
late_corr = late_pos + late_neg
early_corr = early_pos + early_neg
if early_corr > 0.01 and late_corr < early_corr * 0.1:
issues.append(
f"TREND: Correction signal nearly vanishes at step {late} "
f"({late_corr:.4f}) vs step {early} ({early_corr:.4f})."
)
# Warning: if IoU is very low overall (teacher and physics barely agree)
all_ious = [iou for s in steps for iou in step_data[s]["iou"]]
overall_iou = np.mean(all_ious)
if overall_iou < 0.1:
issues.append(
f"LOW IoU: mean IoU across all steps is {overall_iou:.4f}"
f"teacher_mark and physics_mark barely overlap."
)
# ──────────────────────────────────────────────────────────────────────────────
# 4. Spatial visualization
# ──────────────────────────────────────────────────────────────────────────────
def plot_spatial_samples(samples, output_dir: Path, n_plot: int = 5):
"""Randomly select n_plot files and save 4-panel spatial maps."""
print("\n" + "=" * 70)
print(" [4/5] Spatial Visualization")
print("=" * 70)
output_dir.mkdir(parents=True, exist_ok=True)
rng = np.random.RandomState(1234)
indices = rng.choice(len(samples), size=min(n_plot, len(samples)), replace=False)
selected = [samples[i] for i in sorted(indices)]
for s in selected:
name = s["_name"].replace(".npz", "")
features = s["features"]
x = features[:, 0]
y = features[:, 1]
physics_score = s["physics_score"]
teacher_eta = s["teacher_eta"]
teacher_mark = s["teacher_mark"]
correction_label = s["correction_label"]
fig, axes = plt.subplots(1, 4, figsize=(20, 5))
fig.suptitle(
f"{name} (elem={int(s['elements'])}, k={float(s['k']):.2f}, "
f"step={int(s['step'])})",
fontsize=12,
)
# Panel 1: physics_score
sc0 = axes[0].scatter(x, y, c=physics_score, s=8, cmap="viridis")
axes[0].set_title("physics_score")
axes[0].set_aspect("equal")
plt.colorbar(sc0, ax=axes[0], shrink=0.8)
# Panel 2: teacher_eta
sc1 = axes[1].scatter(x, y, c=teacher_eta, s=8, cmap="magma")
axes[1].set_title("teacher_eta")
axes[1].set_aspect("equal")
plt.colorbar(sc1, ax=axes[1], shrink=0.8)
# Panel 3: teacher_mark
sc2 = axes[2].scatter(x, y, c=teacher_mark, s=8, cmap="Reds", vmin=0, vmax=1)
axes[2].set_title("teacher_mark")
axes[2].set_aspect("equal")
plt.colorbar(sc2, ax=axes[2], shrink=0.8)
# Panel 4: correction_label
sc3 = axes[3].scatter(
x, y, c=correction_label, s=8, cmap="coolwarm", vmin=-1, vmax=1
)
axes[3].set_title("correction_label")
axes[3].set_aspect("equal")
plt.colorbar(sc3, ax=axes[3], shrink=0.8)
for ax in axes:
ax.set_xlabel("x")
ax.set_ylabel("y")
plt.tight_layout()
out_path = output_dir / f"{name}.png"
fig.savefig(out_path, dpi=150, bbox_inches="tight")
plt.close(fig)
print(f" ✓ Saved: {out_path}")
# ──────────────────────────────────────────────────────────────────────────────
# 5. Verdict
# ──────────────────────────────────────────────────────────────────────────────
def print_verdict(issues):
"""Print final PASS / WARNING / FAIL."""
print("\n" + "=" * 70)
print(" [5/5] Verdict")
print("=" * 70)
if not issues:
print("\n ✅ PASS — Data is consistent and ready for training.")
return "PASS"
# Classify issues
errors = [i for i in issues if any(k in i.upper() for k in ["MISSING", "BAD ", "MISMATCH"])]
warnings = [i for i in issues if i not in errors]
if errors:
print(f"\n ❌ FAIL — {len(errors)} error(s) found:")
for e in errors:
print(f"{e}")
if warnings:
print(f"\n ⚠️ Also {len(warnings)} warning(s):")
for w in warnings:
print(f"{w}")
return "FAIL"
print(f"\n ⚠️ WARNING — {len(warnings)} suspicious pattern(s):")
for w in warnings:
print(f"{w}")
return "WARNING"
# ──────────────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────────────
def main():
parser = argparse.ArgumentParser(
description="Check step-wise correction dataset quality"
)
parser.add_argument(
"--data-dir", type=str, required=True,
help="Directory containing sampleXXXX_stepYYY.npz files",
)
parser.add_argument(
"--output-dir", type=str, default="outlook/data_correction_check",
help="Directory to save visualization plots",
)
args = parser.parse_args()
data_dir = Path(args.data_dir)
output_dir = Path(args.output_dir)
print("=" * 70)
print(" Correction Data Quality Check")
print("=" * 70)
print(f" Data dir: {data_dir}")
print(f" Output dir: {output_dir}")
# Load
print("\n Loading files...")
samples = load_all_samples(data_dir)
if not samples:
print(" ✗ No sample files found!")
sys.exit(1)
print(f" Loaded {len(samples)} files.")
issues = []
# 1. Field consistency
check_fields(samples, issues)
# 2. Global stats
compute_global_stats(samples)
# 3. Per-step trends
check_per_step_trends(samples, issues)
# 4. Visualization
plot_spatial_samples(samples, output_dir)
# 5. Verdict
verdict = print_verdict(issues)
print()
return 0 if verdict == "PASS" else 1
if __name__ == "__main__":
sys.exit(main())