Updating rotation matrices

Hello,
In a dynamic transversely isotropic material problem, I have defined a rotation matrix for each element. I suppose the rotation matrices need to be updated after each time step. But I have problems in doing so. I’ll show what I mean in the example below:

import dolfinx as fe
from dolfinx.fem.petsc import NonlinearProblem, LinearProblem, assemble_vector, assemble_matrix, create_vector, apply_lifting, set_bc
from dolfinx.nls.petsc import NewtonSolver
import numpy as np
import os
from mpi4py import MPI
import ufl
from petsc4py import PETSc


def main(linear = True, time_dependent = True):
    
    L = 1
    W = 0.2
    mesh = fe.mesh.create_box(MPI.COMM_WORLD, [np.array([0, 0, 0]), np.array([L, W, W])],
                         [20, 6, 6], cell_type=fe.mesh.CellType.hexahedron)
    element = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 1)
    V = fe.fem.functionspace(mesh, element)

    dt = fe.default_scalar_type(0.01)
    num_steps = int(1/float(dt))
    maximum_pressure = 15
    minimum_pressure = 0
    pressure = np.append(np.linspace(minimum_pressure, maximum_pressure, num=int(num_steps/2)),\
                         np.linspace(maximum_pressure, minimum_pressure, num=int(num_steps/2)))
    p = fe.fem.Constant(mesh, pressure[0])

    def clamped_boundary(x):
        return np.isclose(x[0], 0)

    fdim = mesh.topology.dim - 1
    clamped_boundary_facets = fe.mesh.locate_entities_boundary(mesh, fdim, clamped_boundary)
    u_D = np.array([0, 0, 0], dtype=fe.default_scalar_type)
    bc = fe.fem.dirichletbc(u_D, fe.fem.locate_dofs_topological(V, fdim, clamped_boundary_facets), V)

    boundaries = [(1, lambda x: np.isclose(x[0], L))]

    facet_indices, facet_markers = [], []
    fdim = mesh.topology.dim - 1
    for (marker, locator) in boundaries:
        facets = fe.mesh.locate_entities(mesh, fdim, locator)
        facet_indices.append(facets)
        facet_markers.append(np.full_like(facets, marker))
    facet_indices = np.hstack(facet_indices).astype(np.int32)
    facet_markers = np.hstack(facet_markers).astype(np.int32)
    sorted_facets = np.argsort(facet_indices)
    facet_tag = fe.mesh.meshtags(mesh, fdim, facet_indices[sorted_facets], facet_markers[sorted_facets])

    ##########################Defining the rotation matrix for each cell##################################
    ######################################################################################################
    el_r = ufl.TensorElement("DG", mesh.ufl_cell(), 0, shape=(3, 3))
    Q_r = fe.fem.FunctionSpace(mesh, el_r)
    bs_r = Q_r.dofmap.bs
    rotation_matrix = fe.fem.Function(Q_r)

    def get_rotation_matrix(x):
        values = np.zeros((bs_r, x.shape[1]), dtype=np.float64)
        values[0] = 1
        values[4] = 1
        values[8] = 1

        return values

    rotation_matrix.interpolate(get_rotation_matrix)
    ######################################################################################################

    ######################################################################################################
    #########################################Formulation##################################################
    u = ufl.TrialFunction(V)            
    v  = ufl.TestFunction(V)          
    old_u  = fe.fem.Function(V)
    old_velocity  = fe.fem.Function(V)
    old_acceleration  = fe.fem.Function(V)

    d = len(u)
    I = ufl.variable(ufl.Identity(d))             # Identity tensor
    F = ufl.variable(I + ufl.grad(u))             # Deformation gradient
    C = ufl.variable(F.T*F)                   # Right Cauchy-Green tensor
    J = ufl.variable(ufl.det(F))

    ####Check out for the metadata
    metadata = {"quadrature_degree": 4}
    ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_tag, metadata=metadata)
    dx = ufl.Measure("dx", domain=mesh, metadata=metadata)

    # Generalized-alpha method parameters
    alpha_m = fe.fem.Constant(mesh, 0.2)
    alpha_f = fe.fem.Constant(mesh, 0.4)
    gamma   = fe.default_scalar_type(0.5+alpha_f-alpha_m)
    beta    = fe.default_scalar_type((gamma+0.5)**2/4.)

    # Update formula for acceleration
    # a = 1/(2*beta)*((u - u0 - v0*dt)/(0.5*dt*dt) - (1-2*beta)*a0)
    def update_a(u, u_old, v_old, a_old, ufl=True):
        if ufl:
            dt_ = dt
            beta_ = beta
        else:
            dt_ = float(dt)
            beta_ = float(beta)
        return (u-u_old-dt_*v_old)/beta_/dt_**2 - (1-2*beta_)/2/beta_*a_old

    # Update formula for velocity
    # v = dt * ((1-gamma)*a0 + gamma*a) + v0
    def update_v(a, u_old, v_old, a_old, ufl=True):
        if ufl:
            dt_ = dt
            gamma_ = gamma
        else:
            dt_ = float(dt)
            gamma_ = float(gamma)
        return v_old + dt_*((1-gamma_)*a_old + gamma_*a)

    def update_rotation(R, F):  
        new_rotation = ufl.dot(F,R)
  
        return new_rotation  


    def update_fields(u_new, u_old, v_old, a_old, rotation_matrix, F):
        """Update fields at the end of each time step."""
        # Get vectors (references)
        u_vec, u0_vec  = u_new.x.array[:], u_old.x.array[:]
        v0_vec, a0_vec = v_old.x.array[:], a_old.x.array[:]

        # use update functions using vector arguments
        a_vec = update_a(u_vec, u0_vec, v0_vec, a0_vec, ufl=False)
        v_vec = update_v(a_vec, u0_vec, v0_vec, a0_vec, ufl=False)

        R_vec = update_rotation(rotation_matrix, F)
        new_rotation_expression = fe.fem.Expression(R_vec, Q_r.element.interpolation_points())  
        rotation_matrix.interpolate(new_rotation_expression)

        R_vec = rotation_matrix.x.array[:].reshape(num_cells_local, 3, 3)

        # Update (u_old <- u)
        v_old.x.array[:], a_old.x.array[:] = v_vec, a_vec
        u_old.x.array[:] = u_new.x.array[:]

        rotation_matrix.x.array[:] = R_vecs

    
    def avg(x_old, x_new, alpha):
        return alpha*x_old + (1-alpha)*x_new

    normal = -ufl.FacetNormal(mesh)
    
    rho = 1
    E = fe.default_scalar_type(1.0e4)
    nu = fe.default_scalar_type(0.3)
    mu = fe.fem.Constant(mesh, E / (2 * (1 + nu)))
    lmbda = fe.fem.Constant(mesh, E * nu / ((1 + nu) * (1 - 2 * nu)))
    
    def S(u):
        return 2.0 * mu * ufl.sym(ufl.grad(u)) + lmbda * ufl.tr(ufl.sym(ufl.grad(u))) * I

    acceleration = update_a(u, old_u, old_velocity, old_acceleration, ufl=False)
    velocity = update_v(acceleration, old_u, old_velocity, old_acceleration, ufl=True)

    formulation = rho * ufl.inner(alpha_m*old_acceleration + (1-alpha_m)* acceleration, v) * dx \
      + ufl.inner(ufl.grad(v), S(avg(old_u, u, alpha_f))) * dx \
      - ufl.inner(v, p * normal) * ds(1)

    bilinear_form = fe.fem.form(ufl.lhs(formulation))
    linear_form = fe.fem.form(ufl.rhs(formulation))

    ######################################################################################################
    ######################################################################################################

    ######################################################################################################
    ###############################################Solving################################################
    A = assemble_matrix(bilinear_form, bcs=[bc])
    A.assemble()
    b = create_vector(linear_form)

    solver = PETSc.KSP().create(mesh.comm)
    solver.setInitialGuessNonzero(True)
    solver.setOperators(A)
    solver.getPC().setType(PETSc.PC.Type.SOR)
    solver.view()

    displacement = fe.fem.Function(V)

    for i in range(num_steps):
        p.value = pressure[i]

        # Update the right hand side reusing the initial vector
        with b.localForm() as loc_b:
            loc_b.set(0)
        assemble_vector(b, linear_form)

        # Apply Dirichlet boundary condition to the vector
        apply_lifting(b, [bilinear_form], [[bc]])
        b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
        set_bc(b, [bc])

        # Solve linear problem
        solver.solve(b, displacement.vector)
        displacement.x.scatter_forward()

        # Update old fields with new quantities
        update_fields(displacement, old_u, old_velocity, old_acceleration, rotation_matrix, F)



if __name__ == "__main__":
    main()

The rotation matrix (which I’ve set to the identity matrix initially for keeping it simple) get updated in the for loop at the end in the ‘update_fields’ function. But I get the error:
RuntimeError: Cannot interpolate Expression with Argument.
I will appreciate any hint.

There is alot going on in this example, and I would strong advice you to try to minimize complexity for the chance of anyone having to parse it all.

For the context of varying properties, I would advice you to look at: Form compilation — FEniCS Tutorial @ Sorbonne
as it shows how to put form-compilation outside of loops.