fol.loss_functions

Contents

fol.loss_functions#

Loss functions provided by FoLax.

Finite element loss base class#

Authors: Reza Najian Asl, RezaNajian Date: April, 2024 License: FOL/LICENSE

class fol.loss_functions.fe_loss.FiniteElementLoss(name, loss_settings, fe_mesh)[source]#

Bases: Loss

Base class for finite element loss functionals.

This class provides the common infrastructure required to define loss functionals on a finite element mesh. Concrete loss functions in fol.loss_functions are typically derived from this class and implement the abstract element routine ComputeElement().

The base class handles:

  • DOF bookkeeping and global indexing.

  • Dirichlet boundary condition indexing and application.

  • Element integration rule selection (Gauss points).

  • Element-wise evaluation and vectorized mapping over elements.

  • Assembly of a global residual vector and sparse Jacobian matrix.

  • Consistent interfaces for energy-based and weighted-residual losses.

A derived class must implement ComputeElement(), which returns an element scalar contribution, an element residual vector, and an element Jacobian/tangent matrix. The scalar contribution may represent an element energy (energy-based losses) or a scalar constructed from the element residual and element DOFs (weighted-residual losses). The base class does not enforce a specific interpretation of the scalar contribution; it only provides consistent aggregation across the mesh.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. The following keys are required: "ordered_dofs", "element_type", "compute_dims", and "dirichlet_bc_dict". Optional keys include "num_gp", "loss_function_exponent", and "parametric_boundary_learning".

  • fe_mesh (Mesh) – Finite element mesh over which the loss is defined.

ApplyDirichletBCOnDofVector(full_dof_vector, load_increment=1.0)[source]#

Apply Dirichlet boundary conditions on a global DOF vector.

Parameters:

  • full_dof_vector (jax.numpy.ndarray) – Global DOF vector to be modified in-place (functional update).

  • load_increment (float, optional) – Scaling factor applied to the prescribed Dirichlet values. Default is 1.0.

Returns:

Global DOF vector with Dirichlet entries updated.

Return type:

jax.numpy.ndarray

ApplyDirichletBCOnElementResidualAndJacobian(elem_res, elem_jac, elem_BC_vec, elem_mask_BC_vec)[source]#

Apply Dirichlet boundary conditions to an element residual and Jacobian.

This method modifies the element residual and Jacobian so that Dirichlet DOFs are enforced at the element level. The inputs elem_BC_vec and elem_mask_BC_vec are diagonal vectors constructed from global BC vectors restricted to the element.

Parameters:

  • elem_res (jax.numpy.ndarray) – Element residual vector.

  • elem_jac (jax.numpy.ndarray) – Element Jacobian/tangent matrix.

  • elem_BC_vec (jax.numpy.ndarray) – Element BC diagonal vector used to enforce prescribed DOFs.

  • elem_mask_BC_vec (jax.numpy.ndarray) – Element mask diagonal vector used to preserve diagonal entries for constrained DOFs.

Returns:

Element residual and Jacobian after applying Dirichlet modifications.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

ComputeBatchLoss(batch_params, batch_dofs)[source]#

Compute the finite element loss over a batch of parameter samples and their associated solution fields.

This method evaluates the finite element loss for multiple input samples in parallel. Each sample consists of a parameter (control) vector and its corresponding primal solution (DOF) vector. For every sample, Dirichlet boundary conditions are applied to construct a full DOF vector, and the scalar element contributions returned by ComputeElement are summed over all elements to form the sample loss value.

If loss_function_exponent is specified in loss_settings, the summed scalar value is raised to that exponent for each sample. The returned batch loss is the mean of the per-sample loss values.

Parameters:

  • batch_params (jax.numpy.ndarray) – Batch of parameter or control vectors, one per sample. These vectors are mapped to element-level control fields using GetParametersVectors. The input is reshaped internally to (batch_size, -1).

  • batch_dofs (jax.numpy.ndarray) – Batch of primal solution vectors (DOF fields), one per sample. Dirichlet DOFs are inserted using GetFullDofVector. The input is reshaped internally to (batch_size, -1).

Returns:

A tuple containing the mean loss value across the batch and a tuple with the minimum, maximum, and mean per-sample loss values.

Return type:

Tuple[jax.numpy.ndarray, Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]]

abstract ComputeElement(elem_xyz, elem_controls, elem_dofs)[source]#

Compute element-level scalar, residual, and Jacobian contributions.

Derived classes must implement this method. The returned scalar is used by the base class to build element energies and batch losses, while the residual and Jacobian are used to assemble global vectors and matrices.

Parameters:

  • elem_xyz (jax.numpy.ndarray) – Element nodal coordinates.

  • elem_controls (jax.numpy.ndarray) – Element control variables (for example material field values) associated with the element nodes.

  • elem_dofs (jax.numpy.ndarray) – Element DOF vector arranged consistently with the element type and ordered degrees of freedom.

Returns:

  • Scalar element contribution (energy or scalar loss term).

  • Element residual vector.

  • Element Jacobian/tangent matrix.

Return type:

