def read_fiesta_equilibrium(filepath, first_wall=None): """ Current versions of the equilibria are stored in `/compass/Shared/Common/COMPASS-UPGRADE/RP1 Design/Equilibria/v3.1` :param filepath: Path to fiesta g-file equilibria :param first_wall: Path to datafile with limiter line. (Fiesta doesn't store limiter contour into g-file). If `None` IBA limiter v 3.1 is taken. :return: Equilibrium: Instance of `Equilibrium` """ resource_package = 'pleque' with open(filepath, 'r') as f: data = read(f) ds = data_as_ds(data) # If there are some limiter data. Use them as and limiter. if 'r_lim' in ds and 'z_lim' in ds and ds.r_lim.size > 3 and first_wall is None: first_wall = np.stack((ds.r_lim.values, ds.z_lim.values)).T if first_wall is None: print('--- No limiter specified. The IBA v3.1 limiter will be used.') first_wall = 'resources/limiter_v3_1_iba_v2.dat' first_wall = pkg_resources.resource_filename(resource_package, first_wall) if isinstance(first_wall, str): first_wall = np.loadtxt(first_wall) eq = Equilibrium(ds, first_wall=first_wall) eq._Ip = ds.attrs['cpasma'] return eq
def test_cocos_consistency(geqdsk_file, cocos): print('COCOS: {}'.format(cocos)) equilibrium = read_geqdsk(geqdsk_file, cocos=cocos) with open(geqdsk_file, 'r') as f: eq_dict = read_geqdsk_as_dict(f) eq_xr = data_as_ds(eq_dict) eq_xr.psi.values = eq_xr.psi.values * (2 * np.pi) eq_xr.pprime.values /= (2 * np.pi) eq_xr.FFprime.values /= (2 * np.pi) fw = np.stack( (eq_xr['r_lim'].values, eq_xr['z_lim'].values)).T # first wall equilibrium2 = Equilibrium(eq_xr, fw, cocos=cocos + 10) sigma_Ip = np.sign(equilibrium.I_plasma) sigma_B0 = np.sign(equilibrium.F0) cocos_dict = equilibrium._cocosdic Rax = equilibrium.magnetic_axis.R[0] Zax = equilibrium.magnetic_axis.Z[0] assert sigma_B0 * sigma_Ip == np.sign( equilibrium.q(psi_n=0.5)) * cocos_dict['sigma_pol'] assert sigma_Ip == cocos_dict['sigma_Bp'] * equilibrium._psi_sign assert np.sign(equilibrium.F(R=Rax, Z=Zax)) == sigma_B0 assert np.sign(equilibrium.B_tor(R=Rax, Z=Zax)) == sigma_B0 assert np.sign(equilibrium.j_tor(r=0.1, theta=0)) == sigma_Ip assert np.sign(equilibrium.psi(psi_n=1) - equilibrium.psi( psi_n=0)) == sigma_Ip * cocos_dict['sigma_Bp'] assert np.sign(equilibrium.tor_flux(r=0.1, theta=0)) == sigma_B0 assert np.sign( equilibrium.pprime(psi_n=0.5)) == -sigma_Ip * cocos_dict['sigma_Bp'] assert np.sign(equilibrium.q( psi_n=0.5)) == sigma_Ip * sigma_B0 * cocos_dict['sigma_pol'] assert np.isclose( equilibrium.j_tor(R=equilibrium.magnetic_axis.R + 0.5, Z=equilibrium.magnetic_axis.Z), equilibrium2.j_tor(R=equilibrium2.magnetic_axis.R + 0.5, Z=equilibrium2.magnetic_axis.Z)) assert np.isclose(equilibrium.q(psi_n=0.5), equilibrium2.q(psi_n=0.5)) assert np.isclose(equilibrium.I_plasma, equilibrium2.I_plasma, atol=100)
def plot_overview(eq: Equilibrium): import matplotlib.pyplot as plt 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(figsize=(8, 4)) plt.subplot(131) plt.contour(r, z, psi, 20) # plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') plt.plot(eq.lcfs.R, eq.lcfs.Z, label='lcfs') if eq._first_wall is not None: plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') plt.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10) plt.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10) plt.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10) return_axis = plt.gca() plt.title(r'$\psi$') plt.gca().set_aspect('equal') plt.subplot(132) cs = plt.contour(r, z, eq.B_pol(R=r, Z=z), 20) plt.clabel(cs, inline=1) plt.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') plt.title(r'$B_\mathrm{pol}$') plt.gca().set_aspect('equal') plt.subplot(133) cs = plt.contour(r, z, eq.B_tor(R=r, Z=z), 20) plt.clabel(cs, inline=1) plt.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') plt.title(r'$B_\mathrm{tor}$') plt.gca().set_aspect('equal') # number of points n = 100 # module automaticaly identify the type of the input: midplane = eq.coordinates(r=np.linspace(-0.3, 0.3, n), theta=np.zeros(n)) fig, axs = plt.subplots(3, 1, sharex=True) ax = axs[0] # Profile of toroidal field: ax.plot(midplane.r, eq.B_tor(midplane)) # Profile of poloidal field: ax.plot(midplane.r, eq.B_pol(midplane)) return return_axis
def sal_jet(pulse, timex=47.0, time_unit="s"): """ Main loading routine, based on simple access layer, loads ppf data, calculates derivatives :param pulse: JET pulse number :param timex: time of slice :param time_unit: (str) "s" or "ms" :return: equilibrium """ if time_unit.lower() == "s": time_factor = 1. elif time_unit.lower() == "ms": time_factor = 1000. else: raise ValueError("Unknown `time_unit`.") data_path = '/pulse/{}/ppf/signal/jetppf/efit/{}:{}' # default sequence sequence = 0 # obtain psi data (reshape, transpose) and time axis packed_psi = sal.get(data_path.format(pulse, 'psi', sequence)) psi = packed_psi psi.data = packed_psi.data[:, :].reshape(len(packed_psi.dimensions[0]), 33, 33) psi.data = np.swapaxes(psi.data, 1, 2) time = packed_psi.dimensions[0].data # psi grid axis r = sal.get(data_path.format(pulse, 'psir', sequence)).data z = sal.get(data_path.format(pulse, 'psiz', sequence)).data # pressure profile pressure = sal.get(data_path.format(pulse, 'p', sequence)) psi_n = pressure.dimensions[1].data # f-profile f = sal.get(data_path.format(pulse, 'f', sequence)) # q-profile q = sal.get(data_path.format(pulse, 'q', sequence)) # calculate pprime and FFprime deltapsi = deltapsi_calc(pulse) pprime = pprime_calc(pressure, deltapsi, len(psi_n)) FFprime = FFprime_calc(f, deltapsi, len(psi_n)) #create dataset dst = xr.Dataset( { 'psi': (['time', 'R', 'Z'], psi.data), 'pressure': (['time', 'psi_n'], pressure.data), 'pprime': (['time', 'psi_n'], pprime), 'F': (['time', 'psi_n'], f.data), 'FFprime': (['time', 'psi_n'], FFprime), 'q': (['time', 'psi_n'], q.data), 'R': (['R'], r), 'Z': (['Z'], z), }, coords={ 'time': time, 'psi_n': psi_n, }) # select desired time ds = dst.sel(time=timex / time_factor, method='nearest') # try to load limiter from ppfs try: limiter_r = sal.get(data_path.format(pulse, 'rlim', sequence)).data.T limiter_z = sal.get(data_path.format(pulse, 'zlim', sequence)).data.T except NodeNotFound: limiter_r = sal.get(data_path.format(94508, 'rlim', sequence)).data.T limiter_z = sal.get(data_path.format(94508, 'zlim', sequence)).data.T print("Limiter points not present in #{}, loaded from #94508".format( pulse)) limiter = np.column_stack([limiter_r, limiter_z]) # create pleque equilibrium eq = Equilibrium(ds, limiter) return eq
def read(ods: omas.ODS, time=None, time_unit="s"): """ :param ods: OMAS object with description of `wall` and `equilibrium` :param time: (int, float or None) Time of required equilibria. The nearest time slice is taken. If `None` the first time slice is taken. :param time_unit: (str) "s" or "ms" :return: Instance of `Equilibrioum` from the `OMAS` object keeping the IMAS data structure. """ if time_unit.lower() == "s": time_factor = 1. elif time_unit.lower() == "ms": time_factor = 1000. else: raise ValueError("Unknown `time_unit`.") try: shot = ods["info"]["shot"] except: shot = ods["dataset_description"]["data_entry"]["pulse"] if "wall" not in ods: try: # todo: (Need the latest OMAS) from omas.omas_physics import add_wall omas.add_wall(omas) except ImportError: print("The newest OMAS is required to add wall to IDS.") except KeyError: pass # Reading wall: try: r_wall = ods["wall"]["description_2d"][0]["limiter"]["unit"][0][ "outline"]["r"] z_wall = ods["wall"]["description_2d"][0]["limiter"]["unit"][0][ "outline"]["z"] limiter = np.stack((r_wall, z_wall)).T except ValueError: limiter = None # Times ods_times = ods["equilibrium"]["time"] if time is None: time_idx = 0 else: time_idx = np.argmin(np.abs(np.array(ods_times) - time / time_factor)) # Plasma boundary (LCFS) rbnd = ods["equilibrium"]["time_slice"][time_idx]["boundary"]["outline"][ "r"] zbnd = ods["equilibrium"]["time_slice"][time_idx]["boundary"]["outline"][ "z"] psibnd = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][ "psi_boundary"] # Magnetic axis rmgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][ "magnetic_axis"]["r"] zmgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][ "magnetic_axis"]["z"] psimgax = ods["equilibrium"]["time_slice"][time_idx]["global_quantities"][ "psi_axis"] # 1D profiles psi_1d = ods["equilibrium"]["time_slice"][time_idx]["profiles_1d"]["psi"] F = ods['equilibrium.time_slice'][time_idx]['profiles_1d.f'] pressure = ods['equilibrium.time_slice'][time_idx]['profiles_1d.pressure'] FFp = ods['equilibrium.time_slice'][time_idx]['profiles_1d.f_df_dpsi'] pprime = ods['equilibrium.time_slice'][time_idx][ 'profiles_1d.dpressure_dpsi'] q = ods['equilibrium.time_slice'][time_idx]['profiles_1d.q'] psi_n = (psi_1d - psi_1d[0]) / (psibnd - psimgax) # 2D profiles: # todo: resolve this (!) grid_idx = 0 gridtype = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][0][ "grid_type"]["index"] if gridtype != 1: raise ValueError("Only rectangular is supported at the moment!") # if 1 not in gridtype: # raise ValueError("Only rectangular is supported at the moment!") # else: # grid_idx = gridtype.index(1) r_grid = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][ grid_idx]["grid"]["dim1"] z_grid = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][ grid_idx]["grid"]["dim2"] psi = ods["equilibrium"]["time_slice"][time_idx]["profiles_2d"][grid_idx][ "psi"] if psi.shape != (len(r_grid), len(z_grid)): psi = psi.T # x-array Dataset with time data: ds = xr.Dataset( { 'psi': (['R', 'Z'], psi), 'pressure': (['psi_n'], pressure), 'pprime': (['psi_n'], pprime), 'F': (['psi_n'], F), 'FFprime': (['psi_n'], FFp), 'q': (['psi_n'], q), }, coords={ 'time': ods_times[time_idx], 'R': r_grid, 'Z': z_grid, 'psi_n': psi_n, }, attrs={ 'shot': shot, 'time_unit': time_unit, }) eq = Equilibrium(ds, limiter) return eq
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 read_gfile(g_file: str, limiter: str = None): from pleque.io import _geqdsk import numpy as np import xarray as xr from pleque.core import Equilibrium from scipy.interpolate import UnivariateSpline with open(g_file, 'r') as f: eq_gfile = _geqdsk.read(f) psi = eq_gfile['psi'] r = eq_gfile['r'][:, 0] z = eq_gfile['z'][0, :] pressure = eq_gfile['pressure'] F = eq_gfile['F'] q = eq_gfile['q'] psi_n = np.linspace(0, 1, len(F)) eq_ds = xr.Dataset( { 'psi': (['R', 'Z'], psi), 'pressure': ('psi_n', pressure), 'F': ('psi_n', F) }, coords={ 'R': r, 'Z': z, 'psi_n': psi_n }) if limiter is not None: lim = np.loadtxt(limiter) print(lim.shape) else: lim = None eq = Equilibrium(eq_ds, first_wall=lim) eq._geqdsk = eq_gfile eq._q_spl = UnivariateSpline(psi_n, q, s=0, k=3) eq._dq_dpsin_spl = eq._q_spl.derivative() eq._q_anideriv_spl = eq._q_spl.antiderivative() def q(self, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=True, **coords): if R is not None and Z is not None: psi_n = self.psi_n(R=R, Z=Z, grid=grid) return self._q_spl(psi_n) def diff_q(self: eq, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=True, **coords): """ :param self: :param coordinates: :param R: :param Z: :param psi_n: :param coord_type: :param grid: :param coords: :return: Derivative of q with respect to psi. """ if R is not None and Z is not None: psi_n = self.psi_n(R=R, Z=Z, grid=grid) return self._dq_dpsin_spl(psi_n) * self._diff_psiN def tor_flux(self: eq, *coordinates, R=None, Z=None, psi_n=None, coord_type=None, grid=True, **coords): if R is not None and Z is not None: psi_n = self.psi_n(R=R, Z=Z, grid=grid) return eq._q_anideriv_spl(psi_n) * (1 / self._diff_psi_n) # eq.q = q # eq.diff_q = diff_q # eq.tor_flux = tor_flux Equilibrium.q = q Equilibrium.diff_q = diff_q Equilibrium.tor_flux = tor_flux return eq
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'$")
def plot_overview(eq: Equilibrium): import matplotlib.pyplot as plt 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(figsize=(8, 4)) plt.subplot(131) plt.contour(r, z, psi, 20) # plt.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') plt.plot(eq.lcfs.R, eq.lcfs.Z, label='lcfs') if eq._first_wall is not None: plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') plt.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10) plt.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10) plt.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10) return_axis = plt.gca() plt.title(r'$\psi$') plt.gca().set_aspect('equal') plt.subplot(132) cs = plt.contour(r, z, eq.B_pol(R=r, Z=z), 20) plt.clabel(cs, inline=1) plt.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') plt.title(r'$B_\mathrm{pol}$') plt.gca().set_aspect('equal') plt.subplot(133) cs = plt.contour(r, z, eq.B_tor(R=r, Z=z), 20) plt.clabel(cs, inline=1) plt.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') plt.title(r'$B_\mathrm{tor}$') plt.gca().set_aspect('equal') # number of points n = 100 # module automaticaly identify the type of the input: midplane = eq.coordinates(r=np.linspace(-0.3, 0.3, n), theta=np.zeros(n)) fig, axs = plt.subplots(3, 1, sharex=True) ax = axs[0] # Profile of toroidal field: ax.plot(midplane.r, eq.B_tor(midplane)) # Profile of poloidal field: ax.plot(midplane.r, eq.B_pol(midplane)) # # ---- # plt.figure(figsize=(8, 4)) # plt.subplot(131) # plt.pcolormesh(r, z, eq.B_R(R=r, Z=z).T) # if eq._first_wall is not None: # plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') # plt.title(r'$B_\mathrm{R}$') # plt.gca().set_aspect('equal') # # plt.subplot(132) # plt.pcolormesh(r, z, eq.B_Z(R=r, Z=z).T) # if eq._first_wall is not None: # plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') # plt.title(r'$B_\mathrm{Z}$') # plt.gca().set_aspect('equal') # # plt.subplot(133) # plt.pcolormesh(r, z, np.sqrt(eq.B_R(R=r, Z=z) ** 2 + eq.B_Z(R=r, Z=z) ** 2).T) # if eq._first_wall is not None: # plt.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') # plt.title(r'$B_\mathrm{pol}$') # plt.gca().set_aspect('equal') # # psi_n = np.linspace(0, 1, 100) # # plt.figure() # plt.subplot(211) # ax = plt.gca() # ax.plot(psi_n, eq.pressure(psi_n=psi_n), 'C1') # ax.set_xlabel(r'$\psi_\mathrm{N}$') # ax.set_ylabel(r'$p [Pa]$', color='C1') # # ax2 = ax.twinx() # ax2.plot(psi_n, eq.pprime(psi_n=psi_n), color='C2') # ax2.set_ylabel(r"$p'$", color='C2') # # plt.subplot(212) # ax = plt.gca() # ax.plot(psi_n, eq.F(psi_n=psi_n), 'C1') # ax.set_xlabel(r'$\psi_\mathrm{N}$') # ax.set_ylabel(r'$f$ ', color='C1') # # ax2 = ax.twinx() # ax2.plot(psi_n, eq.FFprime(psi_n=psi_n), 'C2') # ax2.set_ylabel(r"$ff'$ ", color='C2') return return_axis
def plot_psi_derivatives(eq: Equilibrium): import matplotlib.pyplot as plt 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) fig, axs = plt.subplots(1, 4, sharex=True, sharey=True, figsize=(8, 4)) ax = axs[0] ax.contour(r, z, psi.T, 20) ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') if eq._first_wall is not None: ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') ax.plot(eq._mg_axis[0], eq._mg_axis[1], 'o', color='b', markersize=10) ax.plot(eq._x_point[0], eq._x_point[1], 'x', color='r', markersize=10) ax.plot(eq._x_point2[0], eq._x_point2[1], 'x', color='r', markersize=10) psi_axis = ax plot_extremes(eq, psi_axis) def psi_xysq_func(r, z): return eq._spl_psi(r, z, dx=1, dy=0, grid=True) ** 2 \ + eq._spl_psi(r, z, dx=0, dy=1, grid=True) ** 2 psi_xysq = psi_xysq_func(r, z) psi_xyopt = (eq._spl_psi(r, z, dx=1, dy=1, grid=True))**2 plt.title(r'$\psi$') ax.set_aspect('equal') ax = axs[1] cl = ax.contour(r, z, psi_xysq.T, np.linspace(0, 0.1, 50)) # plt.colorbar(cl) ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') if eq._first_wall is not None: ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') ax.set_title(r'$\partial_x \psi^2 + \partial_y \psi^2$') ax.set_aspect('equal') ax = axs[2] cl = ax.contour(r, z, psi_xyopt.T, np.linspace(0, 1, 50)) # plt.colorbar(cl) ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') if eq._first_wall is not None: ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') ax.set_title(r'$\partial_{xy} \psi^2$') ax.set_aspect('equal') psi_xy = (eq._spl_psi(r, z, dx=1, dy=1, grid=True))**2 psi_xx = (eq._spl_psi(r, z, dx=2, dy=0, grid=True)) psi_yy = (eq._spl_psi(r, z, dx=0, dy=2, grid=True)) D = psi_xx * psi_yy - psi_xy ax = axs[3] cl = ax.contour(r, z, D.T, np.linspace(-50, 50, 50), cmap='PiYG') plt.colorbar(cl) ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], label='lcfs') if eq._first_wall is not None: ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], 'k') ax.set_title(r'$D$') ax.set_aspect('equal')