Fokker-Planck solver arity error

Hello all. I am experimenting with a solution of the Fokker Planck equation of the following first-order form (arising from a noisy carrier-tracking problem in radio science):
\begin{eqnarray} \frac{\partial p}{\partial t} &=& p + \mathbf{a}\cdot \mathbf{J} - \nabla\cdot (\overline{\mathbf{B}}\cdot\mathbf{J}) \\ \mathbf{J} &=& -\nabla p ,\end{eqnarray}
where \mathbf{J} is a three component probability current, p is the joint probability function in x, y and z and
\mathbf{a} = \left\lbrace\begin{array}{c} (y - x + \sin z)\alpha \\ (x-y)\beta \\ Kx - \omega_0\end{array}\right\rbrace,
with
\overline{\mathbf{B}} = \left\lbrack\begin{array}{ccc} \alpha^2 N_0 & 0 & 0 \\ 0 & 0 & 0 \\ 0 & 0 & 0\end{array}\right \rbrack.

The solution domain \Omega is the rectangular parallelpiped volume (-0.5, -0.5, -\pi) to (0.5, 0.5, \pi). Neumann boundary conditions are enforced on x = -0.5, 0.5 and y=-0.5, 0.5 surfaces and an absorbing boundary p = 0 is assumed on z = -\pi, \pi surfaces. Time discretization is using the implicit Euler method:
p^n = p^{n+1} -\Delta t (p^{n+1} - \mathbf{a}\cdot \mathbf{J}^{n+1} + \nabla\cdot (\overline{\mathbf{B}}\cdot\mathbf{J}^{n+1})).

For solving this, I am using the least-squares finite element approach (similar to what I did here). The residual is constructed as follows:
R = \vert\!\vert p^{n+1}-p^n-\Delta t (p^{n+1} - \mathbf{a}\cdot\mathbf{J}^{n+1}+\nabla\cdot(\overline{\mathbf{B}}\cdot\mathbf{J}^{n+1}))\vert\!\vert^2_\Omega +
\qquad\qquad\vert\!\vert\mathbf{J}^{n+1}+\nabla p^{n+1}\vert\!\vert^2_\Omega + \vert\!\vert\mathbf{n\cdot J}^{n+1}\vert\!\vert^2_{\partial\Omega_N} + \vert\!\vert p^{n+1}\vert\!\vert^2_{\partial\Omega_A}.
Minimization of the residual is achieved using the usual Ritz process on the unknown variables (p^{n+1}, \mathbf{J}^{n+1}). The initial conditions are set in p^n.

I am doing something incorrectly, most likely, in setting up the initial conditions or the way I do the multiplication of the coefficients. The code gets as far as the assembly process and dies as the result of an “arity” error. I’ve had a look at other posts that describe similar errors, but have had no luck in making my code work. The error output is:

