firedrakeproject / firedrake

Firedrake is an automated system for the portable solution of partial differential equations using the finite element method (FEM)
https://firedrakeproject.org
Other
502 stars 158 forks source link

Static Condensation: Slate vs Preconditioner approach #2634

Closed volpatto closed 1 year ago

volpatto commented 1 year ago

Hi firedrakers!

Description

I have been working on a primal hybrid formulation. You can understand it as a primal version of HDG (I think that formulation details are not required, so I will skip most of them). The resulting algebraic system has the following structure:

$$ M x=f $$

with

$$ \mathrm{M}:=\left[\begin{array}{ll} M{00} & M{01} \ M{10} & M{11} \end{array}\right], \quad \mathrm{x}:=\left[\begin{array}{l} \mathrm{p} \ \lambda \end{array}\right], \quad \mathrm{f}:=\left[\begin{array}{c} \mathrm{f}_0 \ \mathrm{f}_1 \end{array}\right] $$

where $p$ is vector associated to the primal variable and $\lambda$ with the Lagrange multiplier.

Condensing the system, we have:

$$ \left(M{11}-M{10}M{00}^{-1} M{01}\right) \lambda=f1 - M{10}M_{00}^{-1} \mathrm{f}_0 $$

and we can recover the solution for the primal variable with:

$$ \mathbf{p}=M_{00}^{-1}\left(f0 - M{01} \boldsymbol{\lambda}\right) $$

I have to solve the problem through Static Condensation.

Issue

I tried two approaches to solve the condensed system. Below I provide the relevant code snippets for each case.

Using Static Condensation as PC in solver parameters.

This approach is fine, works as expected. I set the part to solve the problem as below:

F = a - L

# Solving with Static Condensation
print("*******************************************\nSolving using static condensation.\n")
params = {
    "snes_type": "ksponly",
    "mat_type": "matfree",
    "pmat_type": "matfree",
    "ksp_type": "preonly",
    "pc_type": "python",
    # Use the static condensation PC for hybridized problems
    # and use a direct solve on the reduced system for lambda_h
    "pc_python_type": "firedrake.SCPC",
    "pc_sc_eliminate_fields": "0",
    "condensed_field": {
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
        "ksp_monitor_true_residual": None,
    },
}
problem = NonlinearVariationalProblem(F, solution)
solver = NonlinearVariationalSolver(problem, solver_parameters=params)
solver.solve()
p_h, lambda_h = solution.split()

Solving with Slate

I am not sure with this is the correct term the previous approach also uses Slate under the hood, whatever. The point is I'm obtaining a lot of outputs related to the generated kernels, and this stuff is too much for my shallow understanding of Slate (and Firedrake in general). I tried to follow the Slate paper. Here is the code:

F = a - L

# Local computations (solving explicitly with Slate)
a_form = lhs(F)
b_form = rhs(F)
_A = Tensor(a_form)
_F = Tensor(b_form)
Vector
A = _A.blocks
Fc = _F.blocks
schur_factor = A[1, 0] * A[0, 0].inv
Sexp = A[1, 1] - schur_factor * A[0, 1]
Eexp = Fc[1] - schur_factor * Fc[0]
S = assemble(Sexp)
E = assemble(Eexp)
lambda_local = Function(T)

# Solving for the multiplier
solver_parameters = {
    "ksp_type": "preonly",
    "pc_type": "lu"
}
solve(S, lambda_local, E, solver_parameters=solver_parameters)

# Recovering the solution for the primal variable
Lambda = AssembledVector(lambda_local)
p_h = Function(V)
A00 = A[0, 0]
p_sys = A00.solve(Fc[0] - A[0, 1] * Lambda, decomposition="PartialPivLu")
assemble(p_sys, p_h)

This approach did not work. The output last few lines are:

