def _plot_debug(eq: pleque.Equilibrium, ax: plt.Axes = None, levels=None, colorbar=False): if ax is None: ax = plt.gca() rs = np.linspace(eq.R_min, eq.R_max, 400) zs = np.linspace(eq.Z_min, eq.Z_max, 600) try: if levels is None: levels = 60 cl = ax.contour(rs, zs, eq._spl_psi(rs, zs).T, levels) if colorbar: plt.contour(cl) except Exception: print("WARNING: Something wrong with psi spline.") try: ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], "k+-", label='first wall') except Exception: print("WARNING: No first wall?!") try: ax.plot(eq._lcfs[:, 0], eq._lcfs[:, 1], "C0", label='LCFS') except Exception: print("WARNING: LCFS in troubles?!") try: ax.contour(rs, zs, eq._spl_psi(rs, zs).T, [eq._psi_lcfs], colors="C1", linestyles="--") except Exception: print("WARNING: LCFS contour problem.") try: ax.plot(eq._o_points[:, 0], eq._o_points[:, 1], "C0o", label='o-points') except Exception: print("WARNING: O-points in trouble") try: ax.plot(*eq._mg_axis, "C1o", label='mg axis') except Exception: print("WARNING: mg. axis in trouble") try: ax.plot(eq._x_points[:, 0], eq._x_points[:, 1], "C2x", label='x-points') except Exception: print("WARNING: X-points in trouble") try: ax.plot(eq._x_point[0], eq._x_point[1], "rx", lw=2, label='x-point') except Exception: print("WARNING: THE X-point in trouble") try: ax.plot(eq._limiter_point[0], eq._limiter_point[1], "g+", lw=3, label='limiter point') except: print("WARNING: Limiter point is in trouble.") try: ax.plot(eq._strike_points[:, 0], eq._strike_points[:, 1], "C3+", lw=2, label='strike points') except: print("WARNING: Strike-points in trouble.") ax.legend() ax.set_aspect("equal")
def plot_equilibrium(eq: pleque.Equilibrium, ax: plt.Axes = None, colorbar=False, **kwargs): if ax is None: ax = plt.gca() if eq._first_wall is not None and len(eq._first_wall) > 2: ax.fill_between(eq._first_wall[:, 0], eq._first_wall[:, 1], color='lightgrey') ax.plot(eq._first_wall[:, 0], eq._first_wall[:, 1], color='k', lw=2, label='First wall') sep = eq.separatrix in_fw = eq.in_first_wall(sep) ax.plot(sep.R[in_fw], sep.Z[in_fw], color='C3', lw=2, alpha=0.5) # separatrix ax.plot(eq.lcfs.R, eq.lcfs.Z, color='C1', lw=2, ls='--') coords = eq.grid((400, 600), 'size') rs, zs = coords.mesh() psi = eq.psi(coords) mask_inlcfs = eq.in_lcfs(coords) mask_inlimiter = eq.in_first_wall(coords) mask_insol = mask_inlcfs ^ mask_inlimiter psi_in = np.ma.masked_array(psi, np.logical_not(mask_inlcfs)) psi_out = np.ma.masked_array(psi, mask_inlcfs) # ax.pcolormesh(coords.R, coords.Z, psi_in, shading='gouraud') contour_out = eq.coordinates(r=eq.lcfs.r_mid[0] + 2e-3 * np.arange(1, 11), theta=np.zeros(10), grid=False) cl = ax.contour(coords.R, coords.Z, psi_in, 20, **kwargs) plt.colorbar(cl, ax=ax) # todo: psi should be 1-d (!) resolve this ax.contour(coords.R, coords.Z, psi, np.sort(np.squeeze(contour_out.psi)), colors='C0') # contact = eq.strike_points # ax.plot(contact.R, contact.Z, "C3+") op = eq.magnetic_axis ax.plot(op.R, op.Z, "C0o") psi_lcfs = eq._psi_lcfs z0 = eq._mg_axis[1] # ax.pcolormesh(coords.R, coords.Z, mask_in) # ax.pcolormesh(rs, zs, mask_inlimiter) # ax.pcolormesh(coords.R, coords.Z, mask_insol) ax.set_xlabel('R [m]') ax.set_ylabel('Z [m]') if len(eq._first_wall) > 2: rlim = [np.min(eq.first_wall.R), np.max(eq.first_wall.R)] zlim = [np.min(eq.first_wall.Z), np.max(eq.first_wall.Z)] size = rlim[1] - rlim[0] rlim[0] -= size / 12 rlim[1] += size / 12 size = zlim[1] - zlim[0] zlim[0] -= size / 12 zlim[1] += size / 12 ax.set_xlim(*rlim) ax.set_ylim(*zlim) ax.set_aspect('equal') return ax