Memory Blowup using Nonlinear solver when compiling Form

I am trying to implement an Arruda-Boyce Model and seem to be having a problem compiling. Perhaps I’m missing something stupid here but I just can’t see it.

My MWE is:


from mpi4py import MPI
import dolfinx 
import ufl 
import dolfinx.fem.petsc  

omega = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8, dolfinx.mesh.CellType.triangle) 
gdim = omega.geometry.dim 


V = dolfinx.fem.functionspace(omega, ("Lagrange", 2, (omega.geometry.dim, ))) 


dx = ufl.Measure("dx", domain=omega)
lam_m = dolfinx.fem.Constant(omega, 1.1)

u = dolfinx.fem.Function(V, name="Displacement") 
v = ufl.TestFunction(V)


voldim = len(u)
I =  ufl.Identity(voldim) 

F = ufl.variable(I + ufl.grad(u)) 
J = ufl.variable(ufl.det(F))


def inv_langevin(x):
    """According to R. Jedynak has max relative error 0.16% 
    Cf. https://journals.sagepub.com/doi/10.1177/1081286518811395 """ 
    return (0.756*x**5 +
             -1.383*x**4 +
             1.5733*x**3 +
             -2.9463*x**2 +
             3*x)/(1-x)
    
def inv_langevin(x):
    """According to R. Jedynak has max relative error 0.16% 
    Cf. https://journals.sagepub.com/doi/10.1177/1081286518811395 """ 
    
    return ( 0.88*x**3 +
             -2.88*x**2 +
             3*x)/(1-x)
    
isoF = F / ufl.det(F)**(1/gdim)
bstar = isoF * isoF.T
av_stretch = ufl.sqrt(ufl.tr(bstar)/gdim)
beta = inv_langevin(av_stretch/lam_m) 

psi = lam_m*( av_stretch*beta + lam_m*ufl.ln( beta / ufl.sinh(beta))) 
PK1 = ufl.diff(psi , F) 

Form = ufl.inner(PK1, ufl.grad(v)) * dx 
#res = ufl.extract_blocks(Form)

which compiles if I use

a, L = ufl.system(Form)
a_compiled = dolfinx.fem.compile_form(MPI.COMM_WORLD, a)
L_compiled = dolfinx.fem.compile_form(MPI.COMM_WORLD, L)

but gives a memory error if I do it the way I’d like:

solver = dolfinx.fem.petsc.NonlinearProblem(
    Form,
    u=u,
    J=None,
    bcs=None, 
    petsc_options={
        "snes_monitor": None, 
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
        "mat_mumps_icntl_14": 120,
        "ksp_error_if_not_converged": True,
        "snes_error_if_not_converged": True,
    },
    petsc_options_prefix="arruda_",
)

MemoryError: Unable to allocate 171. GiB for an array with shape (151321, 151321) and data type float64

If you had given the full traceback:

