fluxfem 0.1.1a0__py3-none-any.whl

fluxfem/physics/diffusion.py

@@ -0,0 +1,18 @@
+import jax.numpy as jnp
+
+from ..core.forms import FormContext
+
+
+def diffusion_form(ctx: FormContext, kappa: float) -> jnp.ndarray:
+    """
+    Scalar diffusion bilinear form: kappa * grad_v · grad_u.
+
+    Returns the per-quadrature integrand for a standard Laplacian term.
+    """
+    grad_v = ctx.test.gradN
+    grad_u = ctx.trial.gradN
+    G = jnp.einsum("qia,qja->qij", grad_v, grad_u)  # ∇v_i · ∇u_j
+    return kappa * G
+
+
+__all__ = ["diffusion_form"]

fluxfem/physics/elasticity/__init__.py

@@ -0,0 +1,39 @@
+"""Elasticity-related helpers (linear models, materials, forms)."""
+
+from .materials import lame_parameters, isotropic_3d_D
+from .linear import (
+    linear_elasticity_form,
+    vector_body_force_form,
+    constant_body_force_vector_form,
+    assemble_constant_body_force,
+)
+from .hyperelastic import (
+    right_cauchy_green,
+    green_lagrange_strain,
+    deformation_gradient,
+    pk2_neo_hookean,
+    neo_hookean_residual_form,
+    make_elastic_point_data,
+    write_elastic_vtu,
+)
+from .stress import principal_stresses, principal_sum, max_shear_stress, von_mises_stress
+
+__all__ = [
+    "lame_parameters",
+    "isotropic_3d_D",
+    "linear_elasticity_form",
+    "vector_body_force_form",
+    "constant_body_force_vector_form",
+    "assemble_constant_body_force",
+    "right_cauchy_green",
+    "green_lagrange_strain",
+    "deformation_gradient",
+    "pk2_neo_hookean",
+    "neo_hookean_residual_form",
+    "make_elastic_point_data",
+    "write_elastic_vtu",
+    "principal_stresses",
+    "principal_sum",
+    "max_shear_stress",
+    "von_mises_stress",
+]

fluxfem/physics/elasticity/hyperelastic.py

@@ -0,0 +1,99 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ...core.forms import FormContext
+from ...core.basis import build_B_matrices_finite
+from ..postprocess import make_point_data_displacement, write_point_data_vtu
+
+
+def right_cauchy_green(F: jnp.ndarray) -> jnp.ndarray:
+    """C = F^T F (right Cauchy-Green)."""
+    return jnp.einsum("...ik,...jk->...ij", F, F)
+
+
+def green_lagrange_strain(F: jnp.ndarray) -> jnp.ndarray:
+    """E = 0.5 (C - I)."""
+    I = jnp.eye(F.shape[-1], dtype=F.dtype)
+    C = right_cauchy_green(F)
+    return 0.5 * (C - I)
+
+
+def deformation_gradient(ctx: FormContext, u_elem: jnp.ndarray) -> jnp.ndarray:
+    """
+    Compute deformation gradient F = I + grad_u for a 3D vector displacement.
+
+    ctx: FormContext with test/trial grads (reference configuration)
+    u_elem: (n_ldofs,) element displacement in dof ordering [u0,v0,w0, u1,v1,w1, ...]
+    returns: (n_q, 3, 3) F per quadrature point
+    """
+    u_nodes = u_elem.reshape(-1, 3)  # (n_nodes, 3)
+    grad_u = jnp.einsum("qaj,ai->qij", ctx.trial.gradN, u_nodes)  # ∂u_i/∂X_j
+    I = jnp.eye(3, dtype=u_elem.dtype)
+    return I[None, ...] + grad_u
+
+
+def pk2_neo_hookean(F: jnp.ndarray, mu: float, lam: float) -> jnp.ndarray:
+    """
+    Compressible Neo-Hookean PK2 stress:
+      S = mu * (I - C^{-1}) + lam * ln J * C^{-1}
+      C = F^T F, J = sqrt(det C)
+    """
+    C = right_cauchy_green(F)
+    C_inv = jnp.linalg.inv(C)
+    J = jnp.sqrt(jnp.linalg.det(C))
+    I = jnp.eye(3, dtype=F.dtype)
+    return mu * (I - C_inv) + lam * jnp.log(J)[..., None, None] * C_inv
+
+
+def neo_hookean_residual_form(ctx: FormContext, u_elem: jnp.ndarray, params) -> jnp.ndarray:
+    """
+    Compressible Neo-Hookean residual (Total Lagrangian, PK2).
+    params: dict-like with keys \"mu\", \"lam\" or tuple (mu, lam)
+    returns: (n_q, n_ldofs)
+    """
+    if isinstance(params, dict):
+        mu = params["mu"]
+        lam = params["lam"]
+    else:
+        mu, lam = params
+
+    F = deformation_gradient(ctx, u_elem)          # (n_q, 3, 3)
+    S = pk2_neo_hookean(F, mu, lam)                # (n_q, 3, 3)
+
+    S_voigt = jnp.stack(
+        [
+            S[..., 0, 0],
+            S[..., 1, 1],
+            S[..., 2, 2],
+            S[..., 0, 1],
+            S[..., 1, 2],
+            S[..., 2, 0],
+        ],
+        axis=-1,
+    )  # (n_q, 6)
+
+    B = build_B_matrices_finite(ctx.trial.gradN, F)           # (n_q, 6, n_ldofs)
+    BT = jnp.swapaxes(B, 1, 2)                               # (n_q, n_ldofs, 6)
+    return jnp.einsum("qik,qk->qi", BT, S_voigt)   # (n_q, n_ldofs)
+
+
+__all__ = [
+    "right_cauchy_green",
+    "green_lagrange_strain",
+    "deformation_gradient",
+    "pk2_neo_hookean",
+    "neo_hookean_residual_form",
+    "make_elastic_point_data",
+    "write_elastic_vtu",
+]
+
+
+def make_elastic_point_data(mesh, space, u, *, compute_j: bool = True, deformed_scale: float = 1.0):
+    """Alias to postprocess.make_point_data_displacement for backward compatibility."""
+    return make_point_data_displacement(mesh, space, u, compute_j=compute_j, deformed_scale=deformed_scale)
+
+
+def write_elastic_vtu(mesh, space, u, filepath: str, *, compute_j: bool = True, deformed_scale: float = 1.0):
+    """Alias to postprocess.write_point_data_vtu for backward compatibility."""
+    return write_point_data_vtu(mesh, space, u, filepath, compute_j=compute_j, deformed_scale=deformed_scale)

