PYTHON: Rewrite interpolations.

Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>

python/amunpy/src/amunpy/interpolation.py

@@ -22,14 +22,10 @@
 
 ================================================================================
 
- module: AMUN
+ module: INTERPOLATION
 
-  Python module with subroutines to read AMUN code HDF5 files.
+  Support module for Amun snapshots for the block prolongation with different
-
+  types of interpolation.
-  The only requirements for this package are:
-
-     - numpy
-     - scipy (optional, for data interpolation)
 
 --------------------------------------------------------------------------------
 """
@@ -52,67 +48,44 @@ def rebin(a, newshape):
     m = a.ndim - 1
     if (a.shape[m] > newshape[m]):
         if a.ndim == 3:
-      nn = [newshape[0], int(a.shape[0] / newshape[0]),
+            nn = [newshape[0], a.shape[0] // newshape[0], \
-            newshape[1], int(a.shape[1] / newshape[1]),
+                  newshape[1], a.shape[1] // newshape[1], \
-            newshape[2], int(a.shape[2] / newshape[2])]
+                  newshape[2], a.shape[2] // newshape[2]]
             return a.reshape(nn).mean(5).mean(3).mean(1)
         else:
-      nn = [newshape[0], int(a.shape[0] / newshape[0]),
+            nn = [newshape[0], a.shape[0] // newshape[0], \
-            newshape[1], int(a.shape[1] / newshape[1])]
+                  newshape[1], a.shape[1] // newshape[1]]
             return a.reshape(nn).mean(3).mean(1)
     else:
         for n in range(a.ndim):
-      a = np.repeat(a, int(newshape[n] / a.shape[n]), axis=n)
+            a = np.repeat(a, newshape[n] // a.shape[n], axis=n)
         return(a)
 
 
-def interpolate(a, newshape, nghosts, method='rebin', order=3):
+def interpolate(a, newshape, nghosts=0, method=None, order=1):
     '''
         Subroutine rescales the block by interpolating its values.
     '''
-  if (method == 'rebin' or not scipy_available):
+    if method == None or method == 'rebin' or not scipy_available:
 
-    # calculate the indices in order to remove the ghost zones
+        ng = nghosts
-    #
         if a.ndim == 3:
-      ib = nghosts
+            return rebin(a[ng:-ng,ng:-ng,ng:-ng], newshape)
-      jb = nghosts
-      kb = nghosts
-      ie = a.shape[2] - nghosts
-      je = a.shape[1] - nghosts
-      ke = a.shape[0] - nghosts
-      if (a.shape[0] == 1):
-        kb = 0
-        ke = 1
-
-      return rebin(a[kb:ke,jb:je,ib:ie], newshape)
         else:
-      ib = nghosts
+            return rebin(a[ng:-ng,ng:-ng], newshape)
-      jb = nghosts
-      ie = a.shape[1] - nghosts
-      je = a.shape[0] - nghosts
 
-      return rebin(a[jb:je,ib:ie], newshape)
+    elif method == 'zoom':
 
-  elif (method == 'zoom'):
+        zf = (newshape[1] // (a.shape[1] - 2 * nghosts))
-
+        ng = zf * nghosts
-    fc = int(newshape[1] / (a.shape[1] - 2 * nghosts))
-    ib, ie = fc * nghosts, - fc * nghosts
-    jb, je = fc * nghosts, - fc * nghosts
         if a.ndim == 3:
-      kb, ke = fc * nghosts, - fc * nghosts
+            return zoom(a, zf, order=order, grid_mode=True, mode='nearest')[ng:-ng,ng:-ng,ng:-ng]
-      if (a.shape[0] == 1):
-        kb = 0
-        ke = 1
-
-      return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[kb:ke,jb:je,ib:ie]
         else:
-      return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[jb:je,ib:ie]
+            return zoom(a, zf, order=order, grid_mode=True, mode='nearest')[ng:-ng,ng:-ng]
 
-  elif (method == 'monotonic' or method == 'pchip'):
+    elif method in [ 'monotonic', 'pchip' ]:
 
         dims = np.arange(a.ndim)
-
         q = a
 
         for n in dims:
@@ -125,9 +98,9 @@ def interpolate(a, newshape, nghosts, method='rebin', order=3):
             xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
             xn =  np.arange(0.5, newshape[n]) / newshape[n]
 
-      u = q.reshape([d[0], int(q.size / d[0])])
+            u = q.reshape([d[0], q.size // d[0]])
-      f = np.zeros([newshape[n], int(q.size / d[0])])
+            f = np.zeros([newshape[n], q.size // d[0]])
-      for i in range(int(q.size / d[0])):
+            for i in range(q.size // d[0]):
                 f[:,i] = pchip_interpolate(xo, u[:,i], xn)
 
             d[0] = newshape[n]
@@ -137,10 +110,9 @@ def interpolate(a, newshape, nghosts, method='rebin', order=3):
 
         return q
 
-  elif (method == 'spline'):
+    elif method == 'spline':
 
         dims = np.arange(a.ndim)
-
         q = a
 
         for n in dims:
@@ -153,9 +125,9 @@ def interpolate(a, newshape, nghosts, method='rebin', order=3):
             xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
             xn =  np.arange(0.5, newshape[n]) / newshape[n]
 
-      u = q.reshape([d[0], int(q.size / d[0])])
+            u = q.reshape([d[0], q.size // d[0]])
-      f = np.zeros([newshape[n], int(q.size / d[0])])
+            f = np.zeros([newshape[n], q.size // d[0]])
-      for i in range(int(q.size / d[0])):
+            for i in range(q.size // d[0]):
                 t = splrep(xo, u[:,i], k=5, s=0.0)
                 f[:,i] = splev(xn, t)
 
@@ -169,7 +141,6 @@ def interpolate(a, newshape, nghosts, method='rebin', order=3):
     else:
 
         dims = np.arange(a.ndim)
-
         q = a
 
         for n in dims:
@@ -182,9 +153,9 @@ def interpolate(a, newshape, nghosts, method='rebin', order=3):
             xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
             xn =  np.arange(0.5, newshape[n]) / newshape[n]
 
-      u = q.reshape([d[0], int(q.size / d[0])])
+            u = q.reshape([d[0], q.size // d[0]])
-      f = np.zeros([newshape[n], int(q.size / d[0])])
+            f = np.zeros([newshape[n], q.size // d[0]])
-      for i in range(int(q.size / d[0])):
+            for i in range(q.size // d[0]):
                 t = interp1d(xo, u[:,i], kind=method)
                 f[:,i] = t(xn)