fluxfem 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fluxfem/__init__.py +1 -13
- fluxfem/core/__init__.py +53 -71
- fluxfem/core/assembly.py +41 -32
- fluxfem/core/basis.py +2 -2
- fluxfem/core/context_types.py +36 -12
- fluxfem/core/mixed_space.py +42 -8
- fluxfem/core/mixed_weakform.py +1 -1
- fluxfem/core/space.py +68 -28
- fluxfem/core/weakform.py +95 -77
- fluxfem/mesh/base.py +3 -3
- fluxfem/mesh/contact.py +33 -17
- fluxfem/mesh/io.py +3 -2
- fluxfem/mesh/mortar.py +106 -43
- fluxfem/mesh/supermesh.py +2 -0
- fluxfem/mesh/surface.py +82 -22
- fluxfem/mesh/tet.py +7 -4
- fluxfem/physics/elasticity/hyperelastic.py +32 -3
- fluxfem/physics/elasticity/linear.py +13 -2
- fluxfem/physics/elasticity/stress.py +9 -5
- fluxfem/physics/operators.py +12 -5
- fluxfem/physics/postprocess.py +29 -3
- fluxfem/solver/__init__.py +6 -1
- fluxfem/solver/block_matrix.py +165 -13
- fluxfem/solver/block_system.py +52 -29
- fluxfem/solver/cg.py +43 -30
- fluxfem/solver/dirichlet.py +35 -12
- fluxfem/solver/history.py +15 -3
- fluxfem/solver/newton.py +25 -12
- fluxfem/solver/petsc.py +13 -7
- fluxfem/solver/preconditioner.py +7 -4
- fluxfem/solver/solve_runner.py +42 -24
- fluxfem/solver/solver.py +23 -11
- fluxfem/solver/sparse.py +32 -13
- fluxfem/tools/jit.py +19 -7
- fluxfem/tools/timer.py +14 -12
- fluxfem/tools/visualizer.py +16 -4
- {fluxfem-0.2.0.dist-info → fluxfem-0.2.1.dist-info}/METADATA +18 -7
- fluxfem-0.2.1.dist-info/RECORD +59 -0
- fluxfem-0.2.0.dist-info/RECORD +0 -59
- {fluxfem-0.2.0.dist-info → fluxfem-0.2.1.dist-info}/LICENSE +0 -0
- {fluxfem-0.2.0.dist-info → fluxfem-0.2.1.dist-info}/WHEEL +0 -0
fluxfem/solver/sparse.py
CHANGED
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
from dataclasses import dataclass
|
|
4
|
+
from typing import Any, Iterable, Sequence, TYPE_CHECKING, TypeAlias
|
|
4
5
|
|
|
5
6
|
import numpy as np
|
|
6
7
|
import jax
|
|
@@ -12,7 +13,18 @@ except Exception: # pragma: no cover
|
|
|
12
13
|
sp = None
|
|
13
14
|
|
|
14
15
|
|
|
15
|
-
|
|
16
|
+
if TYPE_CHECKING:
|
|
17
|
+
from jax import Array as JaxArray
|
|
18
|
+
|
|
19
|
+
ArrayLike: TypeAlias = np.ndarray | JaxArray
|
|
20
|
+
else:
|
|
21
|
+
ArrayLike: TypeAlias = np.ndarray
|
|
22
|
+
COOTuple: TypeAlias = tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, int]
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def coalesce_coo(
|
|
26
|
+
rows: ArrayLike, cols: ArrayLike, data: ArrayLike
|
|
27
|
+
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
16
28
|
"""
|
|
17
29
|
Sum duplicate COO entries by sorting (CPU-friendly).
|
|
18
30
|
Returns (rows_u, cols_u, data_u) as NumPy arrays.
|
|
@@ -35,7 +47,7 @@ def coalesce_coo(rows, cols, data):
|
|
|
35
47
|
return r_u, c_u, d_u
|
|
36
48
|
|
|
37
49
|
|
|
38
|
-
def _normalize_flux_mats(mats):
|
|
50
|
+
def _normalize_flux_mats(mats: Sequence["FluxSparseMatrix"]) -> tuple["FluxSparseMatrix", ...]:
|
|
39
51
|
if len(mats) == 1 and isinstance(mats[0], (list, tuple)):
|
|
40
52
|
mats = tuple(mats[0])
|
|
41
53
|
if not mats:
|
|
@@ -43,7 +55,7 @@ def _normalize_flux_mats(mats):
|
|
|
43
55
|
return mats
|
|
44
56
|
|
|
45
57
|
|
|
46
|
-
def concat_flux(*mats, n_dofs: int | None = None):
|
|
58
|
+
def concat_flux(*mats: "FluxSparseMatrix", n_dofs: int | None = None) -> "FluxSparseMatrix":
|
|
47
59
|
"""
|
|
48
60
|
Concatenate COO entries from multiple FluxSparseMatrix objects.
|
|
49
61
|
All matrices must share the same n_dofs unless n_dofs is provided.
|
|
@@ -63,7 +75,7 @@ def concat_flux(*mats, n_dofs: int | None = None):
|
|
|
63
75
|
return FluxSparseMatrix(rows, cols, data, int(n_dofs))
|
|
64
76
|
|
|
65
77
|
|
|
66
|
-
def block_diag_flux(*mats):
|
|
78
|
+
def block_diag_flux(*mats: "FluxSparseMatrix") -> "FluxSparseMatrix":
|
|
67
79
|
"""Block-diagonal concatenation for FluxSparseMatrix objects."""
|
|
68
80
|
mats = _normalize_flux_mats(mats)
|
|
69
81
|
rows_out = []
|
|
@@ -153,7 +165,14 @@ class FluxSparseMatrix:
|
|
|
153
165
|
- data stores the numeric values for the current nonlinear iterate
|
|
154
166
|
"""
|
|
155
167
|
|
|
156
|
-
def __init__(
|
|
168
|
+
def __init__(
|
|
169
|
+
self,
|
|
170
|
+
rows_or_pattern: SparsityPattern | ArrayLike,
|
|
171
|
+
cols: ArrayLike | None = None,
|
|
172
|
+
data: ArrayLike | None = None,
|
|
173
|
+
n_dofs: int | None = None,
|
|
174
|
+
meta: dict | None = None,
|
|
175
|
+
):
|
|
157
176
|
# New signature: FluxSparseMatrix(pattern, data)
|
|
158
177
|
if isinstance(rows_or_pattern, SparsityPattern):
|
|
159
178
|
pattern = rows_or_pattern
|
|
@@ -188,35 +207,35 @@ class FluxSparseMatrix:
|
|
|
188
207
|
self.meta = dict(meta) if meta is not None else None
|
|
189
208
|
|
|
190
209
|
@classmethod
|
|
191
|
-
def from_bilinear(cls, coo_tuple):
|
|
210
|
+
def from_bilinear(cls, coo_tuple: COOTuple) -> "FluxSparseMatrix":
|
|
192
211
|
"""Construct from assemble_bilinear_dense(..., sparse=True)."""
|
|
193
212
|
rows, cols, data, n_dofs = coo_tuple
|
|
194
213
|
return cls(rows, cols, data, n_dofs)
|
|
195
214
|
|
|
196
215
|
@classmethod
|
|
197
|
-
def from_linear(cls, coo_tuple):
|
|
216
|
+
def from_linear(cls, coo_tuple: tuple[jnp.ndarray, jnp.ndarray, int]) -> "FluxSparseMatrix":
|
|
198
217
|
"""Construct from assemble_linear_form(..., sparse=True) (matrix interpretation only)."""
|
|
199
218
|
rows, data, n_dofs = coo_tuple
|
|
200
219
|
cols = jnp.zeros_like(rows)
|
|
201
220
|
return cls(rows, cols, data, n_dofs)
|
|
202
221
|
|
|
203
|
-
def with_data(self, data):
|
|
222
|
+
def with_data(self, data: ArrayLike) -> "FluxSparseMatrix":
|
|
204
223
|
"""Return a new FluxSparseMatrix sharing the same pattern with updated data."""
|
|
205
224
|
return FluxSparseMatrix(self.pattern, data, meta=self.meta)
|
|
206
225
|
|
|
207
|
-
def add_dense(self, dense):
|
|
226
|
+
def add_dense(self, dense: ArrayLike) -> "FluxSparseMatrix":
|
|
208
227
|
"""Return a new FluxSparseMatrix with dense entries added on the pattern."""
|
|
209
228
|
dense_vals = jnp.asarray(dense)[self.pattern.rows, self.pattern.cols]
|
|
210
229
|
return FluxSparseMatrix(self.pattern, self.data + dense_vals)
|
|
211
230
|
|
|
212
|
-
def to_coo(self):
|
|
231
|
+
def to_coo(self) -> COOTuple:
|
|
213
232
|
return self.pattern.rows, self.pattern.cols, self.data, self.pattern.n_dofs
|
|
214
233
|
|
|
215
234
|
@property
|
|
216
235
|
def nnz(self) -> int:
|
|
217
236
|
return int(self.data.shape[0])
|
|
218
237
|
|
|
219
|
-
def coalesce(self):
|
|
238
|
+
def coalesce(self) -> "FluxSparseMatrix":
|
|
220
239
|
"""Return a new FluxSparseMatrix with duplicate entries summed."""
|
|
221
240
|
rows_u, cols_u, data_u = coalesce_coo(self.pattern.rows, self.pattern.cols, self.data)
|
|
222
241
|
return FluxSparseMatrix(rows_u, cols_u, data_u, self.pattern.n_dofs)
|
|
@@ -238,7 +257,7 @@ class FluxSparseMatrix:
|
|
|
238
257
|
d = np.array(self.data, copy=True)
|
|
239
258
|
return sp.csr_matrix((d, (r, c)), shape=(self.pattern.n_dofs, self.pattern.n_dofs))
|
|
240
259
|
|
|
241
|
-
def to_dense(self):
|
|
260
|
+
def to_dense(self) -> jnp.ndarray:
|
|
242
261
|
# small debug helper
|
|
243
262
|
dense = jnp.zeros((self.pattern.n_dofs, self.pattern.n_dofs), dtype=self.data.dtype)
|
|
244
263
|
dense = dense.at[self.pattern.rows, self.pattern.cols].add(self.data)
|
|
@@ -253,7 +272,7 @@ class FluxSparseMatrix:
|
|
|
253
272
|
idx = jnp.stack([self.pattern.rows, self.pattern.cols], axis=-1)
|
|
254
273
|
return jsparse.BCOO((self.data, idx), shape=(self.pattern.n_dofs, self.pattern.n_dofs))
|
|
255
274
|
|
|
256
|
-
def matvec(self, x):
|
|
275
|
+
def matvec(self, x: ArrayLike) -> jnp.ndarray:
|
|
257
276
|
"""Compute y = A x in JAX (iterative solvers)."""
|
|
258
277
|
xj = jnp.asarray(x)
|
|
259
278
|
contrib = self.data * xj[self.pattern.cols]
|
fluxfem/tools/jit.py
CHANGED
|
@@ -1,10 +1,22 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
1
3
|
import jax
|
|
2
4
|
|
|
3
|
-
from
|
|
5
|
+
from typing import Callable, TypeVar
|
|
6
|
+
|
|
7
|
+
from ..core.assembly import JacobianReturn, LinearReturn, ResidualForm, assemble_jacobian, assemble_residual
|
|
4
8
|
from ..core.space import FESpace
|
|
5
9
|
|
|
10
|
+
P = TypeVar("P")
|
|
11
|
+
|
|
6
12
|
|
|
7
|
-
def make_jitted_residual(
|
|
13
|
+
def make_jitted_residual(
|
|
14
|
+
space: FESpace,
|
|
15
|
+
res_form: ResidualForm[P],
|
|
16
|
+
params: P,
|
|
17
|
+
*,
|
|
18
|
+
sparse: bool = False,
|
|
19
|
+
) -> Callable[[jax.Array], LinearReturn]:
|
|
8
20
|
"""
|
|
9
21
|
Create a jitted residual assembler: u -> R(u).
|
|
10
22
|
params and space are closed over.
|
|
@@ -13,7 +25,7 @@ def make_jitted_residual(space: FESpace, res_form, params, *, sparse: bool = Fal
|
|
|
13
25
|
params_jax = params
|
|
14
26
|
|
|
15
27
|
@jax.jit
|
|
16
|
-
def residual(u):
|
|
28
|
+
def residual(u: jax.Array) -> LinearReturn:
|
|
17
29
|
return assemble_residual(space_jax, res_form, u, params_jax, sparse=sparse)
|
|
18
30
|
|
|
19
31
|
return residual
|
|
@@ -21,12 +33,12 @@ def make_jitted_residual(space: FESpace, res_form, params, *, sparse: bool = Fal
|
|
|
21
33
|
|
|
22
34
|
def make_jitted_jacobian(
|
|
23
35
|
space: FESpace,
|
|
24
|
-
res_form,
|
|
25
|
-
params,
|
|
36
|
+
res_form: ResidualForm[P],
|
|
37
|
+
params: P,
|
|
26
38
|
*,
|
|
27
39
|
sparse: bool = False,
|
|
28
40
|
return_flux_matrix: bool = False,
|
|
29
|
-
):
|
|
41
|
+
) -> Callable[[jax.Array], JacobianReturn]:
|
|
30
42
|
"""
|
|
31
43
|
Create a jitted Jacobian assembler: u -> J(u).
|
|
32
44
|
params and space are closed over.
|
|
@@ -35,7 +47,7 @@ def make_jitted_jacobian(
|
|
|
35
47
|
params_jax = params
|
|
36
48
|
|
|
37
49
|
@jax.jit
|
|
38
|
-
def jacobian(u):
|
|
50
|
+
def jacobian(u: jax.Array) -> JacobianReturn:
|
|
39
51
|
return assemble_jacobian(
|
|
40
52
|
space_jax,
|
|
41
53
|
res_form,
|
fluxfem/tools/timer.py
CHANGED
|
@@ -5,13 +5,15 @@ from contextlib import AbstractContextManager
|
|
|
5
5
|
from collections import defaultdict
|
|
6
6
|
from contextlib import contextmanager
|
|
7
7
|
from dataclasses import dataclass
|
|
8
|
-
from typing import Callable, DefaultDict, Dict, Iterator, List, Optional
|
|
8
|
+
from typing import Any, Callable, DefaultDict, Dict, Iterator, List, Optional, TypeAlias
|
|
9
9
|
|
|
10
10
|
import logging
|
|
11
11
|
|
|
12
12
|
logging.basicConfig(level=logging.INFO)
|
|
13
13
|
logger = logging.getLogger(__name__)
|
|
14
14
|
|
|
15
|
+
PlotResult: TypeAlias = tuple[Any, Any]
|
|
16
|
+
|
|
15
17
|
|
|
16
18
|
@dataclass
|
|
17
19
|
class SectionStats:
|
|
@@ -30,7 +32,7 @@ class BaseTimer(ABC):
|
|
|
30
32
|
|
|
31
33
|
class NullTimer(BaseTimer):
|
|
32
34
|
@contextmanager
|
|
33
|
-
def section(self, name: str):
|
|
35
|
+
def section(self, name: str) -> Iterator[None]:
|
|
34
36
|
yield
|
|
35
37
|
|
|
36
38
|
|
|
@@ -85,7 +87,7 @@ class SectionTimer(BaseTimer):
|
|
|
85
87
|
self._records[full_name].append(duration)
|
|
86
88
|
self._stack.pop()
|
|
87
89
|
|
|
88
|
-
def wrap(self, name: str):
|
|
90
|
+
def wrap(self, name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
|
89
91
|
"""
|
|
90
92
|
Decorator form of :meth:`section`.
|
|
91
93
|
|
|
@@ -97,8 +99,8 @@ class SectionTimer(BaseTimer):
|
|
|
97
99
|
...
|
|
98
100
|
"""
|
|
99
101
|
|
|
100
|
-
def _decorator(func):
|
|
101
|
-
def _wrapper(*args, **kwargs):
|
|
102
|
+
def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
103
|
+
def _wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
102
104
|
with self.section(name):
|
|
103
105
|
return func(*args, **kwargs)
|
|
104
106
|
return _wrapper
|
|
@@ -211,7 +213,7 @@ class SectionTimer(BaseTimer):
|
|
|
211
213
|
self,
|
|
212
214
|
sort_by: str = "total",
|
|
213
215
|
descending: bool = True,
|
|
214
|
-
logger_instance=None,
|
|
216
|
+
logger_instance: logging.Logger | None = None,
|
|
215
217
|
) -> str:
|
|
216
218
|
stats = self.summary(sort_by=sort_by, descending=descending)
|
|
217
219
|
if not stats:
|
|
@@ -231,7 +233,7 @@ class SectionTimer(BaseTimer):
|
|
|
231
233
|
|
|
232
234
|
def plot_bar(
|
|
233
235
|
self,
|
|
234
|
-
ax=None,
|
|
236
|
+
ax: Any | None = None,
|
|
235
237
|
sort_by: str = "total",
|
|
236
238
|
value: str = "total",
|
|
237
239
|
descending: bool = True,
|
|
@@ -240,7 +242,7 @@ class SectionTimer(BaseTimer):
|
|
|
240
242
|
stacked_nested: bool = False,
|
|
241
243
|
moving_average: bool = False,
|
|
242
244
|
use_self_time: bool = False,
|
|
243
|
-
):
|
|
245
|
+
) -> PlotResult:
|
|
244
246
|
"""
|
|
245
247
|
Plot timing results as a horizontal bar chart without relying on pyplot state.
|
|
246
248
|
|
|
@@ -417,7 +419,7 @@ class SectionTimer(BaseTimer):
|
|
|
417
419
|
|
|
418
420
|
def plot_pie(
|
|
419
421
|
self,
|
|
420
|
-
ax=None,
|
|
422
|
+
ax: Any | None = None,
|
|
421
423
|
sort_by: str = "total",
|
|
422
424
|
value: str = "total",
|
|
423
425
|
descending: bool = True,
|
|
@@ -430,7 +432,7 @@ class SectionTimer(BaseTimer):
|
|
|
430
432
|
show_total: bool = True,
|
|
431
433
|
moving_average: bool = False,
|
|
432
434
|
use_self_time: bool = False,
|
|
433
|
-
):
|
|
435
|
+
) -> PlotResult:
|
|
434
436
|
"""
|
|
435
437
|
Plot timing results as a pie chart to show relative time share.
|
|
436
438
|
|
|
@@ -552,7 +554,7 @@ class SectionTimer(BaseTimer):
|
|
|
552
554
|
|
|
553
555
|
def plot(
|
|
554
556
|
self,
|
|
555
|
-
ax=None,
|
|
557
|
+
ax: Any | None = None,
|
|
556
558
|
sort_by: str = "total",
|
|
557
559
|
value: str = "total",
|
|
558
560
|
descending: bool = True,
|
|
@@ -563,7 +565,7 @@ class SectionTimer(BaseTimer):
|
|
|
563
565
|
moving_average: bool = False,
|
|
564
566
|
use_self_time: bool = False,
|
|
565
567
|
**kwargs,
|
|
566
|
-
):
|
|
568
|
+
) -> PlotResult:
|
|
567
569
|
"""
|
|
568
570
|
Plot timing results choosing between pie (default) or bar chart.
|
|
569
571
|
|
fluxfem/tools/visualizer.py
CHANGED
|
@@ -3,14 +3,14 @@ from typing import Mapping
|
|
|
3
3
|
|
|
4
4
|
import numpy as np
|
|
5
5
|
|
|
6
|
-
from ..mesh import HexMesh, TetMesh
|
|
6
|
+
from ..mesh import BaseMesh, HexMesh, TetMesh
|
|
7
7
|
|
|
8
8
|
|
|
9
9
|
VTK_HEXAHEDRON = 12
|
|
10
10
|
VTK_TETRA = 10
|
|
11
11
|
|
|
12
12
|
|
|
13
|
-
def _cell_type_for_mesh(mesh):
|
|
13
|
+
def _cell_type_for_mesh(mesh: BaseMesh) -> int:
|
|
14
14
|
if isinstance(mesh, HexMesh):
|
|
15
15
|
return VTK_HEXAHEDRON
|
|
16
16
|
if isinstance(mesh, TetMesh):
|
|
@@ -25,7 +25,13 @@ def _write_dataarray(name: str, data: np.ndarray, ncomp: int = 1) -> str:
|
|
|
25
25
|
return f'<DataArray type="Float32" Name="{name}" format="ascii"{comp_attr}>{values}</DataArray>'
|
|
26
26
|
|
|
27
27
|
|
|
28
|
-
def write_vtu(
|
|
28
|
+
def write_vtu(
|
|
29
|
+
mesh: BaseMesh,
|
|
30
|
+
filepath: str,
|
|
31
|
+
*,
|
|
32
|
+
point_data: Mapping[str, np.ndarray] | None = None,
|
|
33
|
+
cell_data: Mapping[str, np.ndarray] | None = None,
|
|
34
|
+
) -> None:
|
|
29
35
|
"""
|
|
30
36
|
Write an UnstructuredGrid VTU for HexMesh or TetMesh.
|
|
31
37
|
point_data/cell_data: dict name -> ndarray. Point data length must match n_points;
|
|
@@ -87,7 +93,13 @@ def write_vtu(mesh, filepath: str, *, point_data: Mapping[str, np.ndarray] | Non
|
|
|
87
93
|
f.write("\n".join(lines))
|
|
88
94
|
|
|
89
95
|
|
|
90
|
-
def write_displacement_vtu(
|
|
96
|
+
def write_displacement_vtu(
|
|
97
|
+
mesh: BaseMesh,
|
|
98
|
+
u: np.ndarray,
|
|
99
|
+
filepath: str,
|
|
100
|
+
*,
|
|
101
|
+
name: str = "displacement",
|
|
102
|
+
) -> None:
|
|
91
103
|
"""
|
|
92
104
|
Convenience wrapper: reshape displacement vector to point data and write VTU.
|
|
93
105
|
Assumes 3 dof/node ordering [u0,v0,w0, u1,v1,w1, ...].
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: fluxfem
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.1
|
|
4
4
|
Summary: FluxFEM: A weak-form-centric differentiable finite element framework in JAX
|
|
5
5
|
License: Apache-2.0
|
|
6
6
|
Author: Kohei Watanabe
|
|
@@ -94,9 +94,9 @@ import jax.numpy as jnp
|
|
|
94
94
|
|
|
95
95
|
space = ff.make_hex_space(mesh, dim=1, intorder=2)
|
|
96
96
|
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
97
|
+
# bilinear: kernel(ctx) -> (n_ldofs, n_ldofs)
|
|
98
|
+
ker_K = ff.make_element_bilinear_kernel(ff.diffusion_form, 1.0, jit=True)
|
|
99
|
+
K = space.assemble(ff.diffusion_form, 1.0, kernel=ker_K)
|
|
100
100
|
|
|
101
101
|
# linear: kernel(ctx) -> (n_ldofs,)
|
|
102
102
|
def linear_kernel(ctx):
|
|
@@ -104,8 +104,8 @@ def linear_kernel(ctx):
|
|
|
104
104
|
wJ = ctx.w * ctx.test.detJ
|
|
105
105
|
return (integrand * wJ[:, None]).sum(axis=0)
|
|
106
106
|
|
|
107
|
-
|
|
108
|
-
|
|
107
|
+
ker_F = jax.jit(linear_kernel)
|
|
108
|
+
F = space.assemble(ff.scalar_body_force_form, 2.0, kernel=ker_F)
|
|
109
109
|
```
|
|
110
110
|
|
|
111
111
|
### tensor-based vs weak-form-based (diffusion example)
|
|
@@ -182,7 +182,7 @@ def neo_hookean_residual_wf(v, u, params):
|
|
|
182
182
|
dE = 0.5 * (h_wf.matmul(h_wf.grad(v), F) + h_wf.transpose(h_wf.matmul(h_wf.grad(v), F)))
|
|
183
183
|
return h_wf.ddot(S, dE) * h_wf.dOmega()
|
|
184
184
|
|
|
185
|
-
|
|
185
|
+
res_form = ff.ResidualForm.volume(neo_hookean_residual_wf).get_compiled()
|
|
186
186
|
```
|
|
187
187
|
|
|
188
188
|
|
|
@@ -203,6 +203,14 @@ loss_theta_jit = jax.jit(loss_theta)
|
|
|
203
203
|
grad_fn = jax.jit(jax.grad(loss_theta))
|
|
204
204
|
```
|
|
205
205
|
|
|
206
|
+
### FESpace vs FESpacePytree
|
|
207
|
+
|
|
208
|
+
Use `FESpace` for standard workflows with a fixed mesh. When you need to carry
|
|
209
|
+
the space through JAX transformations (e.g., shape optimization where mesh
|
|
210
|
+
coordinates are part of the computation), use `FESpacePytree` via
|
|
211
|
+
`make_*_space_pytree(...)`. This keeps the mesh/basis in the pytree so
|
|
212
|
+
`jax.jit`/`jax.grad` can see geometry changes.
|
|
213
|
+
|
|
206
214
|
### Mixed systems
|
|
207
215
|
|
|
208
216
|
Mixed problems can be assembled from residual blocks and solved as a coupled system.
|
|
@@ -243,6 +251,9 @@ blocks = ff_solver.make_block_matrix(
|
|
|
243
251
|
symmetric=True,
|
|
244
252
|
transpose_rule="T",
|
|
245
253
|
)
|
|
254
|
+
|
|
255
|
+
# Lazy container; assemble when you need the global matrix.
|
|
256
|
+
K = blocks.assemble()
|
|
246
257
|
```
|
|
247
258
|
|
|
248
259
|
FluxFEM also provides contact utilities like `ContactSurfaceSpace` to build constraint contributions.
|
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
fluxfem/__init__.py,sha256=hHJEJ4P89uDvziScw-hzs82MIt-K5OBGFgF5ojoI9B4,7846
|
|
2
|
+
fluxfem/core/__init__.py,sha256=av60HQ0K8oaLko0vKc_GNsNs35TayR8BPGp463GKFN4,10898
|
|
3
|
+
fluxfem/core/assembly.py,sha256=CTgI5KZhXx5PnvPHMZQRKg0GdhwqNm4EvPggF_SHwAY,46971
|
|
4
|
+
fluxfem/core/basis.py,sha256=Awru1LcXgukQsP5scWn6i6PsqszvIFa_k7UM-ShbThQ,32435
|
|
5
|
+
fluxfem/core/context_types.py,sha256=67Q_qPG8DZ6B6pYN-RzibhNsHBZfhQtkX0kfnfW46Qc,1405
|
|
6
|
+
fluxfem/core/data.py,sha256=_byAZTasIGGwNRU3gqw_ZZR1HMVLJpk9gix2AI3AJWk,1818
|
|
7
|
+
fluxfem/core/dtypes.py,sha256=FnJ98C3yIibaatN-wmLO3ksI6XabA7G3rI_SOxPmko8,251
|
|
8
|
+
fluxfem/core/forms.py,sha256=uC5WxG3h3G7I9FxNRVvA8JLVpJP-QjNZtj8fjxavqVE,7697
|
|
9
|
+
fluxfem/core/interp.py,sha256=v9HzkYgQnEqf6GnB9tXqfrNk0N7YsmPcZBfNGNUU3b8,2004
|
|
10
|
+
fluxfem/core/mixed_assembly.py,sha256=4lL15AsKV4KA-D_mhwc_zXyIhff1fwMkYj4iPWDbDwk,8513
|
|
11
|
+
fluxfem/core/mixed_space.py,sha256=VbQyz1iwrk5-UfvxNmlVHS1mPy5WnCTSzBQXNUbMA68,14203
|
|
12
|
+
fluxfem/core/mixed_weakform.py,sha256=7WmleuB9-o2PFr4Mbwh7Dk19PZ2VNudWoFXlugACd80,2878
|
|
13
|
+
fluxfem/core/solver.py,sha256=OOKiEsW9j4BPCtwgIYyaUHh0Bw5dAd975aGl2izOYqQ,3772
|
|
14
|
+
fluxfem/core/space.py,sha256=DQ2hvqz226lnrKlIHc617LGiwdup_6pgfCx53DgdjWA,23606
|
|
15
|
+
fluxfem/core/weakform.py,sha256=54ntcT-ysR3sVwRS2bgVeP6AslMgIgxVkwi3SrMAKiE,71229
|
|
16
|
+
fluxfem/helpers_ts.py,sha256=T3HCqXniF116lvznq0b7CDUok5_yZZc8l9XtlE81yvg,221
|
|
17
|
+
fluxfem/helpers_wf.py,sha256=pA8bprvbdXZL_qourm07xaJAsGyNnAP8wXmJf8fOj-k,2158
|
|
18
|
+
fluxfem/mesh/__init__.py,sha256=u2CzTjFxj_CRbi18FezIDh73k_lan_3wGBifTkdJH4w,2418
|
|
19
|
+
fluxfem/mesh/base.py,sha256=CAZwHp4lQzVKQe2TiWKG2-ktXWuC7AAURxIPR9pI7rU,22750
|
|
20
|
+
fluxfem/mesh/contact.py,sha256=8Uy2UGaUB7z_3OuRJ8Jdn7DET1Bq-X0GGEV9bn0LYtA,30059
|
|
21
|
+
fluxfem/mesh/dtypes.py,sha256=FnJ98C3yIibaatN-wmLO3ksI6XabA7G3rI_SOxPmko8,251
|
|
22
|
+
fluxfem/mesh/hex.py,sha256=VhH0THJMHPhNtwAvqBjIaeNMzBUd8SWzkih2FcR1NWc,12455
|
|
23
|
+
fluxfem/mesh/io.py,sha256=j6EM7AQ9xiDGAc8I_q9t5vkLIGuuezF3uc1kTvCXmb0,3119
|
|
24
|
+
fluxfem/mesh/mortar.py,sha256=nfy8TyncDmyAgx71E2XNEu4Iv5joLU9Wu8vWVdV3WeM,151395
|
|
25
|
+
fluxfem/mesh/predicate.py,sha256=7O9kcKYg_OPtwhqYVdA-j9LQYtcG2Nl781UupKm6U-A,1505
|
|
26
|
+
fluxfem/mesh/supermesh.py,sha256=bKHERAB4vMkeiuT6wFAoKPKbv-O-8C2crrvJkaRW1ak,10463
|
|
27
|
+
fluxfem/mesh/surface.py,sha256=FJgIvLS8WJElpA7dL9_D2wMh7NZgE0gtfLCBw3g8NJY,11179
|
|
28
|
+
fluxfem/mesh/tet.py,sha256=8ijEdjxOrO4mghIApKqIkOwgcIJKNSIL65uGuYHF1-s,10191
|
|
29
|
+
fluxfem/physics/__init__.py,sha256=gUrVF1FSiSVZezqRGvKSTcZDH79Ty3opiJ9-H1gycqY,1426
|
|
30
|
+
fluxfem/physics/diffusion.py,sha256=T3jIGqnuaoSMB8q9x6qNwY23V55-lQGb6ome_WQnOS8,537
|
|
31
|
+
fluxfem/physics/elasticity/__init__.py,sha256=4pAqAstT9T20pY5q417c9FQAarq_T8g2hpqDtrUDkcw,1061
|
|
32
|
+
fluxfem/physics/elasticity/hyperelastic.py,sha256=S61yQYzu5Njowx9QF2qsMOr8pY7Mz5dIhsW-GPev-90,3964
|
|
33
|
+
fluxfem/physics/elasticity/linear.py,sha256=mUayoQ4buWCkqyWC7NfanNjYtoNFiVyzaQjHdfzEIjs,2413
|
|
34
|
+
fluxfem/physics/elasticity/materials.py,sha256=f2y7_nY3KMSmD2-G4znondAOvh-XDvJD_-zcE2EZcKw,977
|
|
35
|
+
fluxfem/physics/elasticity/stress.py,sha256=1JY4ZKZglcst3Tot7uOvAsaGydJUuwRFFvaft3Fghg8,1181
|
|
36
|
+
fluxfem/physics/operators.py,sha256=cuMNoUFSnBtyjCbR-VANrJapvwWVObVQZDCm0ot5L4Q,3289
|
|
37
|
+
fluxfem/physics/postprocess.py,sha256=A6VQq_XBKd75rZbP_mqLR--zUd3Dw_XQfjYX4PID1_M,4860
|
|
38
|
+
fluxfem/solver/__init__.py,sha256=KIQ0RTyjer-flCewrTGu80IX6d1OKnO1HwIn2JmQVL0,2492
|
|
39
|
+
fluxfem/solver/bc.py,sha256=IJa9WmAkeLu73fmxppEZRXfB-1o_B-9tQ8uKXZHdxFU,16649
|
|
40
|
+
fluxfem/solver/block_matrix.py,sha256=9egC5bZcgx3sDUJpAt4jEPy_GgQ-hArvPWSFLnFos18,10825
|
|
41
|
+
fluxfem/solver/block_system.py,sha256=Fy1AuALQZg7uxpseGVDrzRGBf9EYAjOi0Sj-3rDsAH0,18620
|
|
42
|
+
fluxfem/solver/cg.py,sha256=scydBIRdHXG9QOtRQfQG9pL4ITxjFaTIkdYPHjQEZgM,11826
|
|
43
|
+
fluxfem/solver/dirichlet.py,sha256=0ERvZd1XBySQAWV9n_QY2VZN8xOdaUNOgL-vn2DwUnI,15560
|
|
44
|
+
fluxfem/solver/history.py,sha256=y-37QPJqPFKFBRdWUrng6OwYyHZRRcPaRtF8yuqpnDc,1034
|
|
45
|
+
fluxfem/solver/newton.py,sha256=Q8y2e0M05fgKVx8VCnths0XtWkLSVEbea2dU2hdLJvQ,24184
|
|
46
|
+
fluxfem/solver/petsc.py,sha256=Kk4vxQrfxD5D_wc3z8dFHg-77KH_FaMY2JB4oCM5wLM,15221
|
|
47
|
+
fluxfem/solver/preconditioner.py,sha256=ijnhSA54rUeN1KCpW_SJFlANsR2dccPnC5-n6VEg0MU,3781
|
|
48
|
+
fluxfem/solver/result.py,sha256=DLLyaHVJ9NYotDqFD0IHlhxg9Hbs726bepqCG_59t94,3256
|
|
49
|
+
fluxfem/solver/solve_runner.py,sha256=uBwZpwIeyYXhXFE6Wr1qlywl4z9ExMCoG4Sagsn_D8U,31413
|
|
50
|
+
fluxfem/solver/solver.py,sha256=pkxYOncGkK1t1L9GcZwW4BwDsc31StUFplIK3QS_cYo,6157
|
|
51
|
+
fluxfem/solver/sparse.py,sha256=8Ji1xkWPqLfy2xZv1ejhJhxILeH8G1peh5Lf_X16oxw,12757
|
|
52
|
+
fluxfem/tools/__init__.py,sha256=S7flze8kfHABc0fdhGgdELzKO8aeBCrDbwRUE3aTGpA,95
|
|
53
|
+
fluxfem/tools/jit.py,sha256=x8MS-RWJfafT9A5bhPLa80quR0dYREvikfmWqNHbxBE,1445
|
|
54
|
+
fluxfem/tools/timer.py,sha256=pCL2T4Vi2M00wi5eVey75QGglZVRiZW3i31rL2ll25U,23145
|
|
55
|
+
fluxfem/tools/visualizer.py,sha256=x7i2_ABpXEP4ixPSiqStNM8lR0HZbaAuViCSF7W-p-U,3854
|
|
56
|
+
fluxfem-0.2.1.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
|
|
57
|
+
fluxfem-0.2.1.dist-info/METADATA,sha256=Mi6PAmMFKYpcv9c0NOrcXx6V-DvMYEF45SF0RX_RlNM,11366
|
|
58
|
+
fluxfem-0.2.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
|
|
59
|
+
fluxfem-0.2.1.dist-info/RECORD,,
|
fluxfem-0.2.0.dist-info/RECORD
DELETED
|
@@ -1,59 +0,0 @@
|
|
|
1
|
-
fluxfem/__init__.py,sha256=KPEgeWjheIAHQ28MWazyzLCN-lR0khSRltq06Q9tlUg,8022
|
|
2
|
-
fluxfem/core/__init__.py,sha256=BKhJR8dK8oZkX9yW0W8m7u5nHDhnwmUg0wWF3iL4Hm8,10964
|
|
3
|
-
fluxfem/core/assembly.py,sha256=rBQ3RPEVxfCuCnKc9Xxz6vLrRZ9ROO8JNggfV4uGWTc,46070
|
|
4
|
-
fluxfem/core/basis.py,sha256=bh1evqiS5tNG7bqyq7ZzipQR8SBEIM-B7MzJ6_niK0A,32397
|
|
5
|
-
fluxfem/core/context_types.py,sha256=QGv1oc_9Crn3vztyPy_uUN0TpmDkgH5HPrhYzZ_kuys,663
|
|
6
|
-
fluxfem/core/data.py,sha256=_byAZTasIGGwNRU3gqw_ZZR1HMVLJpk9gix2AI3AJWk,1818
|
|
7
|
-
fluxfem/core/dtypes.py,sha256=FnJ98C3yIibaatN-wmLO3ksI6XabA7G3rI_SOxPmko8,251
|
|
8
|
-
fluxfem/core/forms.py,sha256=uC5WxG3h3G7I9FxNRVvA8JLVpJP-QjNZtj8fjxavqVE,7697
|
|
9
|
-
fluxfem/core/interp.py,sha256=v9HzkYgQnEqf6GnB9tXqfrNk0N7YsmPcZBfNGNUU3b8,2004
|
|
10
|
-
fluxfem/core/mixed_assembly.py,sha256=4lL15AsKV4KA-D_mhwc_zXyIhff1fwMkYj4iPWDbDwk,8513
|
|
11
|
-
fluxfem/core/mixed_space.py,sha256=8LZU2VibMdHP4tq_1wrf1rY1LVZo6zZiw8UE6cPaqus,13285
|
|
12
|
-
fluxfem/core/mixed_weakform.py,sha256=1RizZ1KuyCGse3LDbHwCcMHP0YjVKBrwE_USm7COIoQ,2848
|
|
13
|
-
fluxfem/core/solver.py,sha256=OOKiEsW9j4BPCtwgIYyaUHh0Bw5dAd975aGl2izOYqQ,3772
|
|
14
|
-
fluxfem/core/space.py,sha256=XHs0NamEBVNyvaXU6FdIY9xgeY3VU8jiy8cX8qZnOyk,22159
|
|
15
|
-
fluxfem/core/weakform.py,sha256=Sgov4zXfEtYtUASDcy5g1bZu0XsnBMUtQuw1bwj6Abg,70163
|
|
16
|
-
fluxfem/helpers_ts.py,sha256=T3HCqXniF116lvznq0b7CDUok5_yZZc8l9XtlE81yvg,221
|
|
17
|
-
fluxfem/helpers_wf.py,sha256=pA8bprvbdXZL_qourm07xaJAsGyNnAP8wXmJf8fOj-k,2158
|
|
18
|
-
fluxfem/mesh/__init__.py,sha256=u2CzTjFxj_CRbi18FezIDh73k_lan_3wGBifTkdJH4w,2418
|
|
19
|
-
fluxfem/mesh/base.py,sha256=ghsvH8kvwpprmR6x1HkSG0H6QTOpl0CVk4HCBDSIdPY,22716
|
|
20
|
-
fluxfem/mesh/contact.py,sha256=WDwwB0s_Zbd7qAEx4eLGr9_zn2KKh5SS0tfTOgnSTRU,28921
|
|
21
|
-
fluxfem/mesh/dtypes.py,sha256=FnJ98C3yIibaatN-wmLO3ksI6XabA7G3rI_SOxPmko8,251
|
|
22
|
-
fluxfem/mesh/hex.py,sha256=VhH0THJMHPhNtwAvqBjIaeNMzBUd8SWzkih2FcR1NWc,12455
|
|
23
|
-
fluxfem/mesh/io.py,sha256=Nlg0urKVDluX1r2guMjvPWNL-SlS2lOc4eLKuKLex9s,3044
|
|
24
|
-
fluxfem/mesh/mortar.py,sha256=-OhyHkaFEYCOoIvf8GpNUJB0lMGQe7hBGzEQ3rV7zYk,147865
|
|
25
|
-
fluxfem/mesh/predicate.py,sha256=7O9kcKYg_OPtwhqYVdA-j9LQYtcG2Nl781UupKm6U-A,1505
|
|
26
|
-
fluxfem/mesh/supermesh.py,sha256=DNDE3-Tck9KSWIbOWGJ7Dduuag9ofZF_B7PqjNecjtA,10407
|
|
27
|
-
fluxfem/mesh/surface.py,sha256=-xNhs3pLnU5-a8ct8hQTd8YzvjETgTRA3tYK9Dwb-Fw,9791
|
|
28
|
-
fluxfem/mesh/tet.py,sha256=Q-2S-b3N4CF3i40u5tt5mbPn_fYcvCJx4K6Xiyq11-A,10077
|
|
29
|
-
fluxfem/physics/__init__.py,sha256=gUrVF1FSiSVZezqRGvKSTcZDH79Ty3opiJ9-H1gycqY,1426
|
|
30
|
-
fluxfem/physics/diffusion.py,sha256=T3jIGqnuaoSMB8q9x6qNwY23V55-lQGb6ome_WQnOS8,537
|
|
31
|
-
fluxfem/physics/elasticity/__init__.py,sha256=4pAqAstT9T20pY5q417c9FQAarq_T8g2hpqDtrUDkcw,1061
|
|
32
|
-
fluxfem/physics/elasticity/hyperelastic.py,sha256=TvclMK050_h8cO97b3vx2Br8QOoKYzzDfM4nbUfOfL8,3461
|
|
33
|
-
fluxfem/physics/elasticity/linear.py,sha256=JU_LiQoYbABMrJzvFRHKnnCT44iK4WQKZxRu1GLxyq8,2255
|
|
34
|
-
fluxfem/physics/elasticity/materials.py,sha256=f2y7_nY3KMSmD2-G4znondAOvh-XDvJD_-zcE2EZcKw,977
|
|
35
|
-
fluxfem/physics/elasticity/stress.py,sha256=i3TypnLB-Hhd3UsV2tftOCjF1-KNNpnVEy8rYqT2wIM,1135
|
|
36
|
-
fluxfem/physics/operators.py,sha256=LgEvtcDvhskd7BokVsX_wtzceoI2zrLYyB5HWsUBFg4,3009
|
|
37
|
-
fluxfem/physics/postprocess.py,sha256=7NKcSJO_Himw1WTJgM4YWxf1KkjGuBwojxY_0Vwy9Ow,4418
|
|
38
|
-
fluxfem/solver/__init__.py,sha256=gv57yXua8RSsubZ5SbTaXnKALnRdsKOSuCBR7nuBiCo,2337
|
|
39
|
-
fluxfem/solver/bc.py,sha256=IJa9WmAkeLu73fmxppEZRXfB-1o_B-9tQ8uKXZHdxFU,16649
|
|
40
|
-
fluxfem/solver/block_matrix.py,sha256=qJ6rOtiJyxYlSHkzYbi-T_BN7Xfd0drNzXrM8mdq8uo,4187
|
|
41
|
-
fluxfem/solver/block_system.py,sha256=r5nvt2OWGS9iYq0n2G2HB3MM--ThwGsGoUwcM-w6ZFw,17659
|
|
42
|
-
fluxfem/solver/cg.py,sha256=ZP4OcFeB-froIXcg_PVEobf5BYKz5gONALcEpKCbr_k,11101
|
|
43
|
-
fluxfem/solver/dirichlet.py,sha256=tmKpXWP_8v_jIb337dTkvDxsgkOcoE8ViKuZajQn8L4,14875
|
|
44
|
-
fluxfem/solver/history.py,sha256=HhKSmo2B1GaBYEsrifxQxeaeRMiaqm5f6g6DNz5Qb0U,764
|
|
45
|
-
fluxfem/solver/newton.py,sha256=fVqASbzXrQb9ZqjKol2lWioEr7coi4sWaYMGvYDztqc,23444
|
|
46
|
-
fluxfem/solver/petsc.py,sha256=PjCRomFU3BWjwCO_ZAa8dS5sh_OhXvNSZykURmSroRY,14910
|
|
47
|
-
fluxfem/solver/preconditioner.py,sha256=PdjaHUVgwHGUSO29Od04aQtr24cq9XX9atgClfKKZ-8,3578
|
|
48
|
-
fluxfem/solver/result.py,sha256=DLLyaHVJ9NYotDqFD0IHlhxg9Hbs726bepqCG_59t94,3256
|
|
49
|
-
fluxfem/solver/solve_runner.py,sha256=OuEQlxSIhp9ahaRzS45ZWqgq1ztkf-oDcvp394nGZ7g,30417
|
|
50
|
-
fluxfem/solver/solver.py,sha256=sU9z4btnX5pPUAzxFLRusKNGsGg3pu1njRqfkRjQE9I,5640
|
|
51
|
-
fluxfem/solver/sparse.py,sha256=Znn50ncFDrqoWNOwxIvXmcK5WpDtLy5ktvBuavkAFFU,11875
|
|
52
|
-
fluxfem/tools/__init__.py,sha256=S7flze8kfHABc0fdhGgdELzKO8aeBCrDbwRUE3aTGpA,95
|
|
53
|
-
fluxfem/tools/jit.py,sha256=WVThoNs2XD8ZVCeyLjBgbFaZXHot5RK_RNzvvAjSaew,1110
|
|
54
|
-
fluxfem/tools/timer.py,sha256=fvMrKqjOT2S_DoGlTYC5L5TtBnoLS4zyiAZtJVJGt0o,22848
|
|
55
|
-
fluxfem/tools/visualizer.py,sha256=aDqSgtX9xOu120-zCzJfwlCOUmxjL_kJ6kAvj27-H14,3733
|
|
56
|
-
fluxfem-0.2.0.dist-info/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
|
|
57
|
-
fluxfem-0.2.0.dist-info/METADATA,sha256=3RicCvbWhQOUTeQoMgmpYz-q6EznEmYcCn9h63B-Dv0,10931
|
|
58
|
-
fluxfem-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
|
|
59
|
-
fluxfem-0.2.0.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|