fluxfem 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

fluxfem/mesh/dtypes.py

@@ -0,0 +1,12 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def default_dtype() -> jnp.dtype:
+    return jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
+
+
+DEFAULT_DTYPE = default_dtype()
+INDEX_DTYPE = jnp.int64
+NP_INDEX_DTYPE = np.int64

fluxfem/mesh/hex.py

@@ -5,9 +5,9 @@ from dataclasses import dataclass
 from typing import Optional, Dict, Tuple, List, Callable
 import jax
 import jax.numpy as jnp
-from ..core.dtypes import DEFAULT_DTYPE
 import numpy as np
 
+from .dtypes import INDEX_DTYPE, default_dtype
 from .base import BaseMesh, BaseMeshPytree
 
 
@@ -17,10 +17,10 @@ class HexMesh(BaseMesh):
     Structured / unstructured hex mesh (8-node linear hex elements).
     
     coords: (n_nodes, 3) float32
-    conn:   (n_elems, 8) int32  # node indices of each element
+    conn:   (n_elems, 8) int64  # node indices of each element
     """
     coords: jnp.ndarray  # shape (n_nodes, 3)
-    conn: jnp.ndarray    # shape (n_elems, 8), int32
+    conn: jnp.ndarray    # shape (n_elems, 8), int64
 
     def face_node_patterns(self):
         return [
@@ -36,6 +36,7 @@ class HexMesh(BaseMesh):
 @jax.tree_util.register_pytree_node_class
 @dataclass(eq=False)
 class HexMeshPytree(BaseMeshPytree):
+    """Hex mesh registered as a JAX pytree."""
     def face_node_patterns(self):
         return [
             (0, 1, 2, 3),  # -z
@@ -58,8 +59,8 @@ def tag_axis_minmax_facets(
     Tag boundary facets on min/max of the given axis.
 
     Returns:
-        facets: (n_facets, 4) int32, quad node ids
-        facet_tags: (n_facets,) int32, dirichlet_tag on min side, neumann_tag on max side
+        facets: (n_facets, 4) int64, quad node ids
+        facet_tags: (n_facets,) int64, dirichlet_tag on min side, neumann_tag on max side
     """
     coords = np.asarray(mesh.coords)
     conn = np.asarray(mesh.conn)
@@ -98,7 +99,7 @@ def tag_axis_minmax_facets(
                 facet_map[key] = (nodes, tag)
 
     if not facet_map:
-        return jnp.empty((0, 4), dtype=jnp.int32), jnp.empty((0,), dtype=jnp.int32)
+        return jnp.empty((0, 4), dtype=INDEX_DTYPE), jnp.empty((0,), dtype=INDEX_DTYPE)
 
     facets = []
     tags = []
@@ -106,7 +107,7 @@ def tag_axis_minmax_facets(
         facets.append(nodes)
         tags.append(tag)
 
-    return jnp.array(facets, dtype=jnp.int32), jnp.array(tags, dtype=jnp.int32)
+    return jnp.array(facets, dtype=INDEX_DTYPE), jnp.array(tags, dtype=INDEX_DTYPE)
 
 
 @dataclass
@@ -134,9 +135,10 @@ class StructuredHexBox:
             raise ValueError("order must be 1, 2, or 3")
 
         ox, oy, oz = self.origin
-        xs = jnp.linspace(ox, ox + self.lx, self.nx + 1, dtype=DEFAULT_DTYPE)
-        ys = jnp.linspace(oy, oy + self.ly, self.ny + 1, dtype=DEFAULT_DTYPE)
-        zs = jnp.linspace(oz, oz + self.lz, self.nz + 1, dtype=DEFAULT_DTYPE)
+        dtype = default_dtype()
+        xs = jnp.linspace(ox, ox + self.lx, self.nx + 1, dtype=dtype)
+        ys = jnp.linspace(oy, oy + self.ly, self.ny + 1, dtype=dtype)
+        zs = jnp.linspace(oz, oz + self.lz, self.nz + 1, dtype=dtype)
 
         if self.order == 1:
             return self._build_hex8(xs, ys, zs)
@@ -150,7 +152,7 @@ class StructuredHexBox:
             for j in range(self.ny + 1):
                 for i in range(self.nx + 1):
                     coords_list.append([xs[i], ys[j], zs[k]])
-        coords = jnp.array(coords_list, dtype=DEFAULT_DTYPE)
+        coords = jnp.array(coords_list, dtype=default_dtype())
 
         def node_id(i: int, j: int, k: int) -> int:
             return k * (self.ny + 1) * (self.nx + 1) + j * (self.nx + 1) + i
@@ -169,7 +171,7 @@ class StructuredHexBox:
                     n011 = node_id(i,     j + 1, k + 1)
                     conn_list.append([n000, n100, n110, n010, n001, n101, n111, n011])
 
-        conn = jnp.array(conn_list, dtype=jnp.int32)
+        conn = jnp.array(conn_list, dtype=INDEX_DTYPE)
         return HexMesh(coords=coords, conn=conn)
 
     def _build_hex20(self, xs, ys, zs) -> HexMesh:
@@ -240,8 +242,8 @@ class StructuredHexBox:
                         ]
                     )
 
