Updating fem pde code to use new manufactured api solutions. Still need to port over transient PDE fem tests

Author: jdjakem

pyapprox/pde/collocation/manufactured_solutions.py

@@ -4,6 +4,7 @@
 import copy
 
 import sympy as sp
+from pyapprox.util.backends.template import BackendMixin, Array
 from pyapprox.util.backends.numpy import NumpyMixin
 from pyapprox.util.sympy import (
     _evaluate_sp_lambda,
@@ -19,7 +20,7 @@ class ManufacturedSolution(ABC):
     def __init__(
         self,
         nvars: int,
-        bkd=NumpyMixin,
+        bkd: BackendMixin = NumpyMixin,
         oned: bool = False,
     ):
         self._bkd = bkd
@@ -53,14 +54,14 @@ def time_symbol(self):
     def all_symbols(self):
         return self.cartesian_symbols() + self.time_symbol()
 
-    def _steady_expression_to_function(self, expr):
+    def _steady_expression_to_function(self, expr: List):
         all_symbs = self.cartesian_symbols()
         expr_lambda = sp.lambdify(all_symbs, expr, "numpy")
         return partial(
             _evaluate_sp_lambda, expr_lambda, bkd=self._bkd, oned=self._oned
         )
 
-    def _steady_expression_list_to_function(self, exprs):
+    def _steady_expression_list_to_function(self, exprs: List):
         all_symbs = self.cartesian_symbols()
         expr_lambda = [sp.lambdify(all_symbs, expr, "numpy") for expr in exprs]
         return partial(
@@ -70,7 +71,7 @@ def _steady_expression_list_to_function(self, exprs):
             oned=False,
         )
 
-    def _steady_expression_list_of_lists_to_function(self, exprs):
+    def _steady_expression_list_of_lists_to_function(self, exprs: List):
         all_symbs = self.cartesian_symbols()
         expr_lambda = [
             [sp.lambdify(all_symbs, expr, "numpy") for expr in row]
@@ -93,7 +94,7 @@ def _transient_expression_to_function(self, expr):
             oned=self._oned,
         )
 
-    def _transient_expression_list_to_function(self, exprs):
+    def _transient_expression_list_to_function(self, exprs: List):
         all_symbs = self.all_symbols()
         expr_lambda = [sp.lambdify(all_symbs, expr, "numpy") for expr in exprs]
         return partial(
@@ -103,7 +104,7 @@ def _transient_expression_list_to_function(self, exprs):
             oned=False,
         )
 
-    def _transient_expression_list_of_lists_to_function(self, exprs):
+    def _transient_expression_list_of_lists_to_function(self, exprs: List):
         all_symbs = self.all_symbols()
         expr_lambda = [
             [sp.lambdify(all_symbs, expr, "numpy") for expr in row]
@@ -116,7 +117,7 @@ def _transient_expression_list_of_lists_to_function(self, exprs):
             oned=False,
         )
 
-    def is_transient(self):
+    def is_transient(self) -> bool:
         return self._bkd.any(self._bkd.array(list(self.transient.values())))
 
     def _expressions_to_functions(self):
@@ -160,7 +161,7 @@ def _expressions_to_functions(self):
                         self._transient_expression_to_function(expr)
                     )
 
-    def nvars(self):
+    def nvars(self) -> int:
         return self._nvars
 
     def _set_expression_from_bool(self, name: str, expr, transient: bool):
@@ -182,7 +183,7 @@ def _set_expression(self, name: str, expr, expr_str: str):
             transient = False
         self._set_expression_from_bool(name, expr, transient)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         fields = f"{self._expressions}".split(",")
         expr_str = ",\n".join(fields)[1:-1].replace("'", "")
         return "{0}(\n {1}\n)".format(self.__class__.__name__, expr_str)
@@ -256,7 +257,9 @@ def _sympy_reaction_expressions(self, react_str, sol_str):
         # react_prime_expr = react_prime_expr.subs(
         #     sol_symb, self._expressions["solution"]
         # )
+        print(react_str)
         react_str = react_str.replace("u", "({0})".format(sol_str))
+        print(react_str)
         return react_str, sp.sympify(react_str)  # , react_prime_expr
 
     def sympy_reaction_expressions(self):
