FEAX API Reference

Complete API documentation for FEAX (Finite Element Analysis with JAX).

Core Modules

Problem Definition

• feax.problem - Base Problem class and problem definition
• feax.mesh - Mesh generation and handling

Finite Elements

• feax.fe - FiniteElement class definitions
• feax.basis - Shape functions and quadrature rules

Assembly & Solving

• feax.assembler - System matrix and residual assembly
• feax.solver - High-level solver creation (create_solver)
• feax.solvers - Solver infrastructure (options, linear, Newton, AMG, reduced)
• feax.solvers.direct - CPU host-side direct backends: spsolve, umfpack, cholmod
• feax.solvers.time_solver - Time stepping (TimePipeline, ImplicitPipeline, ExplicitPipeline, TimeConfig, run)
• feax.solvers.eigen - Linear buckling (create_linear_buckling_solver; backends "sparse" (ARPACK, default), "dense", "matfree", and "cudss" GPU shift-invert)
• feax.solution - Solution return type (.dofs, .field(i), .node_var(component=))

Boundary Conditions

• feax.DCboundary - Dirichlet boundary conditions

Internal Variables

• feax.traced_params - Material properties and parameters

Utilities

• feax.utils - Utility functions (save_sol, XDMFWriter, zero_like_initial_guess)

Mechanics

Constitutive models for finite element simulations:

• feax.mechanics.cohesive - Cohesive zone models for fracture mechanics
• feax.mechanics.shell - Mindlin plates and classical lamination theory
• feax.mechanics.tmc - Third-medium contact
• feax.mechanics.orientation - Fiber orientation fields and rotated stiffness
• feax.mechanics.layered_solid_element - Layered solid elements for laminates
• feax.mechanics.copv - Mesh and model generation for double-wall composite pressure vessels (COPV)

Topology Optimization (gene)

Density-based topology optimization with adaptive remeshing:

• feax.gene.optimizer - MMA pipeline driver (Pipeline, @constraint, run, Continuation, AdaptiveConfig)
• feax.gene.adaptive - Adaptive remeshing and field transfer
• feax.gene.filters - Density filters and Heaviside projection
• feax.gene.responses - Compliance and volume functions
• feax.gene.mdmm - Constrained optimization (MDMM)
• feax.gene.length_scale - Minimum length-scale constraints (geometric and morphological)
• feax.gene.mesh_extraction - Density / SDF → surface and volume mesh extraction (volume meshes need the gmsh extra)

Experimental

APIs that are usable but may change without a deprecation cycle:

• feax.experimental.cmg - Matrix-free O(band) geometric multigrid on a StructuredGrid band (NarrowBandCMG, ScalarMaterial, KernelMaterial, build_geometry)
• feax.experimental.narrowband - Moving-band OC loop and multires bootstrap (oc_update, run_narrowband_oc, run_narrowband_multires, create_sensitivity_filter)
• feax.experimental.mma - Pure-JAX MMA / GCMMA (MMA, GCMMA)

Flat Toolkit

Lattice material analysis and homogenization:

• feax.flat.graph - Graph-based density field generation
• feax.flat.pbc - Periodic boundary conditions
• feax.flat.solver - Homogenization solvers
• feax.flat.unitcell - Unit cell definitions
• feax.flat.utils - Visualization utilities

Browse All

Use the sidebar to explore all classes, functions, and modules in detail.