Merge branch 'main' into feature/asymline

Author: ShawnL00

.basedpyright/baseline.json

@@ -111,7 +111,7 @@ The FAQ is maintained collaboratively on the
 Acknowledgments
 ===============
 
-Work on meshmode was supported in part by
+Work on sumpy was supported in part by
 
 * the US National Science Foundation under grant numbers DMS-1418961,
   DMS-1654756, SHF-1911019, and OAC-1931577.

doc/misc.rst

@@ -25,6 +25,7 @@
 
 import logging
 from dataclasses import dataclass
+from functools import cached_property
 from itertools import accumulate
 from typing import TYPE_CHECKING, TypeAlias
 
@@ -62,6 +63,7 @@
 .. autoclass:: DerivativeIdentifier
 .. autofunction:: make_identity_diff_op
 .. autofunction:: as_scalar_pde
+.. autofunction:: to_fourier_matrix
 """
 
 
@@ -92,37 +94,43 @@ class DerivativeIdentifier:
 class LinearPDESystemOperator:
     r"""
     Represents a constant-coefficient linear differential operator of a
-    vector-valued function with `dim` spatial variables.
+    vector-valued function with :attr:`spatial_dim` spatial variables and
+    additional temporal variables.
 
-    It is represented by a tuple of immutable dictionaries. The dictionary maps a
-    :class:`DerivativeIdentifier` to the coefficient. Optionally supports a
-    time variable as the last variable in the multi-index of the
-    :class:`DerivativeIdentifier`.
+    The operator is given by a tuple of immutable dictionaries. The dictionary
+    maps a :class:`DerivativeIdentifier` to a (time- and space-independent)
+    coefficient. In the :class:`DerivativeIdentifier`, each multi-index has
+    :attr:`spatial_dim` indices for the spatial variables and the remaining
+    ones represent temporal variables.
+
+    The class also supports basic arithmetic, i.e. multiplication and addition
+    with other operators and constants.
 
-    .. autoattribute:: dim
     .. autoattribute:: eqs
+    .. autoattribute:: spatial_dim
+    .. autoproperty:: total_dims
 
     .. autoproperty:: order