---------------------------------------------------------------------------
KERNEL: subkernel0_cell_to__cell_integral_otherwise
---------------------------------------------------------------------------
ARGUMENTS:
A: type: np:dtype('float64'), shape: (3, 3), dim_tags: (stride:3, stride:1) aspace: local
coords: type: np:dtype('float64'), shape: (6), dim_tags: (stride:1) aspace: private
---------------------------------------------------------------------------
DOMAINS:
{ [k] : 0 <= k <= 2 }
{ [j] : 0 <= j <= 2 }
{ [k_0] : 0 <= k_0 <= 2 }
---------------------------------------------------------------------------
INAME TAGS:
j: None
k: None
k_0: None
---------------------------------------------------------------------------
TEMPORARIES:
t0: type: np:dtype('float64'), shape: () aspace:local
t1: type: np:dtype('float64'), shape: () aspace:local
t2: type: np:dtype('float64'), shape: () aspace:local
t3: type: np:dtype('float64'), shape: () aspace:local
t4: type: np:dtype('float64'), shape: () aspace:local
t5: type: np:dtype('float64'), shape: () aspace:local
t6: type: np:dtype('float64'), shape: () aspace:local
t7: type: np:dtype('float64'), shape: () aspace:local
t8: type: np:dtype('float64'), shape: () aspace:local
t9: type: np:dtype('float64'), shape: () aspace:local
t10: type: np:dtype('float64'), shape: () aspace:local
t11: type: np:dtype('float64'), shape: () aspace:local
t12: type: np:dtype('float64'), shape: () aspace:local
t13: type: np:dtype('float64'), shape: () aspace:local
t14: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1) aspace:local
t15: type: np:dtype('float64'), shape: () aspace:local
t16: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1) aspace:local
t17: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1) aspace:local
t18: type: np:dtype('float64'), shape: (3), dim_tags: (N0:stride:1) aspace:local
---------------------------------------------------------------------------
INSTRUCTIONS:
↱  CODE(|)  {id=insn}
│  PetscLogEventBegin(ID_Log_Event_subkernel0_cell_to__cell_integral_otherwise,0,0,0,0);
└↱ t0 = (-0.9999999999999988)*coords[0] + 0.9999999999999988*coords[2]  {id=insn_0}
↱└ t1 = (-1.0)*coords[1] + coords[5]  {id=insn_1}
└↱ t2 = (-1.0)*coords[0] + coords[4]  {id=insn_2}
↱└ t3 = (-0.9999999999999988)*coords[1] + 0.9999999999999988*coords[3]  {id=insn_3}
└↱ t4 = t0*t1 + (-1.0)*t2*t3  {id=insn_4}
↱└ t5 = (-1.0)*0.5*a̲b̲s(t4)  {id=insn_5}
└↱ t6 = 1.0 / t4  {id=insn_6}
↱└ t7 = (-1.0)*t3*t6  {id=insn_7}
└↱ t8 = t1*t6  {id=insn_8}
↱└ t9 = t0*t6  {id=insn_9}
└↱ t10 = (-1.0)*t2*t6  {id=insn_10}
↱└ t11 = t5*((-1.0)*t7*t8 + (-1.0)*t9*t10)  {id=insn_11}
└↱ t12 = t5*((-1.0)*t8*t8 + (-1.0)*t10*t10)  {id=insn_12}
↱└ t13 = t5*((-1.0)*t7*t7 + (-1.0)*t9*t9)  {id=insn_13}
└↱ t15 = t5*((-1.0)*t8*t7 + (-1.0)*t10*t9)  {id=insn_14}
 │ for k
↱└   t17[k] = t16[k]*t15 + t14[k]*t13  {id=insn_15}
└↱   t18[k] = t16[k]*t12 + t14[k]*t11  {id=insn_16}
 │ end k
 │ for j, k_0
