diff --git a/bindings/sundials4py/CMakeLists.txt b/bindings/sundials4py/CMakeLists.txt index 7b607f2073..84fa8c93dc 100644 --- a/bindings/sundials4py/CMakeLists.txt +++ b/bindings/sundials4py/CMakeLists.txt @@ -116,6 +116,11 @@ set(sundials_SOURCES sunnonlinsol/sunnonlinsol_newton.cpp test/sundials4py_test.cpp) +# Conditionally add KLU bindings +if(SUNDIALS_ENABLE_SUNLINSOL_KLU) + list(APPEND sundials_SOURCES sunlinsol/sunlinsol_klu.cpp) +endif() + # Create the Python sundials library nanobind_add_module(sundials4py STABLE_ABI NB_STATIC ${sundials_SOURCES}) @@ -150,5 +155,9 @@ target_link_libraries( sundials_sundomeigestpower sundials_core) +if(SUNDIALS_ENABLE_SUNLINSOL_KLU) + target_link_libraries(sundials4py PRIVATE sundials_sunlinsolklu) +endif() + # Install directive for scikit-build-core install(TARGETS sundials4py LIBRARY DESTINATION .) diff --git a/bindings/sundials4py/sundials4py.cpp b/bindings/sundials4py/sundials4py.cpp index d23b9bd2b1..670a162dc0 100644 --- a/bindings/sundials4py/sundials4py.cpp +++ b/bindings/sundials4py/sundials4py.cpp @@ -53,6 +53,9 @@ void bind_sundomeigest_power(nb::module_& m); void bind_sunlinsol_band(nb::module_& m); void bind_sunlinsol_dense(nb::module_& m); +#if defined(SUNDIALS_KLU_ENABLED) +void bind_sunlinsol_klu(nb::module_& m); +#endif void bind_sunlinsol_pcg(nb::module_& m); void bind_sunlinsol_spbcgs(nb::module_& m); void bind_sunlinsol_spfgmr(nb::module_& m); @@ -130,6 +133,9 @@ NB_MODULE(sundials4py, m) sundials4py::bind_sunlinsol_band(core_m); sundials4py::bind_sunlinsol_dense(core_m); +#if defined(SUNDIALS_KLU_ENABLED) + sundials4py::bind_sunlinsol_klu(core_m); +#endif sundials4py::bind_sunlinsol_pcg(core_m); sundials4py::bind_sunlinsol_spbcgs(core_m); sundials4py::bind_sunlinsol_spfgmr(core_m); diff --git a/bindings/sundials4py/sunlinsol/generate.yaml b/bindings/sundials4py/sunlinsol/generate.yaml index 9df7de1e27..ed3b0c22ea 100644 --- a/bindings/sundials4py/sunlinsol/generate.yaml +++ b/bindings/sundials4py/sunlinsol/generate.yaml @@ -49,6 +49,10 @@ modules: path: sunlinsol/sunlinsol_dense_generated.hpp headers: - ../../include/sunlinsol/sunlinsol_dense.h + sunlinsol_klu: + path: sunlinsol/sunlinsol_klu_generated.hpp + headers: + - ../../include/sunlinsol/sunlinsol_klu.h sunlinsol_spbcgs: path: sunlinsol/sunlinsol_spbcgs_generated.hpp headers: diff --git a/bindings/sundials4py/sunlinsol/sunlinsol_klu.cpp b/bindings/sundials4py/sunlinsol/sunlinsol_klu.cpp new file mode 100644 index 0000000000..bf558ed678 --- /dev/null +++ b/bindings/sundials4py/sunlinsol/sunlinsol_klu.cpp @@ -0,0 +1,33 @@ +/* ----------------------------------------------------------------- + * Programmer(s): Cody J. Balos @ LLNL + * ----------------------------------------------------------------- + * SUNDIALS Copyright Start + * Copyright (c) 2025-2026, Lawrence Livermore National Security, + * University of Maryland Baltimore County, and the SUNDIALS contributors. + * Copyright (c) 2013-2025, Lawrence Livermore National Security + * and Southern Methodist University. + * Copyright (c) 2002-2013, Lawrence Livermore National Security. + * All rights reserved. + * + * See the top-level LICENSE and NOTICE files for details. + * + * SPDX-License-Identifier: BSD-3-Clause + * SUNDIALS Copyright End + * -----------------------------------------------------------------*/ + +#include "sundials4py.hpp" + +#include +#include + +namespace nb = nanobind; +using namespace sundials::experimental; + +namespace sundials4py { + +void bind_sunlinsol_klu(nb::module_& m) +{ +#include "sunlinsol_klu_generated.hpp" +} + +} // namespace sundials4py diff --git a/bindings/sundials4py/sunlinsol/sunlinsol_klu_generated.hpp b/bindings/sundials4py/sunlinsol/sunlinsol_klu_generated.hpp new file mode 100644 index 0000000000..9bfdb3b56d --- /dev/null +++ b/bindings/sundials4py/sunlinsol/sunlinsol_klu_generated.hpp @@ -0,0 +1,52 @@ +// #ifndef _SUNLINSOL_KLU_H +// +// #ifndef _KLU_H +// +// #endif +// +// #ifdef __cplusplus +// #endif +// + +auto pyClass_SUNLinearSolverContent_KLU = + nb::class_<_SUNLinearSolverContent_KLU>(m, "_SUNLinearSolverContent_KLU", "") + .def(nb::init<>()) // implicit default constructor + ; + +m.def( + "SUNLinSol_KLU", + [](N_Vector y, SUNMatrix A, + SUNContext sunctx) -> std::shared_ptr> + { + auto SUNLinSol_KLU_adapt_return_type_to_shared_ptr = + [](N_Vector y, SUNMatrix A, SUNContext sunctx) + -> std::shared_ptr> + { + auto lambda_result = SUNLinSol_KLU(y, A, sunctx); + + return our_make_shared, + SUNLinearSolverDeleter>(lambda_result); + }; + + return SUNLinSol_KLU_adapt_return_type_to_shared_ptr(y, A, sunctx); + }, + nb::arg("y"), nb::arg("A"), nb::arg("sunctx"), "nb::keep_alive<0, 3>()", + nb::keep_alive<0, 3>()); + +m.def("SUNLinSol_KLUReInit", SUNLinSol_KLUReInit, nb::arg("S"), nb::arg("A"), + nb::arg("nnz"), nb::arg("reinit_type")); + +m.def("SUNLinSol_KLUSetOrdering", SUNLinSol_KLUSetOrdering, nb::arg("S"), + nb::arg("ordering_choice")); + +m.def("SUNLinSol_KLUGetSymbolic", SUNLinSol_KLUGetSymbolic, nb::arg("S")); + +m.def("SUNLinSol_KLUGetNumeric", SUNLinSol_KLUGetNumeric, nb::arg("S")); + +m.def("SUNLinSol_KLUGetCommon", SUNLinSol_KLUGetCommon, nb::arg("S")); +// #ifdef __cplusplus +// +// #endif +// +// #endif +// diff --git a/bindings/sundials4py/test/test_sunlinearsolver.py b/bindings/sundials4py/test/test_sunlinearsolver.py index 2de5f8386e..7ddae9ffeb 100644 --- a/bindings/sundials4py/test/test_sunlinearsolver.py +++ b/bindings/sundials4py/test/test_sunlinearsolver.py @@ -16,6 +16,7 @@ # ----------------------------------------------------------------- import pytest +import sys from fixtures import * from sundials4py.core import * @@ -52,6 +53,41 @@ def test_create_sptfqmr(sunctx, nvec): assert LS is not None +def test_create_klu_if_available(sunctx, nvec): + if not hasattr(sys.modules[__name__], "SUNLinSol_KLU"): + pytest.skip("SUNLinSol_KLU is unavailable in this build") + + sunklu_reinit_full = 1 + sunklu_reinit_partial = 2 + n = N_VGetLength(nvec) + A = SUNSparseMatrix(n, n, n, SUN_CSC_MAT, sunctx) + assert A is not None + + data = SUNSparseMatrix_Data(A) + rowvals = SUNSparseMatrix_IndexValues(A) + colptrs = SUNSparseMatrix_IndexPointers(A) + data[:] = [1.0] * n + rowvals[:] = list(range(n)) + colptrs[:] = list(range(n + 1)) + + LS = SUNLinSol_KLU(nvec, A, sunctx) + assert LS is not None + assert SUNLinSolGetType(LS) == SUNLINEARSOLVER_DIRECT + assert SUNLinSolGetID(LS) == SUNLINEARSOLVER_KLU + + assert SUNLinSol_KLUSetOrdering(LS, 0) == SUN_SUCCESS + assert SUNLinSol_KLUGetCommon(LS) is not None + assert SUNLinSolInitialize(LS) == SUN_SUCCESS + assert SUNLinSolSetup(LS, A) == SUN_SUCCESS + assert SUNLinSol_KLUGetSymbolic(LS) is not None + assert SUNLinSol_KLUGetNumeric(LS) is not None + + assert SUNLinSol_KLUReInit(LS, A, n, sunklu_reinit_partial) == SUN_SUCCESS + assert SUNLinSolSetup(LS, A) == SUN_SUCCESS + assert SUNLinSol_KLUReInit(LS, A, n, sunklu_reinit_full) == SUN_SUCCESS + assert SUNSparseMatrix_NNZ(A) == n + + def test_get_type_and_id(sunctx, nvec): A = SUNDenseMatrix(N_VGetLength(nvec), N_VGetLength(nvec), sunctx) LS = SUNLinSol_Dense(nvec, A, sunctx) diff --git a/doc/shared/Python/Usage.rst b/doc/shared/Python/Usage.rst index 46842a62be..b0ac3183fc 100644 --- a/doc/shared/Python/Usage.rst +++ b/doc/shared/Python/Usage.rst @@ -56,6 +56,14 @@ when running pip. For example: Other SUNDIALS options can also be accessed in this way. Review :numref:`Installation.Options` for more information on the available options. +To build sundials4py with KLU support enabled, pass CMake options for KLU (see :numref:`Installation.Options.KLU`) through ``CMAKE_ARGS`` when building from +source: + +.. code-block:: bash + + export CMAKE_ARGS="-DSUNDIALS_ENABLE_KLU=ON -DKLU_ROOT=/path/to/suitesparse/installation" + pip install sundials4py --no-binary=sundials4py + .. _Python.Usage.Modules: Modules @@ -104,7 +112,7 @@ same capabilities plus continuous forward and adjoint sensitivity analysis. .. note:: Not all SUNDIALS features are supported by the Python interfaces. In - particular, third-party libraries are not yet supported. + particular, third-party libraries are only partially supported. .. _Python.Usage.Example: @@ -116,8 +124,7 @@ highlight some of the differences to using SUNDIALS from C/C++. The items highlighted below similarly apply to using other SUNDIALS packages. For more information on usage differences, continue to the :ref:`next section `. Additional examples can be found in the -``examples/python`` directory of the :examples:`SUNDIALS GitHub repository -`. +``examples/python`` directory of the :examples:`SUNDIALS GitHub repository `. This example demonstrates how to use CVODES to solve the Lotka-Volterra equations, a model of predator-prey dynamics in ecology, given by diff --git a/examples/python/arkode/ark_brusselator1D_klu.py b/examples/python/arkode/ark_brusselator1D_klu.py new file mode 100644 index 0000000000..189debc4e4 --- /dev/null +++ b/examples/python/arkode/ark_brusselator1D_klu.py @@ -0,0 +1,361 @@ +#!/usr/bin/env python3 +# ----------------------------------------------------------------- +# Programmer(s): Cody J. Balos +# ----------------------------------------------------------------- +# SUNDIALS Copyright Start +# Copyright (c) 2025-2026, Lawrence Livermore National Security, +# University of Maryland Baltimore County, and the SUNDIALS contributors. +# Copyright (c) 2013-2025, Lawrence Livermore National Security +# and Southern Methodist University. +# Copyright (c) 2002-2013, Lawrence Livermore National Security. +# All rights reserved. +# +# See the top-level LICENSE and NOTICE files for details. +# +# SPDX-License-Identifier: BSD-3-Clause +# SUNDIALS Copyright End +# ----------------------------------------------------------------- +# This is a direct port of the C example, +# examples/arkode/C_klu/ark_brusselator1D_klu.c +# to use sundials4py. +# ----------------------------------------------------------------- + +import sys +import numpy as np +from sundials4py.core import * +from sundials4py.arkode import * + + +def IDX(x, v): + return 3 * x + v + + +class Brusselator1DProblem: + def __init__(self, N, a, b, du, dv, dw, ep, sunctx): + self.N = N + self.dx = 1.0 / (N - 1) + self.a = a + self.b = b + self.du = du + self.dv = dv + self.dw = dw + self.ep = ep + self.sunctx = sunctx + self.R = None + + def set_init_cond(self, yvec): + y = N_VGetArrayPointer(yvec).reshape((self.N, 3)) + x = np.linspace(0.0, 1.0, self.N) + s = 0.1 * np.sin(np.pi * x) + y[:, 0] = self.a + s + y[:, 1] = self.b / self.a + s + y[:, 2] = self.b + s + return 0 + + def f(self, t, yvec, ydotvec, user_data): + y = N_VGetArrayPointer(yvec).reshape((self.N, 3)) + ydot = N_VGetArrayPointer(ydotvec).reshape((self.N, 3)) + ydot[:, :] = 0.0 + + u = y[:, 0] + v = y[:, 1] + w = y[:, 2] + + uconst = self.du / self.dx / self.dx + vconst = self.dv / self.dx / self.dx + wconst = self.dw / self.dx / self.dx + + ydot[1:-1, 0] = ( + (u[0:-2] - 2.0 * u[1:-1] + u[2:]) * uconst + + self.a + - (w[1:-1] + 1.0) * u[1:-1] + + v[1:-1] * u[1:-1] * u[1:-1] + ) + ydot[1:-1, 1] = ( + (v[0:-2] - 2.0 * v[1:-1] + v[2:]) * vconst + + w[1:-1] * u[1:-1] + - v[1:-1] * u[1:-1] * u[1:-1] + ) + ydot[1:-1, 2] = ( + (w[0:-2] - 2.0 * w[1:-1] + w[2:]) * wconst + + (self.b - w[1:-1]) / self.ep + - w[1:-1] * u[1:-1] + ) + return 0 + + def laplace_matrix(self, A): + N = self.N + colptrs = SUNSparseMatrix_IndexPointers(A) + rowvals = SUNSparseMatrix_IndexValues(A) + data = SUNSparseMatrix_Data(A) + SUNMatZero(A) + + nz = 0 + colptrs[IDX(0, 0)] = nz + colptrs[IDX(0, 1)] = nz + colptrs[IDX(0, 2)] = nz + + uconst = self.du / self.dx / self.dx + vconst = self.dv / self.dx / self.dx + wconst = self.dw / self.dx / self.dx + uconst2 = -2.0 * uconst + vconst2 = -2.0 * vconst + wconst2 = -2.0 * wconst + + for i in range(1, N - 1): + colptrs[IDX(i, 0)] = nz + if i > 1: + data[nz] = uconst + rowvals[nz] = IDX(i - 1, 0) + nz += 1 + data[nz] = uconst2 + rowvals[nz] = IDX(i, 0) + nz += 1 + if i < N - 2: + data[nz] = uconst + rowvals[nz] = IDX(i + 1, 0) + nz += 1 + + colptrs[IDX(i, 1)] = nz + if i > 1: + data[nz] = vconst + rowvals[nz] = IDX(i - 1, 1) + nz += 1 + data[nz] = vconst2 + rowvals[nz] = IDX(i, 1) + nz += 1 + if i < N - 2: + data[nz] = vconst + rowvals[nz] = IDX(i + 1, 1) + nz += 1 + + colptrs[IDX(i, 2)] = nz + if i > 1: + data[nz] = wconst + rowvals[nz] = IDX(i - 1, 2) + nz += 1 + data[nz] = wconst2 + rowvals[nz] = IDX(i, 2) + nz += 1 + if i < N - 2: + data[nz] = wconst + rowvals[nz] = IDX(i + 1, 2) + nz += 1 + + colptrs[IDX(N - 1, 0)] = nz + colptrs[IDX(N - 1, 1)] = nz + colptrs[IDX(N - 1, 2)] = nz + colptrs[IDX(N - 1, 2) + 1] = nz + return 0 + + def reaction_jac(self, yvec, A): + N = self.N + y = N_VGetArrayPointer(yvec).reshape((self.N, 3)) + colptrs = SUNSparseMatrix_IndexPointers(A) + rowvals = SUNSparseMatrix_IndexValues(A) + data = SUNSparseMatrix_Data(A) + SUNMatZero(A) + + nz = 0 + colptrs[IDX(0, 0)] = 0 + colptrs[IDX(0, 1)] = 0 + colptrs[IDX(0, 2)] = 0 + + for i in range(1, N - 1): + u = y[i, 0] + v = y[i, 1] + w = y[i, 2] + + colptrs[IDX(i, 0)] = nz + rowvals[nz] = IDX(i, 0) + data[nz] = 2.0 * u * v - w - 1.0 + nz += 1 + rowvals[nz] = IDX(i, 1) + data[nz] = w - 2.0 * u * v + nz += 1 + rowvals[nz] = IDX(i, 2) + data[nz] = -w + nz += 1 + + colptrs[IDX(i, 1)] = nz + rowvals[nz] = IDX(i, 0) + data[nz] = u * u + nz += 1 + rowvals[nz] = IDX(i, 1) + data[nz] = -u * u + nz += 1 + + colptrs[IDX(i, 2)] = nz + rowvals[nz] = IDX(i, 0) + data[nz] = -u + nz += 1 + rowvals[nz] = IDX(i, 1) + data[nz] = u + nz += 1 + rowvals[nz] = IDX(i, 2) + data[nz] = -1.0 / self.ep - u + nz += 1 + + colptrs[IDX(N - 1, 0)] = nz + colptrs[IDX(N - 1, 1)] = nz + colptrs[IDX(N - 1, 2)] = nz + colptrs[IDX(N - 1, 2) + 1] = nz + return 0 + + def jac(self, t, yvec, fyvec, J, tmp1, tmp2, tmp3, user_data): + status = self.laplace_matrix(J) + if status != 0: + return status + if self.R is None: + self.R = SUNSparseMatrix( + SUNSparseMatrix_Rows(J), + SUNSparseMatrix_Columns(J), + SUNSparseMatrix_NNZ(J), + SUN_CSC_MAT, + self.sunctx, + ) + status = self.reaction_jac(yvec, self.R) + if status != 0: + return status + return SUNMatScaleAdd(1.0, J, self.R) + + +def main(): + if not hasattr(sys.modules[__name__], "SUNLinSol_KLU"): + raise RuntimeError("SUNLinSol_KLU is unavailable in this build") + + T0 = 0.0 + Tf = 10.0 + Nt = 10 + N = 201 + a = 0.6 + b = 2.0 + du = 0.025 + dv = 0.025 + dw = 0.025 + ep = 1.0e-5 + reltol = 1.0e-6 + abstol = 1.0e-10 + + status, sunctx = SUNContext_Create(SUN_COMM_NULL) + assert status == SUN_SUCCESS + NEQ = 3 * N + NNZ = 5 * NEQ + + y = N_VNew_Serial(NEQ, sunctx) + assert y is not None + umask = N_VClone(y) + vmask = N_VClone(y) + wmask = N_VClone(y) + + problem = Brusselator1DProblem(N, a, b, du, dv, dw, ep, sunctx) + problem.set_init_cond(y) + + N_VConst(0.0, umask) + N_VConst(0.0, vmask) + N_VConst(0.0, wmask) + umask_data = N_VGetArrayPointer(umask) + vmask_data = N_VGetArrayPointer(vmask) + wmask_data = N_VGetArrayPointer(wmask) + umask_data.reshape((N, 3))[:, 0] = 1.0 + vmask_data.reshape((N, 3))[:, 1] = 1.0 + wmask_data.reshape((N, 3))[:, 2] = 1.0 + + ark = ARKStepCreate(None, problem.f, T0, y, sunctx) + assert ark is not None + status = ARKodeSStolerances(ark.get(), reltol, abstol) + assert status == ARK_SUCCESS + + A = SUNSparseMatrix(NEQ, NEQ, NNZ, SUN_CSC_MAT, sunctx) + assert A is not None + LS = SUNLinSol_KLU(y, A, sunctx) + assert LS is not None + + status = ARKodeSetLinearSolver(ark.get(), LS, A) + assert status == ARK_SUCCESS + status = ARKodeSetJacFn(ark.get(), problem.jac) + assert status == ARK_SUCCESS + status = ARKodeSetAutonomous(ark.get(), 1) + assert status == ARK_SUCCESS + + print("\n1D Brusselator PDE test problem (KLU solver):") + print(f" N = {N}, NEQ = {NEQ}") + print(f" problem parameters: a = {a}, b = {b}, ep = {ep}") + print(f" diffusion coefficients: du = {du}, dv = {dv}, dw = {dw}") + print(f" reltol = {reltol:.1e}, abstol = {abstol:.1e}\n") + + np.savetxt("bruss_mesh.txt", np.linspace(0.0, 1.0, N), fmt="%.16e") + ydata = N_VGetArrayPointer(y).reshape((N, 3)) + with ( + open("bruss_u.txt", "w") as ufid, + open("bruss_v.txt", "w") as vfid, + open("bruss_w.txt", "w") as wfid, + ): + np.savetxt(ufid, ydata[:, 0][None, :], fmt="%.16e") + np.savetxt(vfid, ydata[:, 1][None, :], fmt="%.16e") + np.savetxt(wfid, ydata[:, 2][None, :], fmt="%.16e") + + t = T0 + dt_out = (Tf - T0) / Nt + tout = T0 + dt_out + print(" t ||u||_rms ||v||_rms ||w||_rms") + print(" ----------------------------------------------") + for _ in range(Nt): + status, t = ARKodeEvolve(ark.get(), tout, y, ARK_NORMAL) + ydata = N_VGetArrayPointer(y).reshape((N, 3)) + u = N_VWL2Norm(y, umask) + u = np.sqrt(u * u / N) + v = N_VWL2Norm(y, vmask) + v = np.sqrt(v * v / N) + w = N_VWL2Norm(y, wmask) + w = np.sqrt(w * w / N) + print(f" {t:10.6f} {u:10.6f} {v:10.6f} {w:10.6f}") + + if status >= 0: + tout = min(tout + dt_out, Tf) + else: + print("Solver failure, stopping integration") + break + + np.savetxt(ufid, ydata[:, 0][None, :], fmt="%.16e") + np.savetxt(vfid, ydata[:, 1][None, :], fmt="%.16e") + np.savetxt(wfid, ydata[:, 2][None, :], fmt="%.16e") + print(" ----------------------------------------------") + + status, nst = ARKodeGetNumSteps(ark.get()) + assert status == ARK_SUCCESS + status, nst_a = ARKodeGetNumStepAttempts(ark.get()) + assert status == ARK_SUCCESS + status, nfe = ARKodeGetNumRhsEvals(ark.get(), 0) + assert status == ARK_SUCCESS + status, nfi = ARKodeGetNumRhsEvals(ark.get(), 1) + assert status == ARK_SUCCESS + status, nsetups = ARKodeGetNumLinSolvSetups(ark.get()) + assert status == ARK_SUCCESS + status, netf = ARKodeGetNumErrTestFails(ark.get()) + assert status == ARK_SUCCESS + status, nni = ARKodeGetNumNonlinSolvIters(ark.get()) + assert status == ARK_SUCCESS + status, ncfn = ARKodeGetNumNonlinSolvConvFails(ark.get()) + assert status == ARK_SUCCESS + status, nje = ARKodeGetNumJacEvals(ark.get()) + assert status == ARK_SUCCESS + + print("\nFinal Solver Statistics:") + print(f" Internal solver steps = {nst} (attempted = {nst_a})") + print(f" Total RHS evals: Fe = {nfe}, Fi = {nfi}") + print(f" Total linear solver setups = {nsetups}") + print(f" Total number of Jacobian evaluations = {nje}") + print(f" Total number of nonlinear iterations = {nni}") + print(f" Total number of nonlinear solver convergence failures = {ncfn}") + print(f" Total number of error test failures = {netf}") + + +def test_ark_brusselator1d_klu(): + if not hasattr(sys.modules[__name__], "SUNLinSol_KLU"): + return + main() + + +if __name__ == "__main__": + main()