Tuple[float, jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementEnergy(elem_xyz, elem_controls, elem_dofs)[source]#

Convenience wrapper returning the scalar element contribution.

Parameters:

  • elem_xyz (jax.numpy.ndarray) – Element nodal coordinates.

  • elem_controls (jax.numpy.ndarray) – Element control variables.

  • elem_dofs (jax.numpy.ndarray) – Element DOF vector.

Returns:

Scalar element contribution returned by ComputeElement().

Return type:

float

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

Vmap-compatible wrapper to compute an element scalar contribution.

This helper extracts element nodal coordinates, controls, and DOFs from global arrays using element_id and then calls ComputeElementEnergy().

Parameters:

  • element_id (jax.numpy.integer) – Element index.

  • elements_nodes (jax.numpy.ndarray) – Connectivity array mapping element ids to node ids.

  • xyz (jax.numpy.ndarray) – Nodal coordinates array.

  • full_control_vector (jax.numpy.ndarray) – Global control/parameter values per node.

  • full_dof_vector (jax.numpy.ndarray) – Global DOF vector.

Returns:

Scalar element contribution.

Return type:

float

ComputeElementJacobianIndices(nodes_ids)[source]#

Compute global (row, col) index pairs for an element Jacobian block.

This method maps element node ids to global DOF indices and returns the 2D index pairs needed to assemble an element matrix into a global sparse matrix.

Parameters:

nodes_ids (jax.numpy.ndarray) – Element node ids.

Returns:

Array of shape (ndofs_elem * ndofs_elem, 2) containing global (row, col) index pairs for the element Jacobian block.

Return type:

jax.numpy.ndarray

ComputeElementResidualAndJacobian(elem_xyz, elem_controls, elem_dofs, elem_BC, elem_mask_BC, transpose_jac)[source]#

Vmap-compatible wrapper for residual/Jacobian evaluation of one element.

This helper extracts element nodal coordinates, controls, DOFs, and element-restricted BC vectors from global arrays and calls ComputeElementResidualAndJacobian().

Parameters:

  • element_id (jax.numpy.integer) – Element index.

  • elements_nodes (jax.numpy.ndarray) – Element connectivity array.

  • xyz (jax.numpy.ndarray) – Nodal coordinates.

  • full_control_vector (jax.numpy.ndarray) – Global control values per node.

  • full_dof_vector (jax.numpy.ndarray) – Global DOF vector.

  • full_dirichlet_BC_vec (jax.numpy.ndarray) – Global BC enforcement vector.

  • full_mask_dirichlet_BC_vec (jax.numpy.ndarray) – Global BC mask vector.

  • transpose_jac (bool) – If True, transpose element Jacobian before BC application.

Returns:

Element residual and Jacobian after BC application.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementsEnergies(total_control_vars, total_primal_vars)[source]#

Compute scalar element contributions for all elements.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Global control/parameter vector (nodal values).

  • total_primal_vars (jax.numpy.ndarray) – Global DOF vector (nodal unknowns).

Returns:

Array of per-element scalar contributions.

Return type:

jax.numpy.ndarray

ComputeJacobianMatrixAndResidualVector(total_control_vars, total_primal_vars, transpose_jacobian=False, new_implementation=False)[source]#

Assemble the global sparse Jacobian matrix and residual vector.

This method evaluates element residuals and Jacobians in batches, applies Dirichlet boundary conditions at element level, and assembles a global residual vector and a global sparse Jacobian matrix in BCOO format.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Global control/parameter vector (nodal values).

  • total_primal_vars (jax.numpy.ndarray) – Global DOF vector (nodal unknowns).

  • transpose_jacobian (bool, optional) – If True, element Jacobians are transposed before assembly. Default is False.

  • new_implementation (bool, optional) – Placeholder flag for alternative implementations. Default is False.

Returns:

  • Sparse global Jacobian matrix.

  • Global residual vector.

Return type:

Tuple[jax.experimental.sparse.BCOO, jax.numpy.ndarray]

ComputeTotalEnergy(total_control_vars, total_primal_vars)[source]#

Compute the total scalar loss value by summing element contributions.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Global control/parameter vector (nodal values).

  • total_primal_vars (jax.numpy.ndarray) – Global DOF vector (nodal unknowns).

Returns:

Total scalar loss value assembled over all elements.

Return type:

jax.numpy.ndarray

GetDOFs()[source]#

Return the ordered DOF names per node.

Returns:

DOF names in the order used by the loss.

Return type:

list

GetFullDofVector(known_dofs, unknown_dofs)[source]#

Construct a full DOF vector including Dirichlet values.

Parameters:

  • known_dofs (jax.numpy.ndarray) – Values used for Dirichlet DOFs when parametric boundary learning is enabled.

  • unknown_dofs (jax.numpy.ndarray) – DOF array that will be completed by inserting Dirichlet values.

Returns:

Full DOF vector with Dirichlet entries applied.

Return type:

jax.numpy.ndarray

GetNumberOfUnknowns()[source]#

Return the number of unknown DOFs (non-Dirichlet).

Returns:

Number of non-Dirichlet DOFs.

Return type:

int

GetParametersVectors(param_vector)[source]#

Map an input parameter vector to the control variables used by elements.

By default this is the identity mapping. Derived classes may override the mapping by setting self.get_param_function during initialization.

Parameters:

param_vector (jax.numpy.ndarray) – Input parameter vector.

Returns:

Control vector used by element routines.

Return type:

jax.numpy.ndarray

GetTotalNumberOfDOFs()[source]#

Return the total number of DOFs including Dirichlet DOFs.

Returns:

Total number of DOFs in the FE system.

Return type:

int

Initialize(reinitialize=False)[source]#

Initialize FE bookkeeping, integration rules, and batching configuration.

This method prepares all data required for fast element-wise evaluation and assembly. It is designed to be called once before evaluation but can be forced to re-run by setting reinitialize=True.

The initialization performs:

  • DOF counts and global indexing.

  • Construction of Dirichlet and non-Dirichlet index sets from loss_settings["dirichlet_bc_dict"].

  • Selection and configuration of the FE element implementation from element_type and Gauss integration settings from num_gp.

  • Definition of a full-DOF reconstruction function that applies Dirichlet values, optionally using parametric boundary learning.

  • Determination of an element batch size that evenly divides the number of elements to support efficient batched assembly.

  • Configuration of an exponent applied to the scalar loss aggregation.

Parameters:

reinitialize (bool, optional) – If True, re-run initialization even if the instance has already been initialized. Default is False.

Returns:

None

Raises:

ValueError – If compute_dims is not provided or if an unsupported Gauss integration rule is requested.

Linear mechanical loss functions#

Authors: Reza Najian Asl, RezaNajian Date: May, 2024 License: FOL/LICENSE

class fol.loss_functions.mechanical.MechanicalLoss(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Mechanical loss functional for linear elastic finite element problems.

This class defines a scalar loss functional for small-strain continuum mechanics problems. The total loss value represents a domain-integrated measure of mechanical equilibrium and is assembled by summing element-level contributions over all finite elements in the mesh.

For each element, the scalar contribution to the total loss is defined as the inner product of the element residual vector with the corresponding element degrees of freedom. The global loss is obtained by accumulating these scalar contributions across the computational domain, resulting in a weighted residual formulation.

Element-level stiffness matrices and force vectors are evaluated using Gaussian quadrature. The constitutive response is linear elastic and is defined by Young’s modulus and Poisson’s ratio provided in the loss settings. Both two-dimensional and three-dimensional problems are supported.

The loss formulation is compatible with adjoint-based sensitivity analysis and gradient-based optimization, as the scalar loss is constructed consistently from the element residuals and state variables.

Parameters:

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

  • loss_settings (dict) – Dictionary containing problem and material definitions. Must include material_dict with keys "young_modulus" and "poisson_ratio". May optionally include "body_foce" to specify a constant body force vector.

  • fe_mesh (Mesh) – Finite element mesh over which the loss is defined.

D#

Constitutive matrix in Voigt notation (3x3 in 2D, 6x6 in 3D).

Type:

jax.numpy.ndarray

body_force#

Body force vector with shape (2, 1) in 2D or (3, 1) in 3D.

Type:

jax.numpy.ndarray

CalculateNMatrix#

Dimension-dependent shape-function matrix constructor.

Type:

callable

CalculateBMatrix#

Dimension-dependent strain-displacement matrix constructor.

Type:

callable

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

Compute element-level contributions to the mechanical loss.

This method evaluates the element stiffness matrix and force vector using Gaussian quadrature. The element residual is defined as Se @ uvwe - Fe. The scalar element contribution represents this element’s contribution to the total scalar loss value of the problem.

The total loss for the mechanical problem is obtained by summing the scalar contributions of all elements over the computational domain. Each element contribution is computed as the inner product of the element residual vector with the element degrees of freedom, resulting in a weighted residual formulation in which the element DOFs act as weighting functions.

Parameters:

  • xyze – Element nodal coordinates.

  • de – Element parameter field values at nodes (e.g., density or heterogeneity coefficients).

  • uvwe – Element DOF vector arranged consistently with the element type and ordered degrees of freedom.

Returns:

  • Scalar element loss contribution defined as uvwe.T @ (Se @ uvwe - Fe).

  • Element residual vector Se @ uvwe - Fe.

  • Element stiffness matrix Se.

Return type:

Tuple[float, jax.numpy.ndarray, jax.numpy.ndarray]

Initialize()[source]#

Initialize FE bookkeeping, integration rules, and batching configuration.

This method prepares all data required for fast element-wise evaluation and assembly. It is designed to be called once before evaluation but can be forced to re-run by setting reinitialize=True.

The initialization performs:

  • DOF counts and global indexing.

  • Construction of Dirichlet and non-Dirichlet index sets from loss_settings["dirichlet_bc_dict"].

  • Selection and configuration of the FE element implementation from element_type and Gauss integration settings from num_gp.

  • Definition of a full-DOF reconstruction function that applies Dirichlet values, optionally using parametric boundary learning.

  • Determination of an element batch size that evenly divides the number of elements to support efficient batched assembly.

  • Configuration of an exponent applied to the scalar loss aggregation.

Parameters:

reinitialize (bool, optional) – If True, re-run initialization even if the instance has already been initialized. Default is False.

Returns:

None

Raises:

ValueError – If compute_dims is not provided or if an unsupported Gauss integration rule is requested.

class fol.loss_functions.mechanical.MechanicalLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: MechanicalLoss

Mechanical loss for 2D quadrilateral finite elements.

This class configures MechanicalLoss for two-dimensional problems discretized with quadrilateral elements. The displacement field consists of two components (Ux, Uy) per node, and the mechanical loss is assembled using the weighted residual formulation defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material parameters, optional body-force definitions, and optional integration settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical.MechanicalLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: MechanicalLoss

Mechanical loss for 2D triangular finite elements.

This class configures MechanicalLoss for two-dimensional problems discretized with triangular elements. The displacement field consists of two components (Ux, Uy) per node, and the mechanical loss is assembled using the weighted residual formulation defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material parameters and optional body-force definitions.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical.MechanicalLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: MechanicalLoss

Mechanical loss for 3D hexahedral finite elements.

This class configures MechanicalLoss for three-dimensional problems discretized with hexahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node, and the mechanical loss is assembled using the weighted residual formulation defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material parameters, optional body-force definitions, and optional integration settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical.MechanicalLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: MechanicalLoss

Mechanical loss for 3D tetrahedral finite elements.

This class configures MechanicalLoss for three-dimensional problems discretized with tetrahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node, and the mechanical loss is assembled using the weighted residual formulation defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material parameters and optional body-force definitions. Must include material_dict with elastic constants.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Saint Venant–Kirchhoff mechanical loss functions#

Authors: Kianoosh Taghikhani, kianoosh1989 Date: July, 2024 License: FOL/LICENSE

class fol.loss_functions.mechanical_saint_venant.SaintVenantMechanicalLoss(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Saint Venant–Kirchhoff mechanical energy loss for finite deformation elasticity.

This class defines an energy-based loss functional for mechanical problems governed by a Saint Venant–Kirchhoff constitutive model. The total loss value represents the total strain energy of the structure and is assembled by summing element-level energy contributions over all finite elements in the mesh.

For each element, the strain energy density is evaluated at Gauss points and integrated to obtain an element energy contribution. The global energy is obtained by accumulating these element energies across the computational domain.

In addition to the scalar energy contribution, this loss provides the element residual vector (internal minus external forces) and the element tangent matrix (material stiffness plus geometric stiffness) required for Newton-based solution procedures.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. Must include material_dict with keys "young_modulus" and "poisson_ratio". Optional entries include "body_foce" and parametric heterogeneity settings via "heterogeneity_field_name" and "heterogeneity_default_value". Element discretization settings (dimension, element type, ordered DOFs) are typically provided by specialized subclasses.

  • fe_mesh (Mesh) – Finite element mesh over which the energy functional is defined.

e#

Young’s modulus used by the constitutive model.

Type:

float

v#

Poisson’s ratio used by the constitutive model.

Type:

float

material_model#

Constitutive material model used to compute energy density, stress, and consistent tangent.

Type:

SaintVenant

body_force#

Body force vector with shape (2, 1) in 2D or (3, 1) in 3D.

Type:

jax.numpy.ndarray

CalculateHexaGeometricStiffness3D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a hexahedral element. :param DN_DX_T: (3, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (6,1), stress vector in Voigt notation at Gauss point

Returns:

(3*num_nodes, 3*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateQuadGeometricStiffness2D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a quadratic element. :param DN_DX_T: (2, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (3,1), stress vector in Voigt notation at Gauss point

Returns:

(2*num_nodes, 2*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateTetraGeometricStiffness3D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a tetra element. :param DN_DX_T: (3, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (6,1), stress vector in Voigt notation at Gauss point

Returns:

(3*num_nodes, 3*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateTriangleGeometricStiffness2D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a triangle element. :param DN_DX_T: (2, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (3,1), stress vector in Voigt notation at Gauss point

Returns:

(2*num_nodes, 2*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

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

Compute element-level energy, residual, and tangent contributions.

This method evaluates the element strain energy by Gaussian quadrature using the Saint Venant–Kirchhoff constitutive model. The returned scalar value represents the strain energy contribution of this element to the total energy of the system. The total loss value is obtained by summing these element energies over all elements in the mesh.

The method also computes the element residual vector as the difference between internal and external nodal force vectors and returns the element tangent matrix including both material and geometric stiffness contributions.

Parameters:

  • xyze – Element nodal coordinates.

  • de – Element parameter field values at nodes used to interpolate material parameters to Gauss points.

  • uvwe – Element displacement DOF vector arranged consistently with the element type and ordered degrees of freedom.

Returns:

  • Scalar element strain energy contribution.

  • Element residual vector defined as Fint - Fe.

  • Element tangent matrix including material and geometric stiffness contributions.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

class fol.loss_functions.mechanical_saint_venant.SaintVenantMechanicalLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: SaintVenantMechanicalLoss

Saint Venant–Kirchhoff mechanical energy loss for 2D quadrilateral elements.

This class configures SaintVenantMechanicalLoss for two-dimensional problems discretized with quadrilateral elements. The displacement field consists of two components (Ux, Uy) per node, and the total energy is assembled by summing element energy contributions defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_saint_venant.SaintVenantMechanicalLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: SaintVenantMechanicalLoss

Saint Venant–Kirchhoff mechanical energy loss for 2D triangular elements.

This class configures SaintVenantMechanicalLoss for two-dimensional problems discretized with triangular elements. The displacement field consists of two components (Ux, Uy) per node, and the total energy is assembled by summing element energy contributions defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_saint_venant.SaintVenantMechanicalLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: SaintVenantMechanicalLoss

Saint Venant–Kirchhoff mechanical energy loss for 3D hexahedral elements.

This class configures SaintVenantMechanicalLoss for three-dimensional problems discretized with hexahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node, and the total energy is assembled by summing element energy contributions defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_saint_venant.SaintVenantMechanicalLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: SaintVenantMechanicalLoss

Saint Venant–Kirchhoff mechanical energy loss for 3D tetrahedral elements.

This class configures SaintVenantMechanicalLoss for three-dimensional problems discretized with tetrahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node, and the total energy is assembled by summing element energy contributions defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Neo-Hookean mechanical loss functions#

Authors: Kianoosh Taghikhani, kianoosh1989 Date: July, 2024 License: FOL/LICENSE

class fol.loss_functions.mechanical_neohooke.NeoHookeMechanicalLoss(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Neo-Hookean mechanical energy loss for finite deformation elasticity.

This class defines an energy-based loss functional for hyperelastic Neo-Hookean materials under finite deformation. The total loss value represents the total strain energy of the structure and is assembled by summing element-level energy contributions over all finite elements in the mesh.

For each element, the strain energy density is evaluated at Gauss points using a Neo-Hookean constitutive model and integrated to obtain an element energy contribution. The global energy is obtained by accumulating these element energies across the computational domain.

In addition to the scalar energy contribution, the loss provides the element residual vector (internal minus external forces) and an element tangent matrix including material and geometric stiffness contributions. These outputs are intended for Newton-based solution procedures.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. Must include material_dict with keys "young_modulus" and "poisson_ratio". Optional entries may include "body_foce" specifying a constant body force vector. Element discretization settings (dimension, element type, ordered DOFs) are typically provided by specialized subclasses.

  • fe_mesh (Mesh) – Finite element mesh over which the energy functional is defined.

CalculateHexaGeometricStiffness3D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a hexahedral element. :param DN_DX_T: (3, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (6,1), stress vector in Voigt notation at Gauss point

Returns:

(3*num_nodes, 3*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateQuadGeometricStiffness2D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a quadratic element. :param DN_DX_T: (2, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (3,1), stress vector in Voigt notation at Gauss point

Returns:

(2*num_nodes, 2*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateTetraGeometricStiffness3D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a tetra element. :param DN_DX_T: (3, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (6,1), stress vector in Voigt notation at Gauss point

Returns:

(3*num_nodes, 3*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

CalculateTriangleGeometricStiffness2D(DN_DX_T, S)[source]#

Compute the geometric stiffness matrix for a triangle element. :param DN_DX_T: (2, num_nodes), shape function derivatives w.r.t spatial coordinates at Gauss point :param S: (3,1), stress vector in Voigt notation at Gauss point

Returns:

(2*num_nodes, 2*num_nodes), geometric stiffness matrix

Return type:

gp_geo_stiffness

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

Compute element-level energy, residual, and tangent contributions.

This method evaluates the element strain energy by Gaussian quadrature using a Neo-Hookean constitutive model. The returned scalar value represents the strain energy contribution of this element to the total energy of the system. The total loss value is obtained by summing these element energies over all elements in the mesh.

The method also computes the element residual vector as the difference between internal and external nodal force vectors and returns the element tangent matrix including both material and geometric stiffness contributions.

Parameters:

  • xyze – Element nodal coordinates.

  • de – Element parameter field values at nodes used to interpolate material parameters to Gauss points.

  • uvwe – Element displacement DOF vector arranged consistently with the element type and ordered degrees of freedom.

Returns:

  • Scalar element strain energy contribution.

  • Element residual vector defined as Fint - Fe.

  • Element tangent matrix including material and geometric stiffness contributions.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

class fol.loss_functions.mechanical_neohooke.NeoHookeMechanicalLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: NeoHookeMechanicalLoss

Neo-Hookean mechanical energy loss for 2D quadrilateral elements.

This class configures NeoHookeMechanicalLoss for two-dimensional problems discretized with quadrilateral elements. The displacement field consists of two components (Ux, Uy) per node.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_neohooke.NeoHookeMechanicalLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: NeoHookeMechanicalLoss

Neo-Hookean mechanical energy loss for 2D triangular elements.

This class configures NeoHookeMechanicalLoss for two-dimensional problems discretized with triangular elements. The displacement field consists of two components (Ux, Uy) per node.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_neohooke.NeoHookeMechanicalLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: NeoHookeMechanicalLoss

Neo-Hookean mechanical energy loss for 3D hexahedral elements.

This class configures NeoHookeMechanicalLoss for three-dimensional problems discretized with hexahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_neohooke.NeoHookeMechanicalLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: NeoHookeMechanicalLoss

Neo-Hookean mechanical energy loss for 3D tetrahedral elements.

This class configures NeoHookeMechanicalLoss for three-dimensional problems discretized with tetrahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Elastoplasticity loss functions#

Authors: Rishabh Arora, rishabharora236-cell Date: Oct, 2025 License: FOL/LICENSE

class fol.loss_functions.mechanical_elastoplasticity.ElastoplasticityLoss(name, loss_settings, fe_mesh)[source]#

Bases: MechanicalLoss

Elastoplastic mechanical loss with J2 plasticity and Gauss-point state updates.

This class extends fol.loss_functions.mechanical.MechanicalLoss to model small-strain elastoplasticity using a J2 plasticity constitutive law. The total loss value follows the weighted-residual formulation used by the base mechanical loss, while the element response additionally depends on a set of internal state variables stored at Gauss points.

For each element, the method ComputeElement() computes the element residual vector and a consistent element tangent matrix by automatic differentiation of the element residual with respect to the element DOFs. The Gauss-point state variables are updated during element evaluation and returned so they can be committed by the nonlinear solver upon convergence.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. Must include material_dict with keys "young_modulus", "poisson_ratio", "yield_limit", "iso_hardening_parameter_1", and "iso_hardening_param_2". Optional entries may include "body_foce" specifying a constant body force vector. Element discretization settings (dimension, element type, ordered DOFs) are typically provided by specialized subclasses.

  • fe_mesh (Mesh) – Finite element mesh over which the loss is defined.

ComputeElement(xyze, de, uvwe, element_state_gps)[source]#

Compute element scalar contribution, state update, residual, and tangent.

This method evaluates the element residual for an elastoplastic material with Gauss-point internal variables. Strains are computed from the element DOFs, stresses are obtained from the J2 plasticity model, and the internal nodal force vector is assembled by Gaussian quadrature. External body-force contributions are assembled in the same way.

The element residual is defined as internal minus external force. The element tangent matrix is computed using automatic differentiation of the element residual with respect to the element DOFs. Updated Gauss- point state variables are returned alongside the residual and tangent.

Parameters:

  • xyze – Element nodal coordinates.

  • de – Element control/parameter values at nodes. This argument is included for interface consistency.

  • uvwe – Element DOF vector arranged consistently with the element type and ordered degrees of freedom. Expected shape is (ndofs_elem, 1).

  • element_state_gps – Gauss-point state variables for this element. Expected shape is (n_gauss_points, n_state_vars).

Returns:

  • Scalar element loss contribution computed as uvwe.T @ r_e where r_e is the element residual.

  • Updated Gauss-point state variables with the same shape as the input state array.

  • Element residual vector r_e with shape (ndofs_elem, 1).

  • Element tangent matrix with shape (ndofs_elem, ndofs_elem).

Return type:

Tuple[float, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementResidualAndJacobian(elem_xyz, elem_controls, elem_dofs, elem_BC, elem_mask_BC, transpose_jac, elem_state_gps)[source]#

Compute element residual and Jacobian with Dirichlet boundary conditions.

This method computes the element residual and tangent matrix using ComputeElement(), optionally transposes the Jacobian, applies element-level Dirichlet boundary conditions, and returns the updated Gauss-point state variables.

Parameters:

  • elem_xyz (jax.numpy.ndarray) – Element nodal coordinates.

  • elem_controls (jax.numpy.ndarray) – Element control/parameter values at nodes.

  • elem_dofs (jax.numpy.ndarray) – Element DOF vector of shape (ndofs_elem, 1).

  • elem_BC (jax.numpy.ndarray) – Element Dirichlet boundary condition vector.

  • elem_mask_BC (jax.numpy.ndarray) – Element boundary condition mask.

  • transpose_jac (bool) – If True, the element Jacobian is transposed before applying boundary conditions.

  • elem_state_gps (jax.numpy.ndarray) – Gauss-point state variables for this element.

Returns:

Element residual vector and Jacobian matrix after applying Dirichlet boundary conditions, and the updated Gauss-point state variables.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

ComputeJacobianMatrixAndResidualVector(total_control_vars, total_primal_vars, old_state_gps, transpose_jacobian=False)[source]#

Assemble the global residual vector and sparse Jacobian with state update.

This method evaluates per-element residuals and tangents in batches, applies Dirichlet boundary conditions, and assembles a global residual vector and a sparse Jacobian matrix. Gauss-point state variables are updated element-wise and returned as a separate array.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Global control/parameter vector (nodal values).

  • total_primal_vars (jax.numpy.ndarray) – Global DOF vector (nodal unknowns) used to build element DOFs.

  • old_state_gps (jax.numpy.ndarray) – Gauss-point state variables for all elements. Expected shape is (n_elements, n_gauss_points, n_state_vars).

  • transpose_jacobian (bool, optional) – If True, element Jacobians are transposed before assembly. Default is False.

Returns:

  • Updated Gauss-point state variables with the same shape as old_state_gps.

  • Sparse global Jacobian matrix.

  • Global residual vector.

Return type:

Tuple[jax.numpy.ndarray, jax.experimental.sparse.BCOO, jax.numpy.ndarray]

class fol.loss_functions.mechanical_elastoplasticity.ElastoplasticityLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: ElastoplasticityLoss

Elastoplasticity loss for 2D quadrilateral elements.

This class configures ElastoplasticityLoss for two-dimensional problems discretized with quadrilateral elements. The displacement field consists of two components (Ux, Uy) per node.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.mechanical_elastoplasticity.ElastoplasticityLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: ElastoplasticityLoss

Elastoplasticity loss for 3D tetrahedral elements.

This class configures ElastoplasticityLoss for three-dimensional problems discretized with tetrahedral elements. The displacement field consists of three components (Ux, Uy, Uz) per node.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Kratos small-displacement mechanical loss functions#

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

class fol.loss_functions.kratos_small_displacement.KratosSmallDisplacement3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Kratos-based 3D small-displacement tetrahedral loss accelerated via JAX FFI.

This class provides a finite element loss for linear small-displacement elasticity in 3D using tetrahedral elements. It derives from fol.loss_functions.fe_loss.FiniteElementLoss and reuses the base class infrastructure for DOF bookkeeping, Dirichlet boundary conditions, and sparse assembly.

In contrast to standard Python/JAX element implementations, this class delegates element-level computations to custom CUDA kernels exposed through fol_ffi_functions and registered with jax.ffi. The kernels compute element residuals and element stiffness matrices, which are then processed to enforce Dirichlet boundary conditions and assembled into a global sparse Jacobian matrix and residual vector.

The total loss value is computed in an energy-consistent manner using the dot product of the global displacement vector with the global nodal residual vector produced by the FFI kernel.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. User-provided settings may include "material_dict" with keys "poisson_ratio" and "young_modulus" and the standard FE settings expected by the base class. The element discretization is fixed to 3D tetrahedra with displacement DOFs.

  • fe_mesh (Mesh) – Finite element mesh containing node coordinates and tetrahedral connectivity.

Raises:

RuntimeError – If fol_ffi_functions is not available.

Notes

This class relies on CUDA FFI targets registered at import time. It is expected to run on platforms where the registered targets are available. Element-level routines are not implemented in Python for this class.

ComputeJacobianMatrixAndResidualVector(total_control_vars, total_primal_vars, transpose_jacobian=False)[source]#

Assemble the global sparse Jacobian matrix and residual vector using FFI.

This method calls the FFI kernel compute_elements to compute element stiffness matrices and element residual vectors for all tetrahedral elements. Dirichlet boundary conditions are applied at element level using the base class masking strategy, and the processed contributions are assembled into:

  • a global residual vector, and

  • a global sparse Jacobian matrix in BCOO format.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Control variables. This argument is accepted for API compatibility and is not used by the current FFI kernels.

  • total_primal_vars (jax.numpy.ndarray) – Global displacement (DOF) vector.

  • transpose_jacobian (bool, optional) – If True, element stiffness matrices are transposed before applying boundary conditions and assembly. Default is False.

Returns:

A tuple containing the global sparse Jacobian matrix and the global residual vector.

Return type:

Tuple[jax.experimental.sparse.BCOO, jax.numpy.ndarray]

ComputeTotalEnergy(total_control_vars, total_primal_vars)[source]#

Compute the total energy-consistent scalar loss value for the system.

This method calls the FFI kernel compute_nodal_residuals to evaluate the global nodal residual vector for the provided displacement field. The returned scalar value is computed as the dot product of the global displacement vector with the global residual vector.

Parameters:

  • total_control_vars (jax.numpy.ndarray) – Control variables. This argument is accepted for API compatibility and is not used by the current FFI kernels.

  • total_primal_vars (jax.numpy.ndarray) – Global displacement (DOF) vector with shape (total_number_of_dofs,) or compatible. The vector is reshaped internally to a batch of size 1.

Returns:

Scalar loss value computed from the displacement field and the nodal residuals.

Return type:

jax.numpy.ndarray

Initialize(reinitialize=False)[source]#

Initialize FE bookkeeping and material parameters.

This method calls the base class initialization and then prepares the material parameters used by the CUDA kernels. If material_dict is provided in loss_settings, it is used directly; otherwise default values are applied.

Parameters:

reinitialize (bool, optional) – If True, re-run initialization even if already initialized. Default is False.

Returns:

None

Nonlinear thermal loss functions#

Authors: Reza Najian Asl, RezaNajian Date: June, 2024 License: FOL/LICENSE

class fol.loss_functions.thermal.ThermalLoss(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Thermal loss functional for steady heat conduction finite element problems.

This class defines a scalar loss functional for heat conduction problems solved with finite elements. The total loss value represents a domain-integrated measure of thermal equilibrium and is assembled by summing element-level contributions over all finite elements in the mesh.

For each element, the scalar contribution to the total loss is defined as the inner product of the element residual vector with the corresponding element temperature degrees of freedom. The global loss is obtained by accumulating these scalar contributions across the computational domain, resulting in a weighted residual formulation.

Element-level stiffness matrices and source vectors are evaluated using Gaussian quadrature. The element conductivity can be nonlinear in the temperature field according to the parameters beta and c in the loss settings, using a conductivity of the form:

k(T) = d(x) * (1 + beta * T^c)

where d(x) is the element parameter field interpolated from de at Gauss points.

Parameters:

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

  • loss_settings (dict) – Dictionary containing problem settings and optional nonlinear conductivity parameters. Supported keys include "beta" and "c" for the temperature-dependent conductivity factor. The element discretization settings (dimension, element type, and ordered DOFs) are typically provided by the specialized subclasses.

  • fe_mesh (Mesh) – Finite element mesh over which the loss is defined.

thermal_loss_settings#

Dictionary containing nonlinear conductivity parameters with keys "beta" and "c".

Type:

dict

ComputeElement(xyze, de, te, body_force=0)[source]#

Compute element-level contributions to the thermal loss.

This method evaluates the element stiffness matrix and source vector using Gaussian quadrature. The element residual is defined as Se @ te - Fe. The scalar element contribution represents this element’s contribution to the total scalar loss value of the problem.

The total loss for the thermal problem is obtained by summing the scalar contributions of all elements over the computational domain. Each element contribution is computed as the inner product of the element residual vector with the element temperature degrees of freedom, resulting in a weighted residual formulation.

Parameters:

  • xyze – Element nodal coordinates.

  • de – Element parameter field values at nodes used to construct the conductivity field (interpolated to Gauss points).

  • te – Element temperature DOFs arranged consistently with the element type and ordered DOFs.

  • body_force – Element volumetric heat source term. This value is treated as a constant over the element and integrated against the shape functions. Default is 0.

Returns:

  • Scalar element loss contribution defined as te.T @ (Se @ te - Fe).

  • Element residual vector Se @ te - Fe.

  • Element stiffness matrix Se.

Return type:

Tuple[float, jax.numpy.ndarray, jax.numpy.ndarray]

Initialize(reinitialize=False)[source]#

Initialize FE bookkeeping, integration rules, and batching configuration.

This method prepares all data required for fast element-wise evaluation and assembly. It is designed to be called once before evaluation but can be forced to re-run by setting reinitialize=True.

The initialization performs:

  • DOF counts and global indexing.

  • Construction of Dirichlet and non-Dirichlet index sets from loss_settings["dirichlet_bc_dict"].

  • Selection and configuration of the FE element implementation from element_type and Gauss integration settings from num_gp.

  • Definition of a full-DOF reconstruction function that applies Dirichlet values, optionally using parametric boundary learning.

  • Determination of an element batch size that evenly divides the number of elements to support efficient batched assembly.

  • Configuration of an exponent applied to the scalar loss aggregation.

Parameters:

reinitialize (bool, optional) – If True, re-run initialization even if the instance has already been initialized. Default is False.

Returns:

None

Raises:

ValueError – If compute_dims is not provided or if an unsupported Gauss integration rule is requested.

class fol.loss_functions.thermal.ThermalLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: ThermalLoss

Thermal loss for 2D quadrilateral finite elements.

This class configures ThermalLoss for two-dimensional problems discretized with quadrilateral elements. The temperature field has a single DOF per node (T), and the thermal loss is assembled using the weighted residual formulation defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing thermal settings, optional nonlinear conductivity parameters, and optional integration settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.thermal.ThermalLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: ThermalLoss

Thermal loss for 2D triangular finite elements.

This class configures ThermalLoss for two-dimensional problems discretized with triangular elements. The temperature field has a single DOF per node (T), and the thermal loss is assembled using the weighted residual formulation defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing thermal settings and optional nonlinear conductivity parameters (e.g., "beta" and "c").

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.thermal.ThermalLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: ThermalLoss

Thermal loss for 3D hexahedral finite elements.

This class configures ThermalLoss for three-dimensional problems discretized with hexahedral elements. The temperature field has a single DOF per node (T), and the thermal loss is assembled using the weighted residual formulation defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing thermal settings, optional nonlinear conductivity parameters, and optional integration settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.thermal.ThermalLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: ThermalLoss

Thermal loss for 3D tetrahedral finite elements.

This class configures ThermalLoss for three-dimensional problems discretized with tetrahedral elements. The temperature field has a single DOF per node (T), and the thermal loss is assembled using the weighted residual formulation defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing thermal settings and optional nonlinear conductivity parameters (e.g., "beta" and "c").

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Transient thermal loss functions#

Authors: Yusuke Yamazaki; Reza Najian Asl (RezaNajian) Date: May, 2025 License: FOL/LICENSE

class fol.loss_functions.transient_thermal.TransientThermalLoss(name, loss_settings, fe_mesh)[source]#

Bases: ThermalLoss

Transient thermal energy functional with implicit time integration.

This class defines an energy-based loss functional for transient heat conduction problems. The total loss value represents the total discrete thermal energy of the system at a time step and is assembled by summing element-level energy contributions over all finite elements in the mesh.

For each element, a scalar energy contribution is computed using Gaussian quadrature based on the current and next nodal temperatures. The global energy is obtained by accumulating these element energies across the computational domain.

In addition to the scalar energy contribution, the class provides the element residual vector and Jacobian matrix associated with the discrete system used for implicit time stepping. The default time integration method is implicit Euler and is configured through loss_settings["time_integration_dict"].

The thermal conductivity may depend on temperature through the parameters beta and c and through a nodal heterogeneity field k0.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material parameters and time integration settings. Expected entries include a "material_dict" with material properties and a "time_integration_dict" specifying the time step.

  • fe_mesh (Mesh) – Finite element mesh over which the energy functional is defined.

ComputeElement(xyze, Te_c, Te_n, Ke)[source]#

Compute element-level energy, residual, and Jacobian contributions.

This method evaluates the discrete transient thermal energy contribution of a single finite element using Gaussian quadrature. The returned scalar value represents the energy contribution of this element to the total system energy. The total energy is obtained by summing element energies over all elements in the mesh.

The element residual vector and Jacobian matrix correspond to the discrete system used for implicit time integration. The residual has the form of a backward-Euler update using the element mass and conductivity operators constructed from the current and next temperatures.

Parameters:

  • xyze – Element nodal coordinates.

  • Te_c – Element nodal temperatures at the current (previous) time step. Expected shape is compatible with the element node count.

  • Te_n – Element nodal temperatures at the next (unknown) time step. Expected shape is compatible with the element node count.

  • Ke – Element nodal conductivity/heterogeneity field (e.g., nodal values of k0) used to interpolate conductivity to Gauss points.

Returns:

  • Scalar element thermal energy contribution.

  • Element residual vector for the time-discrete system.

  • Element Jacobian matrix associated with the residual.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementEnergy(elem_xyz, elem_current_temps, elem_next_temps, elem_heterogeneity)[source]#

Convenience wrapper returning the scalar element contribution.

Parameters:

  • elem_xyz (jax.numpy.ndarray) – Element nodal coordinates.

  • elem_controls (jax.numpy.ndarray) – Element control variables.

  • elem_dofs (jax.numpy.ndarray) – Element DOF vector.

Returns:

Scalar element contribution returned by ComputeElement().

Return type:

float

ComputeElementEnergyVmapCompatible(element_id, elements_nodes, xyz, nodal_current_temps, nodal_next_temps, nodal_heterogeneity)[source]#

Vmap-compatible wrapper to compute an element scalar contribution.

This helper extracts element nodal coordinates, controls, and DOFs from global arrays using element_id and then calls ComputeElementEnergy().

Parameters:

  • element_id (jax.numpy.integer) – Element index.

  • elements_nodes (jax.numpy.ndarray) – Connectivity array mapping element ids to node ids.

  • xyz (jax.numpy.ndarray) – Nodal coordinates array.

  • full_control_vector (jax.numpy.ndarray) – Global control/parameter values per node.

  • full_dof_vector (jax.numpy.ndarray) – Global DOF vector.

Returns:

Scalar element contribution.

Return type:

float

ComputeElementResidualAndJacobian(elem_xyz, elem_current_temps, elem_next_temps, elem_heterogeneity, elem_BC, elem_mask_BC, transpose_jac)[source]#

Compute element residual and Jacobian with Dirichlet boundary conditions.

This method computes the element residual and Jacobian from ComputeElement(), optionally transposes the Jacobian, and applies element-level Dirichlet boundary conditions using the provided boundary condition vectors and masks.

Parameters:

  • elem_xyz (jax.numpy.ndarray) – Element nodal coordinates.

  • elem_current_temps (jax.numpy.ndarray) – Element nodal temperatures at the current time step.

  • elem_next_temps (jax.numpy.ndarray) – Element nodal temperatures at the next time step.

  • elem_heterogeneity (jax.numpy.ndarray) – Element nodal heterogeneity/conductivity field.

  • elem_BC (jax.numpy.ndarray) – Element Dirichlet boundary condition vector.

  • elem_mask_BC (jax.numpy.ndarray) – Element boundary condition mask.

  • transpose_jac (bool) – If True, the element Jacobian is transposed before applying boundary conditions.

Returns:

Element residual vector and Jacobian matrix after applying Dirichlet boundary conditions.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementResidualAndJacobianVmapCompatible(element_id, elements_nodes, xyz, nodal_current_temps, nodal_next_temps, full_dirichlet_BC_vec, full_mask_dirichlet_BC_vec, transpose_jac)[source]#

VMAP-compatible wrapper to compute element residual and Jacobian by element id.

Parameters:

  • element_id (jax.numpy.ndarray) – Element index.

  • elements_nodes (jax.numpy.ndarray) – Element connectivity array.

  • xyz (jax.numpy.ndarray) – Nodal coordinate array.

  • nodal_current_temps (jax.numpy.ndarray) – Current-step nodal temperatures.

  • nodal_next_temps (jax.numpy.ndarray) – Next-step nodal temperatures.

  • full_dirichlet_BC_vec (jax.numpy.ndarray) – Global Dirichlet boundary condition vector.

  • full_mask_dirichlet_BC_vec (jax.numpy.ndarray) – Global boundary condition mask vector.

  • transpose_jac (bool) – If True, the element Jacobian is transposed before applying boundary conditions.

Returns:

Element residual vector and Jacobian matrix after applying Dirichlet boundary conditions.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

ComputeElementsEnergies(nodal_current_temps, nodal_next_temps)[source]#

Compute element energy contributions for all elements in parallel.

Parameters:

  • nodal_current_temps (jax.numpy.ndarray) – Current-step nodal temperatures.

  • nodal_next_temps (jax.numpy.ndarray) – Next-step nodal temperatures.

Returns:

Array of per-element energy contributions.

Return type:

jax.numpy.ndarray

Initialize(reinitialize=False)[source]#

Initialize material and time-integration settings.

This method initializes the base thermal loss and then configures material parameters and time integration settings. If the instance is already initialized and reinitialize is False, the method returns without modifying the current configuration.

Parameters:

reinitialize (bool, optional) – If True, forces reinitialization even if the object is already initialized. Default is False.

Raises:

  • ValueError – If the provided nodal conductivity field k0 does not match the number of mesh nodes.

  • ValueError – If time_integration_dict["time_step"] is not provided.

class fol.loss_functions.transient_thermal.TransientThermalLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: TransientThermalLoss

Transient thermal loss for 2D quadrilateral finite elements.

This class configures TransientThermalLoss for two-dimensional problems discretized with quadrilateral elements. The temperature field has a single DOF per node (T), and the total energy is assembled by summing element energy contributions defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and time_integration_dict.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.transient_thermal.TransientThermalLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: TransientThermalLoss

Transient thermal loss for 2D triangular finite elements.

This class configures TransientThermalLoss for two-dimensional problems discretized with triangular elements. The temperature field has a single DOF per node (T), and the total energy is assembled by summing element energy contributions defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and time_integration_dict.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.transient_thermal.TransientThermalLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: TransientThermalLoss

Transient thermal loss for 3D hexahedral finite elements.

This class configures TransientThermalLoss for three-dimensional problems discretized with hexahedral elements. The temperature field has a single DOF per node (T), and the total energy is assembled by summing element energy contributions defined in the base class.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and time_integration_dict.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.transient_thermal.TransientThermalLoss3DTetra(name, loss_settings, fe_mesh)[source]#

Bases: TransientThermalLoss

Transient thermal loss for 3D tetrahedral finite elements.

This class configures TransientThermalLoss for three-dimensional problems discretized with tetrahedral elements. The temperature field has a single DOF per node (T), and the total energy is assembled by summing element energy contributions defined in the base class.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and time_integration_dict.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Allen–Cahn phase-field loss functions#

Authors: Yusuke Yamazaki; Reza Najian Asl (RezaNajian) Date: April, 2024 License: FOL/LICENSE

class fol.loss_functions.phase_field.AllenCahnLoss(name, loss_settings, fe_mesh)[source]#

Bases: FiniteElementLoss

Allen–Cahn phase-field energy functional with implicit time discretization.

This class defines an energy-based loss functional for Allen–Cahn type phase-field evolution problems. The total loss value represents the total discrete free energy of the system and is assembled by summing element-level energy contributions over all finite elements in the mesh.

For each element, a scalar energy contribution is computed using Gaussian quadrature based on the current and previous phase-field values. The total energy of the system is obtained by accumulating these element energies across the computational domain.

In addition to the scalar energy contribution, the class provides the element residual vector and Jacobian matrix corresponding to the discrete nonlinear system derived from the energy functional. These quantities are used for implicit time integration and Newton-based solution procedures.

Material and time-integration parameters are provided via loss_settings["material_dict"]. The interface width is controlled by epsilon, and the time step size is given by dt. The formulation supports both two-dimensional and three-dimensional problems.

Parameters:

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

  • loss_settings (dict) – Dictionary containing problem settings. Must include material_dict with keys "rho", "cp", "dt", and "epsilon". Element discretization settings (dimension, element type, ordered DOFs) are typically provided by the specialized subclasses.

  • fe_mesh (Mesh) – Finite element mesh over which the energy functional is defined.

rho#

Density parameter from material_dict.

Type:

float

cp#

Heat-capacity-like parameter from material_dict.

Type:

float

dt#

Time step size used for implicit time discretization.

Type:

float

epsilon#

Interface-width parameter controlling the phase-field thickness.

Type:

float

body_force#

Body force vector with shape (2, 1) in 2D or (3, 1) in 3D.

Type:

jax.numpy.ndarray

ComputeElement(xyze, phi_e_c, phi_e_n, body_force=0)[source]#

Compute element-level energy, residual, and Jacobian contributions.

This method evaluates the discrete Allen–Cahn free energy contribution of a single finite element using Gaussian quadrature. The returned scalar value represents the energy contribution of this element to the total system energy. The total energy of the problem is obtained by summing these element energy contributions over all elements in the mesh.

In addition to the scalar energy, the method computes the element residual vector and the element Jacobian matrix corresponding to the discrete nonlinear system derived from the energy functional. These quantities are used in implicit time integration and Newton-based solution procedures.

The formulation depends on the phase-field values at the current time step and the previous time step. The time-discrete contribution enforces temporal consistency through the time step size dt.

Parameters:

  • xyze – Element nodal coordinates.

  • phi_e_c – Element phase-field degrees of freedom at the previous time step (committed state). Expected shape is compatible with the element node count.

  • phi_e_n – Element phase-field degrees of freedom at the current time step (unknown state). Expected shape is compatible with the element node count.

  • body_force – Optional element source term. This argument is included for extensibility and defaults to 0.

Returns:

  • Scalar element free energy contribution.

  • Element residual vector for the discrete Allen–Cahn system.

  • Element Jacobian matrix associated with the residual.

Return type:

Tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]

Raises:

ValueError – If invalid numerical values are encountered during the evaluation of the element energy or residual.

Initialize()[source]#

Initialize FE bookkeeping, integration rules, and batching configuration.

This method prepares all data required for fast element-wise evaluation and assembly. It is designed to be called once before evaluation but can be forced to re-run by setting reinitialize=True.

The initialization performs:

  • DOF counts and global indexing.

  • Construction of Dirichlet and non-Dirichlet index sets from loss_settings["dirichlet_bc_dict"].

  • Selection and configuration of the FE element implementation from element_type and Gauss integration settings from num_gp.

  • Definition of a full-DOF reconstruction function that applies Dirichlet values, optionally using parametric boundary learning.

  • Determination of an element batch size that evenly divides the number of elements to support efficient batched assembly.

  • Configuration of an exponent applied to the scalar loss aggregation.

Parameters:

reinitialize (bool, optional) – If True, re-run initialization even if the instance has already been initialized. Default is False.

Returns:

None

Raises:

ValueError – If compute_dims is not provided or if an unsupported Gauss integration rule is requested.

class fol.loss_functions.phase_field.AllenCahnLoss2DQuad(name, loss_settings, fe_mesh)[source]#

Bases: AllenCahnLoss

Allen–Cahn loss for 2D quadrilateral finite elements.

This class configures AllenCahnLoss for two-dimensional problems discretized with quadrilateral elements. The phase-field has a single DOF per node (Phi), and element contributions are assembled using the base formulation.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.phase_field.AllenCahnLoss2DTri(name, loss_settings, fe_mesh)[source]#

Bases: AllenCahnLoss

Allen–Cahn loss for 2D triangular finite elements.

This class configures AllenCahnLoss for two-dimensional problems discretized with triangular elements. The phase-field has a single DOF per node (Phi), and element contributions are assembled using the base formulation.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

class fol.loss_functions.phase_field.AllenCahnLoss3DHexa(name, loss_settings, fe_mesh)[source]#

Bases: AllenCahnLoss

Allen–Cahn loss for 3D hexahedral finite elements.

This class configures AllenCahnLoss for three-dimensional problems discretized with hexahedral elements. The phase-field has a single DOF per node (Phi), and element contributions are assembled using the base formulation.

If the number of Gauss points is not specified in the loss settings, a default value of num_gp = 2 is used.

Parameters:

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

  • loss_settings (dict) – Dictionary containing material_dict and optional settings such as integration parameters.

  • fe_mesh (Mesh) – Finite element mesh associated with the loss.

Regression loss function#

Authors: Reza Najian Asl, RezaNajian Date: November, 2024 License: FOL/LICENSE

class fol.loss_functions.regression_loss.RegressionLoss(name, loss_settings, fe_mesh)[source]#

Bases: Loss

Regression loss for supervised learning on mesh nodal quantities.

This class implements a mean-squared-error (MSE) regression loss between ground-truth values and predicted values. It is designed to integrate with the FoLax loss interface and to work with a finite element mesh context, where nodal unknowns (DOFs) define the structure of the regression targets.

The loss is computed over batches of samples by flattening the input arrays to (batch_size, -1) and evaluating the element-wise squared error. The returned scalar is the mean of the squared error over the full batch and all output components. A simple error summary (min, max, mean) is also returned for monitoring.

Parameters:

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

  • loss_settings (dict) – Configuration dictionary. Must include "nodal_unknows" defining the nodal degrees of freedom used by this loss (key spelling preserved from the current implementation).

  • fe_mesh (Mesh) – Finite element mesh associated with the nodal unknown structure.

loss_settings#

Loss configuration dictionary.

Type:

dict

fe_mesh#

Finite element mesh used for sizing/indexing.

Type:

Mesh

dofs#

Nodal degrees of freedom as provided by loss_settings["nodal_unknows"].

non_dirichlet_indices#

Indices of DOFs treated as trainable/unknown in this loss. For this regression loss, these are initialized as a full range over all nodal DOFs.

Type:

jax.numpy.ndarray

ComputeBatchLoss(gt_values, pred_values)[source]#

Compute batch regression loss and error summary.

The batch loss is computed as the mean squared error between ground truth and predictions. Inputs are converted to at least 2D arrays and flattened to shape (batch_size, -1) prior to evaluation.

Parameters:

  • gt_values (jax.numpy.ndarray) – Ground-truth values. Any additional trailing dimensions are flattened per sample.

  • pred_values (jax.numpy.ndarray) – Predicted values. Must be broadcast-compatible with gt_values after flattening.

Returns:

  • Mean squared error over the batch and all components.

  • Error summary (min_error, max_error, mean_error) computed from the element-wise squared errors.

Return type:

Tuple[float, Tuple[float, float, float]]

Initialize(reinitialize=False)[source]#

Initialize indices and internal state for loss evaluation.

This method sets up indexing for nodal DOFs based on the number of mesh nodes and the configured DOFs. If the instance is already initialized and reinitialize is False, the method returns without modifying the current configuration.

Parameters:

reinitialize (bool, optional) – If True, forces reinitialization even if the object is already initialized. Default is False.

Contents