Source code for UM2N.generator.mesh_generator

import os

import firedrake as fd
import movement as mv

from UM2N.generator.equation_solver import EquationSolver

os.environ["OMP_NUM_THREADS"] = "1"
__all__ = ["MeshGenerator"]


[docs] class MeshGenerator: """ Responsible for generating and moving a mesh based on a given Helmholtz equation. This method is based on Pyroteus/movement. Attributes: - eq: The Helmholtz equation object. - num_grid_x: Number of grid points in the x-dimension. - num_grid_y: Number of grid points in the y-dimension. - mesh: The initial m esh. """ def __init__( self, params={ "num_grid_x": None, "num_grid_y": None, "eq": None, "mesh": None, }, ): self.eq = params["eq"] # self.num_grid_x = params["num_grid_x"] # self.num_grid_y = params["num_grid_y"] self.mesh = params["mesh"] self.monitor_val = None self.uh = None self.n_monitor_smooth = 15
[docs] def move_mesh(self): """ Moves the mesh using the Monge-Ampere equation. Computes and stores the Jacobian and its determinant. Returns: - The moved mesh """ mover = mv.MongeAmpereMover( self.mesh, self.monitor_func, method="relaxation", rtol=1e-3, maxiter=500 ) mover.move() # extract Hessian of the movement sigma = mover.sigma I = fd.Identity(2) # noqa jacobian = I + sigma jacobian_det = fd.Function(mover.P1, name="jacobian_det") jacobian_det.project( jacobian[0, 0] * jacobian[1, 1] - jacobian[0, 1] * jacobian[1, 0] ) self.jacob_det = jacobian_det self.jacob = jacobian # extract phi of the movement self.phi = mover.phi # extract phi_grad self.grad_phi = mover.grad_phi return mover.mesh
[docs] def get_grad_phi(self): """ Returns the gradient of phi of the mesh movement. """ return self.grad_phi
[docs] def get_phi(self): """ Returns the phi of the mesh movement. """ return self.phi
[docs] def get_monitor_val(self): """ Returns the monitor function value used for mesh movement. """ return self.monitor_val
[docs] def get_jacobian(self): """ Returns the Jacobian of the mesh movement. """ return self.jacob
[docs] def get_jacobian_det(self): """ Returns the determinant of the Jacobian of the mesh movement. """ return self.jacob_det
[docs] def get_gradient(self, mesh): res = self.eq.discretise(mesh) # function_space_ten = fd.TensorFunctionSpace(mesh, "CG", 1) solver = EquationSolver( params={ "LHS": res["LHS"], "RHS": res["RHS"], "function_space": res["function_space"], "bc": res["bc"], } ) uh = solver.solve_eq() self.uh = uh func_vec_space = fd.VectorFunctionSpace(mesh, "CG", 1) uh_grad = fd.interpolate(fd.grad(self.uh), func_vec_space) return uh_grad
[docs] def get_grad_norm(self, mesh): res = self.eq.discretise(mesh) uh_grad = self.get_gradient(mesh) grad_norm = fd.Function(res["function_space"]) grad_norm.project(uh_grad[0] ** 2 + uh_grad[1] ** 2) grad_norm /= grad_norm.vector().max() return grad_norm
[docs] def get_hessian(self, mesh): """ Computes and returns the Hessian of the Helmholtz equation on the given mesh. Parameters: - mesh: The mesh on which to compute the Hessian. Returns: - The Hessian as a projection in the function space. """ res = self.eq.discretise(mesh) function_space_ten = fd.TensorFunctionSpace(mesh, "CG", 1) solver = EquationSolver( params={ "LHS": res["LHS"], "RHS": res["RHS"], "function_space": res["function_space"], "bc": res["bc"], } ) uh = solver.solve_eq() self.uh = uh n = fd.FacetNormal(mesh) l2_projection = fd.Function(function_space_ten) H, h = fd.TrialFunction(function_space_ten), fd.TestFunction(function_space_ten) a = fd.inner(h, H) * fd.dx(domain=mesh) L = -fd.inner(fd.div(h), fd.grad(uh)) * fd.dx(domain=mesh) L += fd.dot(fd.grad(uh), fd.dot(h, n)) * fd.ds(domain=mesh) prob = fd.LinearVariationalProblem(a, L, l2_projection) hessian_prob = fd.LinearVariationalSolver(prob) hessian_prob.solve() return l2_projection
[docs] def get_hessian_norm(self, mesh): res = self.eq.discretise(mesh) function_space = res["function_space"] hessian_norm = fd.Function(function_space) l2_projection = self.get_hessian(mesh) hessian_norm.project( l2_projection[0, 0] ** 2 + l2_projection[0, 1] ** 2 + l2_projection[1, 0] ** 2 + l2_projection[1, 1] ** 2 ) hessian_norm /= hessian_norm.vector().max() return hessian_norm
[docs] def monitor_func(self, mesh): """ Computes the monitor function value based on the Hessian of the Helmholtz equation. Parameters: - mesh: The mesh on which to compute the monitor function. Returns: - The monitor function value. """ res = self.eq.discretise(mesh) function_space = res["function_space"] hessian_norm = fd.Function(function_space) l2_projection = self.get_hessian(mesh) hessian_norm.project( l2_projection[0, 0] ** 2 + l2_projection[0, 1] ** 2 + l2_projection[1, 0] ** 2 + l2_projection[1, 1] ** 2 ) hessian_norm /= hessian_norm.vector().max() raw_monitor_val = 1 + 5 * hessian_norm monitor_val = fd.Function(function_space) monitor_val.assign(raw_monitor_val) return monitor_val
# def monitor_func(self, mesh, alpha=10, beta=5): # # self.project_u_() # # self.solve_u(self.t) # # self.u_hess.project(self.u_cur) # # self.hessian_prob.solve() # # self.f_norm.project( # # self.l2_projection[0, 0] ** 2 # # + self.l2_projection[0, 1] ** 2 # # + self.l2_projection[1, 0] ** 2 # # + self.l2_projection[1, 1] ** 2 # # ) # res = self.eq.discretise(mesh) # function_space = res["function_space"] # hessian_norm = fd.Function(function_space) # l2_projection = self.get_hessian(mesh) # hessian_norm.project( # l2_projection[0, 0] ** 2 # + l2_projection[0, 1] ** 2 # + l2_projection[1, 0] ** 2 # + l2_projection[1, 1] ** 2 # ) # func_vec_space = fd.VectorFunctionSpace(self.mesh, "CG", 1) # uh_grad = fd.interpolate(fd.grad(self.uh), func_vec_space) # self.grad_norm = fd.Function(function_space) # self.grad_norm.project(uh_grad[0] ** 2 + uh_grad[1] ** 2) # # Normlize the hessian # self.hessian_norm /= self.hessian_norm.vector().max() # # Normlize the grad # self.grad_norm /= self.grad_norm.vector().max() # self.monitor_val = fd.Function(function_space) # # Choose the max values between grad norm and hessian norm according to # # [Clare et al 2020] Multi-scale hydro-morphodynamic modelling using mesh movement methods # self.monitor_val.dat.data[:] = np.maximum( # beta * self.hessian_norm.dat.data[:], alpha * self.grad_norm.dat.data[:] # ) # # ################# # V = fd.FunctionSpace(mesh, "CG", 1) # u = fd.TrialFunction(V) # v = fd.TestFunction(V) # function_space = V # # Discretised Eq Definition Start # f = self.monitor_val # dx = mesh.cell_sizes.dat.data[:].mean() # N = self.n_monitor_smooth # K = N * dx**2 / 4 # RHS = f * v * fd.dx(domain=mesh) # LHS = (K * fd.dot(fd.grad(v), fd.grad(u)) + v * u) * fd.dx(domain=mesh) # bc = fd.DirichletBC(function_space, f, "on_boundary") # monitor_smoothed = fd.Function(function_space) # fd.solve( # LHS == RHS, # monitor_smoothed, # solver_parameters={"ksp_type": "cg", "pc_type": "none"}, # bcs=bc, # ) # # ################# # monitor_val = 1 + monitor_smoothed # self.monitor_val.assign(monitor_val) # return monitor_val