vmecpp-benchmarks/run_vmec2000.py at main · proximafusion/vmecpp-benchmarks · GitHub

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
# SPDX-FileCopyrightText: 2025-present Proxima Fusion GmbH <info@proximafusion.com>
#
# SPDX-License-Identifier: MIT
"""
Run this under mpirun!
"""

from pathlib import Path
import argparse
import time
import mpi4py

print("Begin importing simsopt.mhd (often slow because of MPI)...")
start_import = time.time()
import simsopt.mhd  # noqa: E402
end_import = time.time()
print(f"End importing simsopt.mhd. Took {end_import - start_import} s")


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("vmec_indata_file", type=Path)
    p.add_argument("n_cores", type=int)
    args = p.parse_args()

    vmec = simsopt.mhd.Vmec(str(args.vmec_indata_file))
    start = time.time()
    vmec.run()
    runtime = time.time() - start

    p.add_argument("n_cores", type=int)

    if mpi4py.MPI.COMM_WORLD.Get_rank() == 0:
        print(f"VMEC2000,{args.vmec_indata_file},{args.n_cores},{runtime}")