Refactoring

Author: maroba

findiff/__init__.py

@@ -22,7 +22,8 @@
 __version__ = "0.12.0"
 
 
-from .operators import Diff, Identity
+from .interface import Diff
+from .operators import Identity
 from .pde import PDE, BoundaryConditions
 from .compatible import Coef, Coefficient, FinDiff, Id
 from .coefs import coefficients

findiff/compatible.py

@@ -1,13 +1,74 @@
-"""Provides an interface to obsolete classes for backward compatibility."""
+"""This module provides an interface to obsolete classes for backward compatibility."""
 
 from findiff import Diff
-from findiff.legacy.operators import _FinDiff
 from findiff.operators import FieldOperator, Identity
 
 
-# @deprecated(reason="Use findiff.Diff instead.")
 def FinDiff(*args, **kwargs):
+    r"""A representation of a general linear differential operator expressed in finite differences.
 
+        FinDiff objects can be added with other FinDiff objects. They can be multiplied by
+        objects of type Coefficient.
+
+        FinDiff is callable, i.e. to apply the derivative, just call the object on the array to
+        differentiate.
+
+        :param args: variable number of tuples. Defines what derivative to take.
+            If only one tuple is given, you can leave away the tuple parentheses.
+
+        Each tuple has the form
+
+               `(axis, spacing, count)`     for uniform grids
+
+               `(axis, count)`              for non-uniform grids.
+
+             `axis` is the dimension along which to take derivative.
+
+             `spacing` is the grid spacing of the uniform grid along that axis.
+
+             `count` is the order of the derivative, which is optional an defaults to 1.
+
+
+        :param kwargs:  variable number of keyword arguments
+
+            Allowed keywords:
+
+            `acc`:    even integer
+                  The desired accuracy order. Default is acc=2.
+
+        This class is actually deprecated and will be replaced by the Diff class in the future.
+
+    **Example**:
+
+
+       For this example, we want to operate on some 3D array f:
+
+       >>> import numpy as np
+       >>> x, y, z = [np.linspace(-1, 1, 100) for _ in range(3)]
+       >>> X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
+       >>> f = X**2 + Y**2 + Z**2
+
+       To create :math:`\\frac{\\partial f}{\\partial x}` on a uniform grid with spacing dx, dy
+       along the 0th axis or 1st axis, respectively, instantiate a FinDiff object and call it:
+
+       >>> d_dx = FinDiff(0, dx)
+       >>> d_dy = FinDiff(1, dx)
+       >>> result = d_dx(f)
+
+       For :math:`\\frac{\\partial^2 f}{\\partial x^2}` or :math:`\\frac{\\partial^2 f}{\\partial y^2}`:
+
+       >>> d2_dx2 = FinDiff(0, dx, 2)
+       >>> d2_dy2 = FinDiff(1, dy, 2)
+       >>> result_2 = d2_dx2(f)
+       >>> result_3 = d2_dy2(f)
+
+       For :math:`\\frac{\\partial^4 f}{\partial x \\partial^2 y \\partial z}`, do:
+
+       >>> op = FinDiff((0, dx), (1, dy, 2), (2, dz))
+       >>> result_4 = op(f)
+
+
+    """
     if len(args) > 3:
         raise ValueError("FinDiff accepts not more than 3 positional arguments.")
 
@@ -31,15 +92,11 @@ def diff_from_tuple(tpl):
     return diff_from_tuple(args)
 
 
-FinDiff.__doc__ = _FinDiff.__doc__
-
-
-"""
-Define aliasses for backward compatibility:
-"""
+###
+### Define aliasses for backward compatibility:
+###
 
 
-# @deprecated(reason="No need to wrap array in Coefficient object any more.")
 class Coefficient(FieldOperator):
     pass
 

findiff/findiff.py

