drivaernet_data.py

"""Real DrivAerNet++ data access — retrieval & statistics.

DrivAerNet++ (Elrefaie et al., MIT, CC-BY-NC-4.0) is a public dataset of
8,150 car designs each with high-fidelity OpenFOAM CFD results.  The full
dataset is 39 TB (meshes + volume fields) hosted on Harvard Dataverse,
but the per-design aerodynamic drag coefficients ship as a small CSV that
we vendor locally under ./DrivAerNet/.

This module mirrors the *Simulation Agent* retrieval idea from the
"AI Agents in Engineering Design" paper: rather than running CFD on a new
shape, look up the closest REAL pre-computed cases and report their
measured drag.  It needs no torch and no mesh download — just the CSV.

Design-ID morphology codes (first token of the ID):
    F = Fastback    N = Notchback (sedan)    E = Estate (wagon)

Functions:
    load_cd()              -> dict {design_id: Cd}
    dataset_stats()        -> overall + per-class statistics
    retrieve_by_cd(target) -> closest real designs to a target Cd
    class_means()          -> mean Cd per body morphology
"""
from __future__ import annotations

import csv
from pathlib import Path
from typing import Dict, List, Optional, Tuple

_DATA_DIR = Path(__file__).parent / "DrivAerNet"
_CD_CSV = _DATA_DIR / "DrivAerNetPlusPlus_Cd_8k.csv"

_CLASS_NAMES = {"F": "Fastback", "N": "Notchback (sedan)", "E": "Estate (wagon)"}

_CACHE: Optional[Dict[str, float]] = None


def available() -> bool:
    return _CD_CSV.exists()


def load_cd() -> Dict[str, float]:
    """Return {design_id: Cd} from the vendored DrivAerNet++ CSV."""
    global _CACHE
    if _CACHE is not None:
        return _CACHE
    out: Dict[str, float] = {}
    if not _CD_CSV.exists():
        return out
    with open(_CD_CSV) as f:
        rd = csv.DictReader(f)
        id_key = rd.fieldnames[0]
        cd_key = next((k for k in rd.fieldnames
                       if k and ("drag" in k.lower() or k.lower() == "cd")),
                      rd.fieldnames[-1])
        for row in rd:
            try:
                out[str(row[id_key]).strip()] = float(row[cd_key])
            except (ValueError, KeyError, TypeError):
                continue
    _CACHE = out
    return out


def _morphology(design_id: str) -> str:
    return design_id.split("_", 1)[0].upper()


def class_means() -> Dict[str, Tuple[str, int, float, float, float]]:
    """Per-morphology (name, n, mean, min, max) of real Cd."""
    cd = load_cd()
    buckets: Dict[str, List[float]] = {}
    for k, v in cd.items():
        buckets.setdefault(_morphology(k), []).append(v)
    res = {}
    for c, vals in sorted(buckets.items()):
        res[c] = (_CLASS_NAMES.get(c, c), len(vals),
                  sum(vals) / len(vals), min(vals), max(vals))
    return res


def dataset_stats() -> dict:
    cd = load_cd()
    if not cd:
        return {"available": False}
    vals = list(cd.values())
    n = len(vals)
    mean = sum(vals) / n
    var = sum((x - mean) ** 2 for x in vals) / n
    return {
        "available": True,
        "n": n,
        "cd_min": min(vals),
        "cd_mean": mean,
        "cd_max": max(vals),
        "cd_std": var ** 0.5,
        "classes": class_means(),
        "source": str(_CD_CSV.name),
    }


def retrieve_by_cd(target_cd: float, n: int = 5,
                   morphology: Optional[str] = None
                   ) -> List[Tuple[str, float, float]]:
    """Return the n real designs whose Cd is closest to target_cd.

    Each item is (design_id, cd, abs_error).  Optionally restrict to a
    morphology class ("F"/"N"/"E").  Mirrors the paper's "show me designs
    with drag near X" retrieval, but over the REAL CFD dataset.
    """
    cd = load_cd()
    items = cd.items()
    if morphology:
        m = morphology.upper()
        items = [(k, v) for k, v in items if _morphology(k) == m]
    else:
        items = list(items)
    ranked = sorted(items, key=lambda kv: abs(kv[1] - target_cd))
    return [(k, v, abs(v - target_cd)) for k, v in ranked[:n]]


if __name__ == "__main__":
    s = dataset_stats()
    if not s["available"]:
        print("DrivAerNet++ CSV not found — run fetch_drivaernet to download it.")
        raise SystemExit(1)
    print(f"DrivAerNet++ real drag data  ({s['source']})")
    print(f"  N={s['n']}  Cd mean={s['cd_mean']:.3f}  "
          f"std={s['cd_std']:.3f}  range [{s['cd_min']:.3f}, {s['cd_max']:.3f}]")
    print("  per body class:")
    for c, (name, cnt, mean, lo, hi) in s["classes"].items():
        print(f"    {c} {name:18s} n={cnt:4d}  mean={mean:.3f}  [{lo:.3f}, {hi:.3f}]")
    print("\n  retrieval demo - closest real designs to Cd=0.25:")
    for did, cd, err in retrieve_by_cd(0.25, n=5):
        print(f"    {did:18s}  Cd={cd:.4f}  (d={err:.4f})")