Computing entropy viscosity with UFL

Hello FEniCS community,

I’m working on implementing an entropy viscosity method for shock capturing in the 1D Inviscid Burgers’ equation using FEniCSx. However, I’m encountering an error when trying to run my code. Here’s a brief overview of what I’m trying to do:

  1. I’m using a DG(1) space on a 1D mesh.
  2. I’ve implemented the entropy viscosity as described in equation 2.3 to 2.6 in https://people.tamu.edu/~guermond//PUBLICATIONS/MS/non_stationnary_jlg_rp_bp.pdf.
  3. The error occurs when trying to create the nonlinear problem.

I am not sure if this is the right way to implement the norm:
\left\|E\left(u_h\right)-\bar{E}\left(u_h\right)\right\|_{\infty, \Omega}
, average:
\bar{E}\left(u_h\right)
and also the min condition.
\nu_h:=S\left(\min \left(\nu_{\max }, \nu_E\right)\right), with S = I the identity.

Attached is my attempt at implementing the entropy viscosity for capturing shock:

import numpy as np
import ufl
import dolfinx
from dolfinx import mesh, fem, io
from dolfinx.fem import Function, FunctionSpace, dirichletbc, locate_dofs_topological
from ufl import FacetNormal, dx, ds, dS, TestFunction, TrialFunction, grad, inner, max_value, jump, avg, sin, cos, exp, as_vector, CellDiameter, conditional
from mpi4py import MPI
import dolfinx.fem.petsc
import dolfinx.nls.petsc
from petsc4py import PETSc

# Domain parameters
x_R, x_L = 1.0, 0.0
N_el = 2**7
L = x_R - x_L
cell_size = L / N_el

# Time marching parameters
T = 0.5
dt = 1e-3

# Create mesh and function space
domain = mesh.create_interval(MPI.COMM_WORLD, N_el, [x_L, x_R])
V = FunctionSpace(domain, ("DG", 1))

dt_ufl = fem.Constant(domain, dt)

# Initial condition
def initial_condition(x):
    return np.maximum(np.sin(2.0 * np.pi * x[0]), 0.0)

u0 = Function(V)
u0.interpolate(initial_condition)

inlet_value = fem.Constant(domain, 0.0)

# Boundary conditions
def inlet(x):
    return np.isclose(x[0], x_L)

def outlet(x):
    return np.isclose(x[0], x_R)

inlet_facets = mesh.locate_entities_boundary(domain, domain.topology.dim - 1, inlet)
outlet_facets = mesh.locate_entities_boundary(domain, domain.topology.dim - 1, outlet)

inlet_dofs = locate_dofs_topological(V, domain.topology.dim - 1, inlet_facets)
outlet_dofs = locate_dofs_topological(V, domain.topology.dim - 1, outlet_facets)

bc_inlet = dirichletbc(inlet_value, inlet_dofs, V)

# Define variational problem
u = TrialFunction(V)
phi = TestFunction(V)
n = FacetNormal(domain)

def Fc(u):
    return u**2 / 2

def E(u):
    return 0.5 * u**2

def F(u):
    return (1 / 3) * u**3

# Entropy viscosity parameters
c_E = 1.0
c_max = 0.5

def entropy_residual(u, u_old):
    return (E(u) - E(u_old)) / dt + ufl.div(F(u))

def entropy_viscosity(u, u_old):
    h = CellDiameter(domain)
    D_h = entropy_residual(u, u_old)
    E_avg = fem.assemble_scalar(fem.form(E(u) * dx(domain))) / fem.assemble_scalar(fem.form(1 * dx(domain)))
    E_diff = E(u) - E_avg
    E_diff_max = fem.assemble_scalar(fem.form(ufl.max_value(abs(E_diff), 1e-15) * dx(domain)))
    nu_E = c_E * h**2 * abs(D_h) / E_diff_max
    nu_max = c_max * h * max_value(abs(u), abs(u_old))
    nu_h = ufl.conditional(ufl.lt(nu_E, nu_max), nu_E, nu_max)
    return nu_h

def numerical_flux(u_p, u_m, n):
    C = max_value(abs(u_p), abs(u_m))
    flux = (Fc(u_p) + Fc(u_m)) / 2 + C / 2 * (u_p - u_m)
    return flux

def domain_flux(u):
    interior_flux = inner(numerical_flux(u("+"), u("-"), n("+")), phi("+") - phi("-")) * dS
    inlet_flux = numerical_flux(u, inlet_value, n) * phi * ds(1)
    outlet_flux = numerical_flux(u, u, n) * phi * ds(2)
    return interior_flux + inlet_flux + outlet_flux

def volume_term(u, u_old):
    nu_h = entropy_viscosity(u, u_old)
    return (inner(u - u_old, phi) - dt_ufl * inner(as_vector([Fc(u)]), grad(phi)) + dt_ufl * inner(nu_h * grad(u), grad(phi))) * dx(domain)

# Time-stepping
u = Function(V)
u_old = Function(V)
u_old.interpolate(initial_condition)