-        coords = jnp.array(coords_list, dtype=DEFAULT_DTYPE)
-        conn = jnp.array(conn_list, dtype=jnp.int32)
+        coords = jnp.array(coords_list, dtype=default_dtype())
+        conn = jnp.array(conn_list, dtype=INDEX_DTYPE)
         return HexMesh(coords=coords, conn=conn)
 
     def _build_hex27(self, xs, ys, zs) -> HexMesh:
@@ -322,6 +324,6 @@ class StructuredHexBox:
                     # order in lexicographic k,j,i -> length 27
                     conn_list.append(nodes)
 
-        coords = jnp.array(coords_list, dtype=DEFAULT_DTYPE)
-        conn = jnp.array(conn_list, dtype=jnp.int32)
+        coords = jnp.array(coords_list, dtype=default_dtype())
+        conn = jnp.array(conn_list, dtype=INDEX_DTYPE)
         return HexMesh(coords=coords, conn=conn)

fluxfem/mesh/io.py

@@ -4,6 +4,8 @@ import numpy as np
 import jax
 import jax.numpy as jnp
 
+from .dtypes import NP_INDEX_DTYPE
+
 DTYPE = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32
 
 try:
@@ -36,10 +38,10 @@ def load_gmsh_mesh(path: str):
 
     mesh = None
     if "hexahedron" in msh.cells_dict:
-        conn = np.asarray(msh.cells_dict["hexahedron"], dtype=np.int32)
+        conn = np.asarray(msh.cells_dict["hexahedron"], dtype=NP_INDEX_DTYPE)
         mesh = HexMesh(jnp.asarray(coords), jnp.asarray(conn))
     elif "tetra" in msh.cells_dict:
-        conn = np.asarray(msh.cells_dict["tetra"], dtype=np.int32)
+        conn = np.asarray(msh.cells_dict["tetra"], dtype=NP_INDEX_DTYPE)
         mesh = TetMesh(jnp.asarray(coords), jnp.asarray(conn))
     else:
         raise ValueError("gmsh mesh does not contain hexahedron or tetra cells")
@@ -50,7 +52,7 @@ def load_gmsh_mesh(path: str):
     cell_dict = msh.cells_dict
     for fkey in ("quad", "triangle", "tri"):
         if fkey in cell_dict:
-            facets = np.asarray(cell_dict[fkey], dtype=np.int32)
+            facets = np.asarray(cell_dict[fkey], dtype=NP_INDEX_DTYPE)
             break
     tags_raw = None
     if facets is not None:
@@ -63,12 +65,13 @@ def load_gmsh_mesh(path: str):
                         tags_raw = data
                         break
         if tags_raw is not None:
-            facet_tags = np.asarray(tags_raw, dtype=np.int32)
+            facet_tags = np.asarray(tags_raw, dtype=NP_INDEX_DTYPE)
 
     return mesh, facets, facet_tags
 
 
 def load_gmsh_hex_mesh(path: str):
+    """Load a Gmsh mesh and return a HexMesh with optional facets/tags."""
     mesh, facets, tags = load_gmsh_mesh(path)
     if not isinstance(mesh, HexMesh):
         raise ValueError("gmsh mesh is not hexahedral")
@@ -76,6 +79,7 @@ def load_gmsh_hex_mesh(path: str):
 
 
 def load_gmsh_tet_mesh(path: str):
+    """Load a Gmsh mesh and return a TetMesh with optional facets/tags."""
     mesh, facets, tags = load_gmsh_mesh(path)
     if not isinstance(mesh, TetMesh):
         raise ValueError("gmsh mesh is not tetrahedral")