fluxfem/physics/elasticity/linear.py

@@ -0,0 +1,58 @@
+import jax.numpy as jnp
+
+from ...core.assembly import assemble_linear_form
+from ...core.forms import FormContext, vector_load_form
+from ...core.basis import build_B_matrices
+from ...physics.operators import sym_grad
+
+# from ...mechanics.kinematics import build_B_matrices
+
+
+# def linear_elasticity_form(ctx: FormContext, D: jnp.ndarray) -> jnp.ndarray:
+#     """3D linear elasticity bilinear form B^T D B."""
+#     grad_v = ctx.test.grad
+#     grad_u = ctx.trial.grad
+#     B = build_B_matrices(grad_u)               # (n_q, 6, 24)
+#     BT = jnp.swapaxes(build_B_matrices(grad_v), 1, 2)  # (n_q, 24, 6)
+#     BDB = jnp.einsum("qik,kl,qlm->qim", BT, D, B)  # (n_q, 24, 24)
+#     return BDB
+
+
+def linear_elasticity_form(ctx: FormContext, D: jnp.ndarray) -> jnp.ndarray:
+    """
+    Linear-elasticity bilinear form in Voigt notation.
+
+    Returns the per-quadrature integrand for Bv^T D Bu, where B is the
+    symmetric-gradient operator for the test/trial fields.
+    """
+    Bu = sym_grad(ctx.trial)                 # (n_q, 6, ndofs_e)
+    Bv = sym_grad(ctx.test)                  # (n_q, 6, ndofs_e)
+    return jnp.einsum("qik,kl,qlm->qim", jnp.swapaxes(Bv, 1, 2), D, Bu)
+
+
+def vector_body_force_form(ctx: FormContext, load_vec: jnp.ndarray) -> jnp.ndarray:
+    """Linear form for 3D vector body force f (constant in space)."""
+    return vector_load_form(ctx.test, load_vec)
+
+
+def assemble_constant_body_force(space, gravity_vec, density: float, *, sparse: bool = False):
+    """
+    Convenience: assemble body force from density * gravity vector.
+    gravity_vec: length-3 array-like (direction and magnitude of g)
+    density: scalar density (consistent with unit system)
+    """
+    g = jnp.asarray(gravity_vec)
+    f_vec = density * g
+    return assemble_linear_form(space, vector_body_force_form, params=f_vec, sparse=sparse)
+
+
+# Backward compatibility alias
+constant_body_force_vector_form = vector_body_force_form
+
+
+__all__ = [
+    "linear_elasticity_form",
+    "vector_body_force_form",
+    "constant_body_force_vector_form",
+    "assemble_constant_body_force",
+]

