Numerical values from ufl.SpatialCoordinate

dolfinx
1

Hi everyone!

I am implementing a Lubrication + Wear model in FeniCSx for my Masters Thesis.

While i was implementing further improvements i noticed, my previous approach will not be suitable to implement something i would need.

What I am trying to do is the following:

In my variational formulation i have a call to a height function, that returns the separation between the two geometries I am simulating.

Inside this function, i would want to obtain the value from a dictionary, which maps each node in the mesh to a certain value which changes over time.

However, i am unable to obtain the numerical coordinates (which I would need, as they are the key to the dictionary), as when solving, this height function is called with my ufl.SpatialCoordinate object.

There are probably many ways of doing such, do you have any suggestions? For context, the reason i am using a dictionary is because there is no analytic expression to calculate the value i map.

Many thanks!

Code:

from abc import ABC, abstractmethod ## Para clases abstractas
import ufl ## Parte de FeniCSx
from dolfinx import fem, io, mesh, plot, geometry ## Parte de FeniCSx
from mpi4py import MPI ## Parte de FeniCSx
from petsc4py.PETSc import ScalarType ## Parte de FeniCSx
import math, random ## Operaciones matemƔticas bƔsicas
import numpy as np ## Numpy para poder operar con arrays
import matplotlib.pyplot as mpl ## Para pintamiento

class Problema: 

    ## Constructor
    def __init__(self, datos):
        self.datos = datos ## Diccionario que contiene los datos del problema
        self.deltat = 0.1 ## Placeholder 
        self.datos["speed"] = ((self.datos["Ut"] - self.datos["Ub"])**2 +  (self.datos["Vt"] - self.datos["Vb"])**2)**(1/2)
        self.BC = []  ## Array que contiene las condiciones de contorno
        ## Diccionario que contiene los valores de diversos parĆ”metros para la iteraciĆ³n inicial. 
        self.timestepvals= {"Zh": 3, 
                                              "DeltaZh": 0,
                                              "DeltaZ0": 0, 
                                              "S1" : 0} 
    
    ## Setter para la variable h0 para el calculo de la altura. Necesario para poder variar la altura en calculos iterativos. 
    def set_h0(self,h0):
        self.datos["h0"] = h0
    
    ## Metodo abstracto para el calculo de la altura. Obligamos a las clases hijas a definirla de forma que represente el problema
    @abstractmethod
    def height(self):
        pass

    ## Clase para problemas tipo esfera - point contact
class ProblemaRodillo(Problema):

    ## Override del constructor para calcular el Ɣrea del contacto de forma especƭfica para cada caso. 
    def __init__(self, datos):
        Problema.__init__(self, datos) # Llamamos primero al constructor por defecto 
        self.datos["Area_nom"] = math.pi* self.datos["rRod"]**2 

    def height(self,x): 
        
        ## THIS VALUE SHOULD BE SOMETHING LIKE
        # worn = self.worn[(x[0], x[1], x[2])].
        worn = 0.001 

        return self.datos["h0"] + (x[0]**2)/(2*self.datos["rRod"]) # + worn
    
datosPruebaPsi={
    "rRod": 0.03, 
    "Load": 10e4,
    "h0min" : 5e-6,
    "h0max": 6e-6,
    "rho": 860,
    "mu": 0.1,
    "Ub":1,
    "Ut":0,
    "Vb":0,
    "Vt":0,
    
    "domain":{
        "xMin":-0.05,
        "xMax":0.05,
        "yMin":-0.5,
        "yMax":0.5,
        "numNodes":50
    }, 
}

Prueba = ProblemaRodillo(datosPruebaPsi)

xMin = Prueba.datos["domain"]["xMin"]
xMax = Prueba.datos["domain"]["xMax"]
yMin = Prueba.datos["domain"]["yMin"]
yMax = Prueba.datos["domain"]["yMax"]
numNodes = Prueba.datos["domain"]["numNodes"]
msh = mesh.create_rectangle(comm=MPI.COMM_WORLD,
                    points=((xMin, yMin), (xMax, yMax)), n=(numNodes, numNodes),
                    cell_type=mesh.CellType.triangle)