Generating mesh with CGAL 3D mesh generator
<Mesh of topological dimension 3 (tetrahedra) with 7421 vertices and 37627 cells, ordered>
Solution vector size =  221284
Calling FFC just-in-time (JIT) compiler, this may take some time.
Traceback (most recent call last):
  File "FP1.py", line 103, in <module>
    A, b = assemble_system(a, L, bcs)
  File "/usr/lib/python3/dist-packages/dolfin/fem/assembling.py", line 360, in assemble_system
    A_dolfin_form = _create_dolfin_form(A_form, form_compiler_parameters)
  File "/usr/lib/python3/dist-packages/dolfin/fem/assembling.py", line 58, in _create_dolfin_form
    function_spaces=function_spaces)
  File "/usr/lib/python3/dist-packages/dolfin/fem/form.py", line 44, in __init__
    mpi_comm=mesh.mpi_comm())
  File "/usr/lib/python3/dist-packages/dolfin/jit/jit.py", line 47, in mpi_jit
    return local_jit(*args, **kwargs)
  File "/usr/lib/python3/dist-packages/dolfin/jit/jit.py", line 97, in ffc_jit
    return ffc.jit(ufl_form, parameters=p)
  File "/usr/lib/python3/dist-packages/ffc/jitcompiler.py", line 217, in jit
    module = jit_build(ufl_object, module_name, parameters)
  File "/usr/lib/python3/dist-packages/ffc/jitcompiler.py", line 133, in jit_build
    generate=jit_generate)
  File "/usr/lib/python3/dist-packages/dijitso/jit.py", line 165, in jit
    header, source, dependencies = generate(jitable, name, signature, params["generator"])
  File "/usr/lib/python3/dist-packages/ffc/jitcompiler.py", line 66, in jit_generate
    prefix=module_name, parameters=parameters, jit=True)
  File "/usr/lib/python3/dist-packages/ffc/compiler.py", line 143, in compile_form
    prefix, parameters, jit)
  File "/usr/lib/python3/dist-packages/ffc/compiler.py", line 185, in compile_ufl_objects
    analysis = analyze_ufl_objects(ufl_objects, kind, parameters)
  File "/usr/lib/python3/dist-packages/ffc/analysis.py", line 90, in analyze_ufl_objects
    for form in forms)
  File "/usr/lib/python3/dist-packages/ffc/analysis.py", line 90, in <genexpr>
    for form in forms)
  File "/usr/lib/python3/dist-packages/ffc/analysis.py", line 174, in _analyze_form
    do_apply_restrictions=True)
  File "/usr/lib/python3/dist-packages/ufl/algorithms/compute_form_data.py", line 418, in compute_form_data
    check_form_arity(preprocessed_form, self.original_form.arguments(), complex_mode)  # Currently testing how fast this is
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 177, in check_form_arity
    check_integrand_arity(itg.integrand(), arguments, complex_mode)
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 159, in check_integrand_arity
    arg_tuples = map_expr_dag(rules, expr, compress=False)
  File "/usr/lib/python3/dist-packages/ufl/corealg/map_dag.py", line 37, in map_expr_dag
    result, = map_expr_dags(function, [expression], compress=compress)
  File "/usr/lib/python3/dist-packages/ufl/corealg/map_dag.py", line 86, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 48, in sum
    raise ArityMismatch("Adding expressions with non-matching form arguments {0} vs {1}.".format(_afmt(a), _afmt(b)))
ufl.algorithms.check_arities.ArityMismatch: Adding expressions with non-matching form arguments () vs ('v_0',).

My model code is:

from dolfin import *
from mshr import *
import numpy as np
import matplotlib.pyplot as plt
import os, sys, traceback

pi = 3.1415926536
tol = 1.0e-12
# Problem parameters
alfa = 1.0
beta = 0.1
N0 = 0.1
K = 1.0
omega = 1.0
sig = 0.3

# time stepping params
dt = 0.01
NumSteps = 100


