Skip to content

Commit

Permalink
black style
Browse files Browse the repository at this point in the history
  • Loading branch information
@enricofacca
enricofacca committed Sep 27, 2024
1 parent 5835df4 commit 548e3b7
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 61 deletions.
92 changes: 45 additions & 47 deletions src/darsia/measure/wasserstein.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -476,7 +476,9 @@ def setup_ksp_solver(
"""
# Define CG solver
self.linear_solver = darsia.linalg.KSP(
matrix, field_ises=field_ises, nullspace=nullspace,
matrix,
field_ises=field_ises,
nullspace=nullspace,
)

# Define solver options
Expand All @@ -503,16 +505,17 @@ def setup_ksp_solver(
"pc_type": "hypre",
}
if self.grid.dim == 3:
self.solver_options.update({
# tuning parameters for the multigrid
# https://mooseframework.inl.gov/releases/moose/2021-09-15/application_development/hypre.html
"pc_hypre_type": "boomeramg",
"pc_hypre_boomeramg_strong_threshold": 0.7,
"pc_hypre_boomeramg_max_iter": 1,
"pc_hypre_boomeramg_agg_nl": 2,
"pc_hypre_boomeramg_interp_type": "ext+i", # "classic" or "ext+i"
}
)
self.solver_options.update(
{
# tuning parameters for the multigrid
# https://mooseframework.inl.gov/releases/moose/2021-09-15/application_development/hypre.html
"pc_hypre_type": "boomeramg",
"pc_hypre_boomeramg_strong_threshold": 0.7,
"pc_hypre_boomeramg_max_iter": 1,
"pc_hypre_boomeramg_agg_nl": 2,
"pc_hypre_boomeramg_interp_type": "ext+i", # "classic" or "ext+i"
}
)
# Include other PETSc options passed from the user
other_petsc_options = linear_solver_options.get("petsc_options", {})
self.solver_options.update(other_petsc_options)
Expand Down Expand Up @@ -1007,8 +1010,7 @@ def _compute_face_weight(self, flat_flux: np.ndarray) -> np.ndarray:
# Determine the l2 norm of the fluxes on the faces
weighted_face_flux = self._product(harm_avg_face_weights, full_face_flux)
norm_weighted_face_flux = np.maximum(
np.linalg.norm(weighted_face_flux, 2, axis=1),
self.regularization
np.linalg.norm(weighted_face_flux, 2, axis=1), self.regularization
)

# Combine weights**2 / |weight * flux| on faces
Expand Down Expand Up @@ -1057,37 +1059,39 @@ def setup_petsc_variables(self, weight: np.ndarray = None):
self.weight_laplacian_vec.set(1.0)
else:
self.weight_laplacian_vec.setArray(weight)



# THIS DOES NOT WORK in the matrix-matrix multiplication
# self.weight_laplacian_matrix = PETSc.Mat().createDiagonal(self.inv_weight_laplacian_vec)
# We need to create a matrix and set the diagonal
self.weight_laplacian_matrix = PETSc.Mat().createAIJ(
size=[n,n],
csr=(np.arange(n + 1, dtype="int32"),
np.arange(n, dtype="int32"),
np.zeros(n))
)
size=[n, n],
csr=(
np.arange(n + 1, dtype="int32"),
np.arange(n, dtype="int32"),
np.zeros(n),
),
)
self.weight_laplacian_matrix.setDiagonal(self.weight_laplacian_vec)

# We store also the grad in order to use the matrix-matrix-matrix multiplication
self.grad_petsc = self.div_petsc.copy()
self.grad_petsc.transpose()



# assign
self.laplacian_matrix = self.div_petsc.matMatMult(self.weight_laplacian_matrix,self.grad_petsc)

self.laplacian_matrix = self.div_petsc.matMatMult(
self.weight_laplacian_matrix, self.grad_petsc
)

def assemble_schur_complement(self, weights: np.ndarray):
"""
Assemble the Schur complement matrix = D^T * W * D
in the Petsc matrix self.weight_laplacian_matrix
"""
self.weight_laplacian_vec.setArray(weights)
self.weight_laplacian_matrix.setDiagonal(self.weight_laplacian_vec)
self.div_petsc.matMatMult(self.weight_laplacian_matrix,self.grad_petsc,result=self.laplacian_matrix)


self.div_petsc.matMatMult(
self.weight_laplacian_matrix, self.grad_petsc, result=self.laplacian_matrix
)

def linear_solve(
self,
Expand Down Expand Up @@ -1175,21 +1179,22 @@ def linear_solve(
# Allocate memory for solution
solution = np.zeros_like(rhs)


# 1. Reduce flux block
# 1. Reduce flux block
if self.linear_solver_type == "ksp":
self.matrix_flux_inv = sps.diags(1.0 / matrix.diagonal()[self.flux_slice])
self.matrix_flux_inv = sps.diags(
1.0 / matrix.diagonal()[self.flux_slice]
)
self.reduced_rhs = rhs[self.pressure_slice].copy()
self.reduced_rhs -= self.div.dot(self.matrix_flux_inv.dot(rhs[self.flux_slice]))
self.reduced_rhs -= self.div.dot(
self.matrix_flux_inv.dot(rhs[self.flux_slice])
)
else:
(
self.reduced_matrix,
self.reduced_rhs,
self.matrix_flux_inv,
) = self.eliminate_flux(matrix, rhs)



# 2. Build linear solver for reduced system
if setup_linear_solver:
if self.linear_solver_type == "direct":
Expand All @@ -1202,16 +1207,15 @@ def linear_solve(
# Update the Schur complement matrix B^T * A^{-1} * B
weight = 1.0 / matrix.diagonal()[self.flux_slice]
if not hasattr(self, "div_petsc"):
self.setup_petsc_variables(weight = weight)
self.setup_petsc_variables(weight=weight)
else:
self.assemble_schur_complement(weights=weight)

# Setup KSP solver
if not hasattr(self, "linear_solver"):
kernel = np.ones(self.grid.num_cells)
kernel /= np.linalg.norm(kernel)
self.setup_ksp_solver(self.laplacian_matrix,
nullspace=[kernel])
self.setup_ksp_solver(self.laplacian_matrix, nullspace=[kernel])
else:
# Just update the matrix
self.linear_solver.ksp.setOperators(self.laplacian_matrix)
Expand All @@ -1222,16 +1226,13 @@ def linear_solve(
# 3. Solve for the pressure and lagrange multiplier
tic = time.time()
if self.linear_solver_type == "ksp":
pot = self.linear_solver.solve(
self.reduced_rhs, **self.solver_options
)
pot = self.linear_solver.solve(self.reduced_rhs, **self.solver_options)
solution[self.pressure_slice] = pot
solution[self.lagrange_multiplier_slice] = 0.0
else:
solution[self.reduced_system_slice] = self.linear_solver.solve(
self.reduced_rhs, **self.solver_options
)


# 4. Compute flux update
solution[self.flux_slice] = self.compute_flux_update(solution, rhs)
Expand Down Expand Up @@ -1267,7 +1268,6 @@ def linear_solve(
"Implementation requires solution satisfy the constraint."
)


# 1. Reduce flux block
(
self.reduced_matrix,
Expand All @@ -1283,7 +1283,7 @@ def linear_solve(
self.reduced_matrix,
self.reduced_rhs,
)

# 3. Build linear solver for pure pressure system
if setup_linear_solver:
if self.linear_solver_type == "direct":
Expand All @@ -1302,10 +1302,8 @@ def linear_solve(
self.setup_ksp_solver(self.fully_reduced_matrix)
else:
self.setup_ksp_solver(self.fully_reduced_matrix)

self.linear_solver.setup(self.solver_options)


self.linear_solver.setup(self.solver_options)

# Stop timer to measure setup time
time_setup = time.time() - tic
Expand Down Expand Up @@ -1354,7 +1352,7 @@ def eliminate_flux(self, jacobian: sps.csc_matrix, residual: np.ndarray) -> tupl

# Gauss eliminiation on matrices
reduced_jacobian = self.jacobian_subblock + schur_complement

# Gauss elimination on vectors
reduced_residual = residual[self.reduced_system_slice].copy()
reduced_residual -= self.D.dot(J_inv.dot(residual[self.flux_slice]))
Expand Down
29 changes: 15 additions & 14 deletions src/darsia/utils/linalg.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,27 +39,28 @@ def numpy_to_petsc(A: sps.csc_matrix) -> PETSc.Mat:
"""Convert a numpy matrix to a PETSc matrix."""
if not sps.isspmatrix_csr(A):
A = A.tocsr()
return PETSc.Mat().createAIJ(
size=A.shape, csr=(A.indptr, A.indices, A.data)
)

def add_diagonal(A: PETSc.Mat, alpha = 0.0) -> None:
return PETSc.Mat().createAIJ(size=A.shape, csr=(A.indptr, A.indices, A.data))

def add_diagonal(A: PETSc.Mat, alpha=0.0) -> None:
"""
Add a zero diagonal to make the matrix the corresponding entries.
Operate in place.
"""
diagonal = PETSc.Mat().createAIJ(
size=A.size, csr=(np.arange(A.size[0] + 1, dtype="int32"),
np.arange(A.size[0], dtype="int32"),
alpha * np.ones(A.size[0]))
size=A.size,
csr=(
np.arange(A.size[0] + 1, dtype="int32"),
np.arange(A.size[0], dtype="int32"),
alpha * np.ones(A.size[0]),
),
)
print(alpha)
A += diagonal

class KSP:
def __init__(
self,
A: Union[sps.csc_matrix,PETSc.Mat],
A: Union[sps.csc_matrix, PETSc.Mat],
field_ises: Optional[
list[tuple[str, Union[PETSc.IS, npt.NDArray[np.int32]]]]
] = None,
Expand Down Expand Up @@ -103,7 +104,7 @@ def __init__(
self.A = A
# convert to petsc matrix
self.A_petsc = numpy_to_petsc(A)
elif isinstance(A, PETSc.Mat):
elif isinstance(A, PETSc.Mat):
self.A_petsc = A
else:
raise ValueError("A must be a scipy or PETSc matrix")
Expand Down Expand Up @@ -242,15 +243,15 @@ def solve(self, b: np.ndarray, **kwargs) -> np.ndarray:
# check if the rhs is orthogonal to the nullspace
rhs_norm = self.rhs_petsc.norm()
rtol = self.petsc_options.get("ksp_rtol", 1e-6)
for i,k in enumerate(self._petsc_kernels):
dot = abs(self.rhs_petsc.dot(k))/rhs_norm
if min(dot,rhs_norm) > rtol:
for i, k in enumerate(self._petsc_kernels):
dot = abs(self.rhs_petsc.dot(k)) / rhs_norm
if min(dot, rhs_norm) > rtol:
raise ValueError(f"RHS not ortogonal. v_{i}^T rhs= {dot=:2e}")

# solve
self.sol_petsc.set(0.0)
self.ksp.solve(self.rhs_petsc, self.sol_petsc)

# check if the solver worked
reason = self.ksp.getConvergedReason()
if reason < 0:
Expand Down

0 comments on commit 548e3b7

Please sign in to comment.