How can I speed up my integrals? ufl.conditional() needed for the integrals to work in the way that I want them to seems to be causing them to take a large amount of time

dolfinx
1

For an iterative procedure I’m using, I need to perform cumulative integrals on the boundaries several times, and ufl.conditional() allows me to perform cumulative integration using Dolfinx and ufl. I noticed that my boundary integrals were very slow, and I think I have narrowed down that what makes them exceptionally slow to the ufl.conditional(). This is needed to simulate changing the bounds of integration by checking for points beyond the current bound of integration, and multiplying by 0. Unless there is a way to drastically speed up the integral, perhaps by finding a way to do this without using ufl.conditional() or finding a way to use it more efficiently, then while the results are extremely accurate, obtaining them becomes rather impractical.

The following code took 18 seconds to finish the interpolation of the cumulative integration:

import dolfinx as df
from petsc4py import PETSc
import ufl
import numpy as np
from mpi4py.MPI import COMM_WORLD as comm
import matplotlib.pyplot as plt

Nx,Ny = 81,61
mesh = df.mesh.create_rectangle(comm,[[0.0,0.0],[np.pi, np.pi]],[Nx,Ny])
tdim = mesh.topology.dim
mesh.topology.create_connectivity(tdim-1,tdim)
subset_marker = 1
family = "Lagrange"
order = 1

# Scalar function space to interpolate scalar result to
S = df.fem.functionspace(mesh, (family,order))

lower_facets = df.mesh.locate_entities_boundary(mesh, mesh.topology.dim-1, lambda x: np.isclose(x[1], 0.0))
V_n = df.fem.Function(S)
V_n.interpolate(lambda x: -1-np.sin(x[0]))

# Tag lower boundary
subset_facet_tags = df.mesh.meshtags(mesh, tdim-1, lower_facets, 
                                     np.full_like(lower_facets, subset_marker))
boundary_submesh = df.mesh.create_submesh(mesh, tdim-1, lower_facets)[0]
ds_subset = ufl.Measure("ds",mesh,subdomain_data=subset_facet_tags) # ds over lower boundary
# Scalar space over lower boundary
S_sub_boundary = df.fem.functionspace(boundary_submesh, (family, order)) 
# cumulative integral defined over southern boundary
cumul_int = df.fem.Function(S_sub_boundary)


x = ufl.SpatialCoordinate(mesh)
############################################################################
# Slow part ################################################################
############################################################################
# cumulative integral
cumul_int_expr = lambda x0: np.hstack([df.fem.assemble_scalar(df.fem.form(
                    V_n * ufl.conditional(x[0]<x_, 1, 0) * ds_subset(subset_marker))) 
                                                      for x_ in x0[0]])
cumul_int.interpolate(cumul_int_expr)

# Plotting
lower_dofs = df.fem.locate_dofs_topological(S,tdim-1,lower_facets)
V_n_S = V_n.x.array[lower_dofs]

boundary_submesh = df.mesh.create_submesh(mesh, mesh.topology.dim-1, lower_facets)[0] # submesh over southern boundary
S_sub_boundary = df.fem.functionspace(boundary_submesh, (family, order)) 
cumul_int_ex = df.fem.Function(S_sub_boundary)
cumul_int_ex.interpolate(lambda x: -np.cos(x[0])+x[0]+1)

lower_x_axis = S.tabulate_dof_coordinates()[df.fem.locate_dofs_topological(S,tdim-1,lower_facets)][:,0]
sort_idx = np.argsort(lower_x_axis)

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(lower_x_axis[sort_idx], V_n_S[sort_idx])
ax.plot(lower_x_axis[sort_idx], cumul_int_ex.x.array[sort_idx], marker = 'o')
ax.plot(lower_x_axis[sort_idx], cumul_int.x.array[sort_idx])
ax.set_xlabel('$x$')
ax.legend(['$V_{n}|_{y=0}$','$int_{cumul_{exact}}$','$int_{cumul_{approx}}$'])
plt.show()

Out of curiosity, I played around with the code a bit, and I managed to recreated only the values of the interpolated expression below without ufl.conditional() inside of ufl expression in dolfinx.fem.form(), and it finished in around 2 seconds. Here is the code I used to do so:

lower_dofs = df.fem.locate_dofs_topological(S,tdim-1,lower_facets)
lower_x_axis = S.tabulate_dof_coordinates()[df.fem.locate_dofs_topological(S,tdim-1,lower_facets)][:,0]
sort_idx = np.argsort(lower_x_axis)

cumul_int2 = []
deltax = lower_x_axis[sort_idx][1] - lower_x_axis[sort_idx][0]
for i in range(1,len(lower_x_axis[sort_idx])+1):
    facets = df.mesh.locate_entities_boundary(mesh, mesh.topology.dim-1, lambda x: (np.isclose(x[1], 0)) & (x[0]<i*deltax))
    facet_tags = df.mesh.meshtags(mesh, tdim-1, facets, 
                                     np.full_like(facets, subset_marker))
    boundary_submesh = df.mesh.create_submesh(mesh, tdim-1, facets)[0]
    ds = ufl.Measure("ds",mesh,subdomain_data=facet_tags) # ds over facets
    cumul_int2.append(df.fem.assemble_scalar(df.fem.form(
        V_n * ds(subset_marker)
    )))
cumul_int2 = np.hstack(cumul_int2)

# Plotting
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(lower_x_axis[sort_idx], cumul_int2, marker = 'o')
ax.plot(lower_x_axis[sort_idx], cumul_int_ex.x.array[sort_idx], linestyle = '--',marker = '+')
ax.plot(lower_x_axis[sort_idx], cumul_int.x.array[sort_idx])
ax.set_xlabel('$x$')
ax.legend(['$int_{cumul_{fast}}$','$int_{cumul_{exact}}$','$int_{cumul_{approx}}$'])
plt.show()

What this code tells me, is that the integrals over the desired range of x-coordinates can be quickly computed, in order.

NOTE: I do realize that I could take my quickly produced array of values and have scipy quickly interpolate the array, and then interpolate that onto a variable defined onto a Dolfinx functionspace, but I’ve had some unsatisfactory experiences with its interpolation of data if it isn’t smooth, so I’d prefer if I could keep interpolation to using only Dolfinx if possible.

2

The code you make is slow because:

You are using a scalar value rather than a dolfinx.fem.Constant within your dolfinx.fem.form, meaning that the code is recompiled for each instance of x_.

Thus, one thing you will observe is that if you run the same code twice, it speeds up drastically
First time

Slow cumulative integral time: 16.8442 seconds

Second time

Slow cumulative integral time: 0.1674 seconds

However, you can make it even faster by moving form compilation out of the interpolation, and use constants (and make it MPI compatible)

import time

x_c = df.fem.Constant(mesh, PETSc.ScalarType(0.0))
compiled_form = df.fem.form(
    V_n * ufl.conditional(x[0] < x_c, 1, 0) * ds_subset(subset_marker)
)


def cumul_int_expr(x):
    out = np.zeros(x.shape[1])
    for i in range(x.shape[1]):
        x_c.value = x[0][i]
        local_cont = df.fem.assemble_scalar(compiled_form)
        cont = mesh.comm.allreduce(local_cont, op=MPI.SUM)
        out[i] = cont
    return out


start = time.perf_counter()
cumul_int.interpolate(cumul_int_expr)
end = time.perf_counter()
print(f"Faster cumulative integral time: {end - start:.4f} seconds")


start = time.perf_counter()
cumul_int_expr = lambda x0: np.hstack(
    [
        df.fem.assemble_scalar(
            df.fem.form(
                V_n * ufl.conditional(x[0] < x_, 1, 0) * ds_subset(subset_marker)
            )
        )
        for x_ in x0[0]
    ]
)
cumul_int2 = df.fem.Function(S_sub_boundary)
cumul_int2.interpolate(cumul_int_expr)
end = time.perf_counter()

print(f"Slow cumulative integral time: {end - start:.4f} seconds")
np.testing.assert_allclose(cumul_int2.x.array, cumul_int.x.array)
2 Likes
3

Perfect! I didn’t know that scalar values slowed down the code in Dolfinx like that or to that degree; I will have to make a force of habit to avoid using them from now on. Thank you for pointing this out!

And yes, I did notice that in subsequent runs, the code was much much faster. I’m guessing this is owed to caching the comparisons. If I were to change the number of points on the grid, or increase the order of the scalar space, the comparisons have to be made from scratch again and it slows down drastically. Given this, the way the code had been set up might not have been as impractical as I had first thought, but the implementation you’ve provided is definitely preferred.

Thanks again!