fol.responses

fol.responses#

Sensitivity analysis and response evaluation provided by FoLax.

This module contains response objects that can be evaluated on top of a finite element (FE) state and a corresponding control field. A response typically represents a scalar objective or constraint functional of the form

\[J(d, u) = \sum_{e=1}^{N_e} \int_{\Omega_e} \phi(d(x), u(x)) \, \mathrm{d}\Omega ,\]

where d denotes the control field, u denotes the FE solution (DOFs), and the integrals are evaluated numerically using Gauss quadrature.

In addition to evaluating the scalar response value, FoLax responses support sensitivity analysis with respect to:

  • State variables (DOFs), via automatic differentiation.

  • Control variables, via automatic differentiation.

  • Shape variables (nodal coordinates), via automatic differentiation.

  • Adjoint-based gradients for large-scale problems where direct differentiation through the solver is expensive.

Finite element response class#

Authors: Reza Najian Asl, RezaNajian Date: January, 2025 License: FOL/LICENSE

class fol.responses.fe_response.FiniteElementResponse(name, response_formula, fe_loss, control)[source]#

Bases: Response

Finite element response evaluator and sensitivity calculator.

This class evaluates a scalar response functional over a finite element mesh using Gaussian quadrature. The response integrand is provided by the user as a Python expression string via response_formula and is compiled during Initialize() into a JAX-jittable callable.

The user provides only the scalar integrand expression as a string; the rest is handled automatically by this class: interpolation of nodal control values to Gauss points, interpolation of nodal DOFs (state) to Gauss points, quadrature weighting using the Jacobian determinant, element-wise accumulation, and global summation.

The compiled integrand callable is evaluated as:

phi(d_gp, u_gp)

where d_gp is the interpolated control value at a Gauss point and u_gp is the vector of interpolated DOFs at that Gauss point.

Parameters:

  • name (str) – Name identifier for the response instance.

  • response_formula (str) – Scalar-valued integrand expression as a string. The expression must be JAX compatible and evaluate to a scalar. The symbol jnp is available inside the expression namespace.

  • fe_loss (FiniteElementLoss) – Finite element loss object that provides mesh, element type, Gauss integration utilities, DOF layout, and residual/Jacobian assembly methods.

  • control (Control) – Control object that defines how optimization variables map to the nodal (or element) control field used by the response.

response_formula#

User-provided integrand expression string.

Type:

str

fe_loss#

Reference to the FE loss object.

Type:

FiniteElementLoss

control#

Reference to the control object.

Type:

Control

jit_response_function#

JAX-jitted callable compiled from response_formula during Initialize(). It is used at Gauss points during integration.

Type:

callable

CalculateNMatrix(N_vec)[source]#

Computes the shape function matrix (N) for finite elements.

This function generates a num_dofsx(num_dofs*N) shape function matrix, where N is the number of shape functions.

Parameters:

N_vec (jnp.array) – The vector of shape function values.

Returns:

The computed shape function matrix.

Return type:

jnp.array

ComputeAdjointJacobianMatrixAndRHSVector(nodal_control_values, nodal_dof_values)[source]#

Computes the adjoint Jacobian matrix and RHS vector for the finite element system.

The RHS vector is computed by summing element-wise contributions, applying Dirichlet boundary conditions, and scaling appropriately. The adjoint Jacobian matrix is obtained from the finite element loss function, which is transpose of the state Jacobian matrix.

Parameters:

  • nodal_control_values (jnp.array) – The global nodal control variable vector.

  • nodal_dof_values (jnp.array) – The global nodal state variable vector.

Returns:

A tuple containing:

  • sparse_jacobian (jnp.array): The computed adjoint Jacobian matrix.

  • rhs_vector (jnp.array): The computed RHS vector for the system.

Return type:

Tuple[jnp.array, jnp.array]

ComputeAdjointLossElementControlDerivativesVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector, full_adj_dof_vector)[source]#

Computes the control derivatives of the loss function for an element in a vectorized-compatible manner.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

  • full_adj_dof_vector (jnp.array) – The global adjoint state variable vector.

Returns:

The computed control derivatives for the given element.

Return type:

jnp.array

ComputeAdjointLossElementShapeDerivativesVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector, full_adj_dof_vector)[source]#

