Skip to content

Commit

Permalink
Refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
maroba committed Dec 17, 2024
1 parent b7266f2 commit 3b89c08
Show file tree
Hide file tree
Showing 10 changed files with 527 additions and 457 deletions.
3 changes: 2 additions & 1 deletion findiff/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
77 changes: 67 additions & 10 deletions findiff/compatible.py
Original file line number Diff line number Diff line change
@@ -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.")

Expand All @@ -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

Expand Down
268 changes: 268 additions & 0 deletions findiff/findiff.py
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 3b89c08

Please sign in to comment.