compiler: Utilise petsctools and start extending solver parameters functionality

examples: Add example using argparse in conjunction with PetscInitialize

Author: ZoeLeibowitz

devito/petsc/utils.py

@@ -4,6 +4,8 @@
 from pathlib import Path
 from devito.tools import memoized_func
 
+from petsctools import get_petscvariables
+
 
 class PetscOSError(OSError):
     pass
@@ -51,30 +53,7 @@ def core_metadata():
     }
 
 
-@memoized_func
-def get_petsc_variables():
-    """
-    Taken from https://www.firedrakeproject.org/_modules/firedrake/petsc.html
-    Get a dict of PETSc environment variables from the file:
-    $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables
-    """
-    try:
-        petsc_dir = get_petsc_dir()
-    except PetscOSError:
-        petsc_variables = {}
-    else:
-        path = [petsc_dir[-1], 'lib', 'petsc', 'conf', 'petscvariables']
-        variables_path = Path(*path)
-
-        with open(variables_path) as fh:
-            # Split lines on first '=' (assignment)
-            splitlines = (line.split("=", maxsplit=1) for line in fh.readlines())
-        petsc_variables = {k.strip(): v.strip() for k, v in splitlines}
-
-    return petsc_variables
-
-
-petsc_variables = get_petsc_variables()
+petsc_variables = get_petscvariables()
 
 
 def get_petsc_type_mappings():

examples/petsc/solver_options.py

@@ -0,0 +1,51 @@
+import os
+import numpy as np
+
+from devito import (Grid, Function, Eq, Operator, configuration,
+                    switchconfig)
+from devito.petsc import PETScSolve
+from devito.petsc.initialize import PetscInitialize
+import petsctools
+from petsctools import get_commandline_options
+configuration['compiler'] = 'custom'
+os.environ['CC'] = 'mpicc'
+
+
+# Initialize petsctools
+# petsctools.init()
+
+from argparse import ArgumentParser
+
+parser = ArgumentParser()
+
+parser.add_argument('--nx', type=int, default=11)
+parser.add_argument('--ny', type=int, default=11)
+
+args, unknown = parser.parse_known_args()
+
+PetscInitialize(unknown)
+
+grid = Grid(shape=(args.nx, args.ny), extent=(2., 2.), dtype=np.float64)
+
+u = Function(name='u', grid=grid, dtype=np.float64, space_order=2)
+v = Function(name='v', grid=grid, dtype=np.float64, space_order=2)
+
+v.data[:] = 5.0
+
+eq = Eq(v, u.laplace, subdomain=grid.interior)
+
+petsc = PETScSolve([eq], u)
+
+with switchconfig(language='petsc'):
+    op = Operator(petsc)
+    op.apply()
+
+# print(op.ccode)
+
+# print(grid.shape)
+
+
+import sys
+petsctools.options._commandline_options = sys.argv[1:]
+tmp = get_commandline_options()
+print("Command line options:", tmp)

requirements.txt

@@ -8,4 +8,5 @@ codepy>=2019.1,<2025
 multidict<6.3
 anytree>=2.4.3,<=2.13.0
 cloudpickle<3.1.2
-packaging<25.1
+packaging<25.1
+petsctools