Coverage for pyxfem/core/smallsize_tensor.py: 65%
540 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-24 07:17 +0000
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# pylint: disable=invalid-name
4"""
5This module i
6s dedicated to implement faster functionnality for tensor on spaces of dim 2 and 3
7than the equivalent numpy call
8Created on May1 2024
9@author: nchevaug
10"""
12from types import NoneType
13from typing import Tuple, Callable, Literal, Optional
14from warnings import warn
15import numpy as np
16import scipy as sp
17from numba import float32, float64, guvectorize
18import numba as nb
19from pyxfem.core.config import FASTMATH, npfloat, FloatArray
22HAS_MYGRAD = False
24# o2 idetity Id2.u = u
25Id2_2D = np.eye(2)
26Id2_3D = np.eye(3)
27Id2 = {2:Id2_2D, 3:Id2_3D}
28# o4 identity Id4:A = A
29Id4_3D = np.eye(9).reshape((3,3,3,3))
30Id4_2D = np.eye(4).reshape((2,2,2,2))
31Id4 = {2:Id4_2D, 3:Id4_3D}
32# o4 Transposition Tr4:A = A^T
33Tr4_2D = Id4_2D.swapaxes(2,3) #O4 transposr
34Tr4_3D = Id4_3D.swapaxes(2,3) #O4 transposr
35Tr4 = {2:Tr4_2D, 3:Tr4_3D}
36# o4 Symetrizer Sy4:A = (A +A^T)/2.
37Sy4_2D = 0.5*(Tr4_2D+Id4_2D) # O4 symetrizer
38Sy4_3D = 0.5*(Tr4_3D+Id4_3D) # O4 symetrizer
39Sy4 = {2:Sy4_2D, 3:Sy4_3D}
40# o4 spherificator Sp4:A = trace(A) I
41Sp4_2D = np.tensordot(Id2_2D , Id2_2D, axes = 0 )
42Sp4_3D = np.tensordot(Id2_3D , Id2_3D, axes = 0 )
43Sp4 = {2:Sp4_2D, 3:Sp4_3D}
46@nb.njit(parallel = True, cache = True)
47def axpy_pnjit(alpha, x, y):
48 """Compute vector scale-and-add operation in parallel using Numba.
50 Evaluates the in-place operation y[j] += alpha * x[j] over the first axis.
52 Parameters
53 ----------
54 alpha : float or scalar
55 Scalar multiplier.
56 x : ndarray
57 Input array.
58 y : ndarray
59 Output array modified in place.
60 """
61 for j in nb.prange(x.shape[0]): # pylint: disable=not-an-iterable
62 y[j] += alpha*x[j]
64def axpy(alpha, x, y, parallel = False):
65 """Compute the scaled vector addition y = y + alpha * x.
67 Parameters
68 ----------
69 alpha : float or scalar
70 Scalar scaling factor.
71 x : ndarray
72 Input array.
73 y : ndarray
74 Target array modified in place.
75 parallel : bool, optional
76 If True, executes the multithreaded Numba implementation `axpy_pnjit`.
77 Defaults to False.
78 """
79 if parallel :
80 axpy_pnjit(alpha, x, y)
81 else:
82 y += alpha*x
84def random_orthtransform(n:int, dim:Literal[2,3],
85 random_state = np.random.default_rng(seed=None)) -> FloatArray:
86 """Generate a batch of random orthogonal transformation matrices with det(Q) >= 0.
88 Parameters
89 ----------
90 n : int
91 Number of orthogonal transformations to generate.
92 dim : {2, 3}
93 Spatial dimension (2 for planar rotations, 3 for 3D rotations via orthonormal bases).
94 random_state : np.random.Generator, optional
95 Random number generator instance. Defaults to `np.random.default_rng()`.
97 Returns
98 -------
99 FloatArray
100 Array of shape `(n, dim, dim)` containing proper orthogonal rotation matrices.
101 """
102 if dim==2 :
103 alpha = random_state.random((n))*2*np.pi
104 sin = np.sin(alpha)
105 cos = np.cos(alpha)
106 Q = np.zeros((n, 2, 2))
107 Q[:,0,0] = cos
108 Q[:,0,1] = -sin
109 Q[:,1,0] = sin
110 Q[:,1,1] = cos
111 return Q
112 if dim == 3 :
113 Q = random_state.random((n, dim, dim))
114 Q[...,:,0] = Q[...,:,0]/np.linalg.norm(Q[...,:,0], axis = -1, keepdims=True)
115 Q[...,:,1] = np.cross(Q[...,:,2], Q[...,:,0])
116 Q[...,:,1] = Q[...,:,1]/np.linalg.norm(Q[...,:,1], axis = -1, keepdims=True)
117 Q[...,:,2] = np.cross(Q[...,:,0], Q[...,:,1])
118 return Q
120def vecprodmat(n:FloatArray) -> FloatArray:
121 """Construct skew-symmetric cross-product matrices from 3D vectors.
123 For each 3D vector n, constructs the tensor Q such that Q @ x = n x x.
125 Parameters
126 ----------
127 n : FloatArray
128 Array of 3D vectors with trailing dimension of length 3 (shape `(..., 3)`).
130 Returns
131 -------
132 FloatArray
133 Array of skew-symmetric 3x3 matrices with shape `(..., 3, 3)`.
134 """
135 assert n.shape[-1] == 3
136 Q = np.zeros(n.shape[:-1] + (3,3))
137 Q[...,0,1] = -n[...,2]
138 Q[...,0,2] = n[...,1]
139 Q[...,1,0] = n[...,2]
140 Q[...,1,2] = -n[...,0]
141 Q[...,2,0] = -n[...,1]
142 Q[...,2,1] = n[...,0]
143 return Q
145@guvectorize([
146 ('float64[:], float64[:]'),
147 ('float32[:], float32[:]')
148], '(d)->()', nopython=True, cache=True)
149def vec_norm2(v, out):
150 """Compute the Euclidean L2 norm of 1D vector slices via Numba Generalized Ufunc.
152 Parameters
153 ----------
154 v : ndarray
155 1D vector or array slice of dimension `(d,)`.
156 out : ndarray
157 Output scalar buffer receiving the computed norm.
158 """
159 s = 0.0
160 for i in range(v.shape[0]):
161 s += v[i]**2
162 out[0] = np.sqrt(s)
164def random_lintransform(n:int, dim:int,
165 random_state = np.random.default_rng(seed=None) ) -> FloatArray:
166 """Generate a batch of random linear transformation tensors with det(F) >= 0.
168 Filters near-singular transformations (|det| < 1e-2) and reflects tensors
169 with negative determinants to enforce orientation preservation.
171 Parameters
172 ----------
173 n : int
174 Number of transformation matrices.
175 dim : int
176 Matrix dimension (e.g., 2 or 3).
177 random_state : np.random.Generator, optional
178 Random number generator instance. Defaults to `np.random.default_rng()`.
180 Returns
181 -------
182 FloatArray
183 Array of transformation matrices of shape `(n, dim, dim)`.
184 """
185 F = np.zeros((n, dim, dim))
186 J = np.zeros(n)
187 while True:
188 indexes = np.where(np.abs(J)< 1.e-2)[0]
189 k = len(indexes)
190 if k == 0 :
191 break
192 F[indexes] = random_state.random((len(indexes), dim, dim))
193 J[indexes] = det(F[indexes])
195 Jneg = np.sign(J) <= 0.
196 F[Jneg,0,:] = -F[Jneg,0,:]
197 return F
199def random_lintransform2dto3d(n:int, random_state = np.random.default_rng(seed=None)) -> FloatArray:
200 """Generate random transformation matrices F mapping 2D vectors to 3D space.
202 Parameters
203 ----------
204 n : int
205 Number of transformations.
206 random_state : np.random.Generator, optional
207 Random number generator instance. Defaults to `np.random.default_rng()`.
209 Returns
210 -------
211 FloatArray
212 Array of matrices of shape `(n, 3, 2)`.
213 """
214 F =random_state.random((n, 3, 2))
215 return F
217def metric(F:FloatArray) ->FloatArray:
218 """Compute the right Cauchy-Green deformation tensor / Riemannian metric C = F^T @ F.
220 Parameters
221 ----------
222 F : FloatArray
223 Deformation gradient or transformation tensor of shape `(..., M, N)`.
225 Returns
226 -------
227 FloatArray
228 Metric tensor array of shape `(..., N, N)`.
229 """
230 return trans(F)@F
232def dmetric(F:FloatArray) -> FloatArray:
233 """Compute the first derivative of the metric tensor m = F^T @ F with respect to F.
235 Returns the fourth-order tensor R[..., a, b, i, c] = dm_{ab} / dF_{ic}.
237 Parameters
238 ----------
239 F : FloatArray
240 Transformation tensor of shape `(..., dim_out, dim_in)`.
242 Returns
243 -------
244 FloatArray
245 Fourth-order gradient tensor collection of shape `(..., dim_in, dim_in, dim_out, dim_in)`.
246 """
247 dim = F.shape[-1]
248 I = np.eye(dim)
249 return np.einsum("...ac,...ib->...abic", I, F) + np.einsum("...bc,...ia->...abic", I, F)
251def d2metric(F:FloatArray) -> FloatArray:
252 """Compute the second derivative of the metric tensor m with respect to F.
254 Returns the constant sixth-order tensor R[..., a, b, i, c, j, d] = d^2 m_{ab} / (dF_{ic} dF_{jd}).
256 Parameters
257 ----------
258 F : FloatArray
259 Transformation tensor of shape `(..., dim_out, dim_in)`.
261 Returns
262 -------
263 FloatArray
264 Sixth-order tensor of shape `(..., dim_in, dim_in, dim_out, dim_in, dim_out, dim_in)`.
265 """
266 I0 = np.eye(F.shape[-2])
267 I1 = np.eye(F.shape[-1])
268 d2m = np.einsum("...ac, ...ij, ...bd ->...abicjd", I1, I0, I1) \
269 + np.einsum("...bc, ...ij, ...ad ->...abicjd", I1, I0, I1)
270 d2m = d2m.reshape((1,)*(F.ndim-2) + d2m.shape, copy=False)
271 return d2m
273def invmetric(F:FloatArray)->FloatArray:
274 """Compute the inverse of the metric tensor m = F^T @ F.
276 Parameters
277 ----------
278 F : FloatArray
279 Transformation tensor of shape `(..., M, N)`.
281 Returns
282 -------
283 FloatArray
284 Inverse metric tensor array of shape `(..., N, N)`.
285 """
286 return inv(metric(F))
288def dinvmetric(F):
289 """Compute the derivative of the inverse metric tensor with respect to F.
291 Evaluates d(m^{-1})_{ab} / dF_{ic}.
293 Parameters
294 ----------
295 F : FloatArray
296 Transformation tensor of shape `(..., M, N)`.
298 Returns
299 -------
300 FloatArray
301 Fourth-order derivative tensor of shape `(..., N, N, M, N)`.
302 """
303 im = invmetric(F)
304 return -np.einsum("...ac,...ib->...abic", im, F@im)-np.einsum("...cb,...ia->...abic", im, F@im)
306def trace(A:FloatArray) -> FloatArray:
307 """Compute the trace of a batch of square second-order tensors.
309 Parameters
310 ----------
311 A : FloatArray
312 Square tensor collection of shape `(..., N, N)`.
314 Returns
315 -------
316 FloatArray
317 Trace values array of shape `(...)`.
319 Raises
320 ------
321 ValueError
322 If the trailing dimensions are not square.
323 """
324 if A.shape[-1] != A.shape[-2]:
325 raise ValueError("in trace, each array in the collection must be square")
326 dim = A.shape[-1]
327 if dim == 1:
328 tr = A[...,0,0]
329 elif dim == 2:
330 tr = A[...,0,0] + A[...,1,1]
331 elif dim == 3:
332 tr = A[...,0,0] + A[...,1,1] + A[...,2,2]
333 else:
334 tr = np.einsum('...ii', A)
335 return tr
337def trace2(A:FloatArray) -> FloatArray:
338 """Compute the Frobenius squared norm / trace(A^T @ A) for a batch of tensors.
340 Parameters
341 ----------
342 A : FloatArray
343 Tensor collection of shape `(..., N, N)`.
345 Returns
346 -------
347 FloatArray
348 Array containing the sum of squared components for each tensor.
350 Raises
351 ------
352 ValueError
353 If the trailing dimensions are not square.
354 """
355 if A.shape[-1] != A.shape[-2]:
356 raise ValueError("in trace, each array in the collection must be square")
357 return np.sum(np.square(A), axis = (-1,-2))
359def trans(A:FloatArray)-> FloatArray:
360 """ return the transposes of the input collection of second order tensors """
361 """Transpose the last two dimensions of an array of matrices.
363 Parameters
364 ----------
365 A : FloatArray
366 Tensor array of shape `(..., M, N)`.
368 Returns
369 -------
370 FloatArray
371 Transposed tensor array of shape `(..., N, M)`.
372 """
373 return A.swapaxes(-1,-2)
375def contract(A:FloatArray, B:FloatArray)-> FloatArray:
376 """Compute the double contraction (Frobenius inner product) A : B = A_{ij} B_{ij}.
378 Parameters
379 ----------
380 A : FloatArray
381 First batch of tensors of shape `(..., M, N)`.
382 B : FloatArray
383 Second batch of tensors of matching shape `(..., M, N)`.
385 Returns
386 -------
387 FloatArray
388 Contracted scalar field of shape `(...)`.
389 """
390 assert A.shape == B.shape
391 assert A.ndim >= 2
392 if A.shape[:-2] == (2,2):
393 return A[...,0,0]*B[...,0,0]+A[...,0,1]*B[...,0,1] \
394 + A[...,1,0]*B[...,1,0]+ A[...,1,1]*B[...,1,1]
395 return np.einsum("...ij, ...ij-> ...", A, B)
397def _cofac(A:FloatArray) -> FloatArray:
398 """compute cofac matrix for stack of 2x2 or 3x3 matrices """
399 """Compute the cofactor matrix for batches of 2x2 or 3x3 matrices.
401 Parameters
402 ----------
403 A : FloatArray
404 Tensor array of shape `(..., 2, 2)` or `(..., 3, 3)`.
406 Returns
407 -------
408 FloatArray
409 Cofactor matrix array with identical shape to `A`.
411 Raises
412 ------
413 ValueError
414 If the trailing dimensions are neither (2, 2) nor (3, 3).
415 """
416 if A.shape[-2:] == (2,2):
417 cA = np.empty_like(A)
418 cA[..., 0, 0] = A[..., 1, 1]
419 cA[..., 0, 1] = -A[..., 1, 0]
420 cA[..., 1, 0] = -A[..., 0, 1]
421 cA[..., 1, 1] = A[..., 0, 0]
422 return cA
423 if A.shape[-2:] == (3,3):
424 cA = np.empty_like(A)
425 cA[...,0,:] = cross(A[...,1,:], A[...,2,:])
426 cA[...,1,:] = cross(A[...,2,:], A[...,0,:])
427 cA[...,2,:] = cross(A[...,0,:], A[...,1,:])
428 return cA
429 raise ValueError("A must be a (..,2,2) or (...,3,3) array_like")
431@guvectorize([
432 (float64[:,:], float64[:]),
433 (float32[:,:], float32[:])
434], '(m,m)->()', nopython=True, cache=True)
435def det22_numba(A, out):
436 """Compute the determinant of a single 2x2 matrix via Numba Generalized Ufunc.
438 Parameters
439 ----------
440 A : ndarray
441 2x2 input matrix slice of shape `(2, 2)`.
442 out : ndarray
443 Output buffer receiving the determinant scalar.
444 """
445 out[0] = A[0,0]*A[1,1]- A[0,1]*A[1,0]
447def _det(A:FloatArray, out:NoneType|FloatArray=None) -> FloatArray|npfloat:
448 """Compute the determinant for batches of 2x2, 3x3, or 3x2 (metric surface) matrices.
450 Parameters
451 ----------
452 A : FloatArray
453 Matrix or batch of matrices of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
454 out : FloatArray or None, optional
455 Pre-allocated output destination buffer. Defaults to None.
457 Returns
458 -------
459 FloatArray or npfloat
460 Determinants array of shape `(...)`. For 3x2, returns sqrt(det(F^T F)).
462 Raises
463 ------
464 ValueError
465 If matrix dimensions are not supported.
466 """
467 if out is None :
468 out = np.empty_like(A[...,0,0])
469 if A.shape[-2:] == (3,2):
470 m = metric(A)
471 _det(m, out)
472 return np.sqrt(out, out)
473 if A.shape[-2:] == (2,2):
474 det22_numba(A, out)
475 return out
476 #return A[...,0,0]*A[...,1,1]- A[...,0,1]*A[...,1,0]
477 if A.shape[-2:] == (3,3):
478 a= A[...,0,0]
479 b= A[...,0,1]
480 c= A[...,0,2]
481 d= A[...,1,0]
482 e= A[...,1,1]
483 f= A[...,1,2]
484 g= A[...,2,0]
485 h= A[...,2,1]
486 i= A[...,2,2]
487 if np.isscalar(out):
488 out = a*e*i +b*f*g +c*d*h - c*e*g -b*d*i -a*f*h
489 else:
490 out[...] = a*e*i +b*f*g +c*d*h - c*e*g -b*d*i -a*f*h
491 return out
492 raise ValueError(f"Size error {str(A.shape)}")
494def det_invt(A:FloatArray) -> Tuple[FloatArray|npfloat, FloatArray]:
495 """Compute both the determinant and the inverse-transpose A^{-T} simultaneously.
497 Also handles the 3x2 membrane configuration where A^{-T} is computed as A @ (A^T A)^{-1}.
499 Parameters
500 ----------
501 A : FloatArray
502 Tensor array of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
504 Returns
505 -------
506 tuple of (FloatArray or npfloat, FloatArray)
507 Tuple (detA, invAT) containing determinants and inverse-transposed matrices.
509 Raises
510 ------
511 ValueError
512 If trailing dimensions are not (2, 2), (3, 3), or (3, 2).
513 """
514 if A.shape[-2:] == (3,2):
515 m = metric(A)
516 detm = det(m)
517 invAT = A@inv(m)
518 return np.sqrt(detm), invAT
519 if A.shape[-2:] in [(2,2), (3,3)]:
520 detA = det(A)
521 adjA = cofac(A)
522 if A.shape[:-2]== ():
523 invAT = adjA/detA
524 else :
525 invAT = adjA / np.atleast_1d(detA)[..., None, None]
526 return detA, invAT
527 raise ValueError("A.shape[-2,:] should be (2, 2) or (3, 3) or (3,2)"+str(A.shape) )
529det = _det
530def ddet(A:FloatArray)->FloatArray:
531 """Compute the derivative of the determinant with respect to the matrix A.
533 Returns d(det(A)) / dA = cofac(A) for square matrices, or det(A) * A^{-T}.
535 Parameters
536 ----------
537 A : FloatArray
538 Input tensor collection of shape `(..., M, N)`.
540 Returns
541 -------
542 FloatArray
543 Derivative array having the identical shape of `A`.
544 """
545 if A.shape[-1]==A.shape[-2]:
546 return cofac(A)
547 J, invAT = det_invt(A)
548 return ((np.atleast_1d(J)[:, np.newaxis, np.newaxis])*invAT).reshape(A.shape)
550cofac = _cofac
552def cross(A:FloatArray, B:FloatArray) -> FloatArray:
553 """Compute the 3D vector cross product between arrays of vectors.
555 Parameters
556 ----------
557 A : FloatArray
558 First collection of 3D vectors of shape `(..., 3)`.
559 B : FloatArray
560 Second collection of 3D vectors of shape `(..., 3)`.
562 Returns
563 -------
564 FloatArray
565 Cross product vector array of shape `(..., 3)`.
567 Raises
568 ------
569 ValueError
570 If the trailing dimension of A or B is not 3.
571 """
572 if isinstance(A, np.ndarray) and isinstance(B, np.ndarray):
573 return np.cross(A,B)
574 # if A.shape[-1] != 3 or B.shape[-1] !=3 :
575 # raise ValueError("A and be most be of shape (..., 3) ")
576 # C = np.empty_like(A)
577 # C[...,0] = A[..., 1] * B[...,2] - A[..., 2] * B[...,1]
578 # C[...,1] = -A[..., 0] * B[...,2] + A[..., 2] * B[...,0]
579 # C[...,2] = A[..., 0] * B[...,1] - A[..., 1] * B[...,0]
580 # return C
582def invariants(A:FloatArray) -> FloatArray:
583 """Compute principal invariants for a batch of 2D or 3D second-order tensors.
585 For 2D: [tr(A), det(A)].
586 For 3D: [tr(A), 0.5 * (tr(A)^2 - tr(A^T A)), det(A)].
588 Parameters
589 ----------
590 A : FloatArray
591 Square tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
593 Returns
594 -------
595 FloatArray
596 Invariants array with trailing dimension 2 (for 2D) or 3 (for 3D).
598 Raises
599 ------
600 ValueError
601 If trailing dimensions are not (2, 2) or (3, 3).
602 """
603 if A.shape[-2:] == (2,2):
604 r = np.empty( A.shape[:-2]+(2,))
605 r[...,0] = trace(A)
606 r[...,1] = det(A)
607 return r
608 if A.shape[-2:] == (3,3):
609 r = np.empty( A.shape[:-2]+(3,))
610 trA = trace(A)
611 r[...,0] = trA
612 r[...,1] = 0.5*( np.square(trA) - trace(trans(A)@A) )
613 r[...,2] = det(A)
614 return r
615 raise ValueError("A must be of size (..., 2, 2) or (..., 3, 3)")
617def dinvariants(A:FloatArray) -> FloatArray:
618 """Compute derivatives of principal invariants with respect to tensor A.
620 Parameters
621 ----------
622 A : FloatArray
623 Square tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
625 Returns
626 -------
627 FloatArray
628 Derivatives tensor of shape `(..., 2, 2, 2)` for 2D or `(..., 3, 3, 3)` for 3D.
630 Raises
631 ------
632 ValueError
633 If trailing dimensions are not (2, 2) or (3, 3).
634 """
635 if A.shape[-2:] == (2,2):
636 r = np.empty( A.shape[:-2]+(2,2,2))
637 I = np.eye(2)[(np.newaxis,)*len(A.shape[:-2]) ]
638 r[...,0, :, :] = I
639 r[...,1, : ,:] = cofac(A)
640 return r
641 if A.shape[-2:] == (3,3):
642 nfirstdim = len(A.shape[:-2])
643 r = np.empty( A.shape[:-2]+(3,3,3))
644 I = np.eye(3)[(np.newaxis,)*nfirstdim ]
645 trA = trace(A)[..., None, None]
646 r[...,0,:,:] = I
647 r[...,1,:,:] = trA*I-A
648 r[...,2,:,:] = cofac(A)
649 return r
650 raise ValueError("A must be of size (..., 2, 2) or (..., 3, 3)")
652def invariants2Dto3D(invariant2D):
653 """ from the invariant of a "2.5d tensor", compute the invariants of the equivalent 3D tensor
654 what I mean by a 2.5d tensor is for example the gradiant of the
655 transformation in plan stress case :
656 F = F2D_{a,b} e_a x e_b + l2 e_2 x e_2 , where a, b in [0,1] and e_2 = e_0 vec e_1
657 the invarariant 2D are strored in a (..., 3) ndarray, where
658 invariant2D[..., 0] is an array containing I(F2D) = trace(F2D)
659 invariant2D[..., 1] is an array containing
660 II(F2D) = 0.5(trace(F2D)^2 - trace(F2D.T@F2D)) = det(F2D)
661 invariant2D[..., 2] is an array containing l2
662 Parameters
663 ----------
664 invariant2D : ndarray
665 Array of shape `(..., 3)` containing:
666 - [..., 0]: I_2D = trace(F_2D)
667 - [..., 1]: II_2D = det(F_2D)
668 - [..., 2]: lambda_2 (out-of-plane stretch)
670 Returns
671 -------
672 ndarray
673 Array of shape `(..., 3)` with the three corresponding 3D invariants.
674 """
675 invariant3D = np.empty(invariant2D.shape)
676 I2D = invariant2D[...,0]
677 II2D = invariant2D[...,1]
678 l2 = invariant2D[...,2]
679 invariant3D[..., 0] = I2D + l2
680 invariant3D[..., 1] = II2D + l2*I2D
681 invariant3D[..., 2] = II2D * l2
682 return invariant3D
684def dinvariants2Dto3D(invariant2D:FloatArray) -> FloatArray:
685 """ Compute the matrix of the derivatives of the 3D invariants of tensor A
686 computed as a function of the 2D invariant.
687 input invariant2D an FloatArray of shape[-1] = 3 where :
688 invariant2D[..., 0] = I2D = trace(A2D)
689 invariant2D[..., 1] = II2D = det(A2D) = 0.5*(trace(A2D)**2 - trace(A2D@A2D**2)
690 invariant2D[..., 2] = A33
691 return dinvariant3D an FNDarray of shape (..., 3, 3)
692 where dinvariant3D[..., i, j] contain the derivative of the ith 3d invarariant with regard
693 to the jth 2d invariant
694 Parameters
695 ----------
696 invariant2D : FloatArray
697 Array of shape `(..., 3)` holding [I_2D, II_2D, lambda_2].
699 Returns
700 -------
701 FloatArray
702 Jacobian tensor of shape `(..., 3, 3)` where entry `[..., i, j]` is d(I_3D, i) / d(I_2D, j).
703 """
704 I2D = invariant2D[...,0]
705 II2D = invariant2D[...,1]
706 l2 = invariant2D[...,2]
707 dinvariant3D = np.ones(invariant2D.shape + (3,))
708 dinvariant3D[..., 0, 1] = 0.
709 dinvariant3D[..., 1, 0] = l2
710 dinvariant3D[..., 1, 2] = I2D
711 dinvariant3D[..., 2, 0] = 0.
712 dinvariant3D[..., 2, 1] = l2
713 dinvariant3D[..., 2, 2] = II2D
714 return dinvariant3D
716def d2invariants2Dto3D(invariant2D):
717 """ Compute the matrix of the second derivatives of the 3D invariants of tensor A (Hessian)
718 computed as a function of the 2D invariants.
719 input invariant2D an FloatArray of shape[-1] = 3 where :
720 invariant2D[..., 0] = I2D = trace(A2D)
721 invariant2D[..., 1] = II2D = det(A2D) = 0.5*(trace(A2D)**2 - trace(A2D@A2D**2)
722 invariant2D[..., 2] = A33
723 return dinvariant3D an FNDarray of shape (..., 3, 3, 3)
724 where dinvariant3D[..., i, j, k] contain the derivative of the ith 3d invarariant
725 with regard to the jth 2d invariant and the kth 2d invarariant.
726 Note that the resuluting 27 component matrix contain only 4 non zero terms ...
728 Parameters
729 ----------
730 invariant2D : ndarray
731 Array of shape `(..., 3)` holding [I_2D, II_2D, lambda_2].
733 Returns
734 -------
735 ndarray
736 Second derivative tensor of shape `(..., 3, 3, 3)` containing d^2(I_3D, i) / (d(I_2D, j) d(I_2D, k)).
737 """
738 d2invariant3D = np.zeros(invariant2D.shape + (3, 3))
739 d2invariant3D[..., 1, 0, 2] = 1.
740 d2invariant3D[..., 1, 2, 0] = 1.
741 d2invariant3D[..., 2, 1, 2] = 1.
742 d2invariant3D[..., 2, 2, 1] = 1.
743 return d2invariant3D
745def inv(A:FloatArray, detA: Optional[FloatArray|npfloat] = None,
746 out: Optional[FloatArray] =None ) -> FloatArray:
747 """compute inverse matrix for stack of 2x2 or 3x3 matrices """
748 """Compute matrix inverse for batches of 2x2, 3x3, or pseudo-inverse for 3x2 matrices.
750 Parameters
751 ----------
752 A : FloatArray
753 Tensor collection of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
754 detA : FloatArray or npfloat, optional
755 Precomputed determinant of A to accelerate evaluation. Defaults to None.
756 out : FloatArray, optional
757 Output buffer receiving the inverted matrix. Defaults to None.
759 Returns
760 -------
761 FloatArray
762 Inverted matrix collection.
763 """
764 if A.shape[-2:] == (3,2):
765 out = np.empty(A.shape[:-2] + (2,3)) if out is None else out
766 m = metric(A)
767 np.matmul(inv(m), trans(A), out)
768 return out
769 if detA is None :
770 detA = det(A)
771 if not isinstance(detA, (float, int)):
772 detA = detA.reshape(A.shape[:-2]+(1,1))
773 iA = np.empty_like(A) if out is None else out
774 if A.shape[-2:] == (2,2):
775 iA[..., 0,0] = A[...,1,1]
776 iA[..., 0,1] = -A[...,0,1]
777 iA[..., 1,0] = -A[...,1,0]
778 iA[..., 1,1] = A[...,0,0]
779 elif A.shape[-2:] == (3,3):
780 warn("3d case not fully optimized in inv")
781 iA[...,:,0] = cross(A[...,1,:], A[...,2,:])
782 iA[...,:,1] = cross(A[...,2,:], A[...,0,:])
783 iA[...,:,2] = cross(A[...,0,:], A[...,1,:])
784 if A.shape[:-2] == ():
785 iA = iA.squeeze()
786 iA /= detA
787 return iA
789def dinv(A:FloatArray, iA:FloatArray|None = None )-> FloatArray:
790 """Compute the derivative of the inverse matrix with respect to A.
792 Returns R[..., i, j, k, l] = d(A^{-1}_{ij}) / dA_{kl}.
794 Parameters
795 ----------
796 A : FloatArray
797 Tensor collection of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
798 iA : FloatArray or None, optional
799 Precalculated inverse of A. Defaults to None.
801 Returns
802 -------
803 FloatArray
804 Fourth-order derivative tensor.
805 """
806 if A.shape[-2:] == (3,2):
807 im = inv(metric(A))
808 diA = np.empty(A.shape[:-2] + (2,3,3,2))
809 diA = -np.einsum("...ab,...de,...jd,...ie->...aijb", im, im, A, A )
810 diA -= np.einsum("...ad,...be,...jd,...ie->...aijb", im, im, A, A )
811 diA += np.einsum("...ab, ...ij -> ...aijb", im, np.eye(3))
812 return diA
813 iA = inv(A)
814 return -np.einsum("...ac,...db->...abcd", iA, iA)
816def invt(A:FloatArray, detA:FloatArray|npfloat|None = None) -> FloatArray:
817 """Compute the inverse-transpose (or pseudo-inverse transpose) of a batch of matrices.
819 Parameters
820 ----------
821 A : FloatArray
822 Tensor array of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
823 detA : FloatArray or npfloat, optional
824 Precomputed determinant. Defaults to None.
826 Returns
827 -------
828 FloatArray
829 Inverse transposed matrices of shape `(..., N, M)`.
830 """
831 return np.swapaxes(inv(A, detA), -1,-2)
833def dinv_sym(A, iA = None) -> FloatArray:
834 """Compute the derivative of A^{-1} with respect to A assuming A is symmetric.
836 Parameters
837 ----------
838 A : FloatArray
839 Symmetric tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
840 iA : FloatArray or None, optional
841 Precomputed inverse of A. Defaults to None.
843 Returns
844 -------
845 FloatArray
846 Symmetric fourth-order derivative tensor.
848 Raises
849 ------
850 ValueError
851 If trailing dimensions are not (2, 2) or (3, 3).
852 """
853 if A.shape[-2:] in [(3,3), (2,2)]:
854 if iA is None :
855 iA = inv(A)
856 return -0.5*cross_iklj(iA, iA) -0.5*cross_ilkj(iA, iA)
857 raise ValueError(" A.shape[-2:] must be (2,2) or (3,3)")
859def dinvt(F, iFT =None, im = None, dimdF = None)->FloatArray:
860 """Compute the derivative of the inverse-transpose tensor F^{-T} with respect to F.
862 Parameters
863 ----------
864 F : FloatArray
865 Input deformation tensor of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
866 iFT : FloatArray, optional
867 Precomputed F^{-T}. Defaults to None.
868 im : FloatArray, optional
869 Precomputed inverse metric (F^T F)^{-1} for 3x2 tensors. Defaults to None.
870 dimdF : FloatArray, optional
871 Precomputed derivative of the inverse metric. Defaults to None.
873 Returns
874 -------
875 FloatArray
876 Fourth-order derivative tensor.
877 """
878 if F.shape[-2:] == (3,2):
879 I = np.eye(3)
880 im = invmetric(F) if im is None else im
881 dimdF = dinvmetric(F) if dimdF is None else dimdF
882 return np.einsum("...ij, ...ba -> ...iajb", I, im) \
883 + np.einsum("...il, ...lajb-> ...iajb", F, dimdF)
884 iFT = invt(F) if iFT is None else iFT
885 return -cross_ilkj(iFT, iFT)
887def sqrt22_numpy(A: np.ndarray, out=None) -> np.ndarray:
888 """Compute tensor square roots of 2x2 matrices using vectorized NumPy views.
890 Parameters
891 ----------
892 A : np.ndarray
893 Array of 2x2 matrices of shape `(..., 2, 2)`.
894 out : np.ndarray, optional
895 Output buffer modified in place. Defaults to None.
897 Returns
898 -------
899 np.ndarray
900 Tensor square roots of shape `(..., 2, 2)`.
902 Raises
903 ------
904 ValueError
905 If trailing dimensions are not (2, 2).
906 """
907 if A.shape[-2:] != (2, 2):
908 raise ValueError("A.shape[-2:] should be (2, 2)")
909 if out is None:
910 out = np.empty_like(A)
911 A00, A01 = A[..., 0, 0], A[..., 0, 1]
912 A10, A11 = A[..., 1, 0], A[..., 1, 1]
914 detA = A00 * A11 - A01 * A10
915 traA = A00 + A11
917 s = np.sqrt(detA)
918 t = np.sqrt(traA + 2.0 * s)
920 # Ajout d'axes pour le broadcasting de la division
921 # s et t passent de shape (Batch,) à (Batch, 1, 1)
922 t_expand = t[..., np.newaxis, np.newaxis]
924 # Calcul en place
925 # On commence par copier A dans out pour ne pas l'altérer
926 out[...] = A
927 # On ajoute s sur la diagonale
928 out[..., 0, 0] += s
929 out[..., 1, 1] += s
930 # On divise tout le tenseur par t
931 out[...] /= t_expand
932 return out
934@nb.njit(cache=True, fastmath = FASTMATH)
935def sqrt22_kernel(A22, out):
936 """Numba scalar kernel computing the square root of an individual 2x2 matrix.
938 Parameters
939 ----------
940 A22 : ndarray
941 Single 2x2 matrix slice.
942 out : ndarray
943 Output 2x2 buffer receiving the result.
944 """
945 A00 = A22[0, 0]
946 A01 = A22[0, 1]
947 A10 = A22[1, 0]
948 A11 = A22[1, 1]
949 # 1. Déterminant et s
950 detA = A00 * A11 - A01 * A10
951 s = np.sqrt(detA)
952 # 2. Trace et t
953 traA = A00 + A11
954 t = np.sqrt(traA + 2.0 * s)
955 # 3. Calcul final : B = (A + s * I) / t
956 out[0, 0] = (A00 + s) / t
957 out[0, 1] = A01 / t
958 out[1, 0] = A10 / t
959 out[1, 1] = (A11 + s) / t
961@guvectorize([
962 (float64[:,:], float64[:,:]),
963 (float32[:,:], float32[:,:])
964], '(m,m)->(m,m)', nopython=True, cache=True)
965def _sqrt22_guvec(A, out):
966 """Numba Generalized Ufunc wrapper for batch 2x2 tensor square root.
968 Parameters
969 ----------
970 A : ndarray
971 Input matrix slices of shape `(m, m)`.
972 out : ndarray
973 Output matrix slices of shape `(m, m)`.
974 """
975 sqrt22_kernel(A, out)
977def sqrt22_guvec(A: np.ndarray, out=None) -> np.ndarray:
978 """Compute tensor square roots of 2x2 matrices using Numba guvectorize.
980 Parameters
981 ----------
982 A : np.ndarray
983 Array of matrices of shape `(..., 2, 2)`.
984 out : np.ndarray, optional
985 Preallocated destination array. Defaults to None.
987 Returns
988 -------
989 np.ndarray
990 Tensor square roots of shape `(..., 2, 2)`.
992 Raises
993 ------
994 ValueError
995 If trailing dimensions are not (2, 2).
996 """
997 if A.shape[-2:] != (2, 2):
998 raise ValueError("A.shape[-2:] should be (2, 2)")
999 if out is None:
1000 out = np.empty_like(A)
1001 _sqrt22_guvec(A, out=out) # type: ignore
1002 return out
1005def sqrt22_njit(A: np.ndarray, out=None) -> np.ndarray:
1006 """
1007 Calcule la racine carrée d'un batch de tenseurs 2x2.
1008 A = B.B
1009 """"""Compute tensor square roots of 2x2 matrices using multithreaded Numba NJIT.
1011 Parameters
1012 ----------
1013 A : np.ndarray
1014 C-contiguous array of matrices of shape `(..., 2, 2)`.
1015 out : np.ndarray, optional
1016 Preallocated destination array. Defaults to None.
1018 Returns
1019 -------
1020 np.ndarray
1021 Tensor square roots of shape `(..., 2, 2)`.
1023 Raises
1024 ------
1025 ValueError
1026 If A is not C-contiguous or its shape is not `(..., 2, 2)`.
1027 """
1028 if not A.flags.c_contiguous :
1029 raise ValueError("A not contiguous")
1030 if A.shape[-2:] != (2, 2):
1031 raise ValueError("A.shape[-2:] should be (2, 2)")
1033 if out is None:
1034 out = np.empty_like(A)
1035 Aview = A.reshape((-1, 2, 2))
1036 outview = out.reshape((-1, 2, 2))
1037 _sqrt22_pnjit(Aview, out=outview) # type: ignore
1038 return out
1040@nb.njit(parallel=True, cache = True, fastmath = FASTMATH)
1041def _sqrt22_pnjit(A, out):
1042 """Numba parallel loop over flattened batch of 2x2 matrices.
1044 Parameters
1045 ----------
1046 A : ndarray
1047 Array of shape `(N, 2, 2)`.
1048 out : ndarray
1049 Output array of shape `(N, 2, 2)`.
1051 Returns
1052 -------
1053 ndarray
1054 Output buffer `out`.
1055 """
1056 for i in nb.prange(len(A)): # type ignore # pylint: disable=E1133
1057 sqrt22_kernel(A[i], out[i])
1058 return out
1060def sqrt22(A, out = None):
1061 """Compute square roots for a batch of 2x2 matrices (dispatcher).
1063 Parameters
1064 ----------
1065 A : np.ndarray
1066 Array of matrices of shape `(..., 2, 2)`.
1067 out : np.ndarray, optional
1068 Destination buffer. Defaults to None.
1070 Returns
1071 -------
1072 np.ndarray
1073 Tensor square root array.
1074 """
1075 return sqrt22_njit(A, out)
1076 #return sqrt22_guvec(A, out)
1078@nb.njit(cache =True)
1079def sysmv22_kernel(A: FloatArray, b: FloatArray, out : FloatArray):
1080 """Solve a single symmetric 2x2 linear system A @ x = b via Cramer's rule.
1082 Parameters
1083 ----------
1084 A : FloatArray
1085 Symmetric 2x2 matrix.
1086 b : FloatArray
1087 Right-hand side 2D vector.
1088 out : FloatArray
1089 Output 2D vector receiving solution x.
1090 """
1091 A00, A01, A11 = A[0,0], A[0,1], A[1,1]
1092 b0, b1 = b
1093 detA = A00*A11 - A01*A01
1094 out[0] = (A11*b0 - A01*b1)/detA
1095 out[1] = (A00*b1 - A01*b0)/detA
1097@guvectorize([
1098 (float64[:,:], float64[:], float64[:]),
1099 (float32[:,:], float32[:], float32[:])
1100], '(dim, dim),(dim), (dim)', nopython=True, cache=True)
1101def sysmv22_guvec(A, b, out : FloatArray):
1102 """Generalized Ufunc solving batches of symmetric 2x2 systems A @ x = b.
1104 Parameters
1105 ----------
1106 A : ndarray
1107 Batch of 2x2 matrices of shape `(..., 2, 2)`.
1108 b : ndarray
1109 Batch of vectors of shape `(..., 2)`.
1110 out : ndarray
1111 Buffer receiving solutions of shape `(..., 2)`.
1112 """
1113 sysmv22_kernel(A, b, out)
1115@nb.njit(parallel = True, cache =True)
1116def sysmv22_pnjit(A: FloatArray, b: FloatArray, out : FloatArray):
1117 """Parallel Numba kernel solving batches of symmetric 2x2 systems A @ x = b.
1119 Parameters
1120 ----------
1121 A : FloatArray
1122 Batch of 2x2 matrices of shape `(N, 2, 2)`.
1123 b : FloatArray
1124 Batch of right-hand side vectors of shape `(N, 2)`.
1125 out : FloatArray
1126 Output buffer receiving solutions of shape `(N, 2)`.
1127 """
1128 for j in nb.prange(A.shape[0]): # pylint: disable=not-an-iterable
1129 sysmv22_kernel(A[j], b[j], out[j])
1131def sysmv22(A: FloatArray, b: FloatArray,
1132 out:Optional[FloatArray] = None, parallel = False) -> FloatArray:
1133 r"""
1134 Solves the linear system $A x = b$ for a symmetric 2x2 matrix
1135 using explicit Cramer's rule inversion.
1136 Math: $x = A^{-1} b$
1138 Parameters
1139 ----------
1140 A : FloatArray
1141 Batch of symmetric 2x2 matrices of shape `(..., 2, 2)`.
1142 b : FloatArray
1143 Right-hand side vectors of shape `(..., 2)`.
1144 out : FloatArray, optional
1145 Preallocated output buffer. Defaults to None.
1146 parallel : bool, optional
1147 If True, evaluates using multithreaded OpenMP-style execution. Defaults to False.
1149 Returns
1150 -------
1151 FloatArray
1152 Solution vectors of shape `(..., 2)`.
1154 Raises
1155 ------
1156 ValueError
1157 If trailing dimensions of A are not (2, 2).
1158 """
1159 if A.shape[-2:] != (2,2) :
1160 raise ValueError('sysmv22 is for 2x2 batch of tensors' )
1161 out = np.empty( b.shape, npfloat) if out is None else out
1162 if parallel :
1163 sysmv22_pnjit(A, b, out)
1164 else :
1165 sysmv22_guvec(A, b, out)
1166 return out
1168def gemm22(A: FloatArray, B: FloatArray, out:Optional[FloatArray] = None) -> FloatArray:
1169 r"""
1170 Compute matrix-matrix products C = A @ B for batches of 2x2 matrices.
1172 Optimized for zero heap reallocations using views and in-place accumulations.
1174 Parameters
1175 ----------
1176 A : FloatArray
1177 First batch of 2x2 matrices of shape `(..., 2, 2)`.
1178 B : FloatArray
1179 Second batch of 2x2 matrices of shape `(..., 2, 2)`.
1180 out : FloatArray, optional
1181 Preallocated array receiving the result. Defaults to None.
1183 Returns
1184 -------
1185 FloatArray
1186 Matrix product array of shape `(..., 2, 2)`.
1188 Raises
1189 ------
1190 ValueError
1191 If trailing dimensions are not (2, 2).
1192 """
1193 if A.shape[-2:] != (2,2) :
1194 raise ValueError('gemm22 is for 2x2 batch of tensors' )
1195 C = np.empty_like(A) if out is None else out
1196 # 1. Create Views (0 allocations)
1197 A00, A01 = A[..., 0, 0], A[..., 0, 1]
1198 A10, A11 = A[..., 1, 0], A[..., 1, 1]
1199 B00, B01 = B[..., 0, 0], B[..., 0, 1]
1200 B10, B11 = B[..., 1, 0], B[..., 1, 1]
1201 # 2. Allocate EXACTLY ONE temporary buffer for the batch
1202 # This shape matches the batch dimensions (e.g., if A is 1000x2x2, temp is 1000)
1203 temp = np.empty_like(A00)
1204 # --- Compute C00 ---
1205 # C00 = A00*B00
1206 np.multiply(A00, B00, out=C[..., 0, 0])
1207 # temp = A01*B10
1208 np.multiply(A01, B10, out=temp)
1209 # C00 += temp (In-place addition!)
1210 C[..., 0, 0] += temp
1211 # --- Compute C01 ---
1212 np.multiply(A00, B01, out=C[..., 0, 1])
1213 np.multiply(A01, B11, out=temp)
1214 C[..., 0, 1] += temp
1215 # --- Compute C10 ---
1216 np.multiply(A10, B00, out=C[..., 1, 0])
1217 np.multiply(A11, B10, out=temp)
1218 C[..., 1, 0] += temp
1219 # --- Compute C11 ---
1220 np.multiply(A10, B01, out=C[..., 1, 1])
1221 np.multiply(A11, B11, out=temp)
1222 C[..., 1, 1] += temp
1223 return C
1225@nb.njit(cache=True, fastmath = True, inline='always')
1226def gemm22_kernel_njit(A:FloatArray, B: FloatArray, out: FloatArray):
1227 """Inlined Numba scalar kernel computing the product of two 2x2 matrices.
1229 Parameters
1230 ----------
1231 A : FloatArray
1232 2x2 left-hand matrix slice.
1233 B : FloatArray
1234 2x2 right-hand matrix slice.
1235 out : FloatArray
1236 2x2 destination buffer.
1237 """
1238 A00, A01, A10, A11 = A[0,0], A[0,1], A[1,0], A[1,1]
1239 B00, B01, B10, B11 = B[0,0], B[0,1], B[1,0], B[1,1]
1240 out[0,0] = A00*B00 + A01*B10
1241 out[0,1] = A00*B01 + A01*B11
1242 out[1,0] = A10*B00 + A11*B10
1243 out[1,1] = A10*B01 + A11*B11
1245def AAT22(A, out = None):
1246 """Compute the symmetric product B = A @ A^T for a batch of 2x2 matrices.
1248 Parameters
1249 ----------
1250 A : ndarray
1251 Batch of 2x2 matrices of shape `(..., 2, 2)`.
1252 out : ndarray, optional
1253 Preallocated array receiving the result. Defaults to None.
1255 Returns
1256 -------
1257 ndarray
1258 Symmetric matrix product array of shape `(..., 2, 2)`.
1259 """
1260 if out is None:
1261 B = np.empty_like(A)
1262 else:
1263 B = out
1264 #B = np.empty_like(A)
1265 A00 = A[...,0,0]
1266 A01 = A[...,0,1]
1267 A10 = A[...,1,0]
1268 A11 = A[...,1,1]
1269 B[...,0,0] = A00*A00 + A01*A01
1270 B[...,0,1] = A00*A10 + A01*A11
1271 B[...,1,0] = A10*A00 + A11*A01
1272 B[...,1,1] = A10*A10 + A11*A11
1273 return B
1275@nb.njit(cache=True, fastmath = False)
1276def AAT22_kernel(A22, out):
1277 """Numba kernel computing A @ A^T for an individual 2x2 matrix.
1279 Parameters
1280 ----------
1281 A22 : ndarray
1282 Single 2x2 matrix slice.
1283 out : ndarray
1284 Output 2x2 buffer.
1285 """
1286 A00 = A22[0, 0]
1287 A01 = A22[0, 1]
1288 A10 = A22[1, 0]
1289 A11 = A22[1, 1]
1290 out[0, 0] = A00*A00 + A01*A01
1291 out[0, 1] = A00*A10 + A01*A11
1292 out[1, 0] = A10*A00 + A11*A01
1293 out[1, 1] = A10*A10 + A11*A11
1295@guvectorize([
1296 (float64[:,:], float64[:,:]),
1297 (float32[:,:], float32[:,:])
1298], '(m,m)->(m,m)', nopython=True, cache=True)
1299def _AAT22_guvec(A, out):
1300 """Numba Generalized Ufunc computing A @ A^T across matrix slices.
1302 Parameters
1303 ----------
1304 A : ndarray
1305 Input matrix slices of shape `(m, m)`.
1306 out : ndarray
1307 Output matrix slices of shape `(m, m)`.
1308 """
1309 AAT22_kernel(A, out)
1311@nb.njit(parallel = True, cache =True, fastmath = FASTMATH)
1312def _AAT22_pnjit(A, out):
1313 """Parallel Numba kernel computing A @ A^T over an array of matrices.
1315 Parameters
1316 ----------
1317 A : ndarray
1318 Batch of matrices of shape `(N, 2, 2)`.
1319 out : ndarray
1320 Output buffer of shape `(N, 2, 2)`.
1321 """
1322 for j in nb.prange(A.shape[0]): # pylint: disable=not-an-iterable
1323 AAT22_kernel(A[j], out[j])
1326# Wrapper Python propre
1327def AAT22_guvec(A: np.ndarray, out=None) -> np.ndarray:
1328 """Compute A @ A^T for batches of 2x2 matrices via Numba guvectorize.
1330 Parameters
1331 ----------
1332 A : np.ndarray
1333 Array of matrices of shape `(..., 2, 2)`.
1334 out : np.ndarray, optional
1335 Destination buffer. Defaults to None.
1337 Returns
1338 -------
1339 np.ndarray
1340 Symmetric matrix products of shape `(..., 2, 2)`.
1341 """
1342 if out is None:
1343 out = np.empty_like(A)
1344 _AAT22_guvec(A, out=out) # type: ignore
1345 return out
1347AAT22 = AAT22_guvec
1349def _cross_ijkl(A:FloatArray, B: FloatArray) -> FloatArray:
1350 """Compute the fourth-order tensor outer product C_{ijkl} = A_{ij} * B_{kl}.
1352 Parameters
1353 ----------
1354 A : FloatArray
1355 First second-order tensor batch of shape `(..., N, N)`.
1356 B : FloatArray
1357 Second second-order tensor batch of shape `(..., N, N)`.
1359 Returns
1360 -------
1361 FloatArray
1362 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1363 """
1364 return np.einsum("...ij, ...kl-> ...ijkl", A, B)
1366cross_ijkl = _cross_ijkl
1368def cross_ikjl(A:FloatArray, B: FloatArray) -> FloatArray:
1369 """Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{ik} * B_{jl}.
1371 Parameters
1372 ----------
1373 A : FloatArray
1374 First second-order tensor batch of shape `(..., N, N)`.
1375 B : FloatArray
1376 Second second-order tensor batch of shape `(..., N, N)`.
1378 Returns
1379 -------
1380 FloatArray
1381 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1382 """
1383 return np.einsum("...ik, ...jl-> ...ijkl", A, B)
1385def cross_iklj(A:FloatArray, B: FloatArray) -> FloatArray:
1386 """Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{ik} * B_{lj}.
1388 Parameters
1389 ----------
1390 A : FloatArray
1391 First second-order tensor batch of shape `(..., N, N)`.
1392 B : FloatArray
1393 Second second-order tensor batch of shape `(..., N, N)`.
1395 Returns
1396 -------
1397 FloatArray
1398 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1399 """
1400 return np.einsum("...ik, ...lj-> ...ijkl", A, B)
1402def cross_ilkj(A:FloatArray, B: FloatArray) -> FloatArray:
1403 """Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{il} * B_{kj}.
1405 Parameters
1406 ----------
1407 A : FloatArray
1408 First second-order tensor batch of shape `(..., N, N)`.
1409 B : FloatArray
1410 Second second-order tensor batch of shape `(..., N, N)`.
1412 Returns
1413 -------
1414 FloatArray
1415 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1416 """
1417 return np.einsum("...il, ...kj-> ...ijkl", A, B)
1419def cross_likj(A:FloatArray, B: FloatArray) -> FloatArray:
1420 """Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{li} * B_{kj}.
1422 Parameters
1423 ----------
1424 A : FloatArray
1425 First second-order tensor batch of shape `(..., N, N)`.
1426 B : FloatArray
1427 Second second-order tensor batch of shape `(..., N, N)`.
1429 Returns
1430 -------
1431 FloatArray
1432 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1433 """
1434 AijBkl= cross_ijkl(A,B)
1435 axis = list(range(AijBkl.ndim-4)) + [-3,-1,-2,-4]
1436 AliBkj = AijBkl.transpose(axis)
1437 return AliBkj
1439def cross_iljk(A:FloatArray, B: FloatArray) -> FloatArray:
1440 """Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{il} * B_{jk}.
1442 Parameters
1443 ----------
1444 A : FloatArray
1445 First second-order tensor batch of shape `(..., N, N)`.
1446 B : FloatArray
1447 Second second-order tensor batch of shape `(..., N, N)`.
1449 Returns
1450 -------
1451 FloatArray
1452 Fourth-order tensor batch of shape `(..., N, N, N, N)`.
1453 """
1454 AijBkl = cross_ijkl(A,B)
1455 axis = list(range(AijBkl.ndim-4)) + [-4,-2,-1,-3]
1456 AilBjk = AijBkl.transpose(axis)
1457 return AilBjk
1459def df_num( f:Callable[[FloatArray], FloatArray|float], A_in:FloatArray|float,
1460 shape_iA:tuple[int,...], shape_ifA:tuple[int,...])-> FloatArray:
1461 """Compute numerical Jacobian derivatives of tensor functions via finite differences.
1463 Uses `scipy.optimize.approx_fprime` iteratively over batch dimensions.
1465 Parameters
1466 ----------
1467 f : Callable[[FloatArray], FloatArray or float]
1468 Tensor function to differentiate.
1469 A_in : FloatArray or float
1470 Evaluation point or batch of evaluation points.
1471 shape_iA : tuple of int
1472 Shape of a single input argument to `f`.
1473 shape_ifA : tuple of int
1474 Shape of a single output evaluation from `f`.
1476 Returns
1477 -------
1478 FloatArray
1479 Numerical derivative tensor of shape `shape_coll + shape_ifA + shape_iA`.
1480 """
1481 A = np.atleast_1d(A_in)
1482 dim_coll = A.ndim - len(shape_iA)
1483 shape_coll = A.shape[:dim_coll]
1484 shape_idfA = shape_ifA + shape_iA
1485 df_num_res = np.zeros(shape_coll+shape_idfA)
1486 def fflat(iAflat):
1487 fiA = f(iAflat.reshape(shape_iA))
1488 if isinstance(fiA, np.ndarray) :
1489 fiA = fiA.flatten()
1490 return fiA
1491 for k in np.ndindex(A.shape[:dim_coll]):
1492 Ak = A[k]
1493 df_num_res[k] = sp.optimize.approx_fprime(Ak.flatten(), fflat).reshape(shape_idfA)
1494 if not isinstance(A_in, np.ndarray) :
1495 df_num_res = df_num_res.squeeze()
1496 return df_num_res
1498def check_relerr(ref:FloatArray|npfloat, val:FloatArray|npfloat, innershape,
1499 reltol:float = 1.e-4, abstol:float = 1.e-5,
1500 checkname:str="", verbose:bool = True )-> bool:
1501 """Compare two tensor arrays elementwise using absolute and relative error tolerances.
1503 Parameters
1504 ----------
1505 ref : FloatArray or npfloat
1506 Reference ground-truth values.
1507 val : FloatArray or npfloat
1508 Computed values to validate.
1509 innershape : tuple of int
1510 Shape of each tensor element inside the collection.
1511 reltol : float, optional
1512 Maximum allowed relative error. Defaults to 1e-4.
1513 abstol : float, optional
1514 Maximum allowed absolute error. Defaults to 1e-5.
1515 checkname : str, optional
1516 Name or label describing the test. Defaults to "".
1517 verbose : bool, optional
1518 If True, prints detailed diagnostics upon failure. Defaults to True.
1520 Returns
1521 -------
1522 bool
1523 True if all slices pass within error tolerances, False otherwise.
1524 """
1525 ref = np.asarray(ref).reshape((-1,)+ innershape)
1526 val = np.asarray(val).reshape((-1,)+ innershape)
1527 n = ref.shape[0]
1528 for i in range(n):
1529 err = np.linalg.norm((ref[i]-val[i]))
1530 refnorm = np.linalg.norm(ref)
1531 if err > abstol and err > refnorm*reltol :
1532 if verbose:
1533 print(f"check {checkname} Failed, for abstol = {abstol:.4e}, reltol {reltol:.4e}" )
1534 print(f" index {i}, error norm : {err :.4e}" + \
1535 f" relative error norm : {err/refnorm:.4e}")
1536 return False
1537 return True
1539def check_df(f, df, x, ifshape, ixshape, verbose = False):
1540 """Verify an analytical derivative against finite-difference approximation.
1542 Parameters
1543 ----------
1544 f : Callable
1545 Original function.
1546 df : Callable
1547 Analytical derivative function.
1548 x : FloatArray
1549 Evaluation inputs.
1550 ifshape : tuple of int
1551 Output shape of f for a single input item.
1552 ixshape : tuple of int
1553 Input shape for a single item.
1554 verbose : bool, optional
1555 If True, prints verification mismatch details. Defaults to False.
1557 Returns
1558 -------
1559 bool
1560 True if the analytical and numerical gradients agree, False otherwise.
1561 """
1562 dfx_exa = df(x)
1563 dfx_num = df_num(f, x, ixshape, ifshape)
1564 return check_relerr(dfx_exa, dfx_num, ifshape+ixshape, verbose = verbose)
1566# The following block needs mygrad library ...
1567# def df_vector_auto(Fun, X):
1568# """ compute the auto derivative of fun(A) at A where A is a vector and fun(A) a vector"""
1569# if not HAS_MYGRAD:
1570# frame = inspect.currentframe()
1571# fun_name = frame.f_code.co_name if frame else "df_vector_auto"
1572# raise ImportError(f"mygrad library is needed to run {fun_name}")
1573# else:
1574#
1575# dim = X.shape[-1]
1576# dFundX_shape = X.shape + (dim,)
1577# dFundX = np.zeros(dFundX_shape)
1579# for i in range(dim):
1580# mgX = mgTensor(X)
1581# fxi = Fun(mgX)[..., i]
1582# fxi.backward()
1583# dFundX[...,i,:] = np.array(mgX.grad)
1584# return dFundX
1586# def df_tensor_auto(Fun, X, shape_fun=None):
1587#
1588# """ compute the auto derivative of fun(A) at A where A is a tensor and fun(A) a tensor"""
1589# if not HAS_MYGRAD:
1590# frame = inspect.currentframe()
1591# fun_name = frame.f_code.co_name if frame else "df_tensor_auto"
1592# raise ImportError(f"mygrad library is needed to run {fun_name}")
1593# xdim0 = X.shape[-2]
1594# xdim1 = X.shape[-1]
1595# if shape_fun is None:
1596# shape_fun=(xdim0, xdim1)
1597# fdim0 = shape_fun[0]
1598# fdim1 = shape_fun[1]
1599# shapedfdx = X.shape[:-2] + ( fdim0, fdim1, xdim0, xdim1)
1600# dFdX = np.empty(shapedfdx)
1601# mgX = mgTensor(X)
1602# for i in range(fdim0) :
1603# for j in range(fdim1):
1604# FX = Fun(mgX)[..., i, j]
1605# FX.backward()
1606# dFdX[..., i,j, :,: ] = mgX.grad
1607# return dFdX
1609def eigval22sym(A:FloatArray):
1610 """Compute analytical eigenvalues for a batch of symmetric 2x2 matrices.
1612 Parameters
1613 ----------
1614 A : FloatArray
1615 Collection of symmetric 2x2 matrices of shape `(..., 2, 2)`.
1617 Returns
1618 -------
1619 FloatArray
1620 Eigenvalues array of shape `(..., 2)` with entries `[..., 0]` (max) and `[..., 1]` (min).
1622 Raises
1623 ------
1624 ValueError
1625 If trailing dimensions are not (2, 2) or matrices have negative discriminant.
1626 """
1627 if A.shape[-2:] != (2,2) :
1628 raise ValueError('A must be of shape[-2:] == (2,2)')
1629 a = A[..., 0, 0]
1630 b = A[..., 0, 1]
1631 c = A[..., 1, 1]
1632 disc = (a-c)**2 + 4*b**2
1633 if np.any(disc < 0):
1634 raise ValueError('all 2*2 matrix in a must be positive definite')
1635 eigsvals = np.empty(A.shape[:-2]+(2,))
1636 disc = np.sqrt(disc)
1637 eigsvals[..., 0] = (a+c + disc )/2.
1638 eigsvals[..., 1] = (a+c - disc )/2.
1639 return eigsvals
1641def max_eigval22sym(A:FloatArray, outlmax : Optional[FloatArray] = None):
1642 """Compute the maximum eigenvalue for a batch of symmetric 2x2 matrices.
1644 Parameters
1645 ----------
1646 A : FloatArray
1647 Batch of symmetric 2x2 matrices of shape `(..., 2, 2)`.
1648 outlmax : FloatArray, optional
1649 Output buffer receiving eigenvalues. Defaults to None.
1651 Returns
1652 -------
1653 FloatArray
1654 Maximum eigenvalues array of shape `(...)`.
1656 Raises
1657 ------
1658 ValueError
1659 If trailing dimensions are not (2, 2) or matrices have negative discriminant.
1660 """
1661 if A.shape[-2:] != (2,2) :
1662 raise ValueError('A must be of shape[-2:] == (2,2)')
1663 if outlmax is None:
1664 outlmax = np.empty(A.shape[:-2])
1665 a = A[..., 0, 0]
1666 b = A[..., 0, 1]
1667 c = A[..., 1, 1]
1668 disc = (a-c)**2 + 4*b**2
1669 if np.any(disc < 0):
1670 raise ValueError('all 2*2 matrix in a must be positive definite')
1672 disc = np.sqrt(disc)
1673 outlmax[...] = (a+c + disc )/2.
1674 return outlmax
1676@nb.njit(cache=True, fastmath = FASTMATH)
1677def max_eigval22sym_kernel_njit(A:FloatArray):
1678 """Numba scalar kernel computing the maximum eigenvalue of a single symmetric 2x2 matrix.
1680 Parameters
1681 ----------
1682 A : FloatArray
1683 2x2 symmetric matrix slice.
1685 Returns
1686 -------
1687 float
1688 Maximum eigenvalue.
1690 Raises
1691 ------
1692 ValueError
1693 If the discriminant is negative.
1694 """
1695 a = A[0, 0]
1696 b = A[0, 1]
1697 c = A[1, 1]
1698 disc = (a-c)**2 + 4*b**2
1699 if disc < 0:
1700 raise ValueError('2*2 matrix in a must be positive definite')
1701 return (a+c + np.sqrt(disc ))/2.