Add Cython bindings for the PETSc API

.gitignore

@@ -1,7 +1,12 @@
 petsctools/config.ini
 
-*/__pycache__
+**/__pycache__
 *.egg-info
 
 docs/build
 docs/source/generated
+
+docs/source/_static/cython-demo/build
+docs/source/_static/cython-demo/*.c
+docs/source/_static/cython-demo/*.so
+docs/source/_static/cython-demo/*.html

docs/source/_static/cython-demo/fast.pyx

@@ -0,0 +1,36 @@
+import time
+
+import cython
+from petsc4py import PETSc
+
+from petsctools cimport cpetsc
+
+
+def medium():
+    N: cython.int = int(1e8)
+    section: PETSc.Section = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    i: cython.int
+    for i in range(N):
+        if i % 2 == 0:
+            section.setDof(i, 1)
+    print(f"Time elapsed: {time.time() - start}")
+
+
+def fast():
+    N: cython.int = int(1e8)
+    section: cpetsc.PetscSection_py = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    i: cython.int
+    for i in range(N):
+        if i % 2 == 0:
+            cpetsc.CHKERR(cpetsc.PetscSectionSetDof(section.sec, i, 1))
+    print(f"Time elapsed: {time.time() - start}")
+
+
+medium()
+fast()

docs/source/_static/cython-demo/setup.py

@@ -0,0 +1,63 @@
+from Cython.Build import cythonize
+from setuptools import setup, Extension
+
+import petsctools
+import os
+import petsc4py
+
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ExternalDependency:
+    ''' This dataclass stores the relevant information for the compiler as fields
+    that correspond to the keyword arguments of `Extension`. For convenience it
+    also implements addition and `**` unpacking.
+    '''
+    include_dirs: list[str] = field(default_factory=list, init=True)
+    extra_compile_args: list[str] = field(default_factory=list, init=True)
+    libraries: list[str] = field(default_factory=list, init=True)
+    library_dirs: list[str] = field(default_factory=list, init=True)
+    extra_link_args: list[str] = field(default_factory=list, init=True)
+    runtime_library_dirs: list[str] = field(default_factory=list, init=True)
+
+    def __add__(self, other):
+        combined = {}
+        for f in self.__dataclass_fields__.keys():
+            combined[f] = getattr(self, f) + getattr(other, f)
+        return self.__class__(**combined)
+
+    def keys(self):
+        return self.__dataclass_fields__.keys()
+
+    def __getitem__(self, key):
+        try:
+            return getattr(self, key)
+        except AttributeError:
+            raise KeyError(f"Key {key} not present")
+
+
+
+petsc_dir = petsctools.get_petsc_dir()
+petsc_arch = petsctools.get_petsc_arch()
+petsc_dirs = [petsc_dir, os.path.join(petsc_dir, petsc_arch)]
+petsc_ = ExternalDependency(
+    libraries=["petsc"],
+    include_dirs=[petsc4py.get_include()] + [os.path.join(d, "include") for d in petsc_dirs],
+    library_dirs=[os.path.join(petsc_dirs[-1], "lib")],
+    runtime_library_dirs=[os.path.join(petsc_dirs[-1], "lib")],
+)
+
+mods = [
+        Extension(
+        name="cython_demo",
+        language="c",
+        sources=[os.path.join("cython_demo.pyx")],
+        **(petsc_),
+        annotate=True,
+    )
+]
+
+
+setup(ext_modules=mods)

docs/source/_static/cython-demo/slow.py

@@ -0,0 +1,18 @@
+import time
+
+from petsc4py import PETSc
+
+
+def slow():
+    N = int(1e8)
+    section = PETSc.Section().create()
+    section.setChart(0, N)
+
+    start = time.time()
+    for i in range(N):
+        if i % 2 == 0:
+            section.setDof(i, 1)
+    print(f"Time elapsed: {time.time() - start}")
+
+
+slow()

petsctools/cpetsc.pxd

@@ -0,0 +1,60 @@
+"""This file basically exposes the PETSc API as a module for use in Cython."""
+
+# IMPORTANT: This file cannot be accessed if petsctools is installed in editable mode.
+
+from petsc4py cimport PETSc as _PETSc
+
+# clearer aliases from petsc4py, so the names here match the C API
+ctypedef _PETSc.PetscMat Mat
+ctypedef _PETSc.Mat Mat_py
+ctypedef _PETSc.PetscSF PetscSF
+ctypedef _PETSc.SF PetscSF_py
+ctypedef _PETSc.PetscSection PetscSection
+ctypedef _PETSc.Section PetscSection_py
+ctypedef _PETSc.PetscIS IS
+ctypedef _PETSc.IS IS_py
+
+# other PETSc imports
+from petsc4py.PETSc cimport (
+    CHKERR,
+    PetscErrorCode,
+)
+
+
+cdef extern from "petsc.h":
+    # fundamental types
+    ctypedef long PetscInt
+    ctypedef double PetscReal
+    ctypedef double PetscScalar
+    ctypedef enum PetscBool:
+        PETSC_TRUE
+        PETSC_FALSE
+    ctypedef enum InsertMode:
+        INSERT_VALUES
+        ADD_VALUES
+    ctypedef enum PetscCopyMode:
+        PETSC_COPY_VALUES
+        PETSC_OWN_POINTER
+        PETSC_USE_POINTER
+
+    # memory management
+    PetscErrorCode PetscCalloc1(size_t,void*)
+    PetscErrorCode PetscMalloc1(size_t,void*)
+    PetscErrorCode PetscFree(void*)
+
+    # Mat
+    PetscErrorCode MatSetValue(Mat,PetscInt,PetscInt,const PetscScalar,InsertMode)
+    PetscErrorCode MatSetValuesBlockedLocal(Mat,PetscInt,const PetscInt[],PetscInt,const PetscInt[],const PetscScalar[],InsertMode)
+
+    # PetscSF
+    ctypedef struct PetscSFNode:
+        PetscInt rank
+        PetscInt index
+
+    PetscErrorCode PetscSFGetGraph(PetscSF,PetscInt*,PetscInt*,const PetscInt**,const PetscSFNode**)
+    PetscErrorCode PetscSFSetGraph(PetscSF,PetscInt,PetscInt,PetscInt*,PetscCopyMode,PetscSFNode*,PetscCopyMode)
+
+    # PetscSection
+    PetscErrorCode PetscSectionGetDof(PetscSection,PetscInt,PetscInt*)
+    PetscErrorCode PetscSectionSetDof(PetscSection,PetscInt,PetscInt)
+    PetscErrorCode PetscSectionGetOffset(PetscSection,PetscInt,PetscInt*)

pyproject.toml

@@ -41,6 +41,9 @@ ci = [
   {include-group = "test"},
 ]
 
+[tool.setuptools.package-data]
+petsctools = ["__init__.pxd", "cpetsc.pxd"]
+
 [tool.ruff]
 line-length = 79