@@ -832,3 +835,59 @@ def sympy_expressions(self):
         self._set_expression("viscosity", visc_expr, self._visc_str)
         self._set_expression("flux", flux_exprs, self._sol_str)
         self._expressions["forcing"] += forc_expr
+
+
+class ManufacturedStokes(VectorSolutionMixin, ManufacturedSolution):
+    def __init__(
+        self,
+        sol_strs: List[str],
+        nvars: int,
+        navier_stokes: bool,
+        bkd: BackendMixin = NumpyMixin,
+        oned: bool = False,
+    ):
+        self._navier_stokes = navier_stokes
+        if len(sol_strs) != nvars + 1:
+            raise ValueError(f"len(sol_strs)!={nvars+1}")
+        self._vel_strs = sol_strs[:nvars]
+        self._pres_str = sol_strs[nvars]
+        super().__init__(sol_strs, nvars, bkd, oned)
+
+    def sympy_expressions(self):
+        cartesian_symbs = self.cartesian_symbols()
+        exprs = self._expressions["solution"]
+        vel_exprs = exprs[: self._nvars]
+        pres_expr = exprs[self._nvars]
+        vel_forc_exprs = []
+        for vel, s1 in zip(vel_exprs, cartesian_symbs):
+            vel_forc_exprs.append(
+                sum([-vel.diff(s2, 2) for s2 in cartesian_symbs])
+                + pres_expr.diff(s1, 1)
+            )
+            if self._navier_stokes:
+                vel_forc_exprs[-1] += sum(
+                    [
+                        u * vel.diff(s2, 1)
+                        for u, s2 in zip(vel_exprs, cartesian_symbs)
+                    ]
+                )
+        pres_forc_expr = sum(
+            [vel.diff(s, 1) for vel, s in zip(vel_exprs, cartesian_symbs)]
+        )
+
+        vel_grad_exprs = [
+            [v.diff(s, 1) for s in cartesian_symbs] for v in vel_exprs
+        ]
+        pres_grad_expr = [pres_expr.diff(s, 1) for s in cartesian_symbs]
+        flux_exprs = vel_grad_exprs + [pres_grad_expr]
+        self._set_expression("flux", flux_exprs, self._sol_strs[0])
+
+        self._set_expression("vel_forcing", vel_forc_exprs, self._sol_strs[0])
+        self._set_expression(
+            "pres_forcing", pres_forc_expr, self._sol_strs[-1]
+        )
+
+        forc_exprs = vel_exprs + [pres_forc_expr]
+        self._expressions["forcing"] = [
+            f + g for f, g in zip(self._expressions["forcing"], forc_exprs)
+        ]

pyapprox/pde/galerkin/physics.py

@@ -630,7 +630,10 @@ def __call__(self, xx):
         # E.g. when manufactured solution, diff etc. string does not have x
         # in it. If not dependent on x then must use 1e-16*x
         if xx.ndim != vals.ndim:
-            raise RuntimeError(f"vals has the incorrect shape {vals.shape}")
+            raise ValueError(
+                f"vals.ndim is incorrect. Was {vals.ndim} "
+                f"should be {xx.ndim}"
+            )
         return vals[:, 0]
 
     def __repr__(self):
@@ -673,14 +676,23 @@ def __repr__(self):
 
 
 class FEMVectorFunctionFromCallable(FromCallableMixin, FEMVectorFunction):
+    def __init__(self, fun: callable, name: str = None, swapaxes: bool = True):
+        super().__init__(fun, name)
+        self._swapeaxes = swapaxes
+
     def __call__(self, xx):
         vals = self._fun(xx)
         # if self._fun is not implemented correctly this will fail
         # E.g. when manufactured solution, diff etc. string does not have x
         # in it. If not dependent on x then must use 1e-16*x
-        assert xx.ndim == vals.ndim
+        if xx.ndim != vals.ndim:
+            raise ValueError(
+                f"vals.ndim is incorrect. Was {vals.ndim} "
+                f"should be {xx.ndim}"
+            )
         # put vals in format required by FEM
-        vals = np.swapaxes(vals, 0, 1)
+        if self._swapeaxes:
+            vals = np.swapaxes(vals, 0, 1)
         return vals
 
 
