From d708adc4b01c1595edaa1379b6ac0973300a274e Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Sat, 12 Apr 2025 18:39:45 +0800 Subject: [PATCH 1/4] Diffusion updates solver updates -Add a diffusion solver -symbolic ddt -Eulerian ddt CM updates -Anisotropic diffusion CM -Flux CM, allowing flux terms to be passed directly into the solver --- src/underworld3/constitutive_models.py | 139 ++++++ src/underworld3/systems/__init__.py | 4 + src/underworld3/systems/ddt.py | 424 +++++++++++++++++++ src/underworld3/systems/diffusion_solvers.py | 327 ++++++++++++++ 4 files changed, 894 insertions(+) create mode 100644 src/underworld3/systems/diffusion_solvers.py diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index eb704b4d1..f8f89904a 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -1465,6 +1465,145 @@ def _object_viewer(self): return +class AnisotropicDiffusionModel(DiffusionModel): + """ + Anisotropic version of DiffusionModel supporting diagonal diffusion tensors. + + Diffusivity must be provided as a `sympy.Matrix` vector of length `dim`: + - 2D example: `Matrix([Dx, Dy])` + - 3D example: `Matrix([Dx, Dy, Dz])` + + This vector is interpreted as the diagonal of the diffusivity tensor: + κ = diag(Dx, Dy, Dz) + + Usage: + ```python + from sympy import symbols, Matrix + kappa = symbols("kappa") + model = AnisotropicDiffusionModel + model.material_properties = model.Parameters(diffusivity=sp.Matrix([kappa, 2*kappa])) + ``` + """ + + class _Parameters(DiffusionModel._Parameters): + def __init__(inner_self, _owning_model): + dim = _owning_model.dim + inner_self._owning_model = _owning_model + inner_self._diffusivity = sympy.eye(dim) + + @property + def diffusivity(inner_self): + return inner_self._diffusivity + + @diffusivity.setter + def diffusivity(inner_self, value): + dim = inner_self._owning_model.dim + + if not isinstance(value, sympy.Matrix): + raise TypeError("Diffusivity must be a sympy.Matrix vector.") + + # Accept shape (dim, 1) or (1, dim) + if value.shape not in [(dim, 1), (1, dim)]: + raise ValueError( + f"Diffusivity must be a vector of length {dim}. " + f"Got shape {value.shape}." + ) + + # Flatten and validate + elements = [value[i] for i in range(dim)] + validated = [] + for i, v in enumerate(elements): + diff = validate_parameters( + fr"\upkappa_{{{i}}}", v, f"Diffusivity in x_{i}", allow_number=True + ) + if diff is not None: + validated.append(diff) + + inner_self._diffusivity = sympy.diag(*validated) + inner_self._reset() + + def _build_c_tensor(self): + """Constructs the anisotropic (diagonal) diffusivity tensor.""" + self._c = self.Parameters.diffusivity + + def _object_viewer(self): + from IPython.display import Latex, display + super()._object_viewer() + + # Extract diagonal entries and render cleanly + diagonal = self.Parameters.diffusivity.diagonal() + latex_entries = ", ".join([sympy.latex(k) for k in diagonal]) + kappa_latex = r"\kappa = \mathrm{diag}\left(" + latex_entries + r"\right)" + + display(Latex(r"$\quad " + kappa_latex + r"$")) + + +class GenericFluxModel(Constitutive_Model): + r""" + A generic constitutive model with symbolic flux expression. + + Example usage: + ```python + grad_phi = sympy.Matrix([sp.Symbol("∂φ/∂x"), sp.Symbol("∂φ/∂y")]) + flux_expr = sympy.Matrix([[kappa_11, kappa_12], [kappa_21, kappa_22]]) * grad_phi + + model = GenericFluxModel(dim=2) + model.flux = flux_expr + scalar_solver.constititutive_model = model + ``` + """ + + class _Parameters: + def __init__(inner_self, _owning_model): + inner_self._flux = sympy.Matrix(sympy.ones(_owning_model.dim, 1)) + inner_self._owning_model = _owning_model + + @property + def flux(inner_self): + return inner_self._flux + + @flux.setter + def flux(inner_self, value): + dim = inner_self._owning_model.dim + + if not isinstance(value, sympy.Matrix): + raise TypeError("Flux must be a sympy.Matrix vector.") + + # Accept shape (dim, 1) or (1, dim) + if value.shape not in [(dim, 1), (1, dim)]: + raise ValueError( + f"Flux must be a symbolic vector of length {dim}. " + f"Got shape {value.shape}." + ) + + # Flatten and validate + elements = [value[i] for i in range(dim)] + validated = [] + for i, v in enumerate(elements): + flux_component = validate_parameters( + fr"q_{{{i}}}", v, f"Flux component in x_{i}", allow_number=True + ) + if flux_component is not None: + validated.append(flux_component) + + inner_self._flux = sympy.Matrix(validated).reshape(dim, 1) + inner_self._reset() + + @property + def flux(self): + # if self._flux is None: + # raise RuntimeError("Flux expression has not been set.") + return self.Parameters.flux + + + def _object_viewer(self): + from IPython.display import display, Latex + super()._object_viewer() + if self.flux is not None: + display(Latex(r"$\vec{q} = " + sympy.latex(self.flux) + "$")) + else: + display(Latex(r"No flux expression set.")) + class DarcyFlowModel(Constitutive_Model): r""" diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index 282a37cb4..5e16a9fca 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -33,3 +33,7 @@ from .ddt import Lagrangian as Lagrangian_DDt from .ddt import SemiLagrangian as SemiLagragian_DDt from .ddt import Lagrangian_Swarm as Lagrangian_Swarm_DDt +from .ddt import Eulerian as Eulerian_DDt + +# import additional diffusion solvers +from .diffusion_solvers import SNES_Diffusion as Diffusion diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index ba1ec3343..38fb84911 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -18,6 +18,430 @@ # class Eulerian(uw_object): # etc etc... +class Symbolic(uw_object): + r""" + Symbolic History Manager: + + This class manages the update of a variable ψ across timesteps. + The history operator stores ψ over several timesteps (given by 'order') + so that it can compute backward differentiation (BDF) or Adams–Moulton expressions. + + The history operator is defined as follows: + $$\quad \psi_p^{t-n\Delta t} \leftarrow \psi_p^{t-(n-1)\Delta t}\quad$$ + $$\quad \psi_p^{t-(n-1)\Delta t} \leftarrow \psi_p^{t-(n-2)\Delta t} \cdots\quad$$ + $$\quad \psi_p^{t-\Delta t} \leftarrow \psi_p^{t}$$ + + + """ + + @timing.routine_timer_decorator + def __init__( + self, + psi_fn: sympy.Basic, # a sympy expression for ψ; can be scalar or matrix + theta: Optional[float] = 0.5, + varsymbol: Optional[str] = r"u", + verbose: Optional[bool] = False, + bcs=[], + order: int = 1, + smoothing: float = 0.0, + ): + super().__init__() + self.theta = theta + self.bcs = bcs + self.verbose = verbose + self.smoothing = smoothing + self.order = order + + # Ensure psi_fn is a sympy Matrix. + if not isinstance(psi_fn, sympy.Matrix): + try: + # If psi_fn has a shape attribute, convert it. + psi_fn = sympy.Matrix(psi_fn) + except Exception: + psi_fn = sympy.Matrix([[psi_fn]]) + + # Do not force a column vector if psi_fn has more than one column. + # Instead, preserve its full shape. + self._psi_fn = psi_fn # stored with its native shape + self._shape = psi_fn.shape # capture the shape to create history storage + + # Create the history list: each element is a Matrix of shape _shape. + self.psi_star = [sympy.zeros(*self._shape) for _ in range(order)] + self.initiate_history_fn() + return + + @property + def psi_fn(self): + return self._psi_fn + + @psi_fn.setter + def psi_fn(self, new_fn): + if not isinstance(new_fn, sympy.Matrix): + try: + new_fn = sympy.Matrix(new_fn) + except Exception: + new_fn = sympy.Matrix([[new_fn]]) + # Optionally, one could check for matching shape; here we update both. + self._psi_fn = new_fn + self._shape = new_fn.shape + return + + def _object_viewer(self): + from IPython.display import Latex, display + display(Latex(rf"$\quad$History steps = {self.order}$")) + + def update_history_fn(self): + # Update the first history element with a copy of the current ψ. + self.psi_star[0] = self.psi_fn.copy() + + def initiate_history_fn(self): + self.update_history_fn() + # Propagate the initial history to all history steps. + for i in range(1, self.order): + self.psi_star[i] = self.psi_star[0].copy() + return + + def update( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + self.update_pre_solve(evalf, verbose) + return + + def update_pre_solve( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + # Default: no action. + return + + def update_post_solve( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + if verbose: + print(f"Updating history for ψ = {self.psi_fn}", flush=True) + # Shift history: copy each element down the chain. + for i in range(self.order - 1, 0, -1): + self.psi_star[i] = self.psi_star[i - 1].copy() + self.update_history_fn() + return + + def bdf(self, order: Optional[int] = None): + r"""Compute the backward differentiation approximation of the time-derivative of ψ. + For order 1: bdf ≡ ψ - psi_star[0] + """ + if order is None: + order = self.order + else: + order = max(1, min(self.order, order)) + + with sympy.core.evaluate(False): + if order == 1: + bdf0 = self.psi_fn - self.psi_star[0] + elif order == 2: + bdf0 = (3 * self.psi_fn / 2 - 2 * self.psi_star[0] + self.psi_star[1] / 2) + elif order == 3: + bdf0 = (11 * self.psi_fn / 6 - 3 * self.psi_star[0] + (3 * self.psi_star[1]) / 2 - self.psi_star[2] / 3) + return bdf0 + + def adams_moulton_flux(self, order: Optional[int] = None): + if order is None: + order = self.order + else: + order = max(1, min(self.order, order)) + + with sympy.core.evaluate(False): + if order == 1: + am = self.theta * self.psi_fn + (1.0 - self.theta) * self.psi_star[0] + elif order == 2: + am = (5 * self.psi_fn + 8 * self.psi_star[0] - self.psi_star[1]) / 12 + elif order == 3: + am = (9 * self.psi_fn + 19 * self.psi_star[0] - 5 * self.psi_star[1] + self.psi_star[2]) / 24 + return am + +class Eulerian(uw_object): + r"""Eulerian (mesh based) History Manager: + This manages the update of a variable, $\psi$ on the mesh across timesteps. + $$\quad \psi_p^{t-n\Delta t} \leftarrow \psi_p^{t-(n-1)\Delta t}\quad$$ + $$\quad \psi_p^{t-(n-1)\Delta t} \leftarrow \psi_p^{t-(n-2)\Delta t} \cdots\quad$$ + $$\quad \psi_p^{t-\Delta t} \leftarrow \psi_p^{t}$$ + """ + + @timing.routine_timer_decorator + def __init__( + self, + mesh: uw.discretisation.Mesh, + psi_fn: Union[ + uw.discretisation.MeshVariable, sympy.Basic + ], # sympy function or mesh variable + vtype: uw.VarType, + degree: int, + continuous: bool, + evalf: Optional[bool] = False, + theta: Optional[float] = 0.5, + varsymbol: Optional[str] = r"u", + verbose: Optional[bool] = False, + bcs=[], + order=1, + smoothing=0.0, + ): + super().__init__() + + self.mesh = mesh + self.theta = theta + self.bcs = bcs + self.verbose = verbose + self.degree = degree + self.vtype = vtype + self.continuous = continuous + self.smoothing = smoothing + self.evalf = evalf + + # meshVariables are required for: + # + # u(t) - evaluation of u_fn at the current time + # u*(t) - u_* evaluated from + + # psi is evaluated/stored at `order` timesteps. We can't + # be sure if psi is a meshVariable or a function to be evaluated + # psi_star is reaching back through each evaluation and has to be a + # meshVariable (storage) + + if isinstance(psi_fn, uw.discretisation._MeshVariable): + self._psi_fn = psi_fn.sym ### get symbolic form of the meshvariable + self._psi_meshVar = psi_fn + else: + self._psi_fn = psi_fn ### already in symbolic form + self._psi_meshVar = None + + self.order = order + + psi_star = [] + self.psi_star = psi_star + + + for i in range(order): + self.psi_star.append( + uw.discretisation.MeshVariable( + f"psi_star_sl_{self.instance_number}_{i}", + self.mesh, + vtype=vtype, + degree=degree, + continuous=continuous, + varsymbol=rf"{varsymbol}^{{ {'*'*(i+1)} }}", + ) + ) + + # print('initiating history fn', flush=True) + ### Initiate first history value in chain + self.initiate_history_fn() + + + return + + @property + def psi_fn(self): + return self._psi_fn + + @psi_fn.setter + def psi_fn(self, new_fn): + self._psi_fn = new_fn + # self._psi_star_projection_solver.uw_function = self.psi_fn + return + + def _object_viewer(self): + from IPython.display import Latex, Markdown, display + + super()._object_viewer() + + ## feedback on this instance + # display(Latex(r"$\quad\psi = $ " + self.psi._repr_latex_())) + # display( + # Latex( + # r"$\quad\Delta t_{\textrm{phys}} = $ " + # + sympy.sympify(self.dt_physical)._repr_latex_() + # ) + # ) + display(Latex(rf"$\quad$History steps = {self.order}")) + + + def _setup_projections(self): + ### using this to store terms that can't be evaluated (e.g. derivatives) + # The projection operator for mapping derivative values to the mesh - needs to be different for each variable type, unfortunately ... + if self.vtype == uw.VarType.SCALAR: + self._psi_star_projection_solver = uw.systems.solvers.SNES_Projection( + self.mesh, self.psi_star[0], verbose=False + ) + elif self.vtype == uw.VarType.VECTOR: + self._psi_star_projection_solver = ( + uw.systems.solvers.SNES_Vector_Projection( + self.mesh, self.psi_star[0], verbose=False + ) + ) + elif self.vtype == uw.VarType.SYM_TENSOR or self.vtype == uw.VarType.TENSOR: + self._WorkVar = uw.discretisation.MeshVariable( + f"W_star_slcn_{self.instance_number}", + self.mesh, + vtype=uw.VarType.SCALAR, + degree=self.degree, + continuous=self.continuous, + varsymbol=r"W^{*}", + ) + self._psi_star_projection_solver = ( + uw.systems.solvers.SNES_Tensor_Projection( + self.mesh, self.psi_star[0], self._WorkVar, verbose=False + ) + ) + + self._psi_star_projection_solver.uw_function = self.psi_fn + self._psi_star_projection_solver.bcs = self.bcs + self._psi_star_projection_solver.smoothing = self.smoothing + + def update_history_fn(self): + ### update first value in history chain + ### avoids projecting if function can be evaluated + try: + with self.mesh.access(self.psi_star[0]): + try: + with self.mesh.access(self._psi_meshVar): + self.psi_star[0].data[...] = self._psi_meshVar.data[...] + # print('copying data', flush=True) + except: + if self.evalf: + self.psi_star[0].data[...] = uw.function.evalf(self.psi_fn, self.psi_star[0].coords).reshape(-1, max(self.psi_fn.shape)) + # print('evalf data', flush=True) + else: + self.psi_star[0].data[...] = uw.function.evaluate(self.psi_fn, self.psi_star[0].coords).reshape(-1, max(self.psi_fn.shape)) + # print('evaluate data', flush=True) + + except: + self._setup_projections() + self._psi_star_projection_solver.solve() + # print('projecting data', flush=True) + + def initiate_history_fn(self): + self.update_history_fn() + + ### set up all history terms to the initial values + for i in range(self.order - 1, 0, -1): + with self.mesh.access(self.psi_star[i]): + self.psi_star[i].data[...] = self.psi_star[0].data[...] + + return + + def update( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + self.update_pre_solve(evalf, verbose) + return + + def update_pre_solve( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + return + + def update_post_solve( + self, + evalf: Optional[bool] = False, + verbose: Optional[bool] = False, + ): + # if average_over_dt: + # phi = min(1.0, dt / self.dt_physical) + # else: + # phi = 1.0 + + if verbose and uw.mpi.rank == 0: + print(f"Update {self.psi_fn}", flush=True) + + + ### copy values down the chain + for i in range(self.order - 1, 0, -1): + with self.mesh.access(self.psi_star[i]): + self.psi_star[i].data[...] = self.psi_star[i-1].data[...] + + ### update the history fn + self.update_history_fn() + # ### update the first value in the chain + # if self.evaluateable: + # with self.mesh.access(self.psi_star[0]): + # try: + # with self.mesh.access(self._psi_meshVar): + # self.psi_star[0].data[...] = self._psi_meshVar.data[...] + # except: + # if self.evalf: + # self.psi_star[0].data[...] = uw.function.evalf(self.psi_fn, self.psi_star[0].coords).reshape(-1, max(self.psi_fn.shape)) + # else: + # self.psi_star[0].data[...] = uw.function.evaluate(self.psi_fn, self.psi_star[0].coords).reshape(-1, max(self.psi_fn.shape)) + # else: + # self._psi_star_projection_solver.uw_function = self.psi_fn + # self._psi_star_projection_solver.solve() + + return + + def bdf(self, order=None): + r"""Backwards differentiation form for calculating DuDt + Note that you will need `bdf` / $\delta t$ in computing derivatives""" + + if order is None: + order = self.order + else: + order = max(1, min(self.order, order)) + + with sympy.core.evaluate(False): + if order == 1: + bdf0 = self.psi_fn - self.psi_star[0].sym + + elif order == 2: + bdf0 = ( + 3 * self.psi_fn / 2 + - 2 * self.psi_star[0].sym + + self.psi_star[1].sym / 2 + ) + + elif order == 3: + bdf0 = ( + 11 * self.psi_fn / 6 + - 3 * self.psi_star[0].sym + + 3 * self.psi_star[1].sym / 2 + - self.psi_star[2].sym / 3 + ) + + return bdf0 + + def adams_moulton_flux(self, order=None): + if order is None: + order = self.order + else: + order = max(1, min(self.order, order)) + + with sympy.core.evaluate(False): + if order == 1: + # am = (self.psi_fn + self.psi_star[0].sym) / 2 + am = self.theta*self.psi_fn + ((1.-self.theta)*self.psi_star[0].sym) + + elif order == 2: + am = ( + 5 * self.psi_fn + 8 * self.psi_star[0].sym - self.psi_star[1].sym + ) / 12 + + elif order == 3: + am = ( + 9 * self.psi_fn + + 19 * self.psi_star[0].sym + - 5 * self.psi_star[1].sym + + self.psi_star[2].sym + ) / 24 + + return am + class SemiLagrangian(uw_object): r"""Nodal-Swarm Lagrangian History Manager: diff --git a/src/underworld3/systems/diffusion_solvers.py b/src/underworld3/systems/diffusion_solvers.py new file mode 100644 index 000000000..d5475ba71 --- /dev/null +++ b/src/underworld3/systems/diffusion_solvers.py @@ -0,0 +1,327 @@ +import sympy +from sympy import sympify +import numpy as np + +from typing import Optional, Callable, Union + +import underworld3 as uw +from underworld3.systems import SNES_Scalar, SNES_Vector, SNES_Stokes_SaddlePt +from underworld3 import VarType +import underworld3.timing as timing +from underworld3.utilities._api_tools import uw_object + + +from .ddt import SemiLagrangian as SemiLagrangian_DDt +from .ddt import Lagrangian as Lagrangian_DDt +from .ddt import Eulerian as Eulerian_DDt +from .ddt import Symbolic as Symbolic_DDt + +import sympy +import numpy as np +from typing import Optional, Union +import underworld3 as uw +from underworld3.systems import SNES_Vector, SNES_Scalar +from underworld3 import VarType +import underworld3.timing as timing +from underworld3.utilities._api_tools import uw_object + +from .ddt import SemiLagrangian as SemiLagrangian_DDt +from .ddt import Lagrangian as Lagrangian_DDt +from .ddt import Eulerian as Eulerian_DDt +from .ddt import Symbolic as Symbolic_DDt + + +class SNES_Diffusion(SNES_Scalar): + r""" + This class provides a solver for the scalar Diffusion equation using mesh-based finite elements. + + $$ + \color{Green}{\underbrace{ \Bigl[ \frac{\partial u}{\partial t} - \left( \mathbf{v} \cdot \nabla \right) u \Bigr]}_{\dot{\mathbf{f}}}} - + \nabla \cdot + \color{Blue}{\underbrace{\Bigl[ \boldsymbol\kappa \nabla u \Bigr]}_{\mathbf{F}}} = + \color{Maroon}{\underbrace{\Bigl[ f \Bigl] }_{\mathbf{f}}} + $$ + + The term $\mathbf{F}$ relates diffusive fluxes to gradients in the unknown $u$. + + ## Properties + + - The unknown is $u$. + + - The diffusivity tensor, $\kappa$ is provided by setting the `constitutive_model` property to + one of the scalar `uw.constitutive_models` classes and populating the parameters. + It is usually a constant or a function of position / time and may also be non-linear + or anisotropic. + + - Volumetric sources of $u$ are specified using the $f$ property and can be any valid combination of `sympy` functions of position and + `meshVariable` or `swarmVariable` types. + + + + """ + + def _object_viewer(self): + from IPython.display import Latex, Markdown, display + + super()._object_viewer() + + ## feedback on this instance + display(Latex(r"$\quad\mathrm{u} = $ " + self.u.sym._repr_latex_())) + display(Latex(r"$\quad\Delta t = $ " + self.delta_t._repr_latex_())) + + @timing.routine_timer_decorator + def __init__( + self, + mesh: uw.discretisation.Mesh, + u_Field: uw.discretisation.MeshVariable, + order: int = 1, + theta: float = 0., + evalf: Optional[bool] = False, + verbose=False, + DuDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, + DFDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, + ): + ## Parent class will set up default values etc + super().__init__( + mesh, + u_Field, + u_Field.degree, + verbose, + DuDt=DuDt, + DFDt=DFDt, + ) + + # default values for properties + self.f = sympy.Matrix.zeros(1, 1) + + self._constitutive_model = None + + self.theta = theta + + # These are unique to the advection solver + self._delta_t = uw.function.expression( + R"\Delta t", 0, "Physically motivated timestep" + ) + self.is_setup = False + + ### Setup the history terms ... This version should not build anything + ### by default - it's the template / skeleton + + ## NB - Smoothing is generally required for stability. 0.0001 is effective + ## at the various resolutions tested. + + if DuDt is None: + self.Unknowns.DuDt = Eulerian_DDt( + self.mesh, + u_Field, + vtype=uw.VarType.SCALAR, + degree=u_Field.degree, + continuous=u_Field.continuous, + varsymbol=u_Field.symbol, + verbose=verbose, + evalf=evalf, + bcs=self.essential_bcs, + order=order, + smoothing=0.0, + ) + + else: + # validation + if order is None: + order = DuDt.order + + else: + if DuDt.order < order: + raise RuntimeError( + f"DuDt supplied is order {DuDt.order} but order requested is {order}" + ) + + self.Unknowns.DuDt = DuDt + + if DFDt is None: + self.Unknowns.DFDt = Symbolic_DDt( + sympy.Matrix( + [[0] * self.mesh.dim] ), + varsymbol=rf"{{F[ {self.u.symbol} ] }}", + theta=theta, + bcs=None, + order=order, + ) + + # self.Unknowns.DFDt = Eulerian_DDt( + # self.mesh, + # sympy.Matrix( + # [[0] * self.mesh.dim] + # ), # Actual function is not defined at this point + # vtype=uw.VarType.VECTOR, + # degree=u_Field.degree, + # continuous=u_Field.continuous, + # varsymbol=rf"{{F[ {self.u.symbol} ] }}", + # theta=theta, + # evalf=evalf, + # verbose=verbose, + # bcs=None, + # order=order, + # smoothing=0.0, + # ) + + return + + @property + def F0(self): + + f0 = uw.function.expression( + r"f_0 \left( \mathbf{u} \right)", + -self.f + sympy.simplify( self.DuDt.bdf() ) / self.delta_t, + "Diffusion pointwise force term: f_0(u)", + ) + + # backward compatibility + self._f0 = f0 + + return f0 + + @property + def F1(self): + + F1_val = uw.function.expression( + r"\mathbf{F}_1\left( \mathbf{u} \right)", + self.DFDt.adams_moulton_flux().T, + "Diffusion pointwise flux term: F_1(u)", + ) + + # backward compatibility + self._f1 = F1_val + + return F1_val + + @property + def f(self): + return self._f + + @f.setter + def f(self, value): + self.is_setup = False + self._f = sympy.Matrix((value,)) + + + @property + def delta_t(self): + return self._delta_t + + @delta_t.setter + def delta_t(self, value): + self.is_setup = False + self._delta_t.sym = value + + @timing.routine_timer_decorator + def estimate_dt(self): + r""" + Calculates an appropriate timestep for the given + mesh and diffusivity configuration. This is an implicit solver + so the $\delta_t$ should be interpreted as: + + - ${\delta t}_\textrm{diff}: a typical time for the diffusion front to propagate across an element + - ${\delta t}_\textrm{adv}: a typical element-crossing time for a fluid parcel + + returns (${\delta t}_\textrm{diff}, ${\delta t}_\textrm{adv}) + """ + + if isinstance(self.constitutive_model.diffusivity.sym, sympy.Expr): + if uw.function.fn_is_constant_expr( + self.constitutive_model.diffusivity.sym + ): + max_diffusivity = uw.function.evaluate( + self.constitutive_model.diffusivity.sym, + np.zeros((1, self.mesh.dim)), + ) + + else: + k = uw.function.evaluate( + sympy.sympify(self.constitutive_model.diffusivity.sym), + self.mesh._centroids, + self.mesh.N, + ) + + max_diffusivity = k.max() + else: + k = self.constitutive_model.diffusivity.sym + max_diffusivity = k + + ### required modules + from mpi4py import MPI + + ## get global max dif value + comm = uw.mpi.comm + diffusivity_glob = comm.allreduce(max_diffusivity, op=MPI.MAX) + + ## get radius + min_dx = self.mesh.get_min_radius() + + ## estimate dt of adv and diff components + self.dt_diff = 0.0 + + + dt_diff = (min_dx**2) / diffusivity_glob + self.dt_diff = dt_diff + + dt_estimate = dt_diff + + return dt_estimate + + @timing.routine_timer_decorator + def solve( + self, + zero_init_guess: bool = True, + timestep: float = None, + evalf: bool = False, + _force_setup: bool = False, + verbose=False, + ): + """ + Generates solution to constructed system. + + Params + ------ + zero_init_guess: + If `True`, a zero initial guess will be used for the + system solution. Otherwise, the current values of `self.u` will be used. + """ + + if timestep is not None and timestep != self.delta_t: + self.delta_t = timestep # this will force an initialisation because the functions need to be updated + + if _force_setup: + self.is_setup = False + + if not self.constitutive_model._solver_is_setup: + self.is_setup = False + self.DFDt.psi_fn = self.constitutive_model.flux.T + # self._flux = self.constitutive_model.flux.T + # self._flux_star = self._flux.copy() + + + + if not self.is_setup: + self._setup_pointwise_functions(verbose) + self._setup_discretisation(verbose) + self._setup_solver(verbose) + + # Update History / Flux History terms + # SemiLagrange and Lagrange may have different sequencing. + # self.DuDt.update_pre_solve(verbose=verbose) + # self.DFDt.update_pre_solve(verbose=verbose) + + super().solve(zero_init_guess, _force_setup) + + self.DuDt.update_post_solve(evalf=evalf, verbose=verbose) + self.DFDt.update_post_solve(evalf=evalf, verbose=verbose) + + # self._flux_star = self._flux.copy() + + self.is_setup = True + self.constitutive_model._solver_is_setup = True + + return + + From 3d0f920920c6c654cc6445cdf4ddc335bc524a58 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Mon, 14 Apr 2025 12:48:40 +0800 Subject: [PATCH 2/4] Change to symbolic notation Functions now appear correctly in the view() function --- src/underworld3/constitutive_models.py | 84 ++++++++++---------- src/underworld3/systems/ddt.py | 30 ++++--- src/underworld3/systems/diffusion_solvers.py | 9 ++- 3 files changed, 66 insertions(+), 57 deletions(-) diff --git a/src/underworld3/constitutive_models.py b/src/underworld3/constitutive_models.py index f8f89904a..5c02a39bb 100644 --- a/src/underworld3/constitutive_models.py +++ b/src/underworld3/constitutive_models.py @@ -214,8 +214,10 @@ def c(self): if not self._is_setup: self._build_c_tensor() - - return self._c.as_immutable() + if hasattr(self._c, "sym"): + return sympy.Matrix(self._c.sym).as_immutable() + else: + return self._c.as_immutable() @property def flux(self): @@ -1465,51 +1467,39 @@ def _object_viewer(self): return +# AnisotropicDiffusionModel: expects a diffusivity vector and builds a diagonal tensor. class AnisotropicDiffusionModel(DiffusionModel): - """ - Anisotropic version of DiffusionModel supporting diagonal diffusion tensors. - - Diffusivity must be provided as a `sympy.Matrix` vector of length `dim`: - - 2D example: `Matrix([Dx, Dy])` - - 3D example: `Matrix([Dx, Dy, Dz])` - - This vector is interpreted as the diagonal of the diffusivity tensor: - κ = diag(Dx, Dy, Dz) - - Usage: - ```python - from sympy import symbols, Matrix - kappa = symbols("kappa") - model = AnisotropicDiffusionModel - model.material_properties = model.Parameters(diffusivity=sp.Matrix([kappa, 2*kappa])) - ``` - """ - - class _Parameters(DiffusionModel._Parameters): + class _Parameters: def __init__(inner_self, _owning_model): dim = _owning_model.dim inner_self._owning_model = _owning_model - inner_self._diffusivity = sympy.eye(dim) - + # Set default diffusivity as an identity matrix wrapped in an expression + default_diffusivity = sympy.ones(_owning_model.dim, 1) + elements = [default_diffusivity[i] for i in range(dim)] + validated = [] + for i, v in enumerate(elements): + comp = validate_parameters( + fr"\upkappa_{{{i}}}", v, f"Diffusivity in x_{i}", allow_number=True + ) + if comp is not None: + validated.append(comp) + # Store the validated diffusivity as a diagonal matrix + inner_self._diffusivity = sympy.diag(*validated) + @property def diffusivity(inner_self): return inner_self._diffusivity - + @diffusivity.setter - def diffusivity(inner_self, value): + def diffusivity(inner_self, value: sympy.Matrix): dim = inner_self._owning_model.dim - if not isinstance(value, sympy.Matrix): - raise TypeError("Diffusivity must be a sympy.Matrix vector.") - # Accept shape (dim, 1) or (1, dim) if value.shape not in [(dim, 1), (1, dim)]: raise ValueError( - f"Diffusivity must be a vector of length {dim}. " - f"Got shape {value.shape}." + f"Diffusivity must be a vector of length {dim}. Got shape {value.shape}." ) - - # Flatten and validate + # Validate each component using validate_parameters elements = [value[i] for i in range(dim)] validated = [] for i, v in enumerate(elements): @@ -1518,23 +1508,23 @@ def diffusivity(inner_self, value): ) if diff is not None: validated.append(diff) - + # Store the validated diffusivity as a diagonal matrix inner_self._diffusivity = sympy.diag(*validated) inner_self._reset() def _build_c_tensor(self): - """Constructs the anisotropic (diagonal) diffusivity tensor.""" + """Constructs the anisotropic (diagonal) tensor from the diffusivity vector.""" self._c = self.Parameters.diffusivity + self._is_setup = True def _object_viewer(self): from IPython.display import Latex, display + super()._object_viewer() - - # Extract diagonal entries and render cleanly + diagonal = self.Parameters.diffusivity.diagonal() latex_entries = ", ".join([sympy.latex(k) for k in diagonal]) kappa_latex = r"\kappa = \mathrm{diag}\left(" + latex_entries + r"\right)" - display(Latex(r"$\quad " + kappa_latex + r"$")) @@ -1555,20 +1545,28 @@ class GenericFluxModel(Constitutive_Model): class _Parameters: def __init__(inner_self, _owning_model): - inner_self._flux = sympy.Matrix(sympy.ones(_owning_model.dim, 1)) inner_self._owning_model = _owning_model + default_flux = sympy.zeros(_owning_model.dim, 1) + elements = [default_flux[i] for i in range(_owning_model.dim)] + validated = [] + for i, v in enumerate(elements): + flux_component = validate_parameters( + fr"q_{{{i}}}", v, f"Flux component in x_{i}", allow_number=True + ) + if flux_component is not None: + validated.append(flux_component) + + inner_self._flux = sympy.Matrix(validated) + @property def flux(inner_self): return inner_self._flux @flux.setter - def flux(inner_self, value): + def flux(inner_self, value: sympy.Matrix): dim = inner_self._owning_model.dim - if not isinstance(value, sympy.Matrix): - raise TypeError("Flux must be a sympy.Matrix vector.") - # Accept shape (dim, 1) or (1, dim) if value.shape not in [(dim, 1), (1, dim)]: raise ValueError( diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 38fb84911..c41c4869c 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -39,7 +39,7 @@ def __init__( self, psi_fn: sympy.Basic, # a sympy expression for ψ; can be scalar or matrix theta: Optional[float] = 0.5, - varsymbol: Optional[str] = r"u", + varsymbol: Optional[str] = r"\psi", verbose: Optional[bool] = False, bcs=[], order: int = 1, @@ -55,16 +55,16 @@ def __init__( # Ensure psi_fn is a sympy Matrix. if not isinstance(psi_fn, sympy.Matrix): try: - # If psi_fn has a shape attribute, convert it. psi_fn = sympy.Matrix(psi_fn) except Exception: psi_fn = sympy.Matrix([[psi_fn]]) - - # Do not force a column vector if psi_fn has more than one column. - # Instead, preserve its full shape. self._psi_fn = psi_fn # stored with its native shape - self._shape = psi_fn.shape # capture the shape to create history storage - + self._shape = psi_fn.shape # capture the shape + + # Set the display symbol for psi_fn and for the history variable. + self._psi_fn_symbol = varsymbol # e.g. "\psi" + self._psi_star_symbol = varsymbol + r"^\ast" # e.g. "\psi^\ast" + # Create the history list: each element is a Matrix of shape _shape. self.psi_star = [sympy.zeros(*self._shape) for _ in range(order)] self.initiate_history_fn() @@ -88,7 +88,11 @@ def psi_fn(self, new_fn): def _object_viewer(self): from IPython.display import Latex, display - display(Latex(rf"$\quad$History steps = {self.order}$")) + # Display the primary variable + display(Latex(rf"$\quad {self._psi_fn_symbol} = {sympy.latex(self._psi_fn)}$")) + # Display the history variable using the different symbol. + history_latex = ", ".join([sympy.latex(elem) for elem in self.psi_star]) + display(Latex(rf"$\quad {self._psi_star_symbol} = \left[{history_latex}\right]$")) def update_history_fn(self): # Update the first history element with a copy of the current ψ. @@ -124,6 +128,8 @@ def update_post_solve( ): if verbose: print(f"Updating history for ψ = {self.psi_fn}", flush=True) + + # Shift history: copy each element down the chain. for i in range(self.order - 1, 0, -1): self.psi_star[i] = self.psi_star[i - 1].copy() @@ -227,7 +233,7 @@ def __init__( for i in range(order): self.psi_star.append( uw.discretisation.MeshVariable( - f"psi_star_sl_{self.instance_number}_{i}", + f"psi_star_Eulerian_{self.instance_number}_{i}", self.mesh, vtype=vtype, degree=degree, @@ -284,7 +290,7 @@ def _setup_projections(self): ) elif self.vtype == uw.VarType.SYM_TENSOR or self.vtype == uw.VarType.TENSOR: self._WorkVar = uw.discretisation.MeshVariable( - f"W_star_slcn_{self.instance_number}", + f"W_star_Eulerian_{self.instance_number}", self.mesh, vtype=uw.VarType.SCALAR, degree=self.degree, @@ -424,8 +430,8 @@ def adams_moulton_flux(self, order=None): with sympy.core.evaluate(False): if order == 1: - # am = (self.psi_fn + self.psi_star[0].sym) / 2 - am = self.theta*self.psi_fn + ((1.-self.theta)*self.psi_star[0].sym) + am = (self.psi_fn + self.psi_star[0].sym) / 2 + # am = self.theta*self.psi_fn + ((1.-self.theta)*self.psi_star[0].sym) elif order == 2: am = ( diff --git a/src/underworld3/systems/diffusion_solvers.py b/src/underworld3/systems/diffusion_solvers.py index d5475ba71..33556ac0a 100644 --- a/src/underworld3/systems/diffusion_solvers.py +++ b/src/underworld3/systems/diffusion_solvers.py @@ -147,7 +147,7 @@ def __init__( bcs=None, order=order, ) - + ### solution unable to solve after n timesteps, due to the projection of flux term (???) # self.Unknowns.DFDt = Eulerian_DDt( # self.mesh, # sympy.Matrix( @@ -186,7 +186,7 @@ def F1(self): F1_val = uw.function.expression( r"\mathbf{F}_1\left( \mathbf{u} \right)", - self.DFDt.adams_moulton_flux().T, + self.DFDt.adams_moulton_flux(), "Diffusion pointwise flux term: F_1(u)", ) @@ -317,6 +317,11 @@ def solve( self.DuDt.update_post_solve(evalf=evalf, verbose=verbose) self.DFDt.update_post_solve(evalf=evalf, verbose=verbose) + # if isinstance(self.DFDt, Eulerian_DDt): + # for i in range(order): + # ### have to substitute the unknown history term into the symbolic flux term + # self.DFDt.psi_star[i].subs({self.DuDt.psi_fn:self.DuDt.psi_star[i]}) + # self._flux_star = self._flux.copy() self.is_setup = True From 17abdadf1089bf2cb3abfd916829e063c780dc22 Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Thu, 1 May 2025 13:37:13 +0800 Subject: [PATCH 3/4] Add L2_norm calculation --- src/underworld3/maths/__init__.py | 1 + src/underworld3/maths/functions.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/src/underworld3/maths/__init__.py b/src/underworld3/maths/__init__.py index 3878a5221..ab2afd5d4 100644 --- a/src/underworld3/maths/__init__.py +++ b/src/underworld3/maths/__init__.py @@ -5,6 +5,7 @@ ) from .functions import delta as delta_function +from .functions import L2_norm as L2_norm # from .vector_calculus import ( # mesh_vector_calculus_spherical_lonlat as vector_calculus_spherical_lonlat, diff --git a/src/underworld3/maths/functions.py b/src/underworld3/maths/functions.py index 599230cdd..d1e0939e9 100644 --- a/src/underworld3/maths/functions.py +++ b/src/underworld3/maths/functions.py @@ -4,6 +4,7 @@ from sympy import sympify from typing import Optional, Callable from underworld3 import function +from underworld3 import maths def delta( @@ -15,3 +16,32 @@ def delta( delta_fn = sympy.exp(-(x**2) / (2 * epsilon**2)) / (epsilon * sqrt_2pi) return delta_fn + +def L2_norm(n_s, a_s, mesh): + """ + Compute the L2 norm (Euclidean norm) of the difference between numerical and analytical solutions. + + Parameters: + n_s : Numeric solution (scalar or vector field) + a_s : Analytic solution (scalar or vector field) + mesh : The mesh used for integration + + Returns: + L2 norm value (for scalar or vector) + """ + # Check if the input is a vector (SymPy Matrix) + if isinstance(n_s, sympy.Matrix) and n_s.shape[1] > 1: + # Compute squared difference using dot product for vectors + squared_difference = (n_s - a_s).dot(n_s - a_s) + else: + # Compute squared difference for scalars + squared_difference = (n_s - a_s) ** 2 + + + # Integral over the domain + I = maths.Integral(mesh, squared_difference) + + # Compute the L2 norm + L2 = sympy.sqrt( I.evaluate() ) + + return L2 From 7b54ec6809f70281d9715e0d861d2c4f0de9c1ab Mon Sep 17 00:00:00 2001 From: Ben Knight <55677727+bknight1@users.noreply.github.com> Date: Thu, 1 May 2025 13:49:33 +0800 Subject: [PATCH 4/4] Remove separate diffusion_solver file Folded into solvers.py --- src/underworld3/systems/__init__.py | 6 +- src/underworld3/systems/diffusion_solvers.py | 332 ------------------- src/underworld3/systems/solvers.py | 301 +++++++++++++++++ 3 files changed, 305 insertions(+), 334 deletions(-) delete mode 100644 src/underworld3/systems/diffusion_solvers.py diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index 5e16a9fca..ef69e670c 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -24,6 +24,9 @@ from .solvers import SNES_AdvectionDiffusion as AdvDiffusionSLCN from .solvers import SNES_AdvectionDiffusion as AdvDiffusion +# import diffusion-only solver +from .solvers import SNES_Diffusion as Diffusion + # These are now implemented the same way using the ddt module from .solvers import SNES_NavierStokes as NavierStokesSwarm from .solvers import SNES_NavierStokes as NavierStokesSLCN @@ -35,5 +38,4 @@ from .ddt import Lagrangian_Swarm as Lagrangian_Swarm_DDt from .ddt import Eulerian as Eulerian_DDt -# import additional diffusion solvers -from .diffusion_solvers import SNES_Diffusion as Diffusion + diff --git a/src/underworld3/systems/diffusion_solvers.py b/src/underworld3/systems/diffusion_solvers.py deleted file mode 100644 index 33556ac0a..000000000 --- a/src/underworld3/systems/diffusion_solvers.py +++ /dev/null @@ -1,332 +0,0 @@ -import sympy -from sympy import sympify -import numpy as np - -from typing import Optional, Callable, Union - -import underworld3 as uw -from underworld3.systems import SNES_Scalar, SNES_Vector, SNES_Stokes_SaddlePt -from underworld3 import VarType -import underworld3.timing as timing -from underworld3.utilities._api_tools import uw_object - - -from .ddt import SemiLagrangian as SemiLagrangian_DDt -from .ddt import Lagrangian as Lagrangian_DDt -from .ddt import Eulerian as Eulerian_DDt -from .ddt import Symbolic as Symbolic_DDt - -import sympy -import numpy as np -from typing import Optional, Union -import underworld3 as uw -from underworld3.systems import SNES_Vector, SNES_Scalar -from underworld3 import VarType -import underworld3.timing as timing -from underworld3.utilities._api_tools import uw_object - -from .ddt import SemiLagrangian as SemiLagrangian_DDt -from .ddt import Lagrangian as Lagrangian_DDt -from .ddt import Eulerian as Eulerian_DDt -from .ddt import Symbolic as Symbolic_DDt - - -class SNES_Diffusion(SNES_Scalar): - r""" - This class provides a solver for the scalar Diffusion equation using mesh-based finite elements. - - $$ - \color{Green}{\underbrace{ \Bigl[ \frac{\partial u}{\partial t} - \left( \mathbf{v} \cdot \nabla \right) u \Bigr]}_{\dot{\mathbf{f}}}} - - \nabla \cdot - \color{Blue}{\underbrace{\Bigl[ \boldsymbol\kappa \nabla u \Bigr]}_{\mathbf{F}}} = - \color{Maroon}{\underbrace{\Bigl[ f \Bigl] }_{\mathbf{f}}} - $$ - - The term $\mathbf{F}$ relates diffusive fluxes to gradients in the unknown $u$. - - ## Properties - - - The unknown is $u$. - - - The diffusivity tensor, $\kappa$ is provided by setting the `constitutive_model` property to - one of the scalar `uw.constitutive_models` classes and populating the parameters. - It is usually a constant or a function of position / time and may also be non-linear - or anisotropic. - - - Volumetric sources of $u$ are specified using the $f$ property and can be any valid combination of `sympy` functions of position and - `meshVariable` or `swarmVariable` types. - - - - """ - - def _object_viewer(self): - from IPython.display import Latex, Markdown, display - - super()._object_viewer() - - ## feedback on this instance - display(Latex(r"$\quad\mathrm{u} = $ " + self.u.sym._repr_latex_())) - display(Latex(r"$\quad\Delta t = $ " + self.delta_t._repr_latex_())) - - @timing.routine_timer_decorator - def __init__( - self, - mesh: uw.discretisation.Mesh, - u_Field: uw.discretisation.MeshVariable, - order: int = 1, - theta: float = 0., - evalf: Optional[bool] = False, - verbose=False, - DuDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, - DFDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, - ): - ## Parent class will set up default values etc - super().__init__( - mesh, - u_Field, - u_Field.degree, - verbose, - DuDt=DuDt, - DFDt=DFDt, - ) - - # default values for properties - self.f = sympy.Matrix.zeros(1, 1) - - self._constitutive_model = None - - self.theta = theta - - # These are unique to the advection solver - self._delta_t = uw.function.expression( - R"\Delta t", 0, "Physically motivated timestep" - ) - self.is_setup = False - - ### Setup the history terms ... This version should not build anything - ### by default - it's the template / skeleton - - ## NB - Smoothing is generally required for stability. 0.0001 is effective - ## at the various resolutions tested. - - if DuDt is None: - self.Unknowns.DuDt = Eulerian_DDt( - self.mesh, - u_Field, - vtype=uw.VarType.SCALAR, - degree=u_Field.degree, - continuous=u_Field.continuous, - varsymbol=u_Field.symbol, - verbose=verbose, - evalf=evalf, - bcs=self.essential_bcs, - order=order, - smoothing=0.0, - ) - - else: - # validation - if order is None: - order = DuDt.order - - else: - if DuDt.order < order: - raise RuntimeError( - f"DuDt supplied is order {DuDt.order} but order requested is {order}" - ) - - self.Unknowns.DuDt = DuDt - - if DFDt is None: - self.Unknowns.DFDt = Symbolic_DDt( - sympy.Matrix( - [[0] * self.mesh.dim] ), - varsymbol=rf"{{F[ {self.u.symbol} ] }}", - theta=theta, - bcs=None, - order=order, - ) - ### solution unable to solve after n timesteps, due to the projection of flux term (???) - # self.Unknowns.DFDt = Eulerian_DDt( - # self.mesh, - # sympy.Matrix( - # [[0] * self.mesh.dim] - # ), # Actual function is not defined at this point - # vtype=uw.VarType.VECTOR, - # degree=u_Field.degree, - # continuous=u_Field.continuous, - # varsymbol=rf"{{F[ {self.u.symbol} ] }}", - # theta=theta, - # evalf=evalf, - # verbose=verbose, - # bcs=None, - # order=order, - # smoothing=0.0, - # ) - - return - - @property - def F0(self): - - f0 = uw.function.expression( - r"f_0 \left( \mathbf{u} \right)", - -self.f + sympy.simplify( self.DuDt.bdf() ) / self.delta_t, - "Diffusion pointwise force term: f_0(u)", - ) - - # backward compatibility - self._f0 = f0 - - return f0 - - @property - def F1(self): - - F1_val = uw.function.expression( - r"\mathbf{F}_1\left( \mathbf{u} \right)", - self.DFDt.adams_moulton_flux(), - "Diffusion pointwise flux term: F_1(u)", - ) - - # backward compatibility - self._f1 = F1_val - - return F1_val - - @property - def f(self): - return self._f - - @f.setter - def f(self, value): - self.is_setup = False - self._f = sympy.Matrix((value,)) - - - @property - def delta_t(self): - return self._delta_t - - @delta_t.setter - def delta_t(self, value): - self.is_setup = False - self._delta_t.sym = value - - @timing.routine_timer_decorator - def estimate_dt(self): - r""" - Calculates an appropriate timestep for the given - mesh and diffusivity configuration. This is an implicit solver - so the $\delta_t$ should be interpreted as: - - - ${\delta t}_\textrm{diff}: a typical time for the diffusion front to propagate across an element - - ${\delta t}_\textrm{adv}: a typical element-crossing time for a fluid parcel - - returns (${\delta t}_\textrm{diff}, ${\delta t}_\textrm{adv}) - """ - - if isinstance(self.constitutive_model.diffusivity.sym, sympy.Expr): - if uw.function.fn_is_constant_expr( - self.constitutive_model.diffusivity.sym - ): - max_diffusivity = uw.function.evaluate( - self.constitutive_model.diffusivity.sym, - np.zeros((1, self.mesh.dim)), - ) - - else: - k = uw.function.evaluate( - sympy.sympify(self.constitutive_model.diffusivity.sym), - self.mesh._centroids, - self.mesh.N, - ) - - max_diffusivity = k.max() - else: - k = self.constitutive_model.diffusivity.sym - max_diffusivity = k - - ### required modules - from mpi4py import MPI - - ## get global max dif value - comm = uw.mpi.comm - diffusivity_glob = comm.allreduce(max_diffusivity, op=MPI.MAX) - - ## get radius - min_dx = self.mesh.get_min_radius() - - ## estimate dt of adv and diff components - self.dt_diff = 0.0 - - - dt_diff = (min_dx**2) / diffusivity_glob - self.dt_diff = dt_diff - - dt_estimate = dt_diff - - return dt_estimate - - @timing.routine_timer_decorator - def solve( - self, - zero_init_guess: bool = True, - timestep: float = None, - evalf: bool = False, - _force_setup: bool = False, - verbose=False, - ): - """ - Generates solution to constructed system. - - Params - ------ - zero_init_guess: - If `True`, a zero initial guess will be used for the - system solution. Otherwise, the current values of `self.u` will be used. - """ - - if timestep is not None and timestep != self.delta_t: - self.delta_t = timestep # this will force an initialisation because the functions need to be updated - - if _force_setup: - self.is_setup = False - - if not self.constitutive_model._solver_is_setup: - self.is_setup = False - self.DFDt.psi_fn = self.constitutive_model.flux.T - # self._flux = self.constitutive_model.flux.T - # self._flux_star = self._flux.copy() - - - - if not self.is_setup: - self._setup_pointwise_functions(verbose) - self._setup_discretisation(verbose) - self._setup_solver(verbose) - - # Update History / Flux History terms - # SemiLagrange and Lagrange may have different sequencing. - # self.DuDt.update_pre_solve(verbose=verbose) - # self.DFDt.update_pre_solve(verbose=verbose) - - super().solve(zero_init_guess, _force_setup) - - self.DuDt.update_post_solve(evalf=evalf, verbose=verbose) - self.DFDt.update_post_solve(evalf=evalf, verbose=verbose) - - # if isinstance(self.DFDt, Eulerian_DDt): - # for i in range(order): - # ### have to substitute the unknown history term into the symbolic flux term - # self.DFDt.psi_star[i].subs({self.DuDt.psi_fn:self.DuDt.psi_star[i]}) - - # self._flux_star = self._flux.copy() - - self.is_setup = True - self.constitutive_model._solver_is_setup = True - - return - - diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 715f6c867..2f8f2973d 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -13,6 +13,8 @@ from .ddt import SemiLagrangian as SemiLagrangian_DDt from .ddt import Lagrangian as Lagrangian_DDt +from .ddt import Eulerian as Eulerian_DDt +from .ddt import Symbolic as Symbolic_DDt # class UW_Scalar_Temple(SNES_Scalar): @@ -1482,6 +1484,305 @@ def solve( self.constitutive_model._solver_is_setup = True return + + +class SNES_Diffusion(SNES_Scalar): + r""" + This class provides a solver for the scalar Diffusion equation using mesh-based finite elements. + + $$ + \color{Green}{\underbrace{ \Bigl[ \frac{\partial u}{\partial t} - \left( \mathbf{v} \cdot \nabla \right) u \Bigr]}_{\dot{\mathbf{f}}}} - + \nabla \cdot + \color{Blue}{\underbrace{\Bigl[ \boldsymbol\kappa \nabla u \Bigr]}_{\mathbf{F}}} = + \color{Maroon}{\underbrace{\Bigl[ f \Bigl] }_{\mathbf{f}}} + $$ + + The term $\mathbf{F}$ relates diffusive fluxes to gradients in the unknown $u$. + + ## Properties + + - The unknown is $u$. + + - The diffusivity tensor, $\kappa$ is provided by setting the `constitutive_model` property to + one of the scalar `uw.constitutive_models` classes and populating the parameters. + It is usually a constant or a function of position / time and may also be non-linear + or anisotropic. + + - Volumetric sources of $u$ are specified using the $f$ property and can be any valid combination of `sympy` functions of position and + `meshVariable` or `swarmVariable` types. + + + + """ + + def _object_viewer(self): + from IPython.display import Latex, Markdown, display + + super()._object_viewer() + + ## feedback on this instance + display(Latex(r"$\quad\mathrm{u} = $ " + self.u.sym._repr_latex_())) + display(Latex(r"$\quad\Delta t = $ " + self.delta_t._repr_latex_())) + + @timing.routine_timer_decorator + def __init__( + self, + mesh: uw.discretisation.Mesh, + u_Field: uw.discretisation.MeshVariable, + order: int = 1, + theta: float = 0., + evalf: Optional[bool] = False, + verbose=False, + DuDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, + DFDt: Union[Eulerian_DDt, SemiLagrangian_DDt, Lagrangian_DDt] = None, + ): + ## Parent class will set up default values etc + super().__init__( + mesh, + u_Field, + u_Field.degree, + verbose, + DuDt=DuDt, + DFDt=DFDt, + ) + + # default values for properties + self.f = sympy.Matrix.zeros(1, 1) + + self._constitutive_model = None + + self.theta = theta + + # These are unique to the advection solver + self._delta_t = uw.function.expression( + R"\Delta t", 0, "Physically motivated timestep" + ) + self.is_setup = False + + ### Setup the history terms ... This version should not build anything + ### by default - it's the template / skeleton + + ## NB - Smoothing is generally required for stability. 0.0001 is effective + ## at the various resolutions tested. + + if DuDt is None: + self.Unknowns.DuDt = Eulerian_DDt( + self.mesh, + u_Field, + vtype=uw.VarType.SCALAR, + degree=u_Field.degree, + continuous=u_Field.continuous, + varsymbol=u_Field.symbol, + verbose=verbose, + evalf=evalf, + bcs=self.essential_bcs, + order=order, + smoothing=0.0, + ) + + else: + # validation + if order is None: + order = DuDt.order + + else: + if DuDt.order < order: + raise RuntimeError( + f"DuDt supplied is order {DuDt.order} but order requested is {order}" + ) + + self.Unknowns.DuDt = DuDt + + if DFDt is None: + self.Unknowns.DFDt = Symbolic_DDt( + sympy.Matrix( + [[0] * self.mesh.dim] ), + varsymbol=rf"{{F[ {self.u.symbol} ] }}", + theta=theta, + bcs=None, + order=order, + ) + ### solution unable to solve after n timesteps, due to the projection of flux term (???) + # self.Unknowns.DFDt = Eulerian_DDt( + # self.mesh, + # sympy.Matrix( + # [[0] * self.mesh.dim] + # ), # Actual function is not defined at this point + # vtype=uw.VarType.VECTOR, + # degree=u_Field.degree, + # continuous=u_Field.continuous, + # varsymbol=rf"{{F[ {self.u.symbol} ] }}", + # theta=theta, + # evalf=evalf, + # verbose=verbose, + # bcs=None, + # order=order, + # smoothing=0.0, + # ) + + return + + @property + def F0(self): + + f0 = uw.function.expression( + r"f_0 \left( \mathbf{u} \right)", + -self.f + sympy.simplify( self.DuDt.bdf() ) / self.delta_t, + "Diffusion pointwise force term: f_0(u)", + ) + + # backward compatibility + self._f0 = f0 + + return f0 + + @property + def F1(self): + + F1_val = uw.function.expression( + r"\mathbf{F}_1\left( \mathbf{u} \right)", + self.DFDt.adams_moulton_flux(), + "Diffusion pointwise flux term: F_1(u)", + ) + + # backward compatibility + self._f1 = F1_val + + return F1_val + + @property + def f(self): + return self._f + + @f.setter + def f(self, value): + self.is_setup = False + self._f = sympy.Matrix((value,)) + + + @property + def delta_t(self): + return self._delta_t + + @delta_t.setter + def delta_t(self, value): + self.is_setup = False + self._delta_t.sym = value + + @timing.routine_timer_decorator + def estimate_dt(self): + r""" + Calculates an appropriate timestep for the given + mesh and diffusivity configuration. This is an implicit solver + so the $\delta_t$ should be interpreted as: + + - ${\delta t}_\textrm{diff}: a typical time for the diffusion front to propagate across an element + - ${\delta t}_\textrm{adv}: a typical element-crossing time for a fluid parcel + + returns (${\delta t}_\textrm{diff}, ${\delta t}_\textrm{adv}) + """ + + if isinstance(self.constitutive_model.diffusivity.sym, sympy.Expr): + if uw.function.fn_is_constant_expr( + self.constitutive_model.diffusivity.sym + ): + max_diffusivity = uw.function.evaluate( + self.constitutive_model.diffusivity.sym, + np.zeros((1, self.mesh.dim)), + ) + + else: + k = uw.function.evaluate( + sympy.sympify(self.constitutive_model.diffusivity.sym), + self.mesh._centroids, + self.mesh.N, + ) + + max_diffusivity = k.max() + else: + k = self.constitutive_model.diffusivity.sym + max_diffusivity = k + + ### required modules + from mpi4py import MPI + + ## get global max dif value + comm = uw.mpi.comm + diffusivity_glob = comm.allreduce(max_diffusivity, op=MPI.MAX) + + ## get radius + min_dx = self.mesh.get_min_radius() + + ## estimate dt of adv and diff components + self.dt_diff = 0.0 + + + dt_diff = (min_dx**2) / diffusivity_glob + self.dt_diff = dt_diff + + dt_estimate = dt_diff + + return dt_estimate + + @timing.routine_timer_decorator + def solve( + self, + zero_init_guess: bool = True, + timestep: float = None, + evalf: bool = False, + _force_setup: bool = False, + verbose=False, + ): + """ + Generates solution to constructed system. + + Params + ------ + zero_init_guess: + If `True`, a zero initial guess will be used for the + system solution. Otherwise, the current values of `self.u` will be used. + """ + + if timestep is not None and timestep != self.delta_t: + self.delta_t = timestep # this will force an initialisation because the functions need to be updated + + if _force_setup: + self.is_setup = False + + if not self.constitutive_model._solver_is_setup: + self.is_setup = False + self.DFDt.psi_fn = self.constitutive_model.flux.T + # self._flux = self.constitutive_model.flux.T + # self._flux_star = self._flux.copy() + + + + if not self.is_setup: + self._setup_pointwise_functions(verbose) + self._setup_discretisation(verbose) + self._setup_solver(verbose) + + # Update History / Flux History terms + # SemiLagrange and Lagrange may have different sequencing. + # self.DuDt.update_pre_solve(verbose=verbose) + # self.DFDt.update_pre_solve(verbose=verbose) + + super().solve(zero_init_guess, _force_setup) + + self.DuDt.update_post_solve(evalf=evalf, verbose=verbose) + self.DFDt.update_post_solve(evalf=evalf, verbose=verbose) + + # if isinstance(self.DFDt, Eulerian_DDt): + # for i in range(order): + # ### have to substitute the unknown history term into the symbolic flux term + # self.DFDt.psi_star[i].subs({self.DuDt.psi_fn:self.DuDt.psi_star[i]}) + + # self._flux_star = self._flux.copy() + + self.is_setup = True + self.constitutive_model._solver_is_setup = True + + return # This one is already updated to work with the Lagrange D_Dt