Using ufl expressions with the form

I am trying to implement a Linear Elastic case using a custom newton nonlinear solver where stresses are passed into the solver at quadrature points. (this is to help me implement a plasticity code for later). I have tried using the form in the following way.

import basix
from ufl import *
from petsc4py.PETSc import ScalarType
from dolfinx import fem

def epsilon(u):
    return 0.5*(nabla_grad(u) + nabla_grad(u).T)
def sigma(u):
    mu = 1.0
    lambda_ = 1.25
    return lambda_ * nabla_div(u) * Identity(u.geometric_dimension()) + 2*mu*epsilon(u)

element = VectorElement("Lagrange", hexahedron, 1, dim = 3)
mesh = Mesh(element) 

# Define quadrature rule
quadrature_rule = "default" 
quadrature_degree = 3

# Get points and weights for custom integration measure
points, weights = basix.make_quadrature(basix.quadrature.string_to_type(quadrature_rule),
                                        basix.cell.string_to_type(str(mesh.ufl_cell())), quadrature_degree)
metadata = {"quadrature_rule": "custom",
            "quadrature_points": points, "quadrature_weights": weights}
dx = Measure("dx", domain=mesh, metadata=metadata)

# Create quadrature space
Q_element = TensorElement("Quadrature", mesh.ufl_cell(), quadrature_degree, quad_scheme=quadrature_rule)
Q = FunctionSpace(mesh, Q_element)

# Define variational form
V = FunctionSpace(mesh, element)
ds = Measure("ds", mesh)
u = Coefficient(V)
v = TestFunction(V)
f = Constant(mesh, shape=(3, ))
P = Constant(mesh, shape=(3, ))
sigma = Coefficient(Q)

F_form = inner(sigma, epsilon(v)) * dx - dot(f, v) * dx - dot(P, v) * ds

J_form = derivative(F_form, u)

# Interpolation error between V and Q
J = inner(sigma(u) * sigma(u) - sigma, sigma(u) * sigma(u) - sigma) * dx

forms = [F_form, J_form, J]

# sigma expression
mu = 1.0
lambda_ = 1.25
expr = lambda_ * nabla_div(u) * Identity(u.geometric_dimension()) + 2*mu*epsilon(u)
elements =[(element)]
expressions = [(expr, points)]

However I am not sure if it possible to link the expression to the form in the C++ file. I am not sure if what I am trying to do is possible as the form no longer contains “u” when using the expression. This is what I have tried to implement in C++ but it isn’t working:

#include "linearelasticity.h"
#include <basix/finite-element.h>
#include <cmath>
#include <dolfinx.h>
#include <dolfinx/common/log.h>
#include <dolfinx/fem/assembler.h>
#include <dolfinx/fem/petsc.h>
#include <dolfinx/fem/Function.h>
#include <dolfinx/fem/Constant.h>
#include <dolfinx/io/XDMFFile.h>
#include <dolfinx/la/Vector.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/cell_types.h>
#include <dolfinx/nls/NewtonSolver.h>
#include <dolfinx/mesh/generation.h>
#include <vector>
#include <utility>

using namespace dolfinx;
using T = PetscScalar;