fluxfem/physics/elasticity/materials.py

@@ -0,0 +1,32 @@
+import jax
+import jax.numpy as jnp
+
+DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+
+def lame_parameters(E: float, nu: float) -> tuple[float, float]:
+    """Return Lamé parameters (lambda, mu) from Young's modulus and Poisson ratio."""
+    lam = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
+    mu = E / (2.0 * (1.0 + nu))
+    return float(lam), float(mu)
+
+
+def isotropic_3d_D(E: float, nu: float) -> jnp.ndarray:
+    """Return 3D isotropic linear elasticity constitutive matrix in Voigt form."""
+    lam, mu = lame_parameters(E, nu)
+
+    D = jnp.array(
+        [
+            [lam + 2 * mu, lam, lam, 0.0, 0.0, 0.0],
+            [lam, lam + 2 * mu, lam, 0.0, 0.0, 0.0],
+            [lam, lam, lam + 2 * mu, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, mu, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, mu, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, mu],
+        ],
+        dtype=DTYPE,
+    )
+    return D
+
+
+__all__ = ["lame_parameters", "isotropic_3d_D"]

fluxfem/physics/elasticity/stress.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import jax.numpy as jnp
+
+
+def _sym(A: jnp.ndarray) -> jnp.ndarray:
+    return 0.5 * (A + jnp.swapaxes(A, -1, -2))
+
+
+def principal_stresses(S: jnp.ndarray) -> jnp.ndarray:
+    """
+    Return principal stresses (eigvals) for symmetric 3x3 stress tensor.
+    Supports batching over leading dimensions.
+    """
+    S_sym = _sym(S)
+    return jnp.linalg.eigvalsh(S_sym)
+
+
+def principal_sum(S: jnp.ndarray) -> jnp.ndarray:
+    """Sum of principal stresses (trace)."""
+    return jnp.trace(S, axis1=-2, axis2=-1)
+
+
+def max_shear_stress(S: jnp.ndarray) -> jnp.ndarray:
+    """
+    Maximum shear stress = (sigma_max - sigma_min) / 2.
+    """
+    vals = principal_stresses(S)
+    return 0.5 * (vals[..., -1] - vals[..., 0])
+
+
+def von_mises_stress(S: jnp.ndarray) -> jnp.ndarray:
+    """
+    von Mises equivalent stress: sqrt(3/2 * dev(S):dev(S)).
+    """
+    tr = jnp.trace(S, axis1=-2, axis2=-1)[..., None, None]
+    dev = S - tr / 3.0
+    return jnp.sqrt(1.5 * jnp.sum(dev * dev, axis=(-2, -1)))
+
+
+__all__ = [
+    "principal_stresses",
+    "principal_sum",
+    "max_shear_stress",
+    "von_mises_stress",
+]

fluxfem/physics/operators.py

