diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py index 6803e8a75b7b9d4e83974a044af878629561ff44..719801afcace0923a716aa9136fe152092001bd9 100644 --- a/sumpy/expansion/__init__.py +++ b/sumpy/expansion/__init__.py @@ -349,7 +349,7 @@ class LinearPDEBasedExpansionTermsWrangler(ExpansionTermsWrangler): def get_pde_as_diff_op(self): r""" - Returns the PDE as a :class:`sumpy.expansion.diff_op.DifferentialOperator` + Returns the PDE as a :class:`sumpy.expansion.diff_op.LinearPDESystemOperator` object `L` where `L(u) = 0` is the PDE. """ @@ -366,7 +366,9 @@ class LinearPDEBasedExpansionTermsWrangler(ExpansionTermsWrangler): coeff_ident_enumerate_dict = {tuple(mi): i for (i, mi) in enumerate(mis)} - pde_dict = self.get_pde_as_diff_op().mi_to_coeff + diff_op = self.get_pde_as_diff_op() + assert len(diff_op.eqs) == 1 + pde_dict = dict((k.mi, v) for k, v in diff_op.eqs[0].items()) for ident in pde_dict.keys(): if ident not in coeff_ident_enumerate_dict: # Order of the expansion is less than the order of the PDE. diff --git a/sumpy/expansion/diff_op.py b/sumpy/expansion/diff_op.py index 73e85320e3835748112372ad259c0a7c8b50d3e3..15cb3a4669dab6850696e93ee271eade376c5890 100644 --- a/sumpy/expansion/diff_op.py +++ b/sumpy/expansion/diff_op.py @@ -24,77 +24,158 @@ THE SOFTWARE. from pyrsistent import pmap from sumpy.tools import add_mi +from collections import namedtuple __doc__ = """ Differential operator interface ------------------------------- -.. autoclass:: DifferentialOperator +.. autoclass:: LinearPDESystemOperator +.. autoclass:: DerivativeIdentifier +.. autofunction:: make_identity_diff_op """ +DerivativeIdentifier = namedtuple("DerivativeIdentifier", ["mi", "vec_idx"]) -class DifferentialOperator(object): + +class LinearPDESystemOperator(object): r""" - Represents a scalar, constant-coefficient DifferentialOperator of - dimension `dim`. It is represented by a frozen dictionary. - The dictionary maps a multi-index given as a tuple to the coefficient. - This object is immutable. + Represents a constant-coefficient linear differential operator of a + vector-valued function with `dim` variables. It is represented by a tuple of + immutable dictionaries. The dictionary maps a :class:`DerivativeIdentifier` + to the coefficient. This object is immutable. """ - def __init__(self, dim, mi_to_coeff): + def __init__(self, dim, *eqs): """ - :arg dim: dimension of the DifferentialOperator - :arg mi_to_coeff: A dictionary mapping a multi-index to a coefficient + :arg dim: dimension of the LinearPDESystemOperator + :arg eqs: A list of dictionaries mapping a :class:`DerivativeIdentifier` + to a coefficient. """ self.dim = dim - self.mi_to_coeff = mi_to_coeff + self.eqs = tuple(eqs) def __mul__(self, param): - mi_to_coeff = {} - for k, v in self.mi_to_coeff.items(): - mi_to_coeff[k] = v * param - return DifferentialOperator(self.dim, pmap(mi_to_coeff)) + eqs = [] + for eq in self.eqs: + deriv_ident_to_coeff = {} + for k, v in eq.items(): + deriv_ident_to_coeff[k] = v * param + eqs.append(pmap(deriv_ident_to_coeff)) + return LinearPDESystemOperator(self.dim, *eqs) __rmul__ = __mul__ - def __add__(self, other_pde): - assert self.dim == other_pde.dim - res = dict(self.mi_to_coeff) - for k, v in other_pde.mi_to_coeff.items(): - if k in res: - res[k] += v - else: - res[k] = v - return DifferentialOperator(self.dim, pmap(res)) + def __add__(self, other_diff_op): + assert self.dim == other_diff_op.dim + assert len(self.eqs) == len(other_diff_op.eqs) + eqs = [] + for eq, other_eq in zip(self.eqs, other_diff_op.eqs): + res = dict(eq) + for k, v in other_eq.items(): + if k in res: + res[k] += v + else: + res[k] = v + eqs.append(pmap(res)) + return LinearPDESystemOperator(self.dim, *eqs) __radd__ = __add__ - def __sub__(self, other_pde): - return self + (-1)*other_pde + def __sub__(self, other_diff_op): + return self + (-1)*other_diff_op def __repr__(self): - return f"DifferentialOperator({self.dim}, {repr(self.mi_to_coeff)})" - - -def laplacian(pde): - dim = pde.dim - res = DifferentialOperator(dim, pmap()) + return f"LinearPDESystemOperator({self.dim}, {repr(self.eqs)})" + + def __getitem__(self, idx): + item = self.eqs.__getitem__(idx) + if not isinstance(item, tuple): + item = (item,) + return LinearPDESystemOperator(self.dim, *item) + + def to_sym(self, fnames=None): + from sumpy.symbolic import make_sym_vector, Function + x = make_sym_vector("x", self.dim) + + if fnames is None: + noutputs = 0 + for eq in self.eqs: + for deriv_ident in eq.keys(): + noutputs = max(noutputs, deriv_ident.vec_idx) + fnames = [f"f{i}" for i in range(noutputs+1)] + + funcs = [Function(fname)(*x) for fname in fnames] + + res = [] + for eq in self.eqs: + sym_eq = 0 + for deriv_ident, coeff in eq.items(): + expr = funcs[deriv_ident.vec_idx] + for i, val in enumerate(deriv_ident.mi): + for j in range(val): + expr = expr.diff(x[i]) + sym_eq += expr * coeff + res.append(sym_eq) + return res + + +def laplacian(diff_op): + dim = diff_op.dim + empty = [pmap()] * len(diff_op.eqs) + res = LinearPDESystemOperator(dim, *empty) for j in range(dim): mi = [0]*dim mi[j] = 2 - res = res + diff(pde, tuple(mi)) + res = res + diff(diff_op, tuple(mi)) + return res + + +def diff(diff_op, mi): + eqs = [] + for eq in diff_op.eqs: + res = {} + for deriv_ident, v in eq.items(): + new_mi = add_mi(deriv_ident.mi, mi) + res[DerivativeIdentifier(new_mi, deriv_ident.vec_idx)] = v + eqs.append(pmap(res)) + return LinearPDESystemOperator(diff_op.dim, *eqs) + + +def divergence(diff_op): + assert len(diff_op.eqs) == diff_op.dim + dim = diff_op.dim + res = LinearPDESystemOperator(dim, pmap()) + for i in range(dim): + mi = [0]*dim + mi[i] = 1 + res += diff(diff_op[i], tuple(mi)) return res -def diff(pde, mi): - res = {} - for mi_to_coeff_mi, v in pde.mi_to_coeff.items(): - res[add_mi(mi_to_coeff_mi, mi)] = v - return DifferentialOperator(pde.dim, pmap(res)) +def gradient(diff_op): + assert len(diff_op.eqs) == 1 + eqs = [] + dim = diff_op.dim + for i in range(dim): + mi = [0]*dim + mi[i] = 1 + eqs.append(diff(diff_op, tuple(mi)).eqs[0]) + return LinearPDESystemOperator(dim, *eqs) + + +def concat(op1, op2): + assert op1.dim == op2.dim + eqs = list(op1.eqs) + eqs.extend(list(op2.eqs)) + return LinearPDESystemOperator(op1.dim, *eqs) -def make_identity_diff_op(dim): +def make_identity_diff_op(ninput, noutput=1): """ - Returns the identity as a differential operator. + Returns the identity as a linear PDE system operator. + :arg ninput: number of input variables to the function + :arg noutput: number of output values of function """ - mi = tuple([0]*dim) - return DifferentialOperator(dim, pmap({mi: 1})) + mi = tuple([0]*ninput) + eqs = [pmap({DerivativeIdentifier(mi, i): 1}) for i in range(noutput)] + return LinearPDESystemOperator(ninput, *eqs) diff --git a/test/test_misc.py b/test/test_misc.py index 75c77974791ef98a2fab786fcf06f8041bfaa06e..ff382ba3e96151686f3129eb15518ea02a291055 100644 --- a/test/test_misc.py +++ b/test/test_misc.py @@ -308,6 +308,31 @@ def test_cse_matvec(): assert np.allclose(expected_result, actual_result) +def test_diff_op_stokes(): + + from sumpy.expansion.diff_op import (make_identity_diff_op, gradient, + divergence, laplacian, concat) + from sumpy.symbolic import symbols, Function + diff_op = make_identity_diff_op(3, 4) + u = diff_op[:3] + p = diff_op[3] + pde = concat(laplacian(u) - gradient(p), divergence(u)) + + actual_output = pde.to_sym() + x, y, z = syms = symbols("x0, x1, x2") + funcs = symbols("f0, f1, f2, f3", cls=Function) + u, v, w, p = [f(*syms) for f in funcs] + + eq1 = u.diff(x, x) + u.diff(y, y) + u.diff(z, z) - p.diff(x) + eq2 = v.diff(x, x) + v.diff(y, y) + v.diff(z, z) - p.diff(y) + eq3 = w.diff(x, x) + w.diff(y, y) + w.diff(z, z) - p.diff(z) + eq4 = u.diff(x) + v.diff(y) + w.diff(z) + + expected_output = [eq1, eq2, eq3, eq4] + + assert expected_output == actual_output + + # You can test individual routines by typing # $ python test_misc.py 'test_p2p(cl.create_some_context)'