class Wall(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary

class AbsorbingBC(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and (near(x[2], -pi, tol) or near(x[2], pi, tol))


box = Box(dolfin.Point(-0.5,-0.5, -pi), dolfin.Point(0.5, 0.5, pi))

mesh = generate_mesh(box, 70)

info(mesh)

# Mark boundaries
sub_domains = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
sub_domains.set_all(2)
wall = Wall()
wall.mark(sub_domains, 0)
absorb = AbsorbingBC()
absorb.mark(sub_domains, 1)
File("BoxSubDomains.pvd").write(sub_domains)

# PDE Coefficients , C1 is vector, C2 is 3x3 matrix
C1 = Expression(('(x[1] - x[0] + sin(x[3])) * alfa', '(x[0] - x[1]) * beta', 'K * x[0] - omega'), degree = 3, alfa = alfa, beta = beta, K = K, omega = omega)
C2 = Constant(((N0, 0.0, 0.0), (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)))

# Set up function spaces
cell = tetrahedron
ele = FiniteElement('Lagrange', cell, 2) 
vec_ele = VectorElement('Lagrange', cell, 2) 
V4 = FunctionSpace(mesh, MixedElement([ele, vec_ele]))
V = FunctionSpace(mesh, ele)
u1 = Function(V4)
un = Function(V4)
pv, Jv = TestFunctions(V4)

#surface integral definitions from boundaries
ds = Measure('ds', domain = mesh, subdomain_data = sub_domains)
# set initial conditions
g = Expression(('exp(-(x[0]*x[0]+x[1]*x[1]+x[2]*x[2])/(4.0 * pi * pi)) / (2.0 * sig * sig)', '0.0', '0.0', '0.0'), degree = 3, sig = sig)
pp0 = interpolate(g, V4)
un.assign(pp0)

#Boundary condition dictionary
boundary_conditions = {0: {'Wall' : 0.0},
                       1: {'AbsorbingBC': 0.0}}

n = FacetNormal(mesh)

pr, Jr = split(u1) #Split for forms
pn, Jn = split(un)
#Build boundary conditions
bcs = []
integrals_load =[]
#Neumann BC on impermeable walls
for i in boundary_conditions:
    if 'Wall' in boundary_conditions[i]:
        bb1 = inner(dot(n, Jv), dot(n, Jr)) * ds(i)
        integrals_load.append(bb1)

# Build absorbing BC
for i in boundary_conditions:
    if 'AbsorbingBC' in boundary_conditions[i]:
        bb = inner(pv, pr) * ds(i)
        integrals_load.append(bb)
a = (inner(pv - dt * (pv + dot(C1, Jv) - div(C2 * Jv)), pr - dt * (pr + dot(C1, Jr) - div(C2 * Jv))) + inner(Jv + grad(pv), Jr + grad(pr))) * dx + sum(integrals_load)
L = inner(pv - dt * (pv + dot(C1, Jv) - div(C2 * Jv)), pn) * dx

vdim = u1.vector().size()
print ('Solution vector size = ',vdim)
u1.vector().zero()
A, b = assemble_system(a, L, bcs)  #XXXXX Seems to die here XXXX
solver = PETScKrylovSolver("cg","sor")
solver.parameters['absolute_tolerance'] = 1e-5
solver.parameters['relative_tolerance'] = 1e-12
solver.parameters['maximum_iterations'] = 10000
solver.parameters['monitor_convergence'] = True
solver.parameters['error_on_nonconvergence'] = False
solver.parameters['nonzero_initial_guess'] = True
solver.parameters['report'] = True
solver.set_operators(A, A)

fpp = File('Prob.pvd')
fpJ = File('Curr.pvd')
p, J = u1.split()

# Do time stepping
t = 0
for nn in range(NumSteps):
    t += dt
    #Solve
    solver.solve(u1.vector(),b)
    #Update solution
    un.assign(u1)
    fpp << (p, t)
    fpJ << (J, t)

sys.exit(0)

Any hints as to where I am going wrong would be very much appreciated! I am using Fenics 2019 on an Ubuntu 18 PC where Fenics was installed from PPA. This version of Fenics has worked well with other problems, so the installation seems OK.

An approach that should be immune to arity errors (or other pencil-and-paper mistakes) would be to specify the least squares functional directly, as an arity-0 form (i.e., where the solution is just a Function), then obtain the linearized problem via derivative, rather than deriving it on paper. The corresponding linear problem could then be obtained like

lsqf = # least-squares functional, depending on Function u1
res = derivative(lsqf,u1) # linear form with one TestFunction argument
a = derivative(res,u1) # bilinear form with TestFunction and TrialFunction
L = -res

This is effectively one step of Newton iteration, which exactly solves a linear problem.

Thanks for the tip @kamensky. The program got further along than the initial try, but I tried your suggestion and got an even more cryptic error.

Generating mesh with CGAL 3D mesh generator
<Mesh of topological dimension 3 (tetrahedra) with 7421 vertices and 37627 cells, ordered>
Solution vector size =  221284
Calling FFC just-in-time (JIT) compiler, this may take some time.
Calling FFC just-in-time (JIT) compiler, this may take some time.
python3: /usr/include/eigen3/Eigen/src/Core/DenseCoeffsBase.h:162: Eigen::DenseCoeffsBase<Derived, 0>::CoeffReturnType Eigen::DenseCoeffsBase<Derived, 0>::operator[](Eigen::Index) const [with Derived = Eigen::Ref<const Eigen::Matrix<double, -1, 1> >; Eigen::DenseCoeffsBase<Derived, 0>::CoeffReturnType = double; Eigen::Index = long int]: Assertion `index >= 0 && index < size()' failed.
[gondor:13328] *** Process received signal ***
[gondor:13328] Signal: Aborted (6)
[gondor:13328] Signal code:  (-6)
[gondor:13328] [ 0] /lib/x86_64-linux-gnu/libc.so.6(+0x3ef20)[0x7f16027eff20]
[gondor:13328] [ 1] /lib/x86_64-linux-gnu/libc.so.6(gsignal+0xc7)[0x7f16027efe97]
[gondor:13328] [ 2] /lib/x86_64-linux-gnu/libc.so.6(abort+0x141)[0x7f16027f1801]
[gondor:13328] [ 3] /lib/x86_64-linux-gnu/libc.so.6(+0x3039a)[0x7f16027e139a]
[gondor:13328] [ 4] /lib/x86_64-linux-gnu/libc.so.6(+0x30412)[0x7f16027e1412]
[gondor:13328] [ 5] /home/bill/.cache/dijitso/lib/libdijitso-dolfin_expression_49cc6209e2b861c629affe0678f81f51.so(_ZNSt6vectorImSaImEE17_M_realloc_insertIJmEEEvN9__gnu_cxx17__normal_iteratorIPmS1_EEDpOT_+0x0)[0x7f15ca3c3080]
[gondor:13328] [ 6] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZNK6dolfin10Expression4evalERNS_5ArrayIdEERKS2_+0x59)[0x7f1600681e59]
[gondor:13328] [ 7] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZNK6dolfin10Expression4evalEN5Eigen3RefINS1_6MatrixIdLin1ELi1ELi0ELin1ELi1EEELi0ENS1_11InnerStrideILi1EEEEENS2_IKS4_Li0ES6_EERKN3ufc4cellE+0x4e)[0x7f1600681eee]
[gondor:13328] [ 8] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZNK6dolfin10Expression4evalERNS_5ArrayIdEERKS2_RKN3ufc4cellE+0x59)[0x7f1600681db9]
[gondor:13328] [ 9] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZNK6dolfin15GenericFunction8evaluateEPdPKdRKN3ufc4cellE+0x60)[0x7f160069a120]
[gondor:13328] [10] /home/bill/.cache/dijitso/lib/libdijitso-ffc_element_a1e6f37c8b37556938e54a0536f64f40ca895aec.so(_ZNK72ffc_element_a1e6f37c8b37556938e54a0536f64f40ca895aec_finite_element_main13evaluate_dofsEPdRKN3ufc8functionEPKdiRKNS1_4cellEPKNS1_18coordinate_mappingE+0x69)[0x7f15c8190fa9]
[gondor:13328] [11] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZNK6dolfin15GenericFunction24restrict_as_ufc_functionEPdRKNS_13FiniteElementERKNS_4CellEPKdRKN3ufc4cellE+0x21)[0x7f160069a1d1]
[gondor:13328] [12] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZN6dolfin3UFC6updateERKNS_4CellERKSt6vectorIdSaIdEERKN3ufc4cellERKS4_IbSaIbEE+0x93)[0x7f1600679b63]
[gondor:13328] [13] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZN6dolfin15SystemAssembler18cell_wise_assemblyERSt5arrayIPNS_13GenericTensorELm2EERS1_IPNS_3UFCELm2EERNS0_7ScratchERKSt6vectorISt13unordered_mapImdSt4hashImESt8equal_toImESaISt4pairIKmdEEESaISM_EESt10shared_ptrIKNS_12MeshFunctionImEEESV_+0xa87)[0x7f1600672bd7]
[gondor:13328] [14] /usr/lib/x86_64-linux-gnu/libdolfin.so.2019.1(_ZN6dolfin15SystemAssembler8assembleEPNS_13GenericMatrixEPNS_13GenericVectorEPKS3_+0x1c4f)[0x7f160067839f]
[gondor:13328] [15] /usr/lib/python3/dist-packages/dolfin/cpp.cpython-36m-x86_64-linux-gnu.so(+0x70e4e)[0x7f1600c9be4e]
[gondor:13328] [16] /usr/lib/python3/dist-packages/dolfin/cpp.cpython-36m-x86_64-linux-gnu.so(+0x56b42)[0x7f1600c81b42]
[gondor:13328] [17] python3(_PyCFunction_FastCallDict+0x35c)[0x5674fc]
[gondor:13328] [18] python3[0x50abb3]
[gondor:13328] [19] python3(_PyEval_EvalFrameDefault+0x449)[0x50c5b9]
[gondor:13328] [20] python3[0x508245]
[gondor:13328] [21] python3[0x50a080]
[gondor:13328] [22] python3[0x50aa7d]
[gondor:13328] [23] python3(_PyEval_EvalFrameDefault+0x449)[0x50c5b9]
[gondor:13328] [24] python3[0x508245]
[gondor:13328] [25] python3(PyEval_EvalCode+0x23)[0x50b403]
[gondor:13328] [26] python3[0x635222]
[gondor:13328] [27] python3(PyRun_FileExFlags+0x97)[0x6352d7]
[gondor:13328] [28] python3(PyRun_SimpleFileExFlags+0x17f)[0x638a8f]
[gondor:13328] [29] python3(Py_Main+0x591)[0x639631]
[gondor:13328] *** End of error message ***
Aborted (core dumped)

Before this, I had to specify explicitly the quadrature order because I got a warning

Calling FFC just-in-time (JIT) compiler, this may take some time.
  WARNING: The number of integration points for each cell will be: 216
           Consider using the option 'quadrature_degree' to reduce the number of points
Calling FFC just-in-time (JIT) compiler, this may take some time.
Calling FFC just-in-time (JIT) compiler, this may take some time.
Calling FFC just-in-time (JIT) compiler, this may take some time.
  WARNING: The number of integration points for each cell will be: 216
           Consider using the option 'quadrature_degree' to reduce the number of points

before crapping out with a core dump. I added the following lines:

metadata = {"quadrature_degree": 4, "quadrature_scheme": "default"}

and

dxm = dx(metadata=metadata)

The full code is below:

from dolfin import *
from mshr import *
import numpy as np
import matplotlib.pyplot as plt
import os, sys, traceback

pi = 3.1415926536
tol = 1.0e-12
# Problem parameters
alfa = 1.0
beta = 0.1
N0 = 0.1
K = 1.0
omega = 1.0
sig = 0.3

# time stepping params
dt = 0.01
NumSteps = 100


class Wall(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary

class AbsorbingBC(SubDomain):
    def inside(self, x, on_boundary):
        return on_boundary and (near(x[2], -pi, tol) or near(x[2], pi, tol))


box = Box(dolfin.Point(-0.5,-0.5, -pi), dolfin.Point(0.5, 0.5, pi))

mesh = generate_mesh(box, 70)

info(mesh)

# Output mesh for inspection
File("BoxMesh.pvd").write(mesh)

# Mark boundaries
sub_domains = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
sub_domains.set_all(2)
wall = Wall()
wall.mark(sub_domains, 0)
absorb = AbsorbingBC()
absorb.mark(sub_domains, 1)
File("BoxSubDomains.pvd").write(sub_domains)

# PDE Coefficients 
C1 = Expression(('(x[1] - x[0] + sin(x[3])) * alfa', '(x[0] - x[1]) * beta', 'K * x[0] - omega'), degree = 3, alfa = alfa, beta = beta, K = K, omega = omega)
C2 = Constant(((N0, 0.0, 0.0), (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)))

# Set up function spaces
cell = tetrahedron
ele = FiniteElement('Lagrange', cell, 2) 
vec_ele = VectorElement('Lagrange', cell, 2) 
V4 = FunctionSpace(mesh, MixedElement([ele, vec_ele]))
V = FunctionSpace(mesh, ele)
u1 = Function(V4)
un = Function(V4)
metadata = {"quadrature_degree": 4, "quadrature_scheme": "default"}
#surface integral definitions from boundaries
ds = Measure('ds', domain = mesh, subdomain_data = sub_domains)
# dx
dxm = dx(metadata=metadata)
# initial conditions
g = Expression(('exp(-(x[0]*x[0]+x[1]*x[1]+x[2]*x[2])/(4.0 * pi * pi)) / (2.0 * sig * sig)', '0.0', '0.0', '0.0'), degree = 3, sig = sig)
pp0 = interpolate(g, V4)
un.assign(pp0)

#Boundary condition dictionary
boundary_conditions = {0: {'Wall' : 0.0},
                       1: {'AbsorbingBC': 0.0}}

n = FacetNormal(mesh)

pr, Jr = split(u1) #Split for forms
pn, Jn = split(un)
#Build boundary conditions
bcs = []
integrals_load =[]
#Neumann BC on impermeable walls
for i in boundary_conditions:
    if 'Wall' in boundary_conditions[i]:
        bb1 = inner(dot(n, Jr), dot(n, Jr)) * ds(i)

# Build absorbing BC
for i in boundary_conditions:
    if 'AbsorbingBC' in boundary_conditions[i]:
        bb2 = inner(pr, pr) * ds(i)

R = (inner(pr - pn - dt * (pr - dot(C1, Jr) + div(C2 * Jr)), pr - pn - dt * (pr - dot(C1, Jr) + div(C2 * Jr))) + inner(Jr + grad(pr), Jr + grad(pr))) * dxm + bb1 + bb2
res = derivative(R, u1)
a = derivative(res, u1)
L = -res
vdim = u1.vector().size()
print ('Solution vector size = ',vdim)
u1.vector().zero()
A, b = assemble_system(a, L, bcs)
solver = PETScKrylovSolver("cg","sor")
solver.parameters['absolute_tolerance'] = 1e-5
solver.parameters['relative_tolerance'] = 1e-12
solver.parameters['maximum_iterations'] = 10000
solver.parameters['monitor_convergence'] = True
solver.parameters['error_on_nonconvergence'] = False
solver.parameters['nonzero_initial_guess'] = True
solver.parameters['report'] = True
solver.set_operators(A, A)

fpp = File('Prob.pvd')
fpJ = File('Curr.pvd')
p, J = u1.split()

# Do time stepping
t = 0
for nn in range(NumSteps):
    t += dt
    #Solve
    solver.solve(u1.vector(),b)
    #Update solution
    un.assign(u1)
    fpp << (p, t)
    fpJ << (J, t)

sys.exit(0)

Did you perhaps mean sin(x[2]) instead of sin(x[3]) in the definition of C1? Getting rid of the invalid x[3] there avoids this error.

@kamensky indeed that works. After fixing a couple of other small errors, the time stepping works as well. Now I need to verify my formulation.

Thanks for the help! I never would have guessed the trick using the “derivative” operator.