Computes the shape derivatives of the loss function for an element in a vectorized-compatible manner.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

  • full_adj_dof_vector (jnp.array) – The global adjoint state variable vector.

Returns:

The computed shape derivatives for the given element.

Return type:

jnp.array

ComputeAdjointNodalControlDerivatives(nodal_control_values, nodal_dof_values, nodal_adj_dof_values)[source]#

Computes the adjoint-based nodal control derivatives for the entire finite element mesh.

This function calculates local control derivatives for each element using automatic differentiation, then assembles the global derivative vector.

Parameters:

  • nodal_control_values (jnp.array) – The global nodal control variable vector.

  • nodal_dof_values (jnp.array) – The global nodal state variable vector.

  • nodal_adj_dof_values (jnp.array) – The global adjoint state variable vector.

Returns:

The assembled global control derivative vector.

Return type:

jnp.array

ComputeAdjointNodalShapeDerivatives(nodal_control_values, nodal_dof_values, nodal_adj_dof_values)[source]#

Computes the adjoint-based nodal shape derivatives for the entire finite element mesh.

This function calculates local shape derivatives for each element using automatic differentiation, then assembles the global derivative vector.

Parameters:

  • nodal_control_values (jnp.array) – The global nodal control variable vector.

  • nodal_dof_values (jnp.array) – The global nodal state variable vector.

  • nodal_adj_dof_values (jnp.array) – The global adjoint state variable vector.

Returns:

The assembled global shape derivative vector.

Return type:

jnp.array

ComputeElementRHSVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector)[source]#

Computes the RHS vector for a single element in a vectorized-compatible manner.

The element RHS vector is obtained as the gradient of the response with respect to the element’s state variables.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

Returns:

The computed RHS vector for the given element.

Return type:

jnp.array

ComputeFDNodalControlDerivatives(nodal_control_values, fe_solver, fd_step_size=0.0001, fd_mode='FWD')[source]#

Compute finite-difference sensitivities with respect to nodal control values.

This routine perturbs one entry of nodal_control_values at a time, re-solves the FE problem using the provided solver, and estimates the derivative of the total response value.

Supported modes: - "FWD": forward difference, - "CD": central difference.

Parameters:

  • nodal_control_values (jax.numpy.ndarray) – Global nodal control field (1D array).

  • fe_solver (FiniteElementSolver) – Solver used to compute state DOFs for each perturbed control.

  • fd_step_size (float, optional) – Perturbation step size. Default is 1e-4.

  • fd_mode (str, optional) – Finite-difference scheme. Supported values are "FWD" and "CD". Default is "FWD".

Returns:

Finite-difference gradient vector with the same shape as nodal_control_values.

Return type:

jax.numpy.ndarray

Raises:

Exception – If an unsupported fd_mode is requested.

ComputeFDNodalShapeDerivatives(nodal_control_values, fe_solver, fd_step_size=0.0001, fd_mode='FWD')[source]#

Compute finite-difference sensitivities with respect to nodal coordinates (shape).

This routine perturbs nodal coordinates of the FE mesh (in-place) and recomputes the response to estimate shape derivatives. The FE state is re-solved for each perturbation using the provided solver.

Supported modes: - "FWD": forward difference, - "CD": central difference.

Parameters:

  • nodal_control_values (jax.numpy.ndarray) – Global nodal control field used for all perturbations.

  • fe_solver (FiniteElementSolver) – Solver used to compute state DOFs for each perturbed mesh.

  • fd_step_size (float, optional) – Perturbation step size. Default is 1e-4.

  • fd_mode (str, optional) – Finite-difference scheme. Supported values are "FWD" and "CD". Default is "FWD".

Returns:

Flattened shape derivative vector with shape (num_nodes*3,) (the implementation stores 3 components per node, even if the FE problem dimension is smaller).

Return type:

jax.numpy.ndarray

Raises:

Exception – If an unsupported fd_mode is requested.

ComputeLossElementControlGrad(xyze, de, uvwe, adj_uvwe)[source]#

Computes the adjoint-based control gradient of the loss function for a given finite element.

This function calculates the sensitivity of the loss function with respect to control variables using automatic differentiation (jacobian of the residual) and element adjoint vector.

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

  • adj_uvwe (jnp.array) – The adjoint state variables.

