diff --git a/quagmire/equation_systems/diffusion.py b/quagmire/equation_systems/diffusion.py index 93beded..d1b5265 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,15 @@ def __init__(self, self.dirichlet_mask = dirichlet_mask + # 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 + + @property def mesh(self): return self._mesh @@ -146,13 +159,132 @@ 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) + + 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() + + 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 + 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() + + 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, rhs): + + 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) + + + # 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. + + + 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 +319,108 @@ 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 = 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, 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) + + self._mesh.dm.localToGlobal(self.phi._ldata, PHI) + + elapsed_time = 0.0 + + for step in range(0, int(steps)): + + # 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 + + 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() + RHS = self._RHS + PHI = self._PHI + + + # set boundary conditions + self._set_boundary_conditions(mat_LHS, RHS) + + # initialise solver + self._initialise_solver(mat_LHS) + + + self._mesh.dm.localToGlobal(self.phi._ldata, PHI) + + + self.ksp.solve(RHS, PHI) + self._mesh.dm.globalToLocal(PHI, 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: