Skip to content

API Reference

This section contains the automatically generated API reference. Each module is documented with mkdocstrings and the examples included in the docstrings are executed as doctests during continuous integration to ensure they remain correct.

Light-weight, differentiable finite element analysis package.

Finite element primitives used by :mod:tofea.

Element dataclass

Base dataclass for finite elements.

Source code in tofea/elements.py 
10
11
12
13
14
15
16
17
18
@dataclass(frozen=True, slots=True, kw_only=True)
class Element:
    """Base dataclass for finite elements."""

    dx: float = 0.5
    dy: float = 0.5
    dz: float = 0.5
    eps: float = 1e-6
    dtype: type = np.float64

Q4Element dataclass

Bases: Element

Four-node quadrilateral element.

This class provides utilities shared by all quadrilateral elements implemented in :mod:tofea, namely Gauss quadrature points and the gradients of the bilinear shape functions. Having these helpers in a common base class avoids code duplication in the concrete element implementations and gives this abstraction a clear purpose.

Source code in tofea/elements.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class Q4Element(Element):
    """Four-node quadrilateral element.

    This class provides utilities shared by all quadrilateral elements
    implemented in :mod:`tofea`, namely Gauss quadrature points and the
    gradients of the bilinear shape functions.  Having these helpers in a
    common base class avoids code duplication in the concrete element
    implementations and gives this abstraction a clear purpose.
    """

    @staticmethod
    def gauss_points() -> list[tuple[float, float]]:
        """Return the points for 2x2 Gauss quadrature."""

        gp = 1 / np.sqrt(3)
        return [(-gp, -gp), (gp, -gp), (gp, gp), (-gp, gp)]

    def grad_shape_funcs(self, xi: float, eta: float) -> tuple[NDArray, NDArray]:
        """Return shape function gradients for local coordinates."""

        dN_dxi: NDArray = np.array(
            [
                -0.25 * (1 - eta),
                0.25 * (1 - eta),
                0.25 * (1 + eta),
                -0.25 * (1 + eta),
            ],
            dtype=self.dtype,
        )
        dN_deta: NDArray = np.array(
            [
                -0.25 * (1 - xi),
                -0.25 * (1 + xi),
                0.25 * (1 + xi),
                0.25 * (1 - xi),
            ],
            dtype=self.dtype,
        )
        return dN_dxi, dN_deta

gauss_points() staticmethod

Return the points for 2x2 Gauss quadrature.

Source code in tofea/elements.py 
31
32
33
34
35
36
@staticmethod
def gauss_points() -> list[tuple[float, float]]:
    """Return the points for 2x2 Gauss quadrature."""

    gp = 1 / np.sqrt(3)
    return [(-gp, -gp), (gp, -gp), (gp, gp), (-gp, gp)]

grad_shape_funcs(xi, eta)

Return shape function gradients for local coordinates.

Source code in tofea/elements.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def grad_shape_funcs(self, xi: float, eta: float) -> tuple[NDArray, NDArray]:
    """Return shape function gradients for local coordinates."""

    dN_dxi: NDArray = np.array(
        [
            -0.25 * (1 - eta),
            0.25 * (1 - eta),
            0.25 * (1 + eta),
            -0.25 * (1 + eta),
        ],
        dtype=self.dtype,
    )
    dN_deta: NDArray = np.array(
        [
            -0.25 * (1 - xi),
            -0.25 * (1 + xi),
            0.25 * (1 + xi),
            0.25 * (1 - xi),
        ],
        dtype=self.dtype,
    )
    return dN_dxi, dN_deta

Q4Element_K dataclass

Bases: Q4Element

Plane stress elasticity element.

Parameters:

Name Type Description Default
e float

Young's modulus of the material.

 1.0
nu float

Poisson's ratio of the material.

 1 / 3
Source code in tofea/elements.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
@dataclass(frozen=True, slots=True)
class Q4Element_K(Q4Element):
    """Plane stress elasticity element.

    Parameters
    ----------
    e : float
        Young's modulus of the material.
    nu : float
        Poisson's ratio of the material.
    """

    e: float = 1.0
    nu: float = 1 / 3

    @cached_property
    def element(self) -> NDArray:
        """Return the 8x8 stiffness matrix for the element.

        The matrix is assembled using 2x2 Gauss quadrature instead of the
        symbolic integration previously used.  This avoids the costly SymPy
        computations and speeds up object construction significantly while
        maintaining numerical precision.

        Returns
        -------
        numpy.ndarray
            Stiffness matrix of shape ``(8, 8)``.

        Examples
        --------
        >>> from tofea.elements import Q4Element_K
        >>> Q4Element_K().element.shape
        (8, 8)
        """

        C: NDArray = (self.e / (1 - self.nu**2)) * np.array(
            [[1, self.nu, 0], [self.nu, 1, 0], [0, 0, (1 - self.nu) / 2]],
            dtype=self.dtype,
        )

        K: NDArray = np.zeros((8, 8), dtype=self.dtype)

        for xi, eta in self.gauss_points():
            dN_dxi, dN_deta = self.grad_shape_funcs(xi, eta)

            B: NDArray = np.zeros((3, 8), dtype=self.dtype)
            for i in range(4):
                B[0, 2 * i] = dN_dxi[i] / self.dx
                B[1, 2 * i + 1] = dN_deta[i] / self.dy
                B[2, 2 * i] = dN_deta[i] / self.dy
                B[2, 2 * i + 1] = dN_dxi[i] / self.dx

            K += (B.T @ C @ B) * self.dx * self.dy

        K[np.abs(K) < self.eps] = 0
        return K

element cached property

Return the 8x8 stiffness matrix for the element.

The matrix is assembled using 2x2 Gauss quadrature instead of the symbolic integration previously used. This avoids the costly SymPy computations and speeds up object construction significantly while maintaining numerical precision.

Returns:

Type Description
ndarray

Stiffness matrix of shape (8, 8).

Examples:

>>> from tofea.elements import Q4Element_K
>>> Q4Element_K().element.shape
(8, 8)

Q4Element_T dataclass

Bases: Q4Element

Heat conductivity element.

Parameters:

Name Type Description Default
k float

Thermal conductivity of the material.

 1.0
Source code in tofea/elements.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@dataclass(frozen=True, slots=True)
class Q4Element_T(Q4Element):
    """Heat conductivity element.

    Parameters
    ----------
    k : float
        Thermal conductivity of the material.
    """

    k: float = 1.0

    @cached_property
    def element(self) -> NDArray:
        """Return the 4x4 conductivity matrix for the element.

        Similar to :class:`Q4Element_K`, the matrix is assembled numerically
        using 2x2 Gauss quadrature.  This removes the dependency on SymPy for
        runtime calculations and considerably reduces initialization time.

        Returns
        -------
        numpy.ndarray
            Conductivity matrix of shape ``(4, 4)``.

        Examples
        --------
        >>> from tofea.elements import Q4Element_T
        >>> Q4Element_T().element.shape
        (4, 4)
        """

        C: NDArray = np.array([[self.k, 0], [0, self.k]], dtype=self.dtype)

        K: NDArray = np.zeros((4, 4), dtype=self.dtype)

        for xi, eta in self.gauss_points():
            dN_dxi, dN_deta = self.grad_shape_funcs(xi, eta)

            B: NDArray = np.zeros((2, 4), dtype=self.dtype)
            for i in range(4):
                B[0, i] = dN_dxi[i] / self.dx
                B[1, i] = dN_deta[i] / self.dy

            K += (B.T @ C @ B) * self.dx * self.dy

        K[np.abs(K) < self.eps] = 0
        return K

element cached property

Return the 4x4 conductivity matrix for the element.

Similar to :class:Q4Element_K, the matrix is assembled numerically using 2x2 Gauss quadrature. This removes the dependency on SymPy for runtime calculations and considerably reduces initialization time.

Returns:

Type Description
ndarray

Conductivity matrix of shape (4, 4).

Examples:

>>> from tofea.elements import Q4Element_T
>>> Q4Element_T().element.shape
(4, 4)

Finite element solvers for 2D problems.

FEA2D dataclass

Bases: ABC

Abstract base class for 2D finite element models.

Parameters:

Name Type Description Default
fixed NDArray[bool_]

Boolean mask indicating which degrees of freedom are fixed. Shape for thermal problems is (nx + 1, ny + 1) and for elasticity problems (nx + 1, ny + 1, 2).

 required
solver str

Name of the backend solver to use.

 DEFAULT_SOLVER
dx float

Element dimensions in x and y direction.

 0.5
dy float

Element dimensions in x and y direction.

 0.5
Notes

Subclasses must provide the element matrix and a mapping from element degrees of freedom to system degrees of freedom.

Source code in tofea/fea2d.py 
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@dataclass
class FEA2D(ABC):
    """Abstract base class for 2D finite element models.

    Parameters
    ----------
    fixed
        Boolean mask indicating which degrees of freedom are fixed.
        Shape for thermal problems is ``(nx + 1, ny + 1)`` and for
        elasticity problems ``(nx + 1, ny + 1, 2)``.
    solver
        Name of the backend solver to use.
    dx, dy
        Element dimensions in ``x`` and ``y`` direction.

    Notes
    -----
    Subclasses must provide the element matrix and a mapping from element
    degrees of freedom to system degrees of freedom.
    """

    fixed: NDArray[np.bool_]
    solver: str = DEFAULT_SOLVER
    dx: float = 0.5
    dy: float = 0.5

    @property
    @abstractmethod
    def dof_dim(self) -> int: ...

    @property
    @abstractmethod
    def element(self) -> NDArray: ...

    @property
    @abstractmethod
    def dofmap(self) -> NDArray[np.uint32]: ...

    @property
    def shape(self) -> tuple[int, int]:
        """Number of elements in :math:`x` and :math:`y` direction."""

        nx, ny = self.fixed.shape[:2]
        return nx - 1, ny - 1

    @cached_property
    def dofs(self) -> NDArray[np.uint32]:
        """Indices of all degrees of freedom."""

        return np.arange(self.fixed.size, dtype=np.uint32)

    @cached_property
    def fixdofs(self) -> NDArray[np.uint32]:
        """Indices of fixed degrees of freedom."""

        return self.dofs[self.fixed.ravel()]

    @cached_property
    def freedofs(self) -> NDArray[np.uint32]:
        """Indices of free degrees of freedom."""

        return self.dofs[~self.fixed.ravel()]

    @cached_property
    def _solver(self) -> Solver:
        """Backend solver instance."""

        return get_solver(self.solver)

    @cached_property
    def index_map(self) -> NDArray[np.uint32]:
        """Permutation that places free DOFs first."""

        indices = np.concatenate([self.freedofs, self.fixdofs])
        imap = np.zeros_like(self.dofs)
        imap[indices] = self.dofs
        return imap

    @cached_property
    def e2sdofmap(self) -> NDArray[np.uint32]:
        """Map each element to its system DOF indices."""

        nx, ny = self.shape
        x, y = np.unravel_index(np.arange(nx * ny), (nx, ny))
        idxs = self.dof_dim * (y + x * (ny + 1))
        return np.add(self.dofmap[None], idxs[:, None].astype(np.uint32))

    @cached_property
    def e2sdofmap_reshaped(self) -> NDArray[np.uint32]:
        """Element mapping reshaped for tensor operations."""

        return np.reshape(self.e2sdofmap.T, (-1, *self.shape))

    @cached_property
    def keep_indices(
        self,
    ) -> tuple[NDArray[np.bool_], NDArray[np.uint32]]:
        """Indices for assembling the global matrix."""

        i, j = np.meshgrid(range(len(self.dofmap)), range(len(self.dofmap)))
        ix = self.e2sdofmap[:, i].ravel()
        iy = self.e2sdofmap[:, j].ravel()
        keep = np.isin(ix, self.freedofs) & np.isin(iy, self.freedofs)
        indices = np.stack([self.index_map[ix][keep], self.index_map[iy][keep]])
        return keep, indices

    def global_mat(self, x: NDArray) -> tuple[NDArray, NDArray]:
        """Assemble global matrix from element data."""

        x = np.reshape(x, (-1, 1, 1)) * self.element[None]
        x = x.ravel()
        keep, indices = self.keep_indices
        return x[keep], indices

    def solve(self, x: NDArray, b: NDArray) -> NDArray:
        """Solve ``K(x) u = b`` for ``u``."""

        data, indices = self.global_mat(x)
        u_nz = solve_coo(data, indices, b.ravel()[self.freedofs], self._solver)
        u = anp.concatenate([u_nz, np.zeros(len(self.fixdofs))])[self.index_map]
        return u

    def _self_adjoint_objective(self, x: NDArray, b: NDArray) -> NDArray:
        """Internal method to compute a self-adjoint objective."""

        data, indices = self.global_mat(x)
        free_rhs = b.ravel()[self.freedofs]
        objective, _ = solve_and_compute_self_adjoint_objective(
            data, indices, free_rhs, self._solver
        )
        return objective

dofs cached property

Indices of all degrees of freedom.

e2sdofmap cached property

Map each element to its system DOF indices.

e2sdofmap_reshaped cached property

Element mapping reshaped for tensor operations.

fixdofs cached property

Indices of fixed degrees of freedom.

freedofs cached property

Indices of free degrees of freedom.

index_map cached property

Permutation that places free DOFs first.

keep_indices cached property

Indices for assembling the global matrix.

shape property

Number of elements in :math:x and :math:y direction.

global_mat(x)

Assemble global matrix from element data.

Source code in tofea/fea2d.py 
123
124
125
126
127
128
129
def global_mat(self, x: NDArray) -> tuple[NDArray, NDArray]:
    """Assemble global matrix from element data."""

    x = np.reshape(x, (-1, 1, 1)) * self.element[None]
    x = x.ravel()
    keep, indices = self.keep_indices
    return x[keep], indices

solve(x, b)

Solve K(x) u = b for u.

Source code in tofea/fea2d.py 
131
132
133
134
135
136
137
def solve(self, x: NDArray, b: NDArray) -> NDArray:
    """Solve ``K(x) u = b`` for ``u``."""

    data, indices = self.global_mat(x)
    u_nz = solve_coo(data, indices, b.ravel()[self.freedofs], self._solver)
    u = anp.concatenate([u_nz, np.zeros(len(self.fixdofs))])[self.index_map]
    return u

FEA2D_K dataclass

Bases: FEA2D

Finite element model for compliance problems.

This model solves small deformation elasticity problems using bilinear quadrilateral elements.

Parameters:

Name Type Description Default
e float

Young's modulus of the material.

 1.0
nu float

Poisson's ratio of the material.

 1 / 3

Examples:

>>> import numpy as np
>>> from tofea.fea2d import FEA2D_K
>>> fixed = np.zeros((1, 1, 2), dtype=bool)
>>> fem = FEA2D_K(fixed, e=210e9, nu=0.3)
>>> fem.element.shape
(8, 8)
Source code in tofea/fea2d.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
@dataclass
class FEA2D_K(FEA2D):
    """Finite element model for compliance problems.

    This model solves small deformation elasticity problems using bilinear
    quadrilateral elements.

    Parameters
    ----------
    e : float
        Young's modulus of the material.
    nu : float
        Poisson's ratio of the material.

    Examples
    --------
    >>> import numpy as np
    >>> from tofea.fea2d import FEA2D_K
    >>> fixed = np.zeros((1, 1, 2), dtype=bool)
    >>> fem = FEA2D_K(fixed, e=210e9, nu=0.3)
    >>> fem.element.shape
    (8, 8)
    """

    dof_dim: int = 2
    e: float = 1.0
    nu: float = 1 / 3

    @cached_property
    def element(self) -> NDArray:
        """Element stiffness matrix."""
        return Q4Element_K(e=self.e, nu=self.nu, dx=self.dx, dy=self.dy).element

    @cached_property
    def dofmap(self) -> NDArray[np.uint32]:
        """Mapping of element DOFs to system DOFs."""
        _, nely = self.shape
        b = np.arange(2 * (nely + 1), 2 * (nely + 1) + 2)
        a = b + 2
        return np.r_[2, 3, a, b, 0, 1].astype(np.uint32)

    def displacement(self, x: NDArray, b: NDArray) -> NDArray:
        """Return displacement field for density ``x`` and load ``b``.

        Parameters
        ----------
        x : numpy.ndarray
            Density field of shape ``(nx, ny)``.
        b : numpy.ndarray
            Load array of shape ``(nx + 1, ny + 1, 2)``.

        This is a general-purpose method that returns the full displacement field,
        suitable for constructing arbitrary objective functions. If your
        objective is compliance minimization, using the ``compliance()`` method is
        significantly more performant.
        """

        return self.solve(x, b)

    def compliance(self, x: NDArray, b: NDArray) -> NDArray:
        """Computes the compliance objective using a highly efficient self-adjoint pathway.

        This method is optimized for compliance minimization and avoids a redundant
        adjoint solve during gradient computation. For constructing arbitrary objective
        functions based on the displacement field, use the :meth:`displacement` method
        instead.

        Parameters
        ----------
        x : numpy.ndarray
            Density field of shape ``(nx, ny)``.
        b : numpy.ndarray
            Load array of shape ``(nx + 1, ny + 1, 2)``.

        Returns
        -------
        numpy.ndarray
            The scalar compliance value.
        """

        return self._self_adjoint_objective(x, b)

dofmap cached property

Mapping of element DOFs to system DOFs.

element cached property

Element stiffness matrix.

compliance(x, b)

Computes the compliance objective using a highly efficient self-adjoint pathway.

This method is optimized for compliance minimization and avoids a redundant adjoint solve during gradient computation. For constructing arbitrary objective functions based on the displacement field, use the :meth:displacement method instead.

Parameters:

Name Type Description Default
x ndarray

Density field of shape (nx, ny).

 required
b ndarray

Load array of shape (nx + 1, ny + 1, 2).

 required

Returns:

Type Description
ndarray

The scalar compliance value.
Source code in tofea/fea2d.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def compliance(self, x: NDArray, b: NDArray) -> NDArray:
    """Computes the compliance objective using a highly efficient self-adjoint pathway.

    This method is optimized for compliance minimization and avoids a redundant
    adjoint solve during gradient computation. For constructing arbitrary objective
    functions based on the displacement field, use the :meth:`displacement` method
    instead.

    Parameters
    ----------
    x : numpy.ndarray
        Density field of shape ``(nx, ny)``.
    b : numpy.ndarray
        Load array of shape ``(nx + 1, ny + 1, 2)``.

    Returns
    -------
    numpy.ndarray
        The scalar compliance value.
    """

    return self._self_adjoint_objective(x, b)

displacement(x, b)

Return displacement field for density x and load b.

Parameters:

Name Type Description Default
x ndarray

Density field of shape (nx, ny).

 required
b ndarray

Load array of shape (nx + 1, ny + 1, 2).

 required
This
required
suitable
required
objective
required
significantly
required
Source code in tofea/fea2d.py 
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def displacement(self, x: NDArray, b: NDArray) -> NDArray:
    """Return displacement field for density ``x`` and load ``b``.

    Parameters
    ----------
    x : numpy.ndarray
        Density field of shape ``(nx, ny)``.
    b : numpy.ndarray
        Load array of shape ``(nx + 1, ny + 1, 2)``.

    This is a general-purpose method that returns the full displacement field,
    suitable for constructing arbitrary objective functions. If your
    objective is compliance minimization, using the ``compliance()`` method is
    significantly more performant.
    """

    return self.solve(x, b)

FEA2D_T dataclass

Bases: FEA2D

Finite element model for heat conduction problems.

Parameters:

Name Type Description Default
k float

Thermal conductivity of the material.

 1.0

Examples:

>>> import numpy as np
>>> from tofea.fea2d import FEA2D_T
>>> fixed = np.zeros((1, 1), dtype=bool)
>>> fem = FEA2D_T(fixed, k=200.0)
>>> fem.element.shape
(4, 4)
Source code in tofea/fea2d.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
@dataclass
class FEA2D_T(FEA2D):
    """Finite element model for heat conduction problems.

    Parameters
    ----------
    k : float
        Thermal conductivity of the material.

    Examples
    --------
    >>> import numpy as np
    >>> from tofea.fea2d import FEA2D_T
    >>> fixed = np.zeros((1, 1), dtype=bool)
    >>> fem = FEA2D_T(fixed, k=200.0)
    >>> fem.element.shape
    (4, 4)
    """

    dof_dim: int = 1
    k: float = 1.0

    @cached_property
    def element(self) -> NDArray:
        """Element conductivity matrix."""
        return Q4Element_T(k=self.k, dx=self.dx, dy=self.dy).element

    @cached_property
    def dofmap(self) -> NDArray[np.uint32]:
        """Mapping of element DOFs to system DOFs."""
        _, nely = self.shape
        return np.r_[1, (nely + 2), (nely + 1), 0].astype(np.uint32)

    def temperature(self, x: NDArray, b: NDArray) -> NDArray:
        """Return temperature field for density ``x`` and load ``b``.

        Parameters
        ----------
        x : numpy.ndarray
            Density field of shape ``(nx, ny)``.
        b : numpy.ndarray
            Heat load array of shape ``(nx + 1, ny + 1)``.

        This is a general-purpose method that returns the full temperature field,
        suitable for constructing arbitrary objective functions. If your objective
        is thermal compliance minimization, using the ``thermal_compliance()``
        method is significantly more performant.
        """

        return self.solve(x, b)

    def thermal_compliance(self, x: NDArray, b: NDArray) -> NDArray:
        """Computes the thermal compliance objective using a highly efficient self-adjoint pathway.

        This method is optimized for thermal compliance minimization and avoids a
        redundant adjoint solve during gradient computation. For constructing arbitrary
        objective functions based on the temperature field, use the :meth:`temperature`
        method instead.

        Parameters
        ----------
        x : numpy.ndarray
            Density field of shape ``(nx, ny)``.
        b : numpy.ndarray
            Heat load array of shape ``(nx + 1, ny + 1)``.

        Returns
        -------
        numpy.ndarray
            The scalar thermal compliance value.
        """

        return self._self_adjoint_objective(x, b)

dofmap cached property

Mapping of element DOFs to system DOFs.

element cached property

Element conductivity matrix.

temperature(x, b)

Return temperature field for density x and load b.

Parameters:

Name Type Description Default
x ndarray

Density field of shape (nx, ny).

 required
b ndarray

Heat load array of shape (nx + 1, ny + 1).

 required
This
required
suitable
required
is
required
method
required
Source code in tofea/fea2d.py 
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def temperature(self, x: NDArray, b: NDArray) -> NDArray:
    """Return temperature field for density ``x`` and load ``b``.

    Parameters
    ----------
    x : numpy.ndarray
        Density field of shape ``(nx, ny)``.
    b : numpy.ndarray
        Heat load array of shape ``(nx + 1, ny + 1)``.

    This is a general-purpose method that returns the full temperature field,
    suitable for constructing arbitrary objective functions. If your objective
    is thermal compliance minimization, using the ``thermal_compliance()``
    method is significantly more performant.
    """

    return self.solve(x, b)

thermal_compliance(x, b)

Computes the thermal compliance objective using a highly efficient self-adjoint pathway.

This method is optimized for thermal compliance minimization and avoids a redundant adjoint solve during gradient computation. For constructing arbitrary objective functions based on the temperature field, use the :meth:temperature method instead.

Parameters:

Name Type Description Default
x ndarray

Density field of shape (nx, ny).

 required
b ndarray

Heat load array of shape (nx + 1, ny + 1).

 required

Returns:

Type Description
ndarray

The scalar thermal compliance value.
Source code in tofea/fea2d.py 
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def thermal_compliance(self, x: NDArray, b: NDArray) -> NDArray:
    """Computes the thermal compliance objective using a highly efficient self-adjoint pathway.

    This method is optimized for thermal compliance minimization and avoids a
    redundant adjoint solve during gradient computation. For constructing arbitrary
    objective functions based on the temperature field, use the :meth:`temperature`
    method instead.

    Parameters
    ----------
    x : numpy.ndarray
        Density field of shape ``(nx, ny)``.
    b : numpy.ndarray
        Heat load array of shape ``(nx + 1, ny + 1)``.

    Returns
    -------
    numpy.ndarray
        The scalar thermal compliance value.
    """

    return self._self_adjoint_objective(x, b)

Helper utilities for defining boundary conditions and loads.

BoundaryConditions dataclass

Container for fixed degrees of freedom and loads.

Parameters:

Name Type Description Default
shape tuple[int, int]

Number of elements in x and y direction.

 required
dof_dim int

Degrees of freedom per node.

 1

Examples:

>>> bc = BoundaryConditions((1, 1))
>>> bc.fix_edge("left")
>>> bool(bc.fixed[0, 0])
True
>>> bc.apply_point_load(1, 1, 2.0)
>>> float(bc.load[1, 1])
2.0
Source code in tofea/boundary_conditions.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
@dataclass
class BoundaryConditions:
    """Container for fixed degrees of freedom and loads.

    Parameters
    ----------
    shape:
        Number of elements in ``x`` and ``y`` direction.
    dof_dim:
        Degrees of freedom per node.

    Examples
    --------
    >>> bc = BoundaryConditions((1, 1))
    >>> bc.fix_edge("left")
    >>> bool(bc.fixed[0, 0])
    True
    >>> bc.apply_point_load(1, 1, 2.0)
    >>> float(bc.load[1, 1])
    2.0
    """

    shape: tuple[int, int]
    dof_dim: int = 1
    fixed: NDArray[np.bool_] = field(init=False)
    load: NDArray[np.float64] = field(init=False)

    def __post_init__(self) -> None:
        node_shape = (self.shape[0] + 1, self.shape[1] + 1)
        if self.dof_dim == 1:
            self.fixed = np.zeros(node_shape, dtype="?")
            self.load = np.zeros(node_shape, dtype=float)
        else:
            self.fixed = np.zeros((*node_shape, self.dof_dim), dtype="?")
            self.load = np.zeros((*node_shape, self.dof_dim), dtype=float)

    def fix_edge(self, edge: str) -> None:
        """Fix all nodes along an edge."""
        idx = _edge_indices(edge)
        if self.dof_dim == 1:
            self.fixed[idx] = True
        else:
            self.fixed[(*idx, slice(None))] = True

    def apply_point_load(
        self, node_x: int, node_y: int, load_vector: float | Iterable[float]
    ) -> None:
        """Apply a load vector at a single node."""
        if self.dof_dim == 1:
            self.load[node_x, node_y] = float(cast(float, load_vector))
        else:
            vec = np.asarray(tuple(cast(Iterable[float], load_vector)), dtype=float)
            if vec.size != self.dof_dim:
                raise ValueError("load_vector has wrong size")
            self.load[node_x, node_y] = vec

    def apply_uniform_load_on_edge(self, edge: str, load_vector: float | Iterable[float]) -> None:
        """Apply the same load to all nodes on an edge."""
        idx = _edge_indices(edge)
        if self.dof_dim == 1:
            self.load[idx] = float(cast(float, load_vector))
        else:
            vec = np.asarray(tuple(cast(Iterable[float], load_vector)), dtype=float)
            if vec.size != self.dof_dim:
                raise ValueError("load_vector has wrong size")
            self.load[(*idx, slice(None))] = vec

apply_point_load(node_x, node_y, load_vector)

Apply a load vector at a single node.

Source code in tofea/boundary_conditions.py 
73
74
75
76
77
78
79
80
81
82
83
def apply_point_load(
    self, node_x: int, node_y: int, load_vector: float | Iterable[float]
) -> None:
    """Apply a load vector at a single node."""
    if self.dof_dim == 1:
        self.load[node_x, node_y] = float(cast(float, load_vector))
    else:
        vec = np.asarray(tuple(cast(Iterable[float], load_vector)), dtype=float)
        if vec.size != self.dof_dim:
            raise ValueError("load_vector has wrong size")
        self.load[node_x, node_y] = vec

apply_uniform_load_on_edge(edge, load_vector)

Apply the same load to all nodes on an edge.

Source code in tofea/boundary_conditions.py 
85
86
87
88
89
90
91
92
93
94
def apply_uniform_load_on_edge(self, edge: str, load_vector: float | Iterable[float]) -> None:
    """Apply the same load to all nodes on an edge."""
    idx = _edge_indices(edge)
    if self.dof_dim == 1:
        self.load[idx] = float(cast(float, load_vector))
    else:
        vec = np.asarray(tuple(cast(Iterable[float], load_vector)), dtype=float)
        if vec.size != self.dof_dim:
            raise ValueError("load_vector has wrong size")
        self.load[(*idx, slice(None))] = vec

fix_edge(edge)

Fix all nodes along an edge.

Source code in tofea/boundary_conditions.py 
65
66
67
68
69
70
71
def fix_edge(self, edge: str) -> None:
    """Fix all nodes along an edge."""
    idx = _edge_indices(edge)
    if self.dof_dim == 1:
        self.fixed[idx] = True
    else:
        self.fixed[(*idx, slice(None))] = True

Autograd primitives used by the finite element routines.

solve_and_compute_self_adjoint_objective(entries, indices, rhs, solver)

Solve K u = b and return the self-adjoint objective and solution.

Source code in tofea/primitives.py 
145
146
147
148
149
150
151
152
153
154
155
156
@primitive
def solve_and_compute_self_adjoint_objective(
    entries: NDArray,
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],
    rhs: NDArray,
    solver: Solver,
) -> tuple[NDArray, NDArray]:
    """Solve ``K u = b`` and return the self-adjoint objective and solution."""

    u_nz = solve_coo(entries, indices, rhs, solver)
    objective = anp.dot(u_nz, rhs)
    return objective, u_nz

solve_and_compute_self_adjoint_objective_vjp(ans, entries, indices, rhs, solver)

Reverse-mode derivative for solve_and_compute_self_adjoint_objective.

Source code in tofea/primitives.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def solve_and_compute_self_adjoint_objective_vjp(
    ans: tuple[NDArray, NDArray],
    entries: NDArray,  # noqa: ARG001
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],
    rhs: NDArray,  # noqa: ARG001
    solver: Solver,  # noqa: ARG001
) -> Callable[[NDArray | tuple[NDArray, NDArray]], NDArray]:
    """Reverse-mode derivative for ``solve_and_compute_self_adjoint_objective``."""

    _, u_nz = ans

    def vjp(g: NDArray | tuple[NDArray, NDArray]) -> NDArray:
        if not isinstance(g, np.ndarray):
            g = g[0]
        i, j = indices
        return -g * u_nz[i] * u_nz[j]

    return vjp

solve_coo(entries, indices, rhs, solver)

Solve a sparse linear system in COO format.

Parameters:

Name Type Description Default
entries array - like

Non-zero matrix entries.

 required
indices tuple[array - like, array - like]

Row and column indices for entries.

 required