-    .. autoproperty:: total_dims
+    .. autoproperty:: nvariables
+    .. autoproperty:: is_time_dependent
+
     .. automethod:: to_sym
     """
 
-    dim: int
+    spatial_dim: int
+    """The number of spatial dimensions of the PDE (use :attr:`total_dims`
+    to include time).
+    """
     eqs: tuple[Mapping[DerivativeIdentifier, sym.Expr], ...]
+    """A tuple of all the equations in the system."""
 
     if __debug__:
 
         def __post_init__(self) -> None:
             # NOTE: this will raise a TypeError if it's not hashable
             _ = hash(self)
 
-    @property
-    def order(self) -> int:
-        deg = 0
-        for eq in self.eqs:
-            deg = max(deg, max(sum(ident.mi) for ident in eq))
-
-        return deg
+    # {{{ arithmetic
 
     def __mul__(self, other: Number | sym.Expr) -> LinearPDESystemOperator:
         import numbers
@@ -137,7 +145,7 @@ def __mul__(self, other: Number | sym.Expr) -> LinearPDESystemOperator:
 
             eqs.append(constantdict(deriv_ident_to_coeff))
 
-        return LinearPDESystemOperator(self.dim, tuple(eqs))
+        return LinearPDESystemOperator(self.spatial_dim, tuple(eqs))
 
     def __rmul__(self, param: Number | sym.Expr) -> LinearPDESystemOperator:
         return self.__mul__(param)
@@ -146,7 +154,7 @@ def __add__(self, other: LinearPDESystemOperator) -> LinearPDESystemOperator:
         if not isinstance(other, LinearPDESystemOperator):
             return NotImplemented
 
-        assert self.dim == other.dim
+        assert self.spatial_dim == other.spatial_dim
         assert len(self.eqs) == len(other.eqs)
 
         eqs: list[Mapping[DerivativeIdentifier, sym.Expr]] = []
@@ -159,7 +167,7 @@ def __add__(self, other: LinearPDESystemOperator) -> LinearPDESystemOperator:
                     res[k] = v
             eqs.append(constantdict(res))
 
-        return LinearPDESystemOperator(self.dim, tuple(eqs))
+        return LinearPDESystemOperator(self.spatial_dim, tuple(eqs))
 
     def __radd__(self, other: LinearPDESystemOperator) -> LinearPDESystemOperator:
         return self.__add__(other)
@@ -170,39 +178,79 @@ def __sub__(self, other: LinearPDESystemOperator) -> LinearPDESystemOperator:
     def __neg__(self) -> LinearPDESystemOperator:
         return (-1) * self
 
+    # }}}
+
     @override
     def __repr__(self) -> str:
-        return f"LinearPDESystemOperator({self.dim}, {self.eqs!r})"
+        return f"LinearPDESystemOperator({self.spatial_dim}, {self.eqs!r})"
 
     def __getitem__(self, idx: int | slice) -> LinearPDESystemOperator:
         item = self.eqs.__getitem__(idx)
         eqs = item if isinstance(item, tuple) else (item,)
-        return LinearPDESystemOperator(self.dim, eqs)
+        return LinearPDESystemOperator(self.spatial_dim, eqs)
 
     @property
+    def is_time_dependent(self) -> bool:
+        """Is *True* if the PDE operator has a time component."""
+        return self.spatial_dim != self.total_dims
+
+    @cached_property
+    def order(self) -> int:
+        """The order of the PDE operator (maximum order of all derivatives)."""
+        deg = 0
+        for eq in self.eqs:
+            deg = max(deg, max(sum(ident.mi) for ident in eq))
+
+        return deg
+
+    @cached_property
     def total_dims(self) -> int:
-        """
-        Returns the total number of dimensions including time
-        """
-        did = next(iter(self.eqs[0].keys()))
+        """The total number of dimensions (including time)."""
+        did = next(iter(self.eqs[0]))
         return len(did.mi)
 
-    def to_sym(self, fnames: Sequence[str] | None = None) -> list[sym.Expr]:
-        x: list[sym.Expr] = list(sym.make_sym_vector("x", self.dim))
-        x.extend(sym.make_sym_vector("t", self.total_dims - self.dim))
+    @cached_property
+    def nvariables(self) -> int:
+        """Number of variables in the system."""
+        max_vec_idx = max((did.vec_idx for eq in self.eqs for did in eq), default=-1)
+        return max_vec_idx + 1
+
+    def to_sym(
+            self,
+            fnames: Sequence[str] | None = None,
+            *,
+            x_var_name: str = "x",
+            t_var_name: str = "t",
+        ) -> list[sym.Expr]:
+        """Transform the system to a list of :mod:`sympy` expressions.
+
+        :arg fnames: the names of the variables in the system.
+            (defaults to `["f0", "f1", ....]`)
+        :arg x_var_name: the name of the spatial variables.
+            (defaults to `["x0", "x1", ....]`)
+        :arg t_var_name: the name of the temporal variables.
+        """
+        x: list[sym.Expr] = list(sym.make_sym_vector(x_var_name, self.spatial_dim))
+        x.extend(sym.make_sym_vector(t_var_name, self.total_dims - self.spatial_dim))
 
         if fnames is None:
             noutputs = 0
             for eq in self.eqs:
                 for deriv_ident in eq:
                     noutputs = max(noutputs, deriv_ident.vec_idx)
+
             fnames = [f"f{i}" for i in range(noutputs+1)]
 
         funcs = [sym.Function(fname)(*x) for fname in fnames]
+        if len(funcs) < self.nvariables:
+            raise ValueError(
+                f"'fnames' does not match system: {len(fnames)} names "
+                f"(for a system of {self.nvariables} variables)"
+            )
 
         res: list[sym.Expr] = []
         for eq in self.eqs:
-            sym_eq: sym.Expr = sym.sympify(0)
+            sym_eq: sym.Expr = sym.Integer(0)
             for deriv_ident, coeff in eq.items():
                 expr = funcs[deriv_ident.vec_idx]
                 for i, val in enumerate(deriv_ident.mi):
@@ -237,8 +285,8 @@ def _get_all_scalar_pdes(pde: LinearPDESystemOperator) -> list[LinearPDESystemOp
     import sympy as sp
     from sympy.polys.orderings import grevlex
 
-    gens = [sp.Symbol(f"_x{i}") for i in range(pde.dim)]
-    gens += [sp.Symbol(f"_t{i}") for i in range(pde.total_dims - pde.dim)]
+    gens = [sp.Symbol(f"_x{i}") for i in range(pde.spatial_dim)]
+    gens += [sp.Symbol(f"_t{i}") for i in range(pde.total_dims - pde.spatial_dim)]
 
     max_vec_idx = max(deriv_ident.vec_idx for eq in pde.eqs
                       for deriv_ident in eq)
@@ -303,7 +351,8 @@ def intersect(a: SubModulePolyRing, b: SubModulePolyRing) -> SubModulePolyRing:
             for (mi, coeff) in zip(scalar_pde.monoms(),
                                    scalar_pde.coeffs(), strict=True)
         }
-        results.append(LinearPDESystemOperator(pde.dim, (constantdict(pde_dict),)))
+        results.append(LinearPDESystemOperator(pde.spatial_dim,
+                                               (constantdict(pde_dict),)))
 
     return results
 
@@ -382,8 +431,44 @@ def as_scalar_pde(
     return _get_all_scalar_pdes(pde)[comp_idx]
 
 
+def to_fourier_matrix(
+        pde: LinearPDESystemOperator,
+        ks: sym.Matrix,
+    ) -> sym.Matrix:
+    r"""Return the Fourier (symbol) matrix of a constant-coefficient PDE system.
+
+    Each spatial derivative :math:`\partial / \partial x_j` is replaced by
+    multiplication by :math:`i\,k_j`.  The result is a
+    :obj:`sympy.matrices.dense.Matrix` whose ``(row, col)`` entry is the
+    polynomial in the frequency variables contributed by equation *row* acting
+    on component *col*.
+
+    :returns: a matrix of size ``(len(pde.eqs), nvariables)``.
+    """
+    if pde.is_time_dependent:
+        raise ValueError("cannot compute Fourier symbol for time-dependent PDEs")
+
+    ncols = pde.nvariables
+
+    mat = []
+    for eq in pde.eqs:
+        row = [sym.Integer(0)] * ncols
+
+        for deriv_ident, coeff in eq.items():
+            factor: sym.Expr = sym.Integer(1)
+            for j, power in enumerate(deriv_ident.mi):
+                factor *= (sym.I * ks[j]) ** power
+
+            row[deriv_ident.vec_idx] += coeff * factor
+
+        assert len(row) == ncols
+        mat.append(row)
+
+    return sym.Matrix(mat)
+
+
 def laplacian(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
-    dim = diff_op.dim
+    dim = diff_op.spatial_dim
     empty: tuple[Mapping[DerivativeIdentifier, sym.Expr], ...] = (
         (constantdict(),) * len(diff_op.eqs))
 
@@ -409,17 +494,17 @@ def diff(
 
         eqs.append(constantdict(res))
 
-    return LinearPDESystemOperator(diff_op.dim, tuple(eqs))
+    return LinearPDESystemOperator(diff_op.spatial_dim, tuple(eqs))
 
 
 def divergence(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
-    if len(diff_op.eqs) != diff_op.dim:
+    if len(diff_op.eqs) != diff_op.spatial_dim:
         raise ValueError(
             "number of equations does not match system dimension: "
-            f"got {len(diff_op.eqs)} equations for {diff_op.dim}d system")
+            f"got {len(diff_op.eqs)} equations for {diff_op.spatial_dim}d system")
 
-    res = LinearPDESystemOperator(diff_op.dim, (constantdict(),))
-    for i in range(diff_op.dim):
+    res = LinearPDESystemOperator(diff_op.spatial_dim, (constantdict(),))
+    for i in range(diff_op.spatial_dim):
         mi = [0]*diff_op.total_dims
         mi[i] = 1
         res += diff(diff_op[i], tuple(mi))
@@ -433,7 +518,7 @@ def gradient(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
             f"can only take gradient of scalar system: got {len(diff_op.eqs)}d")
 
     eqs: list[Mapping[DerivativeIdentifier, sym.Expr]] = []
-    dim = diff_op.dim
+    dim = diff_op.spatial_dim
     for i in range(dim):
         mi = [0]*diff_op.total_dims
         mi[i] = 1
@@ -443,13 +528,13 @@ def gradient(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
 
 
 def curl(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
-    if len(diff_op.eqs) != diff_op.dim:
+    if len(diff_op.eqs) != diff_op.spatial_dim:
         raise ValueError(
             "number of equations does not match system dimension: "
-            f"got {len(diff_op.eqs)} equations for {diff_op.dim}d system")
+            f"got {len(diff_op.eqs)} equations for {diff_op.spatial_dim}d system")
 
-    if diff_op.dim != 3:
-        raise ValueError(f"can only take curl of 3d system: got {diff_op.dim}d")
+    if diff_op.spatial_dim != 3:
+        raise ValueError(f"can only take curl of 3d system: got {diff_op.spatial_dim}d")
 
     eqs: list[Mapping[DerivativeIdentifier, sym.Expr]] = []
     mis: list[MultiIndex] = []
@@ -464,7 +549,7 @@ def curl(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
             - diff(diff_op[(i+1) % 3], mis[(i+2) % 3]))
         eqs.append(new_pde.eqs[0])
 
-    return LinearPDESystemOperator(diff_op.dim, tuple(eqs))
+    return LinearPDESystemOperator(diff_op.spatial_dim, tuple(eqs))
 
 
 def concat(*ops: LinearPDESystemOperator) -> LinearPDESystemOperator:
@@ -474,8 +559,8 @@ def concat(*ops: LinearPDESystemOperator) -> LinearPDESystemOperator:
     if len(ops) == 1:
         return ops[0]
 
-    dim = ops[0].dim
-    if not all(op.dim == dim for op in ops):
+    dim = ops[0].spatial_dim
+    if not all(op.spatial_dim == dim for op in ops):
         raise ValueError(f"operators must have the same dimension (expected {dim}d)")
 
     eqs = list(ops[0].eqs)