upgrade_tensor_baselines.py

#!/usr/bin/env python3
"""Official TensorMesh vs DOLFINx external-baseline benchmark.

Runs the exact solver classes from camlab-ethz/tensormesh-bench on the same
CPU container, same Gmsh characteristic lengths, BiCGSTAB+Jacobi tolerances,
and repeated 3-D Poisson / hollow-cube linear-elasticity solves.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import io
import json
import os
from pathlib import Path
import platform
import sys
import tarfile
import time
import urllib.request

import numpy as np


TM_COMMIT = "c7d099224720b6a0e49c2c9a77b0163258847880"
BENCH_BRANCH = "main"


def download_extract(url: str, target: Path) -> None:
    target.mkdir(parents=True, exist_ok=True)
    data = urllib.request.urlopen(url, timeout=180).read()
    with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tf:
        root = tf.getmembers()[0].name.split("/")[0]
        for member in tf.getmembers():
            if member.name == root:
                continue
            member.name = member.name[len(root) + 1 :]
            if member.name:
                tf.extract(member, target)


def timed(fn, repeats: int) -> tuple[list[float], list[float]]:
    times, residuals = [], []
    for _ in range(repeats):
        start = time.perf_counter()
        out = fn()
        times.append(time.perf_counter() - start)
        residuals.append(float(out))
    return times, residuals


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--output", type=Path, default=Path("/outputs/external-baselines-v4"))
    ap.add_argument("--repeats", type=int, default=5)
    args = ap.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    work = Path("/tmp/tg_external")
    tm_root, bench_root = work / "TensorMesh", work / "bench"
    download_extract(
        f"https://github.com/camlab-ethz/TensorMesh/archive/{TM_COMMIT}.tar.gz", tm_root
    )
    download_extract(
        f"https://github.com/camlab-ethz/tensormesh-bench/archive/refs/heads/{BENCH_BRANCH}.tar.gz",
        bench_root,
    )
    pipeline = bench_root / "pipeline"
    sys.path[:0] = [
        str(tm_root),
        str(bench_root),
        str(pipeline),
        str(pipeline / "utils"),
        str(pipeline / "benchmarks" / "poisson"),
        str(pipeline / "benchmarks" / "elasticity"),
    ]
    import torch
    import tensormesh as tm
    from mpi4py import MPI
    from benchmark_fenics import DOLFINxPoisson, create_gmsh_mesh as fenics_poisson_mesh
    from benchmark_tensormesh import TmFEM
    from benchmark_fenics_elasticity import (
        DOLFINxElasticity,
        create_gmsh_mesh as fenics_elastic_mesh,
    )
    from benchmark_tensormesh_elasticity import TmElasticityFEM

    torch.set_num_threads(1)
    rows = []
    for equation, lengths in (("poisson", (0.20, 0.15, 0.10, 0.07)), ("elasticity", (0.20, 0.15, 0.10))):
        for h in lengths:
            if equation == "poisson":
                fmesh = fenics_poisson_mesh(MPI.COMM_WORLD, 3, h)
                ffem = DOLFINxPoisson(fmesh)
                ffem()
                ffem()
                ft, fr = timed(lambda: (lambda u: ffem.compute_residual(u))(ffem()), args.repeats)

                cache = f"/tmp/tm_poisson_{h}.msh"
                tmesh = tm.Mesh.gen_cube(chara_length=h, cache_path=cache).double()
                tfem = TmFEM(tmesh)
                tfem()
                tfem()
                tt, tr = timed(lambda: (lambda u: tfem.compute_residual(u))(tfem()), args.repeats)
                fdofs = int(fmesh.geometry.index_map().size_global)
                tdofs = int(tmesh.n_points)
            else:
                fmesh = fenics_elastic_mesh(MPI.COMM_WORLD, 3, h, hollow=True)
                ffem = DOLFINxElasticity(fmesh, 3)
                ffem()
                ffem(return_residual=True)
                ft, fr = timed(lambda: ffem(return_residual=True)[1], args.repeats)

                cache = f"/tmp/tm_elastic_{h}.msh"
                tmesh = tm.Mesh.gen_hollow_cube(chara_length=h, cache_path=cache).double()
                tfem = TmElasticityFEM(tmesh)
                tfem()
                tfem(return_residual=True)
                tt, tr = timed(lambda: tfem(return_residual=True)[1], args.repeats)
                fdofs = int(fmesh.geometry.index_map().size_global * 3)
                tdofs = int(tmesh.n_points * 3)
            row = {
                "equation": equation,
                "dimension": 3,
                "characteristic_length": h,
                "fenics_dofs": fdofs,
                "tensormesh_dofs": tdofs,
                "dof_relative_difference": abs(fdofs - tdofs) / max(fdofs, tdofs),
                "repeats": args.repeats,
                "fenics_median_s": float(np.median(ft)),
                "tensormesh_median_s": float(np.median(tt)),
                "fenics_times_s": json.dumps(ft),
                "tensormesh_times_s": json.dumps(tt),
                "fenics_median_residual": float(np.median(fr)),
                "tensormesh_median_residual": float(np.median(tr)),
                "speedup_fenics_over_tensormesh": float(np.median(ft) / np.median(tt)),
            }
            rows.append(row)
            print(json.dumps(row, sort_keys=True), flush=True)
    with (args.output / "external_baselines.csv").open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]))
        w.writeheader()
        w.writerows(rows)
    result = {
        "paper": "Learning, Solving and Optimizing PDEs with TensorGalerkin",
        "openreview_id": "xpIUvw8ANa",
        "tensormesh_commit": TM_COMMIT,
        "benchmark_repository": "https://github.com/camlab-ethz/tensormesh-bench",
        "rows": len(rows),
        "repeats_per_framework_cell": args.repeats,
        "poisson_speedups": [r["speedup_fenics_over_tensormesh"] for r in rows if r["equation"] == "poisson"],
        "elasticity_speedups": [r["speedup_fenics_over_tensormesh"] for r in rows if r["equation"] == "elasticity"],
        "environment": {
            "python": sys.version,
            "platform": platform.platform(),
            "torch": torch.__version__,
            "dolfinx": __import__("dolfinx").__version__,
            "cpu_threads": torch.get_num_threads(),
        },
        "scope": "Single-rank, one-thread CPU control; paper used 8 MPI ranks, so exact paper timing ratios are not claimed.",
    }
    (args.output / "results.json").write_text(json.dumps(result, indent=2) + "\n")
    manifest = {}
    for p in sorted(args.output.glob("*")):
        if p.name != "SHA256SUMS.json":
            manifest[p.name] = hashlib.sha256(p.read_bytes()).hexdigest()
    (args.output / "SHA256SUMS.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps(result, indent=2), flush=True)


if __name__ == "__main__":
    main()