From 7b2760c239602b4d44c1ed2e3b60d493ae4f3807 Mon Sep 17 00:00:00 2001 From: Ben Mather Date: Wed, 1 Jul 2020 21:08:45 +1000 Subject: [PATCH 1/3] preliminary crank nicholson scheme --- quagmire/equation_systems/diffusion.py | 208 ++++++++++++++++++++++++- 1 file changed, 206 insertions(+), 2 deletions(-) diff --git a/quagmire/equation_systems/diffusion.py b/quagmire/equation_systems/diffusion.py index 93beded..1465bcd 100644 --- a/quagmire/equation_systems/diffusion.py +++ b/quagmire/equation_systems/diffusion.py @@ -22,6 +22,8 @@ from mpi4py import MPI comm = MPI.COMM_WORLD +from petsc4py import PETSc + # To do ... an interface for (iteratively) dealing with # boundaries with normals that are not aligned to the coordinates @@ -47,7 +49,8 @@ def __init__(self, dirichlet_mask=None, neumann_x_mask=None, neumann_y_mask=None, - non_linear=False ): + non_linear=False, + theta=0.5 ): self.__id = "diffusion_solver_{}".format(self._count()) self.mesh = mesh @@ -58,6 +61,7 @@ def __init__(self, self._neumann_y_mask = None self._dirichlet_mask = None self._non_linear = False + self.theta = theta if diffusivity_fn is not None: self.diffusivity = diffusivity_fn @@ -72,6 +76,14 @@ def __init__(self, self.dirichlet_mask = dirichlet_mask + # create some global work vectors + self._RHS = mesh.dm.createGlobalVec() + self._PHI = mesh.dm.createGlobalVec() + + # store this flag to make sure verify has been called before timestepping + self._verified = False + + @property def mesh(self): return self._mesh @@ -146,13 +158,120 @@ def non_linear(self, bool_object): return + @property + def theta(self): + return self._theta + @theta.setter + def theta(self, val): + assert val <= 1 and val >= 0, "theta must be within the range [0,1]" + self._theta = float(val) + + # trigger an update to the LHS / RHS matrices? + + + def construct_matrix(self): + + mesh = self._mesh + kappa = self.diffusivity.evaluate(mesh) + + nnz = mesh.natural_neighbours_count.astype(PETSc.IntType) + + # initialise matrix object + mat = PETSc.Mat().create(comm=comm) + mat.setType('aij') + mat.setSizes(mesh.sizes) + mat.setLGMap(mesh.lgmap_row, mesh.lgmap_col) + mat.setFromOptions() + mat.setPreallocationNNZ(nnz) + # mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, False) + + + # vectorise along the matrix stencil + mat.assemblyBegin() + + nodes = np.arange(0, mesh.npoints+1, dtype=PETSc.IntType) + + # need this to prevent mallocs + mat.setValuesLocalCSR(nodes, nodes[:-1], np.zeros(mesh.npoints)) + + for col in range(1, mesh.natural_neighbours.shape[1]): + node_neighbours = mesh.natural_neighbours[:,col].astype(PETSc.IntType) + + node_distance = np.linalg.norm(mesh.data - mesh.data[node_neighbours], axis=1) + node_distance[node_distance == 0] += 0.000001 # do not divide by zero + delta = 1.0/(2.0*node_distance**2) + + vals = delta*(kappa + kappa[node_neighbours]) + + mat.setValuesLocalCSR(nodes, node_neighbours, vals) + + mat.assemblyEnd() + + diag = mat.getRowSum() + diag.scale(-1.0) + mat.setDiagonal(diag) + + ## NOTE: This matrix can be used to solve steady state diffusion + ## we augment it for timestepping (and to specify Dirichlet BCs) + + return mat - def verify(self): + def _set_boundary_conditions(self, matrix): + + mesh = self._mesh + + dirichlet_mask = self.dirichlet_mask.evaluate(mesh).astype(bool) + # neumann_mask = self.neumann_x_mask.evaluate(mesh).astype(bool) + # neumann_mask += self.neumann_y_mask.evaluate(mesh).astype(bool) + + if dirichlet_mask.any(): + # augment the matrix + + # put zeros where we have the Dirichlet mask + mesh.lvec.setArray(np.invert(dirichlet_mask).astype(np.float)) + mesh.dm.localToGlobal(mesh.lvec, mesh.gvec) + matrix.diagonalScale(mesh.gvec) + + # now put ones along the diagonal + diag = matrix.getDiagonal() + mesh.dm.globalToLocal(diag, mesh.lvec) + lvec_data = mesh.lvec.array + lvec_data[dirichlet_mask] = 1.0 + mesh.lvec.setArray(lvec_data) + mesh.dm.localToGlobal(mesh.lvec, diag) + matrix.setDiagonal(diag) + + + # neumann BCs are by design of the matrix, + # so nothing needs to be done with those masks. + + + def _initialise_solver(self, matrix, ksp_type='gmres', atol=1e-20, rtol=1e-50, pc_type=None, **kwargs): + """ + Initialise linear solver object + """ + + ksp = PETSc.KSP().create(comm) + ksp.setType(ksp_type) + ksp.setOperators(matrix) + ksp.setTolerances(atol, rtol) + if pc_type is not None: + pc = ksp.getPC() + pc.setType(pc_type) + ksp.setFromOptions() + self.ksp = ksp + + + def verify(self, **kwargs): """Verify solver is ready for launch""" # Check to see we have provided everything in the correct form + # store the matrix + self.matrix = self.construct_matrix() + + self._verified = True return @@ -187,9 +306,94 @@ def diffusion_timestep(self): return global_diff_timestep + def _get_scaled_matrix(self, scale): + mat = self.matrix.copy() + mat.scale(scale) + diag = mat.getDiagonal() + diag += 1.0 + mat.setDiagonal(diag) + return mat + + def time_integration(self, timestep, steps=1, Delta_t=None, feedback=None): + if Delta_t is not None: + steps = Delta_t // timestep + timestep = Delta_t / steps + + if not self._verified: + self.verify() + + + # Create LHS and RHS matrices + scale_LHS = timestep*self.theta + scale_RHS = timestep*(1.0 - self.theta) + + mat_LHS = self._get_scaled_matrix(-scale_LHS) # LHS matrix + mat_RHS = self._get_scaled_matrix(scale_RHS) # RHS matrix + + + # set boundary conditions + self._set_boundary_conditions(mat_LHS) + self._set_boundary_conditions(mat_RHS) + + + # initialise solver + self._initialise_solver(mat_LHS) + + + RHS = self._RHS + PHI = self._PHI + self._mesh.dm.localToGlobal(self.phi._ldata, PHI) + + + elapsed_time = 0.0 + + for step in range(0, int(steps)): + + mat_RHS.mult(PHI, RHS) # multiply RHS + self.ksp.solve(RHS, PHI) # solve SLE + + elapsed_time += timestep + + if feedback is not None and step%feedback == 0 or step == steps: + print("{:05d} - t = {:.3g}".format(step, elapsed_time)) + + + self._mesh.dm.globalToLocal(PHI, self.phi._ldata) + + return steps, elapsed_time + + + def steady_state(self, feedback=None): + + if not self._verified: + self.verify() + + mat_LHS = self.matrix.copy() + + + # set boundary conditions + self._set_boundary_conditions(mat_LHS) + + # initialise solver + self._initialise_solver(mat_LHS) + + + RHS = self._RHS + PHI = self._PHI + self._mesh.dm.localToGlobal(self.phi._ldata, PHI) + + + self.ksp.solve(PHI, RHS) + self._mesh.dm.globalToLocal(RHS, self.phi._ldata) + return self.phi.data + + + + def time_integration_explicit(self, timestep, steps=1, Delta_t=None, feedback=None): + from quagmire import function as fn if Delta_t is not None: From 4a24223094b08a279bb5c1ca4aab05b4d1a39764 Mon Sep 17 00:00:00 2001 From: Ben Mather Date: Thu, 2 Jul 2020 15:40:06 +1000 Subject: [PATCH 2/3] enforce BCs and fix timestep size --- quagmire/equation_systems/diffusion.py | 54 ++++++++++++++++++-------- 1 file changed, 38 insertions(+), 16 deletions(-) diff --git a/quagmire/equation_systems/diffusion.py b/quagmire/equation_systems/diffusion.py index 1465bcd..8affd72 100644 --- a/quagmire/equation_systems/diffusion.py +++ b/quagmire/equation_systems/diffusion.py @@ -79,6 +79,7 @@ def __init__(self, # create some global work vectors self._RHS = mesh.dm.createGlobalVec() self._PHI = mesh.dm.createGlobalVec() + self._RHS_outer = mesh.dm.createGlobalVec() # store this flag to make sure verify has been called before timestepping self._verified = False @@ -217,7 +218,7 @@ def construct_matrix(self): return mat - def _set_boundary_conditions(self, matrix): + def _set_boundary_conditions(self, matrix, rhs): mesh = self._mesh @@ -243,6 +244,13 @@ def _set_boundary_conditions(self, matrix): matrix.setDiagonal(diag) + # set the RHS vector + lvec_data.fill(0.0) + lvec_data[dirichlet_mask] = self.phi.data[dirichlet_mask] + mesh.lvec.setArray(lvec_data) + mesh.dm.localToGlobal(mesh.lvec, rhs) + + # neumann BCs are by design of the matrix, # so nothing needs to be done with those masks. @@ -327,33 +335,47 @@ def time_integration(self, timestep, steps=1, Delta_t=None, feedback=None): # Create LHS and RHS matrices - scale_LHS = timestep*self.theta - scale_RHS = timestep*(1.0 - self.theta) + scale_LHS = 0.5*timestep*self.theta + scale_RHS = 0.5*timestep*(1.0 - self.theta) mat_LHS = self._get_scaled_matrix(-scale_LHS) # LHS matrix mat_RHS = self._get_scaled_matrix(scale_RHS) # RHS matrix + PHI = self._PHI + RHS = self._RHS + BCS = self._RHS_outer + # set boundary conditions - self._set_boundary_conditions(mat_LHS) - self._set_boundary_conditions(mat_RHS) + self._set_boundary_conditions(mat_LHS, BCS) + self._set_boundary_conditions(mat_RHS, BCS) + BCS.set(0.0) + + dirichlet_mask = self.dirichlet_mask.evaluate(self._mesh).astype(bool) + dirichlet_BC = self.phi.data[dirichlet_mask].copy() # initialise solver self._initialise_solver(mat_LHS) - - RHS = self._RHS - PHI = self._PHI self._mesh.dm.localToGlobal(self.phi._ldata, PHI) - elapsed_time = 0.0 for step in range(0, int(steps)): - mat_RHS.mult(PHI, RHS) # multiply RHS - self.ksp.solve(RHS, PHI) # solve SLE + # enforce Dirichlet BCs + self._mesh.dm.globalToLocal(PHI, self._mesh.lvec) + self._mesh.lvec.array[dirichlet_mask] = dirichlet_BC + self._mesh.dm.localToGlobal(self._mesh.lvec, PHI) + + + # evaluate right hand side + mat_RHS.mult(PHI, RHS) + RHS += BCS # add BCs + + # find solution inside domain + self.ksp.solve(RHS, PHI) elapsed_time += timestep @@ -372,22 +394,22 @@ def steady_state(self, feedback=None): self.verify() mat_LHS = self.matrix.copy() + RHS = self._RHS + PHI = self._PHI # set boundary conditions - self._set_boundary_conditions(mat_LHS) + self._set_boundary_conditions(mat_LHS, RHS) # initialise solver self._initialise_solver(mat_LHS) - RHS = self._RHS - PHI = self._PHI self._mesh.dm.localToGlobal(self.phi._ldata, PHI) - self.ksp.solve(PHI, RHS) - self._mesh.dm.globalToLocal(RHS, self.phi._ldata) + self.ksp.solve(RHS, PHI) + self._mesh.dm.globalToLocal(PHI, self.phi._ldata) return self.phi.data From d9148b3713c022080faba8e94e500eb640ab41b1 Mon Sep 17 00:00:00 2001 From: Ben Mather Date: Wed, 8 Jul 2020 18:04:19 +1000 Subject: [PATCH 3/3] improvements to area weights - still not perfect..? --- quagmire/equation_systems/diffusion.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/quagmire/equation_systems/diffusion.py b/quagmire/equation_systems/diffusion.py index 8affd72..d1b5265 100644 --- a/quagmire/equation_systems/diffusion.py +++ b/quagmire/equation_systems/diffusion.py @@ -186,6 +186,11 @@ def construct_matrix(self): mat.setPreallocationNNZ(nnz) # mat.setOption(mat.Option.NEW_NONZERO_ALLOCATION_ERR, False) + area = mesh.pointwise_area.evaluate(mesh) + neighbour_area = area[mesh.natural_neighbours] + neighbour_area[~mesh.natural_neighbours_mask] = 0.0 # mask non-neighbours + total_neighbour_area = neighbour_area.sum(axis=1) + # vectorise along the matrix stencil mat.assemblyBegin() @@ -197,13 +202,13 @@ def construct_matrix(self): for col in range(1, mesh.natural_neighbours.shape[1]): node_neighbours = mesh.natural_neighbours[:,col].astype(PETSc.IntType) - + node_distance = np.linalg.norm(mesh.data - mesh.data[node_neighbours], axis=1) node_distance[node_distance == 0] += 0.000001 # do not divide by zero - delta = 1.0/(2.0*node_distance**2) - - vals = delta*(kappa + kappa[node_neighbours]) - + weight = 3*area[node_neighbours]/(total_neighbour_area[node_neighbours]/3) + + vals = weight*0.5*(kappa + kappa[node_neighbours])/node_distance**2 + mat.setValuesLocalCSR(nodes, node_neighbours, vals) mat.assemblyEnd()