Core#

Spaces#

For new public-facing mixed setups, prefer MixedSpaces(...).to_fe_space(). MixedFESpace remains available as a lower-level constructor for advanced or internal usage.

class fluxfem.core.FESpaceBase(*args, **kwargs)#

Protocol for FE space objects used by assembly.

This defines the minimal interface required by the core assembly routines: element-to-DOF connectivity, value dimension, and the ability to build per-element FormContext objects (test/trial fields plus quadrature data).

fluxfem.core.FESpace#: alias of FESpaceClosure

class fluxfem.core.FESpacePytree(mesh: BaseMesh, basis: Basis3D, elem_dofs: jnp.ndarray, value_dim: int = 1, _n_dofs: int | None = None, _n_ldofs: int | None = None, data: SpaceData | None = None, _pattern_cache: PatternCache = <factory>, _kernel_cache: KernelCache = <factory>, _form_context_cache: dict[ContextCacheKey, FormContext]=<factory>, _elem_rows_cache: jnp.ndarray | None = None)#

FESpaceClosure with JAX pytree support.

Use this when a space must be carried through JAX transformations (jit/vmap), or stored inside other pytrees. Only mesh, basis, and elem_dofs are treated as children; metadata is preserved as auxiliary data.

class fluxfem.core.MixedFESpace(fields: dict[str, FESpaceClosure], field_order: Sequence[str] | None = None, field_to_space_key: Mapping[str, str] | None = None, field_alias_space_keys: Mapping[str, Sequence[str]] | None = None)#

Mixed FE space composed of multiple scalar/vector spaces.

Field DOFs are concatenated in field order:: [field0 dofs | field1 dofs | …]

class fluxfem.core.MixedProblem(space: MixedFESpace, residuals: dict[str, Callable] | MixedWeakForm, params: object | None = None, pattern: object | None = None, n_chunks: int | None = None, pad_trace: bool = False)#: Lightweight wrapper for mixed residual assembly with cached compilation.

class fluxfem.core.MixedBlockSystem(mixed: 'MixedFESpace', K: 'object', F: 'object', free_dofs: 'np.ndarray', dirichlet: 'MixedDirichletBC')#

fluxfem.core.make_space(mesh: BaseMeshClosure, basis: Basis3D, element: ElementVector | None = None) → FESpaceClosure#

Build an FE space from a mesh and basis.

element=None → scalar dof per node (elem_dofs = mesh.conn), value_dim=1 element=ElementVector(dim) → vector dof per node, value_dim=dim

fluxfem.core.make_space_pytree(mesh: jax.tree_util.register_pytree_node_class, basis: Basis3D, element: ElementVector | None = None) → jax.tree_util.register_pytree_node_class#: Build a pytree-compatible FE space.

fluxfem.core.make_hex_space(mesh: HexMesh, dim: int = 1, intorder: int = 2) → FESpaceClosure#: Create a trilinear hex space (8-node elements).

fluxfem.core.make_hex_space_pytree(mesh: HexMesh, dim: int = 1, intorder: int = 2) → jax.tree_util.register_pytree_node_class#: Create a pytree trilinear hex space (8-node elements).

fluxfem.core.make_tet_space(mesh: TetMesh, dim: int = 1, intorder: int = 2) → FESpaceClosure#: Create a linear or quadratic tet space based on mesh nodes.

fluxfem.core.make_tet_space_pytree(mesh: TetMesh, dim: int = 1, intorder: int = 2) → jax.tree_util.register_pytree_node_class#: Create a pytree linear or quadratic tet space based on mesh nodes.

Forms#

class fluxfem.core.FormContext(test: FormFieldLike, trial: FormFieldLike, x_q: jnp.ndarray, w: jnp.ndarray, elem_id: jnp.ndarray | int = 0, spaces: dict[str, 'FieldPair'] | None = None, default_space: str | None = None)#: Bundle test/trial fields and quadrature data for element assembly.

class fluxfem.core.MixedFormContext(bindings: dict[str, FieldPair], x_q: jnp.ndarray, w: jnp.ndarray, elem_id: jnp.ndarray | int = 0, unknown: FormFieldLike | None = None, spaces: dict[str, FieldPair] | None = None, default_space: str | None = None)#: FormContext for mixed formulations keyed by binding/field name.