@@ -0,0 +1,109 @@
+
+
+# fluxfem/mechanics/operators.py
+from __future__ import annotations
+import jax
+import jax.numpy as jnp
+
+
+def dot(a, b):
+    """
+    Batched matrix product on the last two axes.
+
+    If the first argument is a FormField, dispatch to vector_load_form to build
+    the linear form contribution for a vector load.
+    """
+    if hasattr(a, "N") and getattr(a, "value_dim", None) is not None:
+        from ..core.forms import vector_load_form
+        return vector_load_form(a, b)
+    return jnp.matmul(a, b)
+
+
+def ddot(a, b, c=None):
+    """
+    Double contraction on the last two axes.
+
+    - ddot(a, b): sum_ij a_ij * b_ij
+    - ddot(a, b, c): a^T b c (Voigt-style linear elasticity blocks)
+    """
+    if c is None:
+        return jnp.einsum("...ij,...ij->...", a, b)
+    a_t = jnp.swapaxes(a, -1, -2)
+    return jnp.einsum("...ik,kl,...lm->...im", a_t, b, c)
+
+
+def transpose_last2(a):
+    """Swap the last two axes (batched transpose)."""
+    return jnp.swapaxes(a, -1, -2)
+
+
+def sym_grad(field) -> jnp.ndarray:
+    """
+    Symmetric gradient operator for vector mechanics (small strain).
+
+    Parameters
+    ----------
+    field : FormField-like
+        Must provide:
+          - field.gradN : (n_q, n_nodes, 3)
+          - field.basis.dofs_per_node (usually 3)
+
+    Returns
+    -------
+    B : jnp.ndarray
+        (n_q, 6, dofs_per_node*n_nodes) Voigt order [xx, yy, zz, xy, yz, zx]
+        Such that eps_voigt(q,:) = B(q,:,:) @ u_elem
+    """
+    gradN = field.gradN                      # (n_q, n_nodes, 3)
+    dofs = getattr(field.basis, "dofs_per_node", 3)
+    n_q, n_nodes, _ = gradN.shape
+    n_dofs = dofs * n_nodes
+
+    def B_single(dN):
+        B = jnp.zeros((6, n_dofs), dtype=dN.dtype)
+
+        def node_fun(a, B):
+            dNdx, dNdy, dNdz = dN[a, 0], dN[a, 1], dN[a, 2]
+            col = dofs * a
+
+            # eps_xx, eps_yy, eps_zz
+            B = B.at[0, col + 0].set(dNdx)  # dux/dx
+            B = B.at[1, col + 1].set(dNdy)  # duy/dy
+            B = B.at[2, col + 2].set(dNdz)  # duz/dz
+
+            # eps_xy = 1/2(dux/dy + duy/dx)
+            B = B.at[3, col + 0].set(dNdy)
+            B = B.at[3, col + 1].set(dNdx)
+
+            # eps_yz = 1/2(duy/dz + duz/dy)
+            B = B.at[4, col + 1].set(dNdz)
+            B = B.at[4, col + 2].set(dNdy)
+
+            # eps_zx = 1/2(duz/dx + dux/dz)
+            B = B.at[5, col + 0].set(dNdz)
+            B = B.at[5, col + 2].set(dNdx)
+            return B
+
+        return jax.lax.fori_loop(0, n_nodes, node_fun, B)
+
+    return jax.vmap(B_single)(gradN)
+
+
+def sym_grad_u(field, u_elem: jnp.ndarray) -> jnp.ndarray:
+    """
+    Apply sym_grad(field) to a local displacement vector.
+
+    Parameters
+    ----------
+    field : FormField-like
+        Vector field basis data.
+    u_elem : jnp.ndarray
+        Element displacement vector (dofs_per_node*n_nodes,).
+
+    Returns
+    -------
+    jnp.ndarray
+        Symmetric strain in Voigt form with shape (n_q, 6).
+    """
+    B = sym_grad(field)
+    return jnp.einsum("qik,k->qi", B, u_elem)

