Testing EQDSK and Boozer with DATA

.vscode/settings.json

@@ -34,6 +34,8 @@
     "python.testing.pytestEnabled": true,
     "python.testing.autoTestDiscoverOnSaveEnabled": true,
     "github.codespaces.devcontainerChangedNotificationStyle": "none",
+    "files.insertFinalNewline": true,
+    "files.trimTrailingWhitespace": true,
     "testing.automaticallyOpenTestResults": "neverOpen",
     "fortran.formatting.formatter": "fprettify",
     "fortran.formatting.fprettifyArgs": [

requirements.txt

@@ -64,6 +64,7 @@ fprettify
 # Misc
 tqdm
 pytest
+pytest-time
 mpi4py
 meson
 meson-python

tests/NEO-2/python/test_get_fluxsurface_area.py

@@ -4,15 +4,16 @@
 
 
 @pytest.fixture
-def run_files():
+def run_files(data_path):
     from libneo import read_eqdsk
 
-    filename_eqdsk = "/proj/plasma/DATA/AUG/EQUILIRBRIUM/QCE_discharge_Neon_shot_39561/eqdsk_39461_5.38s"
+    filename_eqdsk = data_path / "AUG/EQUILIRBRIUM/QCE_discharge_Neon_shot_39561/eqdsk_39461_5.38s"
     eqdsk = read_eqdsk(filename_eqdsk)
     neo2_output = "/temp/grassl_g/Neon_discharge_axisymmetric_gleiter/RUN_lag8/neon_discharge.out"
     yield eqdsk, neo2_output
 
 
+@pytest.mark.xfail(reason="TODO: Update paths")
 def test_get_fluxsurface_area_visual_check(run_files):
     from neo2_ql import get_fluxsurface_area
     from libneo import plot_fluxsurfaces, FluxConverter
@@ -38,4 +39,4 @@ def test_get_fluxsurface_area_visual_check(run_files):
 
 
 if __name__ == "__main__":
-    pytest.main()
+    pytest.main([__file__, "-s"])

tests/NEO-RT/benchmark_with_NEO_2/benchmark_with_NEO_2.ipynb

@@ -49,7 +49,7 @@
    "source": [
     "## Load NEO-2 output\n",
     "\n",
-    "neo2_dir = '/proj/plasma/DATA/DEMO/NEO-2/rerun_DEMO_MARS_PERTURBATION_N1_NEGATIVE_new_pert/'\n",
+    "neo2_dir = os.path.expandvars('$DATA/DEMO/NEO-2/rerun_DEMO_MARS_PERTURBATION_N1_NEGATIVE_new_pert/')\n",
     "neo_input = f90nml.read(os.path.join(neo2_dir, 'neo.in'))\n",
     "neo2_input = f90nml.read(os.path.join(neo2_dir, 'neo2.in'))\n",
     "neo2_output_file = os.path.join(neo2_dir, 'neo2_multispecies_out.h5')\n",

tests/conftest.py

@@ -0,0 +1,12 @@
+import pytest
+
+import os
+from pathlib import Path
+
+@pytest.fixture
+def code_path():
+    return Path(os.environ["CODE"])
+
+@pytest.fixture
+def data_path():
+    return Path(os.environ["DATA"])

tests/libneo/python/test_boozer.py

@@ -1,29 +1,46 @@
-# %%
-from libneo import BoozerFile
-from libneo import read_eqdsk, FluxConverter
+import pytest
+
 import numpy as np
 import matplotlib.pyplot as plt
 
-trial_boozer_file = "/temp/AG-plasma/archive/2023/andfmar/Boozer_files_perturbation_field/ASDEX_U/30835_rmpm90/out_neo-2_rmp_m90-n0"
-trial_eqdsk_file = "/proj/plasma/DATA/AUG/EQDSK/g30835.3200_ed6"
-bc = BoozerFile(trial_boozer_file)
-eqdsk = read_eqdsk(trial_eqdsk_file)
-converter = FluxConverter(eqdsk["qprof"])
-no_edges = slice(10,-10)
+from libneo import BoozerFile
+from libneo import read_eqdsk, FluxConverter
+
+no_edges = slice(10, -10)
+
 
-def test_iota():
-    stor_bc = bc.s[no_edges]
-    iota_bc = np.array(bc.iota[no_edges])
-    stor_eqdsk = converter.sqrtspol2stor(np.sqrt(eqdsk["rho_poloidal"]))
-    iota_eqdsk = np.interp(stor_bc, stor_eqdsk, 1/eqdsk["qprof"])
-    assert np.allclose(iota_eqdsk, -iota_bc, atol=1e-2) # .bc file has only positive itota
+@pytest.fixture
+def boozer(data_path):
+    trial_boozer_file = data_path / "AUG/BOOZER/30835/out_neo-2_rmp_90-n0"
+    return BoozerFile(str(trial_boozer_file))
 
 
-def test_get_rho_poloidal():
-    sqrtspol = bc.get_rho_poloidal()[no_edges]
-    sqrtspol_control = converter.stor2sqrtspol(bc.s)[no_edges]
+@pytest.fixture
+def eqdsk(data_path):
+    trial_eqdsk_file = data_path / "AUG/EQDSK/g30835.3200_ed6"
+    return read_eqdsk(str(trial_eqdsk_file))
+
+
+@pytest.fixture
+def converter(eqdsk):
+    return FluxConverter(eqdsk["qprof"])
+
+
+def test_iota(boozer, eqdsk, converter):
+    stor_bc = boozer.s[no_edges]
+    iota_bc = np.array(boozer.iota[no_edges])
+    stor_eqdsk = converter.sqrtspol2stor(np.sqrt(eqdsk["s_pol"]))
+    iota_eqdsk = np.interp(stor_bc, stor_eqdsk, 1 / eqdsk["qprof"])
+    assert np.allclose(
+        iota_eqdsk, -iota_bc, atol=1e-2
+    )  # .bc file has only positive itota
+
+
+def test_get_rho_poloidal(boozer, converter):
+    sqrtspol = boozer.get_rho_poloidal()[no_edges]
+    sqrtspol_control = converter.stor2sqrtspol(boozer.s)[no_edges]
     assert np.allclose(sqrtspol, sqrtspol_control, atol=1e-2)
 
-if __name__=="__main__":
-    test_iota()
-    test_get_rho_poloidal()
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-s"])

tests/libneo/python/test_efit_to_boozer.py

@@ -0,0 +1,135 @@
+import pytest
+
+import os
+import tempfile
+from pathlib import Path
+import numpy as np
+from numpy.testing import assert_allclose
+import matplotlib.pyplot as plt
+
+import _efit_to_boozer as efit_to_boozer
+from libneo.boozer import get_angles_and_transformation
+from libneo import read_eqdsk, FluxConverter
+
+efit_to_boozer_input = """3600      nstep    - number of integration steps
+500       nlabel   - grid size over radial variable
+500       ntheta   - grid size over poloidal angle
+10000     nsurfmax - number of starting points for separarix search
+2000      nsurf    - number of flux surfaces in Boozer file
+12        mpol     - number of poloidal modes in Boozer file
+2119063900.d0 psimax - psi of plasma boundary
+"""
+
+field_divB0_input = """0                                 ipert        ! 0=eq only, 1=vac, 2,3=vac+plas
+1                                 iequil       ! 0=pert. alone, 1=with equil.
+1.00                              ampl         ! amplitude of perturbation, a.u.
+72                                ntor         ! number of toroidal harmonics
+0.99                              cutoff       ! inner cutoff in psi/psi_a units
+4                                 icftype      ! type of coil file
+'{gfile}'                         gfile        ! equilibrium file
+'unused_field.dat'                pfile        ! coil        file
+'unused_convexwall.dat'           convexfile   ! convex file for stretchcoords
+'UNUSED_FLUXDATA'                 fluxdatapath ! directory with data in flux coord.
+0                                 window size for filtering of psi array over R
+0                                 window size for filtering of psi array over Z
+"""
+
+@pytest.fixture
+def test_files(code_path, data_path):
+    return {
+        "local": code_path / "libneo/test/resources/input_efit_file.dat",
+        "DEMO CHEASE": data_path / "DEMO/EQDSK/Equilibrium_DEMO2019_CHEASE/MOD_Qprof_Test/EQDSK_DEMO2019_q1_COCOS_02.OUT",
+        "AUG": data_path / "AUG/EQDSK/g30835.3200_ed6",
+        "MASTU": data_path / "MASTU/EQDSK/MAST_47051_450ms.geqdsk",
+        # TODO: "DEMO PROCESS": data_path / "DEMO/EQDSK/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF.eqdsk",
+        # TODO: "DEMO standardized": data_path / "DEMO/EQDSK/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF_std.eqdsk",
+    }
+
+
+@pytest.mark.slow
+def test_angles_and_transformation(test_files):
+    stor = 0.4
+    num_theta = 64
+    for key, test_file in test_files.items():
+        print(f"Testing {key} EQDSK file: {test_file}")
+        tmp_path = init_run_path(test_file)
+        print(f"Running in {tmp_path}")
+        os.chdir(tmp_path)
+        efit_to_boozer.efit_to_boozer.init()
+        th_boozers, th_geoms, Gs = get_angles_and_transformation(stor, num_theta)
+
+        plt.figure()
+        plt.plot(th_boozers, th_geoms)
+        plt.xlabel(r"$\theta_{\mathrm{Boozer}}$")
+        plt.ylabel(r"$\theta_{\mathrm{geom}}$")
+        plt.title(f"{key}")
+
+
+@pytest.mark.slow
+def test_q_profile_eqdsk(test_files):
+    for key in ["DEMO CHEASE", "MASTU"]:
+        # TODO: local, AUG, MASTU, PROCESS, standardized
+        test_file = test_files[key]
+        print(f"Testing {key} EQDSK file: {test_file}")
+        tmp_path = init_run_path(test_file)
+        print(f"Running in {tmp_path}")
+        os.chdir(tmp_path)
+        eqdsk_data = read_eqdsk(str(test_file))
+
+        # The data in EQDSK is writting in poloidal flux label,
+        # but efit_to_boozer uses toroidal flux label. Therefore, we need a mapping
+        q_profile_eqdsk = eqdsk_data["qprof"]
+        spol_eqdsk = np.linspace(0.0, 1.0, q_profile_eqdsk.shape[0])
+        converter = FluxConverter(q_profile_eqdsk)
+        stor_eqdsk = converter.spol2stor(spol_eqdsk)
+
+        # The safety factor can be calculated alternatively using field line integration
+        # this is done as part of the symmetry flux transformation and given as output here
+
+        q_profile_field_line_integration = []
+
+        # inp_label = 1 makes it that efit_to_boozer.magdata_in_symfluxcoord_ext uses the
+        # flux label si instead of the actual flux psi to determine which
+        # flux surface one is on. Therefore psi is just set as a dummy variable here.
+        inp_label = 1
+        psi = np.array(0.0)
+        theta = np.array(0.0)
+
+        efit_to_boozer.efit_to_boozer.init()
+        for si in stor_eqdsk:
+            (q, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _) = (
+                efit_to_boozer.magdata_in_symfluxcoord_ext(inp_label, si, psi, theta)
+            )
+            q_profile_field_line_integration.append(q)
+
+        q_profile_field_line_integration = np.array(q_profile_field_line_integration)
+
+        plt.figure()
+        plt.plot(spol_eqdsk, q_profile_eqdsk, "-r", label=r"q profile from EQDSK")
+        plt.plot(spol_eqdsk, q_profile_field_line_integration, "--b", label=r"q profile from field line integration")
+        plt.title(f"{key}")
+        plt.legend()
+
+        assert_allclose(q_profile_field_line_integration, q_profile_eqdsk, rtol=1e-2)
+        print("Alternative safety factor calculation agrees with EQDSK file within 1%")
+
+
+def init_run_path(gfile):
+    tmp_path = Path(tempfile.mkdtemp())
+    write_input_files(tmp_path, gfile)
+    return tmp_path
+
+
+def write_input_files(tmp_path, gfile):
+    efit_to_boozer_input_file = tmp_path / "efit_to_boozer.inp"
+    efit_to_boozer_input_file.write_text(efit_to_boozer_input)
+
+    field_divB0_input_file = tmp_path / "field_divB0.inp"
+    field_divB0_input_file.write_text(field_divB0_input.format(gfile=gfile))
+
+
+if __name__ == "__main__":
+    try:
+        pytest.main([__file__, "-s"])
+    finally:
+        plt.show()