rhs array - like

Right-hand side vector.

 required
solver Solver

Factorization backend used to solve the system.

 required

Returns:

Type Description
ndarray

Solution vector.

Examples:

>>> import numpy as np
>>> from scipy.sparse import coo_matrix
>>> from tofea.primitives import solve_coo
>>> from tofea.solvers import get_solver
>>> m = coo_matrix([[4, 1], [1, 3]])
>>> b = np.array([1, 2])
>>> solve_coo(m.data, (m.row, m.col), b, get_solver("SuperLU"))
array([0.09090909, 0.63636364])
Source code in tofea/primitives.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@primitive
def solve_coo(
    entries: NDArray,
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],
    rhs: NDArray,
    solver: Solver,
) -> NDArray:
    """Solve a sparse linear system in COO format.

    Parameters
    ----------
    entries : array-like
        Non-zero matrix entries.
    indices : tuple[array-like, array-like]
        Row and column indices for ``entries``.
    rhs : array-like
        Right-hand side vector.
    solver : tofea.solvers.Solver
        Factorization backend used to solve the system.

    Returns
    -------
    numpy.ndarray
        Solution vector.

    Examples
    --------
    >>> import numpy as np
    >>> from scipy.sparse import coo_matrix
    >>> from tofea.primitives import solve_coo
    >>> from tofea.solvers import get_solver
    >>> m = coo_matrix([[4, 1], [1, 3]])
    >>> b = np.array([1, 2])
    >>> solve_coo(m.data, (m.row, m.col), b, get_solver("SuperLU"))
    array([0.09090909, 0.63636364])
    """

    a = coo_matrix((entries, indices)).tocsc()
    ctx = getattr(solver, "_ctx", {})
    pattern = ctx.get("pattern")
    if pattern is not None and _same_pattern(a, pattern):
        solver.refactor_numerical(a)
    else:
        solver.clear()
        solver.factor_full(a)
        ctx["pattern"] = _pattern_signature(a)
    return solver.solve(rhs)

solve_coo_b_jvp(g, x, entries, indices, rhs, solver)

Forward-mode derivative of :func:solve_coo with respect to rhs.

Source code in tofea/primitives.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def solve_coo_b_jvp(
    g: NDArray,
    x: NDArray,  # noqa: ARG001
    entries: NDArray,  # noqa: ARG001
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],  # noqa: ARG001
    rhs: NDArray,  # noqa: ARG001
    solver: Solver,
) -> NDArray:
    """Forward-mode derivative of :func:`solve_coo` with respect to ``rhs``."""

    return solver.solve(g)

solve_coo_b_vjp(ans, entries, indices, rhs, solver)

Reverse-mode derivative of :func:solve_coo for rhs.

Source code in tofea/primitives.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
def solve_coo_b_vjp(
    ans: NDArray,  # noqa: ARG001
    entries: NDArray,  # noqa: ARG001
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],  # noqa: ARG001
    rhs: NDArray,  # noqa: ARG001
    solver: Solver,
) -> Callable[[NDArray], NDArray]:
    """Reverse-mode derivative of :func:`solve_coo` for ``rhs``."""

    def vjp(g: NDArray) -> NDArray:
        return solver.solve(g, transpose=True)

    return vjp

solve_coo_entries_jvp(g, x, entries, indices, rhs, solver)

Forward-mode derivative of :func:solve_coo with respect to entries.

Source code in tofea/primitives.py 
80
81
82
83
84
85
86
87
88
89
90
91
def solve_coo_entries_jvp(
    g: NDArray,
    x: NDArray,
    entries: NDArray,  # noqa: ARG001
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],
    rhs: NDArray,  # noqa: ARG001
    solver: Solver,
) -> NDArray:
    """Forward-mode derivative of :func:`solve_coo` with respect to ``entries``."""

    a = coo_matrix((g, indices)).tocsc()
    return solver.solve(-(a @ x))

solve_coo_entries_vjp(ans, entries, indices, rhs, solver)

Reverse-mode derivative of :func:solve_coo for entries.

Source code in tofea/primitives.py 
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def solve_coo_entries_vjp(
    ans: NDArray,
    entries: NDArray,  # noqa: ARG001
    indices: tuple[NDArray[np.int_], NDArray[np.int_]],
    rhs: NDArray,  # noqa: ARG001
    solver: Solver,
) -> Callable[[NDArray], NDArray]:
    """Reverse-mode derivative of :func:`solve_coo` for ``entries``."""

    def vjp(g: NDArray) -> NDArray:
        x = solver.solve(g, transpose=True)
        i, j = indices
        return -x[i] * ans[j]

    return vjp

Linear system solver abstractions.

Solver

Bases: ABC

Abstract interface for linear solvers.

Implementations typically perform a two phase factorization of a sparse matrix: a symbolic analysis that depends only on the sparsity pattern and a numerical factorization that depends on the actual values. factor_full is responsible for both steps while refactor_numerical reuses the symbolic analysis and only updates the numeric values.