Prueba.FUNCSPACE = fem.FunctionSpace(msh, ("Lagrange", 2))
Prueba.PSPACE = ufl.TrialFunction(Prueba.FUNCSPACE)
Prueba.VSPACE = ufl.TestFunction(Prueba.FUNCSPACE)


# Definimos las fronteras para aplicar la condicion de contorno
facets = mesh.locate_entities_boundary(msh, dim=1, marker=lambda x: np.logical_or(np.logical_or(np.isclose(x[0],Prueba.datos["domain"]["xMin"]),np.isclose(x[0], Prueba.datos["domain"]["xMax"])),
                            np.logical_or(np.isclose(x[1], Prueba.datos["domain"]["yMin"]),np.isclose(x[1],Prueba.datos["domain"]["yMax"]))))
# Localizamos los grados de libertad en las fronteras                                                                  
dofs = fem.locate_dofs_topological(V=Prueba.FUNCSPACE, entity_dim=1, entities=facets)
# Aplicamos condiciones de contorno tipo dirichlet
bc = fem.dirichletbc(value=ScalarType(0), dofs=dofs, V=Prueba.FUNCSPACE)
## Introducimos en el array que contiene las condiciones de contorno
Prueba.set_h0(Prueba.datos["h0max"])
rho = Prueba.datos["rho"]
mu = Prueba.datos["mu"]
Um = 0.5*(Prueba.datos["Ub"]+Prueba.datos["Ut"])
Vm = 0.5*(Prueba.datos["Vb"]+Prueba.datos["Vt"])

x = ufl.SpatialCoordinate(msh)
p = Prueba.PSPACE
v = Prueba.VSPACE

## FormulaciĆ³n variacional
a = (rho*(Prueba.height(x)**3)/(12*mu))*ufl.inner(ufl.grad(p),ufl.grad(v))*ufl.dx
L = rho*Prueba.height(x)*ufl.inner(ufl.tensors.as_vector([Um,Vm]),ufl.grad(v))*ufl.dx 

## Asignamos el problema a nuestra clase
problem = fem.petsc.LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "ibcgs", "pc_type":"lu"})
Prueba.solution = problem.solve()

points = np.zeros((3, len(msh.geometry.x[:,0])))
points[0] = msh.geometry.x[:,0]

bb_tree = geometry.BoundingBoxTree(msh, msh.topology.dim)
cell_candidates = geometry.compute_collisions(bb_tree, points.T)
colliding_cells = geometry.compute_colliding_cells(msh, cell_candidates, points.T)
cells = [colliding_cells.links(i)[0] for i, point in enumerate(points.T)]

p_values = Prueba.solution.eval(points.T, cells)

mpl.scatter(points[0,:], p_values, label="Numerical solution")
2

ufl.SpatialCoordinate does not imply the ā€œmesh-nodesā€ (meaning the vertices for linear meshes). The x=ufl.SpatialCoordinate(mesh) represents any quadrature point used during integration in its physical space.

What I would suggest to do, is to interpolate your function (non-analytical) into a suitable function space. Say, if you want to actually use the mesh nodes (vertices), you can create a first order Lagrange function and interpolate into that, i.e.

V  = dolfinx.fem.FunctionSpace(mesh, ("Lagrange", 1))
u = dolfinx.fem.Function(V)

def your_expression(x):
     # map your values here
     return value
u.interpolate(your_expression)

If you instead you want to use the quadrature points (as described by x=ufl.SpatialCoordinate(mesh), I suggest something like the following:

from IPython import embed
import dolfinx
from mpi4py import MPI
import ufl
import basix
import numpy as np
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 5)


def height(x):
    # Add whatever function you want to call here, for now choose identity
    return x[0]


quadrature_degree = 4
dx = ufl.Measure("dx", domain=mesh, metadata={
                 "quadrature_degree": quadrature_degree})
quadrature_points, wts = basix.make_quadrature(
    basix.cell.string_to_type(mesh.ufl_cell().cellname()), quadrature_degree)
Q_element = ufl.FiniteElement(
    "Quadrature", mesh.ufl_cell(), quadrature_degree, quad_scheme="default")
Q = dolfinx.fem.FunctionSpace(mesh, Q_element)
height_function = dolfinx.fem.Function(Q)

x = ufl.SpatialCoordinate(mesh)
x_expr = dolfinx.fem.Expression(x, quadrature_points)
num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
integration_points = x_expr.eval(np.arange(num_cells, dtype=np.int32))

num_points_per_cell = quadrature_points.shape[0]
gdim = mesh.geometry.dim
assert(integration_points.shape[1] == num_points_per_cell*gdim)
assert(integration_points.shape[0] == num_cells)
for i in range(num_cells):
    for j in range(num_points_per_cell):
        coord = integration_points[i, gdim*j:gdim*(j+1)]
        value = height(coord)
        height_function.x.array[i*num_points_per_cell +
                                j:i*num_points_per_cell + (j+1)] = value

for i in range(mesh.geometry.dim):
    # verify that height_function maps to x[0]
    print(dolfinx.fem.assemble_scalar(
        dolfinx.fem.form((height_function-x[0])**2*dx)))


# Use heigh function in general()
V = dolfinx.fem.FunctionSpace(mesh, ("Lagrange", 2))
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
a = dolfinx.fem.form(height_function*ufl.inner(ufl.grad(u), ufl.grad(v))*dx)
A = dolfinx.fem.petsc.assemble_matrix(a)
1 Like
3

Good evening!

Many thanks for the response, @dokken! With your hints I have been able to achieve what I wanted.

I tried the interpolation route, but i was getting nonsensical values ( e-311, i suppose because of the mapping i was doing)

While testing this, i tried to manually assign the fem.Function.vector.array, and realized that that is exactly the structure that i need to save the data. I need a way to assign a value to each DOF, thus the vector of a Function fits perfectly

It is probably not the intended usage, but after some quick tests i am confident I will be able to work like that.

Again, thanks a lot!!

4

Brief follow up on this: Is there any way to differentiate between MeshCoordinates living in different function spaces?

To be more precise, I have a function defined in the function space of the problem (ā€˜CGā€™, 2), and Iā€™m interested in finding out how this function changes when changing the mesh coordinates (the order of the mesh is unknown). Is there a straightforward way to do this?

5

Are you here talking about a shape derivative of a function?
If so, i would suggest reading

6

Thanks for the reference, this is indeed the paper Iā€™m looking at. The issue is, looking for example at Listing 1:

mesh = UnitSquareMesh (10, 10)
X, y = X = SpatialCoordinate (mesh)
J = (X*X + y*y - 1) * dx
dJ = assemble (derivative (J, X))

the functions all live in the ā€˜geometryā€™ function space. Iā€™m basically trying to replicate this but with functions that live in a high-order CG space, to differentiate them wrt the geometry.

EDIT: One option is to interpolate the SpatialCoordinate into the higher order field, but this destroys the ā€˜differentiablityā€™, to the best of my understanding.

EDIT2: For some extra background:

I simulate a system under an external load, which is a spatial function, and lives in a high-order CG space. Iā€™m interested in the derivative of a misfit (itself, a function of the solution of a FEM problem that involves this potential) with respect to the geometry. I first compute the derivative of the misfit wrt the explicit dependence on the mesh, and also wrt the load function (which is an interpolated function). However, since the load function is quite complex, I like to implement it in JAX. Then I use JAX to backpropagate the gradient of the misfit wrt the gradient of the interpolation points. Now I have a field ā€“ with as many DOFs as the high-order CG space, that contains the gradient of the misfit wrt the interpolation points in the high-order space. Now I would like to move this back into the mesh space to feed it into my gmsh-based sensitivity analysis.

7

You can simply do something along the lines of

mesh = UnitSquareMesh (10, 10)
V=FunctionSpace(mesh, "Lagrange", 2)
u=Function(V)
X = SpatialCoordinate (mesh)
J = (u*u) * dx
dJ = assemble (derivative (J, X))

which computes the shape derivative of J for any given function u

8

I struggle a bit getting the behavior I need from this, because I already have the gradient of the misfit wrt the higher-order mesh coordinates df/dx_i, and what Iā€™m trying to do is propagate it to a gradient of the misfit with respect to the mesh coordinates. One initial idea was to set up a form equal to dot(X, df/dx), where df/dx is the derivative obtained from JAX, and then differentiate wrt X, but then the result is multiplied by a mass matrix.