↱└     A[j + k_0 // 3, k_0 + (-3)*(k_0 // 3)] = A[j + k_0 // 3, k_0 + (-3)*(k_0 // 3)] + t18[k_0]*t16[j] + t17[k_0]*t14[j]  {id=insn_17}
│  end j, k_0
└  CODE(|)  {id=insn_18}
   PetscLogEventEnd(ID_Log_Event_subkernel0_cell_to__cell_integral_otherwise,0,0,0,0);
---------------------------------------------------------------------------
===========================================================================
Traceback (most recent call last):
  File "/home/diego/firedrake/src/PyOP2/pyop2/global_kernel.py", line 308, in __call__
    func = self._func_cache[key]
KeyError: 139767866506032

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/diego/firedrake/src/PyOP2/pyop2/compilation.py", line 370, in get_so
    return ctypes.CDLL(soname)
  File "/usr/lib/python3.8/ctypes/__init__.py", line 373, in __init__
    self._handle = _dlopen(self._name, mode)
OSError: /home/diego/firedrake/.cache/pyop2/dc/559c8192ff9d3794baf3637e9a83da.so: cannot open shared object file: No such file or directory

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/diego/firedrake_scripts/scripts/2D/primal_hdg_poisson_2D_slate.py", line 100, in <module>
    assemble(p_sys, p_h)
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/diego/firedrake/src/firedrake/firedrake/adjoint/assembly.py", line 19, in wrapper
    output = assemble(*args, **kwargs)
  File "/home/diego/firedrake/src/firedrake/firedrake/assemble.py", line 101, in assemble
    return _assemble_form(expr, *args, **kwargs)
  File "/home/diego/firedrake/src/firedrake/firedrake/assemble.py", line 290, in _assemble_form
    return assembler.assemble()
  File "/home/diego/firedrake/src/firedrake/firedrake/assemble.py", line 399, in assemble
    parloop()
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/diego/firedrake/src/PyOP2/pyop2/parloop.py", line 213, in __call__
    self._compute(self.iterset.core_part)
  File "/home/diego/firedrake/src/PyOP2/pyop2/parloop.py", line 195, in _compute
    self.global_kernel(self.comm, part.offset, part.offset+part.size, *self.arglist)
  File "/home/diego/firedrake/src/PyOP2/pyop2/global_kernel.py", line 310, in __call__
    func = self.compile(comm)
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/diego/firedrake/src/PyOP2/pyop2/global_kernel.py", line 380, in compile
    return compilation.load(self, extension, self.name,
  File "/home/diego/firedrake/src/PyOP2/pyop2/compilation.py", line 641, in load
    dll = compiler(cppargs, ldargs, cpp=cpp, comm=comm).get_so(code, extension)
  File "/home/diego/firedrake/src/PyOP2/pyop2/compilation.py", line 380, in get_so
    f.write(jitmodule.code_to_compile)
  File "/home/diego/firedrake/src/PyOP2/pyop2/utils.py", line 62, in __get__
    obj.__dict__[self.__name__] = result = self.fget(obj)
  File "/home/diego/firedrake/src/PyOP2/pyop2/global_kernel.py", line 350, in code_to_compile
    code = lp.generate_code_v2(wrapper)
  File "/home/diego/firedrake/src/loopy/loopy/codegen/__init__.py", line 754, in generate_code_v2
    program = linearize(program)
  File "/home/diego/firedrake/src/loopy/loopy/schedule/__init__.py", line 2191, in linearize
    pre_schedule_checks(t_unit)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 1133, in pre_schedule_checks
    check_bounds(t_unit)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 766, in check_bounds
    _check_bounds_inner_rec(t_unit[epoint],
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 751, in _check_bounds_inner_rec
    _check_bounds_inner(kernel, callables_table)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 733, in _check_bounds_inner
    insn.with_transformed_expressions(run_acm)
  File "/home/diego/firedrake/src/loopy/loopy/kernel/instruction.py", line 1021, in with_transformed_expressions
    expression = f(self.expression)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 730, in run_acm
    acm(expr, domain_with_assumptions, insn.id)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 707, in map_call
    _check_bounds_inner_rec(subkernel, self.callables_table)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 751, in _check_bounds_inner_rec
    _check_bounds_inner(kernel, callables_table)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 733, in _check_bounds_inner
    insn.with_transformed_expressions(run_acm)
  File "/home/diego/firedrake/src/loopy/loopy/kernel/instruction.py", line 1021, in with_transformed_expressions
    expression = f(self.expression)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 730, in run_acm
    acm(expr, domain_with_assumptions, insn.id)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 707, in map_call
    _check_bounds_inner_rec(subkernel, self.callables_table)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 751, in _check_bounds_inner_rec
    _check_bounds_inner(kernel, callables_table)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 733, in _check_bounds_inner
    insn.with_transformed_expressions(run_acm)
  File "/home/diego/firedrake/src/loopy/loopy/kernel/instruction.py", line 856, in with_transformed_expressions
    expression = f(self.expression)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 730, in run_acm
    acm(expr, domain_with_assumptions, insn.id)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 651, in map_sum
    self.rec(child, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 651, in map_sum
    self.rec(child, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 651, in map_sum
    self.rec(child, *args, **kwargs)
  File "/home/diego/firedrake/lib/python3.8/site-packages/pymbolic/mapper/__init__.py", line 129, in __call__
    return method(expr, *args, **kwargs)
  File "/home/diego/firedrake/src/loopy/loopy/check.py", line 619, in map_subscript
    raise LoopyIndexError("'%s' in instruction '%s' "
loopy.diagnostic.LoopyIndexError: 'w_0[2]' in instruction 'insn_11' accesses out-of-bounds array element (could not establish '{ [2] }' is a subset of '{ [i0] : i0 = 0 }').

So what I have missed? Is this expected? Any suggestion to fix it? Thanks in advance!

Supporting material

The codes and output.

Firedrake status

Firedrake Configuration:
    package_manager: True
    minimal_petsc: False
    mpicc: None
    mpicxx: None
    mpif90: None
    mpiexec: None
    disable_ssh: True
    honour_petsc_dir: False
    with_parmetis: False
    slepc: True
    packages: []
    honour_pythonpath: False
    opencascade: False
    tinyasm: False
    petsc_int_type: int64
    cache_dir: /home/diego/firedrake/.cache
    complex: False
    remove_build_files: False
    with_blas: None
Additions:
    None
Environment:
    PYTHONPATH: None
    PETSC_ARCH: None
    PETSC_DIR: None
Status of components:
---------------------------------------------------------------------------
|Package             |Branch                        |Revision  |Modified  |
---------------------------------------------------------------------------
|COFFEE              |master                        |70c1e66   |False     |
|FInAT               |master                        |4112fb8   |False     |
|PyOP2               |master                        |1ba576ef  |False     |
|fiat                |master                        |1a446cc   |False     |
|firedrake           |master                        |4bb713b4  |False     |
|h5py                |firedrake                     |78531f08  |False     |
|libspatialindex     |master                        |4768bf3   |True      |
|libsupermesh        |master                        |69012e5   |False     |
|loopy               |main                          |3988272b  |False     |
|petsc               |firedrake                     |729b7b4f1b|False     |
|pyadjoint           |master                        |0dcfe22   |False     |
|slepc               |firedrake                     |759057d1d |False     |
|tsfc                |master                        |351994d   |False     |
|ufl                 |master                        |0c592ec5  |False     |
---------------------------------------------------------------------------
sv2518 commented 1 year ago

If you change the line where you build _F to _F = AssembledVector(assemble(b_form)) this will work fine. Just for reference, we have one test in our test suite which is doing something similar (handwrite the SCPC, but for a 3x3 system) here https://github.com/firedrakeproject/firedrake/blob/2946e2fb6309e5b71fb90a8a2512a12d79b7d06b/tests/slate/test_variational_prb.py#L76 I marked the relevant line in the link.

volpatto commented 1 year ago

Oh, so that is it! Thanks, @sv2518! And thanks also for indicating the relevant test, I checked the tests but missed this one.