import warnings
from typing import Any
from mpi4py import MPI
from petsc4py import PETSc
import dolfinx
import numpy as np
import tqdm.auto
import ufl
from dolfinx import fem
from dolfinx.nls.petsc import NewtonSolver
from packaging.version import Version
import festim as F
from festim.mesh.mesh import Mesh as _Mesh
from festim.source import SourceBase as _SourceBase
from festim.subdomain.volume_subdomain import (
VolumeSubdomain as _VolumeSubdomain,
)
[docs]
class ProblemBase:
"""Base class for :py:class:`HeatTransferProblem
<festim.heat_transfer_problem.HeatTransferProblem>` and
:py:class:`HydrogenTransportProblem
<festim.hydrogen_transport_problem.HydrogenTransportProblem>`.
Attributes:
show_progress_bar: If `True` a progress bar is displayed during the simulation
progress_bar: the progress bar
"""
mesh: _Mesh
sources: list[_SourceBase]
exports: list[Any]
subdomains: list[_VolumeSubdomain]
show_progress_bar: bool
progress_bar: None | tqdm.auto.tqdm
timesteps: list[float]
def __init__(
self,
mesh: _Mesh = None,
sources=None,
exports=None,
subdomains=None,
boundary_conditions=None,
settings=None,
petsc_options=None,
) -> None:
self.mesh = mesh
# for arguments to initialise as empty list
# if arg not None, assign arg, else assign empty list
self.subdomains = subdomains or []
self.boundary_conditions = boundary_conditions or []
self.sources = sources or []
self.exports = exports or []
self.settings = settings
self.dx = None
self.ds = None
self.function_space = None
self.facet_meshtags = None
self.volume_meshtags = None
self.formulation = None
self.bc_forms = []
self.show_progress_bar = True
self.petsc_options = petsc_options
self._timesteps = []
@property
def volume_subdomains(self):
return [s for s in self.subdomains if isinstance(s, F.VolumeSubdomain)]
@property
def surface_subdomains(self):
return [s for s in self.subdomains if isinstance(s, F.SurfaceSubdomain)]
@property
def dt(self):
return self._dt
@property
def timesteps(self):
return self._timesteps
[docs]
def define_boundary_conditions(self):
"""Defines the dirichlet boundary conditions of the model."""
for bc in self.boundary_conditions:
if isinstance(bc, F.DirichletBCBase):
form = self.create_dirichletbc_form(bc)
self.bc_forms.append(form)
[docs]
def get_petsc_options(self) -> dict[str, Any]:
"""Gets the PETSc options to pass to the NewtonProblem solver. Default options
are updated with user-provided options, if any.
Returns:
the petsc options to pass to the NewtonProblem solver.
"""
petsc_options = get_default_petsc_options()
# Update default PETSc options with user-provided options, if any
if self.petsc_options:
petsc_options.update(self.petsc_options)
if self.petsc_options:
if (
"snes_atol" in self.petsc_options
or "snes_rtol" in self.petsc_options
or "snes_max_it" in self.petsc_options
):
warnings.warn(
"You have set one of the following PETSc options: snes_atol, "
"snes_rtol or snes_max_it. These options will be overwritten by "
"the values in festim.Settings (atol, rtol and max_iterations) to "
"ensure consistency between different versions of dolfinx. If you "
"want to set these options manually, please set them in "
"festim.Settings and not in the petsc_options dictionary."
)
petsc_options.update(
{
"snes_atol": self.settings.atol,
"snes_rtol": self.settings.rtol,
"snes_max_it": self.settings.max_iterations,
}
)
return petsc_options
[docs]
def create_solver(self):
"""Creates the solver of the model."""
if Version(dolfinx.__version__) == Version("0.9.0"):
problem = fem.petsc.NonlinearProblem(
self.formulation,
self.u,
bcs=self.bc_forms,
)
self.solver = NewtonSolver(MPI.COMM_WORLD, problem)
self.solver.atol = (
self.settings.atol
if not callable(self.settings.rtol)
else self.settings.rtol(float(self.t))
)
self.solver.rtol = (
self.settings.rtol
if not callable(self.settings.rtol)
else self.settings.rtol(float(self.t))
)
self.solver.max_it = self.settings.max_iterations
ksp = self.solver.krylov_solver
if self.petsc_options is None:
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setErrorIfNotConverged(True)
else:
# Set PETSc options
opts = PETSc.Options()
option_prefix = ksp.getOptionsPrefix()
for k, v in self.petsc_options.items():
opts[f"{option_prefix}{k}"] = v
ksp.setFromOptions()
elif Version(dolfinx.__version__) > Version("0.9.0"):
from dolfinx.fem.petsc import NonlinearProblem
petsc_options = self.get_petsc_options()
self.solver = NonlinearProblem(
self.formulation,
self.u,
bcs=self.bc_forms,
petsc_options=petsc_options,
petsc_options_prefix="festim_solver",
)
self.solver.solver.setMonitor(F.helpers.SnesMonitor)
self.solver.solver.getKSP().setMonitor(F.helpers.KSPMonitor)
self.solver.solver.setConvergenceTest(F.helpers.convergenceTest)
# Delete PETSc options post setting them, ref:
# https://gitlab.com/petsc/petsc/-/issues/1201
snes = self.solver.solver
prefix = snes.getOptionsPrefix()
opts = PETSc.Options()
for k in petsc_options.keys():
del opts[f"{prefix}{k}"]
[docs]
def run(self):
"""Runs the model."""
if self.settings.transient:
# Solve transient
if self.show_progress_bar:
self.progress_bar = tqdm.auto.tqdm(
desc=f"Solving {self.__class__.__name__}",
total=self.settings.final_time,
unit_scale=True,
)
while self.t.value < self.settings.final_time:
self.iterate()
if self.show_progress_bar:
self.progress_bar.refresh() # refresh progress bar to show 100%
self.progress_bar.close()
else:
# Solve steady-state
if Version(dolfinx.__version__) == Version("0.9.0"):
self.solver.solve(self.u)
elif Version(dolfinx.__version__) > Version("0.9.0"):
self.solver.solve()
self.post_processing()
[docs]
def iterate(self):
"""Iterates the model for a given time step."""
self._timesteps.append(float(self.t))
if self.show_progress_bar:
self.progress_bar.update(
min(self.dt.value, abs(self.settings.final_time - self.t.value))
)
# update rtol if it's callable
if callable(self.settings.rtol):
self.solver.rtol = self.settings.rtol(self.t.value)
# update rtol if it's callable
if callable(self.settings.atol):
self.solver.atol = self.settings.atol(self.t.value)
self.t.value += self.dt.value
self.update_time_dependent_values()
# solve main problem
if Version(dolfinx.__version__) == Version("0.9.0"):
nb_its, _converged = self.solver.solve(self.u)
elif Version(dolfinx.__version__) > Version("0.9.0"):
_ = self.solver.solve()
converged_reason = self.solver.solver.getConvergedReason()
assert converged_reason > 0, (
f"Non-linear solver did not converge. Reason code: {converged_reason}. \n See https://petsc.org/release/manualpages/SNES/SNESConvergedReason/ for more information." # noqa: E501
)
nb_its = self.solver.solver.getIterationNumber()
# post processing
self.post_processing()
# update previous solution
self.u_n.x.array[:] = self.u.x.array[:]
# adapt stepsize
if self.settings.stepsize.adaptive:
new_stepsize = self.settings.stepsize.modify_value(
value=self.dt.value, nb_iterations=nb_its, t=self.t.value
)
self.dt.value = new_stepsize
def update_time_dependent_values(self):
t = float(self.t)
for bc in self.boundary_conditions:
if bc.time_dependent:
bc.update(t=t)
for source in self.sources:
if source.value.explicit_time_dependent:
source.value.update(t=t)
# DEFAULT PETSC OPTIONS
# taken from https://github.com/FEniCS/dolfinx/blob/5fcb988c5b0f46b8f9183bc844d8f533a2130d6a/python/demo/demo_cahn-hilliard.py#L279C1-L286C28
use_superlu = PETSc.IntType == np.int64 # or PETSc.ScalarType == np.complex64
sys = PETSc.Sys() # type: ignore
if sys.hasExternalPackage("mumps") and not use_superlu:
linear_solver = "mumps"
elif sys.hasExternalPackage("superlu_dist"):
linear_solver = "superlu_dist"
else:
linear_solver = "petsc"
_DEFAULT_PETSC_OPTS = {
"snes_type": "newtonls",
"snes_linesearch_type": "none",
"snes_stol": np.sqrt(np.finfo(dolfinx.default_real_type).eps) * 1e-2,
"snes_divergence_tolerance": "PETSC_UNLIMITED",
"ksp_type": "preonly",
"pc_type": "lu",
"pc_factor_mat_solver_type": linear_solver,
"snes_error_if_not_converged": True,
"ksp_error_if_not_converged": True,
}
def get_default_petsc_options():
return _DEFAULT_PETSC_OPTS.copy()