@@ -0,0 +1,268 @@
+import itertools
+
+import numpy as np
+from scipy import sparse
+
+from findiff.coefs import coefficients_non_uni, coefficients
+from findiff.grids import GridAxis, EquidistantAxis, NonEquidistantAxis
+from findiff.utils import long_indices_as_ndarray, to_long_index
+
+
+def build_differentiator(order: int, axis: GridAxis, acc):
+    if isinstance(axis, EquidistantAxis):
+        if not axis.periodic:
+            return _FinDiffUniform(axis.dim, order, axis.spacing, acc)
+        else:
+            return _FinDiffUniformPeriodic(axis.dim, order, axis.spacing, acc)
+    elif isinstance(axis, NonEquidistantAxis):
+        if not axis.periodic:
+            return _FinDiffNonUniform(axis.dim, order, axis.coords, acc)
+        else:
+            raise NotImplementedError("Periodic nonuniform axes not yet implemented")
+    else:
+        raise TypeError("Unknown axis type.")
+
+
+class _FinDiffBase:
+
+    def __init__(self, axis, order):
+        self.axis = axis
+        self.order = order
+
+    def validate_f(self, f):
+        try:
+            f.shape[self.axis]
+        except AttributeError as err:
+            raise ValueError(
+                "Diff objects can only be applied to arrays or evaluated(!) functions returning arrays"
+            ) from err
+
+    def apply_to_array(self, yd, y, weights, off_slices, ref_slice, dim):
+        """Applies the finite differences only to slices along a given axis"""
+
+        ndims = len(y.shape)
+
+        all = slice(None, None, 1)
+
+        ref_multi_slice = [all] * ndims
+        ref_multi_slice[dim] = ref_slice
+
+        for w, s in zip(weights, off_slices):
+            off_multi_slice = [all] * ndims
+            off_multi_slice[dim] = s
+            if abs(1 - w) < 1.0e-14:
+                yd[tuple(ref_multi_slice)] += y[tuple(off_multi_slice)]
+            else:
+                yd[tuple(ref_multi_slice)] += w * y[tuple(off_multi_slice)]
+
+    def shift_slice(self, sl, off, max_index):
+
+        if sl.start + off < 0 or sl.stop + off > max_index:
+            raise IndexError("Shift slice out of bounds")
+
+        return slice(sl.start + off, sl.stop + off, sl.step)
+
+
+class _FinDiffUniform(_FinDiffBase):
+
+    def __init__(self, axis, order, spacing, acc):
+        super().__init__(axis, order)
+        self.spacing = spacing
+        self.acc = acc
+        coef_schemes = coefficients(self.order, acc)
+        self.forward = coef_schemes["forward"]
+        self.backward = coef_schemes["backward"]
+        self.center = coef_schemes["center"]
+
+    def __call__(self, f):
+        self.validate_f(f)
+        npts = f.shape[self.axis]
+        weights = self.center["coefficients"]
+        offsets = self.center["offsets"]
+
+        num_bndry_points = len(weights) // 2
+        ref_slice = slice(num_bndry_points, npts - num_bndry_points, 1)
+        off_slices = [
+            self.shift_slice(ref_slice, offsets[k], npts) for k in range(len(offsets))
+        ]
+
+        fd = np.zeros_like(f)
+
+        self.apply_to_array(fd, f, weights, off_slices, ref_slice, self.axis)
+
+        weights = self.forward["coefficients"]
+        offsets = self.forward["offsets"]
+
+        ref_slice = slice(0, num_bndry_points, 1)
+        off_slices = [
+            self.shift_slice(ref_slice, offsets[k], npts) for k in range(len(offsets))
+        ]
+
+        self.apply_to_array(fd, f, weights, off_slices, ref_slice, self.axis)
+
+        weights = self.backward["coefficients"]
+        offsets = self.backward["offsets"]
+
+        ref_slice = slice(npts - num_bndry_points, npts, 1)
+        off_slices = [
+            self.shift_slice(ref_slice, offsets[k], npts) for k in range(len(offsets))
+        ]
+
+        self.apply_to_array(fd, f, weights, off_slices, ref_slice, self.axis)
+
+        h_inv = 1.0 / self.spacing**self.order
+        return fd * h_inv
+
+    def matrix(self, shape):
+
+        h = self.spacing
+
+        ndims = len(shape)
+        siz = np.prod(shape)
+        long_indices_nd = long_indices_as_ndarray(shape)
+
+        axis, order = self.axis, self.order
+        mat = sparse.lil_matrix((siz, siz))
+
+        for scheme in ["center", "forward", "backward"]:
+
+            offsets_1d = getattr(self, scheme)["offsets"]
+            coeffs = getattr(self, scheme)["coefficients"]
+
+            # translate offsets of given scheme to long format
+            offsets_long = []
+            for o_1d in offsets_1d:
+                o_nd = np.zeros(ndims)
+                o_nd[axis] = o_1d
+                o_long = to_long_index(o_nd, shape)
+                offsets_long.append(o_long)
+
+            # determine points where to evaluate current scheme in long format
+            nside = len(self.center["coefficients"]) // 2
+            if scheme == "center":
+                multi_slice = [slice(None, None)] * ndims
+                multi_slice[axis] = slice(nside, -nside)
+                Is = long_indices_nd[tuple(multi_slice)].reshape(-1)
+            elif scheme == "forward":
+                multi_slice = [slice(None, None)] * ndims
+                multi_slice[axis] = slice(0, nside)
+                Is = long_indices_nd[tuple(multi_slice)].reshape(-1)
+            else:
+                multi_slice = [slice(None, None)] * ndims
+                multi_slice[axis] = slice(-nside, None)
+                Is = long_indices_nd[tuple(multi_slice)].reshape(-1)
+
+            for o, c in zip(offsets_long, coeffs):
+                v = c / h**order
+                mat[Is, Is + o] = v
+
+        return mat
+
+
+class _FinDiffUniformPeriodic(_FinDiffBase):
+
+    def __init__(self, axis, order, spacing, acc):
+        super().__init__(axis, order)
+        self.spacing = spacing
+        self.acc = acc
+        self.coefs = coefficients(self.order, acc)["center"]
+
+    def __call__(self, f):
+        self.validate_f(f)
+        fd = np.zeros_like(f)
+        for off, coef in zip(self.coefs["offsets"], self.coefs["coefficients"]):
+            fd += coef * np.roll(f, -off, axis=self.axis)
+        h_inv = 1.0 / self.spacing**self.order
+        return fd * h_inv
+
+    def matrix(self, shape):
+        h = self.spacing
+
+        ndims = len(shape)
+        siz = np.prod(shape)
+        long_indices_nd = long_indices_as_ndarray(shape)
+
+        axis, order = self.axis, self.order
+        mat = sparse.lil_matrix((siz, siz))
+
+        offsets = self.coefs["offsets"]
+        coefs = self.coefs["coefficients"]
+
+        multi_slice = [slice(None, None)] * ndims
+        Is = long_indices_nd[tuple(multi_slice)].reshape(-1)
+
+        idxs_short = [np.arange(n) for n in shape]
+
+        for o, c in zip(offsets, coefs):
+            v = c / h**order
+
+            idxs_short[self.axis] = np.roll(np.arange(shape[self.axis]), -o)
+            grid = np.meshgrid(*idxs_short, indexing="ij")
+            index_tuples = np.stack(grid, axis=-1).reshape(-1, ndims)
+
+            Is_off = np.ravel_multi_index(index_tuples.T, shape)
+
+            mat[Is, Is_off] = v
+
+        return mat
+
+
+class _FinDiffNonUniform(_FinDiffBase):
+    def __init__(self, axis, order, coords, acc):
+        super().__init__(axis, order)
+        self.coords = coords
+        self.acc = acc
+        self.coef_list = []
+        for i in range(len(self.coords)):
+            self.coef_list.append(coefficients_non_uni(order, self.acc, self.coords, i))
+
+    def __call__(self, y):
+        """The core function to take a partial derivative on a non-uniform grid"""
+
+        order, dim = self.order, self.axis
+        yd = np.zeros_like(y)
+
+        ndims = len(y.shape)
+        multi_slice = [slice(None, None)] * ndims
+        ref_multi_slice = [slice(None, None)] * ndims
+
+        for i, x in enumerate(self.coords):
+
+            coefs = self.coef_list[i]
+            weights = coefs["coefficients"]
+            offsets = coefs["offsets"]
+            ref_multi_slice[dim] = i
+
+            for off, w in zip(offsets, weights):
+                multi_slice[dim] = i + off
+                yd[tuple(ref_multi_slice)] += w * y[tuple(multi_slice)]
+
+        return yd
+
+    def matrix(self, shape):
+
+        coords = self.coords
+
+        siz = np.prod(shape)
+        long_inds = np.arange(siz).reshape(shape)
+        short_inds = [np.arange(shape[k]) for k in range(len(shape))]
+        short_inds = list(itertools.product(*short_inds))
+
+        coef_dicts = []
+        for i in range(len(coords)):
+            coef_dicts.append(coefficients_non_uni(self.order, self.acc, coords, i))
+
+        mat = sparse.lil_matrix((siz, siz))
+
+        for base_ind_long, base_ind_short in enumerate(short_inds):
+            cd = coef_dicts[base_ind_short[self.axis]]
+            cs, os = cd["coefficients"], cd["offsets"]
+            for c, o in zip(cs, os):
+                off_short = np.zeros(len(shape), dtype=int)
+                off_short[self.axis] = int(o)
+                off_ind_short = np.array(base_ind_short, dtype=int) + off_short
+                off_long = long_inds[tuple(off_ind_short)]
+
+                mat[base_ind_long, off_long] += c
+
+        return mat