t = 0
while t < T - 0.5 * dt:
    t += float(dt)
    F = volume_term(u, u_old) + dt_ufl * domain_flux(u)
    J = ufl.derivative(F, u)
    problem = fem.petsc.NonlinearProblem(F, u, bcs=[bc_inlet])
    solver = dolfinx.nls.petsc.NewtonSolver(MPI.COMM_WORLD, problem)
    solver.solve(u)
    u.x.scatter_forward()
    u_old.x.array[:] = u.x.array[:].copy()
    print(f"t = {t:.3f}")

print("Simulation completed.")

The error I’m getting is:

ValueError: Component and shape length don’t match.

The full traceback is:

Traceback (most recent call last):
  File "/home/jy384/projects/UnimodalSROB/examples/burgers/burgersFEniCSx_Inviscid_TimeStep.py", line 180, in <module>
    problem = fem.petsc.NonlinearProblem(F, u, bcs=[bc_inlet])
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/fem/petsc.py", line 897, in __init__
    self._L = _create_form(
              ^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/fem/forms.py", line 249, in form
    return _create_form(form)
           ^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/fem/forms.py", line 244, in _create_form
    return _form(form)
           ^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/fem/forms.py", line 186, in _form
    ufcx_form, module, code = jit.ffcx_jit(
                              ^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/jit.py", line 51, in mpi_jit
    return local_jit(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/dolfinx/jit.py", line 201, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], options=p_ffcx, **p_jit)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/codegeneration/jit.py", line 276, in compile_forms
    raise e
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/codegeneration/jit.py", line 256, in compile_forms
    impl = _compile_objects(
           ^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/codegeneration/jit.py", line 383, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/compiler.py", line 108, in compile_ufl_objects
    analysis = analyze_ufl_objects(ufl_objects, options["scalar_type"])  # type: ignore
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/analysis.py", line 94, in analyze_ufl_objects
    form_data = tuple(_analyze_form(form, scalar_type) for form in forms)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/analysis.py", line 94, in <genexpr>
    form_data = tuple(_analyze_form(form, scalar_type) for form in forms)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ffcx/analysis.py", line 180, in _analyze_form
    form_data = ufl.algorithms.compute_form_data(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/compute_form_data.py", line 278, in compute_form_data
    form = preprocess_form(form, complex_mode)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/compute_form_data.py", line 232, in preprocess_form
    form = apply_algebra_lowering(form)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/apply_algebra_lowering.py", line 150, in apply_algebra_lowering
    return map_integrand_dags(LowerCompoundAlgebra(), expr)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/map_integrands.py", line 86, in map_integrand_dags
    return map_integrands(
           ^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/map_integrands.py", line 29, in map_integrands
    mapped_integrals = [
                       ^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/map_integrands.py", line 30, in <listcomp>
    map_integrands(function, itg, only_integral_type) for itg in form.integrals()
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/map_integrands.py", line 39, in map_integrands
    return itg.reconstruct(function(itg.integrand()))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/map_integrands.py", line 87, in <lambda>
    lambda expr: map_expr_dag(function, expr, compress), form, only_integral_type
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/corealg/map_dag.py", line 35, in map_expr_dag
    (result,) = map_expr_dags(
                ^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/corealg/map_dag.py", line 103, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/algorithms/apply_algebra_lowering.py", line 111, in div
    return a[..., i].dx(i)
           ~^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/exproperators.py", line 383, in _getitem
    all_indices, slice_indices, repeated_indices = create_slice_indices(
                                                   ^^^^^^^^^^^^^^^^^^^^^
  File "/home/jy384/miniconda3/envs/fenicsx/lib/python3.11/site-packages/ufl/index_combination_utils.py", line 167, in create_slice_indices
    raise ValueError("Component and shape length don't match.")
ValueError: Component and shape length don't match.

Has anyone encountered a similar issue or can provide guidance on how to properly implement the entropy viscosity in this context?

Thank you in advance for any help or insights!

Can you try simplifying your form to a point were you identify which term is causing the error? i.e., start dropping terms until you don’t get any ufl error, and then add back a single term that makes the error pop up again.

Of course any one of us could do the same exercise, but it’s more instructive if you try it yourself.

Your issue is in the ufl.div operator:

F(u) is a scalar value, thus the divergence doesn’t make sense.
You could either do F(u).dx(0) or

def entropy_residual(u, u_old):
    return (E(u) - E(u_old)) / dt + ufl.div(ufl.as_vector([F(u)]))

Please note that in your code you are re-defining F at a later stage, which is not advisable.
I would do something along the lines of

    F_ = volume_term(u, u_old) + dt_ufl * domain_flux(u)
    J = ufl.derivative(F_, u)
    problem = fem.petsc.NonlinearProblem(F_, u, bcs=[bc_inlet])

Note that you should also move the form and solver outside the time loop, as you only need to define them once when using re-assigning of data like:

    u.x.scatter_forward()
    u_old.x.array[:] = u.x.array[:].copy()