cleaner implementation of theta derivative

easier handling of 3-D/2-D helper arrays, like th2D, rr2D, th3D, rr3D

Author: tgastine

python/magic/checkpoint.py

@@ -745,14 +745,9 @@ def graph2rst(self, gr, filename='checkpoint_ave.from_chk'):
         # Calculate the toroidal potential using wr
         self.ztor = np.zeros_like(self.wpol)
 
-        th3D = np.zeros_like(gr.vr)
-        rr3D = np.zeros_like(th3D)
-        for i in range(self.n_theta_max):
-            th3D[:, i, :] = gr.colatitude[i]
-        for i in range(self.n_r_max):
-            rr3D[:, :, i] = self.radius[i]
-        s3D = rr3D*np.sin(th3D)
-        omr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.vphi, order=4) -
+        th3D = gr.colatitude[None, :, None]
+        s3D = self.radius[None, None, :]*np.sin(gr.colatitude[None, :, None])
+        omr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.vphi, colat=gr.colatitude, order=4) -
                       phideravg(gr.vtheta, minc=self.minc))
 
         for i in range(self.n_r_max):
@@ -790,7 +785,7 @@ def graph2rst(self, gr, filename='checkpoint_ave.from_chk'):
                                    self.radius[i]**2
 
             self.btor = np.zeros_like(self.ztor)
-            jr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.Bphi, order=4) -
+            jr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.Bphi, colat=gr.colatitude, order=4) -
                          phideravg(gr.Btheta, minc=self.minc))
 
             for i in range(self.n_r_max):
@@ -819,16 +814,12 @@ def graph2rst(self, gr, filename='checkpoint_ave.from_chk'):
             # Calculate the toroidal potential using jr
             self.btor_ic = np.zeros_like(self.bpol_ic)
 
-            th3D = np.zeros_like(gr.Br_ic)
-            rr3D = np.zeros_like(th3D)
-            for i in range(self.n_theta_max):
-                th3D[:, i, :] = gr.colatitude[i]
-            for i in range(self.n_r_ic_max-1):
-                rr3D[:, :, i] = self.radius_ic[i]
-            rr3D[:, :, -1] = 1e-4
-            s3D = rr3D*np.sin(th3D)
+            th3D = gr.colatitude[None, :, None]
+            s3D = self.radius_ic[None, None, :]*np.sin(gr.colatitude[None, :, None])
+            s3D[:, :, -1] = 1e-4
             jr_ic = np.zeros_like(th3D)
-            jr_ic = 1./s3D*(thetaderavg(np.sin(th3D)*gr.Bphi_ic, order=4) -
+            jr_ic = 1./s3D*(thetaderavg(np.sin(th3D)*gr.Bphi_ic,
+                                        colat=gr.colatitude, order=4) -
                             phideravg(gr.Btheta_ic, minc=self.minc))
 
             for i in range(self.n_r_ic_max):

python/magic/graph2vtk.py

@@ -367,14 +367,10 @@ def __init__(self, gr, filename='out', scals=['vr', 'emag', 'tfluct'],
                                                         gr.Btheta[..., ::-1]**2+\
                                                         gr.Bphi[..., ::-1]**2)
             elif index == 3 or index == 7: # z-vorticity
-                th3D = np.zeros_like(gr.vphi)
-                rr3D = np.zeros_like(th3D)
-                for i in range(gr.ntheta):
-                    th3D[:, i, :] = gr.colatitude[i]
-                for i in range(gr.nr):
-                    rr3D[:, :, i] = gr.radius[i]
+                th3D = gr.colatitude[None, :, None]
+                rr3D = gr.radius[None, None, :]
                 s3D = rr3D * np.sin(th3D)
-                dtheta = thetaderavg(gr.vphi*s3D)
+                dtheta = thetaderavg(gr.vphi*s3D, colat=gr.colatitude)
                 dr = rderavg(gr.vphi*s3D, gr.radius)
                 vs = gr.vr * np.sin(th3D) + gr.vtheta * np.cos(th3D) # 'vs'
                 ds = np.sin(th3D)*dr + np.cos(th3D)/rr3D*dtheta
@@ -408,15 +404,11 @@ def __init__(self, gr, filename='out', scals=['vr', 'emag', 'tfluct'],
                 self.scals[k, ...] = self.scals[k, ...]-\
                                      self.scals[k, ...].mean(axis=0)
             elif index == 8: # r-vorticity
-                th3D = np.zeros_like(gr.vphi)
-                rr3D = np.zeros_like(th3D)
-                for i in range(gr.ntheta):
-                    th3D[:, i, :] = gr.colatitude[i]
-                for i in range(gr.nr):
-                    rr3D[:, :, i] = gr.radius[i]
+                th3D = gr.radius[None, None, :]
+                rr3D = gr.colatitude[None, :, None]
                 s3D = rr3D * np.sin(th3D)
-                vortr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.vphi)-\
-                        phideravg(gr.vtheta, gr.minc))
+                vortr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.vphi, colat=gr.colatitude)-\
+                                phideravg(gr.vtheta, gr.minc))
                 if deminc:
                     self.scals[k, :, :, 0:gr.nr] = symmetrize(vortr[..., ::-1],
                                                               gr.minc)