fluxfem/physics/postprocess.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+# from ..core.assembly import build_form_contexts
+from ..tools.visualizer import write_vtu
+from ..core.interp import interpolate_field_at_element_points
+
+
+def make_point_data_displacement(mesh, space, u, *, compute_j: bool = True, deformed_scale: float = 1.0):
+    """
+    Common postprocess helper to build point data dictionaries:
+      - displacement
+      - deformed_coords = X + scale * u
+      - optional J (nodal average of det(F))
+
+    Assumes 3 dof/node ordering [u0,v0,w0, u1,v1,w1, ...].
+    """
+    coords = np.asarray(mesh.coords)
+    u_np = np.asarray(u)
+    n_nodes = coords.shape[0]
+    if u_np.shape[0] != 3 * n_nodes:
+        raise ValueError(f"Expected 3 dof/node vector; got {u_np.shape[0]} entries for {n_nodes} nodes")
+
+    u_nodes = u_np.reshape(n_nodes, 3)
+    disp = u_nodes
+    deformed_coords = coords + deformed_scale * disp
+
+    point_data = {
+        "displacement": disp.astype(np.float64),
+        "deformed_coords": deformed_coords.astype(np.float64),
+    }
+
+    if compute_j:
+        elem_conns = np.asarray(space.elem_dofs)
+        # ctxs = build_form_contexts(space)
+        ctxs = space.build_form_contexts()
+        u_arr = jnp.asarray(u)
+
+        from .elasticity.hyperelastic import deformation_gradient  # local import to avoid circular
+
+        def elem_J(ctx, conn):
+            F = deformation_gradient(ctx, u_arr[conn])
+            return jnp.mean(jnp.linalg.det(F))
+
+        J_elem = np.asarray(jax.vmap(elem_J)(ctxs, elem_conns), dtype=float)
+        J_sum = np.zeros(n_nodes, dtype=float)
+        J_cnt = np.zeros(n_nodes, dtype=float)
+        for conn, J_mean in zip(elem_conns, J_elem):
+            node_ids = np.unique(np.asarray(conn) // 3)
+            J_sum[node_ids] += J_mean
+            J_cnt[node_ids] += 1.0
+        J_nodal = J_sum / np.maximum(J_cnt, 1.0)
+        point_data["J"] = J_nodal.astype(np.float64)
+
+    return point_data
+
+
+def write_point_data_vtu(mesh, space, u, filepath: str, *, compute_j: bool = True, deformed_scale: float = 1.0):
+    """Write VTU with displacement/deformed_coords and optional J."""
+    pdata = make_point_data_displacement(mesh, space, u, compute_j=compute_j, deformed_scale=deformed_scale)
+    write_vtu(mesh, filepath, point_data=pdata)
+
+
+__all__ = ["make_point_data_displacement", "write_point_data_vtu", "interpolate_at_points"]
+def interpolate_at_points(space, u, points: np.ndarray):
+    """
+    Interpolate displacement field at given physical points (Hex8 only, structured search).
+    - points: (m,3) array of physical coordinates.
+    Returns: (m,3) interpolated displacement.
+    """
+    pts = np.asarray(points, dtype=float)
+    mesh = space.mesh
+    # Only support StructuredHexBox-backed HexMesh (regular grid)
+    if not hasattr(mesh, "origin") and not hasattr(mesh, "lx"):
+        raise NotImplementedError("interpolate_at_points currently supports StructuredHexBox meshes.")
+    coords = np.asarray(mesh.coords)
+    # grid dimensions from origin/extents and nx,ny,nz inferred from spacing
+    xs = np.unique(coords[:, 0])
+    ys = np.unique(coords[:, 1])
+    zs = np.unique(coords[:, 2])
+    dx, dy, dz = xs[1] - xs[0], ys[1] - ys[0], zs[1] - zs[0]
+    ox, oy, oz = xs.min(), ys.min(), zs.min()
+    nx, ny, nz = len(xs) - 1, len(ys) - 1, len(zs) - 1
+
+    # map physical point to element indices and local coords
+    def phys_to_elem_local(p):
+        x, y, z = p
+        i = min(max(int(np.floor((x - ox) / dx)), 0), nx - 1)
+        j = min(max(int(np.floor((y - oy) / dy)), 0), ny - 1)
+        k = min(max(int(np.floor((z - oz) / dz)), 0), nz - 1)
+        # local coords in [-1,1]
+        xi = 2 * ((x - (ox + i * dx)) / dx) - 1
+        eta = 2 * ((y - (oy + j * dy)) / dy) - 1
+        zeta = 2 * ((z - (oz + k * dz)) / dz) - 1
+        elem_idx = k * (nx * ny) + j * nx + i
+        return elem_idx, np.array([xi, eta, zeta], dtype=float)
+
+    elem_indices = []
+    locals = []
+    for p in pts:
+        e, loc = phys_to_elem_local(p)
+        elem_indices.append(e)
+        locals.append(loc)
+    locals = np.stack(locals, axis=0)  # (m,3)
+
+    vals_per_elem = interpolate_field_at_element_points(space, u, locals)  # (n_elem,m,3) but m is same locals; pick per elem
+    vals = np.zeros((len(pts), 3), dtype=float)
+    for idx, e in enumerate(elem_indices):
+        vals[idx] = vals_per_elem[e, idx]
+    return vals

fluxfem/solver/__init__.py

@@ -0,0 +1,47 @@
+from .sparse import SparsityPattern, FluxSparseMatrix
+from .dirichlet import (
+    enforce_dirichlet_dense,
+    enforce_dirichlet_sparse,
+    free_dofs,
+    condense_dirichlet_fluxsparse,
+    condense_dirichlet_dense,
+    expand_dirichlet_solution,
+)
+from .cg import cg_solve, cg_solve_jax
+from .newton import newton_solve
+from .solve_runner import (
+    NonlinearAnalysis,
+    NewtonLoopConfig,
+    LoadStepResult,
+    NewtonSolveRunner,
+    solve_nonlinear,
+    LinearAnalysis,
+    LinearSolveConfig,
+    LinearStepResult,
+    LinearSolveRunner,
+)
+from .solver import LinearSolver, NonlinearSolver
+
+__all__ = [
+    "SparsityPattern",
+    "FluxSparseMatrix",
+    "enforce_dirichlet_dense",
+    "enforce_dirichlet_sparse",
+    "free_dofs",
+    "condense_dirichlet_fluxsparse",
+    "condense_dirichlet_dense",
+    "expand_dirichlet_solution",
+    "cg_solve",
+    "cg_solve_jax",
+    "newton_solve",
+    "LinearAnalysis",
+    "LinearSolveConfig",
+    "LinearStepResult",
+    "NonlinearAnalysis",
+    "NewtonLoopConfig",
+    "LoadStepResult",
+    "NewtonSolveRunner",
+    "solve_nonlinear",
+    "LinearSolver",
+    "NonlinearSolver",
+]