Returns:

The control gradient of the loss function for the element.

Return type:

jnp.array

ComputeLossElementShapeGrad(xyze, de, uvwe, adj_uvwe)[source]#

Computes the adjoint-based shape gradient of the loss function for a given finite element.

This function calculates the sensitivity of the loss function with respect to nodal coordinates using automatic differentiation (jacobian of the residual) and adjoint vars.

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

  • adj_uvwe (jnp.array) – The adjoint state variables.

Returns:

The shape gradient of the loss function for the element.

Return type:

jnp.array

ComputeResponseElementValue(xyze, de, uvwe)[source]#

Computes the response value for a single finite element.

This method calculates the response contribution from a single element by integrating over the element’s Gauss points.

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

Returns:

The computed response value for the element.

Return type:

jnp.array

ComputeResponseElementValueControlGrad(xyze, de, uvwe)[source]#

Computes the gradient of the response’s element with respect to the control variables.

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

Returns:

The gradient of the response with respect to the control variables.

Return type:

jnp.array

ComputeResponseElementValueShapeGrad(xyze, de, uvwe)[source]#

Computes the gradient of the response’s element with respect to the shape (nodal coordinates).

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

Returns:

The gradient of the response with respect to the nodal coordinates, flattened.

Return type:

jnp.array

ComputeResponseElementValueStateGrad(xyze, de, uvwe)[source]#

Computes the gradient of the response’s element with respect to the state variables.

Parameters:

  • xyze (jnp.array) – The nodal coordinates of the element.

  • de (jnp.array) – The control variables associated with the element.

  • uvwe (jnp.array) – The state variables (displacements) associated with the element.

Returns:

The gradient of the response with respect to the state variables.

Return type:

jnp.array

ComputeResponseElementValueVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector)[source]#

Computes the response value for a single element in a vectorized-compatible manner.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

Returns:

The computed response value for the given element.

Return type:

jnp.array

ComputeResponseLocalNodalControlDerivativesVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector)[source]#

Computes the local nodal control derivatives of the response function for a given element in a vectorized-compatible manner.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

Returns:

The computed control derivatives for the given element.

Return type:

jnp.array

ComputeResponseLocalNodalShapeDerivativesVmapCompatible(element_id, elements_nodes, xyz, full_control_vector, full_dof_vector)[source]#

Computes the local nodal shape derivatives of the response function for a given element in a vectorized-compatible manner.

Parameters:

  • element_id (jnp.integer) – The ID of the element.

  • elements_nodes (jnp.array) – The connectivity matrix of elements to nodes.

  • xyz (jnp.array) – The coordinates of all nodes.

  • full_control_vector (jnp.array) – The global control variable vector.

  • full_dof_vector (jnp.array) – The global state variable vector.

Returns:

The computed shape derivatives for the given element.

Return type:

jnp.array

ComputeValue(nodal_control_values, nodal_dof_values)[source]#

Computes the total response value by summing the contributions from all elements.

Parameters:

  • nodal_control_values (jnp.array) – The global nodal control variable vector.

  • nodal_dof_values (jnp.array) – The global nodal state variable vector.

Returns:

The total computed response value.

Return type:

jnp.array

Finalize()[source]#

Finalizes the response computation.

This method must be implemented by subclasses.

Initialize(reinitialize=False)[source]#

Initializes the finite element response by setting up necessary computations.

If the response is already initialized, it will not be reinitialized unless explicitly requested.

Parameters:

reinitialize (bool, optional) – If True, forces reinitialization. Defaults to False.

Notes#

  • Element-wise response evaluation is performed using the FE element available in the associated fol.loss_functions.fe_loss.FiniteElementLoss instance. The response integrand is evaluated at Gauss points and summed over all elements.

  • Adjoint sensitivities are computed by building the adjoint right-hand side from the response derivative with respect to the state, and by using the transpose of the FE Jacobian provided by the loss function.

  • Finite-difference utilities are provided for verification of both control and shape sensitivities. Forward difference (FWD) and central difference (CD) are supported.

  • The response formula is provided as a string and compiled into a JAX-jitted function during initialization. The formula is evaluated at Gauss points using the interpolated control value and interpolated DOF values.

Contents