fluxfem 0.1.1a0__py3-none-any.whl → 0.1.4__py3-none-any.whl

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/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,
@@ -86,6 +88,8 @@ from .weakform import (
     transpose,
     log,
     transpose_last2,
+    matmul,
+    matmul_std,
     einsum,
 )
 from ..mesh import (
@@ -121,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",
@@ -161,6 +182,10 @@ __all__ = [
     "FESpacePytree",
     "FormContext",
     "MixedFormContext",
+    "VolumeContext",
+    "SurfaceContext",
+    "VolumeContext",
+    "SurfaceContext",
     "FieldPair",
     "ElementVector",
     "vector_load_form",
@@ -229,6 +254,7 @@ __all__ = [
     "compile_mixed_residual",
     "grad",
     "sym_grad",
+    "outer",
     "dot",
     "sdot",
     "ddot",
@@ -244,6 +270,8 @@ __all__ = [
     "transpose",
     "log",
     "transpose_last2",
+    "matmul",
+    "matmul_std",
     "einsum",
     "HexMesh",
     "HexMeshPytree",

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/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: