How to define 1d element embedded in 3d space?

I am trying to define nonlinear timoshenko beam by following code

from dolfinx import log, default_scalar_type
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver
import pyvista
import numpy as np
import ufl
from basix.ufl import element, mixed_element
from dolfinx import fem, io, mesh
from ufl import (Measure, SpatialCoordinate, TestFunctions, TrialFunctions, div, exp, inner, grad, Identity, dot, as_vector, outer)
from mpi4py import MPI
from dolfinx import fem, mesh, plot

L = 1.0
points = np.array([[0.0, 0.0, 0.0], [L, 0.0, 0.0]])
# Create a 1D mesh for the beam
beam_mesh = mesh.create_interval(MPI.COMM_WORLD, 5, points[:,0])

z = element("CG", beam_mesh.basix_cell(), degree=2, shape=(3,))
theta = element("CG", beam_mesh.basix_cell(), degree=1, shape=(3,))

element = mixed_element([z, theta])

V = fem.functionspace(beam_mesh, element)

# Create the function space and functions from the mixed element
z_theta = fem.Function(V)
z, theta = ufl.split(z_theta)  # Split the mixed function into z and theta

# Define the cross-sectional position vector X_S
# For example, at (0, y, z) where y and z are offsets from the beam axis
X_S = ufl.as_vector([0.0, 0.1, 0.1])

def rotation_matrix(theta):
    """Constructs a rotation matrix using the vector theta."""
    theta_x, theta_y, theta_z = theta[0], theta[1], theta[2]
    # Approximate rotation matrix for small rotations
    return as_matrix([[1, -theta_z, theta_y],
                      [theta_z, 1, -theta_x],
                      [-theta_y, theta_x, 1]])

# Compute the rotation matrix R from theta
R = rotation_matrix(theta)

# Calculate u based on Equation (25)
u = z + dot((R - Identity(3)), X_S)

# Now u can be used as a function or expression in FEniCSx
print("Displacement field u defined as:", u)

print(Identity(3))

def compute_strain_components(u, theta, t0):
    """
    Computes the strain components based on Equation (36) for a 1D beam in 3D.
    
    Parameters:
    ----------
    u : Function
        Displacement vector function in 3D.
    theta : Function
        Rotation vector function in 3D.
    t0 : UFL vector
        Tangent vector along the beam's axis (typically x-axis).
        
    Returns:
    --------
    e : UFL vector
        Strain vector with components [epsilon_11, gamma_12, gamma_13]
    """
    # Construct grad(u) as a 3x3 matrix for the 1D beam in 3D
    grad_u = as_matrix([[grad(u[0])[0], grad(u[0])[1], grad(u[0])[2]],
                        [grad(u[1])[0], grad(u[1])[1], grad(u[1])[2]],
                        [grad(u[2])[0], grad(u[2])[1], grad(u[2])[2]]])

    # Deformation gradient F based on the definition F = I + grad(u)
    F = Identity(3) + grad_u  # F is a 3x3 matrix

    # Extract columns for strain calculations
    f1 = F[:, 0]  # First column of F (axial strain)
    f2 = F[:, 1]  # Second column of F (shear strain in xy-plane)
    f3 = F[:, 2]  # Third column of F (shear strain in xz-plane)

    # Compute strain components based on Equation (36)
    epsilon_11 = 0.5 * (dot(f1, f1) - 1)   # Axial strain
    gamma_12 = 0.5 * dot(f2, f1)           # Shear strain in xy-plane
    gamma_13 = 0.5 * dot(f3, f1)           # Shear strain in xz-plane

    # Combine strains into a single vector
    e = as_vector([epsilon_11, gamma_12, gamma_13])
    
    return e

# Example tangent vector along the x-axis (beam length direction)
t0 = as_vector([1.0, 0.0, 0.0])

# Compute strain vector based on displacement and rotation
strain_vector = compute_strain_components(u, theta, t0)

but I am getting following issues

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[25], line 97
     94 t0 = as_vector([1.0, 0.0, 0.0])
     96 # Compute strain vector based on displacement and rotation
---> 97 strain_vector = compute_strain_components(u, theta, t0)

Cell In[25], line 71, in compute_strain_components(u, theta, t0)
     53 """
     54 Computes the strain components based on Equation (36) for a 1D beam in 3D.
     55 
   (...)
     68     Strain vector with components [epsilon_11, gamma_12, gamma_13]
     69 """
     70 # Construct grad(u) as a 3x3 matrix for the 1D beam in 3D
---> 71 grad_u = as_matrix([[grad(u[0])[0], grad(u[0])[1], grad(u[0])[2]],
     72                     [grad(u[1])[0], grad(u[1])[1], grad(u[1])[2]],
     73                     [grad(u[2])[0], grad(u[2])[1], grad(u[2])[2]]])
     75 # Deformation gradient F based on the definition F = I + grad(u)
     76 F = Identity(3) + grad_u  # F is a 3x3 matrix

File /opt/tljh/user/lib/python3.10/site-packages/ufl/exproperators.py:383, in _getitem(self, component)
    380 shape = self.ufl_shape
    382 # Analyse slices (:) and Ellipsis (...)
--> 383 all_indices, slice_indices, repeated_indices = create_slice_indices(
    384     component, shape, self.ufl_free_indices
    385 )
    387 # Check that we have the right number of indices for a tensor with
    388 # this shape
    389 if len(shape) != len(all_indices):

File /opt/tljh/user/lib/python3.10/site-packages/ufl/index_combination_utils.py:150, in create_slice_indices(component, shape, fi)
    148 elif isinstance(ind, int):
    149     if int(ind) >= shape[len(all_indices)]:
--> 150         raise ValueError("Index out of bounds.")
    151     all_indices.append(FixedIndex(ind))
    152 elif isinstance(ind, slice):

ValueError: Index out of bounds.

how may I define 1d element embedded in 3d space, so that I can take the gradient of displacement properly