def write(equilibrium: Equilibrium, grid_1d=None, grid_2d=None, gridtype=1, ods=None, cocosio=3): """ Function saving contents of equilibrium into the omas data structure. :param equilibrium: Equilibrium object :param grid_1d: Coordinate object with 1D grid (linearly spaced array over psi_n). It is used to save 1D equilibrium characteristics into the omas data object. If is None, linearly spaced vector over psi_n inrange [0,1] with 200 points is used. :param grid_2d: Coordinate object with 2D grid (can be the whole reconstruction space). It is used to save 2D equilibrium profiles. If None, default Equilibrium.grid() is used to generate the parameter. :param gridtype: Grid type specification for omas data structure. 1...rectangular. (Any other grid type is not supported at them moment!) :param ods: ods object to save the equilibrium into. If None, new ods object is created. :return: """ # todo: Well for now it is just what ASCOT eats, we will extend this when we pile up more usage.... if ods is None: ods = omas.ODS(cocosio=cocosio) if grid_1d is None: grid_1d = equilibrium.coordinates(psi_n=np.linspace(0, 1, 200)) if grid_2d is None: grid_2d = equilibrium.grid(resolution=(1e-3, 1e-3), dim="step") shot_time = equilibrium.time if equilibrium.time_unit == "ms": shot_time *= 1e3 elif equilibrium.time_unit != "s": print( "WARNING: unknown time unit ({}) is used. Seconds will be used insted for saving." .format(equilibrium.time_unit)) # ods["info"]["shot"] = equilibrium.shot ods["dataset_description"]["data_entry"]["pulse"] = equilibrium.shot # fill the wall part ods["wall"]["ids_properties"]["homogeneous_time"] = 1 # to MT: is this change correct? ods["wall"]["time"] = np.array(shot_time, ndmin=1) ods["wall"]["description_2d"][0]["limiter"]["unit"][0]["outline"][ "r"] = equilibrium.first_wall.R ods["wall"]["description_2d"][0]["limiter"]["unit"][0]["outline"][ "z"] = equilibrium.first_wall.Z ############################ # The Equilibrium Part ############################ # time slices ods["equilibrium"]["ids_properties"]["homogeneous_time"] = 1 # to MT: is this change correct? ods["equilibrium"]["time"] = np.array(shot_time, ndmin=1) # Vacuum # todo: add vacuum Btor, not in equilibrium # As R0 use magnetic axis. Isn't better other definition of R0? # R0 = np.array((np.max(equilibrium.first_wall.R) - np.min(equilibrium.first_wall.R))/2, ndim=1) R0 = equilibrium.magnetic_axis.R[0] F0 = equilibrium.BvacR B0 = F0 / R0 * np.ones_like(ods["equilibrium"]["time"]) ods["equilibrium"]["vacuum_toroidal_field"][ "b0"] = B0 # vacuum B tor at Rmaj ods["equilibrium"]["vacuum_toroidal_field"][ "r0"] = R0 # vacuum B tor at Rmaj # time slice time ods["equilibrium"]["time_slice"][0]["time"] = shot_time # plasma boundary (lcfs) ods["equilibrium"]["time_slice"][0]["boundary"]["outline"][ "r"] = equilibrium.lcfs.R ods["equilibrium"]["time_slice"][0]["boundary"]["outline"][ "z"] = equilibrium.lcfs.Z ods["equilibrium"]["time_slice"][0]["global_quantities"][ "psi_boundary"] = equilibrium._psi_lcfs # Magnetic axis ods["equilibrium"]["time_slice"][0]["global_quantities"]["magnetic_axis"][ "r"] = equilibrium.magnetic_axis.R[0] ods["equilibrium"]["time_slice"][0]["global_quantities"]["magnetic_axis"][ "z"] = equilibrium.magnetic_axis.Z[0] ods["equilibrium"]["time_slice"][0]["global_quantities"][ "psi_axis"] = equilibrium._psi_axis # define the 1 and 2d grids # 1d profiles ods["equilibrium"]["time_slice"][0]["profiles_1d"][ "psi"] = equilibrium.psi(grid_1d) ods["equilibrium"]["time_slice"][0]["profiles_1d"][ "rho_tor"] = equilibrium.tor_flux(grid_1d) ods['equilibrium.time_slice'][0]['profiles_1d.f'] = equilibrium.F(grid_1d) ods['equilibrium.time_slice'][0][ 'profiles_1d.pressure'] = equilibrium.pressure(grid_1d) ods['equilibrium.time_slice'][0][ 'profiles_1d.f_df_dpsi'] = equilibrium.FFprime(grid_1d) ods['equilibrium.time_slice'][0][ 'profiles_1d.dpressure_dpsi'] = equilibrium.pprime(grid_1d) ods['equilibrium.time_slice'][0]['profiles_1d.q'] = equilibrium.q(grid_1d) # get surface volumes and areas surface_volume = np.zeros_like(grid_1d.psi) surface_area = np.zeros_like(grid_1d.psi) for i in range(grid_1d.psi.shape[0]): psin_tmp = grid_1d.psi_n[i] surface = 0 if not psin_tmp == 1: coord_tmp = equilibrium.coordinates(psi_n=psin_tmp) surface = equilibrium._flux_surface(coord_tmp) elif psin_tmp == 1: surface = [equilibrium._as_fluxsurface(equilibrium.lcfs)] # todo: really 0 if open? # todo: fix this 'surface[0].closed' ... maybe calculate it from boundary if len(surface) > 0 and surface[0].closed: surface_volume[i] = surface[0].volume surface_area[i] = surface[0].area else: surface_volume[i] = 0 surface_area[i] = 0 ods["equilibrium"]["time_slice"][0]["profiles_1d"][ "volume"] = surface_volume ods["equilibrium"]["time_slice"][0]["profiles_1d"]["area"] = surface_area # 2D profiles ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid_type"][ "index"] = gridtype ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid"][ "dim1"] = grid_2d.R ods["equilibrium"]["time_slice"][0]["profiles_2d"][0]["grid"][ "dim2"] = grid_2d.Z ods["equilibrium"]["time_slice"][0]["profiles_2d"][0][ "psi"] = equilibrium.psi(grid_2d).T ods["equilibrium"]["time_slice"][0]["profiles_2d"][0][ "b_field_tor"] = equilibrium.B_tor(grid_2d).T ods['equilibrium.time_slice'][0][ 'global_quantities.ip'] = equilibrium.I_plasma return ods
def show_qprofiles(g_file: str, eq: Equilibrium): # from tokamak.formats import geqdsk from pleque.io._geqdsk import read import matplotlib.pyplot as plt with open(g_file, 'r') as f: eq_gfile = read(f) q = eq_gfile['q'] psi_n = np.linspace(0, 1, len(q)) print(eq_gfile.keys()) psin_axis = np.linspace(0, 1, 100) r = np.linspace(eq.R_min, eq.R_max, 100) z = np.linspace(eq.Z_min, eq.Z_max, 120) psi = eq.psi(R=r, Z=z) plt.figure() plt.subplot(121) ax = plt.gca() cs = ax.contour(r, z, psi, 30) ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') if eq._first_wall is not None: plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') ax.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b') ax.plot(eq._x_point[0], eq._x_point[1], 'x', color='r') ax.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r') ax.set_xlabel('R [m]') ax.set_ylabel('Z [m]') ax.set_aspect('equal') ax.set_title(r'$\psi$') plt.colorbar(cs, ax=ax) psi_mod = eq.psi(psi_n=psin_axis) tor_flux = eq.tor_flux(psi_n=psin_axis) q_as_grad = np.gradient(tor_flux, psi_mod) plt.subplot(322) ax = plt.gca() ax.plot(psi_n, np.abs(q), 'x', label='g-file (abs)') ax.plot(psin_axis, q_as_grad, '-', label=r'$\mathrm{d} \Phi/\mathrm{d} \psi$') ax.plot(psin_axis, q_as_grad, '--', label='Pleque') ax.legend() ax.set_xlabel(r'$\psi_\mathrm{N}$') ax.set_ylabel(r'$q$') plt.subplot(324) ax = plt.gca() ax.plot(tor_flux, psi_mod, label='Toroidal flux') ax.set_ylabel(r'$\psi$') ax.set_xlabel(r'$\Phi$') plt.subplot(326) ax = plt.gca() ax.plot(psi_n, eq.pprime(psi_n=psi_n) / 1e3) ax.set_xlabel(r'$\psi_\mathrm{N}$') ax.set_ylabel(r"$p' (\times 10^3)$") ax2 = ax.twinx() ax2.plot(psi_n, eq.FFprime(psi_n=psi_n), 'C1') ax2.set_ylabel(r"$ff'$")