Add finite difference operators.

1 files changed, 174 insertions, 0 deletions

diff --git a/diff.py b/diff.py
new file mode 100644
index 0000000..2430a0e
--- /dev/null
+++ b/diff.py
@@ -0,0 +1,174 @@
+import numpy as np
+
+class FiniteDifferenceUniform(object):
+    """
+    Finite difference operator on a uniform grid.
+
+    :param int diff_order: Order of the derivative to take. Currently only 1 and
+                           2 are supported.
+    :param int acc_order: Accuracy order of the operator. Derivatives of
+                          polynomials up to this order will be computed
+                          exactly.
+    """
+    # public
+    diff_order = None
+    acc_order  = None
+
+    """
+    For each source point, this number of points in each direction is needed
+    to compute the derivative to desired order. Points at the boundary are
+    computed with lower order operators.
+    """
+    stencil = None
+
+    # private
+    _coeffs  = None
+    _div     = None
+
+    _diff_table = [
+        {
+        2 : {
+                'stencil' : 1,
+                'coeffs'  : (-1., 0., 1.),
+                'div'     : 2.,
+            },
+        4 : {
+                'stencil' : 2,
+                'coeffs'  : (1., -8., 0., 8., -1.),
+                'div'     : 12.,
+            },
+        6 : {
+                'stencil' : 3,
+                'coeffs'  : (-1., 9., -45., 0., 45., -9., 1.),
+                'div'     : 60.,
+            },
+        8 : {
+                'stencil' : 4,
+                'coeffs'  : (3., -32., 168., -672., 0., 672., -168., 32., -3.),
+                'div'     : 840.,
+            },
+        },
+        {
+        2 : {
+                'stencil' : 1,
+                'coeffs'  : (1., -2., 1.),
+                'div'     : 1.,
+            },
+        4 : {
+                'stencil' : 2,
+                'coeffs'  : (-1., 16., -30., 16., -1.),
+                'div'     : 12.,
+            },
+        6 : {
+                'stencil' : 3,
+                'coeffs'  : (2., -27., 270., -490., 270., -27., 2.),
+                'div'     : 180.,
+            },
+        8 : {
+                'stencil' : 4,
+                'coeffs'  : (-9., 128., -1008., 8064., -14350., 8064., -1008., 128., -9.),
+                'div'     : 5040.,
+            },
+        },
+    ]
+
+    def __init__(self, diff_order, acc_order):
+        self.diff_order   = diff_order
+        self.acc_order    = acc_order
+        self._coeffs      = self._diff_table[diff_order - 1][acc_order]['coeffs']
+        self._div         = self._diff_table[diff_order - 1][acc_order]['div']
+        self.stencil      = self._diff_table[diff_order - 1][acc_order]['stencil']
+
+    def __call__(self, arr, axis, dx, bound = 'both'):
+        """
+        Evaluate the derivative for an array.
+
+        :param array_like arr: Array to differentiate.
+        :param int axis: Axis along which to differentiate.
+        :param float dx: Spacing between array elements.
+        """
+        coeff   = self._coeffs
+        div     = self._div
+        stencil = self.stencil
+
+        bound_lower = bound == 'lower' or bound == 'both'
+        bound_upper = bound == 'upper' or bound == 'both'
+
+        slices_src = []
+        for i in range(stencil * 2 + 1):
+            sl    = [slice(None) for _ in arr.shape]
+            start = i
+            end   = -stencil * 2 + i
+            if end == 0:
+                end = None
+            sl[axis] = slice(start, end)
+            slices_src.append(tuple(sl))
+
+        slicelist_ret       = [slice(None)] * len(arr.shape)
+        slicelist_ret[axis] = slice(stencil, -stencil)
+
+        ret = np.zeros_like(arr)
+
+        for s, c in zip(slices_src, coeff):
+            ret[tuple(slicelist_ret)] += c * arr[tuple(s)]
+        ret[tuple(slicelist_ret)] /= (div * dx)
+
+        if bound_lower:
+            if self.acc_order > 2:
+                fd_lower = FiniteDifferenceUniform(self.diff_order, self.acc_order - 2)
+
+                slicelist_edge_src       = [slice(None) for _ in arr.shape]
+                slicelist_edge_src[axis] = slice(0, 2 * stencil)
+
+                diff_edge = fd_lower(arr[tuple(slicelist_edge_src)], axis, dx, bound = 'lower')
+
+                slicelist_edge_dst             = [slice(None) for _ in arr.shape]
+                slicelist_edge_dst[axis]       = slice(0, stencil)
+
+                ret[tuple(slicelist_edge_dst)] = diff_edge[tuple(slicelist_edge_dst)]
+            else:
+                slicelist_edge_dst       = [slice(None) for _ in arr.shape]
+                slicelist_edge_dst[axis] = slice(0, 1)
+
+                slicelist_edge_p1       = [slice(None) for _ in arr.shape]
+                slicelist_edge_p1[axis] = slice(1, 2)
+
+                slicelist_edge_m1       = [slice(None) for _ in arr.shape]
+                slicelist_edge_m1[axis] = slice(0, 1)
+
+                ret[tuple(slicelist_edge_dst)] = (arr[tuple(slicelist_edge_p1)] - arr[tuple(slicelist_edge_m1)]) / dx
+
+        if bound_upper:
+            if self.acc_order > 2:
+                fd_lower = FiniteDifferenceUniform(self.diff_order, self.acc_order - 2)
+
+                slicelist_edge_src       = [slice(None) for _ in arr.shape]
+                slicelist_edge_src[axis] = slice(-2 * stencil, None)
+
+                diff_edge = fd_lower(arr[tuple(slicelist_edge_src)], axis, dx, bound = 'upper')
+
+                slicelist_edge_dst             = [slice(None) for _ in arr.shape]
+                slicelist_edge_dst[axis]       = slice(-stencil, None)
+
+                ret[tuple(slicelist_edge_dst)] = diff_edge[tuple(slicelist_edge_dst)]
+            else:
+                slicelist_edge_dst       = [slice(None) for _ in arr.shape]
+                slicelist_edge_dst[axis] = slice(-1, None)
+
+                slicelist_edge_p1       = [slice(None) for _ in arr.shape]
+                slicelist_edge_p1[axis] = slice(-1, None)
+
+                slicelist_edge_m1       = [slice(None) for _ in arr.shape]
+                slicelist_edge_m1[axis] = slice(-2, -1)
+
+                ret[tuple(slicelist_edge_dst)] = (arr[tuple(slicelist_edge_p1)] - arr[tuple(slicelist_edge_m1)]) / dx
+
+        return ret
+
+fd2 = FiniteDifferenceUniform(1, 2)
+fd4 = FiniteDifferenceUniform(1, 4)
+fd8 = FiniteDifferenceUniform(1, 8)
+
+fd22 = FiniteDifferenceUniform(2, 2)
+fd24 = FiniteDifferenceUniform(2, 4)
+fd28 = FiniteDifferenceUniform(2, 8)