File size: 4,479 Bytes
3e77c56 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Download + verify the Geo-FNO Elasticity dataset.
Source: the Geo-FNO Google Drive collection (license MIT, ``neuraloperator/Geo-FNO``).
IMPORTANT: the canonical Drive *folder*
``https://drive.google.com/drive/folders/1YBuaoTdOSr_qzaow-G-iwvbUI7fiUzu8`` is the **entire**
neuraloperator dataset collection (~8 GB: airfoil, car-cfd, channel-shocks, elasticity, ...).
For this project we only need three files (~47 MB total), so we fetch them **by file id**
rather than downloading the whole folder. The ids below were resolved once from that folder
(``gdown.download_folder(..., skip_download=True)``) — see ``docs/RECONCILIATION.md``.
Usage:
python -m stress_operator.data.download --out data
python -m stress_operator.data.download --out data --report-only # just verify + print shapes
"""
from __future__ import annotations
import argparse
import hashlib
import os
import sys
import numpy as np
# Resolved file ids for elasticity/Meshes/* inside the Geo-FNO Drive collection.
DRIVE_FILE_IDS = {
"Random_UnitCell_sigma_10.npy": "1Ia5izgUum-IQLdO6PW70HO8AdAqA_IVb",
"Random_UnitCell_XY_10.npy": "1I-fO-RsFvD3nqBuFrg67R0yqTFdD_gpA",
"Random_UnitCell_rr_10.npy": "1Pjliqhxegoe5VpoLrpBa9n3P4gX9MfTt",
}
# Mirror Transolver's expected layout: <out>/elasticity/Meshes/<file>.
SUBDIR = os.path.join("elasticity", "Meshes")
REQUIRED_FILES = ["Random_UnitCell_sigma_10.npy", "Random_UnitCell_XY_10.npy"]
# Authoritative raw shapes (sample axis LAST). See docs/RECONCILIATION.md.
EXPECTED_RAW_SHAPES = {
"Random_UnitCell_sigma_10.npy": (972, 2000),
"Random_UnitCell_XY_10.npy": (972, 2, 2000),
"Random_UnitCell_rr_10.npy": (42, 2000),
}
def _find(data_dir: str, name: str):
"""Locate ``name`` anywhere under ``data_dir`` (handles nested layouts)."""
direct = os.path.join(data_dir, name)
if os.path.isfile(direct):
return direct
for root, _dirs, files in os.walk(data_dir):
if name in files:
return os.path.join(root, name)
return None
def have_all(data_dir: str) -> bool:
return all(_find(data_dir, f) is not None for f in REQUIRED_FILES)
def download(out_dir: str, force: bool = False) -> None:
import gdown
dest_dir = os.path.join(out_dir, SUBDIR)
os.makedirs(dest_dir, exist_ok=True)
for name, file_id in DRIVE_FILE_IDS.items():
existing = _find(out_dir, name)
if existing is not None and not force:
print(f"[download] {name}: already present ({existing}); skipping.")
continue
out_path = os.path.join(dest_dir, name)
print(f"[download] fetching {name} (id={file_id}) -> {out_path}")
gdown.download(id=file_id, output=out_path, quiet=False)
def _md5(path: str, chunk: int = 1 << 20) -> str:
h = hashlib.md5()
with open(path, "rb") as f:
for block in iter(lambda: f.read(chunk), b""):
h.update(block)
return h.hexdigest()
def report(data_dir: str) -> int:
"""Print shape/dtype/stats and check against expected raw shapes. Returns exit code."""
ok = True
print(f"[report] scanning {data_dir!r}")
for name in DRIVE_FILE_IDS:
path = _find(data_dir, name)
required = name in REQUIRED_FILES
if path is None:
print(f" - {name}: {'MISSING (required)' if required else 'missing (optional)'}")
if required:
ok = False
continue
arr = np.load(path)
exp = EXPECTED_RAW_SHAPES.get(name)
flag = "" if (exp is None or tuple(arr.shape) == tuple(exp)) else f" <-- WARNING: expected {exp}"
print(
f" - {name}: shape={tuple(arr.shape)} dtype={arr.dtype} "
f"min={float(arr.min()):.4g} max={float(arr.max()):.4g} md5={_md5(path)[:8]}{flag}"
)
print("[report] OK" if ok else "[report] FAILED: required files missing")
return 0 if ok else 1
def main() -> int:
ap = argparse.ArgumentParser(description="Download/verify the Geo-FNO Elasticity dataset (3 files).")
ap.add_argument("--out", default="data", help="output directory")
ap.add_argument("--report-only", action="store_true", help="skip download, just verify + print shapes")
ap.add_argument("--force", action="store_true", help="re-download even if files exist")
args = ap.parse_args()
if not args.report_only:
download(args.out, force=args.force)
return report(args.out)
if __name__ == "__main__":
sys.exit(main())
|