Of course, the culprit is that my function ā€˜Jā€™ is in reality a JAX function instead of a ufl function like in your example. But this is kind of necessary given that the problem requires the flexibility of implementing nonlinear functions beyond those available in ufl.

9

Its quite hard for me to understand what you want to achieve without an minimal example.

It is unclear to me how your load

Is not something similar to

Where u for instance has x^2+y^2-1 interpolated into itself prior to assembly.

10

Thanks for your patience, thatā€™s an example of what Iā€™m trying to do. It is possible that I misunderstood something fundamental:

# Here X is mesh coordinate, x is high-order space coordinate
def complexJaxFunction(x):
    return jnp.sin(x[:,0]) # Only as example
mesh = UnitSquareMesh (10, 10)
V = FunctionSpace(mesh, "Lagrange", 2)
u = Function(V)
u.interpolate(lambda x: np.array(complexJaxFunction(x))
J = (u*u) * dx
X = SpatialCoordinate (mesh)
dJ = assemble (derivative (J, X)) # Explicit dependence of J on mesh, but not total derivative

# total derivative is dJ/dX + (dJ/du) (du/dX) =  dJ/dX + (dJ/du)(du/dx)(dx/dX)
# Compute (dJ/du)
v = TestFunction(V)
dJdu = assemble (derivative (J, u, v)) # Dependence of J on u

# Propagate dJ/du to dJ/dx
q = V.tabulate_dof_coordinates()
dJdx = jax.vjp(lambda x: complexJaxFunction(x), q)(dJdu.vector.array[:]) # derivative of J wrt interpolation points

# How do I go from dJ/dx to dJ/dX???
dJdX = f(dJdx) ?????

result = dJ + dJdX

Edit: An unsuccessful attempt at this is to bring the misfit back to FEniCSx:

uu = Function(V)
uu.vector.array[:] = dJdx
J1 = ufl.inner(uu, X) * dx
dJdX = assemble (derivative (J1, X))

Now this does not work because dx scales every component in nontrivial ways, would work if there was a measure that gave a standard scalar product.

11

In the end, I figured out a workaround, in case someone else encounters the same problem. The trick is to interpolate the spatial gradient of the function too:

# Non-UFL (JAX)force function
def complexFunction(x):
    return jnp.array([jnp.sin(x[1]), jnp.sin(x[2]), x[0]])

# Interpolate function
functionSpace = fem.functionspace(domain, ("Lagrange", 2, (3,)))
f = jax.vmap(complexFunction, in_axes=0, out_axes=0)
v = fem.Function(functionSpace)
v.interpolate(lambda x: np.array(f(x.T)).T)

# Interpolate gradient
row_jacobian = jax.jacrev(complexFunction)
jf = jax.vmap(row_jacobian, in_axes=0, out_axes=0)
gradientFunctionSpace = fem.functionspace(domain, ("Lagrange", 2, (3,3)))
dofCoordinates = gradientFunctionSpace.tabulate_dof_coordinates()
spatialGradient = fem.Function(gradientFunctionSpace)
spatialGradient.vector.array[:] = np.array(jf(dofCoordinates)).ravel()

# Create corrected functional
originalSpatialCoordinate = fem.Function(functionSpace)
originalSpatialCoordinate.vector.array[:] = functionSpace.tabulate_dof_coordinates().ravel()
qt = q + ufl.dot(spatialGradient, spatialCoordinate - originalSpatialCoordinate)

# Define form to differentiate
formToDifferentiate = originalForm = ufl.inner(qt,qt)

# Compute gradient
gradient = ufl.derivative(formToDifferentiate*dx, spatialCoordinate, dX)    
derivative = fem.assemble_vector(fem.form(gradient))

This works nicely, I get a gradient of -0.00027387460839545713. With the direct approach from the paper, I was getting -0.0015833329663521318 (almost 500% difference).

if I numerically perturb the mesh using central differences:

domain.geometry.x[1][0] = x0 +/- eps

And then recompute the whole thing, I get a gradient of -0.00027443042688313213 that matches the new results.

1 Like