fluxfem 0.1.3__tar.gz → 0.1.4__tar.gz

{fluxfem-0.1.3 → fluxfem-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fluxfem
-Version: 0.1.3
+Version: 0.1.4
 Summary: FluxFEM: A weak-form-centric differentiable finite element framework in JAX
 License: Apache-2.0
 Author: Kohei Watanabe
@@ -10,8 +10,8 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: jax (==0.8.2)
-Requires-Dist: jaxlib (==0.8.2)
+Requires-Dist: jax (>=0.8.2,<0.9.0)
+Requires-Dist: jaxlib (>=0.8.2,<0.9.0)
 Requires-Dist: matplotlib (>=3.10.7,<4.0.0)
 Requires-Dist: meshio (>=5.3.5,<6.0.0)
 Requires-Dist: pyvista (>=0.46.4,<0.47.0)
@@ -22,6 +22,8 @@ Description-Content-Type: text/markdown
 [![Python Version](https://img.shields.io/pypi/pyversions/fluxfem.svg)](https://pypi.org/project/fluxfem/)
 ![CI](https://github.com/kevin-tofu/fluxfem/actions/workflows/python-tests.yml/badge.svg)
 ![CI](https://github.com/kevin-tofu/fluxfem/actions/workflows/sphinx.yml/badge.svg)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.18055465.svg)](https://doi.org/10.5281/zenodo.18055465)
+
 
 # FluxFEM
  A weak-form-centric differentiable finite element framework in JAX
@@ -44,13 +46,10 @@ Description-Content-Type: text/markdown
 
 
 ## Features
-- Built on JAX, enabling automatic differentiation and high-performance execution via grad, jit, vmap, and related transformations.
-
-- A FEM framework with a weak-form–centric API, emphasizing a smooth transition from theoretical formulations to practical code implementations.
-
-- Supports two assembly approaches: weak-form-based assembly and a tensor-based (scikit-fem–style) assembly.
-
-- enables to handle both Linear / Non-Linear analysis with AD in JAX
+- Built on JAX, enabling automatic differentiation with grad, jit, vmap, and related transformations.
+- Weak-form–centric API that keeps formulations close to code; weak forms are represented as expression trees and compiled to element kernels.
+- Two assembly approaches: weak-form-based assembly and a tensor-based (scikit-fem–style) assembly.
+- Handles both linear and nonlinear analyses with AD in JAX.
 
 ## Usage 
 
@@ -62,6 +61,9 @@ The first approach offers simplicity and convenience, as mathematical expression
 However, for more complex operations, the second approach can be easier to implement in practice.
 This is because the weak-form-based assembly is ultimately transformed into the tensor-based representation internally during computation.
 
+## Weak Form Compile Flow
+Weak-form expressions are compiled into an evaluation plan and then executed per element.
+
 ### weak-form-based assembly
 ```Python
 import fluxfem as ff
@@ -70,7 +72,7 @@ import fluxfem.helpers_wf as h_wf
 space = ff.make_hex_space(mesh, dim=3, intorder=2)
 D = ff.isotropic_3d_D(1.0, 0.3)
 bilinear_form = ff.BilinearForm.volume(
-    lambda u, v, D: h_wf.ddot(v.sym_grad, D @ u.sym_grad) * h_wf.dOmega()
+    lambda u, v, D: h_wf.ddot(v.sym_grad, h_wf.matmul_std(D, u.sym_grad)) * h_wf.dOmega()
 )
 K_wf = space.assemble_bilinear_form(
     bilinear_form.get_compiled(),

{fluxfem-0.1.3 → fluxfem-0.1.4}/README.md

@@ -3,6 +3,8 @@
 [![Python Version](https://img.shields.io/pypi/pyversions/fluxfem.svg)](https://pypi.org/project/fluxfem/)
 ![CI](https://github.com/kevin-tofu/fluxfem/actions/workflows/python-tests.yml/badge.svg)
 ![CI](https://github.com/kevin-tofu/fluxfem/actions/workflows/sphinx.yml/badge.svg)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.18055465.svg)](https://doi.org/10.5281/zenodo.18055465)
+
 
 # FluxFEM
  A weak-form-centric differentiable finite element framework in JAX
@@ -25,13 +27,10 @@
 
 
 ## Features
-- Built on JAX, enabling automatic differentiation and high-performance execution via grad, jit, vmap, and related transformations.
-
-- A FEM framework with a weak-form–centric API, emphasizing a smooth transition from theoretical formulations to practical code implementations.
-
-- Supports two assembly approaches: weak-form-based assembly and a tensor-based (scikit-fem–style) assembly.
-
-- enables to handle both Linear / Non-Linear analysis with AD in JAX
+- Built on JAX, enabling automatic differentiation with grad, jit, vmap, and related transformations.
+- Weak-form–centric API that keeps formulations close to code; weak forms are represented as expression trees and compiled to element kernels.
+- Two assembly approaches: weak-form-based assembly and a tensor-based (scikit-fem–style) assembly.
+- Handles both linear and nonlinear analyses with AD in JAX.
 
 ## Usage 
 
@@ -43,6 +42,9 @@ The first approach offers simplicity and convenience, as mathematical expression
 However, for more complex operations, the second approach can be easier to implement in practice.
 This is because the weak-form-based assembly is ultimately transformed into the tensor-based representation internally during computation.
 
+## Weak Form Compile Flow
+Weak-form expressions are compiled into an evaluation plan and then executed per element.
+
 ### weak-form-based assembly
 ```Python
 import fluxfem as ff
@@ -51,7 +53,7 @@ import fluxfem.helpers_wf as h_wf
 space = ff.make_hex_space(mesh, dim=3, intorder=2)
 D = ff.isotropic_3d_D(1.0, 0.3)
 bilinear_form = ff.BilinearForm.volume(
-    lambda u, v, D: h_wf.ddot(v.sym_grad, D @ u.sym_grad) * h_wf.dOmega()
+    lambda u, v, D: h_wf.ddot(v.sym_grad, h_wf.matmul_std(D, u.sym_grad)) * h_wf.dOmega()
 )
 K_wf = space.assemble_bilinear_form(
     bilinear_form.get_compiled(),

{fluxfem-0.1.3 → fluxfem-0.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fluxfem"
-version = "0.1.3"
+version = "0.1.4"
 description = ""
 authors = [
     {name = "Kohei Watanabe",email = "koheitech001@gmail.com"}
@@ -18,13 +18,13 @@ dependencies = [
     "meshio (>=5.3.5,<6.0.0)",
     "matplotlib (>=3.10.7,<4.0.0)",
     # Note: jax versioning currently around 0.8.x. Allow up to next minor.
-    "jax (==0.8.2)",
-    "jaxlib (==0.8.2)"
+    "jax (>=0.8.2,<0.9.0)",
+    "jaxlib (>=0.8.2,<0.9.0)"
 ]
 
 [tool.poetry]
 name = "fluxfem"
-version = "0.1.3"
+version = "0.1.4"
 description = "FluxFEM: A weak-form-centric differentiable finite element framework in JAX"
 authors = ["Kohei Watanabe <koheitech001@gmail.com>"]
 readme = "README.md"
@@ -39,8 +39,8 @@ python = ">=3.11,<3.14"
 pyvista = ">=0.46.4,<0.47.0"
 meshio = ">=5.3.5,<6.0.0"
 matplotlib = ">=3.10.7,<4.0.0"
-jax = "==0.8.2"
-jaxlib = "==0.8.2"
+jax = ">=0.8.2,<0.9.0"
+jaxlib = ">=0.8.2,<0.9.0"
 
 
 [tool.poetry.group.dev]

{fluxfem-0.1.3 → fluxfem-0.1.4}/src/fluxfem/__init__.py

@@ -1,159 +1,6 @@
-from .core import (
-    FESpaceBase,
-    FESpace,
-    FESpacePytree,
-    Expr,
-    FieldRef,
-    ParamRef,
-    trial_ref,
-    test_ref,
-    unknown_ref,
-    Params,
-    LinearForm,
-    BilinearForm,
-    ResidualForm,
-    MixedWeakForm,
-    compile_bilinear,
-    compile_linear,
-    compile_residual,
-    compile_surface_linear,
-    compile_mixed_residual,
-    sdot,
-    dOmega,
-    FormContext,
-    MixedFormContext,
-    FieldPair,
-    ElementVector,
-    vector_load_form,
-    make_space,
-    make_space_pytree,
-    make_hex_basis,
-    make_hex_basis_pytree,
-    make_hex_space,
-    make_hex_space_pytree,
-    make_hex20_basis,
-    make_hex20_basis_pytree,
-    make_hex20_space,
-    make_hex20_space_pytree,
-    make_hex27_basis,
-    make_hex27_basis_pytree,
-    make_hex27_space,
-    make_hex27_space_pytree,
-    make_tet_basis,
-    make_tet_basis_pytree,
-    make_tet_space,
-    make_tet_space_pytree,
-    make_tet10_basis,
-    make_tet10_basis_pytree,
-    make_tet10_space,
-    make_tet10_space_pytree,
-    make_element_residual_kernel,
-    make_element_jacobian_kernel,
-    element_residual,
-    element_jacobian,
-    make_sparsity_pattern,
-    assemble_functional,
-    assemble_mass_matrix,
-    scalar_body_force_form,
-    make_scalar_body_force_form,
-    constant_body_force_form,
-    vector_body_force_form,
-    BaseMeshPytree,
-    bbox_predicate,
-    plane_predicate,
-    axis_plane_predicate,
-    slab_predicate,
-    HexMesh,
-    HexMeshPytree,
-    StructuredHexBox,
-    SurfaceMesh,
-    SurfaceMeshPytree,
-    load_gmsh_mesh,
-    load_gmsh_hex_mesh,
-    load_gmsh_tet_mesh,
-    make_surface_from_facets,
-    TetMesh,
-    TetMeshPytree,
-    StructuredTetBox,
-    StructuredTetTensorBox,
-    tag_axis_minmax_facets,
-    HexTriLinearBasis,
-    HexTriLinearBasisPytree,
-    HexSerendipityBasis20,
-    HexSerendipityBasis20Pytree,
-    HexTriQuadraticBasis27,
-    HexTriQuadraticBasis27Pytree,
-    TetLinearBasis,
-    TetLinearBasisPytree,
-    TetQuadraticBasis10,
-    TetQuadraticBasis10Pytree,
-    MeshData,
-    BasisData,
-    SpaceData,
-    MeshData,
-    BasisData,
-    SpaceData,
-    lame_parameters,
-    isotropic_3d_D,
-    spdirect_solve_cpu,
-    spdirect_solve_gpu,
-    spdirect_solve_jax,
-    coo_to_csr,
-    SparsityPattern,
-    FluxSparseMatrix,
-    LinearSolver,
-    NonlinearSolver,
-    enforce_dirichlet_dense,
-    enforce_dirichlet_sparse,
-    free_dofs,
-    condense_dirichlet_fluxsparse,
-    condense_dirichlet_dense,
-    expand_dirichlet_solution,
-    cg_solve,
-    cg_solve_jax,
-    NonlinearAnalysis,
-    NewtonLoopConfig,
-    LoadStepResult,
-    NewtonSolveRunner,
-    solve_nonlinear,
-    LinearAnalysis,
-    LinearSolveConfig,
-    LinearStepResult,
-    LinearSolveRunner,
-    newton_solve,
-    SurfaceFormField,
-    SurfaceFormContext,
-    vector_surface_load_form,
-    make_vector_surface_load_form,
-    assemble_surface_linear_form,
-)
-from .tools.visualizer import write_vtu, write_displacement_vtu
-from .tools.jit import make_jitted_residual, make_jitted_jacobian
-from .physics import (
-    linear_elasticity_form,
-    vector_body_force_form,
-    constant_body_force_vector_form,
-    diffusion_form,
-    dot,
-    ddot,
-    transpose_last2,
-    sym_grad,
-    sym_grad_u,
-    right_cauchy_green,
-    green_lagrange_strain,
-    deformation_gradient,
-    pk2_neo_hookean,
-    neo_hookean_residual_form,
-    make_elastic_point_data,
-    write_elastic_vtu,
-    make_point_data_displacement,
-    write_point_data_vtu,
-    interpolate_at_points,
-    principal_stresses,
-    principal_sum,
-    max_shear_stress,
-    von_mises_stress,
-)
+from __future__ import annotations
+
+import importlib
 
 __all__ = [
     "FESpaceBase",
@@ -175,10 +22,13 @@ __all__ = [
     "compile_residual",
     "compile_surface_linear",
     "compile_mixed_residual",
+    "outer",
     "sdot",
     "dOmega",
     "FormContext",
     "MixedFormContext",
+    "VolumeContext",
+    "SurfaceContext",
     "FieldPair",
     "ElementVector",
     "vector_load_form",
@@ -243,6 +93,7 @@ __all__ = [
     "vector_surface_load_form",
     "make_vector_surface_load_form",
     "assemble_surface_linear_form",
+    "tag_axis_minmax_facets",
     "load_gmsh_mesh",
     "load_gmsh_hex_mesh",
     "load_gmsh_tet_mesh",
@@ -307,11 +158,71 @@ __all__ = [
     "von_mises_stress",
 ]
 
+_PHYSICS_EXPORTS = {
+    "lame_parameters",
+    "isotropic_3d_D",
+    "linear_elasticity_form",
+    "vector_body_force_form",
+    "constant_body_force_vector_form",
+    "diffusion_form",
+    "dot",
+    "ddot",
+    "transpose_last2",
+    "sym_grad",
+    "sym_grad_u",
+    "right_cauchy_green",
+    "green_lagrange_strain",
+    "deformation_gradient",
+    "pk2_neo_hookean",
+    "neo_hookean_residual_form",
+    "make_elastic_point_data",
+    "write_elastic_vtu",
+    "make_point_data_displacement",
+    "write_point_data_vtu",
+    "interpolate_at_points",
+    "principal_stresses",
+    "principal_sum",
+    "max_shear_stress",
+    "von_mises_stress",
+}
+
+_TOOLS_VIS_EXPORTS = {
+    "write_vtu",
+    "write_displacement_vtu",
+}
+
+_TOOLS_JIT_EXPORTS = {
+    "make_jitted_residual",
+    "make_jitted_jacobian",
+}
+
+_CORE_EXPORTS = set(__all__) - _PHYSICS_EXPORTS - _TOOLS_VIS_EXPORTS - _TOOLS_JIT_EXPORTS
+
+
+def __getattr__(name: str):
+    if name in _PHYSICS_EXPORTS:
+        module = importlib.import_module(".physics", __name__)
+    elif name in _TOOLS_VIS_EXPORTS:
+        module = importlib.import_module(".tools.visualizer", __name__)
+    elif name in _TOOLS_JIT_EXPORTS:
+        module = importlib.import_module(".tools.jit", __name__)
+    elif name in _CORE_EXPORTS:
+        module = importlib.import_module(".core", __name__)
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    value = getattr(module, name)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))
+
 
 try:
     from importlib.metadata import version, PackageNotFoundError
 except ImportError:
-    # old python
     from importlib_metadata import version, PackageNotFoundError
 
 
@@ -320,21 +231,17 @@ def read_version_from_pyproject():
     import re
 
     root = pathlib.Path(__file__).resolve().parent.parent
-    # while root != root.parent and depth < max_depth:
     pyproject_path = root / "pyproject.toml"
     if pyproject_path.exists():
         content = pyproject_path.read_text()
         match = re.search(r'version\s*=\s*"(.*?)"', content)
         if match:
             return match.group(1)
-    # root = root.parent
-    #     depth += 1
     return "0.0.0"
 
 
 def get_version(package_name):
     try:
-        # print("version-", version(package_name))
         return version(package_name)
     except PackageNotFoundError:
         return read_version_from_pyproject()

{fluxfem-0.1.3 → fluxfem-0.1.4}/src/fluxfem/core/__init__.py

@@ -1,4 +1,5 @@
 from .forms import FormContext, MixedFormContext, FieldPair, ElementVector, vector_load_form
+from .context_types import VolumeContext, SurfaceContext
 from .space import (
     FESpaceBase,
     FESpace,
@@ -71,6 +72,7 @@ from .weakform import (
     compile_mixed_residual,
     grad,
     sym_grad,
+    outer,
     dot,
     sdot,
     ddot,
@@ -87,6 +89,7 @@ from .weakform import (
     log,
     transpose_last2,
     matmul,
+    matmul_std,
     einsum,
 )
 from ..mesh import (
@@ -122,39 +125,56 @@ from .basis import (
     TetQuadraticBasis10,
     TetQuadraticBasis10Pytree,
 )
+import importlib
+
 from ..physics import lame_parameters, isotropic_3d_D
 from .solver import spdirect_solve_cpu, spdirect_solve_gpu, spdirect_solve_jax, coo_to_csr
-from ..solver import (
-    SparsityPattern,
-    FluxSparseMatrix,
-    LinearSolver,
-    NonlinearSolver,
-    enforce_dirichlet_dense,
-    enforce_dirichlet_sparse,
-    free_dofs,
-    condense_dirichlet_fluxsparse,
-    condense_dirichlet_dense,
-    expand_dirichlet_solution,
-    cg_solve,
-    cg_solve_jax,
-    NonlinearAnalysis,
-    NewtonLoopConfig,
-    LoadStepResult,
-    NewtonSolveRunner,
-    solve_nonlinear,
-    LinearAnalysis,
-    LinearSolveConfig,
-    LinearStepResult,
-    LinearSolveRunner,
-    newton_solve,
-)
-from ..solver.bc import (
-    SurfaceFormField,
-    SurfaceFormContext,
-    vector_surface_load_form,
-    make_vector_surface_load_form,
-    assemble_surface_linear_form,
-)
+
+_SOLVER_EXPORTS = {
+    "SparsityPattern",
+    "FluxSparseMatrix",
+    "LinearSolver",
+    "NonlinearSolver",
+    "enforce_dirichlet_dense",
+    "enforce_dirichlet_sparse",
+    "free_dofs",
+    "condense_dirichlet_fluxsparse",
+    "condense_dirichlet_dense",
+    "expand_dirichlet_solution",
+    "cg_solve",
+    "cg_solve_jax",
+    "NonlinearAnalysis",
+    "NewtonLoopConfig",
+    "LoadStepResult",
+    "NewtonSolveRunner",
+    "solve_nonlinear",
+    "LinearAnalysis",
+    "LinearSolveConfig",
+    "LinearStepResult",
+    "LinearSolveRunner",
+    "newton_solve",
+}
+
+_SOLVER_BC_EXPORTS = {
+    "SurfaceFormField",
+    "SurfaceFormContext",
+    "vector_surface_load_form",
+    "make_vector_surface_load_form",
+    "assemble_surface_linear_form",
+}
+
+
+def __getattr__(name: str):
+    if name in _SOLVER_EXPORTS:
+        module = importlib.import_module("..solver", __name__)
+    elif name in _SOLVER_BC_EXPORTS:
+        module = importlib.import_module("..solver.bc", __name__)
+    else:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    value = getattr(module, name)
+    globals()[name] = value
+    return value
 
 __all__ = [
     "FESpace",
@@ -162,6 +182,10 @@ __all__ = [
     "FESpacePytree",
     "FormContext",
     "MixedFormContext",
+    "VolumeContext",
+    "SurfaceContext",
+    "VolumeContext",
+    "SurfaceContext",
     "FieldPair",
     "ElementVector",
     "vector_load_form",
@@ -230,6 +254,7 @@ __all__ = [
     "compile_mixed_residual",
     "grad",
     "sym_grad",
+    "outer",
     "dot",
     "sdot",
     "ddot",
@@ -246,6 +271,7 @@ __all__ = [
     "log",
     "transpose_last2",
     "matmul",
+    "matmul_std",
     "einsum",
     "HexMesh",
     "HexMeshPytree",

fluxfem-0.1.4/src/fluxfem/core/context_types.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Any, Protocol, TypeAlias, runtime_checkable
+
+
+@runtime_checkable
+class VolumeContext(Protocol):
+    """Minimum interface for volume weak-form evaluation."""
+
+    test: Any
+    trial: Any
+    w: Any
+
+
+@runtime_checkable
+class SurfaceContext(Protocol):
+    """Minimum interface for surface weak-form evaluation."""
+
+    v: Any
+    w: Any
+    detJ: Any
+    normal: Any
+
+
+@runtime_checkable
+class FormFieldLike(Protocol):
+    """Minimum interface for form fields used in weak-form evaluation."""
+
+    N: Any
+    gradN: Any
+    detJ: Any
+    value_dim: int
+
+
+UElement: TypeAlias = Any
+ParamsLike: TypeAlias = Any

{fluxfem-0.1.3 → fluxfem-0.1.4}/src/fluxfem/core/forms.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import jax
 import jax.numpy as jnp
 from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeAlias
 
 from .basis import Basis3D
 
@@ -119,7 +120,10 @@ class VectorFormField:
         return cls(N, elem_coords, aux["basis"], aux["value_dim"], gradN, detJ)
 
 
-FormFieldLike = ScalarFormField | VectorFormField
+if TYPE_CHECKING:
+    FormFieldLike: TypeAlias = ScalarFormField | VectorFormField
+else:
+    FormFieldLike = object
 
 
 def vector_load_form(field: FormFieldLike, load_vec: jnp.ndarray) -> jnp.ndarray: