Assign local data to a vector-valued Function in FEniCSx (MPI) for time dependent problem

dolfinx
1

Hi all,

I’m trying to solve a time-dependent equation using FEniCSx with MPI. At each time step, I need to evaluate the right-hand side of the equation, which takes as input a vector-valued function vel_0. For example:

mesh = mesh.create_box(MPI.COMM_WORLD, [np.array([0, 0, 0]), np.array([1, 1, 1])], [10, 10, 10])
V = fem.functionspace(mesh, ("Lagrange", 1, (mesh.geometry.dim,)))

vel_0 = fem.Function(V)
vel_0.x.array[:] = 1.0  # initialize with constant vector value

# vector output function (RHS)
H_vel = fem.Function(V)
H_vel.x.array[:] = velocity_field(vel_0)

To use this in a time integrator, I need to extract the local part of vel_0.x.array (without ghosts). For the first time step, I do:

start, end = V.dofmap.index_map.local_range
vel_numeric = vel_0.x.array[start:end].copy()
H_vel_numeric = H_vel.x.array[start:end].copy()

My question arises for the next time steps, when uses vel_numeric and H_vel_numeric the time integrator will return a new vector vel_numeric, and I need to map it back into a dolfinx.Function to evaluate velocity_field again:

v_fem = fem.Function(V)
v_fem.x.array[start:end] = vel_numeric
v_fem.x.scatter_forward()

velocity_field(v_fem)

However, this leads to incorrect results when evaluating velocity_field(v_fem), even if vel_numeric is the same as the initial. I’m a little confused as to how to proceed.

Thanks

2

Without a reproducible code It is quite hard to see what is wrong with your code.

3

Hi dokken, sure here is a minimal example

import dolfinx
import ufl
import numpy as np
from mpi4py import MPI
from dolfinx import fem
from petsc4py import PETSc
from dolfinx.fem import Function, Constant, functionspace, form
from dolfinx.fem.petsc import assemble_matrix, assemble_vector
from dolfinx import mesh

mesh = mesh.create_box(MPI.COMM_WORLD, [np.array([0, 0, 0]), np.array([10, 10, 10])], [10, 10, 10])

V_scalar =  fem.functionspace(mesh, ("Lagrange", 1))
V_vec =     fem.functionspace(mesh, ("Lagrange", 1, (mesh.geometry.dim,)))

vel_0 = fem.Function(V_vec)


PHI = fem.Function(V_scalar )
v = ufl.TestFunction(V_scalar )
u = ufl.TrialFunction(V_scalar )

def RHS(vel_0):

    H_vel = fem.Function(V_vec)

    rhs = fem.form(ufl.inner(ufl.grad(v),vel_0 ) * ufl.dx)
    a = fem.form(ufl.inner(ufl.grad(v), ufl.grad(u)) * ufl.dx)
    A = fem.petsc.assemble_matrix(a)
    A.assemble()
    b = fem.petsc.create_vector(rhs)
    fem.petsc.assemble_vector(b, rhs)
    b.ghostUpdate(PETSc.InsertMode.ADD_VALUES, PETSc.ScatterMode.REVERSE)

    solver = PETSc.KSP().create(mesh.comm)
    solver.setOperators(A)
    solver.solve(b, PHI.x.petsc_vec)
    PHI.x.scatter_forward()
    # for simplicity we use interpolate
    Gradient= ufl.grad(PHI)
    expr = fem.Expression(Gradient, V_vec.element.interpolation_points())
    H_vel.interpolate(expr)
    return H_vel


# Initial condicion:

xyz = mesh.geometry.x
dim = mesh.geometry.x.shape[0]
vel0 = np.zeros((dim,3))
for i in range(dim):
    xx = xyz[i,0]
    yy = xyz[i,1]
    zz = xyz[i,2]

    norm = np.sqrt(xx*xx+yy*yy+0.2*0.2)

    vel0[i,0] = xx/norm
    vel0[i,1] = yy/norm
    vel0[i,2] = 0.5/norm



vel_0.x.array[:] = vel0.flatten()


# 1. The idea now is to use this function in some external time integrator like SciPy. For this, I need to use just the local part of the solution (without ghost), the first evaluation we have that:

start, end = V_vec.dofmap.index_map.local_range

rhs_initial = RHS(vel_0).x.array[start:end]
vel_0_initial = vel_0.x.array[start:end]


# 2. Now, we can suppose that obtain a new vector vel_0 from the integrator routines, we need to use this to evaluate the RHS function again. To do that, I think that is necesary to contruct a new vector function that contains also the local and the ghost part. For simplicity, suppose that 
vel_0_new = vel_0 (considered concatenate and flatten).

vel_0_new = fem.Function(V_vec)

vel_0_new.x.array[start:end] = vel_0.x.array[start:end]
vel_0_new.x.scatter_forward()

# then for the second evaluation,we should use

vel_0_second = vel_0_new.x.array[start:end]
rhs_second = RHS(vel_0_new).x.array[start:end]


print("initial", mesh.comm.rank, np.sort(rhs_initial)[0:10])
print("second step", mesh.comm.rank,  np.sort(rhs_second)[0:10])

# We can see that rhs_initial and rhs_second are differents.

Aditionally, I noticed that vel_0_new.x.array[start:end] is not a multiple of 3. Does this mean different ranks share different components of the vector associated with a node?. If you wanted to perform vector operations like inner or vector products, how could you do it?

thanks