Traceback (most recent call last):
  File "/root/shared/mwe.py", line 55, in <module>
    solver = dolfinx.fem.petsc.NonlinearProblem(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/petsc.py", line 1249, in __init__
    self._J = _create_form(
              ^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 449, in form
    return _create_form(form)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 441, in _create_form
    return _form(form)
           ^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/forms.py", line 361, in _form
    ufcx_form, module, code = jit.ffcx_jit(
                              ^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/jit.py", line 60, in mpi_jit
    return local_jit(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/jit.py", line 215, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], options=p_ffcx, **p_jit)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/codegeneration/jit.py", line 244, in compile_forms
    raise e
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/codegeneration/jit.py", line 224, in compile_forms
    impl = _compile_objects(
           ^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/codegeneration/jit.py", line 349, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/compiler.py", line 113, in compile_ufl_objects
    ir = compute_ir(analysis, _object_names, _prefix, options, visualise)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/representation.py", line 151, in compute_ir
    _compute_integral_ir(
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/representation.py", line 402, in _compute_integral_ir
    integral_ir = compute_integral_ir(
                  ^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/integral.py", line 170, in compute_integral_ir
    mt_table_reference = build_optimized_tables(
                         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/elementtables.py", line 528, in build_optimized_tables
    tabletype = analyse_table_type(tbl)
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/elementtables.py", line 682, in analyse_table_type
    elif is_quadrature_table(table, rtol=rtol, atol=atol):
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ffcx/ir/elementtables.py", line 644, in is_quadrature_table
    Id = np.eye(num_points)
         ^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/numpy/lib/_twodim_base_impl.py", line 222, in eye
    m = zeros((N, M), dtype=dtype, order=order, device=device)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 171. GiB for an array with shape (151321, 151321) and data type float64
Exception ignored in: <function NonlinearProblem.__del__ at 0x7bbb72dd45e0>
Traceback (most recent call last):
  File "/usr/local/dolfinx-real/lib/python3.12/dist-packages/dolfinx/fem/petsc.py", line 1361, in __del__
    lambda obj: obj is not None, (self._snes, self._A, self._b, self._x, self._P_mat)
                                  ^^^^^^^^^^

it becomes clear that this is due to the quadrature degree estimation of UFL.
I’ve tried to explain how to get this estimate in:

which you can get by adding:

pulled_back = ufl.algorithms.compute_form_data(
    Form,
    do_apply_function_pullbacks=True,
    do_apply_integral_scaling=True,
    do_apply_geometry_lowering=True,
    preserve_geometry_types=(ufl.classes.Jacobian,),
)
print(list(itg.metadata() for itg in pulled_back.integral_data[0].integrals))

to your code.
And you will get:

[{'estimated_polynomial_degree': 388}]

By setting the quadrature degree in the construction of dx:

dx = ufl.Measure("dx", domain=omega, metadata={"quadrature_degree": 25})

the code will run.
It is of course up to you to select an appropriate degree.
25 is just an example. In 3D, we only have specialized quadrature rules up to degree 16.

2 Likes

Thanks for the response. I presume here that the automatic estimation of the needed quadrature rule is based on the inverse_langevin function I used, which blows up for x→1 (but this should basically never actually happen). Is that a reasonable assumption? The quadrature degree indeed changes depending on which approximant I use.

Is there some way to tell the compiler what I expect as reasonable ranges?

I’m a bit scared to second-guess the estimated necessary quadrature degree but at degree 388 I would get a pretty good result even with Monte-Carlo integration using one tenth the number of sample points element-wise.

I’m curious why your traceback is so much shorter than mine:

---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
Cell In[3], line 1
----> 1 solver = dolfinx.fem.petsc.NonlinearProblem(
      2     Form,
      3     u=u,
      4     J=None,
      5     bcs=None,
      6     #entity_maps=[gamma_to_omega],
      7     petsc_options={
      8         "snes_monitor": None, 
      9         "ksp_type": "preonly",
     10         "pc_type": "lu",
     11         "pc_factor_mat_solver_type": "mumps",
     12         "mat_mumps_icntl_14": 120,
     13         "ksp_error_if_not_converged": True,
     14         "snes_error_if_not_converged": True,
     15     },
     16     petsc_options_prefix="contact_",
     17 )

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/fem/petsc.py:1249, in NonlinearProblem.__init__(self, F, u, petsc_options_prefix, bcs, J, P, kind, petsc_options, form_compiler_options, jit_options, entity_maps)
   1246 if J is None:
   1247     J = derivative_block(F, u)
-> 1249 self._J = _create_form(
   1250     J,
   1251     form_compiler_options=form_compiler_options,
   1252     jit_options=jit_options,
   1253     entity_maps=entity_maps,
   1254 )
   1256 if P is not None:
   1257     self._preconditioner = _create_form(
   1258         P,
   1259         form_compiler_options=form_compiler_options,
   1260         jit_options=jit_options,
   1261         entity_maps=entity_maps,
   1262     )

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/fem/forms.py:449, in form(form, dtype, form_compiler_options, jit_options, entity_maps)
    446     else:
    447         return form
--> 449 return _create_form(form)

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/fem/forms.py:441, in form.<locals>._create_form(form)
    431 """Recursively convert ufl.Forms to dolfinx.fem.Form.
    432 
    433 Args:
   (...)    438     A ``dolfinx.fem.Form`` or a list of ``dolfinx.fem.Form``.
    439 """
    440 if isinstance(form, ufl.Form):
--> 441     return _form(form)
    442 elif isinstance(form, ufl.ZeroBaseForm):
    443     return _zero_form(form)

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/fem/forms.py:361, in form.<locals>._form(form)
    359 if msh is None:
    360     raise RuntimeError("Expecting to find a Mesh in the form.")
--> 361 ufcx_form, module, code = jit.ffcx_jit(
    362     msh.comm, form, form_compiler_options=form_compiler_options, jit_options=jit_options
    363 )
    365 # For each argument in form extract its function space
    366 V = [arg.ufl_function_space()._cpp_object for arg in form.arguments()]

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/jit.py:60, in mpi_jit_decorator.<locals>.mpi_jit(comm, *args, **kwargs)
     56 @functools.wraps(local_jit)
     57 def mpi_jit(comm, *args, **kwargs):
     58     # Just call JIT compiler when running in serial
     59     if comm.size == 1:
---> 60         return local_jit(*args, **kwargs)
     62     # Default status (0 == ok, 1 == fail)
     63     status = 0

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/dolfinx/jit.py:215, in ffcx_jit(ufl_object, form_compiler_options, jit_options)
    213 # Switch on type and compile, returning cffi object
    214 if isinstance(ufl_object, ufl.Form):
--> 215     r = ffcx.codegeneration.jit.compile_forms([ufl_object], options=p_ffcx, **p_jit)
    216 elif isinstance(ufl_object, tuple) and isinstance(ufl_object[0], ufl.core.expr.Expr):
    217     r = ffcx.codegeneration.jit.compile_expressions([ufl_object], options=p_ffcx, **p_jit)

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/codegeneration/jit.py:244, in compile_forms(forms, options, cache_dir, timeout, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries, visualise)
    242     except Exception:
    243         pass
--> 244     raise e
    246 obj, module = _load_objects(cache_dir, module_name, form_names)
    247 return obj, module, (decl, impl)

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/codegeneration/jit.py:224, in compile_forms(forms, options, cache_dir, timeout, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries, visualise)
    221     for name in form_names:
    222         decl += form_template.format(name=name)
--> 224     impl = _compile_objects(
    225         decl,
    226         forms,
    227         form_names,
    228         module_name,
    229         p,
    230         cache_dir,
    231         cffi_extra_compile_args,
    232         cffi_verbose,
    233         cffi_debug,
    234         cffi_libraries,
    235         visualise=visualise,
    236     )
    237 except Exception as e:
    238     try:
    239         # remove c file so that it will not timeout next time

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/codegeneration/jit.py:349, in _compile_objects(decl, ufl_objects, object_names, module_name, options, cache_dir, cffi_extra_compile_args, cffi_verbose, cffi_debug, cffi_libraries, visualise)
    345 libraries = _libraries + cffi_libraries if cffi_libraries is not None else _libraries
    347 # JIT uses module_name as prefix, which is needed to make names of all struct/function
    348 # unique across modules
--> 349 _, code_body = ffcx.compiler.compile_ufl_objects(
    350     ufl_objects, prefix=module_name, options=options, visualise=visualise
    351 )
    353 # Raise error immediately prior to compilation if no support for C99
    354 # _Complex. Doing this here allows FFCx to be used for complex codegen on
    355 # Windows.
    356 if sys.platform.startswith("win32"):

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/compiler.py:113, in compile_ufl_objects(ufl_objects, options, object_names, prefix, visualise)
    111 # Stage 2: intermediate representation
    112 cpu_time = time()
--> 113 ir = compute_ir(analysis, _object_names, _prefix, options, visualise)
    114 _print_timing(2, time() - cpu_time)
    116 # Stage 3: code generation

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/representation.py:151, in compute_ir(analysis, object_names, prefix, options, visualise)
    141     for itg_index, itg_data in enumerate(fd.integral_data):  # type: ignore
    142         integral_names[(fd_index, itg_index)] = naming.integral_name(
    143             fd.original_form,  # type: ignore
    144             itg_data.integral_type,
   (...)    147             prefix,
    148         )
    150 irs = [
--> 151     _compute_integral_ir(
    152         fd,
    153         i,
    154         analysis.element_numbers,
    155         integral_names,
    156         options,
    157         visualise,
    158     )
    159     for (i, fd) in enumerate(analysis.form_data)
    160 ]
    161 ir_integrals = list(itertools.chain(*irs))
    163 integral_domains = {
    164     i.expression.name: set(j[0] for j in i.expression.integrand.keys()) for a in irs for i in a
    165 }

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/representation.py:402, in _compute_integral_ir(form_data, form_index, element_numbers, integral_names, options, visualise)
    396 integrand_map: dict[basix.CellType, dict[QuadratureRule, ufl.core.expr.Expr]] = {
    397     cell_type: {rule: integral.integrand() for rule, integral in cell_integrals.items()}
    398     for cell_type, cell_integrals in sorted_integrals.items()
    399 }
    401 # Build more specific intermediate representation
--> 402 integral_ir = compute_integral_ir(
    403     itg_data.domain.ufl_cell(),
    404     itg_data.integral_type,
    405     expression_ir["entity_type"],
    406     integrand_map,
    407     expression_ir["tensor_shape"],
    408     options,
    409     visualise,
    410 )
    412 expression_ir.update(integral_ir)
    414 # Fetch name

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/integral.py:170, in compute_integral_ir(cell, integral_type, entity_type, integrands, argument_shape, p, visualise)
    167     if domain.topological_dimension() != cell.topological_dimension():
    168         is_mixed_dim = True
--> 170 mt_table_reference = build_optimized_tables(
    171     quadrature_rule,
    172     cell,
    173     integral_type,
    174     entity_type,
    175     initial_terminals.values(),
    176     ir["unique_tables"][integral_domain],
    177     use_sum_factorization=p["sum_factorization"],
    178     is_mixed_dim=is_mixed_dim,
    179     rtol=p["table_rtol"],
    180     atol=p["table_atol"],
    181 )
    183 # Fetch unique tables for this quadrature rule
    184 table_types = {v.name: v.ttype for v in mt_table_reference.values()}

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/elementtables.py:528, in build_optimized_tables(quadrature_rule, cell, integral_type, entity_type, modified_terminals, existing_tables, use_sum_factorization, is_mixed_dim, rtol, atol)
    526 # Clean up table
    527 tbl = clamp_table_small_numbers(t["array"], rtol=rtol, atol=atol)
--> 528 tabletype = analyse_table_type(tbl)
    530 if tabletype in piecewise_ttypes:
    531     # Reduce table to dimension 1 along num_points axis in generated code
    532     tbl = tbl[:, :, :1, :]

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/elementtables.py:682, in analyse_table_type(table, rtol, atol)
    679 elif is_ones_table(table, rtol=rtol, atol=atol):
    680     # All values are 1.0
    681     ttype = "ones"
--> 682 elif is_quadrature_table(table, rtol=rtol, atol=atol):
    683     # Identity matrix mapping points to dofs (separately on each entity)
    684     ttype = "quadrature"
    685 else:
    686     # Equal for all points on a given entity

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/ffcx/ir/elementtables.py:644, in is_quadrature_table(table, rtol, atol)
    642 """Check if table is a quadrature table."""
    643 _, num_entities, num_points, num_dofs = table.shape
--> 644 Id = np.eye(num_points)
    645 return num_points == num_dofs and all(
    646     np.allclose(table[0, i, :, :], Id, rtol=rtol, atol=atol) for i in range(num_entities)
    647 )

File ~/miniforge3/envs/fenx2510/lib/python3.14/site-packages/numpy/lib/_twodim_base_impl.py:235, in eye(N, M, k, dtype, order, device, like)
    233 if M is None:
    234     M = N
--> 235 m = zeros((N, M), dtype=dtype, order=order, device=device)
    236 if k >= M:
    237     return m

MemoryError: Unable to allocate 171. GiB for an array with shape (151321, 151321) and data type float64

I didn’t supply it for that reason but I will in future so that one can inspect without running.

It is not the compiler that estimates the quadrature, it is ufl itself. You can use the code:

to see how it interacts with different parts of your form. The ruleset for estimation is found in the file: ufl/ufl/algorithms/estimate_degrees.py at c0af0f3661f85ab601dfa852301a690de08095d3 · FEniCS/ufl · GitHub

OK checking the energy before calculating its Gateaux derivative gives

pulled_back = ufl.algorithms.compute_form_data(
    lam_m*( av_stretch*beta + lam_m*ufl.ln( beta / ufl.sinh(beta))) *dx,
    do_apply_function_pullbacks=True,
    do_apply_integral_scaling=True,
    do_apply_geometry_lowering=True,
    preserve_geometry_types=(ufl.classes.Jacobian,),
)
print(list(itg.metadata() for itg in pulled_back.integral_data[0].integrals))

gives

[{'quadrature_degree': 12, 'estimated_polynomial_degree': 76}]

which tells me a lot about what is happening. Interestingly its not the singularity at x=1 in the inverse_langevin function because even the Taylor Series approximations get high estimations of polynomial degree.

The function and its inverse are smooth over their domains so I will have to make some assumptions about the maximum variation over a single element check the behaviour of the solution when refining.

Thanks for your help again.