How to scale vector element for specific subdomain

General
1

Hi there,

I am trying to scale my vector field by some constant. But, I cannot get the correct dofs of the subdomain. Here is the MWE;

import gmsh
import numpy as np
import pyvista
from dolfinx.fem import (Constant, dirichletbc, Function, FunctionSpace, assemble_scalar,
                         form, locate_dofs_geometrical, locate_dofs_topological)
from dolfinx.io import XDMFFile
from dolfinx.mesh import create_unit_square, locate_entities

from ufl import VectorElement

from mpi4py import MPI
from petsc4py.PETSc import ScalarType
# pyvista.start_xvfb()

mesh = create_unit_square(MPI.COMM_WORLD, 10, 10)

# Q = FunctionSpace(mesh, ("DG", 0))

v_cg = VectorElement("CG", mesh.ufl_cell(), 1)
Q = FunctionSpace(mesh, v_cg)

def Omega_0(x):
    return x[1] <= 0.3

def Omega_1(x):
    return x[1] >= 0.7

kappa = Function(Q)
cells_0 = locate_entities(mesh, mesh.topology.dim, Omega_0)
cells_1 = locate_entities(mesh, mesh.topology.dim, Omega_1)

dofs_0 = locate_dofs_geometrical(Q, Omega_0)
dofs_1 = locate_dofs_geometrical(Q, Omega_1)

kappa.x.array[dofs_0] = 1
kappa.x.array[dofs_1] = 0.1


with XDMFFile(MPI.COMM_WORLD, "test.xdmf", "w", encoding=XDMFFile.Encoding.HDF5 ) as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(kappa)

I import subdomains from gmsh and I can access the cell tags by;

cell_tags = subdomains.find(tag)

How can I access the dofs of the cell_tags?

I get this;
Screenshot from 2023-03-16 19-05-42

whereas I should get this;
Screenshot from 2023-03-16 19-06-50

Any help would be much appreciated!

2

One way to do it is by unrolling the dofmap (as you are using a vector function space):

import numba
import gmsh
import numpy as np
from dolfinx.fem import (Function, FunctionSpace,
                         locate_dofs_geometrical)
from dolfinx.io import XDMFFile
from dolfinx.mesh import create_unit_square, locate_entities

from ufl import VectorElement

from mpi4py import MPI

# pyvista.start_xvfb()

mesh = create_unit_square(MPI.COMM_WORLD, 10, 10)

# Q = FunctionSpace(mesh, ("DG", 0))

v_cg = VectorElement("CG", mesh.ufl_cell(), 1)
Q = FunctionSpace(mesh, v_cg)


def Omega_0(x):
    return x[1] <= 0.3


def Omega_1(x):
    return x[1] >= 0.7


kappa = Function(Q)
cells_0 = locate_entities(mesh, mesh.topology.dim, Omega_0)
cells_1 = locate_entities(mesh, mesh.topology.dim, Omega_1)


@numba.njit
def unroll_dofmap(dofs, bs):
    dofs_unrolled = np.zeros(bs*len(dofs), dtype=np.int32)
    for i, dof in enumerate(dofs):
        for b in range(bs):
            dofs_unrolled[i*bs+b] = dof*bs+b
    return dofs_unrolled


dofs_0 = locate_dofs_geometrical(Q, Omega_0)
dofs_1 = locate_dofs_geometrical(Q, Omega_1)
dofmap_bs = Q.dofmap.bs
kappa.x.array[unroll_dofmap(dofs_0, dofmap_bs)] = 1
kappa.x.array[unroll_dofmap(dofs_1, dofmap_bs)] = 0.1


with XDMFFile(MPI.COMM_WORLD, "test.xdmf", "w", encoding=XDMFFile.Encoding.HDF5) as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(kappa)

There is another way that should be better (I will post it later when I have time to jot it down).

3

Here is an example using only topological information, and illustrating the speedup compared to using locate_dofs_geometrical:

import time

import numba
import numpy as np
from dolfinx.fem import (Function, FunctionSpace, locate_dofs_geometrical,
                         locate_dofs_topological)
from dolfinx.io import XDMFFile
from dolfinx.mesh import create_unit_square, locate_entities
from mpi4py import MPI
from ufl import VectorElement

N = 1000
mesh = create_unit_square(MPI.COMM_WORLD, N, N)

# Q = FunctionSpace(mesh, ("DG", 0))

v_cg = VectorElement("CG", mesh.ufl_cell(), 1)
Q = FunctionSpace(mesh, v_cg)


def Omega_0(x, tol=1e-13):
    return x[1] <= 0.3 + tol


def Omega_1(x, tol=1e-13):
    return x[1] >= 0.7 - tol


kappa = Function(Q)
cells_0 = locate_entities(mesh, mesh.topology.dim, Omega_0)
cells_1 = locate_entities(mesh, mesh.topology.dim, Omega_1)


@numba.njit
def unroll_dofmap(dofs, bs):
    dofs_unrolled = np.zeros(bs*len(dofs), dtype=np.int32)
    for i, dof in enumerate(dofs):
        for b in range(bs):
            dofs_unrolled[i*bs+b] = dof*bs+b

    return dofs_unrolled


dofmap_bs = Q.dofmap.bs
# cold start to compile code
unroll_dofmap(np.array([0], dtype=np.int32), dofmap_bs)
MPI.COMM_WORLD.Barrier()
start = time.perf_counter()
dofs_0 = locate_dofs_topological(Q, mesh.topology.dim, cells_0)
dofs_1 = locate_dofs_topological(Q, mesh.topology.dim, cells_1)
kappa.x.array[unroll_dofmap(dofs_0, dofmap_bs)] = 1
kappa.x.array[unroll_dofmap(dofs_1, dofmap_bs)] = 0.1
kappa.x.scatter_forward()
end = time.perf_counter()
total_time = MPI.COMM_WORLD.gather(end-start, root=0)
if MPI.COMM_WORLD.rank == 0:
    print(
        f"Topological: {min(total_time)=:.2e} {max(total_time)=:.2e} avg(total_time)={sum(total_time)/len(total_time):.2e}")

kappa2 = Function(Q)
MPI.COMM_WORLD.Barrier()
start = time.perf_counter()
dofs_0 = locate_dofs_geometrical(Q, Omega_0)
dofs_1 = locate_dofs_geometrical(Q, Omega_1)
kappa2.x.array[unroll_dofmap(dofs_0, dofmap_bs)] = 1
kappa2.x.array[unroll_dofmap(dofs_1, dofmap_bs)] = 0.1
kappa2.x.scatter_forward()
end = time.perf_counter()
total_time = MPI.COMM_WORLD.gather(end-start, root=0)
if MPI.COMM_WORLD.rank == 0:
    print(
        f"Geometrical: {min(total_time)=:.2e} {max(total_time)=:.2e} avg(total_time)={sum(total_time)/len(total_time):.2e}")

assert np.allclose(kappa2.x.array, kappa.x.array)


with XDMFFile(MPI.COMM_WORLD, "test.xdmf", "w", encoding=XDMFFile.Encoding.HDF5) as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(kappa)

which gives

Serial:
Topological: min(total_time)=1.24e-01 max(total_time)=1.24e-01 avg(total_time)=1.24e-01
Geometrical: min(total_time)=2.15e-01 max(total_time)=2.15e-01 avg(total_time)=2.15e-01
2 proc:
Topological: min(total_time)=9.93e-02 max(total_time)=9.93e-02 avg(total_time)=9.93e-02
Geometrical: min(total_time)=1.08e-01 max(total_time)=1.08e-01 avg(total_time)=1.08e-01
4

And if you want to enter the “danger-zone”/special case of entities=cells, you can set remote=False in

dofs_0 = locate_dofs_topological(Q, mesh.topology.dim, cells_0, remote=False)
dofs_1 = locate_dofs_topological(Q, mesh.topology.dim, cells_1, remote=False)

to improve performance even further

1 Like