Added functionalities and plotting routines

megpy/equilibrium.py

@@ -984,10 +984,111 @@ def refine_equilibrium(self,nw=None,nh=None,nbbbs=None,interp_order=9,verbose=Fa
 
         return self
 
-    def plot_derived(self,):
-
-        return self
+    def plot_derived(self, figsize=(12, 8), save_path=None):
+        """Plot derived quantities from the equilibrium
+
+        Args:
+            figsize (tuple): Figure size (width, height)
+            save_path (str, optional): Path to save the figure
+
+        Returns:
+            matplotlib.figure.Figure: The created figure
+        """
+        if not self.derived:
+            print("No derived quantities available. Run add_derived() first.")
+            return None
+
+        fig, axes = plt.subplots(2, 2, figsize=figsize)
+        fig.suptitle('Equilibrium Derived Quantities', fontsize=16)
+
+        # Plot safety factor
+        axes[0, 0].plot(self.derived['rho_pol'], self.raw['qpsi'], 'b-', linewidth=2)
+        axes[0, 0].set_xlabel('ρ_pol')
+        axes[0, 0].set_ylabel('q (Safety Factor)')
+        axes[0, 0].set_title('Safety Factor Profile')
+        axes[0, 0].grid(True, alpha=0.3)
+
+        # Plot pressure
+        axes[0, 1].plot(self.derived['rho_pol'], self.raw['pres'], 'g-', linewidth=2)
+        axes[0, 1].set_xlabel('ρ_pol')
+        axes[0, 1].set_ylabel('Pressure [Pa]')
+        axes[0, 1].set_title('Pressure Profile')
+        axes[0, 1].grid(True, alpha=0.3)
+
+        # Plot toroidal field function
+        axes[1, 0].plot(self.derived['rho_pol'], self.raw['fpol'], 'm-', linewidth=2)
+        axes[1, 0].set_xlabel('ρ_pol')
+        axes[1, 0].set_ylabel('F = RB_tor [T⋅m]')
+        axes[1, 0].set_title('Toroidal Field Function')
+        axes[1, 0].grid(True, alpha=0.3)
+
+        # Plot current density components
+        if 'j_tor' in self.derived:
+            axes[1, 1].plot(self.derived['rho_pol'], self.derived['j_tor'], 'r-', linewidth=2)
+            axes[1, 1].set_xlabel('ρ_pol')
+            axes[1, 1].set_ylabel('J_tor [A/m²]')
+            axes[1, 1].set_title('Toroidal Current Density')
+            axes[1, 1].grid(True, alpha=0.3)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+
+        return fig
 
-    def plot_fluxsurfaces(self,):
-
-        return self
+    def plot_fluxsurfaces(self, figsize=(10, 8), n_contours=20, save_path=None):
+        """Plot flux surfaces and equilibrium geometry
+
+        Args:
+            figsize (tuple): Figure size (width, height)
+            n_contours (int): Number of flux surface contours to plot
+            save_path (str, optional): Path to save the figure
+
+        Returns:
+            matplotlib.figure.Figure: The created figure
+        """
+        if not self.derived:
+            print("No derived quantities available. Run add_derived() first.")
+            return None
+
+        fig, ax = plt.subplots(figsize=figsize)
+
+        # Create meshgrid for contour plotting
+        R_mesh, Z_mesh = np.meshgrid(self.derived['R'], self.derived['Z'])
+
+        # Plot psi contours
+        levels = np.linspace(self.raw['simag'], self.raw['sibry'], n_contours)
+        contours = ax.contour(R_mesh, Z_mesh, self.derived['psirz'], levels=levels, colors='blue', alpha=0.7)
+        contourf = ax.contourf(R_mesh, Z_mesh, self.derived['psirz'], levels=levels, alpha=0.3, cmap='viridis')
+
+        # Add colorbar
+        cbar = plt.colorbar(contourf, ax=ax)
+        cbar.set_label('ψ [Wb/rad]')
+
+        # Plot boundary if available
+        if 'rbbbs' in self.raw and 'zbbbs' in self.raw:
+            ax.plot(self.raw['rbbbs'], self.raw['zbbbs'], 'r-', linewidth=2, label='LCFS')
+
+        # Plot limiter if available
+        if 'rlim' in self.raw and 'zlim' in self.raw:
+            ax.plot(self.raw['rlim'], self.raw['zlim'], 'k-', linewidth=2, label='Limiter')
+
+        # Mark magnetic axis
+        ax.plot(self.raw['rmaxis'], self.raw['zmaxis'], 'ro', markersize=8, label='Magnetic Axis')
+
+        # Mark X-point if available
+        if 'R_x' in self.derived and 'Z_x' in self.derived:
+            ax.plot(self.derived['R_x'], self.derived['Z_x'], 'rx', markersize=10, label='X-point')
+
+        ax.set_xlabel('R [m]')
+        ax.set_ylabel('Z [m]')
+        ax.set_title('Flux Surfaces (ψ contours)')
+        ax.legend()
+        ax.grid(True, alpha=0.3)
+        ax.set_aspect('equal')
+
+        if save_path:
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+
+        return fig

megpy/poloidal_coordinates.py

@@ -0,0 +1,111 @@
+# Here we reproduce the poloidal coordinates from the paper:
+# "Ideal MHD Stability Calculations in Axisymmetric Toroidal Coordinate Systems"
+# by Grimm, Dewar, and Manickam (1981). [ https://www.sciencedirect.com/science/article/pii/002199918390116X?ref=cra_js_challenge&fr=RR-1 ]
+
+import numpy as np
+import scipy as sp
+import matplotlib.pyplot as plt
+
+
+E = 0.5
+R0 = 1.0
+i=2
+j=0
+
+# calculate the poloidal flux function psi
+def psi(R,Z,R0=R0,E=E):
+    """
+    Calculate the poloidal flux function psi in a cylindrical coordinate system.
+    """
+    term1 = (R**2 - R0**2)**2
+    term2 = E**2 * R**2 * Z**2
+    return term1 + term2
+
+
+# calculate gradient of psi
+def grad_psi(R, Z, R0=R0, E=E):
+    """
+    Calculate the gradient of the poloidal flux function psi.
+    """
+    dpsi_dR = 2 * (R**2 - R0**2) * 2 * R + 2 * E**2 * R * Z**2
+    dpsi_dZ = 2 * E**2 * R**2 * Z
+    return np.array([dpsi_dR, dpsi_dZ])
+
+
+# get contour list (R,Z) for a given psi value
+def get_contour(psi_value, R0=R0, E=E):
+    """
+    Get the contour (R,Z) for a given psi value.
+    """
+    R = np.linspace(0, 2*R0, 100)
+    Z = np.linspace(-2*R0, 2*R0, 100)
+    R, Z = np.meshgrid(R, Z)
+
+    # Calculate psi for the grid
+    psi_grid = psi(R, Z, R0, E)
+
+    # Find the contour where psi equals the specified value
+    R = np.linspace(0, 2*R0, 2000)
+    Z = np.linspace(-2*R0, 2*R0, 2000)
+    R, Z = np.meshgrid(R, Z)
+    psi_grid = psi(R, Z, R0, E)
+    contour = plt.contour(R, Z, psi_grid, levels=[psi_value])
+    contour_data = contour.get_paths()[0].vertices
+    plt.close()  # Close the plot to avoid displaying it
+    return contour_data
+
+
+def compute_poloidal_theta(psi_value=0.5, i=i, j=j, R0=R0, E=E):
+    """
+    Compute the poloidal angle theta along the contour of constant psi.
+    Returns theta, R, Z arrays.
+    """
+    c = get_contour(psi_value, R0=R0, E=E)
+    R = c[:, 0]
+    Z = c[:, 1]
+
+    grad_psi_norm = np.linalg.norm(grad_psi(R, Z, R0=R0, E=E), axis=0)
+
+    Z_left = np.roll(Z, 1)
+    Z_right = np.roll(Z, -1)
+    R_left = np.roll(R, 1)
+    R_right = np.roll(R, -1)
+    ds = np.sqrt((R_left - R_right)**2 + (Z_left - Z_right)**2)
+
+    integrand = grad_psi_norm**(j-1)/R**(i-1)
+    alpha = 2 * np.pi / np.sum(integrand * ds)
+
+    J = R**i/(alpha * grad_psi_norm**j)
+    dtheta_ds = R/(J * grad_psi_norm)
+    theta = np.cumsum(dtheta_ds * ds)
+    return theta, R, Z
+
+# plot mesh of poloidal coordinates for a number of psi values
+# (i.e. plot R,Z, also showing constant theta contours)
+psis = np.linspace(0.1, 0.5, 10)
+plt.figure(figsize=(10, 8))
+Rs = []
+Zs = []
+thetas = []
+for psi_value in psis:
+    theta, R, Z = compute_poloidal_theta(psi_value=psi_value, i=i, j=j, R0=R0, E=E)
+    Rs.append(R)
+    Zs.append(Z)
+    thetas.append(theta)
+
+
+# plot R(theta) and Z(theta) for each psi in a 2x1 subplot
+fig, axs = plt.subplots(2, 1, figsize=(10, 8))
+for i, psi_value in enumerate(psis):
+    axs[0].plot(thetas[i], Rs[i], label=f'psi={psi_value:.2f}')
+    axs[1].plot(thetas[i], Zs[i], label=f'psi={psi_value:.2f}')
+axs[0].set_title('R(theta) for different psi values')
+axs[0].set_xlabel('Theta')
+axs[0].set_ylabel('R')
+axs[1].set_title('Z(theta) for different psi values')
+axs[1].set_xlabel('Theta')
+axs[1].set_ylabel('Z')
+axs[0].legend()
+axs[1].legend()
+plt.tight_layout()
+plt.show()