class fluxfem.core.VolumeContext(*args, **kwargs)#: Minimum interface for volume weak-form evaluation.

class fluxfem.core.SurfaceContext(*args, **kwargs)#: Minimum interface for surface weak-form evaluation.

class fluxfem.core.LinearForm(fn, *, kind: str)#: Linear form wrapper with volume/surface backends.

class fluxfem.core.BilinearForm(fn, *, kind: str)#: Bilinear form wrapper with volume/surface backends.

class fluxfem.core.ResidualForm(fn)#: Residual form wrapper.

class fluxfem.core.MixedWeakForm(*, residuals: Mapping[str, Callable | Expr | MixedResidualBinding])#: Container for mixed weak-form residuals keyed by residual label.

Helper to build mixed residual dictionaries.

Example

res = make_mixed_residuals(u=res_u, p=res_p)

fluxfem.core.bind_mixed_residual(target: str, fn: Callable | Expr, *, space: str | None = None) → MixedResidualBinding#: Create an explicit mixed residual binding.

fluxfem.core.kernel(*, kind: str, domain: str = 'volume')#: Decorator to tag raw kernels with kind/domain metadata for assembly inference.

Kernel metadata (ff.kernel)#

@ff.kernel attaches metadata used by space.assemble to infer the form kind and domain. The following combinations are supported:

kind	domain	expected kernel signature
bilinear	volume	`(ctx, params) -> (n_q, n_ldofs, n_ldofs)`
linear	volume	`(ctx, params) -> (n_q, n_ldofs)`
linear	surface	`(ctx, params) -> (n_q, n_ldofs)`
residual	volume	`(ctx, u_elem, params) -> (n_q, n_ldofs)`
jacobian	volume	`(u_elem, ctx) -> (n_ldofs, n_ldofs)`

Use domain="surface" with a surface-specific assembly helper. For most workflows use SurfaceMesh.assemble_linear_form_on_space(...); lower-level options include assemble_surface_linear_form and assemble_surface_bilinear_form when you need explicit dim/n_total_nodes control.

Example (MixedProblem)#

mixed = ff.MixedSpaces(
    {
        "u": ff.ResidualSpaces(
            test=ff.NamedSpace("V", space),
            unknown=ff.NamedSpace("U", space),
        ),
        "p": ff.ResidualSpaces(
            test=ff.NamedSpace("Qv", space),
            unknown=ff.NamedSpace("Q", space),
        ),
    }
).to_fe_space()
residuals = ff.make_mixed_residuals(
    u=ff.bind_mixed_residual("u", res_u, space="U"),
    p=ff.bind_mixed_residual("p", res_p, space="Q"),
)
params = ff.Params(alpha=1.2, beta=-0.4)
pattern = mixed.get_sparsity_pattern(with_idx=True)
problem = ff.MixedProblem(mixed, residuals, params=params, pattern=pattern)

u0 = jnp.zeros(mixed.n_dofs)
K = problem.assemble_jacobian(u0)
R = problem.assemble_residual(u0)

Mixed naming convention:

ctx.test / ctx.trial for simple single-space code
ctx.bindings["u"] for named mixed-field lookup
ctx.spaces["U"] for explicit space-key lookup
alias keys such as ctx.spaces["V"] can point to the same mixed field when a field was declared through ResidualSpaces / JacobianSpaces

Example (MixedBlockSystem)#

This helper currently applies to MixedSpaces(...).to_fe_space() / MixedFESpace style layouts. It does not yet support MixedRoleSpaces(...).to_fe_space().

blocks = {
    "u": {"u": Kuu, "p": Kup},
    "p": {"u": Kpu, "p": Kpp},
}
rhs = {"u": Fu, "p": Fp}
constraints = {"u": (u_dofs, u_vals)}
system = mixed.build_block_system(blocks=blocks, rhs=rhs, constraints=constraints)

u_free = solver.solve(system.K, system.F)
u_full = system.expand(u_free)
fields = system.split(u_full)

Assembly#

fluxfem.core.make_sparsity_pattern(space: SpaceLike, *, with_idx: bool = True) → Any#