python/magic/libmagic.py

@@ -874,7 +874,7 @@ def rderavg(data, rad, exclude=False):
 
     return der
 
-def thetaderavg(data, order=4):
+def thetaderavg(data, colat=None, order=4):
     """
     Theta-derivative of an input array (finite differences)
 
@@ -883,42 +883,53 @@ def thetaderavg(data, order=4):
 
     :param data: input array
     :type data: numpy.ndarray
+    :param colat: colatitudes (when not specified a regular grid is assumed)
+    :type colat: numpy.ndarray
     :param order: order of the finite-difference scheme (possible values are 2 or 4)
     :type order: int
     :returns: the theta-derivative of the input array
     :rtype: numpy.ndarray
     """
+    if colat is None:
+        if len(data.shape) == 3:
+            ntheta = data.shape[1]
+        elif len(data.shape) == 2:
+            ntheta = data.shape[0]
+        th = np.linspace(0., np.pi, ntheta)
+    else:
+        th = colat
+
     if len(data.shape) == 3: # 3-D
-        ntheta = data.shape[1]
-        dtheta = np.pi/(ntheta-1.)
         if order == 2:
-            der = (np.roll(data, -1,  axis=1)-np.roll(data, 1, axis=1))/(2.*dtheta)
-            der[:, 0, :] = (data[:, 1, :]-data[:, 0, :])/dtheta
-            der[:, -1, :] = (data[:, -1, :]-data[:, -2, :])/dtheta
+            der = (np.roll(data,-1,axis=1)-np.roll(data,1,axis=1)) / \
+                  (np.roll(th[None, :, None],-1)-np.roll(th[None, :, None],1))
+            der[:, 0, :] = (data[:, 1, :]-data[:, 0, :])/(th[1]-th[0])
+            der[:, -1, :] = (data[:, -1, :]-data[:, -2, :])/(th[-1]-th[-2])
         elif order == 4:
-            der = (   -np.roll(data,-2,axis=1) \
-                   +8.*np.roll(data,-1,axis=1) \
-                   -8.*np.roll(data, 1,axis=1)  \
-                      +np.roll(data, 2,axis=1)   )/(12.*dtheta)
-            der[:, 1, :] = (data[:, 2, :]-data[:, 0, :])/(2.*dtheta)
-            der[:, -2, :] = (data[:, -1, :]-data[:, -3, :])/(2.*dtheta)
-            der[:, 0, :] = (data[:, 1, :]-data[:, 0, :])/dtheta
-            der[:, -1, :] = (data[:, -1, :]-data[:, -2, :])/dtheta
+            der = (-np.roll(data,-2,axis=1)+8.*np.roll(data,-1,axis=1) \
+                   -8.*np.roll(data, 1,axis=1)+np.roll(data, 2,axis=1)) / \
+                  (-np.roll(th[None, :, None],-2)+8*np.roll(th[None, :, None],-1) \
+                   -8*np.roll(th[None, :, None],1)+np.roll(th[None, :, None],2))
+            der[:, 1, :] = (data[:, 2, :]-data[:, 0, :])/(th[2]-th[0])
+            der[:, -2, :] = (data[:, -1, :]-data[:, -3, :])/(th[-1]-th[-3])
+            der[:, 0, :] = (data[:, 1, :]-data[:, 0, :])/(th[1]-th[0])
+            der[:, -1, :] = (data[:, -1, :]-data[:, -2, :])/(th[-1]-th[-2])
 
     elif len(data.shape) == 2: #2-D