@@ -781,6 +793,7 @@ def apply_dirichlet_boundary_conditions(
 
 
 class LinearAdvectionDiffusionReaction(AdvectionDiffusionReaction):
+    # Only to be used to get initial guess for nonlinear advection diffusion
     def __init__(
         self,
         mesh: Mesh,
@@ -793,14 +806,16 @@ def __init__(
         react_fun: Optional[FEMNonLinearOperator] = None,
     ):
         if not isinstance(diff_fun, FEMScalarFunction):
-            raise ValueError(
-                "diff_fun must be an instance of FEMScalarFunction"
+            raise TypeError(
+                "diff_fun must be an instance of FEMScalarFunction "
+                f"but was {diff_fun.__class__.__name__}"
             )
         if react_fun is not None and not isinstance(
-            react_fun, FEMVectorFunction
+            react_fun, FEMScalarFunction
         ):
-            raise ValueError(
-                "react_fun must be an instance of FEMVectorFunction"
+            raise TypeError(
+                "react_fun must be an instance of FEMScalarFunction "
+                f"but was {react_fun.__class__.__name__}"
             )
         self._diff_fun = diff_fun
         self._react_fun = react_fun
@@ -816,9 +831,7 @@ def _set_funs(self) -> List:
         return funs + [self._react_fun, self._forc_fun]
 
     def init_guess(self) -> np.ndarray:
-        return np.ones(
-            (self.basis.N),
-        )
+        return np.ones((self.basis.N,))
 
 
 class NonLinearAdvectionDiffusionReaction(AdvectionDiffusionReaction):
@@ -867,7 +880,8 @@ def _set_funs(self) -> List:
         return funs + [self._linear_diff_fun, self._diff_op, self._react_op]
 
 
-class Helmholtz(LinearAdvectionDiffusionReaction):
+# class Helmholtz(LinearAdvectionDiffusionReaction):
+class Helmholtz(NonLinearAdvectionDiffusionReaction):
     def __init__(
         self,
         mesh: Mesh,
@@ -878,34 +892,46 @@ def __init__(
         forc_fun: Optional[Callable[[np.ndarray], np.ndarray]] = None,
     ):
 
+        if len(bndry_conds[1]) > 0 or len(bndry_conds[1]) > 0:
+            raise NotImplementedError(
+                "Currently tests do not pass with Robin or Neumann BCs"
+            )
         if forc_fun is None:
-            forc_fun = self._zero_forcing
+            forc_fun = FEMScalarFunctionFromCallable(
+                self._zero_forcing, "forcing"
+            )
 
         super().__init__(
             mesh,
             element,
             basis,
             bndry_conds,
-            self._unit_diff,
+            FEMScalarFunctionFromCallable(self._unit_diff, "diffusion"),
             forc_fun,
-            self._zero_vel,
+            FEMVectorFunctionFromCallable(self._zero_vel, "velocity", False),
+            FEMNonLinearOperatorFromCallable(
+                lambda x, u: 0 * u, lambda x, u: 0 * u
+            ),
             wave_number_fun,
         )
 
+    def init_guess(self) -> np.ndarray:
+        return np.ones((self.basis.N,))
+
     def _unit_diff(self, x):
         # negative is taken because advection diffusion code solves
         # -k*\nabla^2 u + c*u = f
         # but Helmholtz solves
         # \nabla^2 u + c*u = 0 which is equivalent when k=-1
-        return -1 + 0 * x[0]
+        return -1.0 + 0.0 * x
 
     def _zero_vel(self, x):
         # Helmholtz is a special case of the advection diffusion reaction
         # equation when there is no advection, i.e. velocity field is zero
-        return 0 * x
+        return 0.0 * x
 
     def _zero_forcing(self, x):
-        return 0 * x[0]
+        return 0.0 * x
 
 
 class Stokes(Physics):
@@ -920,6 +946,10 @@ def __init__(
         pres_forc_fun: FEMScalarFunction,
         viscosity: float = 1.0,
     ):
+        if len(bndry_conds[1]) > 0 or len(bndry_conds[1]) > 0:
+            raise NotImplementedError(
+                "Currently tests do not pass with Robin or Neumann BCs"
+            )
         super().__init__(mesh, element, basis, bndry_conds)
 
         # if not isinstance(vel_forc_fun, FEMVectorFunction):