Skip to content

Commit

Permalink
black formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
enricofacca committed Sep 30, 2024
1 parent f71553c commit 2c4f5c1
Showing 1 changed file with 23 additions and 27 deletions.
50 changes: 23 additions & 27 deletions src/darsia/measure/wasserstein.py
Original file line number Diff line number Diff line change
Expand Up @@ -1044,8 +1044,6 @@ def optimality_conditions(
- self.broken_darcy.dot(solution)
- self.flux_embedding.dot(weight.dot(self.mass_matrix_faces.dot(flat_flux)))
)



# ! ---- Solver methods ----
def setup_petsc_variables(self, weight: np.ndarray = None):
Expand All @@ -1055,9 +1053,9 @@ def setup_petsc_variables(self, weight: np.ndarray = None):
# Setup Petsc operators
self.div_petsc = darsia.linalg.numpy_to_petsc(self.div)
n = self.grid.num_faces

# THE FOLLOWING DOES NOT WORK in the matrix-matrix multiplication
# PETSc.Mat().createDiagonal(self.inv_weight_laplacian_vec)
# PETSc.Mat().createDiagonal(self.inv_weight_laplacian_vec)
if weight is None:
ones = np.ones(n)
scipy_sparse = sps.diags(ones)
Expand All @@ -1067,10 +1065,8 @@ def setup_petsc_variables(self, weight: np.ndarray = None):
self.weight_laplacian_vec = darsia.linalg.numpy_to_petsc(weight)
self.weight_laplacian_matrix = darsia.linalg.numpy_to_petsc(scipy_sparse)


# We need to create a matrix and set the diagonal



self.weight_laplacian_matrix.setDiagonal(self.weight_laplacian_vec)

# We store also the grad in order to use the matrix-matrix-matrix multiplication
Expand Down Expand Up @@ -1195,13 +1191,15 @@ def linear_solve(
# Compute the inverse of the flux block
self.matrix_flux_inv = self.compute_invA(matrix)
# form the schur complement
self.reduced_matrix = self.compute_Schur_complement(self.matrix_flux_inv)
self.reduced_matrix = self.compute_Schur_complement(
self.matrix_flux_inv
)

# create the reduced rhs
if self.linear_solver_type == "ksp":
self.reduced_rhs = rhs[self.pressure_slice].copy()
self.reduced_rhs -= self.div.dot(
self.matrix_flux_inv.dot(rhs[self.flux_slice])
self.matrix_flux_inv.dot(rhs[self.flux_slice])
)
else:
self.reduced_rhs = self.compute_reduced_rhs(self.matrix_flux_inv, rhs)
Expand Down Expand Up @@ -1271,21 +1269,22 @@ def linear_solve(
raise NotImplementedError(
"Implementation requires solution satisfy the constraint."
)

if setup_linear_solver:
# 1. Reduce flux block
# Compute the inverse of the flux block
self.matrix_flux_inv = self.compute_invA(matrix)
# form the schur complement
self.reduced_matrix = self.compute_Schur_complement(self.matrix_flux_inv)
self.reduced_matrix = self.compute_Schur_complement(
self.matrix_flux_inv
)

self.reduced_rhs = self.compute_reduced_rhs(self.matrix_flux_inv, rhs)
self.fully_reduced_rhs = self.eliminate_lagrange_multiplier_rhs(
self.reduced_rhs
)

if setup_linear_solver:

# 2. Reduce to pure pressure system
self.fully_reduced_matrix = self.eliminate_lagrange_multiplier_matrix(
self.reduced_matrix
Expand Down Expand Up @@ -1364,30 +1363,31 @@ def eliminate_flux(self, jacobian: sps.csc_matrix, residual: np.ndarray) -> tupl
reduced_residual -= self.D.dot(J_inv.dot(residual[self.flux_slice]))

return reduced_jacobian, reduced_residual, J_inv

def compute_reduced_rhs(self, J_inv: sps.csc_matrix, residual: np.ndarray) -> np.ndarray:

def compute_reduced_rhs(
self, J_inv: sps.csc_matrix, residual: np.ndarray
) -> np.ndarray:
"""
Compute the reduced right hand side.
"""
# Gauss elimination on vectors
reduced_residual = residual[self.reduced_system_slice].copy()
reduced_residual -= self.D.dot(J_inv.dot(residual[self.flux_slice]))

return reduced_residual

def compute_invA(self, jacobian: sps.csc_matrix) -> sps.csc_matrix:
"""
Set the inverse of the flux block
"""
return sps.diags(1.0 / jacobian.diagonal()[self.flux_slice])

def compute_Schur_complement(self, invA: sps.csc_matrix) -> sps.csc_matrix:
"""
Compute the Schur complement, including the lagrange multiplier block
"""
# Build Schur complement wrt flux-block
return self.jacobian_subblock + self.D.dot(invA.dot(self.DT))

return self.jacobian_subblock + self.D.dot(invA.dot(self.DT))

def eliminate_lagrange_multiplier(
self, reduced_jacobian, reduced_residual
Expand Down Expand Up @@ -1427,10 +1427,8 @@ def eliminate_lagrange_multiplier(
].copy()

return self.fully_reduced_jacobian, fully_reduced_residual

def eliminate_lagrange_multiplier_rhs(
self, reduced_residual
) -> tuple:

def eliminate_lagrange_multiplier_rhs(self, reduced_residual) -> tuple:
"""Eliminate the lagrange multiplier from the reduced system.
Employ a Schur complement/block Gauss elimination approach.
Expand All @@ -1452,10 +1450,8 @@ def eliminate_lagrange_multiplier_rhs(
].copy()

return fully_reduced_residual

def eliminate_lagrange_multiplier_matrix(
self, reduced_jacobian
) -> tuple:

def eliminate_lagrange_multiplier_matrix(self, reduced_jacobian) -> tuple:
"""Eliminate the lagrange multiplier from the reduced system.
Employ a Schur complement/block Gauss elimination approach.
Expand Down

0 comments on commit 2c4f5c1

Please sign in to comment.