diff --git a/.vscode/settings.json b/.vscode/settings.json index ed05e65..7531d4f 100644 --- a/.vscode/settings.json +++ b/.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": [ diff --git a/requirements.txt b/requirements.txt index 0e6df08..cb005ba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -64,6 +64,7 @@ fprettify # Misc tqdm pytest +pytest-time mpi4py meson meson-python diff --git a/tests/NEO-2/python/test_get_fluxsurface_area.py b/tests/NEO-2/python/test_get_fluxsurface_area.py index eb1a191..ab2fa13 100644 --- a/tests/NEO-2/python/test_get_fluxsurface_area.py +++ b/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"]) diff --git a/tests/NEO-RT/benchmark_with_NEO_2/benchmark_with_NEO_2.ipynb b/tests/NEO-RT/benchmark_with_NEO_2/benchmark_with_NEO_2.ipynb index 811620d..64f93f2 100644 --- a/tests/NEO-RT/benchmark_with_NEO_2/benchmark_with_NEO_2.ipynb +++ b/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", diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..035d0df --- /dev/null +++ b/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"]) diff --git a/tests/efit_to_boozer/python/convexwall.dat b/tests/efit_to_boozer/python/convexwall.dat deleted file mode 100644 index b855e27..0000000 --- a/tests/efit_to_boozer/python/convexwall.dat +++ /dev/null @@ -1,5 +0,0 @@ -580d0 -620d0 -1220d0 -620d0 -1220d0 640d0 -580d0 640d0 -580d0 -620d0 \ No newline at end of file diff --git a/tests/efit_to_boozer/python/efit_to_boozer.inp b/tests/efit_to_boozer/python/efit_to_boozer.inp deleted file mode 100644 index 974ed82..0000000 --- a/tests/efit_to_boozer/python/efit_to_boozer.inp +++ /dev/null @@ -1,7 +0,0 @@ -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 \ No newline at end of file diff --git a/tests/efit_to_boozer/python/field_divB0.inp b/tests/efit_to_boozer/python/field_divB0.inp deleted file mode 100644 index 7c4a2ec..0000000 --- a/tests/efit_to_boozer/python/field_divB0.inp +++ /dev/null @@ -1,36 +0,0 @@ -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 -'' !'DATA/ASDEX/MESH3D_30835/field.dat' pfile ! coil file -'convexwall.dat' convexfile ! convex file for stretchcoords -'' 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 - - - - - -More detailed comments on switches. -Axisymmetric equilibrium is called in all cases. If some routines are not -called, their data is not required. - -Perturbation switch "ipert": -0 - axisymmetric equilibrium is called alone (can be run withot coild data - and resonance plasma response data) -1 - cylindrical routine is called in addition to equilibrium (plasma response - data is not required -2 - both, cylindrical and flux coordinate routine modelling plasma response - are called, all data is required. Derivatives are not computed -3 - the same as 2 with computation of derivatives (7 times slower) - -Equilibrium switch "iequil": -0 - equilibrium field is not added to the output (needed for computing the - perturbation field alone). The equilibrium routine is called - its - data is needed for plasma response field. For future: this call can be - disabled for stellarator modelling. -1 - normal mode, perturbation is added to the output field \ No newline at end of file diff --git a/tests/efit_to_boozer/python/gfile b/tests/efit_to_boozer/python/gfile deleted file mode 120000 index 44bb0a1..0000000 --- a/tests/efit_to_boozer/python/gfile +++ /dev/null @@ -1 +0,0 @@ -/proj/plasma/DATA/DEMO/teams/Equilibrium_DEMO2019_CHEASE/MOD_Qprof_Test/EQDSK_DEMO2019_q1_COCOS_02.OUT \ No newline at end of file diff --git a/tests/efit_to_boozer/python/test_q_profile_eqdsk.py b/tests/efit_to_boozer/python/test_q_profile_eqdsk.py deleted file mode 100644 index 8dc15dc..0000000 --- a/tests/efit_to_boozer/python/test_q_profile_eqdsk.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Tests for efit_2_boozer and eqdsk_base -""" - -import pytest - -# Import standard modules -import numpy as np -import os - -# Import modules to be tested -import efit2boozer - -from numpy.testing import assert_allclose -from libneo import read_eqdsk, FluxLabelConverter - - -def test_q_profile_eqdsk(): - - # Read data from EQDSK file - basedir_eqdsk = "/proj/plasma/DATA/DEMO/teams/Equilibrium_DEMO2019_CHEASE" - eqfile_eqdsk = "MOD_Qprof_Test/EQDSK_DEMO2019_q1_COCOS_02.OUT" - - eqdsk_data = read_eqdsk(os.path.join(basedir_eqdsk, eqfile_eqdsk)) - - # The data in EQDSK is writting in poloidal flux label, - # but efit2boozer 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 = FluxLabelConverter(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 efit2boozer.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) - - # current_directory = os.path.dirname(os.path.realpath(__file__)) - current_directory = os.getcwd() - test_directory = os.path.join(current_directory, "tests/efit_to_boozer/python") - os.chdir(test_directory) - - efit2boozer.efit_to_boozer.init() - for si in stor_eqdsk: - (q, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _) = ( - efit2boozer.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) - - os.chdir(current_directory) - - assert_allclose(q_profile_field_line_integration, q_profile_eqdsk, rtol=1e-2) - print("Alternative safety factor calculation agrees with EQDSK file within 1%") diff --git a/tests/libneo/python/test_boozer.py b/tests/libneo/python/test_boozer.py index 109fd25..4bf6bbe 100644 --- a/tests/libneo/python/test_boozer.py +++ b/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() \ No newline at end of file + +if __name__ == "__main__": + pytest.main([__file__, "-s"]) diff --git a/tests/libneo/python/test_efit_to_boozer.py b/tests/libneo/python/test_efit_to_boozer.py new file mode 100644 index 0000000..96bde3f --- /dev/null +++ b/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() diff --git a/tests/libneo/python/test_eqdsk.py b/tests/libneo/python/test_eqdsk.py index 6b29320..8272d83 100644 --- a/tests/libneo/python/test_eqdsk.py +++ b/tests/libneo/python/test_eqdsk.py @@ -6,73 +6,73 @@ import json import gzip -import os.path from libneo import eqdsk -# Directory where the golden records are stored for regression tests -golden_record_dir = "/proj/plasma/DATA/TESTS/libneo/eqdsk" +@pytest.fixture +def test_files(code_path, data_path): + return { + "local": code_path / "libneo/test/resources/input_efit_file.dat", + "PROCESS": data_path / "DEMO/EQDSK/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF.eqdsk", + "standardized": data_path / "DEMO/EQDSK/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF_std.eqdsk", + "AUG": data_path / "AUG/EQDSK/g30835.3200_ed6", + # TODO: "MASTU": data_path / "MASTU/EQDSK/MAST_47051_450ms.geqdsk", + # TODO: "CHEASE": data_path / "DEMO/teams/Equilibrium_DEMO2019_CHEASE/MOD_Qprof_Test/EQDSK_DEMO2019_q1_COCOS_02.OUT", + } -test_files = [ - - # Local - "libneo/tests/resources/input_efit_file.dat", - # PROCESS code - "/proj/plasma/DATA/DEMO/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF.eqdsk", - - # Standardized - "/proj/plasma/DATA/DEMO/Equil_2021_PMI_QH_mode_betap_1d04_li_1d02_Ip_18d27MA_SOF_std.eqdsk", - - # TODO: CHEASE - # "/proj/plasma/DATA/DEMO/teams/Equilibrium_DEMO2019_CHEASE/MOD_Qprof_Test/EQDSK_DEMO2019_q1_COCOS_02.OUT" -] - - -def test_eqdsk_read(): - for test_file in test_files: +def test_eqdsk_read(test_files): + for key, test_file in test_files.items(): + print(f"Testing {key} EQDSK file: {test_file}") _ = eqdsk.eqdsk_file(test_file) -def test_eqdsk_golden_records(): - for test_file in test_files: +@pytest.mark.slow +def test_eqdsk_golden_records(data_path, test_files): + golden_record_path = data_path / "TESTS/libneo/eqdsk" + store_golden_records(test_files.values(), golden_record_path) + + for test_file in test_files.values(): eqdsk_object = eqdsk.eqdsk_file(test_file) data = eqdsk_object.__dict__ replace_array_members_by_lists(data) - assert are_dicts_equal(data, get_golden_record(test_file)) + assert are_dicts_equal(data, get_golden_record(test_file, golden_record_path)) -def get_golden_record(filename): +def get_golden_record(file_path, storage_path): """ Returns the golden record for the given test file. """ - file_base, _ = os.path.splitext(os.path.basename(filename)) - with gzip.open(f"{golden_record_dir}/{file_base}.json.gz", "rb") as f: + file_base = file_path.stem + with gzip.open(storage_path / f"{file_base}.json.gz", "rb") as f: compressed_data = f.read() json_data = compressed_data.decode("utf-8") return json.loads(json_data) -def store_golden_records(): +def store_golden_records(file_paths, storage_path, force_overwrite=False): """ Stores reference data for regression tests. Call only manually after fixing a bug and checking that the new results are correct. """ - for file in test_files: - basedir = os.path.dirname(file) - file_base, ext = os.path.splitext(os.path.basename(file)) - data = eqdsk.eqdsk_file(f"{basedir}/{file_base}{ext}").__dict__ + for f in file_paths: + basedir = f.parent + file_base = f.stem + ext = f.suffix + data = eqdsk.eqdsk_file(str(f)).__dict__ replace_array_members_by_lists(data) - outfile = f"{golden_record_dir}/{file_base}.json.gz" + outfile = storage_path / f"{file_base}.json.gz" - if os.path.isfile(outfile): - raise RuntimeError(f"Golden record file {outfile} already exists.") + if outfile.exists(): + print(f"Golden record {outfile} exists.") + if not force_overwrite: + continue with gzip.open(outfile, "wb") as f: json_data = json.dumps(data, indent=4) @@ -137,4 +137,8 @@ def are_float_lists_equal(list1, list2, tolerance=1e-9): elif isinstance(a, list) and isinstance(b, list): if not are_float_lists_equal(a, b, tolerance): return False - return True \ No newline at end of file + return True + + +if __name__ == "__main__": + pytest.main([__file__, "-s"]) diff --git a/tests/libneo/python/test_mgrid.py b/tests/libneo/python/test_mgrid.py index c1ca27e..6407afb 100644 --- a/tests/libneo/python/test_mgrid.py +++ b/tests/libneo/python/test_mgrid.py @@ -1,33 +1,48 @@ +import pytest + import numpy as np from libneo.mgrid import MgridFile -def test_mgrid_read(): - f = MgridFile.from_file('/proj/plasma/DATA/LHD/albert/vmec/makegrid_alternative/mgrid_lhd_nfp10.nc') +@pytest.fixture +def testfile_path(data_path): + return data_path / "LHD/VMEC/makegrid_alternative/mgrid_lhd_nfp10.nc" + + +@pytest.fixture +def mgrid(testfile_path): + return MgridFile.from_file(testfile_path) + + +def test_mgrid_read(testfile_path): + f = MgridFile.from_file(testfile_path) print(f) -def test_mgrid_read_write(): - f = MgridFile.from_file('/proj/plasma/DATA/LHD/albert/vmec/makegrid_alternative/mgrid_lhd_nfp10.nc') - f.write('/tmp/mgrid_test.nc') + +def test_mgrid_read_write(mgrid): + mgrid.write("/tmp/mgrid_test.nc") -def test_mgrid_read_write_read(): - f = MgridFile.from_file('/proj/plasma/DATA/LHD/albert/vmec/makegrid_alternative/mgrid_lhd_nfp10.nc') - f.write('/tmp/mgrid_test.nc') - g = MgridFile.from_file('/tmp/mgrid_test.nc') +def test_mgrid_read_write_read(mgrid): + mgrid.write("/tmp/mgrid_test.nc") + g = MgridFile.from_file("/tmp/mgrid_test.nc") # Assert if all variables are the same - assert f.ir == g.ir - assert f.jz == g.jz - assert f.kp == g.kp - assert f.nfp == g.nfp - assert f.rmin == g.rmin - assert f.zmin == g.zmin - assert f.rmax == g.rmax - assert f.zmax == g.zmax - assert f.mgrid_mode == g.mgrid_mode - assert np.array_equal(f.coil_group, g.coil_group) - assert np.array_equal(f.coil_current, g.coil_current) - assert np.array_equal(f.br, g.br) - assert np.array_equal(f.bp, g.bp) - assert np.array_equal(f.bz, g.bz) + assert mgrid.ir == g.ir + assert mgrid.jz == g.jz + assert mgrid.kp == g.kp + assert mgrid.nfp == g.nfp + assert mgrid.rmin == g.rmin + assert mgrid.zmin == g.zmin + assert mgrid.rmax == g.rmax + assert mgrid.zmax == g.zmax + assert mgrid.mgrid_mode == g.mgrid_mode + assert np.array_equal(mgrid.coil_group, g.coil_group) + assert np.array_equal(mgrid.coil_current, g.coil_current) + assert np.array_equal(mgrid.br, g.br) + assert np.array_equal(mgrid.bp, g.bp) + assert np.array_equal(mgrid.bz, g.bz) + + +if __name__ == "__main__": + pytest.main([__file__, "-s"]) diff --git a/tests/ntv-demo/python/test_vary_vrot_omegae.py b/tests/ntv-demo/python/test_vary_vrot_omegae.py deleted file mode 100644 index 65d9a79..0000000 --- a/tests/ntv-demo/python/test_vary_vrot_omegae.py +++ /dev/null @@ -1,174 +0,0 @@ -# %% Standard imports -import matplotlib.pyplot as plt -import os -import numpy as np - -# Custom imports -from neo2_mars import generate_vrot_for_mars -from neo2_mars import generate_omega_e_for_mars -from neo2_mars import generate_nu_for_mars - -# Import to test -from ntv_demo_parameter_study import read_profiles -from ntv_demo_parameter_study import calculate_omegae, read_omegae_and_vrot -from ntv_demo_parameter_study import get_omegae_linear_coefs - -base_dir = "/proj/plasma/DATA/DEMO/MARS/BASE_INPUT_PROFILES" -negative_vrot_case = "/proj/plasma/DATA/DEMO/MARS/PROFWE/negative_vrot_case" -positive_vrot_case = "/proj/plasma/DATA/DEMO/MARS/PROFWE/positive_vrot_case" - -def test_get_omegae_linear_coefs_visual_check(): - fontsize = 24 - profiles = read_profiles(base_dir) - vrot = profiles["vrot"][1:,1] - sqrtspol = profiles["vrot"][1:,0] - offset, slope = get_omegae_linear_coefs(negative_vrot_case, positive_vrot_case) - omegae_calculated = offset + slope * vrot - fig, ax_left = plt.subplots(1, 1, figsize=(10, 8)) - ax_left.plot(sqrtspol, offset, '--k', label=r'offset $\delta$') - ax_left.plot([],[], '--r', label=r'slope $\kappa$') - ax_left.set_ylabel(r'$\delta$ [rad/s]', color='k', fontsize=fontsize) - ax_left.tick_params(axis='y', labelcolor='k', color='k', labelsize=fontsize) - ax_left.tick_params(axis='x', labelsize=fontsize) - - ax_right = ax_left.twinx() - ax_right.plot(sqrtspol, slope, '--r') - ax_right.set_ylabel(r'$\kappa$ [1]', color='r', fontsize=fontsize) - ax_right.tick_params(axis='y', labelcolor='r', color='r', labelsize=fontsize) - ax_right.spines['right'].set_color('r') - - ax_left.set_xlabel(r'$\rho_\mathrm{pol}$ [1]', fontsize=fontsize) - ax_left.set_ylim([-4e3, 1e4]) - plt.show() - -def test_calculate_omegae_visual_check(): - profiles = read_profiles(base_dir) - vrot = profiles["vrot"][1:,1] - sqrtspol = profiles["vrot"][1:,0] - omegae = calculate_omegae(vrot/2) - half_vrot_case = "/proj/plasma/DATA/DEMO/MARS/PROFWE/half_positive_vrot_case" - omegae_negative_vrot, _ = read_omegae_and_vrot(negative_vrot_case) - omegae_positive_vrot, _ = read_omegae_and_vrot(positive_vrot_case) - omegae_half_vrot, _ = read_omegae_and_vrot(half_vrot_case) - fig, ax = plt.subplots(1, 1, figsize=(10, 8)) - ax.plot(sqrtspol, omegae_negative_vrot, '-b', label='negative vrot') - ax.plot(sqrtspol, omegae_positive_vrot, '-r', label='positive vrot') - ax.plot(sqrtspol, omegae, '--g', label='calculated') - ax.plot(sqrtspol, omegae_half_vrot, '.k', label='control') - ax.set_ylabel(r'$\omega_{e}$ [rad/s]') - ax.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax.set_ylim([-1.2e4, 1e4]) - ax.set_xlim([0.01, 1]) - plt.legend() - plt.show() - -def test_profit_runs_visual_check(): - study = "/proj/plasma/DATA/DEMO/teams/parameter_study_profiles_MARS/vary_vrot_omegae_and_coilwidth/" - base_run = "/proj/plasma/DATA/DEMO/MARS/BASE_INPUT_PROFILES" - runs_color = (0.8, 0.8, 0.8) - base_color = (0.6, 0.2, 0.2) - fig, ax = plt.subplots(4, 2, figsize=(10, 8)) - ax_Te = ax[0, 0] - ax_Ti = ax[0, 1] - ax_n = ax[1, 0] - ax_nue = ax[2, 0] - ax_nui = ax[2, 1] - ax_vrot = ax[3, 0] - ax_omegae = ax[3, 1] - number_of_ploted_runs = 0 - max_number_of_runs = 20 - for folder in os.listdir(study): - if number_of_ploted_runs >= max_number_of_runs: - break - if folder.startswith('run_'): - number_of_ploted_runs += 1 - profiles = read_profiles(os.path.join(study, folder)) - - ax_Te.plot(profiles["Te"][1:,0], profiles["Te"][1:,1]/1000, color=runs_color) - ax_Te.set_ylabel(r'$T_\mathrm{e}$ [keV]') - ax_Te.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_Te.set_xlim([0.01, 1]) - - ax_Ti.plot(profiles["Ti"][1:,0], profiles["Ti"][1:,1]/1000, color=runs_color) - ax_Ti.set_ylabel(r'$T_\mathrm{i}$ [keV]') - ax_Ti.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_Ti.set_xlim([0.01, 1]) - - ax_n.plot(profiles["n"][1:,0], profiles["n"][1:,1], color=runs_color) - ax_n.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_n.set_ylabel(r'$n$ [$\mathrm{m}^{-3}$]') - ax_n.set_xlim([0.01, 1]) - - ax_nue.plot(profiles["nue"][1:,0], profiles["nue"][1:,1], color=runs_color) - ax_nue.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_nue.set_ylabel(r'$\nu_{e}$ [s$^{-1}$]') - ax_nue.set_yscale('log') - ax_nue.set_xlim([0.01, 1]) - - ax_nui.plot(profiles["nui"][1:,0], profiles["nui"][1:,1], color=runs_color) - ax_nui.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_nui.set_ylabel(r'$\nu_{i}$ [s$^{-1}$]') - ax_nui.set_yscale('log') - ax_nui.set_xlim([0.01, 1]) - - ax_vrot.plot(profiles["vrot"][1:,0], profiles["vrot"][1:,1], color=runs_color) - ax_vrot.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_vrot.set_ylabel(r'$V^\varphi$ [rad/s]') - ax_vrot.set_ylim([-3e3, 1.6e4]) - ax_vrot.set_xlim([0.01, 1]) - - ax_omegae.plot(profiles["omegae"][1:,0], profiles["omegae"][1:,1], color=runs_color) - ax_omegae.set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax_omegae.set_ylabel(r'$\Omega_{e}$ [rad/s]') - ax_omegae.set_ylim([-3e3, 1.6e4]) - ax_omegae.set_xlim([0.01, 1]) - - profiles = read_profiles(base_dir) - ax_Te.plot(profiles["Te"][1:,0], profiles["Te"][1:,1]/1000, color=base_color, label='base run') - ax_Ti.plot(profiles["Ti"][1:,0], profiles["Ti"][1:,1]/1000, color=base_color, label='base run') - ax_n.plot(profiles["n"][1:,0], profiles["n"][1:,1], color=base_color, label='base run') - ax_nue.plot(profiles["nue"][1:,0], profiles["nue"][1:,1], color=base_color, label='base run') - ax_nui.plot(profiles["nui"][1:,0], profiles["nui"][1:,1], color=base_color, label='base run') - ax_vrot.plot(profiles["vrot"][1:,0], profiles["vrot"][1:,1], color=base_color, label='base run') - ax_omegae.plot(profiles["omegae"][1:,0], profiles["omegae"][1:,1], color=base_color, label='base run') - ax_Te.legend() - ax_Ti.legend() - ax_n.legend() - ax_nue.legend() - ax_nui.legend() - ax_vrot.legend() - ax_omegae.legend() - fig.suptitle('run profiles', fontsize=16) - plt.tight_layout() - plt.show() - -def test_read_profiles_visual_check(): - profiles = read_profiles(base_dir) - fig, ax = plt.subplots(3, 2, figsize=(10, 8 )) - ax[0, 0].plot(profiles["Te"][1:,0], profiles["Te"][1:,1]/1000) - ax[0, 0].set_ylabel(r'$T_\mathrm{e}$ [keV]') - ax[0, 0].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[0, 1].plot(profiles["Ti"][1:,0], profiles["Ti"][1:,1]/1000) - ax[0, 1].set_ylabel(r'$T_\mathrm{i}$ [keV]') - ax[0, 1].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[1, 0].plot(profiles["n"][1:,0], profiles["n"][1:,1]) - ax[1, 0].set_ylabel(r'$n$ [$\mathrm{m}^{-3}$]') - ax[1, 0].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[1, 1].plot(profiles["vrot"][1:,0], profiles["vrot"][1:,1]) - ax[1, 1].set_ylabel(r'$v_\mathrm{rot}$ [rad/s]') - ax[1, 1].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[2, 0].plot(profiles["nue"][1:,0], profiles["nue"][1:,1]) - ax[2, 0].set_ylabel(r'$\nu_{e}$ [s$^{-1}$]') - ax[2, 0].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[2, 0].set_yscale('log') - ax[2, 1].plot(profiles["nui"][1:,0], profiles["nui"][1:,1]) - ax[2, 1].set_ylabel(r'$\nu_{i}$ [s$^{-1}$]') - ax[2, 1].set_xlabel(r'$\rho_\mathrm{pol}$ [1]') - ax[2, 1].set_yscale('log') - plt.show() - -if __name__ == "__main__": - test_calculate_omegae_visual_check() - test_profit_runs_visual_check() - test_get_omegae_linear_coefs_visual_check() - test_read_profiles_visual_check() \ No newline at end of file diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..7617853 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + slow: marks tests as slow (deselect with '-m "not slow"')