Virtual casing: fix test, save B_external_normal_extended

hiddenSymmetries · Sep 22, 2023 · 8e65c49 · 8e65c49
1 parent 96910db
commit 8e65c49
Showing 1 changed file with 20 additions and 7 deletions.
diff --git a/src/simsopt/mhd/virtual_casing.py b/src/simsopt/mhd/virtual_casing.py
@@ -296,6 +296,7 @@ def save(self, filename="vcasing.nc"):
             f.createDimension('src_nphi', self.src_nphi)
             f.createDimension('trgt_ntheta', self.trgt_ntheta)
             f.createDimension('trgt_nphi', self.trgt_nphi)
+            f.createDimension('trgt_nphi_extended', self.trgt_nphi * 2 * self.nfp)
             f.createDimension('xyz', 3)
 
             src_ntheta = f.createVariable('src_ntheta', 'i', tuple())
@@ -368,6 +369,11 @@ def save(self, filename="vcasing.nc"):
             B_external_normal.description = 'Component of B_external normal to the surface'
             B_external_normal.units = 'Tesla'
 
+            B_external_normal_extended = f.createVariable('B_external_normal_extended', 'd', ('trgt_nphi_extended', 'trgt_ntheta'))
+            B_external_normal_extended[:, :] = self.B_external_normal_extended
+            B_external_normal_extended.description = 'Component of B_external normal to the surface, repeated to cover the full torus'
+            B_external_normal_extended.units = 'Tesla'
+
     @classmethod
     def load(cls, filename):
         """
@@ -410,13 +416,20 @@ def plot(self, ax=None, show=True):
         fig.colorbar(contours)
         fig.tight_layout()
 
-        fig, ax = plt.subplots()
-        contours = ax.contourf(np.linspace(0, 1, self.B_external_normal_extended.T.shape[1]), np.linspace(0, 1, self.B_external_normal_extended.T.shape[0]), self.B_external_normal_extended.T, 25*self.nfp*2)
-        ax.set_xlabel(r'$\phi$')
-        ax.set_ylabel(r'$\theta$')
-        ax.set_title('B_external_normal_extended [Tesla]')
-        fig.colorbar(contours)
-        fig.tight_layout()
+        fig1, ax1 = plt.subplots()
+        shape = self.B_external_normal_extended.T.shape
+        contours = ax1.contourf(
+            np.linspace(0, 1, shape[1]),
+            np.linspace(0, 1, shape[0]),
+            self.B_external_normal_extended.T,
+            25,
+        )
+        ax1.set_xlabel(r'$\phi$')
+        ax1.set_ylabel(r'$\theta$')
+        ax1.set_title('B_external_normal_extended [Tesla]')
+        fig1.colorbar(contours)
+        fig1.tight_layout()
+
         if show:
             plt.show()
         return ax