Source code in tofea/solvers.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solver(ABC):
    """Abstract interface for linear solvers.

    Implementations typically perform a two phase factorization of a sparse
    matrix: a symbolic analysis that depends only on the sparsity pattern and a
    numerical factorization that depends on the actual values. ``factor_full``
    is responsible for both steps while ``refactor_numerical`` reuses the
    symbolic analysis and only updates the numeric values.
    """

    @abstractmethod
    def factor_full(self, m: csc_matrix) -> None:
        """Perform a complete (symbolic and numeric) factorization of ``m``."""

    @abstractmethod
    def refactor_numerical(self, m: csc_matrix) -> None:
        """Update the numeric factorization of ``m`` using a cached symbolic analysis."""

    @abstractmethod
    def solve(self, rhs: NDArray, transpose: bool = False) -> NDArray: ...

    @abstractmethod
    def clear(self) -> None: ...

factor_full(m) abstractmethod

Perform a complete (symbolic and numeric) factorization of m.

Source code in tofea/solvers.py 
23
24
25
@abstractmethod
def factor_full(self, m: csc_matrix) -> None:
    """Perform a complete (symbolic and numeric) factorization of ``m``."""

refactor_numerical(m) abstractmethod

Update the numeric factorization of m using a cached symbolic analysis.

Source code in tofea/solvers.py 
27
28
29
@abstractmethod
def refactor_numerical(self, m: csc_matrix) -> None:
    """Update the numeric factorization of ``m`` using a cached symbolic analysis."""

SuperLU

Bases: Solver

scipy.sparse.linalg.splu wrapper.

Source code in tofea/solvers.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class SuperLU(Solver):
    """`scipy.sparse.linalg.splu` wrapper."""

    def __init__(self, **options: float | int | bool | str | Mapping[str, bool]) -> None:
        """Create a new ``SuperLU`` solver instance."""
        # store solver-specific context on the instance to avoid cross-talk
        self._ctx: dict[str, Any] = {"splu": partial(splu, **options)}

    def factor_full(self, m: csc_matrix) -> None:
        """Compute the complete factorization of ``m``."""
        self._ctx["factorization"] = self._ctx["splu"](m)

    def refactor_numerical(self, m: csc_matrix) -> None:
        """Update numeric values by re-factorizing ``m``.

        ``scipy``'s ``splu`` does not expose symbolic/numeric separation, so we
        simply recompute the full factorization.
        """
        self.factor_full(m)

    def solve(self, rhs: NDArray, transpose: bool = False) -> NDArray:
        """Solve the linear system."""
        return self._ctx["factorization"].solve(rhs, trans="T" if transpose else "N")

    def clear(self) -> None:
        """Remove cached factorization and sparsity pattern."""
        self._ctx.pop("factorization", None)
        self._ctx.pop("pattern", None)

__init__(**options)

Create a new SuperLU solver instance.

Source code in tofea/solvers.py 
41
42
43
44
def __init__(self, **options: float | int | bool | str | Mapping[str, bool]) -> None:
    """Create a new ``SuperLU`` solver instance."""
    # store solver-specific context on the instance to avoid cross-talk
    self._ctx: dict[str, Any] = {"splu": partial(splu, **options)}

clear()

Remove cached factorization and sparsity pattern.

Source code in tofea/solvers.py 
62
63
64
65
def clear(self) -> None:
    """Remove cached factorization and sparsity pattern."""
    self._ctx.pop("factorization", None)
    self._ctx.pop("pattern", None)

factor_full(m)

Compute the complete factorization of m.

Source code in tofea/solvers.py 
46
47
48
def factor_full(self, m: csc_matrix) -> None:
    """Compute the complete factorization of ``m``."""
    self._ctx["factorization"] = self._ctx["splu"](m)

refactor_numerical(m)

Update numeric values by re-factorizing m.

scipy's splu does not expose symbolic/numeric separation, so we simply recompute the full factorization.

Source code in tofea/solvers.py 
50
51
52
53
54
55
56
def refactor_numerical(self, m: csc_matrix) -> None:
    """Update numeric values by re-factorizing ``m``.

    ``scipy``'s ``splu`` does not expose symbolic/numeric separation, so we
    simply recompute the full factorization.
    """
    self.factor_full(m)

solve(rhs, transpose=False)

Solve the linear system.

Source code in tofea/solvers.py 
58
59
60
def solve(self, rhs: NDArray, transpose: bool = False) -> NDArray:
    """Solve the linear system."""
    return self._ctx["factorization"].solve(rhs, trans="T" if transpose else "N")

get_solver(solver)

Return a solver instance by name.

Parameters:

Name Type Description Default
solver str

Name of the solver implementation. Currently only "SuperLU" is available.

 required

Examples:

>>> from tofea.solvers import get_solver, Solver
>>> s = get_solver("SuperLU")
>>> isinstance(s, Solver)
True
Source code in tofea/solvers.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def get_solver(solver: str) -> Solver:
    """Return a solver instance by name.

    Parameters
    ----------
    solver
        Name of the solver implementation. Currently only ``"SuperLU"`` is
        available.

    Examples
    --------
    >>> from tofea.solvers import get_solver, Solver
    >>> s = get_solver("SuperLU")
    >>> isinstance(s, Solver)
    True
    """
    match solver:
        case "SuperLU":
            return SuperLU(
                diag_pivot_thresh=0.1,
                permc_spec="MMD_AT_PLUS_A",
                options={"SymmetricMode": True},
            )
        case _:
            raise ValueError(f"Invalid solver: {solver}")