From fe11dfa192f920e22774e98d4eeaeae8210749b8 Mon Sep 17 00:00:00 2001 From: prasad-sawantdesai Date: Fri, 20 Mar 2026 10:32:03 +0100 Subject: [PATCH 1/4] fix: correct 2D grid orientation for rectangular equilibrium without explicit r/z arrays --- idstools/compute/equilibrium.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/idstools/compute/equilibrium.py b/idstools/compute/equilibrium.py index d103121..4f43ccf 100644 --- a/idstools/compute/equilibrium.py +++ b/idstools/compute/equilibrium.py @@ -86,14 +86,7 @@ def get2d_cartesian_grid(self, time_slice: int, profiles2d_index: int = 0) -> Un ) r1d = profiles2d.grid.dim1 z1d = profiles2d.grid.dim2 - nr = len(r1d) - nz = len(z1d) - r2d = np.empty(shape=(nr, nz)) - z2d = np.empty(shape=(nr, nz)) - for iz in range(nz): - r2d[:, iz] = r1d - for ir in range(nr): - z2d[ir, :] = z1d + r2d, z2d = np.meshgrid(r1d, z1d) if np.all(psi2d == 0.0): logger.error( From 2e641863d92fa290cc50791e7d2c2b6cce089a00 Mon Sep 17 00:00:00 2001 From: prasad-sawantdesai Date: Fri, 20 Mar 2026 15:45:51 +0100 Subject: [PATCH 2/4] use dim1 and dim2 arrays if r an z are not provided --- idstools/compute/equilibrium.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/idstools/compute/equilibrium.py b/idstools/compute/equilibrium.py index 4f43ccf..1713d60 100644 --- a/idstools/compute/equilibrium.py +++ b/idstools/compute/equilibrium.py @@ -82,23 +82,16 @@ def get2d_cartesian_grid(self, time_slice: int, profiles2d_index: int = 0) -> Un if profiles2d.grid_type.index == 1 and np.size(r2d) == 0: logger.warning( - f"profiles_2d[{profiles2d_index}].r is not available and grid type is 1.. Calculating from grid" + f"profiles_2d[{profiles2d_index}].r not available (grid_type=1), using grid.dim1/dim2" ) - r1d = profiles2d.grid.dim1 - z1d = profiles2d.grid.dim2 - r2d, z2d = np.meshgrid(r1d, z1d) + r2d = profiles2d.grid.dim1 + z2d = profiles2d.grid.dim2 if np.all(psi2d == 0.0): logger.error( "All values of psi2d are 0. No contour levels were found within the data range, Can not plot contour" ) return None - if np.size(r2d) != np.size(z2d) or np.size(r2d) != np.size(psi2d): - logger.error( - f"r, z and psi have not the same dimension in \ - equilibrium.time_slice[{time_slice}].profiles_2d[{profiles2d_index}]" - ) - return None return {"r2d": r2d, "z2d": z2d, "psi2d": psi2d} From 884cc0048b3c7f55cb4d9d0603aa2eed1978dece Mon Sep 17 00:00:00 2001 From: prasad-sawantdesai Date: Fri, 20 Mar 2026 15:53:02 +0100 Subject: [PATCH 3/4] fixed formatting --- idstools/compute/equilibrium.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/idstools/compute/equilibrium.py b/idstools/compute/equilibrium.py index 1713d60..d7e6b02 100644 --- a/idstools/compute/equilibrium.py +++ b/idstools/compute/equilibrium.py @@ -81,9 +81,7 @@ def get2d_cartesian_grid(self, time_slice: int, profiles2d_index: int = 0) -> Un psi2d = profiles2d.psi if profiles2d.grid_type.index == 1 and np.size(r2d) == 0: - logger.warning( - f"profiles_2d[{profiles2d_index}].r not available (grid_type=1), using grid.dim1/dim2" - ) + logger.warning(f"profiles_2d[{profiles2d_index}].r not available (grid_type=1), using grid.dim1/dim2") r2d = profiles2d.grid.dim1 z2d = profiles2d.grid.dim2 From 695d9d4f8dfa885423f879265d1ada172a2ce0fe Mon Sep 17 00:00:00 2001 From: prasad-sawantdesai Date: Thu, 26 Mar 2026 09:51:29 +0100 Subject: [PATCH 4/4] fixed function to adhere with IMAS DD and fixed view function --- idstools/compute/equilibrium.py | 52 ++++++++++++++++----------------- idstools/view/equilibrium.py | 10 ++++++- 2 files changed, 35 insertions(+), 27 deletions(-) diff --git a/idstools/compute/equilibrium.py b/idstools/compute/equilibrium.py index d7e6b02..5f2a493 100644 --- a/idstools/compute/equilibrium.py +++ b/idstools/compute/equilibrium.py @@ -41,49 +41,49 @@ def get2d_cartesian_grid(self, time_slice: int, profiles2d_index: int = 0) -> Un This function returns a dictionary containing 2D Cartesian grid coordinates and psi values from an equilibrium IDS object. - Args: time_slice (int): The time slice index of the equilibrium data to be used for generating the - 2D Cartesian grid. Defaults to 0 - profiles2d_index (int): `profiles2d_index` is an integer parameter that represents the index of the - ``profile_2d`` to be used in the calculation. It is used to access the specific 2D profile from the - list of profiles in the `time_slice` object. Defaults to 0 + 2D Cartesian grid. Defaults to 0 + profiles2d_index (int): An integer parameter that represents the index of the + ``profiles_2d`` to be used in the calculation. It is used to access the specific 2D profile from the + list of profiles in the `time_slice` object. Defaults to 0 Returns: A dictionary containing the 2D Cartesian grid coordinates (r2d and z2d) and the corresponding psi - values (psi2d). + values (psi2d), or None if the data is unavailable or invalid. Example: .. code-block:: python - import imas - connection = imas.DBEntry("imas:mdsplus?user=public;pulse=134173;run=106;database=ITER;version=3","r") - idsObj = connection.get('equilibrium') - computeObj = EquilibriumCompute(idsObj) - result = computeObj.get2d_cartesian_grid(time_slice=0) + import imas + connection = imas.DBEntry("imas:mdsplus?user=public;pulse=134173;run=106;database=ITER;version=3","r") + idsObj = connection.get('equilibrium') + computeObj = EquilibriumCompute(idsObj) + result = computeObj.get2d_cartesian_grid(time_slice=0) - {'psi2d': array([[]]), - 'r2d': array([[]]), - 'z2d': array([[]])} + {'r2d': array([...]), 'z2d': array([...]), 'psi2d': array([...])} """ - profiles2d = None + profiles2d = r1d = z1d = None try: - profiles2d = self.ids.time_slice[time_slice].profiles_2d[ - profiles2d_index - ] # using https://docs.python.org/2/glossary.html#term-eafp style + profiles2d = self.ids.time_slice[time_slice].profiles_2d[profiles2d_index] except IndexError: logger.error(f"equilibrium.time_slice[{time_slice}].profiles_2d[{profiles2d_index}] is not available") return None profiles2d = self.ids.time_slice[time_slice].profiles_2d[profiles2d_index] - r2d = profiles2d.r - z2d = profiles2d.z - psi2d = profiles2d.psi - if profiles2d.grid_type.index == 1 and np.size(r2d) == 0: - logger.warning(f"profiles_2d[{profiles2d_index}].r not available (grid_type=1), using grid.dim1/dim2") - r2d = profiles2d.grid.dim1 - z2d = profiles2d.grid.dim2 + if profiles2d.grid_type.index == 1 and profiles2d.grid.dim1 is not None and profiles2d.grid.dim2 is not None: + logger.info( + f"Using equilibrium.time_slice[{time_slice}]" + f".profiles_2d[{profiles2d_index}].grid.dim1/dim2 for the 2D grid" + ) + r1d = profiles2d.grid.dim1 + z1d = profiles2d.grid.dim2 + else: + logger.error("Only rectangular cylindrical grid (grid_type=1) is supported for now") + return None + + psi2d = profiles2d.psi if np.all(psi2d == 0.0): logger.error( @@ -91,7 +91,7 @@ def get2d_cartesian_grid(self, time_slice: int, profiles2d_index: int = 0) -> Un ) return None - return {"r2d": r2d, "z2d": z2d, "psi2d": psi2d} + return {"r2d": r1d, "z2d": z1d, "psi2d": psi2d} def get_rho2d(self, time_slice: int, profiles2d_index: int = 0) -> Union[np.ndarray, None]: """ diff --git a/idstools/view/equilibrium.py b/idstools/view/equilibrium.py index 2416967..f1207fa 100644 --- a/idstools/view/equilibrium.py +++ b/idstools/view/equilibrium.py @@ -84,8 +84,16 @@ def view_magnetic_poloidal_flux( if cartestion_grid is not None: levels = 50 + # As per IMAS data dictionary psi is stored as [R, Z] with shape (N_R, N_Z). + # Check this reference : + # https://imas-data-dictionary.readthedocs.io/en/latest/generated/identifier/poloidal_plane_coordinates_identifier.html + + # For matplotlib contour as per the documentation: + # https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.contour.html + # matplotlib.contour(r, z, Zdata) expects Zdata as [rows=z, cols=r], so the shape must be (N_Z, N_R). + # Therefore we transpose psi before plotting. contour_lines_psi = ax.contour( - cartestion_grid["r2d"], cartestion_grid["z2d"], cartestion_grid["psi2d"], levels, cmap="summer" + cartestion_grid["r2d"], cartestion_grid["z2d"], cartestion_grid["psi2d"].T, levels, cmap="summer" ) # ax.clabel( # contour_lines_psi,