int main(int argc, char* argv[])
{
  dolfinx::init_logging(argc, argv);
  PetscInitialize(&argc, &argv, nullptr, nullptr);
  
  // Set the logging thread name to show the process rank
  int mpi_rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
  std::string thread_name = "RANK " + std::to_string(mpi_rank);
  loguru::set_thread_name(thread_name.c_str());

  {
    auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet);
    auto mesh = std::make_shared<mesh::Mesh>(mesh::create_box(
        MPI_COMM_WORLD, {{{0.0, 0.0, 0.0 },  {0.2, 0.2, 2.0}}}, {6, 6, 20},
        mesh::CellType::hexahedron, part));

    // Get number of cells on process
    const dolfinx::mesh::Topology& topology = mesh->topology();
    const std::size_t tdim = topology.dim();
    const std::size_t num_cells = topology.index_map(tdim)->size_local()
                                  + topology.index_map(tdim)->num_ghosts();
    std::vector<std::int32_t> cells(num_cells);
    std::iota(cells.begin(), cells.end(), 0);

    auto V = std::make_shared<fem::FunctionSpace>(
        fem::create_functionspace(functionspace_form_linearelasticity_F_form, "v", mesh));
    
    auto Q = std::make_shared<fem::FunctionSpace>(
        fem::create_functionspace(functionspace_form_linearelasticity_J, "sigma", mesh));
    
    // Prepare and set Constants for the bilinear form
    auto f = std::make_shared<fem::Constant<T>>(std::vector<T>{0.0, 0.001, 0.0});
    auto P = std::make_shared<fem::Constant<T>>(std::vector<T>{0.0, 0.0, 0.0});
    // Define variational forms
    auto u = std::make_shared<fem::Function<T>>(V);

    auto sigma = std::make_shared<fem::Function<T>>(Q);

    std::span<T> s_arr = sigma->x()->mutable_array();
    int n = s_arr.size();
    for (int i = 0; i < n; i++) {
        s_arr[i] = 0.0;
    }
    auto scalar_form = fem::create_form<PetscScalar>(
        *form_linearelasticity_J, {}, {{"u", u}, {"sigma", sigma}}, {}, {}, mesh);

    auto jacobian = std::make_shared<fem::Form<T>>(fem::create_form<T>(
        *form_linearelasticity_J_form, {V, V}, {{"u", u}, {"sigma", sigma}}, {{"f", f}, {"P", P}}, {}));

    auto residual = std::make_shared<fem::Form<T>>(fem::create_form<T>(
        *form_linearelasticity_F_form, {V}, {{"u", u}, {"sigma", sigma}}, {{"f", f}, {"P", P}}, {}));
    // Define boundary condition
    auto facets = mesh::locate_entities_boundary(
        *mesh, 2,
        [](auto&& x)
        {
            constexpr double eps = 1.0e-8;
            std::vector<std::int8_t> marker(x.extent(1), false);
            for (std::size_t p = 0; p < x.extent(1); ++p)
            {
                double z = x(2, p); // x = x(0, p), y = x(1, p), z = x(2, p)
                if (std::abs(z) < eps) 
                    marker[p] = true;  
            }
            return marker;
        });
    const auto bdofs = fem::locate_dofs_topological({*V}, 2, facets);
    auto bcs = std::vector{std::make_shared<const fem::DirichletBC<T>>(P, bdofs, V)};

    
    Vec dx = nullptr;
    la::petsc::Matrix A = la::petsc::Matrix(fem::petsc::create_matrix(*jacobian, "baij"), false); 
    la::Vector<T> L(residual->function_spaces()[0]->dofmap()->index_map,
           residual->function_spaces()[0]->dofmap()->index_map_bs());


    //initialize solver
    la::petsc::KrylovSolver solver(mesh->comm());
    dolfinx::MPI::Comm comm(mesh->comm());
    solver.set_options_prefix("nls_solve_");
    la::petsc::options::set("nls_solve_ksp_type", "preonly");
    la::petsc::options::set("nls_solve_pc_type", "lu");
    #if PETSC_HAVE_MUMPS
    la::petsc::options::set("nls_solve_pc_factor_mat_solver_type", "mumps");
    #endif
    solver.set_from_options();
    solver.set_operators(A.mat(), A.mat());

    if (!dx)
        MatCreateVecs(A.mat(), &dx, nullptr);

    //ghost initialization
    auto map = residual->function_spaces()[0]->dofmap()->index_map;
    const int bs = residual->function_spaces()[0]->dofmap()->index_map_bs();
    std::int32_t size_local = bs * map->size_local();
    std::vector<PetscInt> ghosts(map->ghosts().begin(), map->ghosts().end());
    std::int64_t size_global = bs * map->size_global();
    Vec _b_petsc = nullptr;
    VecCreateGhostBlockWithArray(map->comm(), bs, size_local, size_global,
                                 ghosts.size(), ghosts.data(),
                                 L.array().data(), &_b_petsc);

    int max_iterations = 25;
    int i = 0;

    std::vector<T> L2_error;
    std::vector<T> dx_norm;

    la::petsc::Vector _u(la::petsc::create_vector_wrap(*u->x()), false);
    while (i < max_iterations)
    {
        // Assemble Jacobian and residual
        std::span<T> b(L.mutable_array());
        std::fill(b.begin(), b.end(), 0.0);
        fem::assemble_vector<T>(b, *residual);
        VecGhostUpdateBegin(_b_petsc, ADD_VALUES, SCATTER_REVERSE);
        VecGhostUpdateEnd(_b_petsc, ADD_VALUES, SCATTER_REVERSE);


        MatZeroEntries(A.mat());
        fem::assemble_matrix(la::petsc::Matrix::set_block_fn(A.mat(), ADD_VALUES), *jacobian,
                           bcs);
        MatAssemblyBegin(A.mat(), MAT_FLUSH_ASSEMBLY);
        MatAssemblyEnd(A.mat(), MAT_FLUSH_ASSEMBLY);
        fem::set_diagonal(la::petsc::Matrix::set_fn(A.mat(), INSERT_VALUES),
                          *jacobian->function_spaces()[0], bcs);
        MatAssemblyBegin(A.mat(), MAT_FINAL_ASSEMBLY);
        MatAssemblyEnd(A.mat(), MAT_FINAL_ASSEMBLY);

        // Compute b - J(u_D-u_(i-1))
        //fem::apply_lifting(b, {jacobian}, {bcs}, {}, 1.0);
        fem::petsc::apply_lifting(_b_petsc, {jacobian}, {bcs}, {_u.vec()} , 1.0);

        // Set bcs
        Vec x_local;
        VecGhostGetLocalForm(_u.vec(), &x_local);
        PetscInt n = 0;
        VecGetSize(x_local, &n);
        const T* array = nullptr;
        VecGetArrayRead(x_local, &array);
        fem::set_bc<T>(b, bcs, std::span<const T>(array, n), -1.0);
        VecRestoreArrayRead(_u.vec(), &array);
        
        VecGhostUpdateBegin(_b_petsc, INSERT_VALUES, SCATTER_FORWARD);
        VecGhostUpdateEnd(_b_petsc, INSERT_VALUES, SCATTER_FORWARD);

        // Solve linear problem
        solver.solve(dx, _b_petsc);
        // Update u_{i+1} = u_i + delta x_i
        VecAXPY(_u.vec(), -1.0, dx);
        ++i;

        //Compute norm of update
        PetscReal correction_norm = 0.0;
        VecNorm(dx, NORM_2, &correction_norm);

        // Compute L2 error comparting to the analytical solution
        //L2_error.push_back(fem::assemble_scalar(*E));
        dx_norm.push_back(correction_norm);
       
        // Interpolate expression into quadrature points
    {
      const auto V_to_Q = fem::create_expression<T>(*expression_linearelasticity_expr,
                                                    {{"u", u}}, {}, mesh);
      auto [eval_points, _shape] = V_to_Q.X();
      std::array<std::size_t, 2> shape
          = {cells.size(), _shape[0] * V_to_Q.value_size()};
      std::vector<T> values(shape[0]*shape[1]);

      V_to_Q.eval(cells, values, shape);

      // Move values into sigma
      dolfinx::la::petsc::Vector _sigma(
          dolfinx::la::petsc::create_vector_wrap(*sigma->x()), false);
      VecSetBlockSize(_sigma.vec(), Q->dofmap()->bs());
      VecSetValuesBlockedLocal(
          _sigma.vec(), sigma->x()->array().size() / Q->dofmap()->bs(),
          Q->dofmap()->list().array().data(), values.data(), INSERT_VALUES);
      sigma->x()->scatter_fwd();
    }
    PetscScalar J = fem::assemble_scalar(scalar_form);
    std::cout << "Interpolation error " << J << "\n";
        if (correction_norm < 1e-10) {
            std::cout <<"converged at i = " << i << std::endl;
            break;
        }
        std::cout << correction_norm << std::endl;

    }
    // Save solution in VTK format
    io::VTKFile file_u(mesh->comm(), "u.pvd", "w");
    file_u.write<T>({*u}, 0.0);
  }

}

For plasticity, I’d strongly suggest to have a look at Rankine demo in dolfiny, dolfiny/rankine.py at master · michalhabera/dolfiny · GitHub
The way it is done in dolfiny would be a preferred way for composition of custom and generated code.

You could also generate code for expression similar to what you do, but then the rest of the code would have to be written manually, including the quadrature loop and basis functions for u etc.

1 Like