diff --git a/idstools/compute/equilibrium.py b/idstools/compute/equilibrium.py index d103121..5f2a493 100644 --- a/idstools/compute/equilibrium.py +++ b/idstools/compute/equilibrium.py @@ -41,73 +41,57 @@ 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 is not available and grid type is 1.. Calculating from grid" + 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 - 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 + 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( "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} + 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,