From 46fc92ebbe964d554e5ff5ce14f2ad45fefe0451 Mon Sep 17 00:00:00 2001 From: lorenzozanisi Date: Mon, 8 Sep 2025 11:18:39 +0100 Subject: [PATCH 1/4] more robust r_minor calculation --- src/pyrokinetics/equilibrium/imas.py | 30 +++++++++++++++++----------- src/pyrokinetics/kinetics/imas.py | 19 ++++++++++++------ 2 files changed, 31 insertions(+), 18 deletions(-) diff --git a/src/pyrokinetics/equilibrium/imas.py b/src/pyrokinetics/equilibrium/imas.py index 3f47140c6..a7e3ffc56 100644 --- a/src/pyrokinetics/equilibrium/imas.py +++ b/src/pyrokinetics/equilibrium/imas.py @@ -1,8 +1,10 @@ from pathlib import Path from typing import Optional + import h5py import numpy as np +import warnings from ..file_utils import FileReader from ..typing import PathLike @@ -134,6 +136,17 @@ def read_from_file( # The number of psi values is the same as the number of r values. The psi grid # uniformly increases from psi_axis to psi_lcfs psi_grid = data["time_slice[]&profiles_1d&psi"][time_index] * psi_units + if psi_n_lcfs==1: + if psi_grid[-1] < psi_grid[0]: + psi_n_lcfs_bk = psi_n_lcfs + if np.min(psi_RZ)!=psi_grid[-1]: + psi_n_lcfs = (np.min(psi_RZ)-psi_grid[0]) / (psi_grid[-1]-psi_grid[0]) + warnings.warn(f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead") + else: + if np.max(psi_RZ)!=psi_grid[-1]: + psi_n_lcfs = (np.max(psi_RZ)-psi_grid[0]) / (psi_grid[-1]-psi_grid[0]) + warnings.warn(f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead") + F = data["time_slice[]&profiles_1d&f"][time_index] * F_units FF_prime = ( @@ -150,8 +163,9 @@ def read_from_file( q = data["time_slice[]&profiles_1d&q"][time_index] * units.dimensionless # Adjust grids if psi_n_lcfs is not 1 + tol = 1.e-5 if psi_n_lcfs != 1.0: - if psi_n_lcfs > 1.0 or psi_n_lcfs < 0.0: + if psi_n_lcfs - 1.0 > tol or psi_n_lcfs < 0.0: raise ValueError( f"psi_n_lcfs={psi_n_lcfs} must be in the range [0,1]" ) @@ -166,7 +180,7 @@ def read_from_file( # Discard elements off the end of the grid, insert new psi_lcfs psi_grid_new = np.concatenate((psi_grid[:lcfs_idx], [psi_lcfs_new])) - + # Linearly interpolate each grid onto the new psi_grid # Need psi to be increasing for np.interp F = np.interp(psi_grid_new, psi_grid[::index], F[::index]) @@ -187,15 +201,8 @@ def read_from_file( r_minor[0] = 0.0 * units.meter Z_mid[0] = Z_axis - Rbdry = data["time_slice[]&boundary&outline&r"][time_index, :] - Zbdry = data["time_slice[]&boundary&outline&z"][time_index, :] - - R_major[-1] = 0.5 * (max(Rbdry) + min(Rbdry)) * units.meter - r_minor[-1] = 0.5 * (max(Rbdry) - min(Rbdry)) * units.meter - Z_mid[-1] = 0.5 * (max(Zbdry) + min(Zbdry)) * units.meter - idx_skipped = [] - for idx, psi in enumerate(psi_grid[1:-1], start=1): + for idx, psi in enumerate(psi_grid[1:], start=1): try: Rc, Zc = _flux_surface_contour(R, Z, psi_RZ, R_axis, Z_axis, psi) R_min, R_max = min(Rc), max(Rc) @@ -203,9 +210,8 @@ def read_from_file( R_major[idx] = 0.5 * (R_max + R_min) r_minor[idx] = 0.5 * (R_max - R_min) Z_mid[idx] = 0.5 * (Z_max + Z_min) - except ValueError: + except NotImplementedError: idx_skipped.append(idx) - for idx in idx_skipped: R_major[idx] = np.interp( abs(psi_grid[idx]), diff --git a/src/pyrokinetics/kinetics/imas.py b/src/pyrokinetics/kinetics/imas.py index d0b7521be..778c64efb 100644 --- a/src/pyrokinetics/kinetics/imas.py +++ b/src/pyrokinetics/kinetics/imas.py @@ -3,6 +3,9 @@ import h5py import numpy as np from periodictable import elements +from typing import Optional + + from ..constants import deuterium_mass, electron_mass from ..equilibrium import Equilibrium @@ -34,7 +37,7 @@ class KineticsReaderIMAS(FileReader, file_type="IMAS", reads=Kinetics): def read_from_file( self, filename: PathLike, - time_index: int = -1, + time_index: Optional[int] = None, time: float = None, eq: Equilibrium = None, ) -> Kinetics: @@ -191,11 +194,15 @@ def read_from_file( # In IMAS file n_ions matches n_ions_fast for i_ion in range(n_ions): - fast_ion_dens_data = ( - data["profiles_1d[]&ion[]&density_fast"][time_index, i_ion] - / units.meter**3 - ) - if np.all(fast_ion_dens_data == 0): + try: + fast_ion_dens_data = ( + data["profiles_1d[]&ion[]&density_fast"][time_index, i_ion] + / units.meter**3 + ) + if np.all(fast_ion_dens_data == 0): + continue + except KeyError: + # No fast ion data available continue fast_ion_dens_func = UnitSpline(psi_n, fast_ion_dens_data) From 07209f3ef232928caf9093d5d20872221c6d2e2d Mon Sep 17 00:00:00 2001 From: lorenzozanisi Date: Wed, 3 Dec 2025 15:22:44 +0000 Subject: [PATCH 2/4] moving equality between psi_n_lcfs_bk and psi_n_lcfs --- src/pyrokinetics/equilibrium/imas.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pyrokinetics/equilibrium/imas.py b/src/pyrokinetics/equilibrium/imas.py index a7e3ffc56..1b162196d 100644 --- a/src/pyrokinetics/equilibrium/imas.py +++ b/src/pyrokinetics/equilibrium/imas.py @@ -137,8 +137,8 @@ def read_from_file( # uniformly increases from psi_axis to psi_lcfs psi_grid = data["time_slice[]&profiles_1d&psi"][time_index] * psi_units if psi_n_lcfs==1: + psi_n_lcfs_bk = psi_n_lcfs if psi_grid[-1] < psi_grid[0]: - psi_n_lcfs_bk = psi_n_lcfs if np.min(psi_RZ)!=psi_grid[-1]: psi_n_lcfs = (np.min(psi_RZ)-psi_grid[0]) / (psi_grid[-1]-psi_grid[0]) warnings.warn(f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead") From 46988bfc5ef556e1f3bc1309415c183a9da9c561 Mon Sep 17 00:00:00 2001 From: lorenzozanisi Date: Wed, 3 Dec 2025 15:23:28 +0000 Subject: [PATCH 3/4] [skip ci] Apply black changes --- src/pyrokinetics/equilibrium/imas.py | 30 +++++++++++++++++----------- src/pyrokinetics/kinetics/imas.py | 4 +--- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/src/pyrokinetics/equilibrium/imas.py b/src/pyrokinetics/equilibrium/imas.py index 1b162196d..cf2654d77 100644 --- a/src/pyrokinetics/equilibrium/imas.py +++ b/src/pyrokinetics/equilibrium/imas.py @@ -1,10 +1,9 @@ +import warnings from pathlib import Path from typing import Optional - import h5py import numpy as np -import warnings from ..file_utils import FileReader from ..typing import PathLike @@ -136,17 +135,24 @@ def read_from_file( # The number of psi values is the same as the number of r values. The psi grid # uniformly increases from psi_axis to psi_lcfs psi_grid = data["time_slice[]&profiles_1d&psi"][time_index] * psi_units - if psi_n_lcfs==1: + if psi_n_lcfs == 1: psi_n_lcfs_bk = psi_n_lcfs if psi_grid[-1] < psi_grid[0]: - if np.min(psi_RZ)!=psi_grid[-1]: - psi_n_lcfs = (np.min(psi_RZ)-psi_grid[0]) / (psi_grid[-1]-psi_grid[0]) - warnings.warn(f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead") + if np.min(psi_RZ) != psi_grid[-1]: + psi_n_lcfs = (np.min(psi_RZ) - psi_grid[0]) / ( + psi_grid[-1] - psi_grid[0] + ) + warnings.warn( + f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead" + ) else: - if np.max(psi_RZ)!=psi_grid[-1]: - psi_n_lcfs = (np.max(psi_RZ)-psi_grid[0]) / (psi_grid[-1]-psi_grid[0]) - warnings.warn(f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead") - + if np.max(psi_RZ) != psi_grid[-1]: + psi_n_lcfs = (np.max(psi_RZ) - psi_grid[0]) / ( + psi_grid[-1] - psi_grid[0] + ) + warnings.warn( + f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead" + ) F = data["time_slice[]&profiles_1d&f"][time_index] * F_units FF_prime = ( @@ -163,7 +169,7 @@ def read_from_file( q = data["time_slice[]&profiles_1d&q"][time_index] * units.dimensionless # Adjust grids if psi_n_lcfs is not 1 - tol = 1.e-5 + tol = 1.0e-5 if psi_n_lcfs != 1.0: if psi_n_lcfs - 1.0 > tol or psi_n_lcfs < 0.0: raise ValueError( @@ -180,7 +186,7 @@ def read_from_file( # Discard elements off the end of the grid, insert new psi_lcfs psi_grid_new = np.concatenate((psi_grid[:lcfs_idx], [psi_lcfs_new])) - + # Linearly interpolate each grid onto the new psi_grid # Need psi to be increasing for np.interp F = np.interp(psi_grid_new, psi_grid[::index], F[::index]) diff --git a/src/pyrokinetics/kinetics/imas.py b/src/pyrokinetics/kinetics/imas.py index 778c64efb..010ae198a 100644 --- a/src/pyrokinetics/kinetics/imas.py +++ b/src/pyrokinetics/kinetics/imas.py @@ -1,11 +1,9 @@ from pathlib import Path +from typing import Optional import h5py import numpy as np from periodictable import elements -from typing import Optional - - from ..constants import deuterium_mass, electron_mass from ..equilibrium import Equilibrium From bebc64aa18f7ea6668f92d55ac73f1aed7a8ac92 Mon Sep 17 00:00:00 2001 From: Bhavin Patel <15210802+bpatel2107@users.noreply.github.com> Date: Fri, 12 Dec 2025 11:54:44 +0000 Subject: [PATCH 4/4] Only adjust psi_n_lcfs in IMAS is new value is close to 1 --- src/pyrokinetics/equilibrium/imas.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/pyrokinetics/equilibrium/imas.py b/src/pyrokinetics/equilibrium/imas.py index cf2654d77..a3dba602e 100644 --- a/src/pyrokinetics/equilibrium/imas.py +++ b/src/pyrokinetics/equilibrium/imas.py @@ -135,24 +135,27 @@ def read_from_file( # The number of psi values is the same as the number of r values. The psi grid # uniformly increases from psi_axis to psi_lcfs psi_grid = data["time_slice[]&profiles_1d&psi"][time_index] * psi_units + # Adjust grids if data outside LCFS is not good and need to reduce psi_n_lcfs + tol = 1.0e-3 if psi_n_lcfs == 1: - psi_n_lcfs_bk = psi_n_lcfs + psi_n_lcfs_potential = 1.0 if psi_grid[-1] < psi_grid[0]: if np.min(psi_RZ) != psi_grid[-1]: - psi_n_lcfs = (np.min(psi_RZ) - psi_grid[0]) / ( + psi_n_lcfs_potential = (np.min(psi_RZ) - psi_grid[0]) / ( psi_grid[-1] - psi_grid[0] ) - warnings.warn( - f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead" - ) else: if np.max(psi_RZ) != psi_grid[-1]: - psi_n_lcfs = (np.max(psi_RZ) - psi_grid[0]) / ( + psi_n_lcfs_potential = (np.max(psi_RZ) - psi_grid[0]) / ( psi_grid[-1] - psi_grid[0] ) - warnings.warn( - f"psi_n_lcfs was {psi_n_lcfs_bk}, {psi_n_lcfs} will be used instead" - ) + if psi_n_lcfs_potential != 1.0 and np.isclose( + psi_n_lcfs_potential, 1.0, atol=tol + ): + warnings.warn( + f"psi_n_lcfs was {psi_n_lcfs}, {psi_n_lcfs_potential} will be used instead" + ) + psi_n_lcfs = psi_n_lcfs_potential F = data["time_slice[]&profiles_1d&f"][time_index] * F_units FF_prime = ( @@ -168,8 +171,6 @@ def read_from_file( ) q = data["time_slice[]&profiles_1d&q"][time_index] * units.dimensionless - # Adjust grids if psi_n_lcfs is not 1 - tol = 1.0e-5 if psi_n_lcfs != 1.0: if psi_n_lcfs - 1.0 > tol or psi_n_lcfs < 0.0: raise ValueError(