diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 22c82ab8..eb3bb827 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -4,11 +4,10 @@ on: push: branches: - main - - development pull_request: branches: - main - - development + # schedule: # Do a nightly run of the tests # - cron: '0 1 * * *' @@ -40,48 +39,12 @@ jobs: architecture: x64 - name: Install dependencies + if: github.event_name == 'push' run: | sudo apt-get update sudo apt-get install -y libhypre-dev libmumps-seq-dev pip install numpy mpi4py - pip install petsc petsc4py - - - name: Set environment variables - run: | - echo "PETSC_DIR=/usr/local/petsc" >> $GITHUB_ENV - echo "PETSC_ARCH=arch-linux-c-opt" >> $GITHUB_ENV - - # Installation of petsc in parallel - - # - name: Install dependencies - # run: | - # sudo apt-get update - # sudo apt-get install -y mpich - # pip install numpy mpi4py - # pip install petsc petsc4py - - # - name: Set environment variables - # run: | - # echo "PETSC_DIR=/usr/local/petsc" >> $GITHUB_ENV - # echo "PETSC_ARCH=arch-linux-c-opt" >> $GITHUB_ENV - - # - name: Install Hypre - # run: | - # git clone https://github.com/hypre-space/hypre.git - # cd hypre/src - # ./configure --with-MPI - # make - # sudo make install - - # - name: Install MUMPS - # run: | - # git clone https://github.com/scivision/mumps.git - # cd mumps - # mkdir build - # cd build - # cmake -G "Ninja" -DMUMPS_parallel=on .. - # cmake --build . - # sudo make install + PETSC_CONFIGURE_OPTIONS="--download-hypre --download-mumps --download-parmetis --download-ml --download-metis --download-scalapack" pip install petsc petsc4py - name: Install DarSIA run: pip install .[dev] @@ -104,4 +67,4 @@ jobs: - name: pytest if: ${{always()}} - run: pytest + run: pytest \ No newline at end of file diff --git a/conda_env.yaml b/conda_env.yaml new file mode 100644 index 00000000..66a09ca6 --- /dev/null +++ b/conda_env.yaml @@ -0,0 +1,34 @@ +# Use the command +# +# conda env create --file conda_env.yaml +# +# to create a conda enviroment named darcia. Using +# conda-force, we get petsc and petsc4py with hypre and mumps. +# +name: darcia +channels: + - conda-forge +dependencies: + - python=3.12 + - pytest + - numpy + - scipy + - pyamg + - matplotlib==3.8.0 + - scikit-image + - future + - colour-science + - ipympl + - pillow + - scikit-learn + - pandas + - plotly + - pydicom + - mpi4py + - petsc + - petsc4py + - pip + - pip: + - largestinteriorrectangle + - opencv-python + - colour-checker-detection \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..6987c038 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[build-system] +requires = [ + "cython >= 3", + "numpy", + "setuptools", +] +build-backend = "setuptools.build_meta" diff --git a/requirements.txt b/requirements.txt index dfa541af..6825448c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ +#Cython>=3.0.0 numpy opencv-python scipy @@ -13,3 +14,5 @@ pandas plotly largestinteriorrectangle pydicom +#petsc4py==3.19.0 # from 3.20. there is requirement for Cython>=3.0.0 +pyamg diff --git a/src/darsia/measure/wasserstein.py b/src/darsia/measure/wasserstein.py index 1c95e25f..97bbaf51 100644 --- a/src/darsia/measure/wasserstein.py +++ b/src/darsia/measure/wasserstein.py @@ -114,6 +114,7 @@ def __init__( - "direct": direct solver - "amg": algebraic multigrid solver - "cg": conjugate gradient solver preconditioned with AMG + - "ksp": PETSc KSP solver - formulation (str): formulation of the linear system. Defaults to "pressure". Supported formulations are: - "full": full system @@ -183,16 +184,16 @@ def _setup_dof_management(self) -> None: num_dofs = num_flux_dofs + num_pressure_dofs + num_lagrange_multiplier_dofs # ! ---- Indices in global system ---- - self.flux_indices = np.arange(num_flux_dofs) + self.flux_indices = np.arange(num_flux_dofs, dtype=np.int32) """np.ndarray: indices of the fluxes""" self.pressure_indices = np.arange( - num_flux_dofs, num_flux_dofs + num_pressure_dofs + num_flux_dofs, num_flux_dofs + num_pressure_dofs, dtype=np.int32 ) """np.ndarray: indices of the pressures""" self.lagrange_multiplier_indices = np.array( - [num_flux_dofs + num_pressure_dofs], dtype=int + [num_flux_dofs + num_pressure_dofs], dtype=np.int32 ) """np.ndarray: indices of the lagrange multiplier""" @@ -314,6 +315,7 @@ def _setup_linear_solver(self) -> None: "direct", "amg", "cg", + "ksp", ], f"Linear solver {self.linear_solver_type} not supported." assert self.formulation in [ "full", @@ -321,6 +323,12 @@ def _setup_linear_solver(self) -> None: "pressure", ], f"Formulation {self.formulation} not supported." + if self.linear_solver_type == "ksp": + if self.formulation == "flux-reduced": + raise ValueError( + "KSP solver only supports for full and pressure formulation." + ) + # Setup inrastructure for Schur complement reduction if self.formulation == "flux_reduced": self.setup_eliminate_flux() @@ -443,6 +451,90 @@ def setup_cg_solver(self, matrix: sps.csc_matrix) -> None: } """dict: options for the iterative linear solver""" + def setup_ksp_solver( + self, + matrix: sps.csc_matrix, + field_ises: Optional[list[tuple[str, np.ndarray]]] = None, + nullspace: Optional[list[np.ndarray]] = None, + ) -> None: + """Setup an KSP solver from PETSc for the given matrix. + + Args: + matrix (sps.csc_matrix): matrix + nullspace (list[np.ndarray]): list of nullspace vectors of the matrix + + Defines: + PETSc.ksp: KSP solver + dict: options for the KSP solver + """ + # Define CG solver + self.linear_solver = darsia.linalg.KSP( + matrix, field_ises=field_ises, nullspace=nullspace + ) + + # Define solver options + linear_solver_options = self.options.get("linear_solver_options", {}) + tol = linear_solver_options.get("tol", 1e-6) + maxiter = linear_solver_options.get("maxiter", 100) + approach = linear_solver_options.get("approach", "direct") + + if field_ises is None: + if approach == "direct": + self.solver_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + } + else: + prec = linear_solver_options.get("pc_type", "hypre") + self.solver_options = { + "ksp_type": approach, + # "ksp_monitor_true_residual": None, + "ksp_rtol": tol, + "ksp_max_it": maxiter, + "pc_type": prec, + } + else: + if approach == "direct": + self.solver_options = { + "ksp_type": "preonly", # do not apply Krylov iterations + "pc_type": "lu", + "pc_factor_shift_type": "inblocks", # for the zero entries + "pc_factor_mat_solver_type": "mumps", + } + else: + prec = linear_solver_options.get("pc_type", "hypre") + # Block preconditioning approach + # the the 0 and 1 in fieldsplit_0 and fieldsplit_1 are the strings + # passed to the field_ises + self.solver_options = { + "ksp_type": approach, + "ksp_rtol": tol, + "ksp_max_it": maxiter, + # "ksp_monitor_true_residual": None, #this is for debugging + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_fact_type": "full", + # use a full factorization of the Schur complement + # other options are "diag","lower","upper" + "pc_fieldsplit_schur_precondition": "selfp", + # selfp -> form an approximate Schur complement + # using S=-B diag(A)^{-1} B^T + # which is the exact Schur complement in our case + # https://petsc.org/release/manualpages/PC/PCFieldSplitSetSchurPre/ + "fieldsplit_flux_ksp_type": "preonly", + "fieldsplit_flux_pc_type": "jacobi", + # use the diagonal of the flux (it is the inverse) + "fieldsplit_pressure": { + "ksp_type": "preonly", + "pc_type": prec, + }, + # an example of the nested dictionary + } + + self.linear_solver.setup(self.solver_options) + """dict: options for the iterative linear solver""" + def setup_eliminate_flux(self) -> None: """Setup the infrastructure for reduced systems through Gauss elimination. @@ -636,7 +728,6 @@ def cell_weighted_flux(self, cell_flux: np.ndarray) -> np.ndarray: raise NotImplementedError("Need to apply matrix vector product.") else: - raise NotImplementedError("Dimension not supported.") def transport_density( @@ -954,21 +1045,52 @@ def linear_solve( """ setup_linear_solver = not (reuse_solver) or not (hasattr(self, "linear_solver")) - if self.formulation == "full": assert ( - self.linear_solver_type == "direct" - ), "Only direct solver supported for full formulation." + self.linear_solver_type == "direct" or self.linear_solver_type == "ksp" + ), "Only direct solver or ksp supported for full formulation." # Setup LU factorization for the full system tic = time.time() if setup_linear_solver: - self.setup_direct_solver(matrix) + if self.linear_solver_type == "direct": + self.setup_direct_solver(matrix) + elif self.linear_solver_type == "ksp": + if hasattr(self, "linear_solver"): + self.linear_solver.kill() + # Extract get the flux-pressure matrix + # TODO: Avoid all these conversions and memory allocations + diag = matrix.diagonal() + A_00 = sps.diags(diag[self.flux_slice]) + flux_pressure_matrix = sps.bmat( + [ + [A_00, -self.div.T], + [self.div, None], + ], + format="csc", + ) + kernel = np.zeros(flux_pressure_matrix.shape[0]) + kernel[self.pressure_slice] = 1.0 + kernel /= np.linalg.norm(kernel) + self.setup_ksp_solver( + flux_pressure_matrix, + field_ises=[ + ("flux", self.flux_indices), + ("pressure", self.pressure_indices), + ], + nullspace=[kernel], + ) time_setup = time.time() - tic # Solve the full system tic = time.time() - solution = self.linear_solver.solve(rhs) + if self.linear_solver_type == "direct": + solution = self.linear_solver.solve(rhs) + elif self.linear_solver_type == "ksp": + solution_flux_pressure = self.linear_solver.solve(rhs[0:-1]) + solution = np.zeros_like(rhs) + solution[0:-1] = solution_flux_pressure + solution[-1] = 0.0 time_solve = time.time() - tic elif self.formulation == "flux_reduced": @@ -1070,6 +1192,21 @@ def linear_solve( elif self.linear_solver_type == "cg": self.setup_cg_solver(self.fully_reduced_matrix) + elif self.linear_solver_type == "ksp": + if not hasattr(self, "linear_solver"): + self.setup_ksp_solver(self.fully_reduced_matrix) + + if reuse_solver: + self.linear_solver.setup(self.solver_options) + + # Free memory if solver needs to be re-setup + if hasattr(self, "linear_solver"): + if not reuse_solver: + self.linear_solver.kill() + self.setup_ksp_solver(self.fully_reduced_matrix) + else: + self.setup_ksp_solver(self.fully_reduced_matrix) + # Stop timer to measure setup time time_setup = time.time() - tic diff --git a/src/darsia/utils/linalg.py b/src/darsia/utils/linalg.py index ba818780..cdbc1143 100644 --- a/src/darsia/utils/linalg.py +++ b/src/darsia/utils/linalg.py @@ -1,6 +1,9 @@ """Wrapper to Krylov subspace methods from SciPy.""" +from typing import Optional, Union + import numpy as np +import numpy.typing as npt import scipy.sparse as sps from scipy.sparse.linalg import cg, gmres @@ -19,3 +22,323 @@ def __init__(self, A: sps.csc_matrix) -> None: def solve(self, b: np.ndarray, **kwargs) -> np.ndarray: return gmres(self.A, b, **kwargs)[0] + + +# Define PETSc solver if petsc4py is available +try: + from petsc4py import PETSc + + def _make_reasons(reasons): + return dict( + [(getattr(reasons, r), r) for r in dir(reasons) if not r.startswith("_")] + ) + + KSPreasons = _make_reasons(PETSc.KSP.ConvergedReason()) + + class KSP: + def __init__( + self, + A: sps.csc_matrix, + field_ises: Optional[ + list[tuple[str, Union[PETSc.IS, npt.NDArray[np.int32]]]] + ] = None, + nullspace: Optional[list[np.ndarray]] = None, + appctx: dict = None, + ) -> None: + """ + KSP solver for PETSc matrices + + Parameters + ---------- + A : sps.csc_matrix + Matrix to solve + field_ises : Optional[list[tuple[str,PETSc.IS]]], optional + Fields index sets, by default None. + This tells how to partition the matrix in blocks for field split + (block-based) preconditioners. + + Example with IS: + is_0 = PETSc.IS().createStride(size=3, first=0, step=1) + is_1 = PETSc.IS().createStride(size=3, first=3, step=1) + [('0',is_0),('1',is_1)] + Example with numpy array: + [('flux',np.array([0,1,2],np.int32)),('pressure',np.array([3,4,5],np.int32))] + nullspace : Optional[list[np.ndarray]], optional + Nullspace vectors, by default None + appctx : dict, optional + Application context, by default None. + It is attached to the KSP object to gather information that can be used + to form the preconditioner. + """ + + # convert csc to csr (if needed) + if not sps.isspmatrix_csr(A): + A = A.tocsr() + + # store (a pointer to) the original matrix + self.A = A + + # convert to petsc matrix + self.A_petsc = PETSc.Mat().createAIJ( + size=A.shape, csr=(A.indptr, A.indices, A.data) + ) + + # This step is needed to avoid a petsc error with LU factorization + # that explicitly needs a zeros along the diagonal + # TODO: use zero_diagonal = PETSc.Mat().createConstantDiagonal(self.A.size,0.0) + zero_diagonal = PETSc.Mat().createAIJ( + size=A.shape, + csr=( + np.arange(A.shape[0] + 1, dtype="int32"), + np.arange(A.shape[0], dtype="int32"), + np.zeros(A.shape[0]), + ), + ) + self.A_petsc += zero_diagonal + + # Convert kernel np array. to petsc nullspace + # TODO: ensure orthogonality of the nullspace vectors + self._petsc_kernels = [] + if nullspace is not None: + # convert to petsc vectors + self._comm = PETSc.COMM_WORLD + for v in nullspace: + p_vec = PETSc.Vec().createWithArray(v, comm=self._comm) + self._petsc_kernels.append(p_vec) + self._nullspace = PETSc.NullSpace().create( + constant=None, vectors=self._petsc_kernels, comm=self._comm + ) + self.A_petsc.setNullSpace(self._nullspace) + + # preallocate petsc vectors + self.sol_petsc = self.A_petsc.createVecLeft() + self.rhs_petsc = self.A_petsc.createVecRight() + + # set field_ises + if field_ises is not None: + if isinstance(field_ises[0][1], PETSc.IS): + self.field_ises = field_ises + if isinstance(field_ises[0][1], np.ndarray): + self.field_ises = [ + (i, PETSc.IS().createGeneral(is_i)) for i, is_i in field_ises + ] + else: + self.field_ises = None + + self.appctx = appctx + + def setup(self, petsc_options: dict) -> None: + petsc_options = flatten_parameters(petsc_options) + + self.prefix = "petsc_solver_" + # TODO: define a unique name in case of multiple problems + + # create ksp solver and assign controls + ksp = PETSc.KSP().create() + ksp.setOperators(self.A_petsc) + ksp.setOptionsPrefix(self.prefix) + ksp.setConvergenceHistory() + ksp.setFromOptions() + opts = PETSc.Options() + opts.prefixPush(self.prefix) + for k, v in petsc_options.items(): + opts[k] = v + opts.prefixPop() + ksp.setConvergenceHistory() + ksp.setFromOptions() + + self.ksp = ksp + + # associate petsc vectors to prefix + self.A_petsc.setOptionsPrefix(self.prefix) + self.A_petsc.setFromOptions() + + self.sol_petsc.setOptionsPrefix(self.prefix) + self.sol_petsc.setFromOptions() + + self.rhs_petsc.setOptionsPrefix(self.prefix) + self.rhs_petsc.setFromOptions() + + # TODO: set field split in __init__ + if (self.field_ises) is not None and ( + "fieldsplit" in petsc_options["pc_type"] + ): + pc = ksp.getPC() + pc.setFromOptions() + # syntax is + # pc.setFieldSplitIS(('0',is_0),('1',is_1)) + pc.setFieldSplitIS(*self.field_ises) + pc.setOptionsPrefix(self.prefix) + pc.setFromOptions() + pc.setUp() + + # split subproblems + ksps = pc.getFieldSplitSubKSP() + for k in ksps: + # Without this, the preconditioner is not set up + # It works for now, but it is not clear why + p = k.getPC() + # This is in order to pass the appctx + # to all preconditioner. TODO: find a better way + p.setAttr("appctx", self.appctx) + p.setUp() + + # set nullspace + if self._nullspace is not None: + if len(self._petsc_kernels) > 1: + raise NotImplementedError( + "Nullspace currently works with one kernel only" + ) + # assign nullspace to each subproblem + for i, local_ksp in enumerate(ksps): + for k in self._petsc_kernels: + sub_vec = k.getSubVector(self.field_ises[i][1]) + if sub_vec.norm() > 0: + local_nullspace = PETSc.NullSpace().create( + constant=False, vectors=[sub_vec] + ) + A_i, _ = local_ksp.getOperators() + A_i.setNullSpace(local_nullspace) + + # attach info to ksp + self.ksp.setAttr("appctx", self.appctx) + + def solve(self, b: np.ndarray, **kwargs) -> np.ndarray: + """ + Solve the linear system + + Parameters + ---------- + b : np.ndarray + Right hand side + + Returns + ------- + np.ndarray + Solution of the linear system + """ + + # convert to petsc vector the rhs + self.rhs_petsc.setArray(b) + + # check if the rhs is orthogonal to the nullspace + for k in self._petsc_kernels: + dot = abs(self.rhs_petsc.dot(k)) + if dot > 1e-10: + raise ValueError("RHS not ortogonal to the nullspace") + + # solve + self.ksp.solve(self.rhs_petsc, self.sol_petsc) + + # check if the solver worked + reason = self.ksp.getConvergedReason() + if reason < 0: + raise ValueError(f"KSP solver failed {reason=} {KSPreasons[reason]}") + + # convert to numpy array + sol = self.sol_petsc.getArray() + + # DEBUG CODE: check convergence in numpy varaibles + # res = self.A.dot(sol) - b + # print('residual norm: ', np.linalg.norm(res)/np.linalg.norm(b)) + return sol + + def kill(self): + """ + Free memory + """ + self.ksp.destroy() + self.A_petsc.destroy() + self.sol_petsc.destroy() + self.rhs_petsc.destroy() + if hasattr(self, "_nullspace"): + self._nullspace.destroy() + for k in self._petsc_kernels: + k.destroy() + self.A_petsc = None + self.sol_petsc = None + self.rhs_petsc = None + self.ksp = None + self._nullspace = None + self._petsc_kernels = None + + # + # Following code is taken from Firedrake + # + def flatten_parameters(parameters, sep="_"): + """Flatten a nested parameters dict, joining keys with sep. + + :arg parameters: a dict to flatten. + :arg sep: separator of keys. + + Used to flatten parameter dictionaries with nested structure to a + flat dict suitable to pass to PETSc. For example: + + .. code-block:: python3 + + flatten_parameters({"a": {"b": {"c": 4}, "d": 2}, "e": 1}, sep="_") + => {"a_b_c": 4, "a_d": 2, "e": 1} + + If a "prefix" key already ends with the provided separator, then + it is not used to concatenate the keys. Hence: + + .. code-block:: python3 + + flatten_parameters({"a_": {"b": {"c": 4}, "d": 2}, "e": 1}, sep="_") + => {"a_b_c": 4, "a_d": 2, "e": 1} + # rather than + => {"a__b_c": 4, "a__d": 2, "e": 1} + """ + new = type(parameters)() + + if not len(parameters): + return new + + def flatten(parameters, *prefixes): + """Iterate over nested dicts, yielding (*keys, value) pairs.""" + sentinel = object() + try: + option = sentinel + for option, value in parameters.items(): + # Recurse into values to flatten any dicts. + for pair in flatten(value, option, *prefixes): + yield pair + # Make sure zero-length dicts come back. + if option is sentinel: + yield (prefixes, parameters) + except AttributeError: + # Non dict values are just returned. + yield (prefixes, parameters) + + def munge(keys): + """Ensure that each intermediate key in keys ends in sep. + + Also, reverse the list.""" + for key in reversed(keys[1:]): + if len(key) and not key.endswith(sep): + yield key + sep + else: + yield key + else: + yield keys[0] + + for keys, value in flatten(parameters): + option = "".join(map(str, munge(keys))) + if option in new: + print( + "Ignoring duplicate option: %s (existing value %s, new value %s)", + option, + new[option], + value, + ) + new[option] = value + return new + +except ImportError: + + class KSP: + def __init__( + self, + **kwargs, + ) -> None: + raise ImportError("petsc4py not found. PETSc solver not available.") diff --git a/tests/unit/test_wasserstein.py b/tests/unit/test_wasserstein.py index ae38473e..b69c428d 100644 --- a/tests/unit/test_wasserstein.py +++ b/tests/unit/test_wasserstein.py @@ -5,6 +5,13 @@ import darsia +try: + import petsc4py + + HAVE_PETSC = True +except ImportError: + HAVE_PETSC = False + # ! ---- 2d version ---- # Coarse src image @@ -111,14 +118,54 @@ lu_options = { # Linear solver "linear_solver": "direct", + "formulation": "full", } amg_options = { "linear_solver": "amg", "linear_solver_options": { "tol": 1e-8, }, + "formulation": "pressure", +} +ksp_direct_options = { + "linear_solver": "ksp", + "linear_solver_options": { + "approach": "direct", + }, + "formulation": "pressure", +} +ksp_krylov_options = { + "linear_solver": "ksp", + "linear_solver_options": { + "tol": 1e-8, + "approach": "cg", + "pc_type": "hypre", + }, + "formulation": "pressure", } -solvers = [lu_options, amg_options] + +ksp_block_krylov_options = { + "linear_solver": "ksp", + "linear_solver_options": { + "tol": 1e-8, + "approach": "gmres", + "pc_type": "hypre", + }, + "formulation": "full", +} + + +solvers = ( + [lu_options, amg_options] + + [ + ksp_direct_options, + ksp_krylov_options, + ksp_block_krylov_options, + ] + if HAVE_PETSC + else [] +) + # General options options = {