-        ntheta = data.shape[0]
-        dtheta = np.pi/(ntheta-1.)
         if order == 2:
-            der = (np.roll(data, -1,  axis=0)-np.roll(data, 1, axis=0))/(2.*dtheta)
-            der[0, :] = (data[1, :]-data[0, :])/dtheta
-            der[-1, :] = (data[-1, :]-data[-2, :])/dtheta
+            der = (np.roll(data,-1,axis=0)-np.roll(data,1,axis=0)) / \
+                  (np.roll(th[:, None],-1)-np.roll(th[:, None],1))
+            der[0, :] = (data[1, :]-data[0, :])/(th[1]-th[0])
+            der[-1, :] = (data[-1, :]-data[-2, :])/(th[-1]-th[-2])
         elif order == 4:
             der = (-np.roll(data,-2,axis=0)+8.*np.roll(data,-1,axis=0)-\
-                  8.*np.roll(data,1,axis=0)+np.roll(data,2,axis=0))/(12.*dtheta)
-            der[1, :] = (data[2, :]-data[0, :])/(2.*dtheta)
-            der[-2, :] = (data[-1, :]-data[-3, :])/(2.*dtheta)
-            der[0, :] = (data[1, :]-data[0, :])/dtheta
-            der[-1, :] = (data[-1, :]-data[-2, :])/dtheta
+                  8.*np.roll(data,1,axis=0)+np.roll(data,2,axis=0)) / \
+                  (-np.roll(th[:, None],-2)+8*np.roll(th[:, None],-1)-\
+                   8.*np.roll(th[:, None],1)+np.roll(th[:, None],2))
+            der[1, :] = (data[2, :]-data[0, :])/(th[2]-th[0])
+            der[-2, :] = (data[-1, :]-data[-3, :])/(th[-1]-th[-3])
+            der[0, :] = (data[1, :]-data[0, :])/(th[1]-th[0])
+            der[-1, :] = (data[-1, :]-data[-2, :])/(th[-1]-th[-2])
 
     return der
 
@@ -953,15 +964,11 @@ def zderavg(data, rad, colat=None, exclude=False):
         th = np.linspace(0., np.pi, ntheta)
 
     if len(data.shape) == 3:  # 3-D
-        thmD = np.zeros_like(data)
-        for i in range(ntheta):
-            thmD[:,i,:] = th[i]
+        thmD = th[None, :, None]
     elif len(data.shape) == 2:  # 2-D
-        thmD = np.zeros((ntheta, nr), np.float64)
-        for i in range(ntheta):
-            thmD[i, :] = th[i]
+        thmD = th[None, :]
 
-    dtheta = thetaderavg(data)
+    dtheta = thetaderavg(data, colat=colat)
     dr = rderavg(data, rad, exclude)
     dz = np.cos(thmD)*dr - np.sin(thmD)/rad*dtheta
 
@@ -997,15 +1004,11 @@ def sderavg(data, rad, colat=None, exclude=False):
         th = np.linspace(0., np.pi, ntheta)
 
     if len(data.shape) == 3:  # 3-D
-        thmD = np.zeros_like(data)
-        for i in range(ntheta):
-            thmD[:,i,:] = th[i]
+        thmD = th[None, :, None]
     elif len(data.shape) == 2:  # 2-D
-        thmD = np.zeros((ntheta, nr), np.float64)
-        for i in range(ntheta):
-            thmD[i, :] = th[i]
+        thmD = th[None, :]
 
-    dtheta = thetaderavg(data)
+    dtheta = thetaderavg(data, colat=colat)
     dr = rderavg(data, rad, exclude)
     ds = np.sin(thmD)*dr + np.cos(thmD)/rad*dtheta