From 92edf688a666ecaffe28fff4bbd8c452e9589e69 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 24 Apr 2025 18:02:30 +0100 Subject: [PATCH 01/53] dsl/compiler: Edits to mixed solver functionality --- devito/passes/iet/definitions.py | 2 +- devito/passes/iet/engine.py | 14 +- devito/passes/iet/misc.py | 3 +- devito/petsc/iet/routines.py | 226 ++++++++++++++++++++++++------- devito/petsc/solve.py | 139 +++++++++++++++---- devito/petsc/types/array.py | 80 ++++++++++- devito/petsc/types/object.py | 7 + devito/petsc/types/types.py | 60 +++++--- devito/types/basic.py | 1 + 9 files changed, 424 insertions(+), 108 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index e56c86645f..f5ae337b9c 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -379,7 +379,7 @@ def place_definitions(self, iet, globs=None, **kwargs): # created by the compiler up to this point (Array, LocalObject, etc.) storage = Storage() defines = FindSymbols('defines-aliases|globals').visit(iet) - + # from IPython import embed; embed() for i in FindSymbols().visit(iet): if i in defines: continue diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index a88cf58d45..992096fd41 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -148,7 +148,10 @@ def apply(self, func, **kwargs): # Minimize code size if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) - efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) + # TODO: fix for petsc bundles + # TODO: somethinng to do with this bug is causing the compiler not to combine loops with petsc bundles + # from IPython import embed; embed() + # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -391,9 +394,9 @@ def abstract_efunc(efunc): - Objects are renamed as "o0", "o1", ... """ functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) - + # from IPython import embed; embed() mapper = abstract_objects(functions) - + # from IPython import embed; embed() efunc = Uxreplace(mapper).visit(efunc) efunc = efunc._rebuild(name='foo') @@ -408,9 +411,10 @@ def abstract_objects(objects0, sregistry=None): objects = [] for i in objects0: if i.is_Bundle: + # from IPython import embed; embed() objects.extend(i.components) objects.append(i) - + # from IPython import embed; embed() # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] @@ -420,7 +424,9 @@ def abstract_objects(objects0, sregistry=None): # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() + # from IPython import embed; embed() for i in objects: + # from IPython import embed; embed() abstract_object(i, mapper, sregistry) return mapper diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index e404a8e373..0f15deec81 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -184,6 +184,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): try: v = tracker[i.base].v subs[i] = v.func(v.base, *i.indices) + # from IPython import embed; embed() continue except KeyError: pass @@ -193,7 +194,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs[i] = v tracker[i.base] = Bunch(header=header, v=v) - + # from IPython import embed; embed() iet = Uxreplace(subs).visit(iet) return iet diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 68765d7801..be76ceace1 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -4,15 +4,16 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, - retrieve_iteration_tree, filter_iterations, Iteration) + retrieve_iteration_tree, filter_iterations, Iteration, + PointerCast, Expression, Transformer) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod) + FieldFromComposite, IntDiv, Deref, Mod, uxreplace) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension from devito.tools import filter_ordered -from devito.petsc.types import PETScArray +from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct @@ -20,7 +21,7 @@ from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, - SubMatrixStruct, CallbackDM) + SubMatrixStruct, CallbackDM, PetscComponentAccess) class CBBuilder: @@ -100,20 +101,20 @@ def user_struct_callback(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata, fielddata.matvecs) + self._make_matvec(fielddata.arrays, fielddata.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) if fielddata.initialguess: self._make_initialguess(fielddata) self._make_user_struct_callback() - def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): + def _make_matvec(self, arrays, matvecs, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), - fielddata) + arrays) objs = self.objs cb = PETScCallable( @@ -125,7 +126,7 @@ def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): self._matvecs.append(cb) self._efuncs[cb.name] = cb - def _create_matvec_body(self, body, fielddata): + def _create_matvec_body(self, body, arrays): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -134,8 +135,8 @@ def _create_matvec_body(self, body, fielddata): ctx = objs['dummyctx'] xlocal = objs['xloc'] ylocal = objs['yloc'] - y_matvec = fielddata.arrays['y'] - x_matvec = fielddata.arrays['x'] + y_matvec = arrays['y'] + x_matvec = arrays['x'] body = self.timedep.uxreplace_time(body) @@ -652,18 +653,16 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - data = all_fielddata.get_field_data(t) - self._make_formfunc(data) - self._make_formrhs(data) - + row_matvecs = all_fielddata.submatrices.submatrices[t] + arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): if mtvs['matvecs']: - self._make_matvec(data, mtvs['matvecs'], prefix=f'{submat}_MatMult') + self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') self._make_user_struct_callback() self._make_whole_matvec() - self._make_whole_formfunc() + self._make_whole_formfunc(all_fielddata) self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -726,56 +725,175 @@ def _whole_matvec_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def _make_whole_formfunc(self): + def _make_whole_formfunc(self, fielddata): + formfuncs = fielddata.formfuncs + # from IPython import embed; embed() + # Compile formfunc `eqns` into an IET via recursive compilation + irs_formfunc, _ = self.rcompile( + formfuncs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body), + fielddata) objs = self.objs - body = self._whole_formfunc_body() - cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), - List(body=body), + body_formfunc, + # List(body=()), retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) self._main_formfunc_callback = cb self._efuncs[cb.name] = cb - def _whole_formfunc_body(self): + def _whole_formfunc_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - ljacctx = objs['ljacctx'] - struct_cast = DummyExpr(ljacctx, JacobianStructCast(objs['dummyptr'])) - X = objs['X'] - F = objs['F'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] - targets = self.injectsolve.expr.rhs.fielddata.targets + body = self.timedep.uxreplace_time(body) - deref_subdms = Dereference(objs['LocalSubdms'], ljacctx) - deref_fields = Dereference(objs['LocalFields'], ljacctx) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) - calls = () - for i, t in enumerate(targets): - field_ptr = FieldFromPointer(objs['LocalFields'].indexed[i], ljacctx) - x_name = f'Xglobal{t.name}' - f_name = f'Fglobal{t.name}' - calls += ( - petsc_call('VecGetSubVector', [X, field_ptr, Byref(sobjs[x_name])]), - petsc_call('VecGetSubVector', [F, field_ptr, Byref(sobjs[f_name])]), - petsc_call(self.formfuncs[i].name, [ - objs['snes'], sobjs[x_name], sobjs[f_name], - VOID(objs['LocalSubdms'].indexed[i], stars='*') - ]), - petsc_call('VecRestoreSubVector', [X, field_ptr, Byref(sobjs[x_name])]), - petsc_call('VecRestoreSubVector', [F, field_ptr, Byref(sobjs[f_name])]), - ) - return CallableBody( - List(body=calls + (BlankLine,)), - init=(objs['begin_user'],), - stacks=(struct_cast, deref_subdms, deref_fields), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + targets = fielddata.targets + arrays = fielddata.arrays + + f_u = arrays[targets[0]]['f'] + f_v = arrays[targets[1]]['f'] + + x_u = arrays[targets[0]]['x'] + x_v = arrays[targets[1]]['x'] + + fbundle = PetscBundle(name='f_vu', components=(f_v, f_u)) + xbundle = PetscBundle(name='x_vu', components=(x_v, x_u)) + + mapper1 = {x_u.base: xbundle.base, x_v.base: xbundle.base} + + indexeds = FindSymbols('indexeds').visit(body) + + subss = {} + for i in indexeds: + if i.base in mapper1: + bundle = mapper1[i.base] + subss[i] = i.func(bundle, *i.indices) + # subss[i] = bundle.indexed.__getitem__((1,)+i.indices) + # subss[i] = subss[i].__getitem__((0,)+i.indices) + subss[i] = PetscComponentAccess(subss[i], 1) + + # subss[i] = xbundle.access_component() + + # from IPython import embed; embed() + + body = Uxreplace(subss).visit(body) + + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, objs['xloc']] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['X'], insert_vals, objs['xloc'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] + ) + + vec_get_array_f = petsc_call( + 'VecGetArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + vec_restore_array_f = petsc_call( + 'VecRestoreArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, objs['floc'], insert_vals, objs['F'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, objs['floc'], insert_vals, objs['F'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] + ) + + body = body._rebuild( + body=body.body + + ( + vec_restore_array_f, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + stacks = ( + # cast, + dm_cast, + dm_get_app_context, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_f, + vec_get_array_x, + dm_get_local_info ) + # Dereference function data in struct + dereference_funcs = [Dereference(i, ctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + + f_soa = PointerCast(fbundle) + x_soa = PointerCast(xbundle) + + formfunc_body = CallableBody( + List(body=body), + init=(objs['begin_user'],), + stacks=stacks+tuple(dereference_funcs), + casts=(f_soa,x_soa), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} + + return Uxreplace(subs).visit(formfunc_body) + def _create_submatrices(self): body = self._submat_callback_body() objs = self.objs @@ -1411,7 +1529,7 @@ def _setup(self): call_coupled_struct_callback, shell_set_ctx, create_submats) + \ - tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + tuple(deref_dms) + tuple(xglobals) return coupled_setup @@ -1501,7 +1619,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callbacks = self.cbbuilder.formrhs + # rhs_callbacks = self.cbbuilder.formrhs xglob = sobjs['xglobal'] bglob = sobjs['bglobal'] @@ -1512,7 +1630,7 @@ def _execute_solve(self): pre_solve = () post_solve = () - for i, (c, t) in enumerate(zip(rhs_callbacks, targets)): + for i, t in enumerate(targets): name = t.name dm = sobjs[f'da{name}'] target_xloc = sobjs[f'xlocal{name}'] @@ -1522,8 +1640,11 @@ def _execute_solve(self): s = sobjs[f'scatter{name}'] pre_solve += ( - petsc_call(c.name, [dm, target_bglob]), + # petsc_call(c.name, [dm, target_bglob]), + # TODO: switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), + + # TODO: need to call reset array self.timedep.place_array(t), petsc_call( 'DMLocalToGlobal', @@ -1549,6 +1670,7 @@ def _execute_solve(self): 'VecScatterEnd', [s, target_bglob, bglob, insert_vals, sreverse] ), + BlankLine, ) post_solve += ( diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index a4917dfe8f..1a3e2bbf95 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,16 +1,17 @@ from functools import singledispatch import sympy +from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions +from devito.symbolics import retrieve_functions, FieldFromComposite from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices) + FieldData, MultipleFieldData, SubMatrices, PetscBundle) __all__ = ['PETScSolve', 'EssentialBC'] @@ -73,14 +74,14 @@ def generate_field_data(self, eqns, target, arrays): ) def build_function_eqns(self, eq, target, arrays): - b, F_target, targets = separate_eqn(eq, target) + b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) def build_matvec_eqns(self, eq, target, arrays): - b, F_target, targets = separate_eqn(eq, target) + b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) if not F_target: return None matvec = self.make_matvec(eq, F_target, arrays, targets) @@ -98,8 +99,11 @@ def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): rhs = 0. else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) + if isinstance(F_target, (int, float)): + rhs = F_target + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) return Eq(arrays['f'], rhs, subdomain=eq.subdomain) def make_rhs(self, eq, b, arrays): @@ -140,21 +144,46 @@ def linear_solve_args(self): funcs = get_funcs(combined_eqns) self.time_mapper = generate_time_mapper(funcs) - targets = list(self.target_eqns.keys()) - jacobian = SubMatrices(targets) + coupled_targets = list(self.target_eqns.keys()) + jacobian = SubMatrices(coupled_targets) + + arrays = self.generate_arrays_combined(*coupled_targets) + + + # from IPython import embed; embed() + # TODO: don't need a 'b' for coupled + # arrays = self.generate_arrays_combined(*coupled_targets) + # all_formfuncs = [] + + # from IPython import embed; embed() + # f_field = [arrays[coupled_targets[t]]['f'] for t in range(len(coupled_targets))] + # f_field = PetscBundle(name='fuv_field', components=f_field) + + + all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) + + # f_v = arrays[coupled_targets[0]]['f'] + # f_u = arrays[coupled_targets[1]]['f'] + - all_data = MultipleFieldData(jacobian) for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) - arrays = self.generate_arrays(target) - self.update_jacobian(eqns, target, jacobian, arrays) - fielddata = self.generate_field_data( - eqns, target, arrays - ) - all_data.add_field_data(fielddata) + # TODO: obvs fix and don't duplicate arrays here + # tmp_arr = self.generate_arrays(target) + self.update_jacobian(eqns, target, jacobian, arrays[target]) + + formfuncs = [self.build_function_eqns(eq, target, coupled_targets, arrays) for eq in eqns] + + all_data.extend_formfuncs(formfuncs) + + # all_formfuncs.extend(formfuncs) + + # all_data. + + # from IPython import embed; embed() return target, tuple(funcs), all_data @@ -164,21 +193,75 @@ def update_jacobian(self, eqns, target, jacobian, arrays): self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) for eq in eqns ] - # Set submatrix only if there's at least one non-zero matvec - if any(m is not None for m in matvecs): + matvecs = [m for m in matvecs if m is not None] + if matvecs: jacobian.set_submatrix(target, submat, matvecs) - def generate_field_data(self, eqns, target, arrays): - formfuncs, formrhs = zip( - *[self.build_function_eqns(eq, target, arrays) for eq in eqns] - ) + # def generate_field_data(self, eqns, target, arrays): + # # from IPython import embed; embed() + # formfuncs, formrhs = zip( + # *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + # ) + + # return FieldData( + # target=target, + # formfuncs=formfuncs, + # formrhs=formrhs, + # arrays=arrays + # ) + + def build_function_eqns(self, eq, main_target, coupled_targets, arrays): + zeroed = eq.lhs - eq.rhs + + zeroed_eqn = Eq(zeroed, 0) + zeroed_eqn = eval_time_derivatives(zeroed) + + mapper = {} + for t in coupled_targets: + target_funcs = generate_targets(Eq(zeroed, 0), t) + mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) + + # from IPython import embed; embed() + + if isinstance(eq, EssentialBC): + rhs = 0. + else: + if isinstance(zeroed, (int, float)): + rhs = zeroed + else: + rhs = zeroed.subs(mapper) + rhs = rhs.subs(self.time_mapper) + + return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) + # return Eq(f_field, rhs, subdomain=eq.subdomain) + + def generate_arrays_combined(self, *targets): + return { + target: { + p: PETScArray( + name=f'{p}_{target.name}', + target=target, + liveness='eager', + localinfo=localinfo + ) + for p in prefixes + } + for target in targets + } + + + # def make_formfunc(self, eq, F_target, arrays, targets): + # if isinstance(eq, EssentialBC): + # rhs = 0. + # else: + # if isinstance(F_target, (int, float)): + # rhs = F_target + # else: + # # from IPython import embed; embed() + # rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + # rhs = rhs.subs(self.time_mapper) + # return Eq(arrays['f'], rhs, subdomain=eq.subdomain) - return FieldData( - target=target, - formfuncs=formfuncs, - formrhs=formrhs, - arrays=arrays - ) class EssentialBC(Eq): @@ -194,7 +277,7 @@ def separate_eqn(eqn, target): zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) target_funcs = set(generate_targets(zeroed_eqn, target)) b, F_target = remove_targets(zeroed_eqn, target_funcs) - return -b, F_target, target_funcs + return -b, F_target, zeroed_eqn, target_funcs def generate_targets(eq, target): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 54bbb4a9f4..3e35e82b2b 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,12 +1,13 @@ from functools import cached_property -from ctypes import POINTER +from ctypes import POINTER, Structure from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic +from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction -from devito.tools import dtype_to_ctype, as_tuple +from devito.types.basic import AbstractFunction, IndexedData +from devito.tools import dtype_to_ctype, as_tuple, dtypes_vector_mapper, CustomDtype from devito.symbolics import FieldFromComposite +from devito.petsc.types.object import PETScStruct class PETScArray(ArrayBasic, Differentiable): @@ -116,3 +117,74 @@ def symbolic_shape(self): FieldFromComposite('g%sm' % d.name, self.localinfo) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + +class PetscBundle(Bundle): + """ + """ + + is_Bundle = True + _data_alignment = False + + @property + def _C_ctype(self): + # TODO: extend to cases with multiple petsc solves... + fields = [(i.name, dtype_to_ctype(i.dtype)) for i in self.components] + return POINTER(type('Field', (Structure,), {'_fields_': fields})) + + @cached_property + def indexed(self): + """The wrapped IndexedData object.""" + return AoSIndexedData(self.name, shape=self._shape, function=self.function) + + @cached_property + def vector(self): + return PETScArray( + name=self.name, + target=self.c0.target, + liveness=self.c0.liveness, + localinfo=self.c0.localinfo, + ) + + @property + def _C_name(self): + return self.vector._C_name + + def __getitem__(self, index): + index = as_tuple(index) + if len(index) == self.ndim: + return super().__getitem__(index) + elif len(index) == self.ndim + 1: + component_index, indices = index[0], index[1:] + # from IPython import embed; embed() + return ComponentAccess(self.indexed[indices], component_index) + else: + raise ValueError("Expected %d or %d indices, got %d instead" + % (self.ndim, self.ndim + 1, len(index))) + + # def access_component(self): + # _component_names = tuple(i.name for i in components) + # class PetscComponentAccess(ComponentAccess): + # _component_names = _component_names + + # # components = self.components + + # # _component_names = tuple(i.name for i in components) + # from IPython import embed; embed() + + # return component + + + +class AoSIndexedData(IndexedData): + @property + def dtype(self): + return self.function._C_ctype + + + + +class PetscComponentAccess(ComponentAccess): + + # _component_names = ('x_u', 'x_v') + pass diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 2e69e4a0fc..62bccc5b2c 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -308,3 +308,10 @@ class ArgvSymbol(DataSymbol): @property def _C_ctype(self): return POINTER(POINTER(c_char)) + + + +# class PetscArrayofStruct(LocalCompositeObject): +# pass +# """ +# """ \ No newline at end of file diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index bfbdd4972b..25da0d5575 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,8 +1,10 @@ import sympy -from devito.tools import Reconstructable, sympy_mutex +from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables +from devito.petsc.types.object import PETScStruct +# from devito.petsc.types.array import PetscArrayofStruct class MetaData(sympy.Function, Reconstructable): @@ -133,7 +135,7 @@ def eval(cls, *args): class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, initialguess=None, arrays=None, **kwargs): - self._target = kwargs.get('target', target) + self._target = kwargs.get('target', None) petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: @@ -189,32 +191,42 @@ def targets(self): return (self.target,) +# TODO: should this acc inherhit from fielddata? maybe not? class MultipleFieldData(FieldData): - def __init__(self, submatrices=None): - self.field_data_list = [] + def __init__(self, targets, arrays, submatrices=None): + self._targets = as_tuple(targets) + self._arrays = arrays + # self.field_data_list = [] self._submatrices = submatrices + self._formfuncs = [] + # self._f_field = f_field - def add_field_data(self, field_data): - self.field_data_list.append(field_data) + # def add_field_data(self, field_data): + # self.field_data_list.append(field_data) - def get_field_data(self, target): - for field_data in self.field_data_list: - if field_data.target == target: - return field_data - raise ValueError(f"FieldData with target {target} not found.") + # def get_field_data(self, target): + # for field_data in self.field_data_list: + # if field_data.target == target: + # return field_data + # raise ValueError(f"FieldData with target {target} not found.") pass - @property - def target(self): - return None + def extend_formfuncs(self, formfuncs): + self._formfuncs.extend(formfuncs) - @property - def targets(self): - return tuple(field_data.target for field_data in self.field_data_list) + # @property + # def target(self): + # return None + + # @property + # def targets(self): + # return tuple(field_data.target for field_data in self.field_data_list) @property def space_dimensions(self): - space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + # space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + space_dims = {t.space_dimensions for t in self.targets} + # from IPython import embed; embed() if len(space_dims) > 1: # TODO: This may not actually have to be the case, but enforcing it for now raise ValueError( @@ -249,6 +261,18 @@ def space_order(self): def submatrices(self): return self._submatrices + # @property + # def formfuncs(self): + # return self._formfuncs + + @property + def targets(self): + return self._targets + + @property + def arrays(self): + return self._arrays + class SubMatrices: def __init__(self, targets): diff --git a/devito/types/basic.py b/devito/types/basic.py index a28ca2f486..9990c55836 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -87,6 +87,7 @@ def _C_typedata(self): if isinstance(_type, CustomDtype): return _type + # from IPython import embed; embed() while issubclass(_type, _Pointer): _type = _type._type_ From f6b38eddb4c3c7c85c93854bcbc76ff41ed1942d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 29 Apr 2025 17:40:00 +0100 Subject: [PATCH 02/53] compiler: Working petscbundle --- devito/petsc/iet/routines.py | 44 ++++++++++++++++-------------------- devito/petsc/types/array.py | 35 ++++++++-------------------- devito/petsc/types/types.py | 3 +-- devito/types/array.py | 8 +++++-- 4 files changed, 36 insertions(+), 54 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index be76ceace1..3daa7efcfc 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -21,7 +21,7 @@ from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, - SubMatrixStruct, CallbackDM, PetscComponentAccess) + SubMatrixStruct, CallbackDM) class CBBuilder: @@ -660,9 +660,9 @@ def _make_core(self): if mtvs['matvecs']: self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') - self._make_user_struct_callback() self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) + self._make_user_struct_callback() self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -758,7 +758,7 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - + # from IPython import embed; embed() targets = fielddata.targets arrays = fielddata.arrays @@ -769,10 +769,16 @@ def _whole_formfunc_body(self, body, fielddata): x_u = arrays[targets[0]]['x'] x_v = arrays[targets[1]]['x'] - fbundle = PetscBundle(name='f_vu', components=(f_v, f_u)) - xbundle = PetscBundle(name='x_vu', components=(x_v, x_u)) - mapper1 = {x_u.base: xbundle.base, x_v.base: xbundle.base} + target_indices = {t: i for i, t in enumerate(targets)} + # from IPython import embed; embed() + + # TODO: to group them, maybe pass in struct name as arg to petscbundle + fbundle = PetscBundle(name='f_vu', components=(f_u, f_v)) + xbundle = PetscBundle(name='x_vu', components=(x_u, x_v)) + + mapper1 = {x_u.base: xbundle, x_v.base: xbundle, f_u.base: fbundle, + f_v.base: fbundle} indexeds = FindSymbols('indexeds').visit(body) @@ -780,14 +786,10 @@ def _whole_formfunc_body(self, body, fielddata): for i in indexeds: if i.base in mapper1: bundle = mapper1[i.base] - subss[i] = i.func(bundle, *i.indices) - # subss[i] = bundle.indexed.__getitem__((1,)+i.indices) - # subss[i] = subss[i].__getitem__((0,)+i.indices) - subss[i] = PetscComponentAccess(subss[i], 1) - - # subss[i] = xbundle.access_component() - - # from IPython import embed; embed() + index = target_indices[i.function.target] + index = (index,)+i.indices + subss[i] = bundle.__getitem__(index) + body = Uxreplace(subss).visit(body) @@ -1373,9 +1375,9 @@ def _create_dmda(self, dmda): class CoupledSetup(BaseSetup): - @property - def snes_ctx(self): - return Byref(self.solver_objs['jacctx']) + # @property + # def snes_ctx(self): + # return Byref(self.solver_objs['jacctx']) def _setup(self): # TODO: minimise code duplication with superclass @@ -1662,14 +1664,6 @@ def _execute_solve(self): 'VecScatterEnd', [s, target_xglob, xglob, insert_vals, sreverse] ), - petsc_call( - 'VecScatterBegin', - [s, target_bglob, bglob, insert_vals, sreverse] - ), - petsc_call( - 'VecScatterEnd', - [s, target_bglob, bglob, insert_vals, sreverse] - ), BlankLine, ) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 3e35e82b2b..39eb2e2347 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,6 +1,8 @@ from functools import cached_property from ctypes import POINTER, Structure +from sympy import Expr + from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess from devito.finite_differences import Differentiable @@ -122,14 +124,14 @@ def symbolic_shape(self): class PetscBundle(Bundle): """ """ - is_Bundle = True + _data_alignment = False @property def _C_ctype(self): # TODO: extend to cases with multiple petsc solves... - fields = [(i.name, dtype_to_ctype(i.dtype)) for i in self.components] + fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type('Field', (Structure,), {'_fields_': fields})) @cached_property @@ -156,35 +158,18 @@ def __getitem__(self, index): return super().__getitem__(index) elif len(index) == self.ndim + 1: component_index, indices = index[0], index[1:] - # from IPython import embed; embed() - return ComponentAccess(self.indexed[indices], component_index) + names = tuple(i.target.name for i in self.components) + return ComponentAccess( + self.indexed[indices], + component_index, + component_names=names + ) else: raise ValueError("Expected %d or %d indices, got %d instead" % (self.ndim, self.ndim + 1, len(index))) - # def access_component(self): - # _component_names = tuple(i.name for i in components) - # class PetscComponentAccess(ComponentAccess): - # _component_names = _component_names - - # # components = self.components - - # # _component_names = tuple(i.name for i in components) - # from IPython import embed; embed() - - # return component - - class AoSIndexedData(IndexedData): @property def dtype(self): return self.function._C_ctype - - - - -class PetscComponentAccess(ComponentAccess): - - # _component_names = ('x_u', 'x_v') - pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 25da0d5575..e470986612 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -135,8 +135,7 @@ def eval(cls, *args): class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, initialguess=None, arrays=None, **kwargs): - self._target = kwargs.get('target', None) - + self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: raise TypeError( diff --git a/devito/types/array.py b/devito/types/array.py index eb1fd0dc4c..32c8ccc97b 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -495,6 +495,7 @@ def _C_get_field(self, region, dim, side=None): return self.c0._C_get_field(region, dim, side=side) def __getitem__(self, index): + # from IPython import embed; embed() index = as_tuple(index) if len(index) == self.ndim: return super().__getitem__(index) @@ -535,19 +536,22 @@ def handles(self): class ComponentAccess(Expr, Pickable): - _component_names = ('x', 'y', 'z', 'w') + _default_component_names = ('x', 'y', 'z', 'w') __rargs__ = ('arg',) __rkwargs__ = ('index',) - def __new__(cls, arg, index=0, **kwargs): + def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) if not is_integer(index) or index > 3: raise ValueError("Expected 0 <= index < 4") + names = component_names or cls._default_component_names + obj = Expr.__new__(cls, arg) obj._index = index + obj._component_names = names return obj From 7300c8f6a896f86d751090a2ed4e7ec13c9ec530 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 30 Apr 2025 18:52:02 +0100 Subject: [PATCH 03/53] compiler: Add PetscMixin to simplify priority destroys --- .github/workflows/pytest-petsc.yml | 1 + devito/passes/iet/definitions.py | 2 +- devito/passes/iet/engine.py | 33 +- devito/passes/iet/misc.py | 3 +- devito/petsc/iet/routines.py | 181 ++++-- devito/petsc/solve.py | 82 +-- devito/petsc/types/array.py | 20 +- devito/petsc/types/object.py | 21 +- devito/petsc/types/types.py | 28 +- devito/tools/utils.py | 1 + devito/types/array.py | 1 - devito/types/basic.py | 1 - .../petsc/random/biharmonic/02_biharmonic.py | 148 +++++ .../random/biharmonic/biharmonic_matfree.c | 556 ++++++++++++++++++ .../biharmonic/biharmonic_matfree_nonscaled.c | 553 +++++++++++++++++ tests/test_petsc.py | 4 + 16 files changed, 1440 insertions(+), 195 deletions(-) create mode 100644 examples/petsc/random/biharmonic/02_biharmonic.py create mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree.c create mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 13177cd7b8..19446642af 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -83,6 +83,7 @@ jobs: ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/biharmonic/02_biharmonic.py - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index f5ae337b9c..e56c86645f 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -379,7 +379,7 @@ def place_definitions(self, iet, globs=None, **kwargs): # created by the compiler up to this point (Array, LocalObject, etc.) storage = Storage() defines = FindSymbols('defines-aliases|globals').visit(iet) - # from IPython import embed; embed() + for i in FindSymbols().visit(iet): if i in defines: continue diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 992096fd41..a90f4ba275 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -13,9 +13,11 @@ from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) +from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension +from devito.petsc.types import PETScArray __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -149,8 +151,6 @@ def apply(self, func, **kwargs): if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) # TODO: fix for petsc bundles - # TODO: somethinng to do with this bug is causing the compiler not to combine loops with petsc bundles - # from IPython import embed; embed() # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -355,7 +355,7 @@ def reuse_efuncs(root, efuncs, sregistry=None): if isinstance(efunc, AsyncCallable): mapper[len(mapper)] = (efunc, [efunc]) continue - + # from IPython import embed; embed() afunc = abstract_efunc(efunc) key = afunc._signature() @@ -394,7 +394,7 @@ def abstract_efunc(efunc): - Objects are renamed as "o0", "o1", ... """ functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) - # from IPython import embed; embed() + mapper = abstract_objects(functions) # from IPython import embed; embed() efunc = Uxreplace(mapper).visit(efunc) @@ -411,24 +411,24 @@ def abstract_objects(objects0, sregistry=None): objects = [] for i in objects0: if i.is_Bundle: - # from IPython import embed; embed() objects.extend(i.components) objects.append(i) - # from IPython import embed; embed() + # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] + # keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() - # from IPython import embed; embed() + for i in objects: - # from IPython import embed; embed() abstract_object(i, mapper, sregistry) + # from IPython import embed; embed() return mapper @@ -472,9 +472,24 @@ def _(i, mapper, sregistry): mapper[i.dmap] = v.dmap +# @abstract_object.register(PETScArray) +# def _(i, mapper, sregistry): +# name = sregistry.make_name(prefix='xx') + +# v = i._rebuild(name=name, initializer=None, alias=True) + +# mapper.update({ +# i: v, +# i.indexed: v.indexed, +# i.dmap: v.dmap, +# i._C_symbol: v._C_symbol, +# }) + + @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') + components = [mapper[f] for f in i.components] v = i._rebuild(name=name, components=components, alias=True) diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index 0f15deec81..e404a8e373 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -184,7 +184,6 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): try: v = tracker[i.base].v subs[i] = v.func(v.base, *i.indices) - # from IPython import embed; embed() continue except KeyError: pass @@ -194,7 +193,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs[i] = v tracker[i.base] = Bunch(header=header, v=v) - # from IPython import embed; embed() + iet = Uxreplace(subs).visit(iet) return iet diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3daa7efcfc..8674dc4df4 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -5,9 +5,9 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, retrieve_iteration_tree, filter_iterations, Iteration, - PointerCast, Expression, Transformer) + PointerCast) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod, uxreplace) + FieldFromComposite, IntDiv, Deref, Mod) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension @@ -20,7 +20,7 @@ from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, - DMCast, JacobianStructCast, JacobianStruct, + DMCast, JacobianStruct, SubMatrixStruct, CallbackDM) @@ -99,6 +99,13 @@ def initialguesses(self): def user_struct_callback(self): return self._user_struct_callback + @property + def zero_memory(self): + """Indicates whether the memory of the output + vector should be set to zero before the computation + in the callback.""" + return True + def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata self._make_matvec(fielddata.arrays, fielddata.matvecs) @@ -148,6 +155,10 @@ def _create_matvec_body(self, body, arrays): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_y_memory = petsc_call( + 'VecSet', [objs['Y'], 0.0] + ) if self.zero_memory else None + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] ) @@ -165,6 +176,10 @@ def _create_matvec_body(self, body, arrays): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) + zero_ylocal_memory = petsc_call( + 'VecSet', [ylocal, 0.0] + ) + vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] ) @@ -186,11 +201,11 @@ def _create_matvec_body(self, body, arrays): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, ylocal, insert_vals, objs['Y'] + dmda, ylocal, add_vals, objs['Y'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, ylocal, insert_vals, objs['Y'] + dmda, ylocal, add_vals, objs['Y'] ]) dm_restore_local_xvec = petsc_call( @@ -221,10 +236,12 @@ def _create_matvec_body(self, body, arrays): stacks = ( mat_get_dm, dm_get_app_context, + zero_y_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, dm_get_local_yvec, + zero_ylocal_memory, vec_get_array_y, vec_get_array_x, dm_get_local_info @@ -290,6 +307,10 @@ def _create_formfunc_body(self, body, fielddata): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_f_memory = petsc_call( + 'VecSet', [objs['F'], 0.0] + ) if self.zero_memory else None + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) @@ -328,11 +349,11 @@ def _create_formfunc_body(self, body, fielddata): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( @@ -356,6 +377,7 @@ def _create_formfunc_body(self, body, fielddata): stacks = ( dm_cast, dm_get_app_context, + zero_f_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, @@ -647,13 +669,19 @@ def main_matvec_callback(self): def main_formfunc_callback(self): return self._main_formfunc_callback + @property + def zero_memory(self): + """Indicates whether the memory of the output + vector should be set to zero before the computation + in the callback.""" + return False + def _make_core(self): injectsolve = self.injectsolve targets = injectsolve.expr.rhs.fielddata.targets all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - row_matvecs = all_fielddata.submatrices.submatrices[t] arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): @@ -688,6 +716,10 @@ def _whole_matvec_body(self): nonzero_submats = self.submatrices.nonzero_submatrix_keys + zero_y_memory = petsc_call( + 'VecSet', [objs['Y'], 0.0] + ) + calls = () for sm in nonzero_submats: idx = self.submatrices.submat_to_index[sm] @@ -720,14 +752,13 @@ def _whole_matvec_body(self): ), ) return CallableBody( - List(body=(ctx_main, BlankLine) + calls), + List(body=(ctx_main, zero_y_memory, BlankLine) + calls), init=(objs['begin_user'],), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) def _make_whole_formfunc(self, fielddata): formfuncs = fielddata.formfuncs - # from IPython import embed; embed() # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, @@ -739,7 +770,6 @@ def _make_whole_formfunc(self, fielddata): cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), body_formfunc, - # List(body=()), retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) @@ -758,40 +788,12 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - # from IPython import embed; embed() - - targets = fielddata.targets - arrays = fielddata.arrays - - f_u = arrays[targets[0]]['f'] - f_v = arrays[targets[1]]['f'] - - x_u = arrays[targets[0]]['x'] - x_v = arrays[targets[1]]['x'] - - - target_indices = {t: i for i, t in enumerate(targets)} - # from IPython import embed; embed() - - # TODO: to group them, maybe pass in struct name as arg to petscbundle - fbundle = PetscBundle(name='f_vu', components=(f_u, f_v)) - xbundle = PetscBundle(name='x_vu', components=(x_u, x_v)) - mapper1 = {x_u.base: xbundle, x_v.base: xbundle, f_u.base: fbundle, - f_v.base: fbundle} - - indexeds = FindSymbols('indexeds').visit(body) - - subss = {} - for i in indexeds: - if i.base in mapper1: - bundle = mapper1[i.base] - index = target_indices[i.function.target] - index = (index,)+i.indices - subss[i] = bundle.__getitem__(index) - - - body = Uxreplace(subss).visit(body) + # Process body for residual callback, including generating bundles etc + bundles = sobjs['bundles'] + fbundle = bundles['f'] + xbundle = bundles['x'] + body = self.bundle_residual(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -799,6 +801,10 @@ def _whole_formfunc_body(self, body, fielddata): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_f_memory = petsc_call( + 'VecSet', [objs['F'], 0.0] + ) + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) @@ -837,11 +843,11 @@ def _whole_formfunc_body(self, body, fielddata): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( @@ -854,8 +860,7 @@ def _whole_formfunc_body(self, body, fielddata): body = body._rebuild( body=body.body + - ( - vec_restore_array_f, + (vec_restore_array_f, vec_restore_array_x, dm_local_to_global_begin, dm_local_to_global_end, @@ -864,9 +869,9 @@ def _whole_formfunc_body(self, body, fielddata): ) stacks = ( - # cast, dm_cast, dm_get_app_context, + zero_f_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, @@ -880,7 +885,6 @@ def _whole_formfunc_body(self, body, fielddata): dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] - f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -888,8 +892,9 @@ def _whole_formfunc_body(self, body, fielddata): List(body=body), init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), - casts=(f_soa,x_soa), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) + casts=(f_soa, x_soa), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} @@ -1041,6 +1046,26 @@ def _submat_callback_body(self): stacks=(get_ctx, deref_subdm), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) + + def bundle_residual(self, body, bundles): + fbundle = bundles['f'] + xbundle = bundles['x'] + + mapper1 = bundles['bundle_mapper'] + + indexeds = FindSymbols('indexeds').visit(body) + + subss = {} + for i in indexeds: + if i.base in mapper1: + bundle = mapper1[i.base] + index = bundles['target_indices'][i.function.target] + index = (index,)+i.indices + subss[i] = bundle.__getitem__(index) + + body = Uxreplace(subss).visit(body) + return body + class BaseObjectBuilder: @@ -1120,6 +1145,7 @@ def _extend_build(self, base_dict): sreg = self.sregistry objs = self.objs targets = self.fielddata.targets + arrays = self.fielddata.arrays base_dict['fields'] = PointerIS( name=sreg.make_name(prefix='fields'), nindices=len(targets) @@ -1154,6 +1180,40 @@ def _extend_build(self, base_dict): base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') base_dict[f'{key}F'] = CallbackVec(f'{key}F') + + # Bundle objects/metadata required by the coupled residual callback + f_components = [] + x_components = [] + bundle_mapper = {} + + target_indices = {t: i for i, t in enumerate(targets)} + + for t in targets: + f_arr = arrays[t]['f'] + x_arr = arrays[t]['x'] + f_components.append(f_arr) + x_components.append(x_arr) + + # TODO: to group them, maybe pass in struct name as arg to petscbundle + fbundle = PetscBundle(name='f_bundle', components=f_components) + xbundle = PetscBundle(name='x_bundle', components=x_components) + + # Build the bundle mapper + for i, t in enumerate(targets): + f_arr = arrays[t]['f'] + x_arr = arrays[t]['x'] + bundle_mapper[f_arr.base] = fbundle + bundle_mapper[x_arr.base] = xbundle + + + base_dict['bundles'] = { + 'f': fbundle, + 'x': xbundle, + 'bundle_mapper': bundle_mapper, + # TODO: maybe this shouldn't be here + 'target_indices': target_indices + } + return base_dict def _target_dependent(self, base_dict): @@ -1410,9 +1470,6 @@ def _setup(self): get_local_size = petsc_call('VecGetSize', [sobjs['xlocal'], Byref(sobjs['localsize'])]) - global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(sobjs['bglobal'])]) - snes_get_ksp = petsc_call('SNESGetKSP', [sobjs['snes'], Byref(sobjs['ksp'])]) @@ -1500,11 +1557,6 @@ def _setup(self): [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] ) for t in targets] - bglobals = [petsc_call( - 'DMCreateGlobalVector', - [sobjs[f'da{t.name}'], Byref(sobjs[f'bglobal{t.name}'])] - ) for t in targets] - coupled_setup = dmda_calls + ( snes_create, snes_set_dm, @@ -1514,7 +1566,6 @@ def _setup(self): global_x, local_x, get_local_size, - global_b, snes_get_ksp, ksp_set_tols, ksp_set_type, @@ -1617,6 +1668,7 @@ def _execute_solve(self): Assigns the required time iterators to the struct and executes the necessary calls to execute the SNES solver. """ + objs = self.objs sobjs = self.solver_objs struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) @@ -1637,13 +1689,11 @@ def _execute_solve(self): dm = sobjs[f'da{name}'] target_xloc = sobjs[f'xlocal{name}'] target_xglob = sobjs[f'xglobal{name}'] - target_bglob = sobjs[f'bglobal{name}'] field = sobjs['fields'].indexed[i] s = sobjs[f'scatter{name}'] pre_solve += ( - # petsc_call(c.name, [dm, target_bglob]), - # TODO: switch to createwitharray and move to setup + # TODO: switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), # TODO: need to call reset array @@ -1682,7 +1732,7 @@ def _execute_solve(self): ) ) - snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], bglob, xglob]),) + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),) return List( body=( @@ -1893,5 +1943,6 @@ def assign_time_iters(self, struct): void = 'void' insert_vals = 'INSERT_VALUES' +add_vals = 'ADD_VALUES' sreverse = 'SCATTER_REVERSE' sforward = 'SCATTER_FORWARD' diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1a3e2bbf95..c34b829798 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,17 +1,16 @@ from functools import singledispatch import sympy -from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions, FieldFromComposite +from devito.symbolics import retrieve_functions from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices, PetscBundle) + FieldData, MultipleFieldData, SubMatrices) __all__ = ['PETScSolve', 'EssentialBC'] @@ -74,14 +73,14 @@ def generate_field_data(self, eqns, target, arrays): ) def build_function_eqns(self, eq, target, arrays): - b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) def build_matvec_eqns(self, eq, target, arrays): - b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, target) if not F_target: return None matvec = self.make_matvec(eq, F_target, arrays, targets) @@ -97,7 +96,9 @@ def make_matvec(self, eq, F_target, arrays, targets): def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - rhs = 0. + # TODO: CHECK THIS + rhs = arrays[main_target]['x'] - eq.rhs + # rhs = 0. else: if isinstance(F_target, (int, float)): rhs = F_target @@ -149,42 +150,19 @@ def linear_solve_args(self): arrays = self.generate_arrays_combined(*coupled_targets) - - # from IPython import embed; embed() - # TODO: don't need a 'b' for coupled - # arrays = self.generate_arrays_combined(*coupled_targets) - # all_formfuncs = [] - - # from IPython import embed; embed() - # f_field = [arrays[coupled_targets[t]]['f'] for t in range(len(coupled_targets))] - # f_field = PetscBundle(name='fuv_field', components=f_field) - - - all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) - - # f_v = arrays[coupled_targets[0]]['f'] - # f_u = arrays[coupled_targets[1]]['f'] - - + all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, + targets=coupled_targets) for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) - - - # TODO: obvs fix and don't duplicate arrays here - # tmp_arr = self.generate_arrays(target) self.update_jacobian(eqns, target, jacobian, arrays[target]) - formfuncs = [self.build_function_eqns(eq, target, coupled_targets, arrays) for eq in eqns] - + formfuncs = [ + self.build_function_eqns(eq, target, coupled_targets, arrays) + for eq in eqns + ] all_data.extend_formfuncs(formfuncs) - # all_formfuncs.extend(formfuncs) - - # all_data. - - # from IPython import embed; embed() - return target, tuple(funcs), all_data def update_jacobian(self, eqns, target, jacobian, arrays): @@ -197,22 +175,10 @@ def update_jacobian(self, eqns, target, jacobian, arrays): if matvecs: jacobian.set_submatrix(target, submat, matvecs) - # def generate_field_data(self, eqns, target, arrays): - # # from IPython import embed; embed() - # formfuncs, formrhs = zip( - # *[self.build_function_eqns(eq, target, arrays) for eq in eqns] - # ) - - # return FieldData( - # target=target, - # formfuncs=formfuncs, - # formrhs=formrhs, - # arrays=arrays - # ) - def build_function_eqns(self, eq, main_target, coupled_targets, arrays): zeroed = eq.lhs - eq.rhs + # TODO: clean up, test coupled with time dependence zeroed_eqn = Eq(zeroed, 0) zeroed_eqn = eval_time_derivatives(zeroed) @@ -221,10 +187,9 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): target_funcs = generate_targets(Eq(zeroed, 0), t) mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) - # from IPython import embed; embed() - if isinstance(eq, EssentialBC): - rhs = 0. + # TODO: CHECK THIS + rhs = arrays[main_target]['x'] - eq.rhs else: if isinstance(zeroed, (int, float)): rhs = zeroed @@ -233,7 +198,6 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): rhs = rhs.subs(self.time_mapper) return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) - # return Eq(f_field, rhs, subdomain=eq.subdomain) def generate_arrays_combined(self, *targets): return { @@ -250,20 +214,6 @@ def generate_arrays_combined(self, *targets): } - # def make_formfunc(self, eq, F_target, arrays, targets): - # if isinstance(eq, EssentialBC): - # rhs = 0. - # else: - # if isinstance(F_target, (int, float)): - # rhs = F_target - # else: - # # from IPython import embed; embed() - # rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - # rhs = rhs.subs(self.time_mapper) - # return Eq(arrays['f'], rhs, subdomain=eq.subdomain) - - - class EssentialBC(Eq): pass diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 39eb2e2347..58ca3e28d7 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,15 +1,12 @@ from functools import cached_property from ctypes import POINTER, Structure -from sympy import Expr - from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess +from devito.types.array import ArrayBasic, Bundle, ComponentAccess from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction, IndexedData -from devito.tools import dtype_to_ctype, as_tuple, dtypes_vector_mapper, CustomDtype +from devito.tools import dtype_to_ctype, as_tuple from devito.symbolics import FieldFromComposite -from devito.petsc.types.object import PETScStruct class PETScArray(ArrayBasic, Differentiable): @@ -125,12 +122,11 @@ class PetscBundle(Bundle): """ """ is_Bundle = True - _data_alignment = False @property def _C_ctype(self): - # TODO: extend to cases with multiple petsc solves... + # TODO: extend to cases with multiple petsc solves...(need diff struct name for each solve) fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type('Field', (Structure,), {'_fields_': fields})) @@ -142,11 +138,11 @@ def indexed(self): @cached_property def vector(self): return PETScArray( - name=self.name, - target=self.c0.target, - liveness=self.c0.liveness, - localinfo=self.c0.localinfo, - ) + name=self.name, + target=self.c0.target, + liveness=self.c0.liveness, + localinfo=self.c0.localinfo, + ) @property def _C_name(self): diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 62bccc5b2c..7821e316d9 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -7,6 +7,12 @@ from devito.petsc.iet.utils import petsc_call +class PetscMixin: + @property + def _C_free_priority(self): + return FREE_PRIORITY[self] + + class CallbackDM(LocalObject): """ PETSc Data Management object (DM). This is the DM instance @@ -16,7 +22,7 @@ class CallbackDM(LocalObject): dtype = CustomDtype('DM') -class DM(LocalObject): +class DM(LocalObject, PetscMixin): """ PETSc Data Management object (DM). This is the primary DM instance created within the main kernel and linked to the SNES @@ -36,11 +42,6 @@ def dofs(self): def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) - # TODO: Switch to an enumeration? - @property - def _C_free_priority(self): - return 4 - DMCast = cast('DM') @@ -310,8 +311,6 @@ def _C_ctype(self): return POINTER(POINTER(c_char)) - -# class PetscArrayofStruct(LocalCompositeObject): -# pass -# """ -# """ \ No newline at end of file +FREE_PRIORITY = { + DM: 4, +} diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index e470986612..973d558f38 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -3,8 +3,6 @@ from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables -from devito.petsc.types.object import PETScStruct -# from devito.petsc.types.array import PetscArrayofStruct class MetaData(sympy.Function, Reconstructable): @@ -195,37 +193,17 @@ class MultipleFieldData(FieldData): def __init__(self, targets, arrays, submatrices=None): self._targets = as_tuple(targets) self._arrays = arrays - # self.field_data_list = [] self._submatrices = submatrices self._formfuncs = [] - # self._f_field = f_field - # def add_field_data(self, field_data): - # self.field_data_list.append(field_data) - - # def get_field_data(self, target): - # for field_data in self.field_data_list: - # if field_data.target == target: - # return field_data - # raise ValueError(f"FieldData with target {target} not found.") - pass + # pass def extend_formfuncs(self, formfuncs): self._formfuncs.extend(formfuncs) - # @property - # def target(self): - # return None - - # @property - # def targets(self): - # return tuple(field_data.target for field_data in self.field_data_list) - @property def space_dimensions(self): - # space_dims = {field_data.space_dimensions for field_data in self.field_data_list} space_dims = {t.space_dimensions for t in self.targets} - # from IPython import embed; embed() if len(space_dims) > 1: # TODO: This may not actually have to be the case, but enforcing it for now raise ValueError( @@ -260,10 +238,6 @@ def space_order(self): def submatrices(self): return self._submatrices - # @property - # def formfuncs(self): - # return self._formfuncs - @property def targets(self): return self._targets diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0a28de16a8..0657146040 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -346,3 +346,4 @@ def key(i): return (v, str(type(i))) return sorted(items, key=key, reverse=True) + diff --git a/devito/types/array.py b/devito/types/array.py index 32c8ccc97b..3c3cedfc88 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -495,7 +495,6 @@ def _C_get_field(self, region, dim, side=None): return self.c0._C_get_field(region, dim, side=side) def __getitem__(self, index): - # from IPython import embed; embed() index = as_tuple(index) if len(index) == self.ndim: return super().__getitem__(index) diff --git a/devito/types/basic.py b/devito/types/basic.py index 9990c55836..a28ca2f486 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -87,7 +87,6 @@ def _C_typedata(self): if isinstance(_type, CustomDtype): return _type - # from IPython import embed; embed() while issubclass(_type, _Pointer): _type = _type._type_ diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py new file mode 100644 index 0000000000..8c15f3d3a2 --- /dev/null +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -0,0 +1,148 @@ +# ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + + +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, switchconfig, + configuration, SubDomain) + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() + +# Subdomains to implement BCs +class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + + +def c(x): + return x**3 * (1 - x)**3 + +def ddc(x): + return 6.0 * x * (1 - x) * (1 - 5.0 * x + 5.0 * x**2) + +def d4c(x): + return -72.0 * (1 - 5.0 * x + 5.0 * x**2) + +def u_exact_fcn(x, y): + return c(x) * c(y) + +def lap_u_exact_fcn(x, y): + return -ddc(x) * c(y) - c(x) * ddc(y) + +def f_fcn(x, y): + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y) + + +sub1 = SubTop() +sub2 = SubBottom() +sub3 = SubLeft() +sub4 = SubRight() + +subdomains = (sub1, sub2, sub3, sub4) + +Lx = np.float64(1.) +Ly = np.float64(1.) + +# n_values = [33, 53, 73, 93, 113] +n_values = [33] +dx = np.array([Lx/(n-1) for n in n_values]) + +u_errors = [] +v_errors = [] + +for n in n_values: + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + u_exact = Function(name='u_exact', grid=grid, space_order=2) + lap_u = Function(name='lap_u', grid=grid, space_order=2) + + eqn1 = Eq(-v.laplace, f, subdomain=grid.interior) + eqn2 = Eq(-u.laplace, v, subdomain=grid.interior) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + X, Y = np.meshgrid(tmpx, tmpy) + + f.data[:] = f_fcn(X, Y) + + # # Create boundary condition expressions using subdomains + # TODO: add initial guess callback for mixed systems + bc_u = [EssentialBC(u, 0., subdomain=sub1)] + bc_u += [EssentialBC(u, 0., subdomain=sub2)] + bc_u += [EssentialBC(u, 0., subdomain=sub3)] + bc_u += [EssentialBC(u, 0., subdomain=sub4)] + bc_v = [EssentialBC(v, 0., subdomain=sub1)] + bc_v += [EssentialBC(v, 0., subdomain=sub2)] + bc_v += [EssentialBC(v, 0., subdomain=sub3)] + bc_v += [EssentialBC(v, 0., subdomain=sub4)] + + # T (see ref) is nonsymmetric so need to set default KSP type to GMRES + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters={'ksp_rtol': 1e-10}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + u_exact.data[:] = u_exact_fcn(X, Y) + lap_u.data[:] = lap_u_exact_fcn(X, Y) + + # Compute infinity norm for u + u_diff = u_exact.data[:] - u.data[:] + u_error = np.linalg.norm(u_diff.ravel(), ord=np.inf) / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) + u_errors.append(u_error) + + # Compute infinity norm for lap_u + v_diff = lap_u.data[:] - v.data[:] + v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) + v_errors.append(v_error) + +# u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) +# v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) + +# assert u_slope > 1.9 +# assert u_slope < 2.1 + +# assert v_slope > 1.9 +# assert v_slope < 2.1 +print(op.ccode) +print("u_errors:", u_errors) +print("v_errors:", v_errors) diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c new file mode 100644 index 0000000000..53121473e7 --- /dev/null +++ b/examples/petsc/random/biharmonic/biharmonic_matfree.c @@ -0,0 +1,556 @@ + +// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + +static char help[] = +"Solve the linear biharmonic equation in 2D. Equation is\n" +" Lap^2 u = f\n" +"where Lap = - grad^2 is the positive Laplacian, equivalently\n" +" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" +"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" +"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" +"2x2 block system with SPD Laplacian blocks on the diagonal:\n" +" | Lap | 0 | | v | | f | \n" +" |-----|-----| |---| = |---| \n" +" | -I | Lap | | u | | 0 | \n" +"Includes manufactured, polynomial exact solution. The discretization is\n" +"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" +"Recommended preconditioning combines fieldsplit:\n" +" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" +"with multigrid as the preconditioner for the diagonal blocks:\n" +" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" +"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" +"monolithic multigrid (-pc_type mg|gamg).\n\n"; + +#include + +typedef struct { + PetscReal v, u; +} Field; + +typedef struct { + PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side +} BiharmCtx; + +struct JacobianCtx +{ + DM * subdms; + IS * fields; + Mat * submats; +} ; + +struct SubMatrixCtx +{ + IS * rows; + IS * cols; +} ; + +static PetscReal c(PetscReal x) { + return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); +} + +static PetscReal ddc(PetscReal x) { + return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal d4c(PetscReal x) { + return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { + return c(x) * c(y); +} + +static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { + return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u +} + +static PetscReal f_fcn(PetscReal x, PetscReal y) { + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u +} + +extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); +extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); +extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); +extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); + +int main(int argc,char **argv) { + DM da; + SNES snes; + Vec w, w_initial, w_exact; + BiharmCtx user; + Field **aW; + PetscReal normv, normu, errv, erru; + DMDALocalInfo info; + IS *fields; + DM *subdms; + PetscInt nfields; + + struct JacobianCtx jctx0; + Mat J; + + PetscCall(PetscInitialize(&argc,&argv,NULL,help)); + + user.f = &f_fcn; + PetscCall(DMDACreate2d(PETSC_COMM_WORLD, + DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, + 33,33,PETSC_DECIDE,PETSC_DECIDE, + 2,1, // degrees of freedom, stencil width + NULL,NULL,&da)); + PetscCall(DMSetApplicationContext(da,&user)); + PetscCall(DMSetFromOptions(da)); + PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates + PetscCall(DMSetMatType(da, MATSHELL)); + PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); + PetscCall(DMDASetFieldName(da,0,"v")); + PetscCall(DMDASetFieldName(da,1,"u")); + PetscCall(DMCreateMatrix(da,&J)); + + PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); + PetscCall(SNESSetDM(snes,da)); + PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); + PetscCall(SNESSetType(snes,SNESKSPONLY)); + PetscCall(SNESSetFromOptions(snes)); + + PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); + PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); + + PetscCall(MatSetDM(J,da)); + PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); + PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); + PetscCall(MatShellSetContext(J,&(jctx0))); + PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); + + PetscCall(DMGetGlobalVector(da,&w_initial)); + PetscCall(VecSet(w_initial,0.0)); + PetscCall(SNESSolve(snes,NULL,w_initial)); + // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(DMRestoreGlobalVector(da,&w_initial)); + PetscCall(DMDestroy(&da)); + + PetscCall(SNESGetSolution(snes,&w)); + PetscCall(SNESGetDM(snes,&da)); + PetscCall(DMDAGetLocalInfo(da,&info)); + + PetscCall(DMCreateGlobalVector(da,&w_exact)); + PetscCall(DMDAVecGetArray(da,w_exact,&aW)); + PetscCall(FormExactWLocal(&info,aW,&user)); + PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); + PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); + PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); + PetscCall(VecAXPY(w,-1.0,w_exact)); + PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); + PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); + PetscCall(PetscPrintf(PETSC_COMM_WORLD, + "done on %d x %d grid ...\n" + " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", + info.mx,info.my,errv/normv,erru/normu)); + + + PetscCall(ISDestroy(&(fields[0]))); + PetscCall(ISDestroy(&(fields[1]))); + PetscCall(PetscFree(fields)); + PetscCall(DMDestroy(&(subdms[0]))); + PetscCall(DMDestroy(&(subdms[1]))); + PetscCall(PetscFree(subdms)); + PetscCall(VecDestroy(&w_exact)); + PetscCall(MatDestroy(&J)); + PetscCall(SNESDestroy(&snes)); + PetscCall(PetscFinalize()); + return 0; +} + +PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, x, y; + PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); + hx = (xymax[0] - xymin[0]) / (info->mx - 1); + hy = (xymax[1] - xymin[1]) / (info->my - 1); + for (j = info->ys; j < info->ys + info->ym; j++) { + y = j * hy; + for (i = info->xs; i < info->xs + info->xm; i++) { + x = i * hx; + aW[j][i].u = u_exact_fcn(x,y); + aW[j][i].v = lap_u_exact_fcn(x,y); + } + } + return 0; +} + + +PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) +{ + Vec xlocal, flocal; + DMDALocalInfo info; + DM da; + PetscScalar *x_vec, *f_vec; + + BiharmCtx *user; + + PetscCall(SNESGetDM(snes,&da)); + + PetscCall(DMGetApplicationContext(da,&user)); + + PetscCall(DMDAGetLocalInfo(da,&info)); + PetscCall(DMGetLocalVector(da,&xlocal)); + PetscCall(DMGetLocalVector(da,&flocal)); + + PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); + PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); + + PetscCall(VecGetArray(xlocal,&x_vec)); + PetscCall(VecGetArray(flocal,&f_vec)); + + Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; + Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + ff[j][i].v = scdiag * xx[j][i].v; + ff[j][i].u = scdiag * xx[j][i].u; + } else { + ve = xx[j][i+1].v; + vw = xx[j][i-1].v; + vn = xx[j+1][i].v; + vs = xx[j-1][i].v; + ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) + - darea * (*(user->f))(x,y); + ue = xx[j][i+1].u; + uw = xx[j][i-1].u; + un = xx[j+1][i].u; + us = xx[j-1][i].u; + ff[j][i].u = - darea * xx[j][i].v + + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); + } + } + } + + PetscCall(VecRestoreArray(xlocal,&x_vec)); + PetscCall(VecRestoreArray(flocal,&f_vec)); + + PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); + PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); + PetscCall(DMRestoreLocalVector(da,&xlocal)); + PetscCall(DMRestoreLocalVector(da,&flocal)); + + return 0; +} + + +PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1./ (info.mx - 1); + hy = 1./ (info.my - 1); + darea = hx * hy; // multiply FD equations by this + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = scdiag * x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = 0.0; + } else { + y_v[j][i] = -darea * x_v[j][i]; + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + + +PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; +// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = scdiag * x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) +{ + Vec J00X; + Vec J00Y; + Vec J10X; + Vec J10Y; + Vec J11X; + Vec J11Y; + + struct SubMatrixCtx * J00ctx; + struct SubMatrixCtx * J10ctx; + struct SubMatrixCtx * J11ctx; + struct JacobianCtx * jctx; + + PetscFunctionBeginUser; + + PetscCall(MatShellGetContext(J,&(jctx))); + + PetscCall(VecSet(Y,0.0)); + + Mat J00 = jctx->submats[0]; + PetscCall(MatShellGetContext(J00,&(J00ctx))); + PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); + PetscCall(MatMult(J00,J00X,J00Y)); + PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); + + Mat J10 = jctx->submats[2]; + PetscCall(MatShellGetContext(J10,&(J10ctx))); + PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); + PetscCall(MatMult(J10,J10X,J10Y)); + PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); + + Mat J11 = jctx->submats[3]; + PetscCall(MatShellGetContext(J11,&(J11ctx))); + PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); + PetscCall(MatMult(J11,J11X,J11Y)); + PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); + + + PetscFunctionReturn(0); +} + + +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) +{ + PetscFunctionBeginUser; + + PetscInt M; + PetscInt N; + Mat block; + DM dm0; + PetscInt dof; + + struct UserCtx0 * ctx0; + struct JacobianCtx * jctx; + struct SubMatrixCtx * subctx; + + PetscCall(MatShellGetContext(J,&(jctx))); + DM * subdms = jctx->subdms; + + PetscInt nsubmats = nfields*nfields; + PetscCall(PetscCalloc1(nsubmats,submats)); + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); + PetscInt subblockrows = M*N; + PetscInt subblockcols = M*N; + Mat * submat_arr = *submats; + + for (int i = 0; i <= nsubmats - 1; i += 1) + { + PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); + PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); + PetscCall(MatSetType(block,MATSHELL)); + PetscCall(PetscMalloc1(1,&(subctx))); + PetscInt rowidx = i / dof; + PetscInt colidx = (i)%(dof); + subctx->rows = &(irow[rowidx]); + subctx->cols = &(icol[colidx]); + PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); + PetscCall(MatSetDM(block,subdms[rowidx])); + PetscCall(MatShellSetContext(block,subctx)); + PetscCall(MatSetUp(block)); + submat_arr[i] = block; + } + PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); + + PetscFunctionReturn(0); +} + + +PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) +{ + PetscFunctionBeginUser; + + jctx->subdms = subdms; + jctx->fields = fields; + + PetscFunctionReturn(0); +} diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c new file mode 100644 index 0000000000..04f2a02927 --- /dev/null +++ b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c @@ -0,0 +1,553 @@ + +// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + +static char help[] = +"Solve the linear biharmonic equation in 2D. Equation is\n" +" Lap^2 u = f\n" +"where Lap = - grad^2 is the positive Laplacian, equivalently\n" +" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" +"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" +"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" +"2x2 block system with SPD Laplacian blocks on the diagonal:\n" +" | Lap | 0 | | v | | f | \n" +" |-----|-----| |---| = |---| \n" +" | -I | Lap | | u | | 0 | \n" +"Includes manufactured, polynomial exact solution. The discretization is\n" +"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" +"Recommended preconditioning combines fieldsplit:\n" +" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" +"with multigrid as the preconditioner for the diagonal blocks:\n" +" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" +"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" +"monolithic multigrid (-pc_type mg|gamg).\n\n"; + +#include + +typedef struct { + PetscReal v, u; +} Field; + +typedef struct { + PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side +} BiharmCtx; + +struct JacobianCtx +{ + DM * subdms; + IS * fields; + Mat * submats; +} ; + +struct SubMatrixCtx +{ + IS * rows; + IS * cols; +} ; + +static PetscReal c(PetscReal x) { + return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); +} + +static PetscReal ddc(PetscReal x) { + return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal d4c(PetscReal x) { + return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { + return c(x) * c(y); +} + +static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { + return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u +} + +static PetscReal f_fcn(PetscReal x, PetscReal y) { + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u +} + +extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); +extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); +extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); +extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); + +int main(int argc,char **argv) { + DM da; + SNES snes; + Vec w, w_initial, w_exact; + BiharmCtx user; + Field **aW; + PetscReal normv, normu, errv, erru; + DMDALocalInfo info; + IS *fields; + DM *subdms; + PetscInt nfields; + + struct JacobianCtx jctx0; + Mat J; + + PetscCall(PetscInitialize(&argc,&argv,NULL,help)); + + user.f = &f_fcn; + PetscCall(DMDACreate2d(PETSC_COMM_WORLD, + DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, + 33,33,PETSC_DECIDE,PETSC_DECIDE, + 2,1, // degrees of freedom, stencil width + NULL,NULL,&da)); + PetscCall(DMSetApplicationContext(da,&user)); + PetscCall(DMSetFromOptions(da)); + PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates + PetscCall(DMSetMatType(da, MATSHELL)); + PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); + PetscCall(DMDASetFieldName(da,0,"v")); + PetscCall(DMDASetFieldName(da,1,"u")); + PetscCall(DMCreateMatrix(da,&J)); + + PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); + PetscCall(SNESSetDM(snes,da)); + PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); + PetscCall(SNESSetType(snes,SNESKSPONLY)); + PetscCall(SNESSetFromOptions(snes)); + + PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); + PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); + + PetscCall(MatSetDM(J,da)); + PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); + PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); + PetscCall(MatShellSetContext(J,&(jctx0))); + PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); + + PetscCall(DMGetGlobalVector(da,&w_initial)); + PetscCall(VecSet(w_initial,0.0)); + PetscCall(SNESSolve(snes,NULL,w_initial)); + // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(DMRestoreGlobalVector(da,&w_initial)); + PetscCall(DMDestroy(&da)); + + PetscCall(SNESGetSolution(snes,&w)); + PetscCall(SNESGetDM(snes,&da)); + PetscCall(DMDAGetLocalInfo(da,&info)); + + PetscCall(DMCreateGlobalVector(da,&w_exact)); + PetscCall(DMDAVecGetArray(da,w_exact,&aW)); + PetscCall(FormExactWLocal(&info,aW,&user)); + + PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); + // PetscCall(VecView(w_exact,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); + PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); + PetscCall(VecAXPY(w,-1.0,w_exact)); + PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); + PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); + PetscCall(PetscPrintf(PETSC_COMM_WORLD, + "done on %d x %d grid ...\n" + " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", + info.mx,info.my,errv/normv,erru/normu)); + + + PetscCall(ISDestroy(&(fields[0]))); + PetscCall(ISDestroy(&(fields[1]))); + PetscCall(PetscFree(fields)); + PetscCall(DMDestroy(&(subdms[0]))); + PetscCall(DMDestroy(&(subdms[1]))); + PetscCall(PetscFree(subdms)); + PetscCall(VecDestroy(&w_exact)); + PetscCall(MatDestroy(&J)); + PetscCall(SNESDestroy(&snes)); + PetscCall(PetscFinalize()); + return 0; +} + +PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, x, y; + PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); + hx = (xymax[0] - xymin[0]) / (info->mx - 1); + hy = (xymax[1] - xymin[1]) / (info->my - 1); + for (j = info->ys; j < info->ys + info->ym; j++) { + y = j * hy; + for (i = info->xs; i < info->xs + info->xm; i++) { + x = i * hx; + aW[j][i].u = u_exact_fcn(x,y); + aW[j][i].v = lap_u_exact_fcn(x,y); + } + } + return 0; +} + + +PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) +{ + Vec xlocal, flocal; + DMDALocalInfo info; + DM da; + PetscScalar *x_vec, *f_vec; + + BiharmCtx *user; + + PetscCall(SNESGetDM(snes,&da)); + + PetscCall(DMGetApplicationContext(da,&user)); + + PetscCall(DMDAGetLocalInfo(da,&info)); + PetscCall(DMGetLocalVector(da,&xlocal)); + PetscCall(DMGetLocalVector(da,&flocal)); + + PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); + PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); + + PetscCall(VecGetArray(xlocal,&x_vec)); + PetscCall(VecGetArray(flocal,&f_vec)); + + Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; + Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + ff[j][i].v = xx[j][i].v; + ff[j][i].u = xx[j][i].u; + } else { + ve = xx[j][i+1].v; + vw = xx[j][i-1].v; + vn = xx[j+1][i].v; + vs = xx[j-1][i].v; + ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) + - (*(user->f))(x,y); + ue = xx[j][i+1].u; + uw = xx[j][i-1].u; + un = xx[j+1][i].u; + us = xx[j-1][i].u; + ff[j][i].u = -xx[j][i].v + + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); + } + } + } + + PetscCall(VecRestoreArray(xlocal,&x_vec)); + PetscCall(VecRestoreArray(flocal,&f_vec)); + + PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); + PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); + PetscCall(DMRestoreLocalVector(da,&xlocal)); + PetscCall(DMRestoreLocalVector(da,&flocal)); + + return 0; +} + + +PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1./ (info.mx - 1); + hy = 1./ (info.my - 1); + darea = hx * hy; // multiply FD equations by this + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + // print info.ys to screen + // PetscCall(PetscPrintf(PETSC_COMM_WORLD,"info.ys = %d\n",info.xm)); + for (j = info.ys+1; j < info.ys + info.ym-1; j++) { + for (i = info.xs+1; i < info.xs + info.xm-1; i++) { + y_v[j][i] = -x_v[j][i]; + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + + +PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; +// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) +{ + Vec J00X; + Vec J00Y; + Vec J10X; + Vec J10Y; + Vec J11X; + Vec J11Y; + + struct SubMatrixCtx * J00ctx; + struct SubMatrixCtx * J10ctx; + struct SubMatrixCtx * J11ctx; + struct JacobianCtx * jctx; + + PetscFunctionBeginUser; + + PetscCall(MatShellGetContext(J,&(jctx))); + + PetscCall(VecSet(Y,0.0)); + + Mat J00 = jctx->submats[0]; + PetscCall(MatShellGetContext(J00,&(J00ctx))); + PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); + PetscCall(MatMult(J00,J00X,J00Y)); + PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); + + Mat J10 = jctx->submats[2]; + PetscCall(MatShellGetContext(J10,&(J10ctx))); + PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); + PetscCall(MatMult(J10,J10X,J10Y)); + PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); + + Mat J11 = jctx->submats[3]; + PetscCall(MatShellGetContext(J11,&(J11ctx))); + PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); + PetscCall(MatMult(J11,J11X,J11Y)); + PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); + + + PetscFunctionReturn(0); +} + + +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) +{ + PetscFunctionBeginUser; + + PetscInt M; + PetscInt N; + Mat block; + DM dm0; + PetscInt dof; + + struct UserCtx0 * ctx0; + struct JacobianCtx * jctx; + struct SubMatrixCtx * subctx; + + PetscCall(MatShellGetContext(J,&(jctx))); + DM * subdms = jctx->subdms; + + PetscInt nsubmats = nfields*nfields; + PetscCall(PetscCalloc1(nsubmats,submats)); + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); + PetscInt subblockrows = M*N; + PetscInt subblockcols = M*N; + Mat * submat_arr = *submats; + + for (int i = 0; i <= nsubmats - 1; i += 1) + { + PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); + PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); + PetscCall(MatSetType(block,MATSHELL)); + PetscCall(PetscMalloc1(1,&(subctx))); + PetscInt rowidx = i / dof; + PetscInt colidx = (i)%(dof); + subctx->rows = &(irow[rowidx]); + subctx->cols = &(icol[colidx]); + PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); + PetscCall(MatSetDM(block,subdms[rowidx])); + PetscCall(MatShellSetContext(block,subctx)); + PetscCall(MatSetUp(block)); + submat_arr[i] = block; + } + PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); + + PetscFunctionReturn(0); +} + + +PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) +{ + PetscFunctionBeginUser; + + jctx->subdms = subdms; + jctx->fields = fields; + + PetscFunctionReturn(0); +} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e89f4c4476..943331f3c3 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1097,3 +1097,7 @@ def test_submatrices(self): # add tests for all new callbacks # def test_create_whole_matmult(): + + +# // add coupled test with 3 targets +# // add coupled test with 1 target? From 1ddc053f3d8565deaecd73dfa1936d36a9edfe43 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 00:24:42 +0100 Subject: [PATCH 04/53] compiler: Working reuse_efunc for petscbundles --- devito/passes/iet/engine.py | 27 ++++++--------------------- devito/petsc/types/object.py | 3 --- devito/types/object.py | 2 ++ 3 files changed, 8 insertions(+), 24 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index a90f4ba275..0c3c7481bf 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -151,7 +151,7 @@ def apply(self, func, **kwargs): if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) # TODO: fix for petsc bundles - # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) + efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -355,7 +355,7 @@ def reuse_efuncs(root, efuncs, sregistry=None): if isinstance(efunc, AsyncCallable): mapper[len(mapper)] = (efunc, [efunc]) continue - # from IPython import embed; embed() + afunc = abstract_efunc(efunc) key = afunc._signature() @@ -396,7 +396,7 @@ def abstract_efunc(efunc): functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) mapper = abstract_objects(functions) - # from IPython import embed; embed() + efunc = Uxreplace(mapper).visit(efunc) efunc = efunc._rebuild(name='foo') @@ -416,8 +416,8 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - # keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] - keys = [Bundle, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + # keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -428,7 +428,6 @@ def abstract_objects(objects0, sregistry=None): for i in objects: abstract_object(i, mapper, sregistry) - # from IPython import embed; embed() return mapper @@ -455,6 +454,7 @@ def _(i, mapper, sregistry): @abstract_object.register(Array) +@abstract_object.register(PETScArray) def _(i, mapper, sregistry): if isinstance(i, Lock): name = sregistry.make_name(prefix='lock') @@ -472,20 +472,6 @@ def _(i, mapper, sregistry): mapper[i.dmap] = v.dmap -# @abstract_object.register(PETScArray) -# def _(i, mapper, sregistry): -# name = sregistry.make_name(prefix='xx') - -# v = i._rebuild(name=name, initializer=None, alias=True) - -# mapper.update({ -# i: v, -# i.indexed: v.indexed, -# i.dmap: v.dmap, -# i._C_symbol: v._C_symbol, -# }) - - @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') @@ -508,7 +494,6 @@ def _(i, mapper, sregistry): name = sregistry.make_name(prefix='o') v = i._rebuild(name) - mapper[i] = v diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 7821e316d9..a285863d14 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -218,9 +218,6 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -JacobianStructCast = cast('struct JacobianCtx *') - - class PETScArrayObject(ArrayObject): _data_alignment = False diff --git a/devito/types/object.py b/devito/types/object.py index 6eea49f6a1..1153ce81eb 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -249,8 +249,10 @@ class LocalCompositeObject(CompositeObject, LocalType): """ __rargs__ = ('name', 'pname', 'fields') + __rkwargs__ = ('modifier', 'liveness') def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): + self.modifier = modifier dtype = CustomDtype(f"struct {pname}", modifier=modifier) Object.__init__(self, name, dtype, None) self._pname = pname From 68346110438eadb1d43e27786f22fe1f333ddbca Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 18:16:54 +0100 Subject: [PATCH 05/53] dsl/compiler: Zero columns for essential bcs --- devito/passes/iet/engine.py | 6 +-- devito/petsc/solve.py | 45 ++++++++++++++++--- devito/petsc/types/types.py | 3 -- devito/tools/utils.py | 2 +- examples/petsc/Poisson/01_poisson.py | 22 +++++---- examples/petsc/Poisson/02_laplace.py | 6 ++- examples/petsc/Poisson/04_poisson.py | 4 +- .../petsc/random/biharmonic/02_biharmonic.py | 18 +++----- tests/test_petsc.py | 16 ++++--- 9 files changed, 79 insertions(+), 43 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 0c3c7481bf..fa34a4946c 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -416,8 +416,7 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] - # keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, ArrayBasic, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -453,8 +452,7 @@ def _(i, mapper, sregistry): }) -@abstract_object.register(Array) -@abstract_object.register(PETScArray) +@abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): name = sregistry.make_name(prefix='lock') diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index c34b829798..f5893fec05 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -52,10 +52,18 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - formfuncs, formrhs = zip( - *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + # TODO: Ensure essential BCs are handled first - this is required to + # maintain the symmetry of the operator when "constructing" the Jacobian. + eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + + # TODO: scaling + + formfuncs, formrhs = map( + lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], + zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) ) - matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + for e in (i if isinstance(i, (tuple)) else [i])] initialguess = [ eq for eq in @@ -88,7 +96,14 @@ def build_matvec_eqns(self, eq, target, arrays): def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): + # TODO: SCALING + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) + # and zero symmetrically. rhs = arrays['x'] + zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) + return (Eq(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) @@ -96,9 +111,16 @@ def make_matvec(self, eq, F_target, arrays, targets): def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # TODO: CHECK THIS - rhs = arrays[main_target]['x'] - eq.rhs - # rhs = 0. + # The initial guess already satisfies the essential boundary conditions, + # so this term will always be zero. It's included here to allow + # testing of the Jacobian via finite differences. + rhs = arrays['x'] - eq.rhs + # Create equation to handle essential boundary conditions - we + # move the essential BCs to the right-hand side + # and zero the corresponding column in the Jacobian. + # TODO: extend this to mixed problems + move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (Eq(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) else: if isinstance(F_target, (int, float)): rhs = F_target @@ -215,6 +237,17 @@ def generate_arrays_combined(self, *targets): class EssentialBC(Eq): + """ + A special equation representing an essential boundary condition. + It is used to handle essential boundary conditions in the PETSc solver. + Until PetscSection + DMDA is supported, we treat essential boundary conditions + as trivial equations in the solver, where we place 1.0 on the diagonal of the jacobian, + zero symmetrically and move the essential boundary condition to the right-hand side. + + NOTE: When users define essential boundary conditions, they need to ensure that + the SubDomains do not overlap. Solver will still run but may see unexpected behaviour + along boundaries. + """ pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 973d558f38..afab9fff63 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -188,7 +188,6 @@ def targets(self): return (self.target,) -# TODO: should this acc inherhit from fielddata? maybe not? class MultipleFieldData(FieldData): def __init__(self, targets, arrays, submatrices=None): self._targets = as_tuple(targets) @@ -196,8 +195,6 @@ def __init__(self, targets, arrays, submatrices=None): self._submatrices = submatrices self._formfuncs = [] - # pass - def extend_formfuncs(self, formfuncs): self._formfuncs.extend(formfuncs) diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0657146040..0920298356 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -337,7 +337,7 @@ def sorted_priority(items, priority): """ def key(i): - for cls in sorted(priority, key=priority.get, reverse=True): + for cls in sorted(priority, key=priority.get): if isinstance(i, cls): v = priority[cls] break diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6d89091cd6..86fdfc6dd3 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -16,6 +16,8 @@ # Subdomains to implement BCs +# NOTE: For essential BCs, we ensure the SubDomains do not overlap + class SubTop(SubDomain): name = 'subtop' @@ -37,7 +39,7 @@ class SubLeft(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('left', 1), y: y} + return {x: ('left', 1), y: ('middle', 1, 1)} class SubRight(SubDomain): @@ -45,7 +47,7 @@ class SubRight(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('right', 1), y: y} + return {x: ('right', 1), y: ('middle', 1, 1)} sub1 = SubTop() @@ -64,6 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) +n_values = [6] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -76,6 +79,7 @@ def analytical(x, y): phi = Function(name='phi', grid=grid, space_order=2, dtype=np.float64) rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + bc = Function(name='bc', grid=grid, space_order=2, dtype=np.float64) eqn = Eq(rhs, phi.laplace, subdomain=grid.interior) @@ -87,11 +91,13 @@ def analytical(x, y): 2.0*X*(Y-1.0)*(Y - 2.0*X + X*Y + 2.0) ) * np.float64(np.exp(X-Y)) + bc.data[:] = 0.0 + # # Create boundary condition expressions using subdomains - bcs = [EssentialBC(phi, np.float64(0.), subdomain=sub1)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub2)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub3)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub4)] + bcs = [EssentialBC(phi, bc, subdomain=sub1)] + bcs += [EssentialBC(phi, bc, subdomain=sub2)] + bcs += [EssentialBC(phi, bc, subdomain=sub3)] + bcs += [EssentialBC(phi, bc, subdomain=sub4)] exprs = [eqn] + bcs petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) @@ -108,5 +114,5 @@ def analytical(x, y): slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -assert slope > 1.9 -assert slope < 2.1 +# assert slope > 1.9 +# assert slope < 2.1 diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py index 780da10ec5..c84855137f 100644 --- a/examples/petsc/Poisson/02_laplace.py +++ b/examples/petsc/Poisson/02_laplace.py @@ -23,12 +23,14 @@ # Subdomains to implement BCs +# NOTE: For essential BCs, we ensure the SubDomains do not overlap + class SubTop(SubDomain): name = 'subtop' def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -36,7 +38,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/Poisson/04_poisson.py b/examples/petsc/Poisson/04_poisson.py index 44f34ec8f9..637ce44076 100644 --- a/examples/petsc/Poisson/04_poisson.py +++ b/examples/petsc/Poisson/04_poisson.py @@ -26,7 +26,7 @@ class SubTop(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -34,7 +34,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index 8c15f3d3a2..6db85d6d96 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -76,8 +76,7 @@ def f_fcn(x, y): Lx = np.float64(1.) Ly = np.float64(1.) -# n_values = [33, 53, 73, 93, 113] -n_values = [33] +n_values = [33, 53, 73, 93, 113] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -135,14 +134,11 @@ def f_fcn(x, y): v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) v_errors.append(v_error) -# u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) -# v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) +u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) +v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) -# assert u_slope > 1.9 -# assert u_slope < 2.1 +assert u_slope > 1.9 +assert u_slope < 2.1 -# assert v_slope > 1.9 -# assert v_slope < 2.1 -print(op.ccode) -print("u_errors:", u_errors) -print("v_errors:", v_errors) +assert v_slope > 1.9 +assert v_slope < 2.1 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 943331f3c3..5c1346287f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -339,7 +339,7 @@ def test_separate_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _ = separate_eqn(eval(eqn), eval(target)) + b, F, _, _= separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -471,7 +471,7 @@ def test_separate_eval_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _ = separate_eqn(eval(eqn), eval(target)) + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -793,6 +793,9 @@ def test_essential_bcs(): Verify that PETScSolve returns the correct output with essential boundary conditions. """ + + # SubDomains used for essential boundary conditions + # should not overlap. class SubTop(SubDomain): name = 'subtop' @@ -814,7 +817,7 @@ class SubLeft(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('left', 1), y: y} + return {x: ('left', 1), y: ('middle', 1, 1)} sub3 = SubLeft() class SubRight(SubDomain): @@ -822,7 +825,7 @@ class SubRight(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('right', 1), y: y} + return {x: ('right', 1), y: ('middle', 1, 1)} sub4 = SubRight() subdomains = (sub1, sub2, sub3, sub4) @@ -913,9 +916,10 @@ def test_coupled_vs_non_coupled(self): callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] # Solving for multiple fields within the same matrix system requires - # additional machinery and more callback functions + # less callback functions than solving them separately. + # TODO: check reuse of callback functions where appropriate assert len(callbacks1) == 8 - assert len(callbacks2) == 11 + assert len(callbacks2) == 6 # Check fielddata type fielddata1 = petsc1[0].rhs.fielddata From 9642278b4c19ec0c0adfc1a2d0ba0208569e5b37 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 18:51:52 +0100 Subject: [PATCH 06/53] dsl: Compatible scaling of jacobian --- devito/petsc/solve.py | 52 +++++++++++++------ examples/petsc/Poisson/01_poisson.py | 6 +-- .../petsc/random/biharmonic/02_biharmonic.py | 4 +- .../random/biharmonic/biharmonic_matfree.c | 2 +- tests/test_petsc.py | 13 ++--- 5 files changed, 49 insertions(+), 28 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index f5893fec05..d675fb0cb9 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,6 +1,7 @@ from functools import singledispatch import sympy +import numpy as np from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative @@ -28,9 +29,11 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns + self.cell_area = None def build_eq(self): target, funcs, fielddata = self.linear_solve_args() + # Placeholder equation for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, @@ -44,6 +47,7 @@ def build_eq(self): def linear_solve_args(self): target, eqns = next(iter(self.target_eqns.items())) eqns = as_tuple(eqns) + self.cell_area = np.prod(target.grid.spacing_symbols) funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) @@ -52,19 +56,17 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - # TODO: Ensure essential BCs are handled first - this is required to - # maintain the symmetry of the operator when "constructing" the Jacobian. + # TODO: Apply essential boundary conditions first to preserve operator symmetry + # during Jacobian "construction". eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - # TODO: scaling - formfuncs, formrhs = map( lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) ) matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] - + from IPython import embed; embed() initialguess = [ eq for eq in (self.make_initial_guess(e, target, arrays) for e in eqns) @@ -96,17 +98,17 @@ def build_matvec_eqns(self, eq, target, arrays): def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # TODO: SCALING # NOTE: Until PetscSection + DMDA is supported, we leave # the essential BCs in the solver. # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) # and zero symmetrically. rhs = arrays['x'] zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) - return (Eq(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) + return (EssentialBC(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays['y'], rhs, subdomain=eq.subdomain) def make_formfunc(self, eq, F_target, arrays, targets): @@ -120,17 +122,19 @@ def make_formfunc(self, eq, F_target, arrays, targets): # and zero the corresponding column in the Jacobian. # TODO: extend this to mixed problems move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (Eq(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) + return (EssentialBC(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) else: if isinstance(F_target, (int, float)): - rhs = F_target + rhs = F_target * self.cell_area else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays['f'], rhs, subdomain=eq.subdomain) def make_rhs(self, eq, b, arrays): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = rhs*self.cell_area return Eq(arrays['b'], rhs, subdomain=eq.subdomain) def make_initial_guess(self, eq, target, arrays): @@ -175,6 +179,8 @@ def linear_solve_args(self): all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) + self.cell_area = np.prod(all_data.grid.spacing_symbols) + for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) self.update_jacobian(eqns, target, jacobian, arrays[target]) @@ -189,11 +195,11 @@ def linear_solve_args(self): def update_jacobian(self, eqns, target, jacobian, arrays): for submat, mtvs in jacobian.submatrices[target].items(): - matvecs = [ - self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) - for eq in eqns - ] + # TODO: maintain symmetry for coupled + matvecs = [e for i in [self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) + for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] matvecs = [m for m in matvecs if m is not None] + if matvecs: jacobian.set_submatrix(target, submat, matvecs) @@ -210,14 +216,14 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): - # TODO: CHECK THIS rhs = arrays[main_target]['x'] - eq.rhs else: if isinstance(zeroed, (int, float)): - rhs = zeroed + rhs = zeroed * self.cell_area else: rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) @@ -235,6 +241,20 @@ def generate_arrays_combined(self, *targets): for target in targets } + def make_matvec(self, eq, F_target, arrays, targets): + if isinstance(eq, EssentialBC): + # TODO: SHOULDN'T NEED this subclass once fixed mixed problems + symmetry + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal + # and zero symmetrically. FIX FOR MIXED PROBLEMS + rhs = arrays['x'] + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area + return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + class EssentialBC(Eq): """ @@ -246,7 +266,7 @@ class EssentialBC(Eq): NOTE: When users define essential boundary conditions, they need to ensure that the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - along boundaries. + at boundaries. """ pass diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 86fdfc6dd3..6f24f74f64 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [6] +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -114,5 +114,5 @@ def analytical(x, y): slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -# assert slope > 1.9 -# assert slope < 2.1 +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index 6db85d6d96..bc7b54654d 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -20,7 +20,7 @@ class SubTop(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -28,7 +28,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c index 53121473e7..2fb432382f 100644 --- a/examples/petsc/random/biharmonic/biharmonic_matfree.c +++ b/examples/petsc/random/biharmonic/biharmonic_matfree.c @@ -97,7 +97,7 @@ int main(int argc,char **argv) { user.f = &f_fcn; PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 33,33,PETSC_DECIDE,PETSC_DECIDE, + 6,6,PETSC_DECIDE,PETSC_DECIDE, 2,1, // degrees of freedom, stencil width NULL,NULL,&da)); PetscCall(DMSetApplicationContext(da,&user)); diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5c1346287f..c136160ecc 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -108,15 +108,15 @@ def test_petsc_solve(): rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) assert str(action_expr[-1].expr.rhs) == ( - 'x_f[x + 1, y + 2]/ctx0->h_x**2' + '(x_f[x + 1, y + 2]/ctx0->h_x**2' ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_x**2' ' + x_f[x + 3, y + 2]/ctx0->h_x**2' ' + x_f[x + 2, y + 1]/ctx0->h_y**2' ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_y**2' - ' + x_f[x + 2, y + 3]/ctx0->h_y**2' + ' + x_f[x + 2, y + 3]/ctx0->h_y**2)*ctx0->h_x*ctx0->h_y' ) - assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' + assert str(rhs_expr[-1].expr.rhs) == 'ctx0->h_x*ctx0->h_y*g[x + 2, y + 2]' # Check the iteration bounds are correct. assert op.arguments().get('x_m') == 0 @@ -836,7 +836,7 @@ def define(self, dimensions): # Solving Ax=b where A is the identity matrix v.data[:] = 5.0 - eqn = Eq(u, v) + eqn = Eq(u, v, subdomain=grid.interior) bcs = [EssentialBC(u, 1., subdomain=sub1)] # top bcs += [EssentialBC(u, 2., subdomain=sub2)] # bottom @@ -1083,11 +1083,12 @@ def test_submatrices(self): j00 = submatrices.get_submatrix(e, 'J00') j11 = submatrices.get_submatrix(g, 'J11') + # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ - + ' Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2)))' + + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ - + ' Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2)))' + + ' h_x*h_y*(Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2))))' # Check the derivative wrt fields assert j00['derivative_wrt'] == e From ce9150f71a2d07513b6f9e9d6787a7e4ab004ffa Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 30 May 2025 17:38:20 +0100 Subject: [PATCH 07/53] dsl: Scale boundary rows of jacobian --- devito/passes/iet/engine.py | 7 +- devito/petsc/iet/routines.py | 44 +-- devito/petsc/solve.py | 255 ++++++++++++------ devito/petsc/types/array.py | 13 +- devito/petsc/types/types.py | 12 + devito/tools/utils.py | 1 - examples/petsc/Poisson/01_poisson.py | 8 +- .../petsc/random/biharmonic/02_biharmonic.py | 18 +- tests/test_petsc.py | 111 +++++++- 9 files changed, 343 insertions(+), 126 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fa34a4946c..ad0ebefd1e 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -10,14 +10,13 @@ from devito.passes import needs_transfer from devito.symbolics import FieldFromComposite, FieldFromPointer from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass -from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, +from devito.types import (Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension -from devito.petsc.types import PETScArray __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -150,7 +149,6 @@ def apply(self, func, **kwargs): # Minimize code size if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) - # TODO: fix for petsc bundles efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -423,7 +421,6 @@ def abstract_objects(objects0, sregistry=None): # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() - for i in objects: abstract_object(i, mapper, sregistry) @@ -473,7 +470,6 @@ def _(i, mapper, sregistry): @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') - components = [mapper[f] for f in i.components] v = i._rebuild(name=name, components=components, alias=True) @@ -492,6 +488,7 @@ def _(i, mapper, sregistry): name = sregistry.make_name(prefix='o') v = i._rebuild(name) + mapper[i] = v diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8674dc4df4..7c824ad638 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -793,7 +793,7 @@ def _whole_formfunc_body(self, body, fielddata): bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] - body = self.bundle_residual(body, bundles) + body = self.residual_bundle(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -1046,13 +1046,9 @@ def _submat_callback_body(self): stacks=(get_ctx, deref_subdm), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - - def bundle_residual(self, body, bundles): - fbundle = bundles['f'] - xbundle = bundles['x'] + def residual_bundle(self, body, bundles): mapper1 = bundles['bundle_mapper'] - indexeds = FindSymbols('indexeds').visit(body) subss = {} @@ -1067,7 +1063,6 @@ def bundle_residual(self, body, bundles): return body - class BaseObjectBuilder: """ A base class for constructing objects needed for a PETSc solver. @@ -1180,7 +1175,6 @@ def _extend_build(self, base_dict): base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') base_dict[f'{key}F'] = CallbackVec(f'{key}F') - # Bundle objects/metadata required by the coupled residual callback f_components = [] x_components = [] @@ -1194,9 +1188,13 @@ def _extend_build(self, base_dict): f_components.append(f_arr) x_components.append(x_arr) - # TODO: to group them, maybe pass in struct name as arg to petscbundle - fbundle = PetscBundle(name='f_bundle', components=f_components) - xbundle = PetscBundle(name='x_bundle', components=x_components) + bundle_pname = sreg.make_name(prefix='Field') + fbundle = PetscBundle( + name='f_bundle', components=f_components, pname=bundle_pname + ) + xbundle = PetscBundle( + name='x_bundle', components=x_components, pname=bundle_pname + ) # Build the bundle mapper for i, t in enumerate(targets): @@ -1205,7 +1203,6 @@ def _extend_build(self, base_dict): bundle_mapper[f_arr.base] = fbundle bundle_mapper[x_arr.base] = xbundle - base_dict['bundles'] = { 'f': fbundle, 'x': xbundle, @@ -1346,6 +1343,10 @@ def _setup(self): self.snes_ctx] ) + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + dmda_calls = self._create_dmda_calls(dmda) mainctx = sobjs['userctx'] @@ -1377,6 +1378,7 @@ def _setup(self): ksp_set_from_ops, matvec_operation, formfunc_operation, + snes_set_options, call_struct_callback, mat_set_dm, calls_set_app_ctx, @@ -1504,6 +1506,10 @@ def _setup(self): self.snes_ctx] ) + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + dmda_calls = self._create_dmda_calls(dmda) mainctx = sobjs['userctx'] @@ -1574,6 +1580,7 @@ def _setup(self): ksp_set_from_ops, matvec_operation, formfunc_operation, + snes_set_options, call_struct_callback, mat_set_dm, calls_set_app_ctx, @@ -1670,17 +1677,12 @@ def _execute_solve(self): """ objs = self.objs sobjs = self.solver_objs - - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - - # rhs_callbacks = self.cbbuilder.formrhs - xglob = sobjs['xglobal'] - bglob = sobjs['bglobal'] + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) targets = self.injectsolve.expr.rhs.fielddata.targets - # TODO: optimise the ccode generated here + # TODO: Optimise the ccode generated here pre_solve = () post_solve = () @@ -1693,10 +1695,10 @@ def _execute_solve(self): s = sobjs[f'scatter{name}'] pre_solve += ( - # TODO: switch to createwitharray and move to setup + # TODO: Switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), - # TODO: need to call reset array + # TODO: Need to call reset array self.timedep.place_array(t), petsc_call( 'DMLocalToGlobal', diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index d675fb0cb9..278a5092b8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -2,6 +2,8 @@ import sympy import numpy as np +from itertools import chain +from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative @@ -29,7 +31,14 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns + # TODO: make this _ self.cell_area = None + # self._centre_stencils = set() + self._diag_scale = defaultdict(set) + + @property + def diag_scale(self): + return self._diag_scale def build_eq(self): target, funcs, fielddata = self.linear_solve_args() @@ -56,17 +65,35 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - # TODO: Apply essential boundary conditions first to preserve operator symmetry - # during Jacobian "construction". + # Apply essential boundary conditions first to preserve + # operator symmetry during Jacobian "construction" eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + matvecs = [e for eq in eqns for e in self.build_matvec_eq(eq, target, arrays)] + formfuncs, formrhs = map( - lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], - zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) + lambda x: [e for i in x for e in i], + zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) ) - matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] - for e in (i if isinstance(i, (tuple)) else [i])] - from IPython import embed; embed() + + # self._centre_stencils[arrays['x']].update( + stencils = set() + for eq in matvecs: + if not isinstance(eq, EssentialBC): + stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) + stencils.add(stencil) + + if len(stencils) > 1: + # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock + # for now just set to trivial 1.0 + scale = 1.0 + else: + scale = next(iter(stencils)) + + # from IPython import embed; embed() + matvecs = self.scale_essential_bcs(matvecs, scale) + formfuncs = self.scale_essential_bcs(formfuncs, scale) + initialguess = [ eq for eq in (self.make_initial_guess(e, target, arrays) for e in eqns) @@ -82,19 +109,18 @@ def generate_field_data(self, eqns, target, arrays): arrays=arrays ) - def build_function_eqns(self, eq, target, arrays): + def build_function_eq(self, eq, target, arrays): b, F_target, _, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) - def build_matvec_eqns(self, eq, target, arrays): + def build_matvec_eq(self, eq, target, arrays): b, F_target, _, targets = separate_eqn(eq, target) - if not F_target: - return None - matvec = self.make_matvec(eq, F_target, arrays, targets) - return matvec + if F_target: + return self.make_matvec(eq, F_target, arrays, targets) + return (None,) def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): @@ -103,39 +129,41 @@ def make_matvec(self, eq, F_target, arrays, targets): # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) # and zero symmetrically. rhs = arrays['x'] - zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) - return (EssentialBC(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) + return (zero_row, zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + rhs = rhs.subs(self.time_mapper) * self.cell_area + # TODO: Average centre stencils if they vary, to scale essential BC rows. + # self.centre = centre_stencil(rhs, arrays['x'], as_coeff=True) + # stencil = centre_stencil(rhs, arrays['x'], as_coeff=True) + # self._centre_stencils[arrays['x']].add(stencil) + # self._centre_stencils.add(stencil) + + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # The initial guess already satisfies the essential boundary conditions, - # so this term will always be zero. It's included here to allow - # testing of the Jacobian via finite differences. + # The initial guess satisfies the essential BCs, so this term is zero. + # Still included to support Jacobian testing via finite differences. rhs = arrays['x'] - eq.rhs - # Create equation to handle essential boundary conditions - we - # move the essential BCs to the right-hand side - # and zero the corresponding column in the Jacobian. - # TODO: extend this to mixed problems - move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (EssentialBC(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) + zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) + # Move essential boundary condition to the right-hand side + zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (zero_row, zero_col) else: if isinstance(F_target, (int, float)): rhs = F_target * self.cell_area else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['f'], rhs, subdomain=eq.subdomain) + rhs = rhs.subs(self.time_mapper) * self.cell_area + return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) def make_rhs(self, eq, b, arrays): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) - rhs = rhs*self.cell_area - return Eq(arrays['b'], rhs, subdomain=eq.subdomain) + rhs = rhs * self.cell_area + return as_tuple(Eq(arrays['b'], rhs, subdomain=eq.subdomain)) def make_initial_guess(self, eq, target, arrays): """ @@ -162,9 +190,28 @@ def generate_arrays(self, target): for p in prefixes } + def scale_essential_bcs(self, equations, scale): + """ + Scale the essential boundary rows so that the Jacobian has a constant diagonal, + thereby reducing its condition number. + """ + # # stencils = self.centre_stencils[arrays['x']] + # if len(stencils) > 1: + # # Scaling of jacobian is therefore ambiguous, potentially could averge across the subblock + # # for now just set to trivial 1.0 + # scale = 1.0 + # else: + # scale = next(iter(stencils)) + return [ + eq._rebuild(rhs=scale * eq.rhs) if isinstance(eq, ZeroRow) else eq + for eq in equations + ] + class InjectSolveNested(InjectSolve): + def linear_solve_args(self): + combined_eqns = [] for eqns in self.target_eqns.values(): combined_eqns.extend(eqns) @@ -181,51 +228,90 @@ def linear_solve_args(self): self.cell_area = np.prod(all_data.grid.spacing_symbols) + all_formfuncs = [] + for target, eqns in self.target_eqns.items(): - eqns = as_tuple(eqns) - self.update_jacobian(eqns, target, jacobian, arrays[target]) - formfuncs = [ - self.build_function_eqns(eq, target, coupled_targets, arrays) - for eq in eqns - ] - all_data.extend_formfuncs(formfuncs) + # Update all rows of the Jacobian for this target + self.update_jacobian(as_tuple(eqns), target, jacobian, arrays[target]) + + formfuncs = chain.from_iterable( + self.build_function_eq(eq, target, coupled_targets, arrays) + for eq in as_tuple(eqns) + ) + # from IPython import embed; embed() + scale, = self._diag_scale[arrays[target]['x']] + all_formfuncs.extend(self.scale_essential_bcs(formfuncs, scale)) + + formfuncs = tuple(sorted( + all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + )) + all_data.extend_formfuncs(formfuncs) return target, tuple(funcs), all_data def update_jacobian(self, eqns, target, jacobian, arrays): + for submat, mtvs in jacobian.submatrices[target].items(): - # TODO: maintain symmetry for coupled - matvecs = [e for i in [self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) - for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) + ] matvecs = [m for m in matvecs if m is not None] + if submat in jacobian.diagonal_submatrix_keys: + stencils = set() + for eq in matvecs: + if not isinstance(eq, EssentialBC): + stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) + stencils.add(stencil) + # from IPython import embed; embed() + if len(stencils) > 1: + # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock + # for now just set to trivial 1.0 + # TODO: doens't need to be a defaultdict, just a dict? + self._diag_scale[arrays['x']].add(1.0) + scale = 1.0 + else: + scale = next(iter(stencils)) + self._diag_scale[arrays['x']].add(scale) + # scale = next(iter(stencils)) + + matvecs = self.scale_essential_bcs(matvecs, scale) + + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + if matvecs: jacobian.set_submatrix(target, submat, matvecs) - def build_function_eqns(self, eq, main_target, coupled_targets, arrays): + def build_function_eq(self, eq, main_target, coupled_targets, arrays): zeroed = eq.lhs - eq.rhs - # TODO: clean up, test coupled with time dependence - zeroed_eqn = Eq(zeroed, 0) - zeroed_eqn = eval_time_derivatives(zeroed) + zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) mapper = {} for t in coupled_targets: - target_funcs = generate_targets(Eq(zeroed, 0), t) + target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): rhs = arrays[main_target]['x'] - eq.rhs + zero_row = ZeroRow( + arrays[main_target]['f'], rhs, subdomain=eq.subdomain + ) + zero_col = ZeroColumn( + arrays[main_target]['x'], eq.rhs, subdomain=eq.subdomain + ) + return (zero_row, zero_col) else: if isinstance(zeroed, (int, float)): rhs = zeroed * self.cell_area else: rhs = zeroed.subs(mapper) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area + rhs = rhs.subs(self.time_mapper)*self.cell_area - return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) + return as_tuple(Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain)) def generate_arrays_combined(self, *targets): return { @@ -241,32 +327,37 @@ def generate_arrays_combined(self, *targets): for target in targets } - def make_matvec(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # TODO: SHOULDN'T NEED this subclass once fixed mixed problems + symmetry - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal - # and zero symmetrically. FIX FOR MIXED PROBLEMS - rhs = arrays['x'] - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['y'], rhs, subdomain=eq.subdomain) - class EssentialBC(Eq): """ - A special equation representing an essential boundary condition. - It is used to handle essential boundary conditions in the PETSc solver. - Until PetscSection + DMDA is supported, we treat essential boundary conditions - as trivial equations in the solver, where we place 1.0 on the diagonal of the jacobian, - zero symmetrically and move the essential boundary condition to the right-hand side. + A special equation used to handle essential boundary conditions + in the PETSc solver. Until PetscSection + DMDA is supported, + we treat essential boundary conditions as trivial equations + in the solver, where we place 1.0 (scaled) on the diagonal of + the jacobian, zero symmetrically and move the boundary + data to the right-hand side. NOTE: When users define essential boundary conditions, they need to ensure that the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. + at boundaries. This will be documented in the PETSc examples. + """ + pass + + +class ZeroRow(EssentialBC): + """ + Equation used to zero the row of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass + + +class ZeroColumn(EssentialBC): + """ + Equation used to zero the column of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. """ pass @@ -354,26 +445,34 @@ def _(expr, targets): @singledispatch -def centre_stencil(expr, target): +def centre_stencil(expr, target, as_coeff=False): """ Extract the centre stencil from an expression. Its coefficient is what would appear on the diagonal of the matrix system if the matrix were formed explicitly. + Parameters + ---------- + expr : the expression to extract the centre stencil from + target : the target function whose centre stencil we want + as_coeff : bool, optional + If True, return the coefficient of the centre stencil """ - return expr if expr == target else 0 + if expr == target: + return 1 if as_coeff else expr + return 0 @centre_stencil.register(sympy.Add) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 - args = [centre_stencil(a, target) for a in expr.args] + args = [centre_stencil(a, target, as_coeff) for a in expr.args] return expr.func(*args, evaluate=False) @centre_stencil.register(Mul) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 @@ -382,16 +481,16 @@ def _(expr, target): if not a.has(target): args.append(a) else: - args.append(centre_stencil(a, target)) + args.append(centre_stencil(a, target, as_coeff)) return expr.func(*args, evaluate=False) @centre_stencil.register(Derivative) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 - args = [centre_stencil(a, target) for a in expr.evaluate.args] + args = [centre_stencil(a, target, as_coeff) for a in expr.evaluate.args] return expr.evaluate.func(*args) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 58ca3e28d7..049785662b 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -124,11 +124,16 @@ class PetscBundle(Bundle): is_Bundle = True _data_alignment = False + __rkwargs__ = Bundle.__rkwargs__ + ('pname',) + + def __init__(self, *args, pname="Field", **kwargs): + super().__init__(*args, **kwargs) + self._pname = pname + @property def _C_ctype(self): - # TODO: extend to cases with multiple petsc solves...(need diff struct name for each solve) fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] - return POINTER(type('Field', (Structure,), {'_fields_': fields})) + return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) @cached_property def indexed(self): @@ -164,6 +169,10 @@ def __getitem__(self, index): raise ValueError("Expected %d or %d indices, got %d instead" % (self.ndim, self.ndim + 1, len(index))) + @property + def pname(self): + return self._pname + class AoSIndexedData(IndexedData): @property diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index afab9fff63..d2c45ca010 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -321,6 +321,18 @@ def get_submatrix(self, field, key): Retrieve a specific submatrix. """ return self.submatrices.get(field, {}).get(key, None) + + @property + def diagonal_submatrix_keys(self): + """ + Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + """ + keys = [] + for i, target in enumerate(self.targets): + diag_key = f'J{i}{i}' + if diag_key in self.submatrices[target]: + keys.append(diag_key) + return keys def __repr__(self): return str(self.submatrices) diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0920298356..1285ffabec 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -346,4 +346,3 @@ def key(i): return (v, str(type(i))) return sorted(items, key=key, reverse=True) - diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6f24f74f64..70f17d94c6 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [9] +n_values = [7] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -112,7 +112,7 @@ def analytical(x, y): error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) errors.append(error) -slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) +# slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -assert slope > 1.9 -assert slope < 2.1 +# assert slope > 1.9 +# assert slope < 2.1 diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index bc7b54654d..973b16499e 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -12,9 +12,12 @@ configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' + PetscInitialize() + # Subdomains to implement BCs + class SubTop(SubDomain): name = 'subtop' @@ -50,18 +53,23 @@ def define(self, dimensions): def c(x): return x**3 * (1 - x)**3 + def ddc(x): return 6.0 * x * (1 - x) * (1 - 5.0 * x + 5.0 * x**2) + def d4c(x): return -72.0 * (1 - 5.0 * x + 5.0 * x**2) + def u_exact_fcn(x, y): return c(x) * c(y) + def lap_u_exact_fcn(x, y): return -ddc(x) * c(y) - c(x) * ddc(y) + def f_fcn(x, y): return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y) @@ -77,6 +85,7 @@ def f_fcn(x, y): Ly = np.float64(1.) n_values = [33, 53, 73, 93, 113] +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -115,7 +124,8 @@ def f_fcn(x, y): bc_v += [EssentialBC(v, 0., subdomain=sub4)] # T (see ref) is nonsymmetric so need to set default KSP type to GMRES - petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters={'ksp_rtol': 1e-10}) + params = {'ksp_rtol': 1e-10} + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters=params) with switchconfig(language='petsc'): op = Operator(petsc) @@ -126,12 +136,14 @@ def f_fcn(x, y): # Compute infinity norm for u u_diff = u_exact.data[:] - u.data[:] - u_error = np.linalg.norm(u_diff.ravel(), ord=np.inf) / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) + u_diff_norm = np.linalg.norm(u_diff.ravel(), ord=np.inf) + u_error = u_diff_norm / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) u_errors.append(u_error) # Compute infinity norm for lap_u v_diff = lap_u.data[:] - v.data[:] - v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) + v_diff_norm = np.linalg.norm(v_diff.ravel(), ord=np.inf) + v_error = v_diff_norm / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) v_errors.append(v_error) u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c136160ecc..382b325e4e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -339,7 +339,7 @@ def test_separate_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _, _= separate_eqn(eval(eqn), eval(target)) + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -793,8 +793,7 @@ def test_essential_bcs(): Verify that PETScSolve returns the correct output with essential boundary conditions. """ - - # SubDomains used for essential boundary conditions + # SubDomains used for essential boundary conditions # should not overlap. class SubTop(SubDomain): name = 'subtop' @@ -917,7 +916,6 @@ def test_coupled_vs_non_coupled(self): # Solving for multiple fields within the same matrix system requires # less callback functions than solving them separately. - # TODO: check reuse of callback functions where appropriate assert len(callbacks1) == 8 assert len(callbacks2) == 6 @@ -972,6 +970,10 @@ def test_coupled_structs(self): assert 'struct UserCtx0\n{' not in ccode assert 'struct UserCtx0\n{' in hcode + # The public struct Field0 only appears in the header file + assert 'struct Field0\n{' not in ccode + assert 'struct Field0\n{' in hcode + @skipif('petsc') def test_coupled_frees(self): grid = Grid(shape=(11, 11), dtype=np.float64) @@ -1083,7 +1085,7 @@ def test_submatrices(self): j00 = submatrices.get_submatrix(e, 'J00') j11 = submatrices.get_submatrix(g, 'J11') - # Compatible scaling to reduce condition number of jacobian + # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' @@ -1096,13 +1098,98 @@ def test_submatrices(self): assert j10['derivative_wrt'] == e assert j11['derivative_wrt'] == g - # TODO: - # @skipif('petsc') - # def test_create_submats(self): + @skipif('petsc') + def test_residual_bundle(self): + grid = Grid(shape=(11, 11), dtype=np.float64) - # add tests for all new callbacks - # def test_create_whole_matmult(): + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) -# // add coupled test with 3 targets -# // add coupled test with 1 target? + petsc1 = PETScSolve({e: [eq1]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2]}) + petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(language='petsc'): + op1 = Operator(petsc1, opt='noop', name='op1') + op2 = Operator(petsc2, opt='noop', name='op2') + op3 = Operator(petsc3, opt='noop', name='op3') + + # Check pointers to array of Field structs. Note this is only + # required when dof>1 when constructing the multi-component DMDA. + f_aos = 'struct Field0 (* f_bundle)[info.gxm] = ' \ + + '(struct Field0 (*)[info.gxm]) f_bundle_vec;' + x_aos = 'struct Field0 (* x_bundle)[info.gxm] = ' \ + + '(struct Field0 (*)[info.gxm]) x_bundle_vec;' + + for op in (op1, op2, op3): + ccode = str(op.ccode) + assert f_aos in ccode + assert x_aos in ccode + + assert 'struct Field0\n{\n PetscScalar e;\n}' \ + in str(op1.ccode) + assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n}' \ + in str(op2.ccode) + assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n ' \ + + 'PetscScalar g;\n}' in str(op3.ccode) + + @skipif('petsc') + def test_essential_bcs(self): + """ + Test mixed problem with SubDomains + """ + class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('right', 1)} + + sub1 = SubTop() + + grid = Grid(shape=(9, 9), subdomains=(sub1,), dtype=np.float64) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + eqn1 = Eq(-v.laplace, f, subdomain=grid.interior) + eqn2 = Eq(-u.laplace, v, subdomain=grid.interior) + + bc_u = [EssentialBC(u, 0., subdomain=sub1)] + bc_v = [EssentialBC(v, 0., subdomain=sub1)] + + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + # Test scaling + J00 = op._func_table['J00_MatMult0'].root + + # Essential BC row + assert 'a1[ix + 2][iy + 2] = (2.0/((o0->h_y*o0->h_y))' \ + ' + 2.0/((o0->h_x*o0->h_x)))*o0->h_x*o0->h_y*a0[ix + 2][iy + 2];' in str(J00) + # Check zeroing of essential BC columns + assert 'a0[ix + 2][iy + 2] = 0.0;' in str(J00) + # Interior loop + assert 'a1[ix + 2][iy + 2] = (2.0*(r0*a0[ix + 2][iy + 2] ' \ + '+ r1*a0[ix + 2][iy + 2]) - (r0*a0[ix + 1][iy + 2] + ' \ + 'r0*a0[ix + 3][iy + 2] + r1*a0[ix + 2][iy + 1] ' \ + '+ r1*a0[ix + 2][iy + 3]))*o0->h_x*o0->h_y;' in str(J00) + + # J00 and J11 are semantically identical so check efunc reuse + assert len(op._func_table.values()) == 7 + # J00_MatMult0 is reused (in replace of J11_MatMult0) + create = op._func_table['MatCreateSubMatrices0'].root + assert 'MatShellSetOperation(submat_arr[0],' \ + + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) + assert 'MatShellSetOperation(submat_arr[3],' \ + + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) + + # TODO: Test mixed, time dependent solvers From f7e00bbc1c1691665d14161a2dbddb947ce50ced Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 21:43:43 +0100 Subject: [PATCH 08/53] types: Edit PetscMixin and use it in other petsc classes --- devito/petsc/iet/routines.py | 18 ++++----- devito/petsc/solve.py | 6 +-- devito/petsc/types/object.py | 72 ++++++++++++++++-------------------- devito/petsc/types/types.py | 10 ++--- tests/test_petsc.py | 24 ++++++------ 5 files changed, 61 insertions(+), 69 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 7c824ad638..00c46f8b7d 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -653,8 +653,8 @@ def submatrices_callback(self): return self._submatrices_callback @property - def submatrices(self): - return self.injectsolve.expr.rhs.fielddata.submatrices + def jacobian(self): + return self.injectsolve.expr.rhs.fielddata.jacobian @property def main_matvec_callback(self): @@ -682,7 +682,7 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - row_matvecs = all_fielddata.submatrices.submatrices[t] + row_matvecs = all_fielddata.jacobian.submatrices[t] arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): if mtvs['matvecs']: @@ -714,7 +714,7 @@ def _whole_matvec_body(self): jctx = objs['ljacctx'] ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) - nonzero_submats = self.submatrices.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrix_keys zero_y_memory = petsc_call( 'VecSet', [objs['Y'], 0.0] @@ -722,7 +722,7 @@ def _whole_matvec_body(self): calls = () for sm in nonzero_submats: - idx = self.submatrices.submat_to_index[sm] + idx = self.jacobian.submat_to_index[sm] ctx = sobjs[f'{sm}ctx'] X = sobjs[f'{sm}X'] Y = sobjs[f'{sm}Y'] @@ -1012,14 +1012,14 @@ def _submat_callback_body(self): upper_bound = objs['nsubmats'] - 1 iteration = Iteration(List(body=iter_body), i, upper_bound) - nonzero_submats = self.submatrices.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrix_keys matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} matmult_op = [ petsc_call( 'MatShellSetOperation', [ - objs['submat_arr'].indexed[self.submatrices.submat_to_index[sb]], + objs['submat_arr'].indexed[self.jacobian.submat_to_index[sb]], 'MATOP_MULT', MatShellSetOp(matvec_lookup[sb].name, void, void), ], @@ -1157,8 +1157,8 @@ def _extend_build(self, base_dict): dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - submatrices = injectsolve.expr.rhs.fielddata.submatrices - submatrix_keys = submatrices.submatrix_keys + jacobian = injectsolve.expr.rhs.fielddata.jacobian + submatrix_keys = jacobian.submatrix_keys base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 278a5092b8..218c146ac5 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -13,7 +13,7 @@ from devito.symbolics import retrieve_functions from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices) + FieldData, MultipleFieldData, Jacobian) __all__ = ['PETScSolve', 'EssentialBC'] @@ -219,11 +219,11 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) coupled_targets = list(self.target_eqns.keys()) - jacobian = SubMatrices(coupled_targets) + jacobian = Jacobian(coupled_targets) arrays = self.generate_arrays_combined(*coupled_targets) - all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, + all_data = MultipleFieldData(jacobian=jacobian, arrays=arrays, targets=coupled_targets) self.cell_area = np.prod(all_data.grid.spacing_symbols) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index a285863d14..e62d125dc2 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -10,10 +10,17 @@ class PetscMixin: @property def _C_free_priority(self): - return FREE_PRIORITY[self] + if type(self) in FREE_PRIORITY: + return FREE_PRIORITY[type(self)] + else: + return super()._C_free_priority -class CallbackDM(LocalObject): +class PetscObject(PetscMixin, LocalObject): + pass + + +class CallbackDM(PetscObject): """ PETSc Data Management object (DM). This is the DM instance accessed within the callback functions via `SNESGetDM` and @@ -22,14 +29,12 @@ class CallbackDM(LocalObject): dtype = CustomDtype('DM') -class DM(LocalObject, PetscMixin): +class DM(CallbackDM): """ PETSc Data Management object (DM). This is the primary DM instance created within the main kernel and linked to the SNES solver using `SNESSetDM`. """ - dtype = CustomDtype('DM') - def __init__(self, *args, dofs=1, **kwargs): super().__init__(*args, **kwargs) self._dofs = dofs @@ -46,7 +51,7 @@ def _C_free(self): DMCast = cast('DM') -class CallbackMat(LocalObject): +class CallbackMat(PetscObject): """ PETSc Matrix object (Mat) used within callback functions. These instances are not destroyed during callback execution; @@ -55,19 +60,13 @@ class CallbackMat(LocalObject): dtype = CustomDtype('Mat') -class Mat(LocalObject): - dtype = CustomDtype('Mat') - +class Mat(CallbackMat): @property def _C_free(self): return petsc_call('MatDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 2 - -class CallbackVec(LocalObject): +class CallbackVec(PetscObject): """ PETSc vector object (Vec). """ @@ -79,12 +78,8 @@ class Vec(CallbackVec): def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 1 - -class PetscMPIInt(LocalObject): +class PetscMPIInt(PetscObject): """ PETSc datatype used to represent `int` parameters to MPI functions. @@ -92,7 +87,7 @@ class PetscMPIInt(LocalObject): dtype = CustomDtype('PetscMPIInt') -class PetscInt(LocalObject): +class PetscInt(PetscObject): """ PETSc datatype used to represent `int` parameters to PETSc functions. @@ -100,7 +95,7 @@ class PetscInt(LocalObject): dtype = CustomIntType('PetscInt') -class KSP(LocalObject): +class KSP(PetscObject): """ PETSc KSP : Linear Systems Solvers. Manages Krylov Methods. @@ -108,7 +103,7 @@ class KSP(LocalObject): dtype = CustomDtype('KSP') -class CallbackSNES(LocalObject): +class CallbackSNES(PetscObject): """ PETSc SNES : Non-Linear Systems Solvers. """ @@ -120,19 +115,15 @@ class SNES(CallbackSNES): def _C_free(self): return petsc_call('SNESDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 3 - -class PC(LocalObject): +class PC(PetscObject): """ PETSc object that manages all preconditioners (PC). """ dtype = CustomDtype('PC') -class KSPConvergedReason(LocalObject): +class KSPConvergedReason(PetscObject): """ PETSc object - reason a Krylov method was determined to have converged or diverged. @@ -140,7 +131,7 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class DMDALocalInfo(LocalObject): +class DMDALocalInfo(PetscObject): """ PETSc object - C struct containing information about the local grid. @@ -148,7 +139,7 @@ class DMDALocalInfo(LocalObject): dtype = CustomDtype('DMDALocalInfo') -class PetscErrorCode(LocalObject): +class PetscErrorCode(PetscObject): """ PETSc datatype used to return PETSc error codes. https://petsc.org/release/manualpages/Sys/PetscErrorCode/ @@ -156,7 +147,7 @@ class PetscErrorCode(LocalObject): dtype = CustomDtype('PetscErrorCode') -class DummyArg(LocalObject): +class DummyArg(PetscObject): """ A void pointer used to satisfy the function signature of the `FormFunction` callback. @@ -164,21 +155,21 @@ class DummyArg(LocalObject): dtype = CustomDtype('void', modifier='*') -class MatReuse(LocalObject): +class MatReuse(PetscObject): dtype = CustomDtype('MatReuse') -class VecScatter(LocalObject): +class VecScatter(PetscObject): dtype = CustomDtype('VecScatter') -class StartPtr(LocalObject): +class StartPtr(PetscObject): def __init__(self, name, dtype): super().__init__(name=name) self.dtype = POINTER(dtype_to_ctype(dtype)) -class SingleIS(LocalObject): +class SingleIS(PetscObject): dtype = CustomDtype('IS') @@ -218,7 +209,8 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -class PETScArrayObject(ArrayObject): + +class PETScArrayObject(PetscMixin, ArrayObject): _data_alignment = False def __init_finalize__(self, *args, **kwargs): @@ -251,10 +243,6 @@ def _C_name(self): def _mem_stack(self): return False - @property - def _C_free_priority(self): - return 0 - class CallbackPointerIS(PETScArrayObject): """ @@ -309,5 +297,9 @@ def _C_ctype(self): FREE_PRIORITY = { + PETScArrayObject: 0, + Vec: 1, + Mat: 2, + SNES: 3, DM: 4, } diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index d2c45ca010..75ee8e8254 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -189,10 +189,10 @@ def targets(self): class MultipleFieldData(FieldData): - def __init__(self, targets, arrays, submatrices=None): + def __init__(self, targets, arrays, jacobian=None): self._targets = as_tuple(targets) self._arrays = arrays - self._submatrices = submatrices + self._jacobian = jacobian self._formfuncs = [] def extend_formfuncs(self, formfuncs): @@ -232,8 +232,8 @@ def space_order(self): return space_orders.pop() @property - def submatrices(self): - return self._submatrices + def jacobian(self): + return self._jacobian @property def targets(self): @@ -244,7 +244,7 @@ def arrays(self): return self._arrays -class SubMatrices: +class Jacobian: def __init__(self, targets): self.targets = targets self.submatrices = self._initialize_submatrices() diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 382b325e4e..8bcaeb3fa6 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1065,25 +1065,25 @@ def test_submatrices(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - submatrices = petsc[0].rhs.fielddata.submatrices + jacobian = petsc[0].rhs.fielddata.jacobian - j00 = submatrices.get_submatrix(e, 'J00') - j01 = submatrices.get_submatrix(e, 'J01') - j10 = submatrices.get_submatrix(g, 'J10') - j11 = submatrices.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(e, 'J00') + j01 = jacobian.get_submatrix(e, 'J01') + j10 = jacobian.get_submatrix(g, 'J10') + j11 = jacobian.get_submatrix(g, 'J11') # Check the number of submatrices - assert len(submatrices.submatrix_keys) == 4 - assert str(submatrices.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + assert len(jacobian.submatrix_keys) == 4 + assert str(jacobian.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 - assert submatrices.nonzero_submatrix_keys == ['J00', 'J11'] - assert submatrices.get_submatrix(e, 'J01')['matvecs'] is None - assert submatrices.get_submatrix(g, 'J10')['matvecs'] is None + assert jacobian.nonzero_submatrix_keys == ['J00', 'J11'] + assert jacobian.get_submatrix(e, 'J01')['matvecs'] is None + assert jacobian.get_submatrix(g, 'J10')['matvecs'] is None - j00 = submatrices.get_submatrix(e, 'J00') - j11 = submatrices.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(e, 'J00') + j11 = jacobian.get_submatrix(g, 'J11') # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ From 6e20d4bad6fa26a8167ce1edf8a1186d859bd66c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:32:56 +0100 Subject: [PATCH 09/53] dsl: Add extraction file --- devito/symbolics/extraction.py | 120 +++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 devito/symbolics/extraction.py diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py new file mode 100644 index 0000000000..dfbc145748 --- /dev/null +++ b/devito/symbolics/extraction.py @@ -0,0 +1,120 @@ +from functools import singledispatch + +import sympy + +from devito.finite_differences.differentiable import Mul +from devito.finite_differences.derivative import Derivative +from devito.types.equation import Eq +from devito.symbolics import retrieve_functions + + +__all__ = ['separate_eqn', 'generate_targets', 'centre_stencil'] + + +def separate_eqn(eqn, target): + """ + Separate the equation into two separate expressions, + where F(target) = b. + """ + zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + from devito.operations.solve import eval_time_derivatives + zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_targets(zeroed_eqn, target_funcs) + return -b, F_target, zeroed_eqn, target_funcs + + +def generate_targets(eq, target): + """ + Extract all the functions that share the same time index as the target + but may have different spatial indices. + """ + funcs = retrieve_functions(eq) + if target.is_TimeFunction: + time_idx = target.indices[target.time_dim] + targets = [ + f for f in funcs if f.function is target.function and time_idx + in f.indices + ] + else: + targets = [f for f in funcs if f.function is target.function] + return targets + + +@singledispatch +def remove_targets(expr, targets): + return (0, expr) if expr in targets else (expr, 0) + + +@remove_targets.register(sympy.Add) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Mul) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) + else (a, a) for a in expr.args]) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Derivative) +def _(expr, targets): + return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) + + +@singledispatch +def centre_stencil(expr, target, as_coeff=False): + """ + Extract the centre stencil from an expression. Its coefficient is what + would appear on the diagonal of the matrix system if the matrix were + formed explicitly. + Parameters + ---------- + expr : the expression to extract the centre stencil from + target : the target function whose centre stencil we want + as_coeff : bool, optional + If True, return the coefficient of the centre stencil + """ + if expr == target: + return 1 if as_coeff else expr + return 0 + + +@centre_stencil.register(sympy.Add) +def _(expr, target, as_coeff=False): + if not expr.has(target): + return 0 + + args = [centre_stencil(a, target, as_coeff) for a in expr.args] + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Mul) +def _(expr, target, as_coeff=False): + if not expr.has(target): + return 0 + + args = [] + for a in expr.args: + if not a.has(target): + args.append(a) + else: + args.append(centre_stencil(a, target, as_coeff)) + + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Derivative) +def _(expr, target, as_coeff=False): + if not expr.has(target): + return 0 + args = [centre_stencil(a, target, as_coeff) for a in expr.evaluate.args] + return expr.evaluate.func(*args) From 8a79fcde140004230b2bfe79323f1db7a17fc2b2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:33:35 +0100 Subject: [PATCH 10/53] dsl: Move symbolic extraction functions outside of the PETSc module --- devito/petsc/solve.py | 136 ++++--------------- devito/symbolics/__init__.py | 1 + tests/test_petsc.py | 251 +---------------------------------- tests/test_symbolics.py | 241 ++++++++++++++++++++++++++++++++- 4 files changed, 271 insertions(+), 358 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 218c146ac5..f3e7cd377b 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -11,6 +11,8 @@ from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions +from devito.symbolics.extraction import (separate_eqn, centre_stencil, + generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian) @@ -76,7 +78,6 @@ def generate_field_data(self, eqns, target, arrays): zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) ) - # self._centre_stencils[arrays['x']].update( stencils = set() for eq in matvecs: if not isinstance(eq, EssentialBC): @@ -362,33 +363,33 @@ class ZeroColumn(EssentialBC): pass -def separate_eqn(eqn, target): - """ - Separate the equation into two separate expressions, - where F(target) = b. - """ - zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) - zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - target_funcs = set(generate_targets(zeroed_eqn, target)) - b, F_target = remove_targets(zeroed_eqn, target_funcs) - return -b, F_target, zeroed_eqn, target_funcs - - -def generate_targets(eq, target): - """ - Extract all the functions that share the same time index as the target - but may have different spatial indices. - """ - funcs = retrieve_functions(eq) - if isinstance(target, TimeFunction): - time_idx = target.indices[target.time_dim] - targets = [ - f for f in funcs if f.function is target.function and time_idx - in f.indices - ] - else: - targets = [f for f in funcs if f.function is target.function] - return targets +# def separate_eqn(eqn, target): +# """ +# Separate the equation into two separate expressions, +# where F(target) = b. +# """ +# zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) +# zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) +# target_funcs = set(generate_targets(zeroed_eqn, target)) +# b, F_target = remove_targets(zeroed_eqn, target_funcs) +# return -b, F_target, zeroed_eqn, target_funcs + + +# def generate_targets(eq, target): +# """ +# Extract all the functions that share the same time index as the target +# but may have different spatial indices. +# """ +# funcs = retrieve_functions(eq) +# if isinstance(target, TimeFunction): +# time_idx = target.indices[target.time_dim] +# targets = [ +# f for f in funcs if f.function is target.function and time_idx +# in f.indices +# ] +# else: +# targets = [f for f in funcs if f.function is target.function] +# return targets def targets_to_arrays(array, targets): @@ -415,85 +416,6 @@ def targets_to_arrays(array, targets): return dict(zip(targets, array_targets)) -@singledispatch -def remove_targets(expr, targets): - return (0, expr) if expr in targets else (expr, 0) - - -@remove_targets.register(sympy.Add) -def _(expr, targets): - if not any(expr.has(t) for t in targets): - return (expr, 0) - - args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) - return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) - - -@remove_targets.register(Mul) -def _(expr, targets): - if not any(expr.has(t) for t in targets): - return (expr, 0) - - args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) - else (a, a) for a in expr.args]) - return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) - - -@remove_targets.register(Derivative) -def _(expr, targets): - return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) - - -@singledispatch -def centre_stencil(expr, target, as_coeff=False): - """ - Extract the centre stencil from an expression. Its coefficient is what - would appear on the diagonal of the matrix system if the matrix were - formed explicitly. - Parameters - ---------- - expr : the expression to extract the centre stencil from - target : the target function whose centre stencil we want - as_coeff : bool, optional - If True, return the coefficient of the centre stencil - """ - if expr == target: - return 1 if as_coeff else expr - return 0 - - -@centre_stencil.register(sympy.Add) -def _(expr, target, as_coeff=False): - if not expr.has(target): - return 0 - - args = [centre_stencil(a, target, as_coeff) for a in expr.args] - return expr.func(*args, evaluate=False) - - -@centre_stencil.register(Mul) -def _(expr, target, as_coeff=False): - if not expr.has(target): - return 0 - - args = [] - for a in expr.args: - if not a.has(target): - args.append(a) - else: - args.append(centre_stencil(a, target, as_coeff)) - - return expr.func(*args, evaluate=False) - - -@centre_stencil.register(Derivative) -def _(expr, target, as_coeff=False): - if not expr.has(target): - return 0 - args = [centre_stencil(a, target, as_coeff) for a in expr.evaluate.args] - return expr.evaluate.func(*args) - - def generate_time_mapper(funcs): """ Replace time indices with `Symbols` in equations used within diff --git a/devito/symbolics/__init__.py b/devito/symbolics/__init__.py index 3f1525297a..47935f789f 100644 --- a/devito/symbolics/__init__.py +++ b/devito/symbolics/__init__.py @@ -4,3 +4,4 @@ from devito.symbolics.search import * # noqa from devito.symbolics.inspection import * # noqa from devito.symbolics.manipulation import * # noqa +from devito.symbolics.extraction import * # noqa diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8bcaeb3fa6..8f9cc4d9a8 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -12,8 +12,7 @@ from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, LinearSolveExpr, FieldData, MultipleFieldData) -from devito.petsc.solve import (PETScSolve, separate_eqn, centre_stencil, - EssentialBC) +from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -276,254 +275,6 @@ def test_cinterface_petsc_struct(): assert 'struct UserCtx0\n{' in hcode -@skipif('petsc') -@pytest.mark.parametrize('eqn, target, expected', [ - ('Eq(f1.laplace, g1)', - 'f1', ('g1(x, y)', 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), - ('Eq(g1, f1.laplace)', - 'f1', ('-g1(x, y)', '-Derivative(f1(x, y), (x, 2)) - Derivative(f1(x, y), (y, 2))')), - ('Eq(g1, f1.laplace)', 'g1', - ('Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))', 'g1(x, y)')), - ('Eq(f1 + f1.laplace, g1)', 'f1', ('g1(x, y)', - 'f1(x, y) + Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), - ('Eq(g1.dx + f1.dx, g1)', 'f1', - ('g1(x, y) - Derivative(g1(x, y), x)', 'Derivative(f1(x, y), x)')), - ('Eq(g1.dx + f1.dx, g1)', 'g1', - ('-Derivative(f1(x, y), x)', '-g1(x, y) + Derivative(g1(x, y), x)')), - ('Eq(f1 * g1.dx, g1)', 'g1', ('0', 'f1(x, y)*Derivative(g1(x, y), x) - g1(x, y)')), - ('Eq(f1 * g1.dx, g1)', 'f1', ('g1(x, y)', 'f1(x, y)*Derivative(g1(x, y), x)')), - ('Eq((f1 * g1.dx).dy, f1)', 'f1', - ('0', '-f1(x, y) + Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), - ('Eq((f1 * g1.dx).dy, f1)', 'g1', - ('f1(x, y)', 'Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), - ('Eq(f2.laplace, g2)', 'g2', - ('-Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', - '-g2(t, x, y)')), - ('Eq(f2.laplace, g2)', 'f2', ('g2(t, x, y)', - 'Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), - ('Eq(f2.laplace, f2)', 'f2', ('0', - '-f2(t, x, y) + Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), - ('Eq(f2*g2, f2)', 'f2', ('0', 'f2(t, x, y)*g2(t, x, y) - f2(t, x, y)')), - ('Eq(f2*g2, f2)', 'g2', ('f2(t, x, y)', 'f2(t, x, y)*g2(t, x, y)')), - ('Eq(g2*f2.laplace, f2)', 'g2', ('f2(t, x, y)', - '(Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2)))*g2(t, x, y)')), - ('Eq(f2.forward, f2)', 'f2.forward', ('f2(t, x, y)', 'f2(t + dt, x, y)')), - ('Eq(f2.forward, f2)', 'f2', ('-f2(t + dt, x, y)', '-f2(t, x, y)')), - ('Eq(f2.forward.laplace, f2)', 'f2.forward', ('f2(t, x, y)', - 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), - ('Eq(f2.forward.laplace, f2)', 'f2', - ('-Derivative(f2(t + dt, x, y), (x, 2)) - Derivative(f2(t + dt, x, y), (y, 2))', - '-f2(t, x, y)')), - ('Eq(f2.laplace + f2.forward.laplace, g2)', 'f2.forward', - ('g2(t, x, y) - Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', - 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), - ('Eq(g2.laplace, f2 + g2.forward)', 'g2.forward', - ('f2(t, x, y) - Derivative(g2(t, x, y), (x, 2)) - Derivative(g2(t, x, y), (y, 2))', - '-g2(t + dt, x, y)')) -]) -def test_separate_eqn(eqn, target, expected): - """ - Test the separate_eqn function. - - This function is called within PETScSolve to decompose the equation - into the form F(x) = b. This is necessary to utilise the SNES - interface in PETSc. - """ - grid = Grid((2, 2)) - - so = 2 - - f1 = Function(name='f1', grid=grid, space_order=so) # noqa - g1 = Function(name='g1', grid=grid, space_order=so) # noqa - - f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa - g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - - b, F, _, _ = separate_eqn(eval(eqn), eval(target)) - expected_b, expected_F = expected - - assert str(b) == expected_b - assert str(F) == expected_F - - -@skipif('petsc') -@pytest.mark.parametrize('eqn, target, expected', [ - ('Eq(f1.laplace, g1).evaluate', 'f1', - ( - 'g1(x, y)', - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' - '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' - )), - ('Eq(g1, f1.laplace).evaluate', 'f1', - ( - '-g1(x, y)', - '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' - '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' - )), - ('Eq(g1, f1.laplace).evaluate', 'g1', - ( - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' - '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', - 'g1(x, y)' - )), - ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', - ( - 'g1(x, y)', - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' - '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' - )), - ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', - ( - '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', - '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' - )), - ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', - ( - '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', - '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' - )), - ('Eq(f1 * g1.dx, g1).evaluate', 'g1', - ( - '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' - )), - ('Eq(f1 * g1.dx, g1).evaluate', 'f1', - ( - 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' - )), - ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', - ( - '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' - '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' - '- f1(x, y)' - )), - ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', - ( - 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' - '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' - )), - ('Eq(f2.laplace, g2).evaluate', 'g2', - ( - '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' - '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' - 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' - )), - ('Eq(f2.laplace, g2).evaluate', 'f2', - ( - 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' - '/h_y**2 + f2(t, x, y + h_y)/h_y**2' - )), - ('Eq(f2.laplace, f2).evaluate', 'f2', - ( - '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' - ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' - )), - ('Eq(g2*f2.laplace, f2).evaluate', 'g2', - ( - 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' - ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' - )), - ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', - ( - 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' - ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' - 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' - )), - ('Eq(f2.forward.laplace, f2).evaluate', 'f2', - ( - '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' - 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' - 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', - '-f2(t, x, y)' - )), - ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', - ( - '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' - 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' - 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' - 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' - 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' - 'f2(t + dt, x, y + h_y)/h_y**2' - )), - ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', - ( - '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' - 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' - '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' - )) -]) -def test_separate_eval_eqn(eqn, target, expected): - """ - Test the separate_eqn function on pre-evaluated equations. - This ensures that evaluated equations can be passed to PETScSolve, - allowing users to modify stencils for specific boundary conditions, - such as implementing free surface boundary conditions. - """ - grid = Grid((2, 2)) - - so = 2 - - f1 = Function(name='f1', grid=grid, space_order=so) # noqa - g1 = Function(name='g1', grid=grid, space_order=so) # noqa - - f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa - g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - - b, F, _, _ = separate_eqn(eval(eqn), eval(target)) - expected_b, expected_F = expected - - assert str(b) == expected_b - assert str(F) == expected_F - - -@skipif('petsc') -@pytest.mark.parametrize('expr, so, target, expected', [ - ('f1.laplace', 2, 'f1', '-2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), - ('f1 + f1.laplace', 2, 'f1', - 'f1(x, y) - 2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), - ('g1.dx + f1.dx', 2, 'f1', '-f1(x, y)/h_x'), - ('10 + f1.dx2', 2, 'g1', '0'), - ('(f1 * g1.dx).dy', 2, 'f1', - '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)'), - ('(f1 * g1.dx).dy', 2, 'g1', '-(-1/h_y)*f1(x, y)*g1(x, y)/h_x'), - ('f2.laplace', 2, 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), - ('f2*g2', 2, 'f2', 'f2(t, x, y)*g2(t, x, y)'), - ('g2*f2.laplace', 2, 'f2', - '(-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2)*g2(t, x, y)'), - ('f2.forward', 2, 'f2.forward', 'f2(t + dt, x, y)'), - ('f2.forward.laplace', 2, 'f2.forward', - '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), - ('f2.laplace + f2.forward.laplace', 2, 'f2.forward', - '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), - ('f2.laplace + f2.forward.laplace', 2, - 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), - ('f2.laplace', 4, 'f2', '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), - ('f2.laplace + f2.forward.laplace', 4, 'f2.forward', - '-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2'), - ('f2.laplace + f2.forward.laplace', 4, 'f2', - '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), - ('f2.forward*f2.forward.laplace', 4, 'f2.forward', - '(-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2)*f2(t + dt, x, y)') -]) -def test_centre_stencil(expr, so, target, expected): - """ - Test extraction of centre stencil from an equation. - """ - grid = Grid((2, 2)) - - f1 = Function(name='f1', grid=grid, space_order=so) # noqa - g1 = Function(name='g1', grid=grid, space_order=so) # noqa - - f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa - g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - - centre = centre_stencil(eval(expr), eval(target)) - - assert str(centre) == expected - - @skipif('petsc') def test_callback_arguments(): """ diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index deedbec950..1d55ce5a59 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -14,7 +14,8 @@ CallFromPointer, Cast, DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, pow_to_mul, - retrieve_derivatives, BaseCast, SizeOf, sympy_dtype) + retrieve_derivatives, BaseCast, SizeOf, sympy_dtype, + separate_eqn, centre_stencil) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, ComponentAccess, StencilDimension, Symbol as dSymbol, @@ -927,3 +928,241 @@ def test_print_div(): b = SizeOf(np.int64) cstr = ccode(a / b) assert cstr == 'sizeof(int)/sizeof(long)' + + +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1)', + 'f1', ('g1(x, y)', 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), + ('Eq(g1, f1.laplace)', + 'f1', ('-g1(x, y)', '-Derivative(f1(x, y), (x, 2)) - Derivative(f1(x, y), (y, 2))')), + ('Eq(g1, f1.laplace)', 'g1', + ('Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))', 'g1(x, y)')), + ('Eq(f1 + f1.laplace, g1)', 'f1', ('g1(x, y)', + 'f1(x, y) + Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), + ('Eq(g1.dx + f1.dx, g1)', 'f1', + ('g1(x, y) - Derivative(g1(x, y), x)', 'Derivative(f1(x, y), x)')), + ('Eq(g1.dx + f1.dx, g1)', 'g1', + ('-Derivative(f1(x, y), x)', '-g1(x, y) + Derivative(g1(x, y), x)')), + ('Eq(f1 * g1.dx, g1)', 'g1', ('0', 'f1(x, y)*Derivative(g1(x, y), x) - g1(x, y)')), + ('Eq(f1 * g1.dx, g1)', 'f1', ('g1(x, y)', 'f1(x, y)*Derivative(g1(x, y), x)')), + ('Eq((f1 * g1.dx).dy, f1)', 'f1', + ('0', '-f1(x, y) + Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), + ('Eq((f1 * g1.dx).dy, f1)', 'g1', + ('f1(x, y)', 'Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')), + ('Eq(f2.laplace, g2)', 'g2', + ('-Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', + '-g2(t, x, y)')), + ('Eq(f2.laplace, g2)', 'f2', ('g2(t, x, y)', + 'Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), + ('Eq(f2.laplace, f2)', 'f2', ('0', + '-f2(t, x, y) + Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')), + ('Eq(f2*g2, f2)', 'f2', ('0', 'f2(t, x, y)*g2(t, x, y) - f2(t, x, y)')), + ('Eq(f2*g2, f2)', 'g2', ('f2(t, x, y)', 'f2(t, x, y)*g2(t, x, y)')), + ('Eq(g2*f2.laplace, f2)', 'g2', ('f2(t, x, y)', + '(Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2)))*g2(t, x, y)')), + ('Eq(f2.forward, f2)', 'f2.forward', ('f2(t, x, y)', 'f2(t + dt, x, y)')), + ('Eq(f2.forward, f2)', 'f2', ('-f2(t + dt, x, y)', '-f2(t, x, y)')), + ('Eq(f2.forward.laplace, f2)', 'f2.forward', ('f2(t, x, y)', + 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), + ('Eq(f2.forward.laplace, f2)', 'f2', + ('-Derivative(f2(t + dt, x, y), (x, 2)) - Derivative(f2(t + dt, x, y), (y, 2))', + '-f2(t, x, y)')), + ('Eq(f2.laplace + f2.forward.laplace, g2)', 'f2.forward', + ('g2(t, x, y) - Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))', + 'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')), + ('Eq(g2.laplace, f2 + g2.forward)', 'g2.forward', + ('f2(t, x, y) - Derivative(g2(t, x, y), (x, 2)) - Derivative(g2(t, x, y), (y, 2))', + '-g2(t + dt, x, y)')) +]) +def test_separate_eqn(eqn, target, expected): + """ + Test the separate_eqn function. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' + )), + ('Eq(g1, f1.laplace).evaluate', 'f1', + ( + '-g1(x, y)', + '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' + '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' + )), + ('Eq(g1, f1.laplace).evaluate', 'g1', + ( + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', + 'g1(x, y)' + )), + ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' + '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', + ( + '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', + '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', + ( + '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', + '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'g1', + ( + '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'f1', + ( + 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', + ( + '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' + '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' + '- f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', + ( + 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' + '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' + )), + ('Eq(f2.laplace, g2).evaluate', 'g2', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' + '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' + )), + ('Eq(f2.laplace, g2).evaluate', 'f2', + ( + 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' + '/h_y**2 + f2(t, x, y + h_y)/h_y**2' + )), + ('Eq(f2.laplace, f2).evaluate', 'f2', + ( + '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' + )), + ('Eq(g2*f2.laplace, f2).evaluate', 'g2', + ( + 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', + ( + 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' + ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2', + ( + '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' + 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', + '-f2(t, x, y)' + )), + ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' + 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' + 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' + 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' + 'f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', + ( + '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' + 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' + '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' + )) +]) +def test_separate_eval_eqn(eqn, target, expected): + """ + Test the separate_eqn function on pre-evaluated equations. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@pytest.mark.parametrize('expr, so, target, expected', [ + ('f1.laplace', 2, 'f1', '-2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), + ('f1 + f1.laplace', 2, 'f1', + 'f1(x, y) - 2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'), + ('g1.dx + f1.dx', 2, 'f1', '-f1(x, y)/h_x'), + ('10 + f1.dx2', 2, 'g1', '0'), + ('(f1 * g1.dx).dy', 2, 'f1', + '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)'), + ('(f1 * g1.dx).dy', 2, 'g1', '-(-1/h_y)*f1(x, y)*g1(x, y)/h_x'), + ('f2.laplace', 2, 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), + ('f2*g2', 2, 'f2', 'f2(t, x, y)*g2(t, x, y)'), + ('g2*f2.laplace', 2, 'f2', + '(-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2)*g2(t, x, y)'), + ('f2.forward', 2, 'f2.forward', 'f2(t + dt, x, y)'), + ('f2.forward.laplace', 2, 'f2.forward', + '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 2, 'f2.forward', + '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 2, + 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'), + ('f2.laplace', 4, 'f2', '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 4, 'f2.forward', + '-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2'), + ('f2.laplace + f2.forward.laplace', 4, 'f2', + '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'), + ('f2.forward*f2.forward.laplace', 4, 'f2.forward', + '(-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2)*f2(t + dt, x, y)') +]) +def test_centre_stencil(expr, so, target, expected): + """ + Test extraction of centre stencil from an equation. + """ + grid = Grid((2, 2)) + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + centre = centre_stencil(eval(expr), eval(target)) + + assert str(centre) == expected From 3000bd8e7d7541c684897324e9f0c4d2287d46e0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:47:00 +0100 Subject: [PATCH 11/53] types: Move EssentialBC to another file --- devito/petsc/__init__.py | 1 + devito/petsc/solve.py | 80 +++++++++++----------------------- devito/petsc/types/__init__.py | 2 + 3 files changed, 29 insertions(+), 54 deletions(-) diff --git a/devito/petsc/__init__.py b/devito/petsc/__init__.py index 2927bc960e..6aab617e8d 100644 --- a/devito/petsc/__init__.py +++ b/devito/petsc/__init__.py @@ -1 +1,2 @@ from devito.petsc.solve import * # noqa +from devito.petsc.types.equation import * # noqa \ No newline at end of file diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index f3e7cd377b..1004a403b6 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -16,9 +16,10 @@ from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian) +from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn -__all__ = ['PETScSolve', 'EssentialBC'] +__all__ = ['PETScSolve'] def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): @@ -329,67 +330,38 @@ def generate_arrays_combined(self, *targets): } -class EssentialBC(Eq): - """ - A special equation used to handle essential boundary conditions - in the PETSc solver. Until PetscSection + DMDA is supported, - we treat essential boundary conditions as trivial equations - in the solver, where we place 1.0 (scaled) on the diagonal of - the jacobian, zero symmetrically and move the boundary - data to the right-hand side. - - NOTE: When users define essential boundary conditions, they need to ensure that - the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. This will be documented in the PETSc examples. - """ - pass - - -class ZeroRow(EssentialBC): - """ - Equation used to zero the row of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. - """ - pass - - -class ZeroColumn(EssentialBC): - """ - Equation used to zero the column of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. - """ - pass +# class EssentialBC(Eq): +# """ +# A special equation used to handle essential boundary conditions +# in the PETSc solver. Until PetscSection + DMDA is supported, +# we treat essential boundary conditions as trivial equations +# in the solver, where we place 1.0 (scaled) on the diagonal of +# the jacobian, zero symmetrically and move the boundary +# data to the right-hand side. + +# NOTE: When users define essential boundary conditions, they need to ensure that +# the SubDomains do not overlap. Solver will still run but may see unexpected behaviour +# at boundaries. This will be documented in the PETSc examples. +# """ +# pass -# def separate_eqn(eqn, target): +# class ZeroRow(EssentialBC): # """ -# Separate the equation into two separate expressions, -# where F(target) = b. +# Equation used to zero the row of the Jacobian corresponding +# to an essential BC. +# This is only used interally by the compiler, not by users. # """ -# zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) -# zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) -# target_funcs = set(generate_targets(zeroed_eqn, target)) -# b, F_target = remove_targets(zeroed_eqn, target_funcs) -# return -b, F_target, zeroed_eqn, target_funcs +# pass -# def generate_targets(eq, target): +# class ZeroColumn(EssentialBC): # """ -# Extract all the functions that share the same time index as the target -# but may have different spatial indices. +# Equation used to zero the column of the Jacobian corresponding +# to an essential BC. +# This is only used interally by the compiler, not by users. # """ -# funcs = retrieve_functions(eq) -# if isinstance(target, TimeFunction): -# time_idx = target.indices[target.time_dim] -# targets = [ -# f for f in funcs if f.function is target.function and time_idx -# in f.indices -# ] -# else: -# targets = [f for f in funcs if f.function is target.function] -# return targets +# pass def targets_to_arrays(array, targets): diff --git a/devito/petsc/types/__init__.py b/devito/petsc/types/__init__.py index ebcceb8d45..f2305a8352 100644 --- a/devito/petsc/types/__init__.py +++ b/devito/petsc/types/__init__.py @@ -1,3 +1,5 @@ from .array import * # noqa from .types import * # noqa from .object import * # noqa +from .equation import * # noqa +from .macros import * # noqa From a0457f057645cd8b034c01d5accb9b233316b2ea Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:48:24 +0100 Subject: [PATCH 12/53] dsl/compiler: Add equation.py to petsc module --- devito/petsc/types/equation.py | 38 ++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 devito/petsc/types/equation.py diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py new file mode 100644 index 0000000000..a5c5bc55ef --- /dev/null +++ b/devito/petsc/types/equation.py @@ -0,0 +1,38 @@ +from devito.types.equation import Eq + + +__all__ = ['EssentialBC'] + + +class EssentialBC(Eq): + """ + A special equation used to handle essential boundary conditions + in the PETSc solver. Until PetscSection + DMDA is supported, + we treat essential boundary conditions as trivial equations + in the solver, where we place 1.0 (scaled) on the diagonal of + the jacobian, zero symmetrically and move the boundary + data to the right-hand side. + + NOTE: When users define essential boundary conditions, they need to ensure that + the SubDomains do not overlap. Solver will still run but may see unexpected behaviour + at boundaries. This will be documented in the PETSc examples. + """ + pass + + +class ZeroRow(EssentialBC): + """ + Equation used to zero the row of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass + + +class ZeroColumn(EssentialBC): + """ + Equation used to zero the column of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass \ No newline at end of file From 7132eb113dd479ae95d2b67324a95606511735fe Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 4 Jun 2025 10:55:51 +0100 Subject: [PATCH 13/53] dsl: Add jacobian class for single fields --- devito/petsc/iet/routines.py | 9 +- devito/petsc/solve.py | 260 ++++---------------------------- devito/petsc/types/types.py | 278 ++++++++++++++++++++++++++++++----- devito/types/grid.py | 7 +- 4 files changed, 280 insertions(+), 274 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 00c46f8b7d..1fc4cd7758 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -108,7 +108,8 @@ def zero_memory(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata.arrays, fielddata.matvecs) + # from IPython.core.debugger import set_trace + self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) if fielddata.initialguess: @@ -266,7 +267,7 @@ def _create_matvec_body(self, body, arrays): return matvec_body def _make_formfunc(self, fielddata): - formfuncs = fielddata.formfuncs + formfuncs = fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, @@ -403,7 +404,7 @@ def _create_formfunc_body(self, body, fielddata): return Uxreplace(subs).visit(formfunc_body) def _make_formrhs(self, fielddata): - formrhs = fielddata.formrhs + formrhs = fielddata.residual.formrhs sobjs = self.solver_objs # Compile formrhs `eqns` into an IET via recursive compilation @@ -758,7 +759,7 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self, fielddata): - formfuncs = fielddata.formfuncs + formfuncs = fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1004a403b6..eb8f171b21 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -15,7 +15,8 @@ generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, Jacobian) + FieldData, MultipleFieldData, Jacobian, Residual, + MixedResidual, MixedJacobian) from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn @@ -26,7 +27,7 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() else: - return InjectSolveNested(solver_parameters, target_eqns).build_eq() + return InjectMixedSolve(solver_parameters, target_eqns).build_eq() class InjectSolve: @@ -34,14 +35,6 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns - # TODO: make this _ - self.cell_area = None - # self._centre_stencils = set() - self._diag_scale = defaultdict(set) - - @property - def diag_scale(self): - return self._diag_scale def build_eq(self): target, funcs, fielddata = self.linear_solve_args() @@ -59,42 +52,20 @@ def build_eq(self): def linear_solve_args(self): target, eqns = next(iter(self.target_eqns.items())) eqns = as_tuple(eqns) - self.cell_area = np.prod(target.grid.spacing_symbols) funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) arrays = self.generate_arrays(target) - return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) - - def generate_field_data(self, eqns, target, arrays): - # Apply essential boundary conditions first to preserve - # operator symmetry during Jacobian "construction" eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - matvecs = [e for eq in eqns for e in self.build_matvec_eq(eq, target, arrays)] - - formfuncs, formrhs = map( - lambda x: [e for i in x for e in i], - zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) - ) - - stencils = set() - for eq in matvecs: - if not isinstance(eq, EssentialBC): - stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) - stencils.add(stencil) + jacobian = Jacobian(target, self.time_mapper, arrays) + jacobian.build_block(eqns) - if len(stencils) > 1: - # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock - # for now just set to trivial 1.0 - scale = 1.0 - else: - scale = next(iter(stencils)) + scale = 1.0 - # from IPython import embed; embed() - matvecs = self.scale_essential_bcs(matvecs, scale) - formfuncs = self.scale_essential_bcs(formfuncs, scale) + residual = Residual(target, self.time_mapper, arrays, scale) + residual.build_equations(eqns) initialguess = [ eq for eq in @@ -102,70 +73,15 @@ def generate_field_data(self, eqns, target, arrays): if eq is not None ] - return FieldData( + field_data = FieldData( target=target, - matvecs=matvecs, - formfuncs=formfuncs, - formrhs=formrhs, + jacobian=jacobian, + residual=residual, initialguess=initialguess, arrays=arrays ) - def build_function_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - formfunc = self.make_formfunc(eq, F_target, arrays, targets) - formrhs = self.make_rhs(eq, b, arrays) - - return (formfunc, formrhs) - - def build_matvec_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - if F_target: - return self.make_matvec(eq, F_target, arrays, targets) - return (None,) - - def make_matvec(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) - # and zero symmetrically. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) * self.cell_area - # TODO: Average centre stencils if they vary, to scale essential BC rows. - # self.centre = centre_stencil(rhs, arrays['x'], as_coeff=True) - # stencil = centre_stencil(rhs, arrays['x'], as_coeff=True) - # self._centre_stencils[arrays['x']].add(stencil) - # self._centre_stencils.add(stencil) - - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) - - def make_formfunc(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # The initial guess satisfies the essential BCs, so this term is zero. - # Still included to support Jacobian testing via finite differences. - rhs = arrays['x'] - eq.rhs - zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) - # Move essential boundary condition to the right-hand side - zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (zero_row, zero_col) - else: - if isinstance(F_target, (int, float)): - rhs = F_target * self.cell_area - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) * self.cell_area - return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) - - def make_rhs(self, eq, b, arrays): - rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) - rhs = rhs * self.cell_area - return as_tuple(Eq(arrays['b'], rhs, subdomain=eq.subdomain)) + return target, tuple(funcs), field_data def make_initial_guess(self, eq, target, arrays): """ @@ -192,25 +108,8 @@ def generate_arrays(self, target): for p in prefixes } - def scale_essential_bcs(self, equations, scale): - """ - Scale the essential boundary rows so that the Jacobian has a constant diagonal, - thereby reducing its condition number. - """ - # # stencils = self.centre_stencils[arrays['x']] - # if len(stencils) > 1: - # # Scaling of jacobian is therefore ambiguous, potentially could averge across the subblock - # # for now just set to trivial 1.0 - # scale = 1.0 - # else: - # scale = next(iter(stencils)) - return [ - eq._rebuild(rhs=scale * eq.rhs) if isinstance(eq, ZeroRow) else eq - for eq in equations - ] - -class InjectSolveNested(InjectSolve): +class InjectMixedSolve(InjectSolve): def linear_solve_args(self): @@ -221,99 +120,24 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) coupled_targets = list(self.target_eqns.keys()) - jacobian = Jacobian(coupled_targets) arrays = self.generate_arrays_combined(*coupled_targets) - all_data = MultipleFieldData(jacobian=jacobian, arrays=arrays, - targets=coupled_targets) - - self.cell_area = np.prod(all_data.grid.spacing_symbols) + jacobian = MixedJacobian(coupled_targets, self.time_mapper, arrays) + jacobian.build_blocks(self.target_eqns) - all_formfuncs = [] - - for target, eqns in self.target_eqns.items(): - - # Update all rows of the Jacobian for this target - self.update_jacobian(as_tuple(eqns), target, jacobian, arrays[target]) - - formfuncs = chain.from_iterable( - self.build_function_eq(eq, target, coupled_targets, arrays) - for eq in as_tuple(eqns) - ) - # from IPython import embed; embed() - scale, = self._diag_scale[arrays[target]['x']] - all_formfuncs.extend(self.scale_essential_bcs(formfuncs, scale)) - - formfuncs = tuple(sorted( - all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) - )) - all_data.extend_formfuncs(formfuncs) - - return target, tuple(funcs), all_data - - def update_jacobian(self, eqns, target, jacobian, arrays): - - for submat, mtvs in jacobian.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) - ] - matvecs = [m for m in matvecs if m is not None] - - if submat in jacobian.diagonal_submatrix_keys: - stencils = set() - for eq in matvecs: - if not isinstance(eq, EssentialBC): - stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) - stencils.add(stencil) - # from IPython import embed; embed() - if len(stencils) > 1: - # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock - # for now just set to trivial 1.0 - # TODO: doens't need to be a defaultdict, just a dict? - self._diag_scale[arrays['x']].add(1.0) - scale = 1.0 - else: - scale = next(iter(stencils)) - self._diag_scale[arrays['x']].add(scale) - # scale = next(iter(stencils)) - - matvecs = self.scale_essential_bcs(matvecs, scale) - - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - jacobian.set_submatrix(target, submat, matvecs) - - def build_function_eq(self, eq, main_target, coupled_targets, arrays): - zeroed = eq.lhs - eq.rhs - - zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) - eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - - mapper = {} - for t in coupled_targets: - target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) - mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) + residual = MixedResidual(coupled_targets, self.time_mapper, arrays) + residual.build_equations(self.target_eqns) - if isinstance(eq, EssentialBC): - rhs = arrays[main_target]['x'] - eq.rhs - zero_row = ZeroRow( - arrays[main_target]['f'], rhs, subdomain=eq.subdomain - ) - zero_col = ZeroColumn( - arrays[main_target]['x'], eq.rhs, subdomain=eq.subdomain - ) - return (zero_row, zero_col) - else: - if isinstance(zeroed, (int, float)): - rhs = zeroed * self.cell_area - else: - rhs = zeroed.subs(mapper) - rhs = rhs.subs(self.time_mapper)*self.cell_area + all_data = MultipleFieldData( + targets=coupled_targets, + arrays=arrays, + jacobian=jacobian, + residual=residual + ) - return as_tuple(Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain)) + # TODO: rethink what target to return here??? + return coupled_targets[0], tuple(funcs), all_data def generate_arrays_combined(self, *targets): return { @@ -330,40 +154,6 @@ def generate_arrays_combined(self, *targets): } -# class EssentialBC(Eq): -# """ -# A special equation used to handle essential boundary conditions -# in the PETSc solver. Until PetscSection + DMDA is supported, -# we treat essential boundary conditions as trivial equations -# in the solver, where we place 1.0 (scaled) on the diagonal of -# the jacobian, zero symmetrically and move the boundary -# data to the right-hand side. - -# NOTE: When users define essential boundary conditions, they need to ensure that -# the SubDomains do not overlap. Solver will still run but may see unexpected behaviour -# at boundaries. This will be documented in the PETSc examples. -# """ -# pass - - -# class ZeroRow(EssentialBC): -# """ -# Equation used to zero the row of the Jacobian corresponding -# to an essential BC. -# This is only used interally by the compiler, not by users. -# """ -# pass - - -# class ZeroColumn(EssentialBC): -# """ -# Equation used to zero the column of the Jacobian corresponding -# to an essential BC. -# This is only used interally by the compiler, not by users. -# """ -# pass - - def targets_to_arrays(array, targets): """ Map each target in `targets` to a corresponding array generated from `array`, diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 75ee8e8254..3897552f7b 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,8 +1,14 @@ import sympy +from itertools import chain + from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables +from devito.symbolics.extraction import separate_eqn, generate_targets +from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn +from devito.types.equation import Eq +from devito.operations.solve import eval_time_derivatives class MetaData(sympy.Function, Reconstructable): @@ -131,7 +137,7 @@ def eval(cls, *args): class FieldData: - def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, + def __init__(self, target=None, jacobian=None, residual=None, initialguess=None, arrays=None, **kwargs): self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] @@ -141,9 +147,8 @@ def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, f"PETSc configuration. " f"Expected {petsc_precision}, but got {self._target.dtype}." ) - self._matvecs = matvecs - self._formfuncs = formfuncs - self._formrhs = formrhs + self._jacobian = jacobian + self._residual = residual self._initialguess = initialguess self._arrays = arrays @@ -152,16 +157,12 @@ def target(self): return self._target @property - def matvecs(self): - return self._matvecs - - @property - def formfuncs(self): - return self._formfuncs + def jacobian(self): + return self._jacobian @property - def formrhs(self): - return self._formrhs + def residual(self): + return self._residual @property def initialguess(self): @@ -189,14 +190,11 @@ def targets(self): class MultipleFieldData(FieldData): - def __init__(self, targets, arrays, jacobian=None): + def __init__(self, targets, arrays, jacobian=None, residual=None): self._targets = as_tuple(targets) self._arrays = arrays self._jacobian = jacobian - self._formfuncs = [] - - def extend_formfuncs(self, formfuncs): - self._formfuncs.extend(formfuncs) + self._residual = residual @property def space_dimensions(self): @@ -239,14 +237,13 @@ def jacobian(self): def targets(self): return self._targets - @property - def arrays(self): - return self._arrays -class Jacobian: - def __init__(self, targets): - self.targets = targets +class BaseJacobian: + def __init__(self, targets, time_mapper, arrays): + self.targets = as_tuple(targets) + self.time_mapper = time_mapper + self.arrays = arrays self.submatrices = self._initialize_submatrices() def _initialize_submatrices(self): @@ -298,6 +295,18 @@ def submat_to_index(self): for key, value in submats.items() } + @property + def diagonal_submatrix_keys(self): + """ + Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + """ + keys = [] + for i, target in enumerate(self.targets): + diag_key = f'J{i}{i}' + if diag_key in self.submatrices[target]: + keys.append(diag_key) + return keys + def set_submatrix(self, field, key, matvecs): """ Set a specific submatrix for a field. @@ -321,18 +330,219 @@ def get_submatrix(self, field, key): Retrieve a specific submatrix. """ return self.submatrices.get(field, {}).get(key, None) - + + def build_matvec_eq(self, eq, target, arrays): + b, F_target, _, targets = separate_eqn(eq, target) + if F_target: + return self.make_matvec(eq, F_target, targets, arrays) + return (None,) + + def make_matvec(self, eq, F_target, targets, arrays): + if isinstance(eq, EssentialBC): + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) + # and zero symmetrically. + rhs = arrays['x'] + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + # rhs = rhs.subs(self.time_mapper) * self.cell_area + rhs = rhs = rhs.subs(self.time_mapper) + + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + + def __repr__(self): + return str(self.submatrices) + + +class Jacobian(BaseJacobian): + @property - def diagonal_submatrix_keys(self): + def target(self): + return self.targets[0] + + @property + def matvecs(self): + return self.submatrices[self.target]['J00']['matvecs'] + + # TODO: use same structure arrays for both jacobian and mixedjacobian + def build_block(self, eqns): + for submat, mtvs in self.submatrices[self.target].items(): + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + if matvecs: + self.set_submatrix(self.target, submat, matvecs) + + +class MixedJacobian(BaseJacobian): + + # TODO: use same structure arrays for both jacobian and mixedjacobian + def build_block(self, target, eqns): + for submat, mtvs in self.submatrices[target].items(): + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays[target]) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + if matvecs: + self.set_submatrix(target, submat, matvecs) + + def build_blocks(self, target_eqns): + for target, eqns in target_eqns.items(): + self.build_block(target, eqns) + + +class BaseResidual: + def scale_essential_bcs(self, equations): """ - Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). """ - keys = [] - for i, target in enumerate(self.targets): - diag_key = f'J{i}{i}' - if diag_key in self.submatrices[target]: - keys.append(diag_key) - return keys + return [ + eq._rebuild(rhs=self.scale * eq.rhs) if isinstance(eq, ZeroRow) else eq + for eq in equations + ] - def __repr__(self): - return str(self.submatrices) + +class Residual(BaseResidual): + """ + """ + + def __init__(self, target, time_mapper, arrays, scale): + self.target = target + self.time_mapper = time_mapper + self.arrays = arrays + self.scale = scale + self.formfuncs = [] + self.formrhs = [] + + def build_equations(self, eqns): + """ + """ + for eq in eqns: + b, F_target, _, targets = separate_eqn(eq, self.target) + F_target = self.make_F_target(eq, F_target, targets) + b = self.make_b(eq, b) + self.formfuncs.extend(F_target) + self.formrhs.extend(b) + + self.formfuncs = self.scale_essential_bcs(self.formfuncs) + + def make_F_target(self, eq, F_target, targets): + arrays = self.arrays + volume = self.target.grid.symbolic_volume_cell + if isinstance(eq, EssentialBC): + # The initial guess satisfies the essential BCs, so this term is zero. + # Still included to support Jacobian testing via finite differences. + rhs = arrays['x'] - eq.rhs + zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) + # Move essential boundary condition to the right-hand side + zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (zero_row, zero_col) + else: + if isinstance(F_target, (int, float)): + rhs = F_target * volume + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) * volume + return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) + + def make_b(self, eq, b): + rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = rhs * self.target.grid.symbolic_volume_cell + return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + + +class MixedResidual(BaseResidual): + """ + """ + # TODO: change default and pass in correct scale + def __init__(self, targets, time_mapper, arrays, scale=1.0): + self.targets = as_tuple(targets) + self.time_mapper = time_mapper + self.arrays = arrays + self.scale = scale + self.formfuncs = [] + + def build_equations(self, eqn_dict): + all_formfuncs = [] + for target, eqns in eqn_dict.items(): + + formfuncs = chain.from_iterable( + self.build_function_eq(eq, target) + for eq in as_tuple(eqns) + ) + + # scale, = self._diag_scale[arrays[target]['x']] + # fix this + # scale = 1.0 + all_formfuncs.extend(self.scale_essential_bcs(formfuncs)) + + self.formfuncs = tuple(sorted( + all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + )) + + + def build_function_eq(self, eq, target): + zeroed = eq.lhs - eq.rhs + + zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + + mapper = {} + for t in self.targets: + target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) + mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) + + if isinstance(eq, EssentialBC): + rhs = self.arrays[target]['x'] - eq.rhs + zero_row = ZeroRow( + self.arrays[target]['f'], rhs, subdomain=eq.subdomain + ) + zero_col = ZeroColumn( + self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain + ) + return (zero_row, zero_col) + else: + if isinstance(zeroed, (int, float)): + # rhs = zeroed * self.cell_area + rhs = zeroed + else: + rhs = zeroed.subs(mapper) + # rhs = rhs.subs(self.time_mapper)*self.cell_area + rhs = rhs.subs(self.time_mapper) + + return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) + + + +def targets_to_arrays(array, targets): + """ + Map each target in `targets` to a corresponding array generated from `array`, + matching the spatial indices of the target. + Example: + -------- + >>> array + vec_u(x, y) + >>> targets + {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} + >>> targets_to_arrays(array, targets) + {u(t + dt, x - h_x, y): vec_u(x - h_x, y), + u(t + dt, x + h_x, y): vec_u(x + h_x, y), + u(t + dt, x, y): vec_u(x, y)} + """ + space_indices = [ + tuple(f.indices[d] for d in f.space_dimensions) for f in targets + ] + array_targets = [ + array.subs(dict(zip(array.indices, i))) for i in space_indices + ] + return dict(zip(targets, array_targets)) diff --git a/devito/types/grid.py b/devito/types/grid.py index f0d4a440d5..404619f3c0 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -280,10 +280,15 @@ def interior(self): """The interior SubDomain of the Grid.""" return self.subdomains['interior'] + @property + def symbolic_volume_cell(self): + """Symbolic volume of a single cell e.g. h_x*h_y*h_z in 3D.""" + return prod(d.spacing for d in self.dimensions) + @property def volume_cell(self): """Volume of a single cell e.g h_x*h_y*h_z in 3D.""" - return prod(d.spacing for d in self.dimensions).subs(self.spacing_map) + return self.symbolic_volume_cell.subs(self.spacing_map) @property def spacing(self): From 6897ab8490638586879bc7e7c907bd6f9a0abe87 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 4 Jun 2025 15:08:36 +0100 Subject: [PATCH 14/53] dsl: Fix compatible scaling for single fields --- .github/workflows/pytest-petsc.yml | 2 +- devito/petsc/iet/routines.py | 12 +- devito/petsc/solve.py | 136 ++--- devito/petsc/types/array.py | 4 + devito/petsc/types/equation.py | 39 +- devito/petsc/types/object.py | 1 - devito/petsc/types/types.py | 370 +++++++----- examples/petsc/Poisson/01_poisson.py | 7 +- examples/petsc/petsc_test.py | 2 +- .../random/{biharmonic => }/02_biharmonic.py | 0 .../random/biharmonic/biharmonic_matfree.c | 556 ------------------ .../biharmonic/biharmonic_matfree_nonscaled.c | 553 ----------------- tests/test_petsc.py | 1 - 13 files changed, 321 insertions(+), 1362 deletions(-) rename examples/petsc/random/{biharmonic => }/02_biharmonic.py (100%) delete mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree.c delete mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 19446642af..9c035a86e7 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -83,7 +83,7 @@ jobs: ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py - ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/biharmonic/02_biharmonic.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/02_biharmonic.py - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 1fc4cd7758..8f580a1b52 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -108,11 +108,10 @@ def zero_memory(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - # from IPython.core.debugger import set_trace self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) - if fielddata.initialguess: + if fielddata.initialguess.equations: self._make_initialguess(fielddata) self._make_user_struct_callback() @@ -516,7 +515,7 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess + initguess = fielddata.initialguess.equations sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -1208,7 +1207,6 @@ def _extend_build(self, base_dict): 'f': fbundle, 'x': xbundle, 'bundle_mapper': bundle_mapper, - # TODO: maybe this shouldn't be here 'target_indices': target_indices } @@ -1438,10 +1436,6 @@ def _create_dmda(self, dmda): class CoupledSetup(BaseSetup): - # @property - # def snes_ctx(self): - # return Byref(self.solver_objs['jacctx']) - def _setup(self): # TODO: minimise code duplication with superclass objs = self.objs @@ -1683,7 +1677,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) targets = self.injectsolve.expr.rhs.fielddata.targets - # TODO: Optimise the ccode generated here + # TODO: optimise the ccode generated here pre_solve = () post_solve = () diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index eb8f171b21..d2f0f39809 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,29 +1,55 @@ -from functools import singledispatch - -import sympy -import numpy as np -from itertools import chain -from collections import defaultdict - -from devito.finite_differences.differentiable import Mul -from devito.finite_differences.derivative import Derivative -from devito.types import Eq, Symbol, SteppingDimension, TimeFunction +from devito.types import Symbol, SteppingDimension from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions -from devito.symbolics.extraction import (separate_eqn, centre_stencil, - generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian, Residual, - MixedResidual, MixedJacobian) -from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn + MixedResidual, MixedJacobian, InitialGuess) +from devito.petsc.types.equation import EssentialBC __all__ = ['PETScSolve'] -def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): +def PETScSolve(target_eqns, target=None, solver_parameters=None): + """ + Returns a symbolic equation representing a linear PETSc solver, + enriched with all the necessary metadata for execution within an `Operator`. + When passed to an `Operator`, this symbolic equation triggers code generation + and lowering to the PETSc backend. + + This function supports both single- and multi-target systems. In the multi-target + (mixed system) case, the solution vector spans all provided target fields. + + Parameters + ---------- + target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq + The targets and symbolic equations defining the system to be solved. + + - **Single-field problem**: + Pass a single Eq or list of Eq, and specify `target` separately: + PETScSolve(Eq1, target) + PETScSolve([Eq1, Eq2], target) + + - **Multi-field (mixed) problem**: + Pass a dictionary mapping each target field to its Eq(s): + PETScSolve({u: Eq1, v: Eq2}) + PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) + + target : Function-like + The function (e.g., `Function`, `TimeFunction`) into which the linear + system solves. This represents the solution vector updated by the solver. + + solver_parameters : dict, optional + PETSc solver options. + + Returns + ------- + Eq + A symbolic equation that wraps the system, solver metadata, + and boundary conditions. This can be passed directly to a Devito Operator. + """ if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() else: @@ -59,19 +85,11 @@ def linear_solve_args(self): eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - jacobian = Jacobian(target, self.time_mapper, arrays) - jacobian.build_block(eqns) - - scale = 1.0 + jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, self.time_mapper, arrays, scale) - residual.build_equations(eqns) + residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = [ - eq for eq in - (self.make_initial_guess(e, target, arrays) for e in eqns) - if eq is not None - ] + initialguess = InitialGuess(target, eqns, arrays) field_data = FieldData( target=target, @@ -83,22 +101,6 @@ def linear_solve_args(self): return target, tuple(funcs), field_data - def make_initial_guess(self, eq, target, arrays): - """ - Enforce initial guess to satisfy essential BCs. - # TODO: For time-stepping, only enforce these once outside the time loop - and use the previous time-step solution as the initial guess for next time step. - # TODO: Extend this to "coupled". - """ - if isinstance(eq, EssentialBC): - assert eq.lhs == target - return Eq( - arrays['x'], eq.rhs, - subdomain=eq.subdomain - ) - else: - return None - def generate_arrays(self, target): return { p: PETScArray(name=f'{p}_{target.name}', @@ -112,7 +114,6 @@ def generate_arrays(self, target): class InjectMixedSolve(InjectSolve): def linear_solve_args(self): - combined_eqns = [] for eqns in self.target_eqns.values(): combined_eqns.extend(eqns) @@ -123,11 +124,14 @@ def linear_solve_args(self): arrays = self.generate_arrays_combined(*coupled_targets) - jacobian = MixedJacobian(coupled_targets, self.time_mapper, arrays) - jacobian.build_blocks(self.target_eqns) + jacobian = MixedJacobian( + self.target_eqns, arrays, self.time_mapper + ) - residual = MixedResidual(coupled_targets, self.time_mapper, arrays) - residual.build_equations(self.target_eqns) + residual = MixedResidual( + self.target_eqns, arrays, + self.time_mapper, jacobian.target_scaler_mapper + ) all_data = MultipleFieldData( targets=coupled_targets, @@ -136,46 +140,10 @@ def linear_solve_args(self): residual=residual ) - # TODO: rethink what target to return here??? return coupled_targets[0], tuple(funcs), all_data def generate_arrays_combined(self, *targets): - return { - target: { - p: PETScArray( - name=f'{p}_{target.name}', - target=target, - liveness='eager', - localinfo=localinfo - ) - for p in prefixes - } - for target in targets - } - - -def targets_to_arrays(array, targets): - """ - Map each target in `targets` to a corresponding array generated from `array`, - matching the spatial indices of the target. - Example: - -------- - >>> array - vec_u(x, y) - >>> targets - {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} - >>> targets_to_arrays(array, targets) - {u(t + dt, x - h_x, y): vec_u(x - h_x, y), - u(t + dt, x + h_x, y): vec_u(x + h_x, y), - u(t + dt, x, y): vec_u(x, y)} - """ - space_indices = [ - tuple(f.indices[d] for d in f.space_dimensions) for f in targets - ] - array_targets = [ - array.subs(dict(zip(array.indices, i))) for i in space_indices - ] - return dict(zip(targets, array_targets)) + return {target: self.generate_arrays(target) for target in targets} def generate_time_mapper(funcs): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 049785662b..8b57aca44c 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -69,6 +69,10 @@ def dimensions(self): def target(self): return self._target + @property + def grid(self): + return self.target.grid + @property def coefficients(self): """Form of the coefficients of the function.""" diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index a5c5bc55ef..ad929c8ed5 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -6,33 +6,38 @@ class EssentialBC(Eq): """ - A special equation used to handle essential boundary conditions - in the PETSc solver. Until PetscSection + DMDA is supported, - we treat essential boundary conditions as trivial equations - in the solver, where we place 1.0 (scaled) on the diagonal of - the jacobian, zero symmetrically and move the boundary - data to the right-hand side. - - NOTE: When users define essential boundary conditions, they need to ensure that - the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. This will be documented in the PETSc examples. + Represents an essential boundary condition for use with PETScSolve. + + Due to ongoing work on PetscSection and DMDA integration (WIP), + these conditions are imposed as trivial equations. The compiler + will automatically zero the corresponding rows/columns in the Jacobian + and lift the boundary terms into the residual RHS. + + Note: + - To define an essential boundary condition, use: + Eq(target, boundary_value, subdomain=...), + where `target` is the Function-like object passed to PETScSolve. + - SubDomains used for multiple EssentialBCs must not overlap. """ pass class ZeroRow(EssentialBC): """ - Equation used to zero the row of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. + Equation used to zero all entries, except the diagonal, + of a row in the Jacobian. + + Note: + This is only used interally by the compiler, not by users. """ pass class ZeroColumn(EssentialBC): """ - Equation used to zero the column of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. + Equation used to zero the column of the Jacobian. + + Note: + This is only used interally by the compiler, not by users. """ - pass \ No newline at end of file + pass diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index e62d125dc2..32892f2aa4 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -209,7 +209,6 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None - class PETScArrayObject(PetscMixin, ArrayObject): _data_alignment = False diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 3897552f7b..a0a9cd2fe5 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -5,7 +5,7 @@ from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables -from devito.symbolics.extraction import separate_eqn, generate_targets +from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn from devito.types.equation import Eq from devito.operations.solve import eval_time_derivatives @@ -186,7 +186,7 @@ def space_order(self): @property def targets(self): - return (self.target,) + return as_tuple(self.target) class MultipleFieldData(FieldData): @@ -229,22 +229,104 @@ def space_order(self): ) return space_orders.pop() - @property - def jacobian(self): - return self._jacobian - @property def targets(self): return self._targets - -class BaseJacobian: - def __init__(self, targets, time_mapper, arrays): - self.targets = as_tuple(targets) +class Jacobian: + def __init__(self, target, eqns, arrays, time_mapper): + self.target = target + self.eqns = eqns + self.arrays = arrays self.time_mapper = time_mapper + self._build_matvecs() + + @property + def matvecs(self): + return self._matvecs + + @property + def scdiag(self): + return self._scdiag + + def _build_matvecs(self): + matvecs = [ + e for eq in self.eqns for e in + self._build_matvec_eq(eq) + if e is not None + ] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + matvecs = self._scale_non_bcs(matvecs) + scdiag = self._compute_scdiag(matvecs) + matvecs = self._scale_bcs(matvecs, scdiag) + + self._matvecs = matvecs + self._scdiag = scdiag + + def _build_matvec_eq(self, eq, target=None, arrays=None): + target = target or self.target + arrays = arrays or self.arrays + + b, F_target, _, targets = separate_eqn(eq, target) + if F_target: + return self._make_matvec(eq, F_target, targets, arrays) + return (None,) + + def _make_matvec(self, eq, F_target, targets, arrays): + if isinstance(eq, EssentialBC): + # NOTE: Essential BCs are trivial equations in the solver. + # See `EssentialBC` for more details. + rhs = arrays['x'] + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0., subdomain=eq.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + + def _scale_non_bcs(self, matvecs, target=None): + target = target or self.target + vol = target.grid.symbolic_volume_cell + + return [ + m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) + for m in matvecs + ] + + def _compute_scdiag(self, matvecs, arrays=None): + """ + """ + arrays = arrays or self.arrays + + centres = { + centre_stencil(m.rhs, arrays['x'], as_coeff=True) + for m in matvecs if not isinstance(m, EssentialBC) + } + # add comments + return centres.pop() if len(centres) == 1 else 1.0 + + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the essential BCs + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + + +class MixedJacobian(Jacobian): + def __init__(self, target_eqns, arrays, time_mapper): + """ + """ + self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays + self.time_mapper = time_mapper self.submatrices = self._initialize_submatrices() + self._build_blocks(target_eqns) def _initialize_submatrices(self): """ @@ -260,11 +342,33 @@ def _initialize_submatrices(self): submatrices[target][key] = { 'matvecs': None, 'derivative_wrt': self.targets[j], - 'index': i * num_targets + j + 'index': i * num_targets + j, + 'scdiag': None } return submatrices + def _build_blocks(self, target_eqns): + for target, eqns in target_eqns.items(): + self._build_block(target, eqns) + + def _build_block(self, target, eqns): + arrays = self.arrays[target] + for submat, mtvs in self.submatrices[target].items(): + matvecs = [ + e for eq in eqns for e in + self._build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + matvecs = self._scale_non_bcs(matvecs, target) + scdiag = self._compute_scdiag(matvecs, arrays) + matvecs = self._scale_bcs(matvecs, scdiag) + + if matvecs: + self.set_submatrix(target, submat, matvecs, scdiag) + @property def submatrix_keys(self): """ @@ -295,21 +399,21 @@ def submat_to_index(self): for key, value in submats.items() } - @property - def diagonal_submatrix_keys(self): + # CHECK/TEST THIS + def is_diagonal_submatrix(self, key): """ - Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + Return True if the given key corresponds to a diagonal + submatrix (e.g., 'J00', 'J11'), else False. """ - keys = [] - for i, target in enumerate(self.targets): + for i, t in enumerate(self.targets): diag_key = f'J{i}{i}' - if diag_key in self.submatrices[target]: - keys.append(diag_key) - return keys + if key == diag_key and diag_key in self.submatrices[t]: + return True + return False - def set_submatrix(self, field, key, matvecs): + def set_submatrix(self, field, key, matvecs, scdiag): """ - Set a specific submatrix for a field. + Set the matrix-vector equations for a submatrix. Parameters ---------- @@ -321,7 +425,8 @@ def set_submatrix(self, field, key, matvecs): The matrix-vector equations forming the submatrix. """ if field in self.submatrices and key in self.submatrices[field]: - self.submatrices[field][key]["matvecs"] = matvecs + self.submatrices[field][key]['matvecs'] = matvecs + self.submatrices[field][key]['scdiag'] = scdiag else: raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') @@ -331,112 +436,63 @@ def get_submatrix(self, field, key): """ return self.submatrices.get(field, {}).get(key, None) - def build_matvec_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - if F_target: - return self.make_matvec(eq, F_target, targets, arrays) - return (None,) - - def make_matvec(self, eq, F_target, targets, arrays): - if isinstance(eq, EssentialBC): - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) - # and zero symmetrically. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - # rhs = rhs.subs(self.time_mapper) * self.cell_area - rhs = rhs = rhs.subs(self.time_mapper) - - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + @property + def target_scaler_mapper(self): + """ + Return a mapping from each target to its diagonal submatrix's scaler. + """ + mapper = {} + for target, submats in self.submatrices.items(): + for key, value in submats.items(): + if self.is_diagonal_submatrix(key): + mapper[target] = value.get('scdiag') + break + return mapper def __repr__(self): + # TODO: edit return str(self.submatrices) -class Jacobian(BaseJacobian): +class Residual: + """ + """ + def __init__(self, target, eqns, arrays, time_mapper, scdiag): + self.target = target + self.eqns = eqns + self.arrays = arrays + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_equations() @property - def target(self): - return self.targets[0] + def formfuncs(self): + """ + """ + return self._formfuncs @property - def matvecs(self): - return self.submatrices[self.target]['J00']['matvecs'] - - # TODO: use same structure arrays for both jacobian and mixedjacobian - def build_block(self, eqns): - for submat, mtvs in self.submatrices[self.target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - self.set_submatrix(self.target, submat, matvecs) - - -class MixedJacobian(BaseJacobian): - - # TODO: use same structure arrays for both jacobian and mixedjacobian - def build_block(self, target, eqns): - for submat, mtvs in self.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays[target]) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - self.set_submatrix(target, submat, matvecs) - - def build_blocks(self, target_eqns): - for target, eqns in target_eqns.items(): - self.build_block(target, eqns) - - -class BaseResidual: - def scale_essential_bcs(self, equations): + def formrhs(self): """ """ - return [ - eq._rebuild(rhs=self.scale * eq.rhs) if isinstance(eq, ZeroRow) else eq - for eq in equations - ] + return self._formrhs - -class Residual(BaseResidual): - """ - """ - - def __init__(self, target, time_mapper, arrays, scale): - self.target = target - self.time_mapper = time_mapper - self.arrays = arrays - self.scale = scale - self.formfuncs = [] - self.formrhs = [] - - def build_equations(self, eqns): + def _build_equations(self): """ """ - for eq in eqns: + funcs = [] + rhs = [] + + for eq in self.eqns: b, F_target, _, targets = separate_eqn(eq, self.target) - F_target = self.make_F_target(eq, F_target, targets) - b = self.make_b(eq, b) - self.formfuncs.extend(F_target) - self.formrhs.extend(b) + funcs.extend(self._make_F_target(eq, F_target, targets)) + # TODO: if b is zero then don't need a rhs vector+callback + rhs.extend(self._make_b(eq, b)) - self.formfuncs = self.scale_essential_bcs(self.formfuncs) + self._formfuncs = [self._scale_bcs(eq) for eq in funcs] + self._formrhs = rhs - def make_F_target(self, eq, F_target, targets): + def _make_F_target(self, eq, F_target, targets): arrays = self.arrays volume = self.target.grid.symbolic_volume_cell if isinstance(eq, EssentialBC): @@ -455,55 +511,64 @@ def make_F_target(self, eq, F_target, targets): rhs = rhs.subs(self.time_mapper) * volume return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) - def make_b(self, eq, b): + def _make_b(self, eq, b): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + def _scale_bcs(self, eq, scdiag=None): + """ + Scale ZeroRow equations using scdiag + """ + scdiag = scdiag or self.scdiag + return eq._rebuild(rhs=scdiag * eq.rhs) if isinstance(eq, ZeroRow) else eq + -class MixedResidual(BaseResidual): +class MixedResidual(Residual): """ """ - # TODO: change default and pass in correct scale - def __init__(self, targets, time_mapper, arrays, scale=1.0): - self.targets = as_tuple(targets) - self.time_mapper = time_mapper + def __init__(self, target_eqns, arrays, time_mapper, scdiag): + self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays - self.scale = scale - self.formfuncs = [] + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_equations(target_eqns) - def build_equations(self, eqn_dict): + @property + def formrhs(self): + """ + """ + return None + + def _build_equations(self, target_eqns): all_formfuncs = [] - for target, eqns in eqn_dict.items(): + for target, eqns in target_eqns.items(): formfuncs = chain.from_iterable( - self.build_function_eq(eq, target) + self._build_function_eq(eq, target) for eq in as_tuple(eqns) ) + all_formfuncs.extend(formfuncs) - # scale, = self._diag_scale[arrays[target]['x']] - # fix this - # scale = 1.0 - all_formfuncs.extend(self.scale_essential_bcs(formfuncs)) - - self.formfuncs = tuple(sorted( + self._formfuncs = tuple(sorted( all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) )) - - def build_function_eq(self, eq, target): + def _build_function_eq(self, eq, target): zeroed = eq.lhs - eq.rhs zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + volume = target.grid.symbolic_volume_cell + mapper = {} for t in self.targets: target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): - rhs = self.arrays[target]['x'] - eq.rhs + rhs = (self.arrays[target]['x'] - eq.rhs)*self.scdiag[target] zero_row = ZeroRow( self.arrays[target]['f'], rhs, subdomain=eq.subdomain ) @@ -513,16 +578,51 @@ def build_function_eq(self, eq, target): return (zero_row, zero_col) else: if isinstance(zeroed, (int, float)): - # rhs = zeroed * self.cell_area - rhs = zeroed + rhs = zeroed * volume else: rhs = zeroed.subs(mapper) - # rhs = rhs.subs(self.time_mapper)*self.cell_area - rhs = rhs.subs(self.time_mapper) + rhs = rhs.subs(self.time_mapper)*volume return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) +class InitialGuess: + """ + Enforce initial guess to satisfy essential BCs. + # TODO: Extend this to "coupled". + """ + def __init__(self, target, eqns, arrays): + self.target = target + self.eqns = as_tuple(eqns) + self.arrays = arrays + self._build_equations() + + @property + def equations(self): + """ + """ + return self._equations + + def _build_equations(self): + """ + Return a list of initial guess equations. + """ + self._equations = [ + eq for eq in + (self._make_initial_guess(e) for e in self.eqns) + if eq is not None + ] + + def _make_initial_guess(self, eq): + if isinstance(eq, EssentialBC): + assert eq.lhs == self.target + return Eq( + self.arrays['x'], eq.rhs, + subdomain=eq.subdomain + ) + else: + return None + def targets_to_arrays(array, targets): """ diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 70f17d94c6..4981566af6 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,6 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [7] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -112,7 +111,7 @@ def analytical(x, y): error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) errors.append(error) -# slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) +slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -# assert slope > 1.9 -# assert slope < 2.1 +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..12d245480e 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -20,7 +20,7 @@ v.data[:] = 5.0 -eq = Eq(v, u.laplace, subdomain=grid.interior) +eq = Eq(0., u.laplace, subdomain=grid.interior) petsc = PETScSolve([eq], u) diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py similarity index 100% rename from examples/petsc/random/biharmonic/02_biharmonic.py rename to examples/petsc/random/02_biharmonic.py diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c deleted file mode 100644 index 2fb432382f..0000000000 --- a/examples/petsc/random/biharmonic/biharmonic_matfree.c +++ /dev/null @@ -1,556 +0,0 @@ - -// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c - -static char help[] = -"Solve the linear biharmonic equation in 2D. Equation is\n" -" Lap^2 u = f\n" -"where Lap = - grad^2 is the positive Laplacian, equivalently\n" -" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" -"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" -"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" -"2x2 block system with SPD Laplacian blocks on the diagonal:\n" -" | Lap | 0 | | v | | f | \n" -" |-----|-----| |---| = |---| \n" -" | -I | Lap | | u | | 0 | \n" -"Includes manufactured, polynomial exact solution. The discretization is\n" -"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" -"Recommended preconditioning combines fieldsplit:\n" -" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" -"with multigrid as the preconditioner for the diagonal blocks:\n" -" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" -"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" -"monolithic multigrid (-pc_type mg|gamg).\n\n"; - -#include - -typedef struct { - PetscReal v, u; -} Field; - -typedef struct { - PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side -} BiharmCtx; - -struct JacobianCtx -{ - DM * subdms; - IS * fields; - Mat * submats; -} ; - -struct SubMatrixCtx -{ - IS * rows; - IS * cols; -} ; - -static PetscReal c(PetscReal x) { - return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); -} - -static PetscReal ddc(PetscReal x) { - return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal d4c(PetscReal x) { - return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { - return c(x) * c(y); -} - -static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { - return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u -} - -static PetscReal f_fcn(PetscReal x, PetscReal y) { - return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u -} - -extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); -extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); -extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); -extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); - -int main(int argc,char **argv) { - DM da; - SNES snes; - Vec w, w_initial, w_exact; - BiharmCtx user; - Field **aW; - PetscReal normv, normu, errv, erru; - DMDALocalInfo info; - IS *fields; - DM *subdms; - PetscInt nfields; - - struct JacobianCtx jctx0; - Mat J; - - PetscCall(PetscInitialize(&argc,&argv,NULL,help)); - - user.f = &f_fcn; - PetscCall(DMDACreate2d(PETSC_COMM_WORLD, - DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 6,6,PETSC_DECIDE,PETSC_DECIDE, - 2,1, // degrees of freedom, stencil width - NULL,NULL,&da)); - PetscCall(DMSetApplicationContext(da,&user)); - PetscCall(DMSetFromOptions(da)); - PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates - PetscCall(DMSetMatType(da, MATSHELL)); - PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); - PetscCall(DMDASetFieldName(da,0,"v")); - PetscCall(DMDASetFieldName(da,1,"u")); - PetscCall(DMCreateMatrix(da,&J)); - - PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); - PetscCall(SNESSetDM(snes,da)); - PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); - PetscCall(SNESSetType(snes,SNESKSPONLY)); - PetscCall(SNESSetFromOptions(snes)); - - PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); - PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); - - PetscCall(MatSetDM(J,da)); - PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); - PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); - PetscCall(MatShellSetContext(J,&(jctx0))); - PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); - - PetscCall(DMGetGlobalVector(da,&w_initial)); - PetscCall(VecSet(w_initial,0.0)); - PetscCall(SNESSolve(snes,NULL,w_initial)); - // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(DMRestoreGlobalVector(da,&w_initial)); - PetscCall(DMDestroy(&da)); - - PetscCall(SNESGetSolution(snes,&w)); - PetscCall(SNESGetDM(snes,&da)); - PetscCall(DMDAGetLocalInfo(da,&info)); - - PetscCall(DMCreateGlobalVector(da,&w_exact)); - PetscCall(DMDAVecGetArray(da,w_exact,&aW)); - PetscCall(FormExactWLocal(&info,aW,&user)); - PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); - PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); - PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); - PetscCall(VecAXPY(w,-1.0,w_exact)); - PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); - PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); - PetscCall(PetscPrintf(PETSC_COMM_WORLD, - "done on %d x %d grid ...\n" - " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", - info.mx,info.my,errv/normv,erru/normu)); - - - PetscCall(ISDestroy(&(fields[0]))); - PetscCall(ISDestroy(&(fields[1]))); - PetscCall(PetscFree(fields)); - PetscCall(DMDestroy(&(subdms[0]))); - PetscCall(DMDestroy(&(subdms[1]))); - PetscCall(PetscFree(subdms)); - PetscCall(VecDestroy(&w_exact)); - PetscCall(MatDestroy(&J)); - PetscCall(SNESDestroy(&snes)); - PetscCall(PetscFinalize()); - return 0; -} - -PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, x, y; - PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); - hx = (xymax[0] - xymin[0]) / (info->mx - 1); - hy = (xymax[1] - xymin[1]) / (info->my - 1); - for (j = info->ys; j < info->ys + info->ym; j++) { - y = j * hy; - for (i = info->xs; i < info->xs + info->xm; i++) { - x = i * hx; - aW[j][i].u = u_exact_fcn(x,y); - aW[j][i].v = lap_u_exact_fcn(x,y); - } - } - return 0; -} - - -PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) -{ - Vec xlocal, flocal; - DMDALocalInfo info; - DM da; - PetscScalar *x_vec, *f_vec; - - BiharmCtx *user; - - PetscCall(SNESGetDM(snes,&da)); - - PetscCall(DMGetApplicationContext(da,&user)); - - PetscCall(DMDAGetLocalInfo(da,&info)); - PetscCall(DMGetLocalVector(da,&xlocal)); - PetscCall(DMGetLocalVector(da,&flocal)); - - PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); - PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); - - PetscCall(VecGetArray(xlocal,&x_vec)); - PetscCall(VecGetArray(flocal,&f_vec)); - - Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; - Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - ff[j][i].v = scdiag * xx[j][i].v; - ff[j][i].u = scdiag * xx[j][i].u; - } else { - ve = xx[j][i+1].v; - vw = xx[j][i-1].v; - vn = xx[j+1][i].v; - vs = xx[j-1][i].v; - ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) - - darea * (*(user->f))(x,y); - ue = xx[j][i+1].u; - uw = xx[j][i-1].u; - un = xx[j+1][i].u; - us = xx[j-1][i].u; - ff[j][i].u = - darea * xx[j][i].v - + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); - } - } - } - - PetscCall(VecRestoreArray(xlocal,&x_vec)); - PetscCall(VecRestoreArray(flocal,&f_vec)); - - PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); - PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); - PetscCall(DMRestoreLocalVector(da,&xlocal)); - PetscCall(DMRestoreLocalVector(da,&flocal)); - - return 0; -} - - -PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1./ (info.mx - 1); - hy = 1./ (info.my - 1); - darea = hx * hy; // multiply FD equations by this - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = scdiag * x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = 0.0; - } else { - y_v[j][i] = -darea * x_v[j][i]; - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - - -PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; -// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = scdiag * x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) -{ - Vec J00X; - Vec J00Y; - Vec J10X; - Vec J10Y; - Vec J11X; - Vec J11Y; - - struct SubMatrixCtx * J00ctx; - struct SubMatrixCtx * J10ctx; - struct SubMatrixCtx * J11ctx; - struct JacobianCtx * jctx; - - PetscFunctionBeginUser; - - PetscCall(MatShellGetContext(J,&(jctx))); - - PetscCall(VecSet(Y,0.0)); - - Mat J00 = jctx->submats[0]; - PetscCall(MatShellGetContext(J00,&(J00ctx))); - PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); - PetscCall(MatMult(J00,J00X,J00Y)); - PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); - - Mat J10 = jctx->submats[2]; - PetscCall(MatShellGetContext(J10,&(J10ctx))); - PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); - PetscCall(MatMult(J10,J10X,J10Y)); - PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); - - Mat J11 = jctx->submats[3]; - PetscCall(MatShellGetContext(J11,&(J11ctx))); - PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); - PetscCall(MatMult(J11,J11X,J11Y)); - PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); - - - PetscFunctionReturn(0); -} - - -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) -{ - PetscFunctionBeginUser; - - PetscInt M; - PetscInt N; - Mat block; - DM dm0; - PetscInt dof; - - struct UserCtx0 * ctx0; - struct JacobianCtx * jctx; - struct SubMatrixCtx * subctx; - - PetscCall(MatShellGetContext(J,&(jctx))); - DM * subdms = jctx->subdms; - - PetscInt nsubmats = nfields*nfields; - PetscCall(PetscCalloc1(nsubmats,submats)); - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); - PetscInt subblockrows = M*N; - PetscInt subblockcols = M*N; - Mat * submat_arr = *submats; - - for (int i = 0; i <= nsubmats - 1; i += 1) - { - PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); - PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); - PetscCall(MatSetType(block,MATSHELL)); - PetscCall(PetscMalloc1(1,&(subctx))); - PetscInt rowidx = i / dof; - PetscInt colidx = (i)%(dof); - subctx->rows = &(irow[rowidx]); - subctx->cols = &(icol[colidx]); - PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); - PetscCall(MatSetDM(block,subdms[rowidx])); - PetscCall(MatShellSetContext(block,subctx)); - PetscCall(MatSetUp(block)); - submat_arr[i] = block; - } - PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); - - PetscFunctionReturn(0); -} - - -PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) -{ - PetscFunctionBeginUser; - - jctx->subdms = subdms; - jctx->fields = fields; - - PetscFunctionReturn(0); -} diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c deleted file mode 100644 index 04f2a02927..0000000000 --- a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c +++ /dev/null @@ -1,553 +0,0 @@ - -// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c - -static char help[] = -"Solve the linear biharmonic equation in 2D. Equation is\n" -" Lap^2 u = f\n" -"where Lap = - grad^2 is the positive Laplacian, equivalently\n" -" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" -"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" -"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" -"2x2 block system with SPD Laplacian blocks on the diagonal:\n" -" | Lap | 0 | | v | | f | \n" -" |-----|-----| |---| = |---| \n" -" | -I | Lap | | u | | 0 | \n" -"Includes manufactured, polynomial exact solution. The discretization is\n" -"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" -"Recommended preconditioning combines fieldsplit:\n" -" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" -"with multigrid as the preconditioner for the diagonal blocks:\n" -" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" -"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" -"monolithic multigrid (-pc_type mg|gamg).\n\n"; - -#include - -typedef struct { - PetscReal v, u; -} Field; - -typedef struct { - PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side -} BiharmCtx; - -struct JacobianCtx -{ - DM * subdms; - IS * fields; - Mat * submats; -} ; - -struct SubMatrixCtx -{ - IS * rows; - IS * cols; -} ; - -static PetscReal c(PetscReal x) { - return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); -} - -static PetscReal ddc(PetscReal x) { - return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal d4c(PetscReal x) { - return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { - return c(x) * c(y); -} - -static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { - return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u -} - -static PetscReal f_fcn(PetscReal x, PetscReal y) { - return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u -} - -extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); -extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); -extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); -extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); - -int main(int argc,char **argv) { - DM da; - SNES snes; - Vec w, w_initial, w_exact; - BiharmCtx user; - Field **aW; - PetscReal normv, normu, errv, erru; - DMDALocalInfo info; - IS *fields; - DM *subdms; - PetscInt nfields; - - struct JacobianCtx jctx0; - Mat J; - - PetscCall(PetscInitialize(&argc,&argv,NULL,help)); - - user.f = &f_fcn; - PetscCall(DMDACreate2d(PETSC_COMM_WORLD, - DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 33,33,PETSC_DECIDE,PETSC_DECIDE, - 2,1, // degrees of freedom, stencil width - NULL,NULL,&da)); - PetscCall(DMSetApplicationContext(da,&user)); - PetscCall(DMSetFromOptions(da)); - PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates - PetscCall(DMSetMatType(da, MATSHELL)); - PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); - PetscCall(DMDASetFieldName(da,0,"v")); - PetscCall(DMDASetFieldName(da,1,"u")); - PetscCall(DMCreateMatrix(da,&J)); - - PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); - PetscCall(SNESSetDM(snes,da)); - PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); - PetscCall(SNESSetType(snes,SNESKSPONLY)); - PetscCall(SNESSetFromOptions(snes)); - - PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); - PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); - - PetscCall(MatSetDM(J,da)); - PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); - PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); - PetscCall(MatShellSetContext(J,&(jctx0))); - PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); - - PetscCall(DMGetGlobalVector(da,&w_initial)); - PetscCall(VecSet(w_initial,0.0)); - PetscCall(SNESSolve(snes,NULL,w_initial)); - // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(DMRestoreGlobalVector(da,&w_initial)); - PetscCall(DMDestroy(&da)); - - PetscCall(SNESGetSolution(snes,&w)); - PetscCall(SNESGetDM(snes,&da)); - PetscCall(DMDAGetLocalInfo(da,&info)); - - PetscCall(DMCreateGlobalVector(da,&w_exact)); - PetscCall(DMDAVecGetArray(da,w_exact,&aW)); - PetscCall(FormExactWLocal(&info,aW,&user)); - - PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); - // PetscCall(VecView(w_exact,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); - PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); - PetscCall(VecAXPY(w,-1.0,w_exact)); - PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); - PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); - PetscCall(PetscPrintf(PETSC_COMM_WORLD, - "done on %d x %d grid ...\n" - " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", - info.mx,info.my,errv/normv,erru/normu)); - - - PetscCall(ISDestroy(&(fields[0]))); - PetscCall(ISDestroy(&(fields[1]))); - PetscCall(PetscFree(fields)); - PetscCall(DMDestroy(&(subdms[0]))); - PetscCall(DMDestroy(&(subdms[1]))); - PetscCall(PetscFree(subdms)); - PetscCall(VecDestroy(&w_exact)); - PetscCall(MatDestroy(&J)); - PetscCall(SNESDestroy(&snes)); - PetscCall(PetscFinalize()); - return 0; -} - -PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, x, y; - PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); - hx = (xymax[0] - xymin[0]) / (info->mx - 1); - hy = (xymax[1] - xymin[1]) / (info->my - 1); - for (j = info->ys; j < info->ys + info->ym; j++) { - y = j * hy; - for (i = info->xs; i < info->xs + info->xm; i++) { - x = i * hx; - aW[j][i].u = u_exact_fcn(x,y); - aW[j][i].v = lap_u_exact_fcn(x,y); - } - } - return 0; -} - - -PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) -{ - Vec xlocal, flocal; - DMDALocalInfo info; - DM da; - PetscScalar *x_vec, *f_vec; - - BiharmCtx *user; - - PetscCall(SNESGetDM(snes,&da)); - - PetscCall(DMGetApplicationContext(da,&user)); - - PetscCall(DMDAGetLocalInfo(da,&info)); - PetscCall(DMGetLocalVector(da,&xlocal)); - PetscCall(DMGetLocalVector(da,&flocal)); - - PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); - PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); - - PetscCall(VecGetArray(xlocal,&x_vec)); - PetscCall(VecGetArray(flocal,&f_vec)); - - Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; - Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - ff[j][i].v = xx[j][i].v; - ff[j][i].u = xx[j][i].u; - } else { - ve = xx[j][i+1].v; - vw = xx[j][i-1].v; - vn = xx[j+1][i].v; - vs = xx[j-1][i].v; - ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) - - (*(user->f))(x,y); - ue = xx[j][i+1].u; - uw = xx[j][i-1].u; - un = xx[j+1][i].u; - us = xx[j-1][i].u; - ff[j][i].u = -xx[j][i].v - + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); - } - } - } - - PetscCall(VecRestoreArray(xlocal,&x_vec)); - PetscCall(VecRestoreArray(flocal,&f_vec)); - - PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); - PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); - PetscCall(DMRestoreLocalVector(da,&xlocal)); - PetscCall(DMRestoreLocalVector(da,&flocal)); - - return 0; -} - - -PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1./ (info.mx - 1); - hy = 1./ (info.my - 1); - darea = hx * hy; // multiply FD equations by this - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - // print info.ys to screen - // PetscCall(PetscPrintf(PETSC_COMM_WORLD,"info.ys = %d\n",info.xm)); - for (j = info.ys+1; j < info.ys + info.ym-1; j++) { - for (i = info.xs+1; i < info.xs + info.xm-1; i++) { - y_v[j][i] = -x_v[j][i]; - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - - -PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; -// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) -{ - Vec J00X; - Vec J00Y; - Vec J10X; - Vec J10Y; - Vec J11X; - Vec J11Y; - - struct SubMatrixCtx * J00ctx; - struct SubMatrixCtx * J10ctx; - struct SubMatrixCtx * J11ctx; - struct JacobianCtx * jctx; - - PetscFunctionBeginUser; - - PetscCall(MatShellGetContext(J,&(jctx))); - - PetscCall(VecSet(Y,0.0)); - - Mat J00 = jctx->submats[0]; - PetscCall(MatShellGetContext(J00,&(J00ctx))); - PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); - PetscCall(MatMult(J00,J00X,J00Y)); - PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); - - Mat J10 = jctx->submats[2]; - PetscCall(MatShellGetContext(J10,&(J10ctx))); - PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); - PetscCall(MatMult(J10,J10X,J10Y)); - PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); - - Mat J11 = jctx->submats[3]; - PetscCall(MatShellGetContext(J11,&(J11ctx))); - PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); - PetscCall(MatMult(J11,J11X,J11Y)); - PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); - - - PetscFunctionReturn(0); -} - - -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) -{ - PetscFunctionBeginUser; - - PetscInt M; - PetscInt N; - Mat block; - DM dm0; - PetscInt dof; - - struct UserCtx0 * ctx0; - struct JacobianCtx * jctx; - struct SubMatrixCtx * subctx; - - PetscCall(MatShellGetContext(J,&(jctx))); - DM * subdms = jctx->subdms; - - PetscInt nsubmats = nfields*nfields; - PetscCall(PetscCalloc1(nsubmats,submats)); - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); - PetscInt subblockrows = M*N; - PetscInt subblockcols = M*N; - Mat * submat_arr = *submats; - - for (int i = 0; i <= nsubmats - 1; i += 1) - { - PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); - PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); - PetscCall(MatSetType(block,MATSHELL)); - PetscCall(PetscMalloc1(1,&(subctx))); - PetscInt rowidx = i / dof; - PetscInt colidx = (i)%(dof); - subctx->rows = &(irow[rowidx]); - subctx->cols = &(icol[colidx]); - PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); - PetscCall(MatSetDM(block,subdms[rowidx])); - PetscCall(MatShellSetContext(block,subctx)); - PetscCall(MatSetUp(block)); - submat_arr[i] = block; - } - PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); - - PetscFunctionReturn(0); -} - - -PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) -{ - PetscFunctionBeginUser; - - jctx->subdms = subdms; - jctx->fields = fields; - - PetscFunctionReturn(0); -} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8f9cc4d9a8..e4cc05e024 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,6 +1,5 @@ import numpy as np import os -import pytest from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, From 8386992147685e923964601c87e11560e26bce1c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 5 Jun 2025 16:35:49 +0100 Subject: [PATCH 15/53] dsl: Improve Jacobian abstraction --- devito/petsc/iet/routines.py | 57 +++---- devito/petsc/solve.py | 6 +- devito/petsc/types/equation.py | 4 +- devito/petsc/types/types.py | 223 ++++++++++++------------- examples/petsc/random/02_biharmonic.py | 3 +- tests/test_petsc.py | 39 ++--- 6 files changed, 164 insertions(+), 168 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8f580a1b52..8907cfe795 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -111,7 +111,7 @@ def _make_core(self): self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) - if fielddata.initialguess.equations: + if fielddata.initialguess.eqs: self._make_initialguess(fielddata) self._make_user_struct_callback() @@ -515,7 +515,7 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess.equations + initguess = fielddata.initialguess.eqs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -678,15 +678,11 @@ def zero_memory(self): def _make_core(self): injectsolve = self.injectsolve - targets = injectsolve.expr.rhs.fielddata.targets all_fielddata = injectsolve.expr.rhs.fielddata - for t in targets: - row_matvecs = all_fielddata.jacobian.submatrices[t] - arrays = all_fielddata.arrays[t] - for submat, mtvs in row_matvecs.items(): - if mtvs['matvecs']: - self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') + for sm in all_fielddata.jacobian.nonzero_submatrices: + arrays = all_fielddata.arrays[sm.row_target] + self._make_matvec(arrays, sm.matvecs, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) @@ -714,7 +710,7 @@ def _whole_matvec_body(self): jctx = objs['ljacctx'] ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) - nonzero_submats = self.jacobian.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrices zero_y_memory = petsc_call( 'VecSet', [objs['Y'], 0.0] @@ -722,17 +718,17 @@ def _whole_matvec_body(self): calls = () for sm in nonzero_submats: - idx = self.jacobian.submat_to_index[sm] - ctx = sobjs[f'{sm}ctx'] - X = sobjs[f'{sm}X'] - Y = sobjs[f'{sm}Y'] + name = sm.name + ctx = sobjs[f'{name}ctx'] + X = sobjs[f'{name}X'] + Y = sobjs[f'{name}Y'] rows = objs['rows'].base cols = objs['cols'].base - sm_indexed = objs['Submats'].indexed[idx] + sm_indexed = objs['Submats'].indexed[sm.linear_idx] calls += ( - DummyExpr(sobjs[sm], FieldFromPointer(sm_indexed, jctx)), - petsc_call('MatShellGetContext', [sobjs[sm], Byref(ctx)]), + DummyExpr(sobjs[name], FieldFromPointer(sm_indexed, jctx)), + petsc_call('MatShellGetContext', [sobjs[name], Byref(ctx)]), petsc_call( 'VecGetSubVector', [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] @@ -741,7 +737,7 @@ def _whole_matvec_body(self): 'VecGetSubVector', [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] ), - petsc_call('MatMult', [sobjs[sm], X, Y]), + petsc_call('MatMult', [sobjs[name], X, Y]), petsc_call( 'VecRestoreSubVector', [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] @@ -1012,19 +1008,19 @@ def _submat_callback_body(self): upper_bound = objs['nsubmats'] - 1 iteration = Iteration(List(body=iter_body), i, upper_bound) - nonzero_submats = self.jacobian.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrices matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} matmult_op = [ petsc_call( 'MatShellSetOperation', [ - objs['submat_arr'].indexed[self.jacobian.submat_to_index[sb]], + objs['submat_arr'].indexed[sb.linear_idx], 'MATOP_MULT', - MatShellSetOp(matvec_lookup[sb].name, void, void), + MatShellSetOp(matvec_lookup[sb.name].name, void, void), ], ) - for sb in nonzero_submats if sb in matvec_lookup + for sb in nonzero_submats if sb.name in matvec_lookup ] body = [ @@ -1158,22 +1154,23 @@ def _extend_build(self, base_dict): }) jacobian = injectsolve.expr.rhs.fielddata.jacobian - submatrix_keys = jacobian.submatrix_keys + submatrices = jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), fields=objs['ljacctx'].fields, ) - for key in submatrix_keys: - base_dict[key] = Mat(name=key) - base_dict[f'{key}ctx'] = SubMatrixStruct( - name=f'{key}ctx', + for sm in submatrices: + name = sm.name + base_dict[name] = Mat(name=name) + base_dict[f'{name}ctx'] = SubMatrixStruct( + name=f'{name}ctx', fields=objs['subctx'].fields, ) - base_dict[f'{key}X'] = CallbackVec(f'{key}X') - base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') - base_dict[f'{key}F'] = CallbackVec(f'{key}F') + base_dict[f'{name}X'] = CallbackVec(f'{name}X') + base_dict[f'{name}Y'] = CallbackVec(f'{name}Y') + base_dict[f'{name}F'] = CallbackVec(f'{name}F') # Bundle objects/metadata required by the coupled residual callback f_components = [] diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index d2f0f39809..7e4f02df12 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -86,9 +86,7 @@ def linear_solve_args(self): eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, eqns, arrays) field_data = FieldData( @@ -129,8 +127,8 @@ def linear_solve_args(self): ) residual = MixedResidual( - self.target_eqns, arrays, - self.time_mapper, jacobian.target_scaler_mapper + self.target_eqns, arrays, self.time_mapper, + jacobian.target_scaler_mapper ) all_data = MultipleFieldData( diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index ad929c8ed5..6d476c12f7 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -26,7 +26,7 @@ class ZeroRow(EssentialBC): """ Equation used to zero all entries, except the diagonal, of a row in the Jacobian. - + Note: This is only used interally by the compiler, not by users. """ @@ -36,7 +36,7 @@ class ZeroRow(EssentialBC): class ZeroColumn(EssentialBC): """ Equation used to zero the column of the Jacobian. - + Note: This is only used interally by the compiler, not by users. """ diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index a0a9cd2fe5..8bc08d804d 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -235,6 +235,16 @@ def targets(self): class Jacobian: + """ + Represents a Jacobian matrix in a matrix-free form. + + This Jacobian is defined implicitly via matrix-vector products + derived from the symbolic equations provided in `matvecs`. + + It assumes the problem is linear, meaning the Jacobian + corresponds to a constant coefficient matrix and does not + require explicit symbolic differentiation. + """ def __init__(self, target, eqns, arrays, time_mapper): self.target = target self.eqns = eqns @@ -318,140 +328,129 @@ def _scale_bcs(self, matvecs, scdiag): ] +class SubMatrixBlock: + def __init__(self, name, matvecs, scdiag, + row_target, col_target, row_idx, col_idx, linear_idx): + self.name = name + self.matvecs = matvecs + self.scdiag = scdiag + self.row_target = row_target + self.col_target = col_target + self.row_idx = row_idx + self.col_idx = col_idx + self.linear_idx = linear_idx + + def is_diag(self): + return self.row_idx == self.col_idx + + def __repr__(self): + return (f"") + + class MixedJacobian(Jacobian): + """ + Represents a Jacobian for a linear system with a solution vector + composed of multiple fields (targets). + + Defines matrix-vector products for each sub-block, + each with its own generated callback. The matrix may be treated + as monolithic or block-structured in PETSc, but sub-block + callbacks are generated in both cases. + + Assumes a **linear** problem, so this Jacobian corresponds to a + coefficient matrix and does not require differentiation. + + # TODO: pcfieldsplit support for each block + """ def __init__(self, target_eqns, arrays, time_mapper): """ """ self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper - self.submatrices = self._initialize_submatrices() + self._submatrices = [] self._build_blocks(target_eqns) - def _initialize_submatrices(self): - """ - Create a dict of submatrices for each target with metadata. - """ - submatrices = {} - num_targets = len(self.targets) - - for i, target in enumerate(self.targets): - submatrices[target] = {} - for j in range(num_targets): - key = f'J{i}{j}' - submatrices[target][key] = { - 'matvecs': None, - 'derivative_wrt': self.targets[j], - 'index': i * num_targets + j, - 'scdiag': None - } - - return submatrices - - def _build_blocks(self, target_eqns): - for target, eqns in target_eqns.items(): - self._build_block(target, eqns) - - def _build_block(self, target, eqns): - arrays = self.arrays[target] - for submat, mtvs in self.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self._build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - matvecs = self._scale_non_bcs(matvecs, target) - scdiag = self._compute_scdiag(matvecs, arrays) - matvecs = self._scale_bcs(matvecs, scdiag) - - if matvecs: - self.set_submatrix(target, submat, matvecs, scdiag) - @property - def submatrix_keys(self): + def submatrices(self): """ - Return a list of all submatrix keys (e.g., ['J00', 'J01', 'J10', 'J11']). + Return a list of all submatrix blocks. + Each block contains metadata about the matrix-vector products. """ - return [key for submats in self.submatrices.values() for key in submats.keys()] + return self._submatrices @property - def nonzero_submatrix_keys(self): + def no_submatrices(self): """ - Returns a list of submats where 'matvecs' is not None. + Return the number of submatrix blocks. """ - return [ - key - for submats in self.submatrices.values() - for key, value in submats.items() - if value['matvecs'] is not None - ] + return len(self._submatrices) @property - def submat_to_index(self): - """ - Returns a dict mapping submatrix keys to their index. - """ - return { - key: value['index'] - for submats in self.submatrices.values() - for key, value in submats.items() - } - - # CHECK/TEST THIS - def is_diagonal_submatrix(self, key): - """ - Return True if the given key corresponds to a diagonal - submatrix (e.g., 'J00', 'J11'), else False. - """ - for i, t in enumerate(self.targets): - diag_key = f'J{i}{i}' - if key == diag_key and diag_key in self.submatrices[t]: - return True - return False - - def set_submatrix(self, field, key, matvecs, scdiag): - """ - Set the matrix-vector equations for a submatrix. - - Parameters - ---------- - field : Function - The target field that the submatrix operates on. - key: str - The identifier for the submatrix (e.g., 'J00', 'J01'). - matvecs: list of Eq - The matrix-vector equations forming the submatrix. - """ - if field in self.submatrices and key in self.submatrices[field]: - self.submatrices[field][key]['matvecs'] = matvecs - self.submatrices[field][key]['scdiag'] = scdiag - else: - raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') - - def get_submatrix(self, field, key): - """ - Retrieve a specific submatrix. - """ - return self.submatrices.get(field, {}).get(key, None) + def nonzero_submatrices(self): + """Return SubMatrixBlock objects that have non-empty matvecs.""" + return [submat for submat in self.submatrices if submat.matvecs] @property def target_scaler_mapper(self): """ - Return a mapping from each target to its diagonal submatrix's scaler. + Map each row_target index to its diagonal submatrix's scaler. """ mapper = {} - for target, submats in self.submatrices.items(): - for key, value in submats.items(): - if self.is_diagonal_submatrix(key): - mapper[target] = value.get('scdiag') - break + for sm in self.submatrices: + if sm.row_idx == sm.col_idx: + mapper[sm.row_target] = sm.scdiag return mapper + def _build_blocks(self, target_eqns): + """ + Build all SubMatrixBlock objects for the Jacobian. + """ + for i, row_target in enumerate(self.targets): + eqns = target_eqns[row_target] + arrays = self.arrays[row_target] + for j, col_target in enumerate(self.targets): + matvecs = [ + e for eq in eqns for e in + self._build_matvec_eq(eq, col_target, arrays) + ] + matvecs = [m for m in matvecs if m is not None] + # Sort to put EssentialBC first if any + matvecs = tuple( + sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) + ) + matvecs = self._scale_non_bcs(matvecs, row_target) + scdiag = self._compute_scdiag(matvecs, arrays) + matvecs = self._scale_bcs(matvecs, scdiag) + + name = f'J{i}{j}' + block = SubMatrixBlock( + name=name, + matvecs=matvecs, + scdiag=scdiag, + row_target=row_target, + col_target=col_target, + row_idx=i, + col_idx=j, + linear_idx=i * len(self.targets) + j + ) + self._submatrices.append(block) + + def get_submatrix(self, row_idx, col_idx): + """ + Return the SubMatrixBlock at (row_idx, col_idx), or None if not found. + """ + for sm in self.submatrices: + if sm.row_idx == row_idx and sm.col_idx == col_idx: + return sm + return None + def __repr__(self): - # TODO: edit - return str(self.submatrices) + summary = ', '.join( + f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})" + for sm in self.submatrices + ) + return f"" class Residual: @@ -486,7 +485,7 @@ def _build_equations(self): for eq in self.eqns: b, F_target, _, targets = separate_eqn(eq, self.target) funcs.extend(self._make_F_target(eq, F_target, targets)) - # TODO: if b is zero then don't need a rhs vector+callback + # TODO: If b is zero then don't need a rhs vector+callback rhs.extend(self._make_b(eq, b)) self._formfuncs = [self._scale_bcs(eq) for eq in funcs] @@ -596,18 +595,18 @@ def __init__(self, target, eqns, arrays): self.eqns = as_tuple(eqns) self.arrays = arrays self._build_equations() - + @property - def equations(self): + def eqs(self): """ """ - return self._equations + return self._eqs def _build_equations(self): """ Return a list of initial guess equations. """ - self._equations = [ + self._eqs = [ eq for eq in (self._make_initial_guess(e) for e in self.eqns) if eq is not None diff --git a/examples/petsc/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py index 973b16499e..214e38afda 100644 --- a/examples/petsc/random/02_biharmonic.py +++ b/examples/petsc/random/02_biharmonic.py @@ -129,7 +129,8 @@ def f_fcn(x, y): with switchconfig(language='petsc'): op = Operator(petsc) - op.apply() + # op.apply() + print(op.ccode) u_exact.data[:] = u_exact_fcn(X, Y) lap_u.data[:] = lap_u_exact_fcn(X, Y) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e4cc05e024..d0f4bdd087 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -817,36 +817,37 @@ def test_submatrices(self): jacobian = petsc[0].rhs.fielddata.jacobian - j00 = jacobian.get_submatrix(e, 'J00') - j01 = jacobian.get_submatrix(e, 'J01') - j10 = jacobian.get_submatrix(g, 'J10') - j11 = jacobian.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(0, 0) + j01 = jacobian.get_submatrix(0, 1) + j10 = jacobian.get_submatrix(1, 0) + j11 = jacobian.get_submatrix(1, 1) # Check the number of submatrices - assert len(jacobian.submatrix_keys) == 4 - assert str(jacobian.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + assert jacobian.no_submatrices == 4 # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 - assert jacobian.nonzero_submatrix_keys == ['J00', 'J11'] - assert jacobian.get_submatrix(e, 'J01')['matvecs'] is None - assert jacobian.get_submatrix(g, 'J10')['matvecs'] is None - - j00 = jacobian.get_submatrix(e, 'J00') - j11 = jacobian.get_submatrix(g, 'J11') + nonzero_submats = jacobian.nonzero_submatrices + assert len(nonzero_submats) == 2 + assert j00 in nonzero_submats + assert j11 in nonzero_submats + assert j01 not in nonzero_submats + assert j10 not in nonzero_submats + assert not j01.matvecs + assert not j10.matvecs # Compatible scaling to reduce condition number of jacobian - assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + assert str(j00.matvecs[0]) == 'Eq(y_e(x, y),' \ + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' - assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ + assert str(j11.matvecs[0]) == 'Eq(y_g(x, y),' \ + ' h_x*h_y*(Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2))))' - # Check the derivative wrt fields - assert j00['derivative_wrt'] == e - assert j01['derivative_wrt'] == g - assert j10['derivative_wrt'] == e - assert j11['derivative_wrt'] == g + # Check the col_targets + assert j00.col_target == e + assert j01.col_target == g + assert j10.col_target == e + assert j11.col_target == g @skipif('petsc') def test_residual_bundle(self): From fb207810eb23a24c4784fa20f97c89d77a178ba3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 5 Jun 2025 18:01:40 +0100 Subject: [PATCH 16/53] examples: Add pressure norm check in stokes example --- devito/petsc/iet/routines.py | 81 ++++---- devito/petsc/solve.py | 9 +- devito/petsc/types/array.py | 17 +- devito/petsc/types/types.py | 57 +++--- devito/types/array.py | 6 +- examples/petsc/Poisson/01_poisson.py | 1 + examples/petsc/cfd/01_navierstokes.py | 34 ++-- tests/test_petsc.py | 256 +++++++++++++++++++++----- 8 files changed, 333 insertions(+), 128 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8907cfe795..2fbdafe1ef 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -106,22 +106,30 @@ def zero_memory(self): in the callback.""" return True + @property + def fielddata(self): + return self.injectsolve.expr.rhs.fielddata + + @property + def arrays(self): + return self.fielddata.arrays + def _make_core(self): - fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) - self._make_formfunc(fielddata) - self._make_formrhs(fielddata) - if fielddata.initialguess.eqs: - self._make_initialguess(fielddata) + self._make_matvec(self.fielddata.jacobian) + self._make_formfunc() + self._make_formrhs() + if self.fielddata.initialguess.eqs: + self._make_initialguess() self._make_user_struct_callback() - def _make_matvec(self, arrays, matvecs, prefix='MatMult'): + def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation + matvecs = jacobian.matvecs irs_matvec, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), - arrays) + jacobian) objs = self.objs cb = PETScCallable( @@ -133,7 +141,7 @@ def _make_matvec(self, arrays, matvecs, prefix='MatMult'): self._matvecs.append(cb) self._efuncs[cb.name] = cb - def _create_matvec_body(self, body, arrays): + def _create_matvec_body(self, body, jacobian): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -142,8 +150,8 @@ def _create_matvec_body(self, body, arrays): ctx = objs['dummyctx'] xlocal = objs['xloc'] ylocal = objs['yloc'] - y_matvec = arrays['y'] - x_matvec = arrays['x'] + y_matvec = self.arrays[jacobian.row_target]['y'] + x_matvec = self.arrays[jacobian.col_target]['x'] body = self.timedep.uxreplace_time(body) @@ -265,15 +273,15 @@ def _create_matvec_body(self, body, arrays): self._struct_params.extend(fields) return matvec_body - def _make_formfunc(self, fielddata): - formfuncs = fielddata.residual.formfuncs + def _make_formfunc(self): + formfuncs = self.fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._create_formfunc_body( - List(body=irs_formfunc.uiet.body), fielddata + List(body=irs_formfunc.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -285,10 +293,11 @@ def _make_formfunc(self, fielddata): self._formfuncs.append(cb) self._efuncs[cb.name] = cb - def _create_formfunc_body(self, body, fielddata): + def _create_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -298,8 +307,8 @@ def _create_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - f_formfunc = fielddata.arrays['f'] - x_formfunc = fielddata.arrays['x'] + f_formfunc = self.fielddata.arrays[target]['f'] + x_formfunc = self.fielddata.arrays[target]['x'] dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -402,8 +411,8 @@ def _create_formfunc_body(self, body, fielddata): return Uxreplace(subs).visit(formfunc_body) - def _make_formrhs(self, fielddata): - formrhs = fielddata.residual.formrhs + def _make_formrhs(self): + formrhs = self.fielddata.residual.formrhs sobjs = self.solver_objs # Compile formrhs `eqns` into an IET via recursive compilation @@ -412,7 +421,7 @@ def _make_formrhs(self, fielddata): concretize_mapper=self.concretize_mapper ) body_formrhs = self._create_form_rhs_body( - List(body=irs_formrhs.uiet.body), fielddata + List(body=irs_formrhs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -424,10 +433,11 @@ def _make_formrhs(self, fielddata): self._formrhs.append(cb) self._efuncs[cb.name] = cb - def _create_form_rhs_body(self, body, fielddata): + def _create_form_rhs_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -446,7 +456,7 @@ def _create_form_rhs_body(self, body, fielddata): sobjs['blocal'] ]) - b_arr = fielddata.arrays['b'] + b_arr = self.fielddata.arrays[target]['b'] vec_get_array = petsc_call( 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] @@ -514,8 +524,8 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) - def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess.eqs + def _make_initialguess(self): + initguess = self.fielddata.initialguess.eqs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -524,7 +534,7 @@ def _make_initialguess(self, fielddata): concretize_mapper=self.concretize_mapper ) body_init_guess = self._create_initial_guess_body( - List(body=irs.uiet.body), fielddata + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -536,15 +546,16 @@ def _make_initialguess(self, fielddata): self._initialguesses.append(cb) self._efuncs[cb.name] = cb - def _create_initial_guess_body(self, body, fielddata): + def _create_initial_guess_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - x_arr = fielddata.arrays['x'] + x_arr = self.fielddata.arrays[target]['x'] vec_get_array = petsc_call( 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] @@ -681,8 +692,8 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for sm in all_fielddata.jacobian.nonzero_submatrices: - arrays = all_fielddata.arrays[sm.row_target] - self._make_matvec(arrays, sm.matvecs, prefix=f'{sm.name}_MatMult') + # arrays = all_fielddata.arrays[sm.row_target] + self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) @@ -1044,18 +1055,18 @@ def _submat_callback_body(self): ) def residual_bundle(self, body, bundles): - mapper1 = bundles['bundle_mapper'] + mapper = bundles['bundle_mapper'] indexeds = FindSymbols('indexeds').visit(body) + subs = {} - subss = {} for i in indexeds: - if i.base in mapper1: - bundle = mapper1[i.base] + if i.base in mapper: + bundle = mapper[i.base] index = bundles['target_indices'][i.function.target] index = (index,)+i.indices - subss[i] = bundle.__getitem__(index) + subs[i] = bundle.__getitem__(index) - body = Uxreplace(subss).visit(body) + body = Uxreplace(subs).visit(body) return body diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 7e4f02df12..8adb9cda1e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -81,7 +81,7 @@ def linear_solve_args(self): funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays(target) + arrays = self.generate_arrays_combined(target) eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) @@ -108,6 +108,9 @@ def generate_arrays(self, target): for p in prefixes } + def generate_arrays_combined(self, *targets): + return {target: self.generate_arrays(target) for target in targets} + class InjectMixedSolve(InjectSolve): @@ -140,8 +143,8 @@ def linear_solve_args(self): return coupled_targets[0], tuple(funcs), all_data - def generate_arrays_combined(self, *targets): - return {target: self.generate_arrays(target) for target in targets} + # def generate_arrays_combined(self, *targets): + # return {target: self.generate_arrays(target) for target in targets} def generate_time_mapper(funcs): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 8b57aca44c..c7a09f8265 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,3 +1,5 @@ +from sympy import Expr + from functools import cached_property from ctypes import POINTER, Structure @@ -164,7 +166,7 @@ def __getitem__(self, index): elif len(index) == self.ndim + 1: component_index, indices = index[0], index[1:] names = tuple(i.target.name for i in self.components) - return ComponentAccess( + return PetscComponentAccess( self.indexed[indices], component_index, component_names=names @@ -178,6 +180,19 @@ def pname(self): return self._pname +class PetscComponentAccess(ComponentAccess): + def __new__(cls, arg, index=0, component_names=None, **kwargs): + if not arg.is_Indexed: + raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) + names = component_names or cls._default_component_names + + obj = Expr.__new__(cls, arg) + obj._index = index + obj._component_names = names + + return obj + + class AoSIndexedData(IndexedData): @property def dtype(self): diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 8bc08d804d..21433b5944 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -260,6 +260,14 @@ def matvecs(self): def scdiag(self): return self._scdiag + @property + def row_target(self): + return self.target + + @property + def col_target(self): + return self.target + def _build_matvecs(self): matvecs = [ e for eq in self.eqns for e in @@ -275,30 +283,35 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_eq(self, eq, target=None, arrays=None): - target = target or self.target - arrays = arrays or self.arrays + def _build_matvec_eq(self, eq, col_target=None, row_target=None): + col_target = col_target or self.target + row_target = row_target or self.target - b, F_target, _, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, col_target) if F_target: - return self._make_matvec(eq, F_target, targets, arrays) + return self._make_matvec( + eq, F_target, targets, col_target, row_target + ) return (None,) - def _make_matvec(self, eq, F_target, targets, arrays): + def _make_matvec(self, eq, F_target, targets, col_target, row_target): + y = self.arrays[row_target]['y'] + x = self.arrays[col_target]['x'] + if isinstance(eq, EssentialBC): # NOTE: Essential BCs are trivial equations in the solver. # See `EssentialBC` for more details. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0., subdomain=eq.subdomain) + zero_row = ZeroRow(y, x, subdomain=eq.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=eq.subdomain) return (zero_row, zero_column) else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + return as_tuple(Eq(y, rhs, subdomain=eq.subdomain)) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target + # TODO: make this a property of the class so don't need target as an arg vol = target.grid.symbolic_volume_cell return [ @@ -306,13 +319,13 @@ def _scale_non_bcs(self, matvecs, target=None): for m in matvecs ] - def _compute_scdiag(self, matvecs, arrays=None): + def _compute_scdiag(self, matvecs, col_target=None): """ """ - arrays = arrays or self.arrays + x = self.arrays[col_target or self.target]['x'] centres = { - centre_stencil(m.rhs, arrays['x'], as_coeff=True) + centre_stencil(m.rhs, x, as_coeff=True) for m in matvecs if not isinstance(m, EssentialBC) } # add comments @@ -329,8 +342,8 @@ def _scale_bcs(self, matvecs, scdiag): class SubMatrixBlock: - def __init__(self, name, matvecs, scdiag, - row_target, col_target, row_idx, col_idx, linear_idx): + def __init__(self, name, matvecs, scdiag, row_target, + col_target, row_idx, col_idx, linear_idx): self.name = name self.matvecs = matvecs self.scdiag = scdiag @@ -408,11 +421,10 @@ def _build_blocks(self, target_eqns): """ for i, row_target in enumerate(self.targets): eqns = target_eqns[row_target] - arrays = self.arrays[row_target] for j, col_target in enumerate(self.targets): matvecs = [ e for eq in eqns for e in - self._build_matvec_eq(eq, col_target, arrays) + self._build_matvec_eq(eq, col_target, row_target) ] matvecs = [m for m in matvecs if m is not None] # Sort to put EssentialBC first if any @@ -420,7 +432,7 @@ def _build_blocks(self, target_eqns): sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) ) matvecs = self._scale_non_bcs(matvecs, row_target) - scdiag = self._compute_scdiag(matvecs, arrays) + scdiag = self._compute_scdiag(matvecs, col_target) matvecs = self._scale_bcs(matvecs, scdiag) name = f'J{i}{j}' @@ -492,7 +504,7 @@ def _build_equations(self): self._formrhs = rhs def _make_F_target(self, eq, F_target, targets): - arrays = self.arrays + arrays = self.arrays[self.target] volume = self.target.grid.symbolic_volume_cell if isinstance(eq, EssentialBC): # The initial guess satisfies the essential BCs, so this term is zero. @@ -511,9 +523,10 @@ def _make_F_target(self, eq, F_target, targets): return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) def _make_b(self, eq, b): + b_arr = self.arrays[self.target]['b'] rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + return as_tuple(Eq(b_arr, rhs, subdomain=eq.subdomain)) def _scale_bcs(self, eq, scdiag=None): """ @@ -616,7 +629,7 @@ def _make_initial_guess(self, eq): if isinstance(eq, EssentialBC): assert eq.lhs == self.target return Eq( - self.arrays['x'], eq.rhs, + self.arrays[self.target]['x'], eq.rhs, subdomain=eq.subdomain ) else: diff --git a/devito/types/array.py b/devito/types/array.py index 3c3cedfc88..b3367f44a9 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -538,7 +538,7 @@ class ComponentAccess(Expr, Pickable): _default_component_names = ('x', 'y', 'z', 'w') __rargs__ = ('arg',) - __rkwargs__ = ('index',) + __rkwargs__ = ('index', 'component_names') def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: @@ -579,6 +579,10 @@ def arg(self): def index(self): return self._index + @property + def component_names(self): + return self._component_names + @property def sindex(self): return self._component_names[self.index] diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 4981566af6..6f24f74f64 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,6 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 01e96e1afe..7797321045 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -1,8 +1,8 @@ import os import numpy as np -from devito import (Grid, TimeFunction, Function, Constant, Eq, - Operator, norm, SubDomain, switchconfig, configuration) +from devito import (Grid, TimeFunction, Constant, Eq, + Operator, SubDomain, switchconfig, configuration) from devito.symbolics import retrieve_functions, INT from devito.petsc import PETScSolve, EssentialBC @@ -232,11 +232,9 @@ def neumann_right(eq, subdomain): u1 = TimeFunction(name='u1', grid=grid, space_order=2, dtype=np.float64) v1 = TimeFunction(name='v1', grid=grid, space_order=2, dtype=np.float64) -pn1 = Function(name='pn1', grid=grid, space_order=2, dtype=np.float64) +pn1 = TimeFunction(name='pn1', grid=grid, space_order=2, dtype=np.float64) -pn1.data[:] = 0. - -eq_pn1 = Eq(pn1.laplace, rho*(1./dt*(u1.forward.dxc+v1.forward.dyc)), +eq_pn1 = Eq(pn1.forward.laplace, rho*(1./dt*(u1.forward.dxc+v1.forward.dyc)), subdomain=grid.interior) @@ -244,21 +242,21 @@ def neumann_right(eq, subdomain): bc_pn1 += [neumann_bottom(eq_pn1, sub2)] bc_pn1 += [neumann_left(eq_pn1, sub3)] bc_pn1 += [neumann_right(eq_pn1, sub4)] -bc_pn1 += [EssentialBC(pn1, 0., subdomain=sub5)] +bc_pn1 += [EssentialBC(pn1.forward, 0., subdomain=sub5)] bc_pn1 += [neumann_right(neumann_bottom(eq_pn1, sub6), sub6)] bc_pn1 += [neumann_left(neumann_top(eq_pn1, sub7), sub7)] bc_pn1 += [neumann_right(neumann_top(eq_pn1, sub8), sub8)] -eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1) +eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1.forward) -eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace) -eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace) +eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace, subdomain=grid.interior) +eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace, subdomain=grid.interior) -update_u = Eq(u1.forward, u1.forward - (dt/rho)*(pn1.dxc), +update_u = Eq(u1.forward, u1.forward - (dt/rho)*(pn1.forward.dxc), subdomain=grid.interior) -update_v = Eq(v1.forward, v1.forward - (dt/rho)*(pn1.dyc), +update_v = Eq(v1.forward, v1.forward - (dt/rho)*(pn1.forward.dyc), subdomain=grid.interior) # TODO: Can drop due to initial guess CB @@ -296,12 +294,6 @@ def neumann_right(eq, subdomain): op = Operator(exprs) op.apply(time_m=0, time_M=ns-1, dt=dt) -u1_norm = norm(u1) -v1_norm = norm(v1) -p1_norm = norm(pn1) - - -# TODO: change these norm checks to array checks (use paper) -assert np.isclose(u1_norm, 13.966067703420883, atol=0, rtol=1e-7) -assert np.isclose(v1_norm, 7.9575677674738285, atol=0, rtol=1e-7) -assert np.isclose(p1_norm, 36.46263134701362, atol=0, rtol=1e-7) +# Pressure norm check +tol = 1e-3 +assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d0f4bdd087..e0707d7b1e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,3 +1,5 @@ +import pytest + import numpy as np import os @@ -615,19 +617,38 @@ class TestCoupledLinear: # TODO: Add more comprehensive tests for fully coupled problems. # TODO: Add subdomain tests, time loop, multiple coupled etc. + @pytest.mark.parametrize('eq1, eq2, so', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '6'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '2'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '4'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '6'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '2'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '4'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '6'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '2'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '4'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '6'), + ]) @skipif('petsc') - def test_coupled_vs_non_coupled(self): + def test_coupled_vs_non_coupled(self, eq1, eq2, so): + """ + Test that solving multiple **uncoupled** equations separately + vs. together with `PETScSolve` yields the same result. + This test is non time-dependent. + """ grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2) + functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions f.data[:] = 5. h.data[:] = 5. - eq1 = Eq(e.laplace, f) - eq2 = Eq(g.laplace, h) + eq1 = eval(eq1) + eq2 = eval(eq2) # Non-coupled petsc1 = PETScSolve(eq1, target=e) @@ -645,8 +666,6 @@ def test_coupled_vs_non_coupled(self): g.data[:] = 0 # Coupled - # TODO: Need more friendly API for coupled - just - # using a dict for now petsc3 = PETScSolve({e: [eq1], g: [eq2]}) with switchconfig(language='petsc'): @@ -658,8 +677,8 @@ def test_coupled_vs_non_coupled(self): print('enorm1:', enorm1) print('enorm2:', enorm2) - assert np.isclose(enorm1, enorm2, rtol=1e-16) - assert np.isclose(gnorm1, gnorm2, rtol=1e-16) + assert np.isclose(enorm1, enorm2, atol=1e-14) + assert np.isclose(gnorm1, gnorm2, atol=1e-14) callbacks1 = [meta_call.root for meta_call in op1._func_table.values()] callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] @@ -724,49 +743,33 @@ def test_coupled_structs(self): assert 'struct Field0\n{' not in ccode assert 'struct Field0\n{' in hcode + @pytest.mark.parametrize('n_fields', [2, 3, 4, 5, 6]) @skipif('petsc') - def test_coupled_frees(self): + def test_coupled_frees(self, n_fields): grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2) - for n in ['e', 'f', 'g', 'h']] - e, f, g, h = functions - - eq1 = Eq(e.laplace, h) - eq2 = Eq(f.laplace, h) - eq3 = Eq(g.laplace, h) + functions = [Function(name=f'u{i}', grid=grid, space_order=2) + for i in range(n_fields + 1)] + *solved_funcs, h = functions - petsc1 = PETScSolve({e: [eq1], f: [eq2]}) - petsc2 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + equations = [Eq(func.laplace, h) for func in solved_funcs] + petsc = PETScSolve({func: [eq] for func, eq in zip(solved_funcs, equations)}) with switchconfig(language='petsc'): - op1 = Operator(petsc1, opt='noop') - op2 = Operator(petsc2, opt='noop') + op = Operator(petsc, opt='noop') + + frees = op.body.frees - frees1 = op1.body.frees - frees2 = op2.body.frees - - # Check solver with two fields - # IS destroys - assert str(frees1[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' - assert str(frees1[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' - assert str(frees1[2]) == 'PetscCall(PetscFree(fields0));' - # Sub DM destroys - assert str(frees1[3]) == 'PetscCall(DMDestroy(&(subdms0[0])));' - assert str(frees1[4]) == 'PetscCall(DMDestroy(&(subdms0[1])));' - assert str(frees1[5]) == 'PetscCall(PetscFree(subdms0));' - - # Check solver with three fields - # IS destroys - assert str(frees2[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' - assert str(frees2[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' - assert str(frees2[2]) == 'PetscCall(ISDestroy(&(fields0[2])));' - assert str(frees2[3]) == 'PetscCall(PetscFree(fields0));' - # Sub DM destroys - assert str(frees2[4]) == 'PetscCall(DMDestroy(&(subdms0[0])));' - assert str(frees2[5]) == 'PetscCall(DMDestroy(&(subdms0[1])));' - assert str(frees2[6]) == 'PetscCall(DMDestroy(&(subdms0[2])));' - assert str(frees2[7]) == 'PetscCall(PetscFree(subdms0));' + # IS Destroy calls + for i in range(n_fields): + assert str(frees[i]) == f'PetscCall(ISDestroy(&(fields0[{i}])));' + assert str(frees[n_fields]) == 'PetscCall(PetscFree(fields0));' + + # DM Destroy calls + for i in range(n_fields): + assert str(frees[n_fields + 1 + i]) == \ + f'PetscCall(DMDestroy(&(subdms0[{i}])));' + assert str(frees[n_fields*2 + 1]) == 'PetscCall(PetscFree(subdms0));' @skipif('petsc') def test_dmda_dofs(self): @@ -803,7 +806,7 @@ def test_dmda_dofs(self): in str(op3) @skipif('petsc') - def test_submatrices(self): + def test_mixed_jacobian(self): grid = Grid(shape=(11, 11), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=2) @@ -849,6 +852,169 @@ def test_submatrices(self): assert j10.col_target == e assert j11.col_target == g + + # TODO: FIX THIS LOGIC - THE OUTPUT FOR J01_MATVEC IS Eq(y_e(x, y), -h_x*h_y*x_e(x, y)) BUT IT SHOULD + # BE Eq(y_e(x, y), -h_x*h_y*x_G(x, y)) + + # @pytest.mark.parametrize('eq1, eq2, so, j01_matvec, j10_matvec', [ + # ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + # ]) + # @skipif('petsc') + # def test_coupling(self, eq1, eq2, so, scale): + # """ + # Test linear coupling between fields, where the off-diagonal + # Jacobian submatrices are nonzero. + # """ + # grid = Grid(shape=(9,9), dtype=np.float64) + + # functions = [Function(name=n, grid=grid, space_order=eval(so)) + # for n in ['e', 'f', 'g', 'h']] + # e, f, g, h = functions + + # eq1 = eval(eq1) + # eq2 = eval(eq2) + + # petsc = PETScSolve({e: [eq1], g: [eq2]}) + + # jacobian = petsc[0].rhs.fielddata.jacobian + + # j01 = jacobian.get_submatrix(0, 1) + # j10 = jacobian.get_submatrix(1, 0) + + # assert len(j01.matvecs) == 1 + # assert len(j10.matvecs) == 1 + + # assert j01.col_target == g + # assert j10.col_target == e + + # assert str(j01.matvecs[0]) == j01_matvec + # assert str(j10.matvecs[0]) == j10_matvec + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', '-2.0*h_x/h_x**2'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', '-2.5*h_x/h_x**2'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*(1 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*(5.0 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*(5.0 - 2.5/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*(1 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*(1 - 5.0/h_x**2)'), + ]) + @skipif('petsc') + def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 1D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), + + ]) + @skipif('petsc') + def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 2D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), + + ]) + @skipif('petsc') + def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 3D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,9,9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + # test coupled residual callback - check the .dot for each field etc + + @skipif('petsc') def test_residual_bundle(self): grid = Grid(shape=(11, 11), dtype=np.float64) From 5c1b686938507d73c43392316fcc15890d2e5133 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 8 Jun 2025 12:46:10 +0100 Subject: [PATCH 17/53] tests: Add more petsc tests for linear coupled --- .github/workflows/examples-mpi.yml | 4 +- .github/workflows/examples.yml | 4 +- .github/workflows/flake8.yml | 4 +- .github/workflows/pytest-core-mpi.yml | 4 +- .github/workflows/pytest-core-nompi.yml | 4 +- .github/workflows/pytest-petsc.yml | 4 +- .github/workflows/tutorials.yml | 4 +- devito/passes/iet/engine.py | 9 +- devito/petsc/iet/routines.py | 24 +-- devito/petsc/solve.py | 11 +- devito/petsc/types/equation.py | 10 +- devito/petsc/types/types.py | 10 +- devito/tools/utils.py | 2 +- examples/petsc/Poisson/01_poisson.py | 3 - examples/petsc/Poisson/02_laplace.py | 2 - examples/petsc/petsc_test.py | 2 +- examples/petsc/random/02_biharmonic.py | 4 +- tests/test_petsc.py | 213 ++++++++++++++++-------- 18 files changed, 190 insertions(+), 128 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index e44a03410d..47a8629c05 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,11 +17,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 9b6d5a9ade..4ca681131f 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 40b57d7706..e207557f10 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index fbee3367f1..07ca0534d6 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 23b5b20344..b00115ba7f 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 9c035a86e7..2bf55b5049 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: pytest: diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index c6e9c2adba..75b7ef3982 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: tutorials: diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index ad0ebefd1e..f141f436d3 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -1,6 +1,7 @@ from collections import OrderedDict, defaultdict from functools import partial, singledispatch, wraps +from devito.finite_differences.differentiable import Differentiable from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable, FindNodes, FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction, FixedArgsCallable, Uxreplace, @@ -10,13 +11,13 @@ from devito.passes import needs_transfer from devito.symbolics import FieldFromComposite, FieldFromPointer from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass -from devito.types import (Bundle, CompositeObject, Lock, IncrDimension, +from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) -from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension +from devito.types.array import ArrayBasic __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -414,7 +415,8 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, ArrayBasic, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, Array, Differentiable, DiscreteFunction, + AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -449,6 +451,7 @@ def _(i, mapper, sregistry): }) +@abstract_object.register(Array) @abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 2fbdafe1ef..4e48a571f2 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -688,15 +688,11 @@ def zero_memory(self): return False def _make_core(self): - injectsolve = self.injectsolve - all_fielddata = injectsolve.expr.rhs.fielddata - - for sm in all_fielddata.jacobian.nonzero_submatrices: - # arrays = all_fielddata.arrays[sm.row_target] + for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() - self._make_whole_formfunc(all_fielddata) + self._make_whole_formfunc() self._make_user_struct_callback() self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -764,15 +760,15 @@ def _whole_matvec_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def _make_whole_formfunc(self, fielddata): - formfuncs = fielddata.residual.formfuncs + def _make_whole_formfunc(self): + formfuncs = self.fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body), - fielddata) + body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) + objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), @@ -783,7 +779,7 @@ def _make_whole_formfunc(self, fielddata): self._main_formfunc_callback = cb self._efuncs[cb.name] = cb - def _whole_formfunc_body(self, body, fielddata): + def _whole_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -796,7 +792,7 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - # Process body for residual callback, including generating bundles etc + # Process body with bundles for residual callback bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] @@ -1143,7 +1139,6 @@ def _extend_build(self, base_dict): class CoupledObjectBuilder(BaseObjectBuilder): def _extend_build(self, base_dict): - injectsolve = self.injectsolve sreg = self.sregistry objs = self.objs targets = self.fielddata.targets @@ -1164,8 +1159,7 @@ def _extend_build(self, base_dict): dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - jacobian = injectsolve.expr.rhs.fielddata.jacobian - submatrices = jacobian.nonzero_submatrices + submatrices = self.fielddata.jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 8adb9cda1e..61d0c187e6 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -27,12 +27,12 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq The targets and symbolic equations defining the system to be solved. - - **Single-field problem**: + - Single-field problem: Pass a single Eq or list of Eq, and specify `target` separately: PETScSolve(Eq1, target) PETScSolve([Eq1, Eq2], target) - - **Multi-field (mixed) problem**: + - Multi-field (mixed) problem: Pass a dictionary mapping each target field to its Eq(s): PETScSolve({u: Eq1, v: Eq2}) PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) @@ -47,8 +47,8 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Returns ------- Eq - A symbolic equation that wraps the system, solver metadata, - and boundary conditions. This can be passed directly to a Devito Operator. + A symbolic equation that wraps the linear solver. + This can be passed directly to a Devito Operator. """ if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() @@ -143,9 +143,6 @@ def linear_solve_args(self): return coupled_targets[0], tuple(funcs), all_data - # def generate_arrays_combined(self, *targets): - # return {target: self.generate_arrays(target) for target in targets} - def generate_time_mapper(funcs): """ diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index 6d476c12f7..e819b48a22 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -27,8 +27,9 @@ class ZeroRow(EssentialBC): Equation used to zero all entries, except the diagonal, of a row in the Jacobian. - Note: - This is only used interally by the compiler, not by users. + Warnings + -------- + Created and managed directly by Devito, not by users. """ pass @@ -37,7 +38,8 @@ class ZeroColumn(EssentialBC): """ Equation used to zero the column of the Jacobian. - Note: - This is only used interally by the compiler, not by users. + Warnings + -------- + Created and managed directly by Devito, not by users. """ pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 21433b5944..44d43d2f97 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -236,7 +236,7 @@ def targets(self): class Jacobian: """ - Represents a Jacobian matrix in a matrix-free form. + Represents a Jacobian matrix. This Jacobian is defined implicitly via matrix-vector products derived from the symbolic equations provided in `matvecs`. @@ -311,7 +311,6 @@ def _make_matvec(self, eq, F_target, targets, col_target, row_target): def _scale_non_bcs(self, matvecs, target=None): target = target or self.target - # TODO: make this a property of the class so don't need target as an arg vol = target.grid.symbolic_volume_cell return [ @@ -328,7 +327,6 @@ def _compute_scdiag(self, matvecs, col_target=None): centre_stencil(m.rhs, x, as_coeff=True) for m in matvecs if not isinstance(m, EssentialBC) } - # add comments return centres.pop() if len(centres) == 1 else 1.0 def _scale_bcs(self, matvecs, scdiag): @@ -370,7 +368,7 @@ class MixedJacobian(Jacobian): as monolithic or block-structured in PETSc, but sub-block callbacks are generated in both cases. - Assumes a **linear** problem, so this Jacobian corresponds to a + Assumes a linear problem, so this Jacobian corresponds to a coefficient matrix and does not require differentiation. # TODO: pcfieldsplit support for each block @@ -407,7 +405,8 @@ def nonzero_submatrices(self): @property def target_scaler_mapper(self): """ - Map each row_target index to its diagonal submatrix's scaler. + Map each row target to the scdiag of its corresponding + diagonal subblock. """ mapper = {} for sm in self.submatrices: @@ -601,7 +600,6 @@ def _build_function_eq(self, eq, target): class InitialGuess: """ Enforce initial guess to satisfy essential BCs. - # TODO: Extend this to "coupled". """ def __init__(self, target, eqns, arrays): self.target = target diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 1285ffabec..0a28de16a8 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -337,7 +337,7 @@ def sorted_priority(items, priority): """ def key(i): - for cls in sorted(priority, key=priority.get): + for cls in sorted(priority, key=priority.get, reverse=True): if isinstance(i, cls): v = priority[cls] break diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6f24f74f64..7ed32e8bbd 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -16,8 +16,6 @@ # Subdomains to implement BCs -# NOTE: For essential BCs, we ensure the SubDomains do not overlap - class SubTop(SubDomain): name = 'subtop' @@ -66,7 +64,6 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py index c84855137f..9df68f9ab9 100644 --- a/examples/petsc/Poisson/02_laplace.py +++ b/examples/petsc/Poisson/02_laplace.py @@ -23,8 +23,6 @@ # Subdomains to implement BCs -# NOTE: For essential BCs, we ensure the SubDomains do not overlap - class SubTop(SubDomain): name = 'subtop' diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 12d245480e..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -20,7 +20,7 @@ v.data[:] = 5.0 -eq = Eq(0., u.laplace, subdomain=grid.interior) +eq = Eq(v, u.laplace, subdomain=grid.interior) petsc = PETScSolve([eq], u) diff --git a/examples/petsc/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py index 214e38afda..f08ffc07de 100644 --- a/examples/petsc/random/02_biharmonic.py +++ b/examples/petsc/random/02_biharmonic.py @@ -85,7 +85,6 @@ def f_fcn(x, y): Ly = np.float64(1.) n_values = [33, 53, 73, 93, 113] -n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -129,8 +128,7 @@ def f_fcn(x, y): with switchconfig(language='petsc'): op = Operator(petsc) - # op.apply() - print(op.ccode) + op.apply() u_exact.data[:] = u_exact_fcn(X, Y) lap_u.data[:] = lap_u_exact_fcn(X, Y) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e0707d7b1e..b1b9c3d187 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -852,63 +852,84 @@ def test_mixed_jacobian(self): assert j10.col_target == e assert j11.col_target == g + @pytest.mark.parametrize('eq1, eq2, j01_matvec, j10_matvec', [ + ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', + 'Eq(y_e(x, y), -h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), -h_x*h_y*x_e(x, y))'), + ('Eq(-e.laplace, 2.*g)', 'Eq(-g.laplace, 2.*e)', + 'Eq(y_e(x, y), -2.0*h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), -2.0*h_x*h_y*x_e(x, y))'), + ('Eq(-e.laplace, g.dx)', 'Eq(-g.laplace, e.dx)', + 'Eq(y_e(x, y), -h_x*h_y*Derivative(x_g(x, y), x))', + 'Eq(y_g(x, y), -h_x*h_y*Derivative(x_e(x, y), x))'), + ('Eq(-e.laplace, g.dx + g)', 'Eq(-g.laplace, e.dx + e)', + 'Eq(y_e(x, y), h_x*h_y*(-x_g(x, y) - Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(-x_e(x, y) - Derivative(x_e(x, y), x)))'), + ('Eq(e, g.dx + g)', 'Eq(g, e.dx + e)', + 'Eq(y_e(x, y), h_x*h_y*(-x_g(x, y) - Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(-x_e(x, y) - Derivative(x_e(x, y), x)))'), + ('Eq(e, g.dx + g.dy)', 'Eq(g, e.dx + e.dy)', + 'Eq(y_e(x, y), h_x*h_y*(-Derivative(x_g(x, y), x) - Derivative(x_g(x, y), y)))', + 'Eq(y_g(x, y), h_x*h_y*(-Derivative(x_e(x, y), x) - Derivative(x_e(x, y), y)))'), + ('Eq(g, -e.laplace)', 'Eq(e, -g.laplace)', + 'Eq(y_e(x, y), h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), h_x*h_y*x_e(x, y))'), + ('Eq(e + g, e.dx + 2.*g.dx)', 'Eq(g + e, g.dx + 2.*e.dx)', + 'Eq(y_e(x, y), h_x*h_y*(x_g(x, y) - 2.0*Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(x_e(x, y) - 2.0*Derivative(x_e(x, y), x)))'), + ]) + @skipif('petsc') + def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): + """ + Test linear coupling between fields, where the off-diagonal + Jacobian submatrices are nonzero. + """ + grid = Grid(shape=(9, 9), dtype=np.float64) - # TODO: FIX THIS LOGIC - THE OUTPUT FOR J01_MATVEC IS Eq(y_e(x, y), -h_x*h_y*x_e(x, y)) BUT IT SHOULD - # BE Eq(y_e(x, y), -h_x*h_y*x_G(x, y)) - - # @pytest.mark.parametrize('eq1, eq2, so, j01_matvec, j10_matvec', [ - # ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), - # ]) - # @skipif('petsc') - # def test_coupling(self, eq1, eq2, so, scale): - # """ - # Test linear coupling between fields, where the off-diagonal - # Jacobian submatrices are nonzero. - # """ - # grid = Grid(shape=(9,9), dtype=np.float64) - - # functions = [Function(name=n, grid=grid, space_order=eval(so)) - # for n in ['e', 'f', 'g', 'h']] - # e, f, g, h = functions - - # eq1 = eval(eq1) - # eq2 = eval(eq2) + e = Function(name='e', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) - # petsc = PETScSolve({e: [eq1], g: [eq2]}) + eq1 = eval(eq1) + eq2 = eval(eq2) - # jacobian = petsc[0].rhs.fielddata.jacobian + petsc = PETScSolve({e: [eq1], g: [eq2]}) - # j01 = jacobian.get_submatrix(0, 1) - # j10 = jacobian.get_submatrix(1, 0) + jacobian = petsc[0].rhs.fielddata.jacobian - # assert len(j01.matvecs) == 1 - # assert len(j10.matvecs) == 1 + j01 = jacobian.get_submatrix(0, 1) + j10 = jacobian.get_submatrix(1, 0) - # assert j01.col_target == g - # assert j10.col_target == e + assert j01.col_target == g + assert j10.col_target == e - # assert str(j01.matvecs[0]) == j01_matvec - # assert str(j10.matvecs[0]) == j10_matvec + assert str(j01.matvecs[0]) == j01_matvec + assert str(j10.matvecs[0]) == j10_matvec @pytest.mark.parametrize('eq1, eq2, so, scale', [ ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', '-2.0*h_x/h_x**2'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', '-2.5*h_x/h_x**2'), ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*(1 - 2.0/h_x**2)'), ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*(1 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*(5.0 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*(5.0 - 2.5/h_x**2)'), - ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', 'h_x*(1 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*(1 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*(1 - 5.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*(5.0 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*(5.0 - 2.5/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', + 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', + 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*(1 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*(1 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 1D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ grid = Grid(shape=(9,), dtype=np.float64) @@ -931,28 +952,37 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): assert str(j11.scdiag) == scale @pytest.mark.parametrize('eq1, eq2, so, scale', [ - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), - + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', + 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', + 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', + 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', + 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', + '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', + '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 2D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ - grid = Grid(shape=(9,9), dtype=np.float64) + grid = Grid(shape=(9, 9), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] @@ -972,28 +1002,40 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): assert str(j11.scdiag) == scale @pytest.mark.parametrize('eq1, eq2, so, scale', [ - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), - + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', + 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', + 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', + 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', + 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', + 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', + 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + ' + + '1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', + 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', + 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 3D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ - grid = Grid(shape=(9,9,9), dtype=np.float64) + grid = Grid(shape=(9, 9, 9), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] @@ -1012,9 +1054,6 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): assert str(j00.scdiag) == scale assert str(j11.scdiag) == scale - # test coupled residual callback - check the .dot for each field etc - - @skipif('petsc') def test_residual_bundle(self): grid = Grid(shape=(11, 11), dtype=np.float64) @@ -1055,6 +1094,42 @@ def test_residual_bundle(self): assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n ' \ + 'PetscScalar g;\n}' in str(op3.ccode) + @skipif('petsc') + def test_residual_callback(self): + """ + Check that the main residual callback correctly accesses the + target fields in the bundle. + """ + grid = Grid(shape=(9, 9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + # Check the residual callback + residual = op._func_table['WholeFormFunc0'].root + + exprs = FindNodes(Expression).visit(residual) + exprs = [str(e) for e in exprs] + + assert 'f_bundle[x + 2][y + 2].e = (r4*x_bundle[x + 1][y + 2].e + ' + \ + 'r4*x_bundle[x + 3][y + 2].e + r5*x_bundle[x + 2][y + 1].e + r5*' + \ + 'x_bundle[x + 2][y + 3].e - 2.0*(r4*x_bundle[x + 2][y + 2].e + r5*' + \ + 'x_bundle[x + 2][y + 2].e) - f[x + 2][y + 2])*ctx0->h_x*ctx0->h_y;' in exprs + + assert 'f_bundle[x + 2][y + 2].g = (r4*x_bundle[x + 1][y + 2].g + ' + \ + 'r4*x_bundle[x + 3][y + 2].g + r5*x_bundle[x + 2][y + 1].g + r5*' + \ + 'x_bundle[x + 2][y + 3].g - 2.0*(r4*x_bundle[x + 2][y + 2].g + r5*' + \ + 'x_bundle[x + 2][y + 2].g) - h[x + 2][y + 2])*ctx0->h_x*ctx0->h_y;' in exprs + @skipif('petsc') def test_essential_bcs(self): """ From b6710be878ddbf51d9d225af4bba31b6284d25bb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 00:03:07 +0100 Subject: [PATCH 18/53] types: Fix PetscBundle symbolic shape --- devito/petsc/types/array.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index c7a09f8265..e16b036c5a 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -140,6 +140,10 @@ def __init__(self, *args, pname="Field", **kwargs): def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) + + @cached_property + def symbolic_shape(self): + return self.c0.symbolic_shape @cached_property def indexed(self): From 01d153ad8f3fb57eb97afcffa04f03456312c47a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 00:50:32 +0100 Subject: [PATCH 19/53] tests: Fix petsc tests --- tests/test_petsc.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 7fd59404d3..d5087c13e6 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -621,12 +621,12 @@ class TestCoupledLinear: ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '6'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '2'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '4'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '6'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '2'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '4'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '6'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '2'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '4'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '6'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '2'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '4'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '6'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '2'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '4'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '6'), From d6bd4f4c96362ce4398c1e8c7606877f152057dd Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 11:57:42 +0100 Subject: [PATCH 20/53] misc: Address comments --- devito/passes/iet/engine.py | 1 - devito/petsc/iet/routines.py | 52 +++++++++++++++++++--------------- devito/petsc/solve.py | 2 +- devito/petsc/types/array.py | 40 ++++++++++++++++++++++---- devito/petsc/types/types.py | 51 ++++++++++++++++++--------------- devito/symbolics/extraction.py | 2 ++ devito/types/array.py | 6 ++-- tests/test_petsc.py | 2 +- 8 files changed, 100 insertions(+), 56 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index ad4377d979..0dd6ad0a56 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -530,7 +530,6 @@ def _(i, mapper, sregistry): }) -@abstract_object.register(Array) @abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4e48a571f2..3ddf78e770 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -114,6 +114,10 @@ def fielddata(self): def arrays(self): return self.fielddata.arrays + @property + def target(self): + return self.fielddata.target + def _make_core(self): self._make_matvec(self.fielddata.jacobian) self._make_formfunc() @@ -163,9 +167,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) if self.zero_memory else None + zero_y_memory = self.zero_vector(objs['Y']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] @@ -297,7 +299,8 @@ def _create_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + arrays = self.arrays + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -307,8 +310,8 @@ def _create_formfunc_body(self, body): fields = self._dummy_fields(body) self._struct_params.extend(fields) - f_formfunc = self.fielddata.arrays[target]['f'] - x_formfunc = self.fielddata.arrays[target]['x'] + f_formfunc = arrays[target]['f'] + x_formfunc = arrays[target]['x'] dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -316,9 +319,7 @@ def _create_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = petsc_call( - 'VecSet', [objs['F'], 0.0] - ) if self.zero_memory else None + zero_f_memory = self.zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] @@ -437,7 +438,7 @@ def _create_form_rhs_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -508,13 +509,15 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + dereference_funcs = tuple( + [Dereference(i, ctx) for i in + fields if isinstance(i.function, AbstractFunction)] + ) formrhs_body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+dereference_funcs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -550,7 +553,7 @@ def _create_initial_guess_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -653,6 +656,12 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper + def zero_vector(self, vec): + """ + Zeros the memory of the output vector before computation + """ + return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None + class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -1059,7 +1068,7 @@ def residual_bundle(self, body, bundles): if i.base in mapper: bundle = mapper[i.base] index = bundles['target_indices'][i.function.target] - index = (index,)+i.indices + index = (index,) + i.indices subs[i] = bundle.__getitem__(index) body = Uxreplace(subs).visit(body) @@ -1178,9 +1187,9 @@ def _extend_build(self, base_dict): base_dict[f'{name}F'] = CallbackVec(f'{name}F') # Bundle objects/metadata required by the coupled residual callback - f_components = [] - x_components = [] + f_components, x_components = [], [] bundle_mapper = {} + pname = sreg.make_name(prefix='Field') target_indices = {t: i for i, t in enumerate(targets)} @@ -1190,18 +1199,15 @@ def _extend_build(self, base_dict): f_components.append(f_arr) x_components.append(x_arr) - bundle_pname = sreg.make_name(prefix='Field') fbundle = PetscBundle( - name='f_bundle', components=f_components, pname=bundle_pname + name='f_bundle', components=f_components, pname=pname ) xbundle = PetscBundle( - name='x_bundle', components=x_components, pname=bundle_pname + name='x_bundle', components=x_components, pname=pname ) # Build the bundle mapper - for i, t in enumerate(targets): - f_arr = arrays[t]['f'] - x_arr = arrays[t]['x'] + for f_arr, x_arr in zip(f_components, x_components): bundle_mapper[f_arr.base] = fbundle bundle_mapper[x_arr.base] = xbundle diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 61d0c187e6..9a85774ef8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -83,7 +83,7 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) arrays = self.generate_arrays_combined(target) - eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + eqns = sorted(eqns, key=lambda e: not isinstance(e, EssentialBC)) jacobian = Jacobian(target, eqns, arrays, self.time_mapper) residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index e16b036c5a..e691a96c33 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -126,6 +126,34 @@ def symbolic_shape(self): class PetscBundle(Bundle): """ + Tensor symbol representing an unrolled vector of PETScArrays. + + This class declares a struct in the generated ccode to represent the + fields defined at each node of the grid. For example: + + typedef struct { + PetscScalar u,v,omega,temperature; + } Field; + + Residual evaluations are then written using: + + f[i][j].omega = ... + + Reference - https://petsc.org/release/manual/vec/#sec-struct + + Parameters + ---------- + name : str + Name of the symbol. + components : tuple of PETScArray + The PETScArrays of the Bundle. + pname : str, optional + The name of the struct in the generated C code. Defaults to "Field". + + Warnings + -------- + PetscBundles are created and managed directly by Devito (IOW, they are not + expected to be used directly in user code). """ is_Bundle = True _data_alignment = False @@ -136,11 +164,11 @@ def __init__(self, *args, pname="Field", **kwargs): super().__init__(*args, **kwargs) self._pname = pname - @property + @cached_property def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) - + @cached_property def symbolic_shape(self): return self.c0.symbolic_shape @@ -176,8 +204,10 @@ def __getitem__(self, index): component_names=names ) else: - raise ValueError("Expected %d or %d indices, got %d instead" - % (self.ndim, self.ndim + 1, len(index))) + raise ValueError( + f"Expected {self.ndim} or {self.ndim + 1} indices, " + f"got {len(index)} instead" + ) @property def pname(self): @@ -187,7 +217,7 @@ def pname(self): class PetscComponentAccess(ComponentAccess): def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: - raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) + raise ValueError(f"Expected Indexed, got `{type(arg)}` instead") names = component_names or cls._default_component_names obj = Expr.__new__(cls, arg) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 44d43d2f97..2e866896cb 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -2,7 +2,7 @@ from itertools import chain -from devito.tools import Reconstructable, sympy_mutex, as_tuple +from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil @@ -269,11 +269,11 @@ def col_target(self): return self.target def _build_matvecs(self): - matvecs = [ - e for eq in self.eqns for e in - self._build_matvec_eq(eq) - if e is not None - ] + matvecs = [] + for eq in self.eqns: + matvecs.extend( + e for e in self._build_matvec_eq(eq) if e is not None + ) matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -287,12 +287,13 @@ def _build_matvec_eq(self, eq, col_target=None, row_target=None): col_target = col_target or self.target row_target = row_target or self.target - b, F_target, _, targets = separate_eqn(eq, col_target) + _, F_target, _, targets = separate_eqn(eq, col_target) if F_target: return self._make_matvec( eq, F_target, targets, col_target, row_target ) - return (None,) + else: + return (None,) def _make_matvec(self, eq, F_target, targets, col_target, row_target): y = self.arrays[row_target]['y'] @@ -307,7 +308,7 @@ def _make_matvec(self, eq, F_target, targets, col_target, row_target): else: rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return as_tuple(Eq(y, rhs, subdomain=eq.subdomain)) + return (Eq(y, rhs, subdomain=eq.subdomain),) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target @@ -376,7 +377,7 @@ class MixedJacobian(Jacobian): def __init__(self, target_eqns, arrays, time_mapper): """ """ - self.targets = as_tuple(target_eqns.keys()) + self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper self._submatrices = [] @@ -391,7 +392,7 @@ def submatrices(self): return self._submatrices @property - def no_submatrices(self): + def n_submatrices(self): """ Return the number of submatrix blocks. """ @@ -421,11 +422,13 @@ def _build_blocks(self, target_eqns): for i, row_target in enumerate(self.targets): eqns = target_eqns[row_target] for j, col_target in enumerate(self.targets): - matvecs = [ - e for eq in eqns for e in - self._build_matvec_eq(eq, col_target, row_target) - ] + matvecs = [] + for eq in eqns: + matvecs.extend( + e for e in self._build_matvec_eq(eq, col_target, row_target) + ) matvecs = [m for m in matvecs if m is not None] + # Sort to put EssentialBC first if any matvecs = tuple( sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) @@ -461,7 +464,7 @@ def __repr__(self): f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})" for sm in self.submatrices ) - return f"" + return f"" class Residual: @@ -499,12 +502,13 @@ def _build_equations(self): # TODO: If b is zero then don't need a rhs vector+callback rhs.extend(self._make_b(eq, b)) - self._formfuncs = [self._scale_bcs(eq) for eq in funcs] - self._formrhs = rhs + self._formfuncs = tuple([self._scale_bcs(eq) for eq in funcs]) + self._formrhs = tuple(rhs) def _make_F_target(self, eq, F_target, targets): arrays = self.arrays[self.target] volume = self.target.grid.symbolic_volume_cell + if isinstance(eq, EssentialBC): # The initial guess satisfies the essential BCs, so this term is zero. # Still included to support Jacobian testing via finite differences. @@ -513,19 +517,20 @@ def _make_F_target(self, eq, F_target, targets): # Move essential boundary condition to the right-hand side zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) return (zero_row, zero_col) + else: if isinstance(F_target, (int, float)): rhs = F_target * volume else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) * volume - return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) + return (Eq(arrays['f'], rhs, subdomain=eq.subdomain),) def _make_b(self, eq, b): b_arr = self.arrays[self.target]['b'] rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return as_tuple(Eq(b_arr, rhs, subdomain=eq.subdomain)) + return (Eq(b_arr, rhs, subdomain=eq.subdomain),) def _scale_bcs(self, eq, scdiag=None): """ @@ -617,11 +622,11 @@ def _build_equations(self): """ Return a list of initial guess equations. """ - self._eqs = [ + self._eqs = tuple([ eq for eq in (self._make_initial_guess(e) for e in self.eqns) if eq is not None - ] + ]) def _make_initial_guess(self, eq): if isinstance(eq, EssentialBC): @@ -655,4 +660,4 @@ def targets_to_arrays(array, targets): array_targets = [ array.subs(dict(zip(array.indices, i))) for i in space_indices ] - return dict(zip(targets, array_targets)) + return frozendict(zip(targets, array_targets)) diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index dfbc145748..fc47e94444 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -17,9 +17,11 @@ def separate_eqn(eqn, target): where F(target) = b. """ zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + from devito.operations.solve import eval_time_derivatives zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_targets(zeroed_eqn, target_funcs) return -b, F_target, zeroed_eqn, target_funcs diff --git a/devito/types/array.py b/devito/types/array.py index 8c7fdb8a3d..aec3bd0626 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -531,8 +531,10 @@ def __getitem__(self, index): component_index, indices = index[0], index[1:] return ComponentAccess(self.indexed[indices], component_index) else: - raise ValueError("Expected %d or %d indices, got %d instead" - % (self.ndim, self.ndim + 1, len(index))) + raise ValueError( + f"Expected {self.ndim} or {self.ndim + 1} indices, " + f"got {len(index)} instead" + ) @property def _C_ctype(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d5087c13e6..10e030b064 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -826,7 +826,7 @@ def test_mixed_jacobian(self): j11 = jacobian.get_submatrix(1, 1) # Check the number of submatrices - assert jacobian.no_submatrices == 4 + assert jacobian.n_submatrices == 4 # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 From 4fad6aa52235df66975eb85902b77988aad916b7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 12:04:04 +0100 Subject: [PATCH 21/53] misc: Merge leftover --- tests/test_symbolics.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index ec90facd13..61c45527f8 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -957,7 +957,6 @@ def test_print_div(): assert cstr == 'sizeof(int)/sizeof(long)' -<<<<<<< HEAD @pytest.mark.parametrize('eqn, target, expected', [ ('Eq(f1.laplace, g1)', 'f1', ('g1(x, y)', 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')), @@ -1194,7 +1193,8 @@ def test_centre_stencil(expr, so, target, expected): centre = centre_stencil(eval(expr), eval(target)) assert str(centre) == expected -======= + + def test_customdtype_complex(): """ Test that `CustomDtype` doesn't brak is_imag @@ -1327,4 +1327,3 @@ def test_conj(self): Operator([Eq(g, Conj(f))])() assert np.all(np.isclose(g.data, np.conj(f.data))) ->>>>>>> petsc From 27dd201b0365e3b07895d2dbc47ae7826692ae3e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 17:33:19 +0100 Subject: [PATCH 22/53] compiler: Edit switch_log_level --- devito/logger.py | 5 +++-- devito/operator/operator.py | 6 +++++- devito/petsc/iet/passes.py | 18 +++++++++--------- devito/petsc/iet/routines.py | 14 +++++++++----- devito/petsc/initialize.py | 7 ++++--- 5 files changed, 30 insertions(+), 20 deletions(-) diff --git a/devito/logger.py b/devito/logger.py index 9c1dfa56d2..c9834dbfe7 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -76,9 +76,10 @@ def set_log_level(level, comm=None): used, for example, if one wants to log to one file per rank. """ from devito import configuration - + from devito.mpi.distributed import MPI + if comm is not None and configuration['mpi']: - if comm.rank != 0: + if comm!= MPI.COMM_NULL and comm.rank != 0: logger.removeHandler(stream_handler) logger.addHandler(logging.NullHandler()) else: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 3a36952a58..5216bde481 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -38,7 +38,7 @@ from devito.petsc.iet.passes import lower_petsc from devito.petsc.clusters import petsc_preprocess -__all__ = ['Operator'] +__all__ = ['Operator', 'SpecialOp'] class Operator(Callable): @@ -1456,3 +1456,7 @@ def parse_kwargs(**kwargs): kwargs['subs'] = {k: sympify(v) for k, v in kwargs.get('subs', {}).items()} return kwargs + + +class SpecialOp(Operator): + pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 3c73bfd391..292f455acf 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -52,9 +52,10 @@ def lower_petsc(iet, **kwargs): # Assumption is that all solves are on the same grid if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") + grid = unique_grids.pop() # Create core PETSc calls (not specific to each PETScSolve) - core = make_core_petsc_calls(objs, **kwargs) + core = make_core_petsc_calls(objs, grid, **kwargs) setup = [] subs = {} @@ -62,7 +63,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, **kwargs) + builder = Builder(injectsolve, objs, iters, grid, **kwargs) setup.extend(builder.solversetup.calls) @@ -108,8 +109,9 @@ def finalize(iet): return iet._rebuild(body=finalize_body) -def make_core_petsc_calls(objs, **kwargs): - call_mpi = petsc_call_mpi('MPI_Comm_size', [objs['comm'], Byref(objs['size'])]) +def make_core_petsc_calls(objs, grid, **kwargs): + comm = grid.distributor._obj_comm + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine @@ -123,16 +125,18 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, **kwargs): + def __init__(self, injectsolve, objs, iters, grid, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters + self.grid = grid self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) self.args = { 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, + 'grid': self.grid, **self.kwargs } self.args['solver_objs'] = self.objbuilder.solver_objs @@ -190,9 +194,6 @@ def populate_matrix_context(efuncs, objs): ) -# TODO: Devito MPI + PETSc testing -# if kwargs['options']['mpi'] -> communicator = grid.distributor._obj_comm -communicator = 'PETSC_COMM_WORLD' subdms = PointerDM(name='subdms') fields = PointerIS(name='fields') submats = PointerMat(name='submats') @@ -208,7 +209,6 @@ def populate_matrix_context(efuncs, objs): # they are semantically identical. objs = frozendict({ 'size': PetscMPIInt(name='size'), - 'comm': communicator, 'err': PetscErrorCode(name='err'), 'block': CallbackMat('block'), 'submat_arr': PointerMat(name='submat_arr'), diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ddf78e770..ae305a5ba0 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -963,7 +963,7 @@ def _submat_callback_body(self): ptr = DummyExpr(objs['submat_arr']._C_symbol, Deref(objs['Submats']), init=True) - mat_create = petsc_call('MatCreate', [self.objs['comm'], Byref(objs['block'])]) + mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])]) mat_set_sizes = petsc_call( 'MatSetSizes', [ @@ -1085,6 +1085,7 @@ def __init__(self, **kwargs): self.injectsolve = kwargs.get('injectsolve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') + self.grid = kwargs.get('grid') self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1124,6 +1125,8 @@ def _build(self): 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } + # TODO: Devito MPI + PETSc testing + base_dict['comm'] = self.grid.distributor._obj_comm self._target_dependent(base_dict) return self._extend_build(base_dict) @@ -1279,7 +1282,7 @@ def _setup(self): solver_params = self.injectsolve.expr.rhs.solver_parameters - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) @@ -1305,7 +1308,7 @@ def _setup(self): v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space ) local_x = petsc_call('VecCreateMPIWithArray', - ['PETSC_COMM_WORLD', 1, local_size, 'PETSC_DECIDE', + [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', field_from_ptr, Byref(sobjs['xlocal'])]) # TODO: potentially also need to set the DM and local/global map to xlocal @@ -1408,11 +1411,12 @@ def _create_dmda_calls(self, dmda): def _create_dmda(self, dmda): objs = self.objs + sobjs = self.solver_objs grid = self.fielddata.grid nspace_dims = len(grid.dimensions) # MPI communicator - args = [objs['comm']] + args = [sobjs['comm']] # Type of ghost nodes args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) @@ -1453,7 +1457,7 @@ def _setup(self): solver_params = self.injectsolve.expr.rhs.solver_parameters - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 80e3c7520c..3fc419c26a 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -3,7 +3,8 @@ from ctypes import POINTER, cast, c_char import atexit -from devito import Operator, switchconfig +from devito import switchconfig +from devito.operator.operator import SpecialOp from devito.types import Symbol from devito.types.equation import PetscEq from devito.petsc.types import Initialize, Finalize @@ -21,11 +22,11 @@ def PetscInitialize(): # This would prevent circular imports when initializing during import # from the PETSc module. with switchconfig(language='petsc'): - op_init = Operator( + op_init = SpecialOp( [PetscEq(dummy, Initialize(dummy))], name='kernel_init', opt='noop' ) - op_finalize = Operator( + op_finalize = SpecialOp( [PetscEq(dummy, Finalize(dummy))], name='kernel_finalize', opt='noop' ) From 30ed5bac8e551b00be9c8485f0bf81702d3e7476 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 18:16:45 +0100 Subject: [PATCH 23/53] tests: Add mpi petsc test --- devito/petsc/iet/passes.py | 32 ++++++++++++++++++++++++++++++-- devito/petsc/iet/routines.py | 8 +++++++- examples/petsc/petsc_test.py | 2 +- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 292f455acf..f77879cf10 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -93,6 +93,17 @@ def initialize(iet): help_string = c.Line(r'static char help[] = "This is help text.\n";') init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) + + # print_comm_size = + + # size = c.Line('PetscMPIInt size;') + # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + # rank = c.Line('PetscMPIInt rank;') + # get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') + # print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') + # print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') + + init_body = CallableBody( body=(petsc_func_begin_user, help_string, init_body), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) @@ -110,10 +121,27 @@ def finalize(iet): def make_core_petsc_calls(objs, grid, **kwargs): - comm = grid.distributor._obj_comm + devito_mpi = kwargs['options'].get('mpi', False) + if devito_mpi: + comm = grid.distributor._obj_comm + else: + comm = 'PETSC_COMM_WORLD' + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) + # from IPython import embed; embed() + # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + rank = c.Line('PetscMPIInt rank;') + get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') + print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') + + print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') + + + flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') + - return call_mpi, BlankLine + return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + # return call_mpi, BlankLine class Builder: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index ae305a5ba0..812e7b9b86 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1086,6 +1086,8 @@ def __init__(self, **kwargs): self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.grid = kwargs.get('grid') + # from IPython import embed; embed() + self.devito_mpi = kwargs['options'].get('mpi', False) self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1126,7 +1128,11 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - base_dict['comm'] = self.grid.distributor._obj_comm + # from IPython import embed; embed() + if self.devito_mpi: + base_dict['comm'] = self.grid.distributor._obj_comm + else: + base_dict['comm'] = 'PETSC_COMM_WORLD' self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..eb74c639db 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -print(op.ccode) +# print(op.ccode) From 51d4d7e9dc91d03b26b15449c9fadae9ecc591b9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 00:19:13 +0100 Subject: [PATCH 24/53] misc: Address more comments and add docstrings --- devito/petsc/iet/routines.py | 33 ++++++++++++++++----------------- devito/petsc/types/types.py | 28 +++++++++++++++++++++++++--- devito/symbolics/extraction.py | 3 ++- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ddf78e770..4efc88c1b5 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -258,13 +258,12 @@ def _create_matvec_body(self, body, jacobian): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) matvec_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -398,13 +397,12 @@ def _create_formfunc_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) formfunc_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) # Replace non-function data with pointer to data in struct @@ -509,15 +507,12 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - dereference_funcs = tuple( - [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] - ) + derefs = self.dereference_funcs(ctx, fields) formrhs_body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+dereference_funcs, + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -590,13 +585,12 @@ def _create_initial_guess_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -662,6 +656,12 @@ def zero_vector(self, vec): """ return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None + def dereference_funcs(self, struct, fields): + return tuple( + [Dereference(i, struct) for i in + fields if isinstance(i.function, AbstractFunction)] + ) + class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -894,8 +894,7 @@ def _whole_formfunc_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -903,7 +902,7 @@ def _whole_formfunc_body(self, body): formfunc_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, casts=(f_soa, x_soa), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 2e866896cb..eee16047cb 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -241,7 +241,7 @@ class Jacobian: This Jacobian is defined implicitly via matrix-vector products derived from the symbolic equations provided in `matvecs`. - It assumes the problem is linear, meaning the Jacobian + The class assumes the problem is linear, meaning the Jacobian corresponds to a constant coefficient matrix and does not require explicit symbolic differentiation. """ @@ -274,6 +274,7 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_eq(eq) if e is not None ) + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -469,6 +470,20 @@ def __repr__(self): class Residual: """ + Gennerates the metadata needed to define the nonlinear residual function + F(target) = 0 for use with PETSc's SNES interface. + + PETSc's SNES interface includes methods for solving nonlinear systems of + equations using Newton-type methods. For linear problems, `SNESKSPONLY` + is used to perform a single Newton iteration, unifying the + interface for both linear and nonlinear problems. + + This class encapsulates the symbolic equations used to construct the + residual function F(target) = F_(target) - b, where b contains all + terms independent of the solution `target`. + + References: + - https://petsc.org/main/manual/snes/ """ def __init__(self, target, eqns, arrays, time_mapper, scdiag): self.target = target @@ -481,12 +496,18 @@ def __init__(self, target, eqns, arrays, time_mapper, scdiag): @property def formfuncs(self): """ + Stores the equations used to build the `FormFunction` + callback generated at the IET level. This function is + passed to PETSc via `SNESSetFunction(..., FormFunction, ...)`. """ return self._formfuncs @property def formrhs(self): """ + Stores the equations used to generate the RHS + vector `b` through the `FormRHS` callback generated at the IET level. + The SNES solver is then called via `SNESSolve(..., b, target)`. """ return self._formrhs @@ -544,7 +565,7 @@ class MixedResidual(Residual): """ """ def __init__(self, target_eqns, arrays, time_mapper, scdiag): - self.targets = as_tuple(target_eqns.keys()) + self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag @@ -592,6 +613,7 @@ def _build_function_eq(self, eq, target): self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain ) return (zero_row, zero_col) + else: if isinstance(zeroed, (int, float)): rhs = zeroed * volume @@ -599,7 +621,7 @@ def _build_function_eq(self, eq, target): rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper)*volume - return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) + return (Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain),) class InitialGuess: diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index fc47e94444..c0a28433cc 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -20,9 +20,10 @@ def separate_eqn(eqn, target): from devito.operations.solve import eval_time_derivatives zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - target_funcs = set(generate_targets(zeroed_eqn, target)) + target_funcs = set(generate_targets(zeroed_eqn, target)) b, F_target = remove_targets(zeroed_eqn, target_funcs) + return -b, F_target, zeroed_eqn, target_funcs From a4a56603fb92c4303df4b4c2f2568877875c552f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 12:21:49 +0100 Subject: [PATCH 25/53] misc: Clean up more docstrings --- devito/petsc/iet/routines.py | 92 +++++++++---------- devito/petsc/solve.py | 91 +++++++++--------- devito/petsc/types/types.py | 172 +++++++++++++++++------------------ tests/test_petsc.py | 10 +- 4 files changed, 181 insertions(+), 184 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4efc88c1b5..f762c3d30b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -46,9 +46,9 @@ def __init__(self, **kwargs): self._user_struct_callback = None # TODO: Test pickling. The mutability of these lists # could cause issues when pickling? - self._matvecs = [] - self._formfuncs = [] - self._formrhs = [] + self._J_efuncs = [] + self._F_efuncs = [] + self._b_efuncs = [] self._initialguesses = [] self._make_core() @@ -73,23 +73,23 @@ def main_matvec_callback(self): is set in the main kernel via `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` """ - return self._matvecs[0] + return self._J_efuncs[0] @property def main_formfunc_callback(self): - return self._formfuncs[0] + return self._F_efuncs[0] @property - def matvecs(self): - return self._matvecs + def J_efuncs(self): + return self._J_efuncs @property - def formfuncs(self): - return self._formfuncs + def F_efuncs(self): + return self._F_efuncs @property - def formrhs(self): - return self._formrhs + def b_efuncs(self): + return self._b_efuncs @property def initialguesses(self): @@ -122,27 +122,27 @@ def _make_core(self): self._make_matvec(self.fielddata.jacobian) self._make_formfunc() self._make_formrhs() - if self.fielddata.initialguess.eqs: + if self.fielddata.initialguess.exprs: self._make_initialguess() self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation matvecs = jacobian.matvecs - irs_matvec, _ = self.rcompile(matvecs, + irs, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) - body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), + body = self._create_matvec_body(List(body=irs.uiet.body), jacobian) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix=prefix), - body_matvec, + body, retval=objs['err'], parameters=(objs['J'], objs['X'], objs['Y']) ) - self._matvecs.append(cb) + self._J_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_matvec_body(self, body, jacobian): @@ -260,7 +260,7 @@ def _create_matvec_body(self, body, jacobian): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - matvec_body = CallableBody( + body = CallableBody( List(body=body), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -269,20 +269,20 @@ def _create_matvec_body(self, body, jacobian): # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - matvec_body = Uxreplace(subs).visit(matvec_body) + body = Uxreplace(subs).visit(body) self._struct_params.extend(fields) - return matvec_body + return body def _make_formfunc(self): - formfuncs = self.fielddata.residual.formfuncs - # Compile formfunc `eqns` into an IET via recursive compilation - irs_formfunc, _ = self.rcompile( - formfuncs, options={'mpi': False}, sregistry=self.sregistry, + F_exprs = self.fielddata.residual.F_exprs + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._create_formfunc_body( - List(body=irs_formfunc.uiet.body) + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -291,7 +291,7 @@ def _make_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._formfuncs.append(cb) + self._F_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): @@ -399,7 +399,7 @@ def _create_formfunc_body(self, body): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - formfunc_body = CallableBody( + body = CallableBody( List(body=body), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -408,28 +408,28 @@ def _create_formfunc_body(self, body): # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - return Uxreplace(subs).visit(formfunc_body) + return Uxreplace(subs).visit(body) def _make_formrhs(self): - formrhs = self.fielddata.residual.formrhs + b_exprs = self.fielddata.residual.b_exprs sobjs = self.solver_objs - # Compile formrhs `eqns` into an IET via recursive compilation - irs_formrhs, _ = self.rcompile( - formrhs, options={'mpi': False}, sregistry=self.sregistry, + # Compile `b_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + b_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formrhs = self._create_form_rhs_body( - List(body=irs_formrhs.uiet.body) + body = self._create_form_rhs_body( + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='FormRHS'), - body_formrhs, + body, retval=objs['err'], parameters=(sobjs['callbackdm'], objs['B']) ) - self._formrhs.append(cb) + self._b_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): @@ -509,7 +509,7 @@ def _create_form_rhs_body(self, body): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - formrhs_body = CallableBody( + body = CallableBody( List(body=[body]), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -520,24 +520,24 @@ def _create_form_rhs_body(self, body): subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields if not isinstance(i.function, AbstractFunction)} - return Uxreplace(subs).visit(formrhs_body) + return Uxreplace(subs).visit(body) def _make_initialguess(self): - initguess = self.fielddata.initialguess.eqs + exprs = self.fielddata.initialguess.exprs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation irs, _ = self.rcompile( - initguess, options={'mpi': False}, sregistry=self.sregistry, + exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_init_guess = self._create_initial_guess_body( + body = self._create_initial_guess_body( List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='FormInitialGuess'), - body_init_guess, + body, retval=objs['err'], parameters=(sobjs['callbackdm'], objs['xloc']) ) @@ -770,10 +770,10 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self): - formfuncs = self.fielddata.residual.formfuncs + F_exprs = self.fielddata.residual.F_exprs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( - formfuncs, options={'mpi': False}, sregistry=self.sregistry, + F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) @@ -1024,7 +1024,7 @@ def _submat_callback_body(self): iteration = Iteration(List(body=iter_body), i, upper_bound) nonzero_submats = self.jacobian.nonzero_submatrices - matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} + matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.J_efuncs} matmult_op = [ petsc_call( @@ -1616,11 +1616,11 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callback = self.cbbuilder.formrhs[0] + b_efunc = self.cbbuilder.b_efuncs[0] dmda = sobjs['dmda'] - rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) + rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) vec_place_array = self.timedep.place_array(target) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 9a85774ef8..949328cf67 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -12,9 +12,9 @@ __all__ = ['PETScSolve'] -def PETScSolve(target_eqns, target=None, solver_parameters=None): +def PETScSolve(target_exprs, target=None, solver_parameters=None): """ - Returns a symbolic equation representing a linear PETSc solver, + Returns a symbolic expression representing a linear PETSc solver, enriched with all the necessary metadata for execution within an `Operator`. When passed to an `Operator`, this symbolic equation triggers code generation and lowering to the PETSc backend. @@ -24,8 +24,8 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Parameters ---------- - target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq - The targets and symbolic equations defining the system to be solved. + target_exprs : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq + The targets and symbolic expressions defining the system to be solved. - Single-field problem: Pass a single Eq or list of Eq, and specify `target` separately: @@ -46,26 +46,26 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Returns ------- - Eq - A symbolic equation that wraps the linear solver. + Eq: + A symbolic expression that wraps the linear solver. This can be passed directly to a Devito Operator. """ if target is not None: - return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() + return InjectSolve(solver_parameters, {target: target_exprs}).build_expr() else: - return InjectMixedSolve(solver_parameters, target_eqns).build_eq() + return InjectMixedSolve(solver_parameters, target_exprs).build_expr() class InjectSolve: - def __init__(self, solver_parameters=None, target_eqns=None): + def __init__(self, solver_parameters=None, target_exprs=None): self.solver_params = solver_parameters self.time_mapper = None - self.target_eqns = target_eqns + self.target_exprs = target_exprs - def build_eq(self): + def build_expr(self): target, funcs, fielddata = self.linear_solve_args() - # Placeholder equation for inserting calls to the solver + # Placeholder expression for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, self.solver_params, @@ -76,18 +76,18 @@ def build_eq(self): return [PetscEq(target, linear_solve)] def linear_solve_args(self): - target, eqns = next(iter(self.target_eqns.items())) - eqns = as_tuple(eqns) + target, exprs = next(iter(self.target_exprs.items())) + exprs = as_tuple(exprs) - funcs = get_funcs(eqns) + funcs = get_funcs(exprs) self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays_combined(target) + arrays = self.generate_arrays(target) - eqns = sorted(eqns, key=lambda e: not isinstance(e, EssentialBC)) + exprs = sorted(exprs, key=lambda e: not isinstance(e, EssentialBC)) - jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, eqns, arrays) + jacobian = Jacobian(target, exprs, arrays, self.time_mapper) + residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag) + initialguess = InitialGuess(target, exprs, arrays) field_data = FieldData( target=target, @@ -98,55 +98,55 @@ def linear_solve_args(self): ) return target, tuple(funcs), field_data - - def generate_arrays(self, target): + + def generate_arrays(self, *targets): return { - p: PETScArray(name=f'{p}_{target.name}', - target=target, - liveness='eager', - localinfo=localinfo) - for p in prefixes + t: { + p: PETScArray(name=f'{p}_{t.name}', + target=t, + liveness='eager', + localinfo=localinfo) + for p in prefixes + } + for t in targets } - def generate_arrays_combined(self, *targets): - return {target: self.generate_arrays(target) for target in targets} - class InjectMixedSolve(InjectSolve): def linear_solve_args(self): - combined_eqns = [] - for eqns in self.target_eqns.values(): - combined_eqns.extend(eqns) - funcs = get_funcs(combined_eqns) - self.time_mapper = generate_time_mapper(funcs) + exprs = [] + for e in self.target_exprs.values(): + exprs.extend(e) - coupled_targets = list(self.target_eqns.keys()) + funcs = get_funcs(exprs) + self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays_combined(*coupled_targets) + targets = list(self.target_exprs.keys()) + arrays = self.generate_arrays(*targets) jacobian = MixedJacobian( - self.target_eqns, arrays, self.time_mapper + self.target_exprs, arrays, self.time_mapper ) residual = MixedResidual( - self.target_eqns, arrays, self.time_mapper, + self.target_exprs, arrays, self.time_mapper, jacobian.target_scaler_mapper ) all_data = MultipleFieldData( - targets=coupled_targets, + targets=targets, arrays=arrays, jacobian=jacobian, residual=residual ) - return coupled_targets[0], tuple(funcs), all_data + return targets[0], tuple(funcs), all_data def generate_time_mapper(funcs): """ - Replace time indices with `Symbols` in equations used within + Replace time indices with `Symbols` in expressions used within PETSc callback functions. These symbols are Uxreplaced at the IET level to align with the `TimeDimension` and `ModuloDimension` objects present in the initial lowering. @@ -176,11 +176,10 @@ def generate_time_mapper(funcs): return dict(zip(time_indices, tau_symbs)) -def get_funcs(eqns): +def get_funcs(exprs): funcs = [ - func - for eq in eqns - for func in retrieve_functions(eval_time_derivatives(eq.lhs - eq.rhs)) + f for e in exprs + for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) ] return filter_ordered(funcs) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index eee16047cb..f4e92f2f2e 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -239,21 +239,26 @@ class Jacobian: Represents a Jacobian matrix. This Jacobian is defined implicitly via matrix-vector products - derived from the symbolic equations provided in `matvecs`. + derived from the symbolic expressions provided in `matvecs`. The class assumes the problem is linear, meaning the Jacobian corresponds to a constant coefficient matrix and does not require explicit symbolic differentiation. """ - def __init__(self, target, eqns, arrays, time_mapper): + def __init__(self, target, exprs, arrays, time_mapper): self.target = target - self.eqns = eqns + self.exprs = exprs self.arrays = arrays self.time_mapper = time_mapper self._build_matvecs() @property def matvecs(self): + """ + Stores the expressions used to generate the `MatMult` + callback generated at the IET level. This function is + passed to PETSc via `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. + """ return self._matvecs @property @@ -270,9 +275,9 @@ def col_target(self): def _build_matvecs(self): matvecs = [] - for eq in self.eqns: + for eq in self.exprs: matvecs.extend( - e for e in self._build_matvec_eq(eq) if e is not None + e for e in self._build_matvec_expr(eq) if e is not None ) matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) @@ -284,32 +289,32 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_eq(self, eq, col_target=None, row_target=None): + def _build_matvec_expr(self, expr, col_target=None, row_target=None): col_target = col_target or self.target row_target = row_target or self.target - _, F_target, _, targets = separate_eqn(eq, col_target) + _, F_target, _, targets = separate_eqn(expr, col_target) if F_target: return self._make_matvec( - eq, F_target, targets, col_target, row_target + expr, F_target, targets, col_target, row_target ) else: return (None,) - def _make_matvec(self, eq, F_target, targets, col_target, row_target): + def _make_matvec(self, expr, F_target, targets, col_target, row_target): y = self.arrays[row_target]['y'] x = self.arrays[col_target]['x'] - if isinstance(eq, EssentialBC): + if isinstance(expr, EssentialBC): # NOTE: Essential BCs are trivial equations in the solver. # See `EssentialBC` for more details. - zero_row = ZeroRow(y, x, subdomain=eq.subdomain) - zero_column = ZeroColumn(x, 0., subdomain=eq.subdomain) + zero_row = ZeroRow(y, x, subdomain=expr.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) return (zero_row, zero_column) else: rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return (Eq(y, rhs, subdomain=eq.subdomain),) + return (Eq(y, rhs, subdomain=expr.subdomain),) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target @@ -376,8 +381,6 @@ class MixedJacobian(Jacobian): # TODO: pcfieldsplit support for each block """ def __init__(self, target_eqns, arrays, time_mapper): - """ - """ self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper @@ -394,15 +397,13 @@ def submatrices(self): @property def n_submatrices(self): - """ - Return the number of submatrix blocks. - """ + """Return the number of submatrix blocks.""" return len(self._submatrices) @property def nonzero_submatrices(self): """Return SubMatrixBlock objects that have non-empty matvecs.""" - return [submat for submat in self.submatrices if submat.matvecs] + return [m for m in self.submatrices if m.matvecs] @property def target_scaler_mapper(self): @@ -411,22 +412,22 @@ def target_scaler_mapper(self): diagonal subblock. """ mapper = {} - for sm in self.submatrices: - if sm.row_idx == sm.col_idx: - mapper[sm.row_target] = sm.scdiag + for m in self.submatrices: + if m.row_idx == m.col_idx: + mapper[m.row_target] = m.scdiag return mapper - def _build_blocks(self, target_eqns): + def _build_blocks(self, target_exprs): """ Build all SubMatrixBlock objects for the Jacobian. """ for i, row_target in enumerate(self.targets): - eqns = target_eqns[row_target] + exprs = target_exprs[row_target] for j, col_target in enumerate(self.targets): matvecs = [] - for eq in eqns: + for expr in exprs: matvecs.extend( - e for e in self._build_matvec_eq(eq, col_target, row_target) + e for e in self._build_matvec_expr(expr, col_target, row_target) ) matvecs = [m for m in matvecs if m is not None] @@ -475,56 +476,56 @@ class Residual: PETSc's SNES interface includes methods for solving nonlinear systems of equations using Newton-type methods. For linear problems, `SNESKSPONLY` - is used to perform a single Newton iteration, unifying the + can be used to perform a single Newton iteration, unifying the interface for both linear and nonlinear problems. - This class encapsulates the symbolic equations used to construct the + This class encapsulates the symbolic expressions used to construct the residual function F(target) = F_(target) - b, where b contains all terms independent of the solution `target`. References: - https://petsc.org/main/manual/snes/ """ - def __init__(self, target, eqns, arrays, time_mapper, scdiag): + def __init__(self, target, exprs, arrays, time_mapper, scdiag): self.target = target - self.eqns = eqns + self.exprs = exprs self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag - self._build_equations() + self._build_exprs() @property - def formfuncs(self): + def F_exprs(self): """ - Stores the equations used to build the `FormFunction` + Stores the expressions used to build the `FormFunction` callback generated at the IET level. This function is passed to PETSc via `SNESSetFunction(..., FormFunction, ...)`. """ - return self._formfuncs + return self._F_exprs @property - def formrhs(self): + def b_exprs(self): """ - Stores the equations used to generate the RHS + Stores the expressions used to generate the RHS vector `b` through the `FormRHS` callback generated at the IET level. The SNES solver is then called via `SNESSolve(..., b, target)`. """ - return self._formrhs + return self._b_exprs - def _build_equations(self): + def _build_exprs(self): """ """ - funcs = [] - rhs = [] + F_exprs = [] + b_exprs = [] - for eq in self.eqns: - b, F_target, _, targets = separate_eqn(eq, self.target) - funcs.extend(self._make_F_target(eq, F_target, targets)) + for e in self.exprs: + b, F_target, _, targets = separate_eqn(e, self.target) + F_exprs.extend(self._make_F_target(e, F_target, targets)) # TODO: If b is zero then don't need a rhs vector+callback - rhs.extend(self._make_b(eq, b)) + b_exprs.extend(self._make_b(e, b)) - self._formfuncs = tuple([self._scale_bcs(eq) for eq in funcs]) - self._formrhs = tuple(rhs) + self._F_exprs = tuple([self._scale_bcs(e) for e in F_exprs]) + self._b_exprs = tuple(b_exprs) def _make_F_target(self, eq, F_target, targets): arrays = self.arrays[self.target] @@ -547,15 +548,15 @@ def _make_F_target(self, eq, F_target, targets): rhs = rhs.subs(self.time_mapper) * volume return (Eq(arrays['f'], rhs, subdomain=eq.subdomain),) - def _make_b(self, eq, b): + def _make_b(self, expr, b): b_arr = self.arrays[self.target]['b'] - rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = 0. if isinstance(expr, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return (Eq(b_arr, rhs, subdomain=eq.subdomain),) + return (Eq(b_arr, rhs, subdomain=expr.subdomain),) def _scale_bcs(self, eq, scdiag=None): """ - Scale ZeroRow equations using scdiag + Scale ZeroRow exprs using scdiag """ scdiag = scdiag or self.scdiag return eq._rebuild(rhs=scdiag * eq.rhs) if isinstance(eq, ZeroRow) else eq @@ -564,37 +565,35 @@ def _scale_bcs(self, eq, scdiag=None): class MixedResidual(Residual): """ """ - def __init__(self, target_eqns, arrays, time_mapper, scdiag): - self.targets = tuple(target_eqns.keys()) + def __init__(self, target_exprs, arrays, time_mapper, scdiag): + self.targets = tuple(target_exprs.keys()) self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag - self._build_equations(target_eqns) + self._build_exprs(target_exprs) @property - def formrhs(self): + def b_exprs(self): """ """ return None - def _build_equations(self, target_eqns): - all_formfuncs = [] - for target, eqns in target_eqns.items(): + def _build_exprs(self, target_exprs): + residual_exprs = [] + for t, exprs in target_exprs.items(): - formfuncs = chain.from_iterable( - self._build_function_eq(eq, target) - for eq in as_tuple(eqns) - ) - all_formfuncs.extend(formfuncs) + residual_exprs.extend( + chain.from_iterable(self._build_residual(e, t) + for e in as_tuple(exprs) + )) - self._formfuncs = tuple(sorted( - all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + self._F_exprs = tuple(sorted( + residual_exprs, key=lambda e: not isinstance(e, EssentialBC) )) - def _build_function_eq(self, eq, target): - zeroed = eq.lhs - eq.rhs - - zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + def _build_residual(self, expr, target): + zeroed = expr.lhs - expr.rhs + zeroed_eqn = Eq(zeroed, 0) eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) volume = target.grid.symbolic_volume_cell @@ -604,13 +603,13 @@ def _build_function_eq(self, eq, target): target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) - if isinstance(eq, EssentialBC): - rhs = (self.arrays[target]['x'] - eq.rhs)*self.scdiag[target] + if isinstance(expr, EssentialBC): + rhs = (self.arrays[target]['x'] - expr.rhs)*self.scdiag[target] zero_row = ZeroRow( - self.arrays[target]['f'], rhs, subdomain=eq.subdomain + self.arrays[target]['f'], rhs, subdomain=expr.subdomain ) zero_col = ZeroColumn( - self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain + self.arrays[target]['x'], expr.rhs, subdomain=expr.subdomain ) return (zero_row, zero_col) @@ -621,41 +620,40 @@ def _build_function_eq(self, eq, target): rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper)*volume - return (Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain),) + return (Eq(self.arrays[target]['f'], rhs, subdomain=expr.subdomain),) class InitialGuess: """ Enforce initial guess to satisfy essential BCs. """ - def __init__(self, target, eqns, arrays): + def __init__(self, target, exprs, arrays): self.target = target - self.eqns = as_tuple(eqns) self.arrays = arrays - self._build_equations() + self._build_exprs(as_tuple(exprs)) @property - def eqs(self): + def exprs(self): """ """ - return self._eqs + return self._exprs - def _build_equations(self): + def _build_exprs(self, exprs): """ - Return a list of initial guess equations. + Return a list of initial guess expressions. """ - self._eqs = tuple([ + self._exprs = tuple([ eq for eq in - (self._make_initial_guess(e) for e in self.eqns) + (self._make_initial_guess(e) for e in exprs) if eq is not None ]) - def _make_initial_guess(self, eq): - if isinstance(eq, EssentialBC): - assert eq.lhs == self.target + def _make_initial_guess(self, expr): + if isinstance(expr, EssentialBC): + assert expr.lhs == self.target return Eq( - self.arrays[self.target]['x'], eq.rhs, - subdomain=eq.subdomain + self.arrays[self.target]['x'], expr.rhs, + subdomain=expr.subdomain ) else: return None diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 10e030b064..c5f9c42b43 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -100,12 +100,12 @@ def test_petsc_solve(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - matvec_callback = [root for root in callable_roots if root.name == 'MatMult0'] + matvec_efunc = [root for root in callable_roots if root.name == 'MatMult0'] - formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS0'] + b_efunc = [root for root in callable_roots if root.name == 'FormRHS0'] - action_expr = FindNodes(Expression).visit(matvec_callback[0]) - rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) + action_expr = FindNodes(Expression).visit(matvec_efunc[0]) + rhs_expr = FindNodes(Expression).visit(b_efunc[0]) assert str(action_expr[-1].expr.rhs) == ( '(x_f[x + 1, y + 2]/ctx0->h_x**2' @@ -127,7 +127,7 @@ def test_petsc_solve(): assert len(retrieve_iteration_tree(op)) == 0 # TODO: Remove pragmas from PETSc callback functions - assert len(matvec_callback[0].parameters) == 3 + assert len(matvec_efunc[0].parameters) == 3 @skipif('petsc') From 312c8266aa14468ccabb0f633a6954dce1c8c2ad Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 17:18:41 +0100 Subject: [PATCH 26/53] misc: Add BaseJacobian --- devito/petsc/types/types.py | 216 ++++++++++++++++++++++-------------- 1 file changed, 131 insertions(+), 85 deletions(-) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index f4e92f2f2e..5e9796a5d8 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -137,6 +137,26 @@ def eval(cls, *args): class FieldData: + """ + Metadata class passed to `LinearSolveExpr`. Encapsulates metadata for a single + `target` field needed to interface with PETSc SNES solvers. + + Parameters + ---------- + + target : Function-like + The target field to solve into, which is a Function-like object. + jacobian : Jacobian + Defines the matrix-vector product for the linear system, where the vector is + the PETScArray representing the `target`. + residual : Residual + Defines the nonlinear residual function F(target) = 0. + initialguess : InitialGuess + Defines the initial guess for the solution, which satisfies + essential boundary conditions. + arrays : dict + A dictionary mapping `target` to its corresponding PETScArrays. + """ def __init__(self, target=None, jacobian=None, residual=None, initialguess=None, arrays=None, **kwargs): self._target = target @@ -190,6 +210,24 @@ def targets(self): class MultipleFieldData(FieldData): + """ + Metadata class passed to `LinearSolveExpr`, for mixed-field problems, + where the solution vector spans multiple `targets`. + + Parameters + ---------- + targets : list of Function-like + The fields to solve into, each represented by a Function-like object. + jacobian : MixedJacobian + Defines the matrix-vector products for the full system Jacobian. + residual : MixedResidual + Defines the nonlinear residual function F(targets) = 0. + initialguess : InitialGuess + Defines the initial guess metadata, which satisfies + essential boundary conditions. + arrays : dict + A dictionary mapping the `targets` to their corresponding PETScArrays. + """ def __init__(self, targets, arrays, jacobian=None, residual=None): self._targets = as_tuple(targets) self._arrays = arrays @@ -208,6 +246,7 @@ def space_dimensions(self): @property def grid(self): + """The unique `Grid` associated with all targets.""" grids = [t.grid for t in self.targets] if len(set(grids)) > 1: raise ValueError( @@ -233,8 +272,70 @@ def space_order(self): def targets(self): return self._targets + +class BaseJacobian: + def __init__(self, arrays, target=None): + self.arrays = arrays + self.target = target + + def _scale_non_bcs(self, matvecs, target=None): + target = target or self.target + vol = target.grid.symbolic_volume_cell + + return [ + m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) + for m in matvecs + ] + + def _compute_scdiag(self, matvecs, col_target=None): + """ + """ + x = self.arrays[col_target or self.target]['x'] + + centres = { + centre_stencil(m.rhs, x, as_coeff=True) + for m in matvecs if not isinstance(m, EssentialBC) + } + return centres.pop() if len(centres) == 1 else 1.0 + + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the essential BCs + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + + def _build_matvec_expr(self, expr, **kwargs): + col_target = kwargs.get('col_target', self.target) + row_target = kwargs.get('row_target', self.target) + + _, F_target, _, targets = separate_eqn(expr, col_target) + if F_target: + return self._make_matvec( + expr, F_target, targets, col_target, row_target + ) + else: + return (None,) + + def _make_matvec(self, expr, F_target, targets, col_target, row_target): + y = self.arrays[row_target]['y'] + x = self.arrays[col_target]['x'] + + if isinstance(expr, EssentialBC): + # NOTE: Essential BCs are trivial equations in the solver. + # See `EssentialBC` for more details. + zero_row = ZeroRow(y, x, subdomain=expr.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(x, targets)) + rhs = rhs.subs(self.time_mapper) + return (Eq(y, rhs, subdomain=expr.subdomain),) + -class Jacobian: +class Jacobian(BaseJacobian): """ Represents a Jacobian matrix. @@ -246,14 +347,14 @@ class Jacobian: require explicit symbolic differentiation. """ def __init__(self, target, exprs, arrays, time_mapper): - self.target = target + super().__init__(arrays=arrays, target=target) self.exprs = exprs - self.arrays = arrays self.time_mapper = time_mapper self._build_matvecs() @property def matvecs(self): + # TODO: add shortcut explanation etc """ Stores the expressions used to generate the `MatMult` callback generated at the IET level. This function is @@ -279,7 +380,6 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_expr(eq) if e is not None ) - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -289,83 +389,8 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_expr(self, expr, col_target=None, row_target=None): - col_target = col_target or self.target - row_target = row_target or self.target - _, F_target, _, targets = separate_eqn(expr, col_target) - if F_target: - return self._make_matvec( - expr, F_target, targets, col_target, row_target - ) - else: - return (None,) - - def _make_matvec(self, expr, F_target, targets, col_target, row_target): - y = self.arrays[row_target]['y'] - x = self.arrays[col_target]['x'] - - if isinstance(expr, EssentialBC): - # NOTE: Essential BCs are trivial equations in the solver. - # See `EssentialBC` for more details. - zero_row = ZeroRow(y, x, subdomain=expr.subdomain) - zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(x, targets)) - rhs = rhs.subs(self.time_mapper) - return (Eq(y, rhs, subdomain=expr.subdomain),) - - def _scale_non_bcs(self, matvecs, target=None): - target = target or self.target - vol = target.grid.symbolic_volume_cell - - return [ - m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) - for m in matvecs - ] - - def _compute_scdiag(self, matvecs, col_target=None): - """ - """ - x = self.arrays[col_target or self.target]['x'] - - centres = { - centre_stencil(m.rhs, x, as_coeff=True) - for m in matvecs if not isinstance(m, EssentialBC) - } - return centres.pop() if len(centres) == 1 else 1.0 - - def _scale_bcs(self, matvecs, scdiag): - """ - Scale the essential BCs - """ - return [ - m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m - for m in matvecs - ] - - -class SubMatrixBlock: - def __init__(self, name, matvecs, scdiag, row_target, - col_target, row_idx, col_idx, linear_idx): - self.name = name - self.matvecs = matvecs - self.scdiag = scdiag - self.row_target = row_target - self.col_target = col_target - self.row_idx = row_idx - self.col_idx = col_idx - self.linear_idx = linear_idx - - def is_diag(self): - return self.row_idx == self.col_idx - - def __repr__(self): - return (f"") - - -class MixedJacobian(Jacobian): +class MixedJacobian(BaseJacobian): """ Represents a Jacobian for a linear system with a solution vector composed of multiple fields (targets). @@ -380,12 +405,12 @@ class MixedJacobian(Jacobian): # TODO: pcfieldsplit support for each block """ - def __init__(self, target_eqns, arrays, time_mapper): - self.targets = tuple(target_eqns.keys()) - self.arrays = arrays + def __init__(self, target_exprs, arrays, time_mapper): + super().__init__(arrays=arrays, target=None) + self.targets = tuple(target_exprs.keys()) self.time_mapper = time_mapper self._submatrices = [] - self._build_blocks(target_eqns) + self._build_blocks(target_exprs) @property def submatrices(self): @@ -427,7 +452,9 @@ def _build_blocks(self, target_exprs): matvecs = [] for expr in exprs: matvecs.extend( - e for e in self._build_matvec_expr(expr, col_target, row_target) + e for e in self._build_matvec_expr( + expr, col_target=col_target, row_target=row_target + ) ) matvecs = [m for m in matvecs if m is not None] @@ -469,6 +496,25 @@ def __repr__(self): return f"" +class SubMatrixBlock: + def __init__(self, name, matvecs, scdiag, row_target, + col_target, row_idx, col_idx, linear_idx): + self.name = name + self.matvecs = matvecs + self.scdiag = scdiag + self.row_target = row_target + self.col_target = col_target + self.row_idx = row_idx + self.col_idx = col_idx + self.linear_idx = linear_idx + + def is_diag(self): + return self.row_idx == self.col_idx + + def __repr__(self): + return (f"") + + class Residual: """ Gennerates the metadata needed to define the nonlinear residual function From 048f693472d9b2b59673c33b93579592105cbc62 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:35:34 +0100 Subject: [PATCH 27/53] misc: Docstrings, stop list output for PETScSolve, tests --- devito/petsc/iet/routines.py | 12 ++-- devito/petsc/solve.py | 14 ++-- devito/petsc/types/types.py | 70 ++++++++++-------- examples/petsc/cfd/01_navierstokes.py | 2 +- .../petsc/seismic/01_staggered_acoustic.py | 4 +- tests/test_petsc.py | 71 +++++++++++++++---- 6 files changed, 119 insertions(+), 54 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f762c3d30b..f87eaab78b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -129,11 +129,13 @@ def _make_core(self): def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation matvecs = jacobian.matvecs - irs, _ = self.rcompile(matvecs, - options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper) - body = self._create_matvec_body(List(body=irs.uiet.body), - jacobian) + irs, _ = self.rcompile( + matvecs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._create_matvec_body( + List(body=irs.uiet.body), jacobian + ) objs = self.objs cb = PETScCallable( diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 949328cf67..1f4423707e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -73,7 +73,7 @@ def build_expr(self): time_mapper=self.time_mapper, localinfo=localinfo ) - return [PetscEq(target, linear_solve)] + return PetscEq(target, linear_solve) def linear_solve_args(self): target, exprs = next(iter(self.target_exprs.items())) @@ -98,14 +98,16 @@ def linear_solve_args(self): ) return target, tuple(funcs), field_data - + def generate_arrays(self, *targets): return { t: { - p: PETScArray(name=f'{p}_{t.name}', - target=t, - liveness='eager', - localinfo=localinfo) + p: PETScArray( + name=f'{p}_{t.name}', + target=t, + liveness='eager', + localinfo=localinfo + ) for p in prefixes } for t in targets diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 5e9796a5d8..b97b6ed406 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -138,8 +138,8 @@ def eval(cls, *args): class FieldData: """ - Metadata class passed to `LinearSolveExpr`. Encapsulates metadata for a single - `target` field needed to interface with PETSc SNES solvers. + Metadata for a single `target` field passed to `LinearSolveExpr`. + Used to interface with PETSc SNES solvers at the IET level. Parameters ---------- @@ -147,10 +147,10 @@ class FieldData: target : Function-like The target field to solve into, which is a Function-like object. jacobian : Jacobian - Defines the matrix-vector product for the linear system, where the vector is + Defines the matrix-vector product for the linear system, where the vector is the PETScArray representing the `target`. residual : Residual - Defines the nonlinear residual function F(target) = 0. + Defines the residual function F(target) = 0. initialguess : InitialGuess Defines the initial guess for the solution, which satisfies essential boundary conditions. @@ -212,7 +212,8 @@ def targets(self): class MultipleFieldData(FieldData): """ Metadata class passed to `LinearSolveExpr`, for mixed-field problems, - where the solution vector spans multiple `targets`. + where the solution vector spans multiple `targets`. Used to interface + with PETSc SNES solvers at the IET level. Parameters ---------- @@ -221,7 +222,7 @@ class MultipleFieldData(FieldData): jacobian : MixedJacobian Defines the matrix-vector products for the full system Jacobian. residual : MixedResidual - Defines the nonlinear residual function F(targets) = 0. + Defines the residual function F(targets) = 0. initialguess : InitialGuess Defines the initial guess metadata, which satisfies essential boundary conditions. @@ -272,13 +273,17 @@ def space_order(self): def targets(self): return self._targets - + class BaseJacobian: def __init__(self, arrays, target=None): self.arrays = arrays self.target = target def _scale_non_bcs(self, matvecs, target=None): + """ + Scale the symbolic expressions `matvecs` by the grid cell volume, + excluding EssentialBCs. + """ target = target or self.target vol = target.grid.symbolic_volume_cell @@ -287,8 +292,20 @@ def _scale_non_bcs(self, matvecs, target=None): for m in matvecs ] + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the EssentialBCs in `matvecs` by `scdiag`. + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + def _compute_scdiag(self, matvecs, col_target=None): """ + Compute the diagonal scaling factor from the symbolic matrix-vector + expressions in `matvecs`. If the centre stencil (i.e. the diagonal term of the + matrix) is not unique, defaults to 1.0. """ x = self.arrays[col_target or self.target]['x'] @@ -298,15 +315,6 @@ def _compute_scdiag(self, matvecs, col_target=None): } return centres.pop() if len(centres) == 1 else 1.0 - def _scale_bcs(self, matvecs, scdiag): - """ - Scale the essential BCs - """ - return [ - m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m - for m in matvecs - ] - def _build_matvec_expr(self, expr, **kwargs): col_target = kwargs.get('col_target', self.target) row_target = kwargs.get('row_target', self.target) @@ -356,9 +364,9 @@ def __init__(self, target, exprs, arrays, time_mapper): def matvecs(self): # TODO: add shortcut explanation etc """ - Stores the expressions used to generate the `MatMult` - callback generated at the IET level. This function is - passed to PETSc via `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. + Stores the expressions used to generate the `MatMult` callback generated + at the IET level. This function is passed to PETSc via + `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. """ return self._matvecs @@ -517,7 +525,7 @@ def __repr__(self): class Residual: """ - Gennerates the metadata needed to define the nonlinear residual function + Generates the metadata needed to define the nonlinear residual function F(target) = 0 for use with PETSc's SNES interface. PETSc's SNES interface includes methods for solving nonlinear systems of @@ -610,6 +618,8 @@ def _scale_bcs(self, eq, scdiag=None): class MixedResidual(Residual): """ + Generates the metadata needed to define the nonlinear residual function + F(targets) = 0 for use with PETSc's SNES interface. """ def __init__(self, target_exprs, arrays, time_mapper, scdiag): self.targets = tuple(target_exprs.keys()) @@ -621,6 +631,9 @@ def __init__(self, target_exprs, arrays, time_mapper, scdiag): @property def b_exprs(self): """ + For mixed solvers, a callback to form the RHS vector `b` is not generated, + a single residual callback is generated to compute F(targets). + TODO: Investigate if this is optimal. """ return None @@ -629,10 +642,10 @@ def _build_exprs(self, target_exprs): for t, exprs in target_exprs.items(): residual_exprs.extend( - chain.from_iterable(self._build_residual(e, t) - for e in as_tuple(exprs) - )) - + chain.from_iterable( + self._build_residual(e, t) for e in as_tuple(exprs) + ) + ) self._F_exprs = tuple(sorted( residual_exprs, key=lambda e: not isinstance(e, EssentialBC) )) @@ -671,7 +684,9 @@ def _build_residual(self, expr, target): class InitialGuess: """ - Enforce initial guess to satisfy essential BCs. + Metadata passed to `LinearSolveExpr` to define the initial guess + symbolic expressions, enforcing the initial guess to satisfy essential + boundary conditions. """ def __init__(self, target, exprs, arrays): self.target = target @@ -680,13 +695,12 @@ def __init__(self, target, exprs, arrays): @property def exprs(self): - """ - """ return self._exprs def _build_exprs(self, exprs): """ - Return a list of initial guess expressions. + Return a list of initial guess symbolic expressions + that satisfy essential boundary conditions. """ self._exprs = tuple([ eq for eq in diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 7797321045..7bd59a0c83 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -288,7 +288,7 @@ def neumann_right(eq, subdomain): tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) -exprs = tentu + tentv + eqn_p + [update_u, update_v] + bc_u1 + bc_v1 +exprs = [tentu, tentv, eqn_p, update_u, update_v, bc_u1, bc_v1] with switchconfig(language='petsc'): op = Operator(exprs) diff --git a/examples/petsc/seismic/01_staggered_acoustic.py b/examples/petsc/seismic/01_staggered_acoustic.py index 4f0c91ed49..fc9e75938d 100644 --- a/examples/petsc/seismic/01_staggered_acoustic.py +++ b/examples/petsc/seismic/01_staggered_acoustic.py @@ -65,7 +65,7 @@ petsc_p_2 = PETScSolve(p_2, target=p2.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): - op_2 = Operator(petsc_v_x_2 + petsc_v_z_2 + petsc_p_2 + src_p_2, opt='noop') + op_2 = Operator([petsc_v_x_2, petsc_v_z_2, petsc_p_2, src_p_2], opt='noop') op_2(time=src.time_range.num-1, dt=dt) norm_p2 = norm(p2) @@ -90,7 +90,7 @@ petsc_p_4 = PETScSolve(p_4, target=p4.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): - op_4 = Operator(petsc_v_x_4 + petsc_v_z_4 + petsc_p_4 + src_p_4, opt='noop') + op_4 = Operator([petsc_v_x_4, petsc_v_z_4, petsc_p_4, src_p_4], opt='noop') op_4(time=src.time_range.num-1, dt=dt) norm_p4 = norm(p4) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c5f9c42b43..2352a83b91 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -150,7 +150,7 @@ def test_multiple_petsc_solves(): petsc2 = PETScSolve(eqn2, f2) with switchconfig(language='petsc'): - op = Operator(petsc1+petsc2, opt='noop') + op = Operator([petsc1, petsc2], opt='noop') callable_roots = [meta_call.root for meta_call in op._func_table.values()] @@ -320,7 +320,7 @@ def test_petsc_struct(): eqn2 = Eq(f1, g1*mu2) with switchconfig(language='petsc'): - op = Operator([eqn2] + petsc1) + op = Operator([eqn2, petsc1]) arguments = op.arguments() @@ -507,7 +507,7 @@ def test_time_loop(): petsc5 = PETScSolve(eq5, v2) with switchconfig(language='petsc'): - op4 = Operator(petsc4 + petsc5) + op4 = Operator([petsc4, petsc5]) op4.apply(time_M=3) body4 = str(op4.body) @@ -609,6 +609,53 @@ def define(self, dimensions): assert np.allclose(u.data[-1, 1:-1], 4.0) # right +@skipif('petsc') +def test_jacobian(): + + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + sub1 = SubLeft() + sub2 = SubRight() + + grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) + + e = Function(name='e', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + bc_1 = EssentialBC(e, 1.0, subdomain=sub1) + bc_2 = EssentialBC(e, 2.0, subdomain=sub2) + + eq1 = Eq(e.laplace + e, f + 2.0) + + petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + + jac = petsc.rhs.fielddata.jacobian + + assert jac.row_target == e + assert jac.col_target == e + + # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). + # NOTE: this is likely to change when PetscSection + DMDA is supported + assert len(jac.matvecs) == 5 + # TODO: I think some internals are preventing symplification here? + assert str(jac.scdiag) == 'h_x*(1 - 2.0/h_x**2)' + + assert all(isinstance(m, EssentialBC) for m in jac.matvecs[:4]) + assert not isinstance(jac.matvecs[-1], EssentialBC) + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. @@ -655,7 +702,7 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): petsc2 = PETScSolve(eq2, target=g) with switchconfig(language='petsc'): - op1 = Operator(petsc1 + petsc2, opt='noop') + op1 = Operator([petsc1, petsc2], opt='noop') op1.apply() enorm1 = norm(e) @@ -689,9 +736,9 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): assert len(callbacks2) == 6 # Check fielddata type - fielddata1 = petsc1[0].rhs.fielddata - fielddata2 = petsc2[0].rhs.fielddata - fielddata3 = petsc3[0].rhs.fielddata + fielddata1 = petsc1.rhs.fielddata + fielddata2 = petsc2.rhs.fielddata + fielddata3 = petsc3.rhs.fielddata assert isinstance(fielddata1, FieldData) assert isinstance(fielddata2, FieldData) @@ -818,7 +865,7 @@ def test_mixed_jacobian(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j01 = jacobian.get_submatrix(0, 1) @@ -894,7 +941,7 @@ def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j01 = jacobian.get_submatrix(0, 1) j10 = jacobian.get_submatrix(1, 0) @@ -943,7 +990,7 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -993,7 +1040,7 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1046,7 +1093,7 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) From 905def2ca9cbdfaaf7dfeb3e8e446e7b39349e4d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:39:12 +0100 Subject: [PATCH 28/53] misc: Flake8 --- tests/test_petsc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 2352a83b91..87e58c2fe0 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -625,12 +625,12 @@ class SubRight(SubDomain): def define(self, dimensions): x, = dimensions return {x: ('right', 1)} - + sub1 = SubLeft() sub2 = SubRight() grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) - + e = Function(name='e', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) @@ -645,7 +645,7 @@ def define(self, dimensions): assert jac.row_target == e assert jac.col_target == e - + # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). # NOTE: this is likely to change when PetscSection + DMDA is supported assert len(jac.matvecs) == 5 From 19fddbe7fbdfa680d5782f3261116c0d02d3d85f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:54:14 +0100 Subject: [PATCH 29/53] misc: Fix exprs in 01_navierstokes.py due to change in PETScSolve output type --- examples/petsc/cfd/01_navierstokes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 7bd59a0c83..1c678d977b 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -288,7 +288,7 @@ def neumann_right(eq, subdomain): tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) -exprs = [tentu, tentv, eqn_p, update_u, update_v, bc_u1, bc_v1] +exprs = [tentu, tentv, eqn_p, update_u, update_v] + bc_u1 + bc_v1 with switchconfig(language='petsc'): op = Operator(exprs) From 8037cf9c002fe9d7db49b52d9667c4a533aa7c88 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 16:04:27 +0100 Subject: [PATCH 30/53] misc: Clean up --- devito/logger.py | 4 ++-- devito/operator/operator.py | 6 +----- devito/petsc/iet/passes.py | 19 +------------------ devito/petsc/iet/routines.py | 2 -- devito/petsc/initialize.py | 7 +++---- examples/petsc/petsc_test.py | 2 +- 6 files changed, 8 insertions(+), 32 deletions(-) diff --git a/devito/logger.py b/devito/logger.py index c9834dbfe7..e5df4bb565 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -77,9 +77,9 @@ def set_log_level(level, comm=None): """ from devito import configuration from devito.mpi.distributed import MPI - + if comm is not None and configuration['mpi']: - if comm!= MPI.COMM_NULL and comm.rank != 0: + if comm != MPI.COMM_NULL and comm.rank != 0: logger.removeHandler(stream_handler) logger.addHandler(logging.NullHandler()) else: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 5216bde481..3a36952a58 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -38,7 +38,7 @@ from devito.petsc.iet.passes import lower_petsc from devito.petsc.clusters import petsc_preprocess -__all__ = ['Operator', 'SpecialOp'] +__all__ = ['Operator'] class Operator(Callable): @@ -1456,7 +1456,3 @@ def parse_kwargs(**kwargs): kwargs['subs'] = {k: sympify(v) for k, v in kwargs.get('subs', {}).items()} return kwargs - - -class SpecialOp(Operator): - pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f77879cf10..bf9a2c12fe 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -93,17 +93,6 @@ def initialize(iet): help_string = c.Line(r'static char help[] = "This is help text.\n";') init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) - - # print_comm_size = - - # size = c.Line('PetscMPIInt size;') - # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') - # rank = c.Line('PetscMPIInt rank;') - # get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') - # print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - # print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - - init_body = CallableBody( body=(petsc_func_begin_user, help_string, init_body), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) @@ -128,20 +117,14 @@ def make_core_petsc_calls(objs, grid, **kwargs): comm = 'PETSC_COMM_WORLD' call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) - # from IPython import embed; embed() - # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + rank = c.Line('PetscMPIInt rank;') get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - - flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine - # return call_mpi, BlankLine class Builder: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f0e6a26756..7f8651e6f3 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1087,7 +1087,6 @@ def __init__(self, **kwargs): self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.grid = kwargs.get('grid') - # from IPython import embed; embed() self.devito_mpi = kwargs['options'].get('mpi', False) self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1129,7 +1128,6 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - # from IPython import embed; embed() if self.devito_mpi: base_dict['comm'] = self.grid.distributor._obj_comm else: diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 3fc419c26a..80e3c7520c 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -3,8 +3,7 @@ from ctypes import POINTER, cast, c_char import atexit -from devito import switchconfig -from devito.operator.operator import SpecialOp +from devito import Operator, switchconfig from devito.types import Symbol from devito.types.equation import PetscEq from devito.petsc.types import Initialize, Finalize @@ -22,11 +21,11 @@ def PetscInitialize(): # This would prevent circular imports when initializing during import # from the PETSc module. with switchconfig(language='petsc'): - op_init = SpecialOp( + op_init = Operator( [PetscEq(dummy, Initialize(dummy))], name='kernel_init', opt='noop' ) - op_finalize = SpecialOp( + op_finalize = Operator( [PetscEq(dummy, Finalize(dummy))], name='kernel_finalize', opt='noop' ) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index eb74c639db..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -# print(op.ccode) +print(op.ccode) From 8c16f2e112743f896613f88b53d521eeb922d02b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 16:51:56 +0100 Subject: [PATCH 31/53] misc: Utilise zero_vector function --- devito/petsc/iet/routines.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f87eaab78b..fe9cce66f5 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -188,9 +188,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) - zero_ylocal_memory = petsc_call( - 'VecSet', [ylocal, 0.0] - ) + zero_ylocal_memory = self.zero_vector(ylocal) vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] @@ -730,9 +728,7 @@ def _whole_matvec_body(self): nonzero_submats = self.jacobian.nonzero_submatrices - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) + zero_y_memory = self.zero_vector(objs['Y']) calls = () for sm in nonzero_submats: From f3d4d2deacd95184d26b6ee0ed3ef21cb7136856 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 22:56:59 +0100 Subject: [PATCH 32/53] misc: Move vecset to function inside iet/utils.py --- devito/petsc/iet/routines.py | 38 +++++---------------------- devito/petsc/iet/utils.py | 4 +++ devito/petsc/types/types.py | 9 +++++-- examples/petsc/cfd/01_navierstokes.py | 5 ++++ 4 files changed, 22 insertions(+), 34 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f87eaab78b..9686c91d5e 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,7 @@ from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) -from devito.petsc.iet.utils import petsc_call, petsc_struct +from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, @@ -99,13 +99,6 @@ def initialguesses(self): def user_struct_callback(self): return self._user_struct_callback - @property - def zero_memory(self): - """Indicates whether the memory of the output - vector should be set to zero before the computation - in the callback.""" - return True - @property def fielddata(self): return self.injectsolve.expr.rhs.fielddata @@ -169,7 +162,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_y_memory = self.zero_vector(objs['Y']) + zero_y_memory = zero_vector(objs['Y']) if jacobian.zero_memory else None dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] @@ -188,9 +181,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) - zero_ylocal_memory = petsc_call( - 'VecSet', [ylocal, 0.0] - ) + zero_ylocal_memory = zero_vector(ylocal) vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] @@ -320,7 +311,7 @@ def _create_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = self.zero_vector(objs['F']) + zero_f_memory = zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] @@ -652,12 +643,6 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper - def zero_vector(self, vec): - """ - Zeros the memory of the output vector before computation - """ - return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None - def dereference_funcs(self, struct, fields): return tuple( [Dereference(i, struct) for i in @@ -691,13 +676,6 @@ def main_matvec_callback(self): def main_formfunc_callback(self): return self._main_formfunc_callback - @property - def zero_memory(self): - """Indicates whether the memory of the output - vector should be set to zero before the computation - in the callback.""" - return False - def _make_core(self): for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') @@ -730,9 +708,7 @@ def _whole_matvec_body(self): nonzero_submats = self.jacobian.nonzero_submatrices - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) + zero_y_memory = zero_vector(objs['Y']) calls = () for sm in nonzero_submats: @@ -815,9 +791,7 @@ def _whole_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = petsc_call( - 'VecSet', [objs['F'], 0.0] - ) + zero_f_memory = zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 99da0468ad..d143bcc8c9 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -18,6 +18,10 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): modifier=modifier) +def zero_vector(vec): + return petsc_call('VecSet', [vec, 0.0]) + + # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. petsc_iet_mapper = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index b97b6ed406..4f235bbbde 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -382,6 +382,10 @@ def row_target(self): def col_target(self): return self.target + @property + def zero_memory(self): + return True + def _build_matvecs(self): matvecs = [] for eq in self.exprs: @@ -516,8 +520,9 @@ def __init__(self, name, matvecs, scdiag, row_target, self.col_idx = col_idx self.linear_idx = linear_idx - def is_diag(self): - return self.row_idx == self.col_idx + @property + def zero_memory(self): + return False def __repr__(self): return (f"") diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 1c678d977b..bd591c1eee 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -297,3 +297,8 @@ def neumann_right(eq, subdomain): # Pressure norm check tol = 1e-3 assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol + +from devito import norm +print(norm(u1)) +print(norm(v1)) +print(norm(pn1)) From 8d10b8718e6702023e84e3d6c3de8e5d59001218 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 22:59:21 +0100 Subject: [PATCH 33/53] misc: Flake8 --- examples/petsc/cfd/01_navierstokes.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index bd591c1eee..1c678d977b 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -297,8 +297,3 @@ def neumann_right(eq, subdomain): # Pressure norm check tol = 1e-3 assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol - -from devito import norm -print(norm(u1)) -print(norm(v1)) -print(norm(pn1)) From 53b58a92e501b5e1db53a343187a84da445bc0f4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 23:50:10 +0100 Subject: [PATCH 34/53] misc: Clean up --- devito/petsc/iet/routines.py | 47 +++++++++++------------------------- devito/petsc/iet/utils.py | 45 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 33 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9686c91d5e..ec78054094 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,8 @@ from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) -from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector +from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, + dereference_funcs, residual_bundle) from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, @@ -120,7 +121,7 @@ def _make_core(self): self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): - # Compile matvec `eqns` into an IET via recursive compilation + # Compile `matvecs` into an IET via recursive compilation matvecs = jacobian.matvecs irs, _ = self.rcompile( matvecs, options={'mpi': False}, sregistry=self.sregistry, @@ -251,7 +252,7 @@ def _create_matvec_body(self, body, jacobian): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=body), @@ -390,7 +391,7 @@ def _create_formfunc_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=body), @@ -500,7 +501,7 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), @@ -578,7 +579,7 @@ def _create_initial_guess_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), @@ -643,12 +644,6 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper - def dereference_funcs(self, struct, fields): - return tuple( - [Dereference(i, struct) for i in - fields if isinstance(i.function, AbstractFunction)] - ) - class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -749,17 +744,17 @@ def _whole_matvec_body(self): def _make_whole_formfunc(self): F_exprs = self.fielddata.residual.F_exprs - # Compile formfunc `eqns` into an IET via recursive compilation - irs_formfunc, _ = self.rcompile( + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) + body = self._whole_formfunc_body(List(body=irs.uiet.body)) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), - body_formfunc, + body, retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) @@ -783,7 +778,8 @@ def _whole_formfunc_body(self, body): bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] - body = self.residual_bundle(body, bundles) + + body = residual_bundle(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -870,7 +866,7 @@ def _whole_formfunc_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -1034,21 +1030,6 @@ def _submat_callback_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def residual_bundle(self, body, bundles): - mapper = bundles['bundle_mapper'] - indexeds = FindSymbols('indexeds').visit(body) - subs = {} - - for i in indexeds: - if i.base in mapper: - bundle = mapper[i.base] - index = bundles['target_indices'][i.function.target] - index = (index,) + i.indices - subs[i] = bundle.__getitem__(index) - - body = Uxreplace(subs).visit(body) - return body - class BaseObjectBuilder: """ diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index d143bcc8c9..3402cfe28f 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,5 +1,7 @@ from devito.petsc.iet.nodes import PetscMetaData, PETScCall from devito.ir.equations import OpPetsc +from devito.ir.iet import Dereference, FindSymbols, Uxreplace +from devito.types.basic import AbstractFunction def petsc_call(specific_call, call_args): @@ -19,9 +21,52 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): def zero_vector(vec): + """ + Set all entries of a PETSc vector to zero. + """ return petsc_call('VecSet', [vec, 0.0]) +def dereference_funcs(struct, fields): + """ + Dereference AbstractFunctions from a struct. + """ + return tuple( + [Dereference(i, struct) for i in + fields if isinstance(i.function, AbstractFunction)] + ) + + +def residual_bundle(body, bundles): + """ + Replaces PetscArrays in `body` with PetscBundle struct field accesses + (e.g., f_v[ix][iy] -> f_bundle[ix][iy].v). + + Example: + f_v[ix][iy] = x_v[ix][iy]; + f_u[ix][iy] = x_u[ix][iy]; + becomes: + f_bundle[ix][iy].v = x_bundle[ix][iy].v; + f_bundle[ix][iy].u = x_bundle[ix][iy].u; + + NOTE: This is used because the data is interleaved for + multi-component DMDAs in PETSc. + """ + mapper = bundles['bundle_mapper'] + indexeds = FindSymbols('indexeds').visit(body) + subs = {} + + for i in indexeds: + if i.base in mapper: + bundle = mapper[i.base] + index = bundles['target_indices'][i.function.target] + index = (index,) + i.indices + subs[i] = bundle.__getitem__(index) + + body = Uxreplace(subs).visit(body) + return body + + # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. petsc_iet_mapper = {OpPetsc: PetscMetaData} From 857eb9cebc711a36b3608e56433ec78f78e46e04 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 15 Jun 2025 20:55:46 +0100 Subject: [PATCH 35/53] misc: Add more tests --- tests/test_petsc.py | 83 +++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 87e58c2fe0..678b1071c8 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -12,7 +12,8 @@ from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr, FieldData, MultipleFieldData) + LinearSolveExpr, FieldData, MultipleFieldData, + SubMatrixBlock) from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -647,7 +648,7 @@ def define(self, dimensions): assert jac.col_target == e # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). - # NOTE: this is likely to change when PetscSection + DMDA is supported + # NOTE: This is likely to change when PetscSection + DMDA is supported assert len(jac.matvecs) == 5 # TODO: I think some internals are preventing symplification here? assert str(jac.scdiag) == 'h_x*(1 - 2.0/h_x**2)' @@ -656,6 +657,48 @@ def define(self, dimensions): assert not isinstance(jac.matvecs[-1], EssentialBC) +@skipif('petsc') +def test_residual(): + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + sub1 = SubLeft() + sub2 = SubRight() + + grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) + + e = Function(name='e', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + bc_1 = EssentialBC(e, 1.0, subdomain=sub1) + bc_2 = EssentialBC(e, 2.0, subdomain=sub2) + + eq1 = Eq(e.laplace + e, f + 2.0) + + petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + + res = petsc.rhs.fielddata.residual + + assert res.target == e + # NOTE: This is likely to change when PetscSection + DMDA is supported + assert len(res.F_exprs) == 5 + assert len(res.b_exprs) == 3 + + assert not res.time_mapper + assert str(res.scdiag) == 'h_x*(1 - 2.0/h_x**2)' + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. @@ -872,6 +915,42 @@ def test_mixed_jacobian(self): j10 = jacobian.get_submatrix(1, 0) j11 = jacobian.get_submatrix(1, 1) + # Check type of each submatrix is a SubMatrixBlock + assert isinstance(j00, SubMatrixBlock) + assert isinstance(j01, SubMatrixBlock) + assert isinstance(j10, SubMatrixBlock) + assert isinstance(j11, SubMatrixBlock) + + assert j00.name == 'J00' + assert j01.name == 'J01' + assert j10.name == 'J10' + assert j11.name == 'J11' + + assert j00.row_target == e + assert j01.row_target == e + assert j10.row_target == g + assert j11.row_target == g + + assert j00.col_target == e + assert j01.col_target == g + assert j10.col_target == e + assert j11.col_target == g + + assert j00.row_idx == 0 + assert j01.row_idx == 0 + assert j10.row_idx == 1 + assert j11.row_idx == 1 + + assert j00.col_idx == 0 + assert j01.col_idx == 1 + assert j10.col_idx == 0 + assert j11.col_idx == 1 + + assert j00.linear_idx == 0 + assert j01.linear_idx == 1 + assert j10.linear_idx == 2 + assert j11.linear_idx == 3 + # Check the number of submatrices assert jacobian.n_submatrices == 4 From 7a5b10fd84275d21b8056ece5bb1af46cb0303c5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 09:44:06 +0100 Subject: [PATCH 36/53] mpi: Start parallel tests --- tests/test_petsc.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 87e58c2fe0..daa6b6d734 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1232,3 +1232,15 @@ def define(self, dimensions): + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) # TODO: Test mixed, time dependent solvers + + +class TestMPI: + + @pytest.mark.parallel(mode=4) + def test_laplacian(self, mode): + """ + """ + grid = Grid(shape=(4,)) + + f = Function(name='f') + From 467c19babb6e28484ea11e75c955a8d77e7174a8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 10:12:26 +0100 Subject: [PATCH 37/53] add to petsc laplacian test --- devito/petsc/iet/passes.py | 3 +- tests/test_petsc.py | 80 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 78 insertions(+), 5 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index bf9a2c12fe..f6af3f0a4f 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -124,7 +124,8 @@ def make_core_petsc_calls(objs, grid, **kwargs): print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + # return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + return call_mpi, BlankLine class Builder: diff --git a/tests/test_petsc.py b/tests/test_petsc.py index daa6b6d734..222de35282 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1236,11 +1236,83 @@ def define(self, dimensions): class TestMPI: - @pytest.mark.parallel(mode=4) - def test_laplacian(self, mode): + # @pytest.mark.parallel(mode=4) + @skipif('petsc') + def test_laplacian(self): """ """ - grid = Grid(shape=(4,)) - f = Function(name='f') + # Subdomains to implement BCs + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + + sub1 = SubLeft() + sub2 = SubRight() + subdomains = (sub1, sub2,) + + + def exact(x): + return -np.float64(np.exp(x)) + + + # n = 9, 17, 33, 65, 129, 257 + n_values = [2**k + 1 for k in range(3, 9)] + n_values = [9] + dx = np.array([1./(n-1) for n in n_values]) + errors = [] + + for n in n_values: + + grid = Grid( + shape=(n,), subdomains=subdomains, dtype=np.float64 + ) + + u = Function(name='u', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + bc = Function(name='bc', grid=grid, space_order=2) + + eqn = Eq(-u.laplace, f, subdomain=grid.interior) + + X = np.linspace(0, 1.0, n).astype(np.float64) + f.data[:] = np.float64(np.exp(X)) + + bc.data[0] = -np.float64(1.0) # u(0) = -1 + bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e + + # Create boundary condition expressions using subdomains + bcs = [EssentialBC(u, bc, subdomain=sub1)] + bcs += [EssentialBC(u, bc, subdomain=sub2)] + + exprs = [eqn] + bcs + petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + + op = Operator(petsc, language='petsc') + op.apply() + + # u_exact = Function(name='u_exact', grid=grid, space_order=2) + # u_exact.data[:] = exact(X) + + # diff = u_exact.data[:] - u.data[:] + # u_diff_norm = np.linalg.norm(diff, ord=np.inf) + # u_error = u_diff_norm / np.linalg.norm(u_exact.data[:], ord=np.inf) + # errors.append(u_error) + + # Expected norms computed "manually" from sequential runs + norm_u = norm(u) + assert norm_u == 2.0 + # Expected norms computed "manually" from sequential runs + # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file From d2e3eb59260fecc9020160db495a3e6ee7a2d110 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 10:28:53 +0100 Subject: [PATCH 38/53] mpi tests --- tests/test_petsc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index b01dcac2f1..c0e36919dc 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1317,7 +1317,8 @@ class TestMPI: # @pytest.mark.parallel(mode=4) @skipif('petsc') - def test_laplacian(self): + @pytest.mark.parallel(mode=[(2)]) + def test_laplacian(self, mode): """ """ @@ -1391,7 +1392,7 @@ def exact(x): # Expected norms computed "manually" from sequential runs norm_u = norm(u) - assert norm_u == 2.0 + assert norm_u == 5.467052700706644 # Expected norms computed "manually" from sequential runs # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file From 6165373e003ee4e17c01a0b7e66c61fd07122f53 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 12:48:13 +0100 Subject: [PATCH 39/53] edit test --- examples/petsc/petsc_test.py | 2 +- tests/test_petsc.py | 71 +++++++++++------------------------- 2 files changed, 23 insertions(+), 50 deletions(-) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..eb74c639db 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -print(op.ccode) +# print(op.ccode) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c0e36919dc..72823c4985 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1314,15 +1314,13 @@ def define(self, dimensions): class TestMPI: + # TODO: Add test for DMDACreate() in parallel - # @pytest.mark.parallel(mode=4) @skipif('petsc') - @pytest.mark.parallel(mode=[(2)]) + @pytest.mark.parallel(mode=1) def test_laplacian(self, mode): """ """ - - # Subdomains to implement BCs class SubLeft(SubDomain): name = 'subleft' @@ -1330,7 +1328,6 @@ def define(self, dimensions): x, = dimensions return {x: ('left', 1)} - class SubRight(SubDomain): name = 'subright' @@ -1343,56 +1340,32 @@ def define(self, dimensions): sub2 = SubRight() subdomains = (sub1, sub2,) + n = 9 - def exact(x): - return -np.float64(np.exp(x)) - - - # n = 9, 17, 33, 65, 129, 257 - n_values = [2**k + 1 for k in range(3, 9)] - n_values = [9] - dx = np.array([1./(n-1) for n in n_values]) - errors = [] - - for n in n_values: - - grid = Grid( - shape=(n,), subdomains=subdomains, dtype=np.float64 - ) + grid = Grid( + shape=(n,), subdomains=subdomains, dtype=np.float64 + ) - u = Function(name='u', grid=grid, space_order=2) - f = Function(name='f', grid=grid, space_order=2) - bc = Function(name='bc', grid=grid, space_order=2) - - eqn = Eq(-u.laplace, f, subdomain=grid.interior) - - X = np.linspace(0, 1.0, n).astype(np.float64) - f.data[:] = np.float64(np.exp(X)) - - bc.data[0] = -np.float64(1.0) # u(0) = -1 - bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e + u = Function(name='u', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + bc = Function(name='bc', grid=grid, space_order=2) - # Create boundary condition expressions using subdomains - bcs = [EssentialBC(u, bc, subdomain=sub1)] - bcs += [EssentialBC(u, bc, subdomain=sub2)] + eqn = Eq(-u.laplace, f, subdomain=grid.interior) - exprs = [eqn] + bcs - petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + X = np.linspace(0, 1.0, n).astype(np.float64) + f.data[:] = np.float64(np.exp(X)) - op = Operator(petsc, language='petsc') - op.apply() + bc.data[0] = -np.float64(1.0) # u(0) = -1 + bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e - # u_exact = Function(name='u_exact', grid=grid, space_order=2) - # u_exact.data[:] = exact(X) + # Create boundary condition expressions using subdomains + bcs = [EssentialBC(u, bc, subdomain=sub1)] + bcs += [EssentialBC(u, bc, subdomain=sub2)] - # diff = u_exact.data[:] - u.data[:] - # u_diff_norm = np.linalg.norm(diff, ord=np.inf) - # u_error = u_diff_norm / np.linalg.norm(u_exact.data[:], ord=np.inf) - # errors.append(u_error) + petsc = PETScSolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) - # Expected norms computed "manually" from sequential runs - norm_u = norm(u) - assert norm_u == 5.467052700706644 + op = Operator(petsc, language='petsc') + op.apply() - # Expected norms computed "manually" from sequential runs - # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file + # Expected norm computed "manually" from sequential run + assert np.isclose(norm(u), 5.467052700706644, rtol=1e-15, atol=1e-15) From 0b43ea68b5c9d6fc31d6b726abc078af729a8e2b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 18:33:55 +0100 Subject: [PATCH 40/53] tests: Add petsc parallel test --- requirements-testing.txt | 1 + tests/test_petsc.py | 33 +++++++++++++++++++++++++-------- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/requirements-testing.txt b/requirements-testing.txt index d10923cb00..2acf0e4f11 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -1,6 +1,7 @@ pytest>=7.2,<8.5 pytest-runner<6.0.2 pytest-cov<6.1.2 +pytest-order flake8-pyproject>=1.2.3,<1.2.4 nbval<0.11.1 scipy<1.15.3 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 72823c4985..1132ebd5fb 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -20,6 +20,7 @@ @skipif('petsc') +@pytest.mark.order(0) def test_petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized @@ -28,6 +29,14 @@ def test_petsc_initialization(): PetscInitialize() +@skipif('petsc') +@pytest.mark.parallel(mode=[1, 2, 4, 6]) +def test_petsc_initialization_parallel(mode): + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() + + @skipif('petsc') def test_petsc_local_object(): """ @@ -1316,11 +1325,22 @@ def define(self, dimensions): class TestMPI: # TODO: Add test for DMDACreate() in parallel + @pytest.mark.parametrize('nx, unorm', [ + (17, 7.441506654790017), + (33, 10.317652759863675), + (65, 14.445123374862874), + (129, 20.32492895656658), + (257, 28.67050632840985) + ]) @skipif('petsc') - @pytest.mark.parallel(mode=1) - def test_laplacian(self, mode): + @pytest.mark.parallel(mode=[1, 2, 4, 8]) + def test_laplacian_1d(self, nx, unorm, mode): """ """ + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() + class SubLeft(SubDomain): name = 'subleft' @@ -1335,15 +1355,12 @@ def define(self, dimensions): x, = dimensions return {x: ('right', 1)} - sub1 = SubLeft() sub2 = SubRight() subdomains = (sub1, sub2,) - n = 9 - grid = Grid( - shape=(n,), subdomains=subdomains, dtype=np.float64 + shape=(nx,), subdomains=subdomains, dtype=np.float64 ) u = Function(name='u', grid=grid, space_order=2) @@ -1352,7 +1369,7 @@ def define(self, dimensions): eqn = Eq(-u.laplace, f, subdomain=grid.interior) - X = np.linspace(0, 1.0, n).astype(np.float64) + X = np.linspace(0, 1.0, nx).astype(np.float64) f.data[:] = np.float64(np.exp(X)) bc.data[0] = -np.float64(1.0) # u(0) = -1 @@ -1368,4 +1385,4 @@ def define(self, dimensions): op.apply() # Expected norm computed "manually" from sequential run - assert np.isclose(norm(u), 5.467052700706644, rtol=1e-15, atol=1e-15) + assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) From 62d004538fc638792c5f3bd4f1e59e842264768a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 18:43:09 +0100 Subject: [PATCH 41/53] clean up --- devito/petsc/iet/passes.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f6af3f0a4f..f358696c7a 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -115,16 +115,8 @@ def make_core_petsc_calls(objs, grid, **kwargs): comm = grid.distributor._obj_comm else: comm = 'PETSC_COMM_WORLD' - - call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) - - rank = c.Line('PetscMPIInt rank;') - get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') - print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - # return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine From bd963794aa1879e50825af431c58fc490312873b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 19:42:05 +0100 Subject: [PATCH 42/53] misc: Simplify comm extraction in lower_petsc --- devito/petsc/iet/passes.py | 18 ++++++++---------- devito/petsc/iet/routines.py | 12 ++++++------ examples/petsc/petsc_test.py | 2 -- 3 files changed, 14 insertions(+), 18 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f358696c7a..f20f30991f 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -53,9 +53,11 @@ def lower_petsc(iet, **kwargs): if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") grid = unique_grids.pop() + devito_mpi = kwargs['options'].get('mpi', False) + comm = grid.distributor._obj_comm if devito_mpi else 'PETSC_COMM_WORLD' # Create core PETSc calls (not specific to each PETScSolve) - core = make_core_petsc_calls(objs, grid, **kwargs) + core = make_core_petsc_calls(objs, comm) setup = [] subs = {} @@ -63,7 +65,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, grid, **kwargs) + builder = Builder(injectsolve, objs, iters, grid, comm, **kwargs) setup.extend(builder.solversetup.calls) @@ -109,13 +111,7 @@ def finalize(iet): return iet._rebuild(body=finalize_body) -def make_core_petsc_calls(objs, grid, **kwargs): - devito_mpi = kwargs['options'].get('mpi', False) - if devito_mpi: - comm = grid.distributor._obj_comm - else: - comm = 'PETSC_COMM_WORLD' - +def make_core_petsc_calls(objs, comm): call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine @@ -129,11 +125,12 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, grid, **kwargs): + def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters self.grid = grid + self.comm = comm self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) self.args = { @@ -141,6 +138,7 @@ def __init__(self, injectsolve, objs, iters, grid, **kwargs): 'objs': self.objs, 'iters': self.iters, 'grid': self.grid, + 'comm': self.comm, **self.kwargs } self.args['solver_objs'] = self.objbuilder.solver_objs diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ecf52462d..f999b13e71 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1041,8 +1041,7 @@ def __init__(self, **kwargs): self.injectsolve = kwargs.get('injectsolve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') - self.grid = kwargs.get('grid') - self.devito_mpi = kwargs['options'].get('mpi', False) + self.comm = kwargs.get('comm') self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1083,10 +1082,11 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - if self.devito_mpi: - base_dict['comm'] = self.grid.distributor._obj_comm - else: - base_dict['comm'] = 'PETSC_COMM_WORLD' + # if self.devito_mpi: + # base_dict['comm'] = self.grid.distributor._obj_comm + # else: + # base_dict['comm'] = 'PETSC_COMM_WORLD' + base_dict['comm'] = self.comm self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index eb74c639db..821a3569fc 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -27,5 +27,3 @@ with switchconfig(language='petsc'): op = Operator(petsc) op.apply() - -# print(op.ccode) From ce5cf82672d006bc07fabfc563a17010fb25a1bf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 19:57:09 +0100 Subject: [PATCH 43/53] clean up --- devito/petsc/iet/passes.py | 6 ++---- devito/petsc/iet/routines.py | 5 ----- examples/petsc/petsc_test.py | 2 ++ tests/test_petsc.py | 1 + 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f20f30991f..b4859dda16 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -65,7 +65,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, grid, comm, **kwargs) + builder = Builder(injectsolve, objs, iters, comm, **kwargs) setup.extend(builder.solversetup.calls) @@ -125,11 +125,10 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): + def __init__(self, injectsolve, objs, iters, comm, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters - self.grid = grid self.comm = comm self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) @@ -137,7 +136,6 @@ def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, - 'grid': self.grid, 'comm': self.comm, **self.kwargs } diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f999b13e71..da6ded1f9a 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1081,11 +1081,6 @@ def _build(self): 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } - # TODO: Devito MPI + PETSc testing - # if self.devito_mpi: - # base_dict['comm'] = self.grid.distributor._obj_comm - # else: - # base_dict['comm'] = 'PETSC_COMM_WORLD' base_dict['comm'] = self.comm self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 821a3569fc..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -27,3 +27,5 @@ with switchconfig(language='petsc'): op = Operator(petsc) op.apply() + +print(op.ccode) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1132ebd5fb..a2f9d8b8ec 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1385,4 +1385,5 @@ def define(self, dimensions): op.apply() # Expected norm computed "manually" from sequential run + # What rtol and atol should be used? assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) From b7c4082554e77ce10321d0532c2b75706da5fa54 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 21:23:23 +0100 Subject: [PATCH 44/53] trigger petsc CI --- .github/workflows/pytest-petsc.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 2bf55b5049..502e9b2139 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -11,10 +11,12 @@ on: branches: - main - petsc + - biharmonic pull_request: branches: - main - petsc + - biharmonic jobs: pytest: From e485c9c002b0f13c0002636e52f5a1f64ad05198 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 12:47:57 +0100 Subject: [PATCH 45/53] misc: Add todo: --- devito/petsc/iet/routines.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index da6ded1f9a..9249c7a557 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1389,7 +1389,10 @@ def _create_dmda(self, dmda): # Number of degrees of freedom per node args.append(dmda.dofs) # "Stencil width" -> size of overlap + # TODO: Instead, this probably should be + # extracted from fielddata.target._size_outhalo? stencil_width = self.fielddata.space_order + args.append(stencil_width) args.extend([objs['Null']]*nspace_dims) From df4e638f738d3c3047ee7d3728447358f7bae363 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 15:56:51 +0100 Subject: [PATCH 46/53] address some of ed's comments --- requirements-testing.txt | 1 - tests/test_petsc.py | 42 +++++++++++++++++++++++++--------------- 2 files changed, 26 insertions(+), 17 deletions(-) diff --git a/requirements-testing.txt b/requirements-testing.txt index 2acf0e4f11..d10923cb00 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -1,7 +1,6 @@ pytest>=7.2,<8.5 pytest-runner<6.0.2 pytest-cov<6.1.2 -pytest-order flake8-pyproject>=1.2.3,<1.2.4 nbval<0.11.1 scipy<1.15.3 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index a2f9d8b8ec..371901be76 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -20,7 +20,7 @@ @skipif('petsc') -@pytest.mark.order(0) +@pytest.fixture(scope='session', autouse=True) def test_petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized @@ -1341,27 +1341,37 @@ def test_laplacian_1d(self, nx, unorm, mode): os.environ['CC'] = 'mpicc' PetscInitialize() - class SubLeft(SubDomain): - name = 'subleft' + # class SubLeft(SubDomain): + # name = 'subleft' - def define(self, dimensions): - x, = dimensions - return {x: ('left', 1)} + # def define(self, dimensions): + # x, = dimensions + # return {x: ('left', 1)} + + # class SubRight(SubDomain): + # name = 'subright' - class SubRight(SubDomain): - name = 'subright' + # def define(self, dimensions): + # x, = dimensions + # return {x: ('right', 1)} + + # grid = Grid(shape=(nx,), dtype=np.float64) + + # sub1 = SubLeft(grid=grid) + # sub2 = SubRight(grid=grid) + + class SubSide(SubDomain): + def __init__(self, side='left', grid=None): + self.side = side + self.name = f'sub{side}' + super().__init__(grid=grid) def define(self, dimensions): x, = dimensions - return {x: ('right', 1)} - - sub1 = SubLeft() - sub2 = SubRight() - subdomains = (sub1, sub2,) + return {x: (self.side, 1)} - grid = Grid( - shape=(nx,), subdomains=subdomains, dtype=np.float64 - ) + grid = Grid(shape=(nx,), dtype=np.float64) + sub1, sub2 = [SubSide(side=s, grid=grid) for s in ('left', 'right')] u = Function(name='u', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) From 72e8222769caaf84382f6aecb92f41e2c1c75d91 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 17:07:23 +0100 Subject: [PATCH 47/53] misc: Address comments --- tests/test_petsc.py | 32 ++++++-------------------------- 1 file changed, 6 insertions(+), 26 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 371901be76..3f64a79c82 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1333,7 +1333,7 @@ class TestMPI: (257, 28.67050632840985) ]) @skipif('petsc') - @pytest.mark.parallel(mode=[1, 2, 4, 8]) + @pytest.mark.parallel(mode=[2, 4, 8]) def test_laplacian_1d(self, nx, unorm, mode): """ """ @@ -1341,25 +1341,6 @@ def test_laplacian_1d(self, nx, unorm, mode): os.environ['CC'] = 'mpicc' PetscInitialize() - # class SubLeft(SubDomain): - # name = 'subleft' - - # def define(self, dimensions): - # x, = dimensions - # return {x: ('left', 1)} - - # class SubRight(SubDomain): - # name = 'subright' - - # def define(self, dimensions): - # x, = dimensions - # return {x: ('right', 1)} - - # grid = Grid(shape=(nx,), dtype=np.float64) - - # sub1 = SubLeft(grid=grid) - # sub2 = SubRight(grid=grid) - class SubSide(SubDomain): def __init__(self, side='left', grid=None): self.side = side @@ -1375,19 +1356,18 @@ def define(self, dimensions): u = Function(name='u', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) - bc = Function(name='bc', grid=grid, space_order=2) + + u0 = Constant(name='u0', value=-1.0, dtype=np.float64) + u1 = Constant(name='u1', value=-np.exp(1.0), dtype=np.float64) eqn = Eq(-u.laplace, f, subdomain=grid.interior) X = np.linspace(0, 1.0, nx).astype(np.float64) f.data[:] = np.float64(np.exp(X)) - bc.data[0] = -np.float64(1.0) # u(0) = -1 - bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e - # Create boundary condition expressions using subdomains - bcs = [EssentialBC(u, bc, subdomain=sub1)] - bcs += [EssentialBC(u, bc, subdomain=sub2)] + bcs = [EssentialBC(u, u0, subdomain=sub1)] + bcs += [EssentialBC(u, u1, subdomain=sub2)] petsc = PETScSolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) From 16be19a9a4b0372bf6078007013a36bdb7137fc2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 20 Jun 2025 12:28:51 +0100 Subject: [PATCH 48/53] dsl/compiler: Add PETSc logging instrastructure and tests --- devito/operator/operator.py | 6 +- devito/operator/profiling.py | 36 +++++- devito/petsc/iet/logging.py | 77 ++++++++++++ devito/petsc/iet/passes.py | 51 +++++--- devito/petsc/iet/routines.py | 42 ++++--- devito/petsc/logging.py | 180 +++++++++++++++++++++++++++++ devito/petsc/solve.py | 16 ++- devito/petsc/types/object.py | 2 + devito/petsc/types/types.py | 10 +- devito/symbolics/extended_sympy.py | 6 +- tests/test_petsc.py | 108 ++++++++++++++++- 11 files changed, 483 insertions(+), 51 deletions(-) create mode 100644 devito/petsc/iet/logging.py create mode 100644 devito/petsc/logging.py diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 3a36952a58..61464120c5 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -194,7 +194,7 @@ def _sanitize_exprs(cls, expressions, **kwargs): @classmethod def _build(cls, expressions, **kwargs): # Python- (i.e., compile-) and C-level (i.e., run-time) performance - profiler = create_profile('timers') + profiler = create_profile('timers', kwargs['language']) # Lower the input expressions into an IET irs, byproduct = cls._lower(expressions, profiler=profiler, **kwargs) @@ -1004,7 +1004,9 @@ def _emit_apply_profiling(self, args): elapsed = fround(self._profiler.py_timers['apply']) info(f"Operator `{self.name}` ran in {elapsed:.2f} s") - summary = self._profiler.summary(args, self._dtype, reduce_over=elapsed) + summary = self._profiler.summary( + args, self._dtype, self.parameters, reduce_over=elapsed + ) if not is_log_enabled_for('PERF'): # Do not waste time diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index ae2f50e171..360d506d17 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -17,6 +17,7 @@ from devito.parameters import configuration from devito.symbolics import subs_op_args from devito.tools import DefaultOrderedDict, flatten +from devito.petsc.logging import PetscSummary __all__ = ['create_profile'] @@ -42,7 +43,7 @@ class Profiler: _attempted_init = False - def __init__(self, name): + def __init__(self, name, language): self.name = name # Operation reductions observed in sections @@ -55,6 +56,9 @@ def __init__(self, name): # Python-level timers self.py_timers = OrderedDict() + # For language specific summaries + self.language = language + self._attempted_init = True def analyze(self, iet): @@ -179,7 +183,7 @@ def record_ops_variation(self, initial, final): def all_sections(self): return list(self._sections) + flatten(self._subsections.values()) - def summary(self, args, dtype, reduce_over=None): + def summary(self, args, dtype, params, reduce_over=None): """ Return a PerformanceSummary of the profiled sections. @@ -277,7 +281,7 @@ def _allgather_from_comm(self, comm, time, ops, points, traffic, sops, itershape return list(zip(times, opss, pointss, traffics, sops, itershapess)) # Override basic summary so that arguments other than runtime are computed. - def summary(self, args, dtype, reduce_over=None): + def summary(self, args, dtype, params, reduce_over=None): grid = args.grid comm = args.comm @@ -338,6 +342,11 @@ def summary(self, args, dtype, reduce_over=None): # data transfers) summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup) + # Add the language specific summary + summary.add_language_summary( + self.language, + language_summary_mapper[self.language](params) + ) return summary @@ -478,6 +487,16 @@ def add_glb_fdlike(self, key, points, time): self.globals[key] = PerfEntry(time, None, gpointss, None, None, None) + def add_language_summary(self, lang, summary): + """ + Register a language specific summary (e.g., PetscSummary) + and dynamically add a property to access it via perf_summary.. + """ + # TODO: Consider renaming `PerformanceSummary` to something more generic + # (e.g., `Summary`), or separating `PetscSummary` entirely from + # `PerformanceSummary`. + setattr(self, lang, summary) + @property def globals_all(self): v0 = self.globals['vanilla'] @@ -503,7 +522,7 @@ def timings(self): return OrderedDict([(k, v.time) for k, v in self.items()]) -def create_profile(name): +def create_profile(name, language): """Create a new Profiler.""" if configuration['log-level'] in ['DEBUG', 'PERF'] and \ configuration['profiling'] == 'basic': @@ -511,13 +530,13 @@ def create_profile(name): level = 'advanced' else: level = configuration['profiling'] - profiler = profiler_registry[level](name) + profiler = profiler_registry[level](name, language) if profiler._attempted_init: return profiler else: warning(f"Couldn't set up `{level}` profiler; reverting to 'advanced'") - profiler = profiler_registry['advanced'](name) + profiler = profiler_registry['advanced'](name, language) # We expect the `advanced` profiler to always initialize successfully assert profiler._attempted_init return profiler @@ -533,3 +552,8 @@ def create_profile(name): 'advisor': AdvisorProfiler } """Profiling levels.""" + + +language_summary_mapper = { + 'petsc': PetscSummary +} diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py new file mode 100644 index 0000000000..6a1c74d767 --- /dev/null +++ b/devito/petsc/iet/logging.py @@ -0,0 +1,77 @@ +from functools import cached_property + +from devito.symbolics import Byref, FieldFromPointer +from devito.ir.iet import DummyExpr +from devito.logger import PERF + +from devito.petsc.iet.utils import petsc_call +from devito.petsc.logging import petsc_return_variable_dict, PetscInfo + + +class PetscLogger: + """ + Class for PETSc loggers that collect solver related statistics. + """ + def __init__(self, level, **kwargs): + self.sobjs = kwargs.get('solver_objs') + self.sreg = kwargs.get('sregistry') + self.section_mapper = kwargs.get('section_mapper', {}) + self.injectsolve = kwargs.get('injectsolve', None) + + self.function_list = [] + + if level <= PERF: + self.function_list.extend([ + 'kspgetiterationnumber', + 'snesgetiterationnumber' + ]) + + # TODO: To be extended with if level <= DEBUG: ... + + name = self.sreg.make_name(prefix='petscinfo') + pname = self.sreg.make_name(prefix='petscprofiler') + + self.statstruct = PetscInfo( + name, pname, self.logobjs, self.sobjs, + self.section_mapper, self.injectsolve, + self.function_list + ) + + @cached_property + def logobjs(self): + """ + Create PETSc objects specifically needed for logging solver statistics. + """ + return { + info.name: info.variable_type( + self.sreg.make_name(prefix=info.output_param) + ) + for func_name in self.function_list + for info in [petsc_return_variable_dict[func_name]] + } + + @cached_property + def calls(self): + """ + Generate the PETSc calls that will be injected into the C code to + extract solver statistics. + """ + struct = self.statstruct + calls = [] + for param in self.function_list: + param = petsc_return_variable_dict[param] + + inputs = [] + for i in param.input_params: + inputs.append(self.sobjs[i]) + + logobj = self.logobjs[param.name] + + calls.append( + petsc_call(param.name, inputs + [Byref(logobj)]) + ) + # TODO: Perform a PetscCIntCast here? + expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol) + calls.append(expr) + + return tuple(calls) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index b4859dda16..8d7042bb6c 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -5,7 +5,7 @@ from devito.passes.iet.engine import iet_pass from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, DummyExpr, CallableBody, List, Call, Callable, - FindNodes) + FindNodes, Section) from devito.symbolics import Byref, Macro, FieldFromPointer from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol @@ -22,8 +22,11 @@ CoupledObjectBuilder, BaseSetup, CoupledSetup, Solver, CoupledSolver, TimeDependent, NonTimeDependent) +from devito.petsc.iet.logging import PetscLogger from devito.petsc.iet.utils import petsc_call, petsc_call_mpi +import devito.logger as dl + @iet_pass def lower_petsc(iet, **kwargs): @@ -63,14 +66,17 @@ def lower_petsc(iet, **kwargs): subs = {} efuncs = {} + # Map PETScSolve to its Section (for logging) + section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) + for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, comm, **kwargs) + builder = Builder(injectsolve, objs, iters, comm, section_mapper, **kwargs) setup.extend(builder.solversetup.calls) # Transform the spatial iteration loop with the calls to execute the solver - subs.update({builder.solve.spatial_body: builder.solve.calls}) + subs.update({builder.solve.spatial_body: builder.calls}) efuncs.update(builder.cbbuilder.efuncs) @@ -78,7 +84,7 @@ def lower_petsc(iet, **kwargs): iet = Transformer(subs).visit(iet) - body = core + tuple(setup) + (BlankLine,) + iet.body.body + body = core + tuple(setup) + iet.body.body body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())} @@ -125,49 +131,64 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, comm, **kwargs): + def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters self.comm = comm + self.section_mapper = section_mapper self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) - self.args = { + self.common_kwargs = { 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, 'comm': self.comm, + 'section_mapper': self.section_mapper, **self.kwargs } - self.args['solver_objs'] = self.objbuilder.solver_objs - self.args['timedep'] = self.timedep - self.args['cbbuilder'] = self.cbbuilder + self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs + self.common_kwargs['timedep'] = self.timedep + self.common_kwargs['cbbuilder'] = self.cbbuilder + self.common_kwargs['logger'] = self.logger @cached_property def objbuilder(self): return ( - CoupledObjectBuilder(**self.args) + CoupledObjectBuilder(**self.common_kwargs) if self.coupled else - BaseObjectBuilder(**self.args) + BaseObjectBuilder(**self.common_kwargs) ) @cached_property def timedep(self): time_mapper = self.injectsolve.expr.rhs.time_mapper timedep_class = TimeDependent if time_mapper else NonTimeDependent - return timedep_class(**self.args) + return timedep_class(**self.common_kwargs) @cached_property def cbbuilder(self): - return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args) + return CCBBuilder(**self.common_kwargs) \ + if self.coupled else CBBuilder(**self.common_kwargs) @cached_property def solversetup(self): - return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args) + return CoupledSetup(**self.common_kwargs) \ + if self.coupled else BaseSetup(**self.common_kwargs) @cached_property def solve(self): - return CoupledSolver(**self.args) if self.coupled else Solver(**self.args) + return CoupledSolver(**self.common_kwargs) \ + if self.coupled else Solver(**self.common_kwargs) + + @cached_property + def logger(self): + log_level = dl.logger.level + return PetscLogger(log_level, **self.common_kwargs) + + @cached_property + def calls(self): + return List(body=self.solve.calls+self.logger.calls) def populate_matrix_context(efuncs, objs): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9249c7a557..d21af4be6b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -7,22 +7,21 @@ retrieve_iteration_tree, filter_iterations, Iteration, PointerCast) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod) + FieldFromComposite, IntDiv, Deref, Mod, String) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension from devito.tools import filter_ordered -from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, dereference_funcs, residual_bundle) from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, - PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, - DMCast, JacobianStruct, - SubMatrixStruct, CallbackDM) +from devito.petsc.types import (PETScArray, PetscBundle, DM, Mat, CallbackVec, Vec, + KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, + VecScatter, DMCast, JacobianStruct, SubMatrixStruct, + CallbackDM) class CBBuilder: @@ -1068,6 +1067,10 @@ def _build(self): """ sreg = self.sregistry targets = self.fielddata.targets + + snes_name = sreg.make_name(prefix='snes') + options_prefix = self.injectsolve.expr.rhs.options_prefix + base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), 'xglobal': Vec(sreg.make_name(prefix='xglobal')), @@ -1076,10 +1079,12 @@ def _build(self): 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), 'pc': PC(sreg.make_name(prefix='pc')), - 'snes': SNES(sreg.make_name(prefix='snes')), + 'snes': SNES(snes_name), 'localsize': PetscInt(sreg.make_name(prefix='localsize')), 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), + 'snesprefix': String((options_prefix or '') + '_'), + 'options_prefix': options_prefix, } base_dict['comm'] = self.comm self._target_dependent(base_dict) @@ -1239,6 +1244,10 @@ def _setup(self): snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] + ) if sobjs['options_prefix'] else None + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) @@ -1327,6 +1336,7 @@ def _setup(self): base_setup = dmda_calls + ( snes_create, + snes_options_prefix, snes_set_dm, create_matrix, snes_set_jac, @@ -1412,11 +1422,14 @@ def _setup(self): sobjs = self.solver_objs dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] + ) if sobjs['options_prefix'] else None + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) @@ -1530,6 +1543,7 @@ def _setup(self): coupled_setup = dmda_calls + ( snes_create, + snes_options_prefix, snes_set_dm, create_matrix, snes_set_jac, @@ -1618,7 +1632,7 @@ def _execute_solve(self): vec_reset_array, BlankLine, ) - return List(body=run_solver_calls) + return run_solver_calls @cached_property def spatial_body(self): @@ -1701,15 +1715,7 @@ def _execute_solve(self): snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),) - return List( - body=( - (struct_assignment,) - + pre_solve - + snes_solve - + post_solve - + (BlankLine,) - ) - ) + return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,) class NonTimeDependent: diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py new file mode 100644 index 0000000000..f552b69969 --- /dev/null +++ b/devito/petsc/logging.py @@ -0,0 +1,180 @@ +from collections import namedtuple, OrderedDict +from dataclasses import dataclass + +from devito.types import CompositeObject + +from devito.petsc.types import PetscInt +from devito.petsc.utils import petsc_type_mappings + + +class PetscEntry: + def __init__(self, **kwargs): + self.kwargs = kwargs + for k, v in self.kwargs.items(): + setattr(self, k, v) + self._properties = {k.lower(): v for k, v in kwargs.items()} + + def __getitem__(self, key): + return self._properties[key.lower()] + + def __repr__(self): + return f"PetscEntry({', '.join(f'{k}={v}' for k, v in self.kwargs.items())})" + + +class PetscSummary(dict): + """ + A summary of PETSc statistics collected for all solver runs + associated with a single operator during execution. + """ + PetscKey = namedtuple('PetscKey', 'name options_prefix') + + def __init__(self, params, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.petscinfos = [i for i in params if isinstance(i, PetscInfo)] + + # Gather all unique PETSc function names across all PetscInfo objects + self._functions = list(dict.fromkeys( + petsc_return_variable_dict[key].name + for struct in self.petscinfos + for key in struct.function_list + )) + self._property_name_map = {} + # Dynamically create a property on this class for each PETSc function + self._add_properties() + + # Initialize the summary by adding PETSc information from each PetscInfo + # object (each corresponding to an individual PETScSolve) + for i in self.petscinfos: + self.add_info(i) + + def add_info(self, petscinfo): + """ + For a given PetscInfo object, create a key + and entry and add it to the PetscSummary. + """ + key = self.PetscKey(*petscinfo.summary_key) + entry = self.petsc_entry(petscinfo) + self[key] = entry + + def petsc_entry(self, petscinfo): + """ + Create a named tuple entry for the given PetscInfo object, + containing the values for each PETSc function call. + """ + funcs = self._functions + values = tuple(getattr(petscinfo, c) for c in funcs) + return PetscEntry(**{k: v for k, v in zip(funcs, values)}) + + def _add_properties(self): + """ + For each function name in `self._functions` (e.g., 'KSPGetIterationNumber'), + dynamically add a property to the class with the same name. + + Each property returns an OrderedDict that maps each PetscKey to the + result of looking up that function on the corresponding PetscEntry, + if the function exists on that entry. + """ + def make_property(function): + def getter(self): + return OrderedDict( + (k, getattr(v, function)) + for k, v in self.items() + # Only include entries that have the function + if hasattr(v, function) + ) + return property(getter) + + for f in self._functions: + # Inject the new property onto the class itself + setattr(self.__class__, f, make_property(f)) + self._property_name_map[f.lower()] = f + + def get_entry(self, name, options_prefix=None): + """ + Retrieve a single PetscEntry for a given name + and options_prefix. + """ + key = self.PetscKey(name, options_prefix) + if key not in self: + raise ValueError( + f"No PETSc information for:" + f" name='{name}'" + f" options_prefix='{options_prefix}'" + ) + return self[key] + + def __getitem__(self, key): + if isinstance(key, str): + # Allow case insensitive key access + original = self._property_name_map.get(key.lower()) + if original: + return getattr(self, original) + raise KeyError(f"No PETSc function named '{key}'") + elif isinstance(key, tuple) and len(key) == 2: + # Allow tuple keys (name, options_prefix) + key = self.PetscKey(*key) + return super().__getitem__(key) + + +class PetscInfo(CompositeObject): + + __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper', + 'injectsolve', 'function_list') + + def __init__(self, name, pname, logobjs, sobjs, section_mapper, + injectsolve, function_list): + + self.logobjs = logobjs + self.sobjs = sobjs + self.section_mapper = section_mapper + self.injectsolve = injectsolve + self.function_list = function_list + + mapper = {v: k for k, v in petsc_type_mappings.items()} + fields = [(str(i), mapper[str(i._C_ctype)]) for i in logobjs.values()] + super().__init__(name, pname, fields) + + @property + def section(self): + section = self.section_mapper.items() + return next((k[0].name for k, v in section if self.injectsolve in v), None) + + @property + def summary_key(self): + return (self.section, self.sobjs['options_prefix']) + + def __getattr__(self, attr): + if attr in self.logobjs.keys(): + return getattr(self.value._obj, self.logobjs[attr].name) + raise AttributeError(f"{attr} not found in PETSc return variables") + + +@dataclass +class PetscReturnVariable: + name: str + variable_type: None + input_params: list + output_param: str + + +# NOTE: +# In the future, this dictionary should be generated automatically from PETSc internals. +# For now, it serves as the reference for PETSc function metadata. +# If any of the PETSc function signatures change (e.g., names, input/output parameters), +# this dictionary must be updated accordingly. + +petsc_return_variable_dict = { + 'kspgetiterationnumber': PetscReturnVariable( + name='KSPGetIterationNumber', + variable_type=PetscInt, + input_params=['ksp'], + output_param='kspiter' + ), + 'snesgetiterationnumber': PetscReturnVariable( + name='SNESGetIterationNumber', + variable_type=PetscInt, + input_params=['snes'], + output_param='snesiter', + ) +} diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1f4423707e..55633d40fd 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -12,7 +12,7 @@ __all__ = ['PETScSolve'] -def PETScSolve(target_exprs, target=None, solver_parameters=None): +def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix=None): """ Returns a symbolic expression representing a linear PETSc solver, enriched with all the necessary metadata for execution within an `Operator`. @@ -51,16 +51,21 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None): This can be passed directly to a Devito Operator. """ if target is not None: - return InjectSolve(solver_parameters, {target: target_exprs}).build_expr() + return InjectSolve( + solver_parameters, {target: target_exprs}, options_prefix + ).build_expr() else: - return InjectMixedSolve(solver_parameters, target_exprs).build_expr() + return InjectMixedSolve( + solver_parameters, target_exprs, options_prefix + ).build_expr() class InjectSolve: - def __init__(self, solver_parameters=None, target_exprs=None): + def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None): self.solver_params = solver_parameters self.time_mapper = None self.target_exprs = target_exprs + self.options_prefix = options_prefix def build_expr(self): target, funcs, fielddata = self.linear_solve_args() @@ -71,7 +76,8 @@ def build_expr(self): self.solver_params, fielddata=fielddata, time_mapper=self.time_mapper, - localinfo=localinfo + localinfo=localinfo, + options_prefix=self.options_prefix ) return PetscEq(target, linear_solve) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 32892f2aa4..1484e1be81 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,9 +1,11 @@ from ctypes import POINTER, c_char + from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject, CustomDimension) from devito.symbolics import Byref, cast from devito.types.basic import DataSymbol + from devito.petsc.iet.utils import petsc_call diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 4f235bbbde..baab036f17 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -60,7 +60,7 @@ class LinearSolveExpr(MetaData): __rargs__ = ('expr',) __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', - 'localinfo') + 'localinfo', 'options_prefix') defaults = { 'ksp_type': 'gmres', @@ -72,7 +72,8 @@ class LinearSolveExpr(MetaData): } def __new__(cls, expr, solver_parameters=None, - fielddata=None, time_mapper=None, localinfo=None, **kwargs): + fielddata=None, time_mapper=None, localinfo=None, + options_prefix=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -88,6 +89,7 @@ def __new__(cls, expr, solver_parameters=None, obj._fielddata = fielddata if fielddata else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo + obj._options_prefix = options_prefix return obj def __repr__(self): @@ -125,6 +127,10 @@ def time_mapper(self): def localinfo(self): return self._localinfo + @property + def options_prefix(self): + return self._options_prefix + @property def grid(self): return self.fielddata.grid diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index a75e1091d3..4016a88466 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -581,7 +581,11 @@ class Keyword(ReservedWord): class String(ReservedWord): - pass + + def __str__(self): + return f'"{self.value}"' + + __repr__ = __str__ class Macro(ReservedWord): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 3f64a79c82..44668092ad 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -2,10 +2,12 @@ import numpy as np import os +from collections import OrderedDict from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, configuration, norm, switchconfig, SubDomain) +from devito.operator.profiling import PerformanceSummary from devito.ir.iet import (Call, ElementalFunction, FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject @@ -17,11 +19,11 @@ from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize +from devito.petsc.logging import PetscSummary -@skipif('petsc') @pytest.fixture(scope='session', autouse=True) -def test_petsc_initialization(): +def petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized configuration['compiler'] = 'custom' @@ -1377,3 +1379,105 @@ def define(self, dimensions): # Expected norm computed "manually" from sequential run # What rtol and atol should be used? assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) + + +class TestLogging: + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_logging(self, log_level): + """Verify PetscSummary output when the log level is 'PERF' or 'DEBUG.""" + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f']] + e, f = functions + f.data[:] = 5.0 + eq = Eq(e.laplace, f) + + petsc = PETScSolve(eq, target=e, options_prefix='poisson') + + with switchconfig(language='petsc', log_level=log_level): + op = Operator(petsc) + summary = op.apply() + + # One PerformanceSummary + assert len(summary) == 1 + + # Access the PetscSummary + petsc_summary = summary.petsc + + assert isinstance(summary, PerformanceSummary) + assert isinstance(petsc_summary, PetscSummary) + + # One section with a single solver + assert len(petsc_summary) == 1 + + entry0 = petsc_summary.get_entry('section0', 'poisson') + entry1 = petsc_summary[('section0', 'poisson')] + assert entry0 == entry1 + assert str(entry0) == \ + 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' + + assert entry0.SNESGetIterationNumber == 1 + + snesiters0 = petsc_summary.SNESGetIterationNumber + snesiters1 = petsc_summary['SNESGetIterationNumber'] + # Check case insensitive key access + snesiters2 = petsc_summary['snesgetiterationnumber'] + snesiters3 = petsc_summary['SNESgetiterationNumber'] + + assert snesiters0 == snesiters1 == snesiters2 == snesiters3 + + assert isinstance(snesiters0, OrderedDict) + assert len(snesiters0) == 1 + key, value = next(iter(snesiters0.items())) + assert str(key) == "PetscKey(name='section0', options_prefix='poisson')" + assert value == 1 + + @skipif('petsc') + def test_logging_multiple_solves(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + e.data[:] = 5.0 + f.data[:] = 6.0 + + eq1 = Eq(g.laplace, e) + eq2 = Eq(h, f + 5.0) + + solver1 = PETScSolve(eq1, target=g, options_prefix='poisson1') + solver2 = PETScSolve(eq2, target=h, options_prefix='poisson2') + + with switchconfig(language='petsc', log_level='DEBUG'): + op = Operator([solver1, solver2]) + summary = op.apply() + + petsc_summary = summary.petsc + # One PetscKey, PetscEntry for each solver + assert len(petsc_summary) == 2 + + entry1 = petsc_summary.get_entry('section0', 'poisson1') + entry2 = petsc_summary.get_entry('section1', 'poisson2') + + assert str(entry1) == \ + 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' + assert str(entry2) == \ + 'PetscEntry(KSPGetIterationNumber=1, SNESGetIterationNumber=1)' + + assert len(petsc_summary.KSPGetIterationNumber) == 2 + assert len(petsc_summary.SNESGetIterationNumber) == 2 + + assert entry1.KSPGetIterationNumber == 16 + assert entry1.SNESGetIterationNumber == 1 + assert entry2.KSPGetIterationNumber == 1 + assert entry2.SNESGetIterationNumber == 1 + + # Test key access to PetscEntry + assert entry1['KSPGetIterationNumber'] == 16 + assert entry1['SNESGetIterationNumber'] == 1 + # Case insensitive key access + assert entry1['kspgetiterationnumber'] == 16 From e3e26f0155be3b7d126ddfd92450775395ba0c46 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Jul 2025 17:27:18 +0100 Subject: [PATCH 49/53] compiler: Fix language summary --- devito/operator/profiling.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 360d506d17..5e30500761 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -342,11 +342,11 @@ def summary(self, args, dtype, params, reduce_over=None): # data transfers) summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup) - # Add the language specific summary - summary.add_language_summary( - self.language, - language_summary_mapper[self.language](params) - ) + # Add the language specific summary if necessary + mapper_func = language_summary_mapper.get(self.language) + if mapper_func: + summary.add_language_summary(self.language, mapper_func(params)) + return summary From 9a182daeb3d914ed92e80fa29b6d265118089a56 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Jul 2025 18:25:22 +0100 Subject: [PATCH 50/53] misc: Fix advisor profiling with new language arg --- devito/operator/profiling.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 5e30500761..04f6779904 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -375,11 +375,11 @@ class AdvisorProfiler(AdvancedProfiler): _default_libs = ['ittnotify'] _ext_calls = [_api_resume, _api_pause] - def __init__(self, name): + def __init__(self, name, language): if self._attempted_init: return - super().__init__(name) + super().__init__(name, language) path = get_advisor_path() if path: From 774af10bb6a63861b480369295e56018ce001d10 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 16 Jul 2025 23:05:04 +0100 Subject: [PATCH 51/53] misc: Clean up and docstrings --- devito/petsc/iet/routines.py | 63 ++++++++++++++-------------------- devito/petsc/types/array.py | 2 +- devito/petsc/types/types.py | 35 +++++++++++-------- devito/symbolics/extraction.py | 7 ++-- devito/types/object.py | 2 +- 5 files changed, 54 insertions(+), 55 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index d21af4be6b..503d773372 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -42,14 +42,12 @@ def __init__(self, **kwargs): self._struct_params = [] self._main_matvec_callback = None - self._main_formfunc_callback = None self._user_struct_callback = None - # TODO: Test pickling. The mutability of these lists - # could cause issues when pickling? + self._F_efunc = None + self._b_efunc = None + self._J_efuncs = [] - self._F_efuncs = [] - self._b_efuncs = [] - self._initialguesses = [] + self._initial_guesses = [] self._make_core() self._efuncs = self._uxreplace_efuncs() @@ -69,31 +67,26 @@ def filtered_struct_params(self): @property def main_matvec_callback(self): """ - This is the matvec callback associated with the whole Jacobian i.e - is set in the main kernel via - `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` + The matrix-vector callback for the full Jacobian. + This is the function set in the main Kernel via: + PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))...)); + The callback has the signature `(Mat, Vec, Vec)`. """ return self._J_efuncs[0] - @property - def main_formfunc_callback(self): - return self._F_efuncs[0] - @property def J_efuncs(self): + """ + List of matrix-vector callbacks. + Each callback has the signature `(Mat, Vec, Vec)`. Typically, this list + contains a single element, but in mixed systems it can include multiple + callbacks, one for each subblock. + """ return self._J_efuncs @property - def F_efuncs(self): - return self._F_efuncs - - @property - def b_efuncs(self): - return self._b_efuncs - - @property - def initialguesses(self): - return self._initialguesses + def initial_guesses(self): + return self._initial_guesses @property def user_struct_callback(self): @@ -284,7 +277,7 @@ def _make_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._F_efuncs.append(cb) + self._F_efunc = cb self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): @@ -422,7 +415,7 @@ def _make_formrhs(self): retval=objs['err'], parameters=(sobjs['callbackdm'], objs['B']) ) - self._b_efuncs.append(cb) + self._b_efunc = cb self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): @@ -534,7 +527,7 @@ def _make_initialguess(self): retval=objs['err'], parameters=(sobjs['callbackdm'], objs['xloc']) ) - self._initialguesses.append(cb) + self._initial_guesses.append(cb) self._efuncs[cb.name] = cb def _create_initial_guess_body(self, body): @@ -660,16 +653,12 @@ def jacobian(self): @property def main_matvec_callback(self): """ - This is the matvec callback associated with the whole Jacobian i.e + This is the matrix-vector callback associated with the whole Jacobian i.e is set in the main kernel via `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` """ return self._main_matvec_callback - @property - def main_formfunc_callback(self): - return self._main_formfunc_callback - def _make_core(self): for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') @@ -757,7 +746,7 @@ def _make_whole_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._main_formfunc_callback = cb + self._F_efunc = cb self._efuncs[cb.name] = cb def _whole_formfunc_body(self, body): @@ -1310,7 +1299,7 @@ def _setup(self): 'MatShellSetOperation', [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - formfunc = self.cbbuilder.main_formfunc_callback + formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), @@ -1477,7 +1466,7 @@ def _setup(self): 'MatShellSetOperation', [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - formfunc = self.cbbuilder.main_formfunc_callback + formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), @@ -1593,7 +1582,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - b_efunc = self.cbbuilder.b_efuncs[0] + b_efunc = self.cbbuilder._b_efunc dmda = sobjs['dmda'] @@ -1601,8 +1590,8 @@ def _execute_solve(self): vec_place_array = self.timedep.place_array(target) - if self.cbbuilder.initialguesses: - initguess = self.cbbuilder.initialguesses[0] + if self.cbbuilder.initial_guesses: + initguess = self.cbbuilder.initial_guesses[0] initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']]) else: initguess_call = None diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index e691a96c33..1bed71ec50 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -169,7 +169,7 @@ def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) - @cached_property + @property def symbolic_shape(self): return self.c0.symbolic_shape diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index baab036f17..4f1f492c8e 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,6 +1,7 @@ import sympy from itertools import chain +from functools import cached_property from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict from devito.tools.dtypes_lowering import dtype_mapper @@ -241,7 +242,7 @@ def __init__(self, targets, arrays, jacobian=None, residual=None): self._jacobian = jacobian self._residual = residual - @property + @cached_property def space_dimensions(self): space_dims = {t.space_dimensions for t in self.targets} if len(space_dims) > 1: @@ -251,7 +252,7 @@ def space_dimensions(self): ) return space_dims.pop() - @property + @cached_property def grid(self): """The unique `Grid` associated with all targets.""" grids = [t.grid for t in self.targets] @@ -261,7 +262,7 @@ def grid(self): ) return grids.pop() - @property + @cached_property def space_order(self): # NOTE: since we use DMDA to create vecs for the coupled solves, # all fields must have the same space order @@ -398,7 +399,8 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_expr(eq) if e is not None ) - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + key = lambda e: not isinstance(e, EssentialBC) + matvecs = tuple(sorted(matvecs, key=key)) matvecs = self._scale_non_bcs(matvecs) scdiag = self._compute_scdiag(matvecs) @@ -425,7 +427,7 @@ class MixedJacobian(BaseJacobian): """ def __init__(self, target_exprs, arrays, time_mapper): super().__init__(arrays=arrays, target=None) - self.targets = tuple(target_exprs.keys()) + self.targets = tuple(target_exprs) self.time_mapper = time_mapper self._submatrices = [] self._build_blocks(target_exprs) @@ -443,12 +445,12 @@ def n_submatrices(self): """Return the number of submatrix blocks.""" return len(self._submatrices) - @property + @cached_property def nonzero_submatrices(self): """Return SubMatrixBlock objects that have non-empty matvecs.""" return [m for m in self.submatrices if m.matvecs] - @property + @cached_property def target_scaler_mapper(self): """ Map each row target to the scdiag of its corresponding @@ -467,13 +469,8 @@ def _build_blocks(self, target_exprs): for i, row_target in enumerate(self.targets): exprs = target_exprs[row_target] for j, col_target in enumerate(self.targets): - matvecs = [] - for expr in exprs: - matvecs.extend( - e for e in self._build_matvec_expr( - expr, col_target=col_target, row_target=row_target - ) - ) + + matvecs = self._build_submatrix_matvecs(exprs, row_target, col_target) matvecs = [m for m in matvecs if m is not None] # Sort to put EssentialBC first if any @@ -497,6 +494,16 @@ def _build_blocks(self, target_exprs): ) self._submatrices.append(block) + def _build_submatrix_matvecs(self, exprs, row_target, col_target): + matvecs = [] + for expr in exprs: + matvecs.extend( + e for e in self._build_matvec_expr( + expr, col_target=col_target, row_target=row_target + ) + ) + return matvecs + def get_submatrix(self, row_idx, col_idx): """ Return the SubMatrixBlock at (row_idx, col_idx), or None if not found. diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index c0a28433cc..d429f83831 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -79,10 +79,13 @@ def centre_stencil(expr, target, as_coeff=False): Extract the centre stencil from an expression. Its coefficient is what would appear on the diagonal of the matrix system if the matrix were formed explicitly. + Parameters ---------- - expr : the expression to extract the centre stencil from - target : the target function whose centre stencil we want + expr : expr-like + The expression to extract the centre stencil from + target : Function + The target function whose centre stencil we want as_coeff : bool, optional If True, return the coefficient of the centre stencil """ diff --git a/devito/types/object.py b/devito/types/object.py index 1153ce81eb..201a087b3d 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -252,10 +252,10 @@ class LocalCompositeObject(CompositeObject, LocalType): __rkwargs__ = ('modifier', 'liveness') def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): - self.modifier = modifier dtype = CustomDtype(f"struct {pname}", modifier=modifier) Object.__init__(self, name, dtype, None) self._pname = pname + self.modifier = modifier assert liveness in ['eager', 'lazy'] self._liveness = liveness self._fields = fields From 910c98731e441a3eebe2e6738354246924516921 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Jul 2025 10:47:56 +0100 Subject: [PATCH 52/53] misc: Clean up --- devito/petsc/clusters.py | 2 +- devito/petsc/iet/logging.py | 4 +- devito/petsc/iet/passes.py | 30 ++++---- devito/petsc/iet/routines.py | 121 ++++++++++++++++----------------- devito/petsc/logging.py | 8 +-- devito/petsc/solve.py | 8 +-- devito/petsc/types/types.py | 24 +++---- devito/symbolics/extraction.py | 4 +- tests/test_petsc.py | 30 ++++---- 9 files changed, 114 insertions(+), 117 deletions(-) diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index e035ccbefc..7171669490 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -20,7 +20,7 @@ def petsc_lift(clusters): processed = [] for c in clusters: if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions) + ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py index 6a1c74d767..af8b5c4851 100644 --- a/devito/petsc/iet/logging.py +++ b/devito/petsc/iet/logging.py @@ -16,7 +16,7 @@ def __init__(self, level, **kwargs): self.sobjs = kwargs.get('solver_objs') self.sreg = kwargs.get('sregistry') self.section_mapper = kwargs.get('section_mapper', {}) - self.injectsolve = kwargs.get('injectsolve', None) + self.inject_solve = kwargs.get('inject_solve', None) self.function_list = [] @@ -33,7 +33,7 @@ def __init__(self, level, **kwargs): self.statstruct = PetscInfo( name, pname, self.logobjs, self.sobjs, - self.section_mapper, self.injectsolve, + self.section_mapper, self.inject_solve, self.function_list ) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 8d7042bb6c..aadb711c76 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -31,10 +31,10 @@ @iet_pass def lower_petsc(iet, **kwargs): # Check if PETScSolve was used - injectsolve_mapper = MapNodes(Iteration, PetscMetaData, + inject_solve_mapper = MapNodes(Iteration, PetscMetaData, 'groupby').visit(iet) - if not injectsolve_mapper: + if not inject_solve_mapper: return iet, {} if kwargs['language'] not in petsc_languages: @@ -51,7 +51,7 @@ def lower_petsc(iet, **kwargs): if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): return finalize(iet), core_metadata() - unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()} + unique_grids = {i.expr.rhs.grid for (i,) in inject_solve_mapper.values()} # Assumption is that all solves are on the same grid if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") @@ -69,9 +69,9 @@ def lower_petsc(iet, **kwargs): # Map PETScSolve to its Section (for logging) section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) - for iters, (injectsolve,) in injectsolve_mapper.items(): + for iters, (inject_solve,) in inject_solve_mapper.items(): - builder = Builder(injectsolve, objs, iters, comm, section_mapper, **kwargs) + builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs) setup.extend(builder.solversetup.calls) @@ -129,18 +129,18 @@ class Builder: and other functionalities as needed. The class will be extended to accommodate different solver types by returning subclasses of the objects initialised in __init__, - depending on the properties of `injectsolve`. + depending on the properties of `inject_solve`. """ - def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): - self.injectsolve = injectsolve + def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs): + self.inject_solve = inject_solve self.objs = objs self.iters = iters self.comm = comm self.section_mapper = section_mapper self.kwargs = kwargs - self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) + self.coupled = isinstance(inject_solve.expr.rhs.field_data, MultipleFieldData) self.common_kwargs = { - 'injectsolve': self.injectsolve, + 'inject_solve': self.inject_solve, 'objs': self.objs, 'iters': self.iters, 'comm': self.comm, @@ -148,7 +148,7 @@ def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): **self.kwargs } self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs - self.common_kwargs['timedep'] = self.timedep + self.common_kwargs['time_dependence'] = self.time_dependence self.common_kwargs['cbbuilder'] = self.cbbuilder self.common_kwargs['logger'] = self.logger @@ -161,10 +161,10 @@ def objbuilder(self): ) @cached_property - def timedep(self): - time_mapper = self.injectsolve.expr.rhs.time_mapper - timedep_class = TimeDependent if time_mapper else NonTimeDependent - return timedep_class(**self.common_kwargs) + def time_dependence(self): + mapper = self.inject_solve.expr.rhs.time_mapper + time_class = TimeDependent if mapper else NonTimeDependent + return time_class(**self.common_kwargs) @cached_property def cbbuilder(self): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 503d773372..4f7bf01e8e 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -33,10 +33,10 @@ def __init__(self, **kwargs): self.rcompile = kwargs.get('rcompile', None) self.sregistry = kwargs.get('sregistry', None) self.concretize_mapper = kwargs.get('concretize_mapper', {}) - self.timedep = kwargs.get('timedep') + self.time_dependence = kwargs.get('time_dependence') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self._efuncs = OrderedDict() self._struct_params = [] @@ -93,23 +93,23 @@ def user_struct_callback(self): return self._user_struct_callback @property - def fielddata(self): - return self.injectsolve.expr.rhs.fielddata + def field_data(self): + return self.inject_solve.expr.rhs.field_data @property def arrays(self): - return self.fielddata.arrays + return self.field_data.arrays @property def target(self): - return self.fielddata.target + return self.field_data.target def _make_core(self): - self._make_matvec(self.fielddata.jacobian) + self._make_matvec(self.field_data.jacobian) self._make_formfunc() self._make_formrhs() - if self.fielddata.initialguess.exprs: - self._make_initialguess() + if self.field_data.initial_guess.exprs: + self._make_initial_guess() self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): @@ -134,7 +134,7 @@ def _make_matvec(self, jacobian, prefix='MatMult'): self._efuncs[cb.name] = cb def _create_matvec_body(self, body, jacobian): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -145,7 +145,7 @@ def _create_matvec_body(self, body, jacobian): y_matvec = self.arrays[jacobian.row_target]['y'] x_matvec = self.arrays[jacobian.col_target]['x'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) @@ -261,7 +261,7 @@ def _create_matvec_body(self, body, jacobian): return body def _make_formfunc(self): - F_exprs = self.fielddata.residual.F_exprs + F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, @@ -281,7 +281,7 @@ def _make_formfunc(self): self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs arrays = self.arrays @@ -290,7 +290,7 @@ def _create_formfunc_body(self, body): dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -397,7 +397,7 @@ def _create_formfunc_body(self, body): return Uxreplace(subs).visit(body) def _make_formrhs(self): - b_exprs = self.fielddata.residual.b_exprs + b_exprs = self.field_data.residual.b_exprs sobjs = self.solver_objs # Compile `b_exprs` into an IET via recursive compilation @@ -419,7 +419,7 @@ def _make_formrhs(self): self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs target = self.target @@ -441,7 +441,7 @@ def _create_form_rhs_body(self, body): sobjs['blocal'] ]) - b_arr = self.fielddata.arrays[target]['b'] + b_arr = self.field_data.arrays[target]['b'] vec_get_array = petsc_call( 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] @@ -451,7 +451,7 @@ def _create_form_rhs_body(self, body): 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -508,8 +508,8 @@ def _create_form_rhs_body(self, body): return Uxreplace(subs).visit(body) - def _make_initialguess(self): - exprs = self.fielddata.initialguess.exprs + def _make_initial_guess(self): + exprs = self.field_data.initial_guess.exprs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -531,7 +531,7 @@ def _make_initialguess(self): self._efuncs[cb.name] = cb def _create_initial_guess_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs target = self.target @@ -539,7 +539,7 @@ def _create_initial_guess_body(self, body): dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - x_arr = self.fielddata.arrays[target]['x'] + x_arr = self.field_data.arrays[target]['x'] vec_get_array = petsc_call( 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] @@ -549,7 +549,7 @@ def _create_initial_guess_body(self, body): 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -648,7 +648,7 @@ def submatrices_callback(self): @property def jacobian(self): - return self.injectsolve.expr.rhs.fielddata.jacobian + return self.inject_solve.expr.rhs.field_data.jacobian @property def main_matvec_callback(self): @@ -660,7 +660,7 @@ def main_matvec_callback(self): return self._main_matvec_callback def _make_core(self): - for sm in self.fielddata.jacobian.nonzero_submatrices: + for sm in self.field_data.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() @@ -731,7 +731,7 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self): - F_exprs = self.fielddata.residual.F_exprs + F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, @@ -750,14 +750,14 @@ def _make_whole_formfunc(self): self._efuncs[cb.name] = cb def _whole_formfunc_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -1026,11 +1026,11 @@ class BaseObjectBuilder: method to support specific use cases. """ def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.comm = kwargs.get('comm') - self.fielddata = self.injectsolve.expr.rhs.fielddata + self.field_data = self.inject_solve.expr.rhs.field_data self.solver_objs = self._build() def _build(self): @@ -1055,10 +1055,10 @@ def _build(self): functions via `SNESGetDM`. """ sreg = self.sregistry - targets = self.fielddata.targets + targets = self.field_data.targets snes_name = sreg.make_name(prefix='snes') - options_prefix = self.injectsolve.expr.rhs.options_prefix + options_prefix = self.inject_solve.expr.rhs.options_prefix base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), @@ -1085,7 +1085,7 @@ def _target_dependent(self, base_dict): that will be updated at each time step. """ sreg = self.sregistry - target = self.fielddata.target + target = self.field_data.target base_dict[f'{target.name}_ptr'] = StartPtr( sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype ) @@ -1102,8 +1102,8 @@ class CoupledObjectBuilder(BaseObjectBuilder): def _extend_build(self, base_dict): sreg = self.sregistry objs = self.objs - targets = self.fielddata.targets - arrays = self.fielddata.arrays + targets = self.field_data.targets + arrays = self.field_data.arrays base_dict['fields'] = PointerIS( name=sreg.make_name(prefix='fields'), nindices=len(targets) @@ -1113,14 +1113,14 @@ def _extend_build(self, base_dict): ) base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) - space_dims = len(self.fielddata.grid.dimensions) + space_dims = len(self.field_data.grid.dimensions) dim_labels = ["M", "N", "P"] base_dict.update({ dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - submatrices = self.fielddata.jacobian.nonzero_submatrices + submatrices = self.field_data.jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), @@ -1174,7 +1174,7 @@ def _extend_build(self, base_dict): def _target_dependent(self, base_dict): sreg = self.sregistry - targets = self.fielddata.targets + targets = self.field_data.targets for t in targets: name = t.name base_dict[f'{name}_ptr'] = StartPtr( @@ -1208,11 +1208,11 @@ def _target_dependent(self, base_dict): class BaseSetup: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') self.cbbuilder = kwargs.get('cbbuilder') - self.fielddata = self.injectsolve.expr.rhs.fielddata + self.field_data = self.inject_solve.expr.rhs.field_data self.calls = self._setup() @property @@ -1229,7 +1229,7 @@ def _setup(self): dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters + solver_params = self.inject_solve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) @@ -1252,7 +1252,7 @@ def _setup(self): global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(sobjs['xglobal'])]) - target = self.fielddata.target + target = self.field_data.target field_from_ptr = FieldFromPointer( target.function._C_field_data, target.function._C_symbol ) @@ -1366,7 +1366,7 @@ def _create_dmda_calls(self, dmda): def _create_dmda(self, dmda): objs = self.objs sobjs = self.solver_objs - grid = self.fielddata.grid + grid = self.field_data.grid nspace_dims = len(grid.dimensions) # MPI communicator @@ -1389,8 +1389,8 @@ def _create_dmda(self, dmda): args.append(dmda.dofs) # "Stencil width" -> size of overlap # TODO: Instead, this probably should be - # extracted from fielddata.target._size_outhalo? - stencil_width = self.fielddata.space_order + # extracted from field_data.target._size_outhalo? + stencil_width = self.field_data.space_order args.append(stencil_width) args.extend([objs['Null']]*nspace_dims) @@ -1411,7 +1411,7 @@ def _setup(self): sobjs = self.solver_objs dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters + solver_params = self.inject_solve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) @@ -1518,7 +1518,7 @@ def _setup(self): Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] ) - targets = self.fielddata.targets + targets = self.field_data.targets deref_dms = [ DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) @@ -1563,13 +1563,12 @@ def _setup(self): class Solver: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') self.iters = kwargs.get('iters') self.cbbuilder = kwargs.get('cbbuilder') - self.timedep = kwargs.get('timedep') - # TODO: Should/could _execute_solve be a cached_property? + self.time_dependence = kwargs.get('time_dependence') self.calls = self._execute_solve() def _execute_solve(self): @@ -1578,9 +1577,9 @@ def _execute_solve(self): the necessary calls to execute the SNES solver. """ sobjs = self.solver_objs - target = self.injectsolve.expr.rhs.fielddata.target + target = self.inject_solve.expr.rhs.field_data.target - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) b_efunc = self.cbbuilder._b_efunc @@ -1588,7 +1587,7 @@ def _execute_solve(self): rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) - vec_place_array = self.timedep.place_array(target) + vec_place_array = self.time_dependence.place_array(target) if self.cbbuilder.initial_guesses: initguess = self.cbbuilder.initial_guesses[0] @@ -1609,7 +1608,7 @@ def _execute_solve(self): dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] ) - vec_reset_array = self.timedep.reset_array(target) + vec_reset_array = self.time_dependence.reset_array(target) run_solver_calls = (struct_assignment,) + ( rhs_call, @@ -1631,7 +1630,7 @@ def spatial_body(self): root = filter_iterations(tree, key=lambda i: i.dim.is_Space) if root: root = root[0] - if self.injectsolve in FindNodes(PetscMetaData).visit(root): + if self.inject_solve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) spatial_body, = spatial_body return spatial_body @@ -1647,8 +1646,8 @@ def _execute_solve(self): sobjs = self.solver_objs xglob = sobjs['xglobal'] - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - targets = self.injectsolve.expr.rhs.fielddata.targets + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) + targets = self.inject_solve.expr.rhs.field_data.targets # TODO: optimise the ccode generated here pre_solve = () @@ -1667,7 +1666,7 @@ def _execute_solve(self): petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), # TODO: Need to call reset array - self.timedep.place_array(t), + self.time_dependence.place_array(t), petsc_call( 'DMLocalToGlobal', [dm, target_xloc, insert_vals, target_xglob] @@ -1709,12 +1708,12 @@ def _execute_solve(self): class NonTimeDependent: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.iters = kwargs.get('iters') self.sobjs = kwargs.get('solver_objs') self.kwargs = kwargs self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) - self.time_idx_to_symb = self.injectsolve.expr.rhs.time_mapper + self.time_idx_to_symb = self.inject_solve.expr.rhs.time_mapper def _origin_to_moddim_mapper(self, iters): return {} @@ -1773,7 +1772,7 @@ class TimeDependent(NonTimeDependent): """ @property def time_spacing(self): - return self.injectsolve.expr.rhs.grid.stepping_dim.spacing + return self.inject_solve.expr.rhs.grid.stepping_dim.spacing @cached_property def symb_to_moddim(self): diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py index f552b69969..40ad2711fa 100644 --- a/devito/petsc/logging.py +++ b/devito/petsc/logging.py @@ -120,15 +120,15 @@ def __getitem__(self, key): class PetscInfo(CompositeObject): __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper', - 'injectsolve', 'function_list') + 'inject_solve', 'function_list') def __init__(self, name, pname, logobjs, sobjs, section_mapper, - injectsolve, function_list): + inject_solve, function_list): self.logobjs = logobjs self.sobjs = sobjs self.section_mapper = section_mapper - self.injectsolve = injectsolve + self.inject_solve = inject_solve self.function_list = function_list mapper = {v: k for k, v in petsc_type_mappings.items()} @@ -138,7 +138,7 @@ def __init__(self, name, pname, logobjs, sobjs, section_mapper, @property def section(self): section = self.section_mapper.items() - return next((k[0].name for k, v in section if self.injectsolve in v), None) + return next((k[0].name for k, v in section if self.inject_solve in v), None) @property def summary_key(self): diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 55633d40fd..77022fc6b8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -68,13 +68,13 @@ def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=Non self.options_prefix = options_prefix def build_expr(self): - target, funcs, fielddata = self.linear_solve_args() + target, funcs, field_data = self.linear_solve_args() # Placeholder expression for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, self.solver_params, - fielddata=fielddata, + field_data=field_data, time_mapper=self.time_mapper, localinfo=localinfo, options_prefix=self.options_prefix @@ -93,13 +93,13 @@ def linear_solve_args(self): jacobian = Jacobian(target, exprs, arrays, self.time_mapper) residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, exprs, arrays) + initial_guess = InitialGuess(target, exprs, arrays) field_data = FieldData( target=target, jacobian=jacobian, residual=residual, - initialguess=initialguess, + initial_guess=initial_guess, arrays=arrays ) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 4f1f492c8e..fa87920571 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -60,7 +60,7 @@ class LinearSolveExpr(MetaData): """ __rargs__ = ('expr',) - __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', + __rkwargs__ = ('solver_parameters', 'field_data', 'time_mapper', 'localinfo', 'options_prefix') defaults = { @@ -73,7 +73,7 @@ class LinearSolveExpr(MetaData): } def __new__(cls, expr, solver_parameters=None, - fielddata=None, time_mapper=None, localinfo=None, + field_data=None, time_mapper=None, localinfo=None, options_prefix=None, **kwargs): if solver_parameters is None: @@ -87,7 +87,7 @@ def __new__(cls, expr, solver_parameters=None, obj._expr = expr obj._solver_parameters = solver_parameters - obj._fielddata = fielddata if fielddata else FieldData() + obj._field_data = field_data if field_data else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo obj._options_prefix = options_prefix @@ -113,8 +113,8 @@ def expr(self): return self._expr @property - def fielddata(self): - return self._fielddata + def field_data(self): + return self._field_data @property def solver_parameters(self): @@ -134,7 +134,7 @@ def options_prefix(self): @property def grid(self): - return self.fielddata.grid + return self.field_data.grid @classmethod def eval(cls, *args): @@ -158,14 +158,14 @@ class FieldData: the PETScArray representing the `target`. residual : Residual Defines the residual function F(target) = 0. - initialguess : InitialGuess + initial_guess : InitialGuess Defines the initial guess for the solution, which satisfies essential boundary conditions. arrays : dict A dictionary mapping `target` to its corresponding PETScArrays. """ def __init__(self, target=None, jacobian=None, residual=None, - initialguess=None, arrays=None, **kwargs): + initial_guess=None, arrays=None, **kwargs): self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: @@ -176,7 +176,7 @@ def __init__(self, target=None, jacobian=None, residual=None, ) self._jacobian = jacobian self._residual = residual - self._initialguess = initialguess + self._initial_guess = initial_guess self._arrays = arrays @property @@ -192,8 +192,8 @@ def residual(self): return self._residual @property - def initialguess(self): - return self._initialguess + def initial_guess(self): + return self._initial_guess @property def arrays(self): @@ -230,7 +230,7 @@ class MultipleFieldData(FieldData): Defines the matrix-vector products for the full system Jacobian. residual : MixedResidual Defines the residual function F(targets) = 0. - initialguess : InitialGuess + initial_guess : InitialGuess Defines the initial guess metadata, which satisfies essential boundary conditions. arrays : dict diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index d429f83831..45a48c80e8 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -76,9 +76,7 @@ def _(expr, targets): @singledispatch def centre_stencil(expr, target, as_coeff=False): """ - Extract the centre stencil from an expression. Its coefficient is what - would appear on the diagonal of the matrix system if the matrix were - formed explicitly. + Extract the centre stencil from an expression. Parameters ---------- diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 44668092ad..f297b7993f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -214,7 +214,7 @@ def test_LinearSolveExpr(): eqn = Eq(f, g.laplace) - linsolveexpr = LinearSolveExpr(eqn.rhs, fielddata=FieldData(target=f)) + linsolveexpr = LinearSolveExpr(eqn.rhs, field_data=FieldData(target=f)) # Check the solver parameters assert linsolveexpr.solver_parameters == \ @@ -653,7 +653,7 @@ def define(self, dimensions): petsc = PETScSolve([eq1, bc_1, bc_2], target=e) - jac = petsc.rhs.fielddata.jacobian + jac = petsc.rhs.field_data.jacobian assert jac.row_target == e assert jac.col_target == e @@ -699,7 +699,7 @@ def define(self, dimensions): petsc = PETScSolve([eq1, bc_1, bc_2], target=e) - res = petsc.rhs.fielddata.residual + res = petsc.rhs.field_data.residual assert res.target == e # NOTE: This is likely to change when PetscSection + DMDA is supported @@ -789,14 +789,14 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): assert len(callbacks1) == 8 assert len(callbacks2) == 6 - # Check fielddata type - fielddata1 = petsc1.rhs.fielddata - fielddata2 = petsc2.rhs.fielddata - fielddata3 = petsc3.rhs.fielddata + # Check field_data type + field0 = petsc1.rhs.field_data + field1 = petsc2.rhs.field_data + field2 = petsc3.rhs.field_data - assert isinstance(fielddata1, FieldData) - assert isinstance(fielddata2, FieldData) - assert isinstance(fielddata3, MultipleFieldData) + assert isinstance(field0, FieldData) + assert isinstance(field1, FieldData) + assert isinstance(field2, MultipleFieldData) @skipif('petsc') def test_coupled_structs(self): @@ -919,7 +919,7 @@ def test_mixed_jacobian(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j01 = jacobian.get_submatrix(0, 1) @@ -1031,7 +1031,7 @@ def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j01 = jacobian.get_submatrix(0, 1) j10 = jacobian.get_submatrix(1, 0) @@ -1080,7 +1080,7 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1130,7 +1130,7 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1183,7 +1183,7 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) From ff4b2d7185e939b8bdc7a3e45be5ad37c008d181 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Jul 2025 11:28:08 +0100 Subject: [PATCH 53/53] misc: flake8 --- devito/petsc/iet/passes.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index aadb711c76..d0118d6863 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -32,7 +32,7 @@ def lower_petsc(iet, **kwargs): # Check if PETScSolve was used inject_solve_mapper = MapNodes(Iteration, PetscMetaData, - 'groupby').visit(iet) + 'groupby').visit(iet) if not inject_solve_mapper: return iet, {} @@ -73,7 +73,7 @@ def lower_petsc(iet, **kwargs): builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs) - setup.extend(builder.solversetup.calls) + setup.extend(builder.solver_setup.calls) # Transform the spatial iteration loop with the calls to execute the solver subs.update({builder.solve.spatial_body: builder.calls}) @@ -147,13 +147,13 @@ def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs): 'section_mapper': self.section_mapper, **self.kwargs } - self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs + self.common_kwargs['solver_objs'] = self.object_builder.solver_objs self.common_kwargs['time_dependence'] = self.time_dependence self.common_kwargs['cbbuilder'] = self.cbbuilder self.common_kwargs['logger'] = self.logger @cached_property - def objbuilder(self): + def object_builder(self): return ( CoupledObjectBuilder(**self.common_kwargs) if self.coupled else @@ -172,7 +172,7 @@ def cbbuilder(self): if self.coupled else CBBuilder(**self.common_kwargs) @cached_property - def solversetup(self): + def solver_setup(self): return CoupledSetup(**self.common_kwargs) \ if self.coupled else BaseSetup(**self.common_kwargs)