#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# pylint: disable=invalid-name
"""
This module i
s dedicated to implement faster functionnality for tensor on spaces of dim 2 and 3
than the equivalent numpy call
Created on May1 2024
@author: nchevaug
"""
from types import NoneType
from typing import Tuple, Callable, Literal, Optional
from warnings import warn
import numpy as np
import scipy as sp
from numba import float32, float64, guvectorize
import numba as nb
from pyxfem.core.config import FASTMATH, npfloat, FloatArray
HAS_MYGRAD = False
# o2 idetity Id2.u = u
Id2_2D = np.eye(2)
Id2_3D = np.eye(3)
Id2 = {2:Id2_2D, 3:Id2_3D}
# o4 identity Id4:A = A
Id4_3D = np.eye(9).reshape((3,3,3,3))
Id4_2D = np.eye(4).reshape((2,2,2,2))
Id4 = {2:Id4_2D, 3:Id4_3D}
# o4 Transposition Tr4:A = A^T
Tr4_2D = Id4_2D.swapaxes(2,3) #O4 transposr
Tr4_3D = Id4_3D.swapaxes(2,3) #O4 transposr
Tr4 = {2:Tr4_2D, 3:Tr4_3D}
# o4 Symetrizer Sy4:A = (A +A^T)/2.
Sy4_2D = 0.5*(Tr4_2D+Id4_2D) # O4 symetrizer
Sy4_3D = 0.5*(Tr4_3D+Id4_3D) # O4 symetrizer
Sy4 = {2:Sy4_2D, 3:Sy4_3D}
# o4 spherificator Sp4:A = trace(A) I
Sp4_2D = np.tensordot(Id2_2D , Id2_2D, axes = 0 )
Sp4_3D = np.tensordot(Id2_3D , Id2_3D, axes = 0 )
Sp4 = {2:Sp4_2D, 3:Sp4_3D}
[docs]
@nb.njit(parallel = True, cache = True)
def axpy_pnjit(alpha, x, y):
"""Compute vector scale-and-add operation in parallel using Numba.
Evaluates the in-place operation y[j] += alpha * x[j] over the first axis.
Parameters
----------
alpha : float or scalar
Scalar multiplier.
x : ndarray
Input array.
y : ndarray
Output array modified in place.
"""
for j in nb.prange(x.shape[0]): # pylint: disable=not-an-iterable
y[j] += alpha*x[j]
[docs]
def axpy(alpha, x, y, parallel = False):
"""Compute the scaled vector addition y = y + alpha * x.
Parameters
----------
alpha : float or scalar
Scalar scaling factor.
x : ndarray
Input array.
y : ndarray
Target array modified in place.
parallel : bool, optional
If True, executes the multithreaded Numba implementation `axpy_pnjit`.
Defaults to False.
"""
if parallel :
axpy_pnjit(alpha, x, y)
else:
y += alpha*x
[docs]
def vecprodmat(n:FloatArray) -> FloatArray:
"""Construct skew-symmetric cross-product matrices from 3D vectors.
For each 3D vector n, constructs the tensor Q such that Q @ x = n x x.
Parameters
----------
n : FloatArray
Array of 3D vectors with trailing dimension of length 3 (shape `(..., 3)`).
Returns
-------
FloatArray
Array of skew-symmetric 3x3 matrices with shape `(..., 3, 3)`.
"""
assert n.shape[-1] == 3
Q = np.zeros(n.shape[:-1] + (3,3))
Q[...,0,1] = -n[...,2]
Q[...,0,2] = n[...,1]
Q[...,1,0] = n[...,2]
Q[...,1,2] = -n[...,0]
Q[...,2,0] = -n[...,1]
Q[...,2,1] = n[...,0]
return Q
@guvectorize([
('float64[:], float64[:]'),
('float32[:], float32[:]')
], '(d)->()', nopython=True, cache=True)
def vec_norm2(v, out):
"""Compute the Euclidean L2 norm of 1D vector slices via Numba Generalized Ufunc.
Parameters
----------
v : ndarray
1D vector or array slice of dimension `(d,)`.
out : ndarray
Output scalar buffer receiving the computed norm.
"""
s = 0.0
for i in range(v.shape[0]):
s += v[i]**2
out[0] = np.sqrt(s)
[docs]
def metric(F:FloatArray) ->FloatArray:
"""Compute the right Cauchy-Green deformation tensor / Riemannian metric C = F^T @ F.
Parameters
----------
F : FloatArray
Deformation gradient or transformation tensor of shape `(..., M, N)`.
Returns
-------
FloatArray
Metric tensor array of shape `(..., N, N)`.
"""
return trans(F)@F
[docs]
def dmetric(F:FloatArray) -> FloatArray:
"""Compute the first derivative of the metric tensor m = F^T @ F with respect to F.
Returns the fourth-order tensor R[..., a, b, i, c] = dm_{ab} / dF_{ic}.
Parameters
----------
F : FloatArray
Transformation tensor of shape `(..., dim_out, dim_in)`.
Returns
-------
FloatArray
Fourth-order gradient tensor collection of shape `(..., dim_in, dim_in, dim_out, dim_in)`.
"""
dim = F.shape[-1]
I = np.eye(dim)
return np.einsum("...ac,...ib->...abic", I, F) + np.einsum("...bc,...ia->...abic", I, F)
[docs]
def d2metric(F:FloatArray) -> FloatArray:
"""Compute the second derivative of the metric tensor m with respect to F.
Returns the constant sixth-order tensor R[..., a, b, i, c, j, d] = d^2 m_{ab} / (dF_{ic} dF_{jd}).
Parameters
----------
F : FloatArray
Transformation tensor of shape `(..., dim_out, dim_in)`.
Returns
-------
FloatArray
Sixth-order tensor of shape `(..., dim_in, dim_in, dim_out, dim_in, dim_out, dim_in)`.
"""
I0 = np.eye(F.shape[-2])
I1 = np.eye(F.shape[-1])
d2m = np.einsum("...ac, ...ij, ...bd ->...abicjd", I1, I0, I1) \
+ np.einsum("...bc, ...ij, ...ad ->...abicjd", I1, I0, I1)
d2m = d2m.reshape((1,)*(F.ndim-2) + d2m.shape, copy=False)
return d2m
[docs]
def invmetric(F:FloatArray)->FloatArray:
"""Compute the inverse of the metric tensor m = F^T @ F.
Parameters
----------
F : FloatArray
Transformation tensor of shape `(..., M, N)`.
Returns
-------
FloatArray
Inverse metric tensor array of shape `(..., N, N)`.
"""
return inv(metric(F))
[docs]
def dinvmetric(F):
"""Compute the derivative of the inverse metric tensor with respect to F.
Evaluates d(m^{-1})_{ab} / dF_{ic}.
Parameters
----------
F : FloatArray
Transformation tensor of shape `(..., M, N)`.
Returns
-------
FloatArray
Fourth-order derivative tensor of shape `(..., N, N, M, N)`.
"""
im = invmetric(F)
return -np.einsum("...ac,...ib->...abic", im, F@im)-np.einsum("...cb,...ia->...abic", im, F@im)
[docs]
def trace(A:FloatArray) -> FloatArray:
"""Compute the trace of a batch of square second-order tensors.
Parameters
----------
A : FloatArray
Square tensor collection of shape `(..., N, N)`.
Returns
-------
FloatArray
Trace values array of shape `(...)`.
Raises
------
ValueError
If the trailing dimensions are not square.
"""
if A.shape[-1] != A.shape[-2]:
raise ValueError("in trace, each array in the collection must be square")
dim = A.shape[-1]
if dim == 1:
tr = A[...,0,0]
elif dim == 2:
tr = A[...,0,0] + A[...,1,1]
elif dim == 3:
tr = A[...,0,0] + A[...,1,1] + A[...,2,2]
else:
tr = np.einsum('...ii', A)
return tr
[docs]
def trace2(A:FloatArray) -> FloatArray:
"""Compute the Frobenius squared norm / trace(A^T @ A) for a batch of tensors.
Parameters
----------
A : FloatArray
Tensor collection of shape `(..., N, N)`.
Returns
-------
FloatArray
Array containing the sum of squared components for each tensor.
Raises
------
ValueError
If the trailing dimensions are not square.
"""
if A.shape[-1] != A.shape[-2]:
raise ValueError("in trace, each array in the collection must be square")
return np.sum(np.square(A), axis = (-1,-2))
[docs]
def trans(A:FloatArray)-> FloatArray:
""" return the transposes of the input collection of second order tensors """
"""Transpose the last two dimensions of an array of matrices.
Parameters
----------
A : FloatArray
Tensor array of shape `(..., M, N)`.
Returns
-------
FloatArray
Transposed tensor array of shape `(..., N, M)`.
"""
return A.swapaxes(-1,-2)
[docs]
def contract(A:FloatArray, B:FloatArray)-> FloatArray:
"""Compute the double contraction (Frobenius inner product) A : B = A_{ij} B_{ij}.
Parameters
----------
A : FloatArray
First batch of tensors of shape `(..., M, N)`.
B : FloatArray
Second batch of tensors of matching shape `(..., M, N)`.
Returns
-------
FloatArray
Contracted scalar field of shape `(...)`.
"""
assert A.shape == B.shape
assert A.ndim >= 2
if A.shape[:-2] == (2,2):
return A[...,0,0]*B[...,0,0]+A[...,0,1]*B[...,0,1] \
+ A[...,1,0]*B[...,1,0]+ A[...,1,1]*B[...,1,1]
return np.einsum("...ij, ...ij-> ...", A, B)
def _cofac(A:FloatArray) -> FloatArray:
"""compute cofac matrix for stack of 2x2 or 3x3 matrices """
"""Compute the cofactor matrix for batches of 2x2 or 3x3 matrices.
Parameters
----------
A : FloatArray
Tensor array of shape `(..., 2, 2)` or `(..., 3, 3)`.
Returns
-------
FloatArray
Cofactor matrix array with identical shape to `A`.
Raises
------
ValueError
If the trailing dimensions are neither (2, 2) nor (3, 3).
"""
if A.shape[-2:] == (2,2):
cA = np.empty_like(A)
cA[..., 0, 0] = A[..., 1, 1]
cA[..., 0, 1] = -A[..., 1, 0]
cA[..., 1, 0] = -A[..., 0, 1]
cA[..., 1, 1] = A[..., 0, 0]
return cA
if A.shape[-2:] == (3,3):
cA = np.empty_like(A)
cA[...,0,:] = cross(A[...,1,:], A[...,2,:])
cA[...,1,:] = cross(A[...,2,:], A[...,0,:])
cA[...,2,:] = cross(A[...,0,:], A[...,1,:])
return cA
raise ValueError("A must be a (..,2,2) or (...,3,3) array_like")
@guvectorize([
(float64[:,:], float64[:]),
(float32[:,:], float32[:])
], '(m,m)->()', nopython=True, cache=True)
def det22_numba(A, out):
"""Compute the determinant of a single 2x2 matrix via Numba Generalized Ufunc.
Parameters
----------
A : ndarray
2x2 input matrix slice of shape `(2, 2)`.
out : ndarray
Output buffer receiving the determinant scalar.
"""
out[0] = A[0,0]*A[1,1]- A[0,1]*A[1,0]
def _det(A:FloatArray, out:NoneType|FloatArray=None) -> FloatArray|npfloat:
"""Compute the determinant for batches of 2x2, 3x3, or 3x2 (metric surface) matrices.
Parameters
----------
A : FloatArray
Matrix or batch of matrices of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
out : FloatArray or None, optional
Pre-allocated output destination buffer. Defaults to None.
Returns
-------
FloatArray or npfloat
Determinants array of shape `(...)`. For 3x2, returns sqrt(det(F^T F)).
Raises
------
ValueError
If matrix dimensions are not supported.
"""
if out is None :
out = np.empty_like(A[...,0,0])
if A.shape[-2:] == (3,2):
m = metric(A)
_det(m, out)
return np.sqrt(out, out)
if A.shape[-2:] == (2,2):
det22_numba(A, out)
return out
#return A[...,0,0]*A[...,1,1]- A[...,0,1]*A[...,1,0]
if A.shape[-2:] == (3,3):
a= A[...,0,0]
b= A[...,0,1]
c= A[...,0,2]
d= A[...,1,0]
e= A[...,1,1]
f= A[...,1,2]
g= A[...,2,0]
h= A[...,2,1]
i= A[...,2,2]
if np.isscalar(out):
out = a*e*i +b*f*g +c*d*h - c*e*g -b*d*i -a*f*h
else:
out[...] = a*e*i +b*f*g +c*d*h - c*e*g -b*d*i -a*f*h
return out
raise ValueError(f"Size error {str(A.shape)}")
[docs]
def det_invt(A:FloatArray) -> Tuple[FloatArray|npfloat, FloatArray]:
"""Compute both the determinant and the inverse-transpose A^{-T} simultaneously.
Also handles the 3x2 membrane configuration where A^{-T} is computed as A @ (A^T A)^{-1}.
Parameters
----------
A : FloatArray
Tensor array of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
Returns
-------
tuple of (FloatArray or npfloat, FloatArray)
Tuple (detA, invAT) containing determinants and inverse-transposed matrices.
Raises
------
ValueError
If trailing dimensions are not (2, 2), (3, 3), or (3, 2).
"""
if A.shape[-2:] == (3,2):
m = metric(A)
detm = det(m)
invAT = A@inv(m)
return np.sqrt(detm), invAT
if A.shape[-2:] in [(2,2), (3,3)]:
detA = det(A)
adjA = cofac(A)
if A.shape[:-2]== ():
invAT = adjA/detA
else :
invAT = adjA / np.atleast_1d(detA)[..., None, None]
return detA, invAT
raise ValueError("A.shape[-2,:] should be (2, 2) or (3, 3) or (3,2)"+str(A.shape) )
det = _det
[docs]
def ddet(A:FloatArray)->FloatArray:
"""Compute the derivative of the determinant with respect to the matrix A.
Returns d(det(A)) / dA = cofac(A) for square matrices, or det(A) * A^{-T}.
Parameters
----------
A : FloatArray
Input tensor collection of shape `(..., M, N)`.
Returns
-------
FloatArray
Derivative array having the identical shape of `A`.
"""
if A.shape[-1]==A.shape[-2]:
return cofac(A)
J, invAT = det_invt(A)
return ((np.atleast_1d(J)[:, np.newaxis, np.newaxis])*invAT).reshape(A.shape)
cofac = _cofac
[docs]
def cross(A:FloatArray, B:FloatArray) -> FloatArray:
"""Compute the 3D vector cross product between arrays of vectors.
Parameters
----------
A : FloatArray
First collection of 3D vectors of shape `(..., 3)`.
B : FloatArray
Second collection of 3D vectors of shape `(..., 3)`.
Returns
-------
FloatArray
Cross product vector array of shape `(..., 3)`.
Raises
------
ValueError
If the trailing dimension of A or B is not 3.
"""
if isinstance(A, np.ndarray) and isinstance(B, np.ndarray):
return np.cross(A,B)
# if A.shape[-1] != 3 or B.shape[-1] !=3 :
# raise ValueError("A and be most be of shape (..., 3) ")
# C = np.empty_like(A)
# C[...,0] = A[..., 1] * B[...,2] - A[..., 2] * B[...,1]
# C[...,1] = -A[..., 0] * B[...,2] + A[..., 2] * B[...,0]
# C[...,2] = A[..., 0] * B[...,1] - A[..., 1] * B[...,0]
# return C
[docs]
def invariants(A:FloatArray) -> FloatArray:
"""Compute principal invariants for a batch of 2D or 3D second-order tensors.
For 2D: [tr(A), det(A)].
For 3D: [tr(A), 0.5 * (tr(A)^2 - tr(A^T A)), det(A)].
Parameters
----------
A : FloatArray
Square tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
Returns
-------
FloatArray
Invariants array with trailing dimension 2 (for 2D) or 3 (for 3D).
Raises
------
ValueError
If trailing dimensions are not (2, 2) or (3, 3).
"""
if A.shape[-2:] == (2,2):
r = np.empty( A.shape[:-2]+(2,))
r[...,0] = trace(A)
r[...,1] = det(A)
return r
if A.shape[-2:] == (3,3):
r = np.empty( A.shape[:-2]+(3,))
trA = trace(A)
r[...,0] = trA
r[...,1] = 0.5*( np.square(trA) - trace(trans(A)@A) )
r[...,2] = det(A)
return r
raise ValueError("A must be of size (..., 2, 2) or (..., 3, 3)")
[docs]
def dinvariants(A:FloatArray) -> FloatArray:
"""Compute derivatives of principal invariants with respect to tensor A.
Parameters
----------
A : FloatArray
Square tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
Returns
-------
FloatArray
Derivatives tensor of shape `(..., 2, 2, 2)` for 2D or `(..., 3, 3, 3)` for 3D.
Raises
------
ValueError
If trailing dimensions are not (2, 2) or (3, 3).
"""
if A.shape[-2:] == (2,2):
r = np.empty( A.shape[:-2]+(2,2,2))
I = np.eye(2)[(np.newaxis,)*len(A.shape[:-2]) ]
r[...,0, :, :] = I
r[...,1, : ,:] = cofac(A)
return r
if A.shape[-2:] == (3,3):
nfirstdim = len(A.shape[:-2])
r = np.empty( A.shape[:-2]+(3,3,3))
I = np.eye(3)[(np.newaxis,)*nfirstdim ]
trA = trace(A)[..., None, None]
r[...,0,:,:] = I
r[...,1,:,:] = trA*I-A
r[...,2,:,:] = cofac(A)
return r
raise ValueError("A must be of size (..., 2, 2) or (..., 3, 3)")
[docs]
def invariants2Dto3D(invariant2D):
""" from the invariant of a "2.5d tensor", compute the invariants of the equivalent 3D tensor
what I mean by a 2.5d tensor is for example the gradiant of the
transformation in plan stress case :
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
the invarariant 2D are strored in a (..., 3) ndarray, where
invariant2D[..., 0] is an array containing I(F2D) = trace(F2D)
invariant2D[..., 1] is an array containing
II(F2D) = 0.5(trace(F2D)^2 - trace(F2D.T@F2D)) = det(F2D)
invariant2D[..., 2] is an array containing l2
Parameters
----------
invariant2D : ndarray
Array of shape `(..., 3)` containing:
- [..., 0]: I_2D = trace(F_2D)
- [..., 1]: II_2D = det(F_2D)
- [..., 2]: lambda_2 (out-of-plane stretch)
Returns
-------
ndarray
Array of shape `(..., 3)` with the three corresponding 3D invariants.
"""
invariant3D = np.empty(invariant2D.shape)
I2D = invariant2D[...,0]
II2D = invariant2D[...,1]
l2 = invariant2D[...,2]
invariant3D[..., 0] = I2D + l2
invariant3D[..., 1] = II2D + l2*I2D
invariant3D[..., 2] = II2D * l2
return invariant3D
[docs]
def dinvariants2Dto3D(invariant2D:FloatArray) -> FloatArray:
""" Compute the matrix of the derivatives of the 3D invariants of tensor A
computed as a function of the 2D invariant.
input invariant2D an FloatArray of shape[-1] = 3 where :
invariant2D[..., 0] = I2D = trace(A2D)
invariant2D[..., 1] = II2D = det(A2D) = 0.5*(trace(A2D)**2 - trace(A2D@A2D**2)
invariant2D[..., 2] = A33
return dinvariant3D an FNDarray of shape (..., 3, 3)
where dinvariant3D[..., i, j] contain the derivative of the ith 3d invarariant with regard
to the jth 2d invariant
Parameters
----------
invariant2D : FloatArray
Array of shape `(..., 3)` holding [I_2D, II_2D, lambda_2].
Returns
-------
FloatArray
Jacobian tensor of shape `(..., 3, 3)` where entry `[..., i, j]` is d(I_3D, i) / d(I_2D, j).
"""
I2D = invariant2D[...,0]
II2D = invariant2D[...,1]
l2 = invariant2D[...,2]
dinvariant3D = np.ones(invariant2D.shape + (3,))
dinvariant3D[..., 0, 1] = 0.
dinvariant3D[..., 1, 0] = l2
dinvariant3D[..., 1, 2] = I2D
dinvariant3D[..., 2, 0] = 0.
dinvariant3D[..., 2, 1] = l2
dinvariant3D[..., 2, 2] = II2D
return dinvariant3D
[docs]
def d2invariants2Dto3D(invariant2D):
""" Compute the matrix of the second derivatives of the 3D invariants of tensor A (Hessian)
computed as a function of the 2D invariants.
input invariant2D an FloatArray of shape[-1] = 3 where :
invariant2D[..., 0] = I2D = trace(A2D)
invariant2D[..., 1] = II2D = det(A2D) = 0.5*(trace(A2D)**2 - trace(A2D@A2D**2)
invariant2D[..., 2] = A33
return dinvariant3D an FNDarray of shape (..., 3, 3, 3)
where dinvariant3D[..., i, j, k] contain the derivative of the ith 3d invarariant
with regard to the jth 2d invariant and the kth 2d invarariant.
Note that the resuluting 27 component matrix contain only 4 non zero terms ...
Parameters
----------
invariant2D : ndarray
Array of shape `(..., 3)` holding [I_2D, II_2D, lambda_2].
Returns
-------
ndarray
Second derivative tensor of shape `(..., 3, 3, 3)` containing d^2(I_3D, i) / (d(I_2D, j) d(I_2D, k)).
"""
d2invariant3D = np.zeros(invariant2D.shape + (3, 3))
d2invariant3D[..., 1, 0, 2] = 1.
d2invariant3D[..., 1, 2, 0] = 1.
d2invariant3D[..., 2, 1, 2] = 1.
d2invariant3D[..., 2, 2, 1] = 1.
return d2invariant3D
[docs]
def inv(A:FloatArray, detA: Optional[FloatArray|npfloat] = None,
out: Optional[FloatArray] =None ) -> FloatArray:
"""compute inverse matrix for stack of 2x2 or 3x3 matrices """
"""Compute matrix inverse for batches of 2x2, 3x3, or pseudo-inverse for 3x2 matrices.
Parameters
----------
A : FloatArray
Tensor collection of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
detA : FloatArray or npfloat, optional
Precomputed determinant of A to accelerate evaluation. Defaults to None.
out : FloatArray, optional
Output buffer receiving the inverted matrix. Defaults to None.
Returns
-------
FloatArray
Inverted matrix collection.
"""
if A.shape[-2:] == (3,2):
out = np.empty(A.shape[:-2] + (2,3)) if out is None else out
m = metric(A)
np.matmul(inv(m), trans(A), out)
return out
if detA is None :
detA = det(A)
if not isinstance(detA, (float, int)):
detA = detA.reshape(A.shape[:-2]+(1,1))
iA = np.empty_like(A) if out is None else out
if A.shape[-2:] == (2,2):
iA[..., 0,0] = A[...,1,1]
iA[..., 0,1] = -A[...,0,1]
iA[..., 1,0] = -A[...,1,0]
iA[..., 1,1] = A[...,0,0]
elif A.shape[-2:] == (3,3):
warn("3d case not fully optimized in inv")
iA[...,:,0] = cross(A[...,1,:], A[...,2,:])
iA[...,:,1] = cross(A[...,2,:], A[...,0,:])
iA[...,:,2] = cross(A[...,0,:], A[...,1,:])
if A.shape[:-2] == ():
iA = iA.squeeze()
iA /= detA
return iA
[docs]
def dinv(A:FloatArray, iA:FloatArray|None = None )-> FloatArray:
"""Compute the derivative of the inverse matrix with respect to A.
Returns R[..., i, j, k, l] = d(A^{-1}_{ij}) / dA_{kl}.
Parameters
----------
A : FloatArray
Tensor collection of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
iA : FloatArray or None, optional
Precalculated inverse of A. Defaults to None.
Returns
-------
FloatArray
Fourth-order derivative tensor.
"""
if A.shape[-2:] == (3,2):
im = inv(metric(A))
diA = np.empty(A.shape[:-2] + (2,3,3,2))
diA = -np.einsum("...ab,...de,...jd,...ie->...aijb", im, im, A, A )
diA -= np.einsum("...ad,...be,...jd,...ie->...aijb", im, im, A, A )
diA += np.einsum("...ab, ...ij -> ...aijb", im, np.eye(3))
return diA
iA = inv(A)
return -np.einsum("...ac,...db->...abcd", iA, iA)
[docs]
def invt(A:FloatArray, detA:FloatArray|npfloat|None = None) -> FloatArray:
"""Compute the inverse-transpose (or pseudo-inverse transpose) of a batch of matrices.
Parameters
----------
A : FloatArray
Tensor array of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
detA : FloatArray or npfloat, optional
Precomputed determinant. Defaults to None.
Returns
-------
FloatArray
Inverse transposed matrices of shape `(..., N, M)`.
"""
return np.swapaxes(inv(A, detA), -1,-2)
[docs]
def dinv_sym(A, iA = None) -> FloatArray:
"""Compute the derivative of A^{-1} with respect to A assuming A is symmetric.
Parameters
----------
A : FloatArray
Symmetric tensor collection of shape `(..., 2, 2)` or `(..., 3, 3)`.
iA : FloatArray or None, optional
Precomputed inverse of A. Defaults to None.
Returns
-------
FloatArray
Symmetric fourth-order derivative tensor.
Raises
------
ValueError
If trailing dimensions are not (2, 2) or (3, 3).
"""
if A.shape[-2:] in [(3,3), (2,2)]:
if iA is None :
iA = inv(A)
return -0.5*cross_iklj(iA, iA) -0.5*cross_ilkj(iA, iA)
raise ValueError(" A.shape[-2:] must be (2,2) or (3,3)")
[docs]
def dinvt(F, iFT =None, im = None, dimdF = None)->FloatArray:
"""Compute the derivative of the inverse-transpose tensor F^{-T} with respect to F.
Parameters
----------
F : FloatArray
Input deformation tensor of shape `(..., 2, 2)`, `(..., 3, 3)`, or `(..., 3, 2)`.
iFT : FloatArray, optional
Precomputed F^{-T}. Defaults to None.
im : FloatArray, optional
Precomputed inverse metric (F^T F)^{-1} for 3x2 tensors. Defaults to None.
dimdF : FloatArray, optional
Precomputed derivative of the inverse metric. Defaults to None.
Returns
-------
FloatArray
Fourth-order derivative tensor.
"""
if F.shape[-2:] == (3,2):
I = np.eye(3)
im = invmetric(F) if im is None else im
dimdF = dinvmetric(F) if dimdF is None else dimdF
return np.einsum("...ij, ...ba -> ...iajb", I, im) \
+ np.einsum("...il, ...lajb-> ...iajb", F, dimdF)
iFT = invt(F) if iFT is None else iFT
return -cross_ilkj(iFT, iFT)
[docs]
def sqrt22_numpy(A: np.ndarray, out=None) -> np.ndarray:
"""Compute tensor square roots of 2x2 matrices using vectorized NumPy views.
Parameters
----------
A : np.ndarray
Array of 2x2 matrices of shape `(..., 2, 2)`.
out : np.ndarray, optional
Output buffer modified in place. Defaults to None.
Returns
-------
np.ndarray
Tensor square roots of shape `(..., 2, 2)`.
Raises
------
ValueError
If trailing dimensions are not (2, 2).
"""
if A.shape[-2:] != (2, 2):
raise ValueError("A.shape[-2:] should be (2, 2)")
if out is None:
out = np.empty_like(A)
A00, A01 = A[..., 0, 0], A[..., 0, 1]
A10, A11 = A[..., 1, 0], A[..., 1, 1]
detA = A00 * A11 - A01 * A10
traA = A00 + A11
s = np.sqrt(detA)
t = np.sqrt(traA + 2.0 * s)
# Ajout d'axes pour le broadcasting de la division
# s et t passent de shape (Batch,) Ã (Batch, 1, 1)
t_expand = t[..., np.newaxis, np.newaxis]
# Calcul en place
# On commence par copier A dans out pour ne pas l'altérer
out[...] = A
# On ajoute s sur la diagonale
out[..., 0, 0] += s
out[..., 1, 1] += s
# On divise tout le tenseur par t
out[...] /= t_expand
return out
[docs]
@nb.njit(cache=True, fastmath = FASTMATH)
def sqrt22_kernel(A22, out):
"""Numba scalar kernel computing the square root of an individual 2x2 matrix.
Parameters
----------
A22 : ndarray
Single 2x2 matrix slice.
out : ndarray
Output 2x2 buffer receiving the result.
"""
A00 = A22[0, 0]
A01 = A22[0, 1]
A10 = A22[1, 0]
A11 = A22[1, 1]
# 1. Déterminant et s
detA = A00 * A11 - A01 * A10
s = np.sqrt(detA)
# 2. Trace et t
traA = A00 + A11
t = np.sqrt(traA + 2.0 * s)
# 3. Calcul final : B = (A + s * I) / t
out[0, 0] = (A00 + s) / t
out[0, 1] = A01 / t
out[1, 0] = A10 / t
out[1, 1] = (A11 + s) / t
@guvectorize([
(float64[:,:], float64[:,:]),
(float32[:,:], float32[:,:])
], '(m,m)->(m,m)', nopython=True, cache=True)
def _sqrt22_guvec(A, out):
"""Numba Generalized Ufunc wrapper for batch 2x2 tensor square root.
Parameters
----------
A : ndarray
Input matrix slices of shape `(m, m)`.
out : ndarray
Output matrix slices of shape `(m, m)`.
"""
sqrt22_kernel(A, out)
[docs]
def sqrt22_guvec(A: np.ndarray, out=None) -> np.ndarray:
"""Compute tensor square roots of 2x2 matrices using Numba guvectorize.
Parameters
----------
A : np.ndarray
Array of matrices of shape `(..., 2, 2)`.
out : np.ndarray, optional
Preallocated destination array. Defaults to None.
Returns
-------
np.ndarray
Tensor square roots of shape `(..., 2, 2)`.
Raises
------
ValueError
If trailing dimensions are not (2, 2).
"""
if A.shape[-2:] != (2, 2):
raise ValueError("A.shape[-2:] should be (2, 2)")
if out is None:
out = np.empty_like(A)
_sqrt22_guvec(A, out=out) # type: ignore
return out
[docs]
def sqrt22_njit(A: np.ndarray, out=None) -> np.ndarray:
"""
Calcule la racine carrée d'un batch de tenseurs 2x2.
A = B.B
""""""Compute tensor square roots of 2x2 matrices using multithreaded Numba NJIT.
Parameters
----------
A : np.ndarray
C-contiguous array of matrices of shape `(..., 2, 2)`.
out : np.ndarray, optional
Preallocated destination array. Defaults to None.
Returns
-------
np.ndarray
Tensor square roots of shape `(..., 2, 2)`.
Raises
------
ValueError
If A is not C-contiguous or its shape is not `(..., 2, 2)`.
"""
if not A.flags.c_contiguous :
raise ValueError("A not contiguous")
if A.shape[-2:] != (2, 2):
raise ValueError("A.shape[-2:] should be (2, 2)")
if out is None:
out = np.empty_like(A)
Aview = A.reshape((-1, 2, 2))
outview = out.reshape((-1, 2, 2))
_sqrt22_pnjit(Aview, out=outview) # type: ignore
return out
@nb.njit(parallel=True, cache = True, fastmath = FASTMATH)
def _sqrt22_pnjit(A, out):
"""Numba parallel loop over flattened batch of 2x2 matrices.
Parameters
----------
A : ndarray
Array of shape `(N, 2, 2)`.
out : ndarray
Output array of shape `(N, 2, 2)`.
Returns
-------
ndarray
Output buffer `out`.
"""
for i in nb.prange(len(A)): # type ignore # pylint: disable=E1133
sqrt22_kernel(A[i], out[i])
return out
[docs]
def sqrt22(A, out = None):
"""Compute square roots for a batch of 2x2 matrices (dispatcher).
Parameters
----------
A : np.ndarray
Array of matrices of shape `(..., 2, 2)`.
out : np.ndarray, optional
Destination buffer. Defaults to None.
Returns
-------
np.ndarray
Tensor square root array.
"""
return sqrt22_njit(A, out)
#return sqrt22_guvec(A, out)
[docs]
@nb.njit(cache =True)
def sysmv22_kernel(A: FloatArray, b: FloatArray, out : FloatArray):
"""Solve a single symmetric 2x2 linear system A @ x = b via Cramer's rule.
Parameters
----------
A : FloatArray
Symmetric 2x2 matrix.
b : FloatArray
Right-hand side 2D vector.
out : FloatArray
Output 2D vector receiving solution x.
"""
A00, A01, A11 = A[0,0], A[0,1], A[1,1]
b0, b1 = b
detA = A00*A11 - A01*A01
out[0] = (A11*b0 - A01*b1)/detA
out[1] = (A00*b1 - A01*b0)/detA
@guvectorize([
(float64[:,:], float64[:], float64[:]),
(float32[:,:], float32[:], float32[:])
], '(dim, dim),(dim), (dim)', nopython=True, cache=True)
def sysmv22_guvec(A, b, out : FloatArray):
"""Generalized Ufunc solving batches of symmetric 2x2 systems A @ x = b.
Parameters
----------
A : ndarray
Batch of 2x2 matrices of shape `(..., 2, 2)`.
b : ndarray
Batch of vectors of shape `(..., 2)`.
out : ndarray
Buffer receiving solutions of shape `(..., 2)`.
"""
sysmv22_kernel(A, b, out)
[docs]
@nb.njit(parallel = True, cache =True)
def sysmv22_pnjit(A: FloatArray, b: FloatArray, out : FloatArray):
"""Parallel Numba kernel solving batches of symmetric 2x2 systems A @ x = b.
Parameters
----------
A : FloatArray
Batch of 2x2 matrices of shape `(N, 2, 2)`.
b : FloatArray
Batch of right-hand side vectors of shape `(N, 2)`.
out : FloatArray
Output buffer receiving solutions of shape `(N, 2)`.
"""
for j in nb.prange(A.shape[0]): # pylint: disable=not-an-iterable
sysmv22_kernel(A[j], b[j], out[j])
[docs]
def sysmv22(A: FloatArray, b: FloatArray,
out:Optional[FloatArray] = None, parallel = False) -> FloatArray:
r"""
Solves the linear system $A x = b$ for a symmetric 2x2 matrix
using explicit Cramer's rule inversion.
Math: $x = A^{-1} b$
Parameters
----------
A : FloatArray
Batch of symmetric 2x2 matrices of shape `(..., 2, 2)`.
b : FloatArray
Right-hand side vectors of shape `(..., 2)`.
out : FloatArray, optional
Preallocated output buffer. Defaults to None.
parallel : bool, optional
If True, evaluates using multithreaded OpenMP-style execution. Defaults to False.
Returns
-------
FloatArray
Solution vectors of shape `(..., 2)`.
Raises
------
ValueError
If trailing dimensions of A are not (2, 2).
"""
if A.shape[-2:] != (2,2) :
raise ValueError('sysmv22 is for 2x2 batch of tensors' )
out = np.empty( b.shape, npfloat) if out is None else out
if parallel :
sysmv22_pnjit(A, b, out)
else :
sysmv22_guvec(A, b, out)
return out
[docs]
def gemm22(A: FloatArray, B: FloatArray, out:Optional[FloatArray] = None) -> FloatArray:
r"""
Compute matrix-matrix products C = A @ B for batches of 2x2 matrices.
Optimized for zero heap reallocations using views and in-place accumulations.
Parameters
----------
A : FloatArray
First batch of 2x2 matrices of shape `(..., 2, 2)`.
B : FloatArray
Second batch of 2x2 matrices of shape `(..., 2, 2)`.
out : FloatArray, optional
Preallocated array receiving the result. Defaults to None.
Returns
-------
FloatArray
Matrix product array of shape `(..., 2, 2)`.
Raises
------
ValueError
If trailing dimensions are not (2, 2).
"""
if A.shape[-2:] != (2,2) :
raise ValueError('gemm22 is for 2x2 batch of tensors' )
C = np.empty_like(A) if out is None else out
# 1. Create Views (0 allocations)
A00, A01 = A[..., 0, 0], A[..., 0, 1]
A10, A11 = A[..., 1, 0], A[..., 1, 1]
B00, B01 = B[..., 0, 0], B[..., 0, 1]
B10, B11 = B[..., 1, 0], B[..., 1, 1]
# 2. Allocate EXACTLY ONE temporary buffer for the batch
# This shape matches the batch dimensions (e.g., if A is 1000x2x2, temp is 1000)
temp = np.empty_like(A00)
# --- Compute C00 ---
# C00 = A00*B00
np.multiply(A00, B00, out=C[..., 0, 0])
# temp = A01*B10
np.multiply(A01, B10, out=temp)
# C00 += temp (In-place addition!)
C[..., 0, 0] += temp
# --- Compute C01 ---
np.multiply(A00, B01, out=C[..., 0, 1])
np.multiply(A01, B11, out=temp)
C[..., 0, 1] += temp
# --- Compute C10 ---
np.multiply(A10, B00, out=C[..., 1, 0])
np.multiply(A11, B10, out=temp)
C[..., 1, 0] += temp
# --- Compute C11 ---
np.multiply(A10, B01, out=C[..., 1, 1])
np.multiply(A11, B11, out=temp)
C[..., 1, 1] += temp
return C
[docs]
@nb.njit(cache=True, fastmath = True, inline='always')
def gemm22_kernel_njit(A:FloatArray, B: FloatArray, out: FloatArray):
"""Inlined Numba scalar kernel computing the product of two 2x2 matrices.
Parameters
----------
A : FloatArray
2x2 left-hand matrix slice.
B : FloatArray
2x2 right-hand matrix slice.
out : FloatArray
2x2 destination buffer.
"""
A00, A01, A10, A11 = A[0,0], A[0,1], A[1,0], A[1,1]
B00, B01, B10, B11 = B[0,0], B[0,1], B[1,0], B[1,1]
out[0,0] = A00*B00 + A01*B10
out[0,1] = A00*B01 + A01*B11
out[1,0] = A10*B00 + A11*B10
out[1,1] = A10*B01 + A11*B11
[docs]
def AAT22(A, out = None):
"""Compute the symmetric product B = A @ A^T for a batch of 2x2 matrices.
Parameters
----------
A : ndarray
Batch of 2x2 matrices of shape `(..., 2, 2)`.
out : ndarray, optional
Preallocated array receiving the result. Defaults to None.
Returns
-------
ndarray
Symmetric matrix product array of shape `(..., 2, 2)`.
"""
if out is None:
B = np.empty_like(A)
else:
B = out
#B = np.empty_like(A)
A00 = A[...,0,0]
A01 = A[...,0,1]
A10 = A[...,1,0]
A11 = A[...,1,1]
B[...,0,0] = A00*A00 + A01*A01
B[...,0,1] = A00*A10 + A01*A11
B[...,1,0] = A10*A00 + A11*A01
B[...,1,1] = A10*A10 + A11*A11
return B
[docs]
@nb.njit(cache=True, fastmath = False)
def AAT22_kernel(A22, out):
"""Numba kernel computing A @ A^T for an individual 2x2 matrix.
Parameters
----------
A22 : ndarray
Single 2x2 matrix slice.
out : ndarray
Output 2x2 buffer.
"""
A00 = A22[0, 0]
A01 = A22[0, 1]
A10 = A22[1, 0]
A11 = A22[1, 1]
out[0, 0] = A00*A00 + A01*A01
out[0, 1] = A00*A10 + A01*A11
out[1, 0] = A10*A00 + A11*A01
out[1, 1] = A10*A10 + A11*A11
@guvectorize([
(float64[:,:], float64[:,:]),
(float32[:,:], float32[:,:])
], '(m,m)->(m,m)', nopython=True, cache=True)
def _AAT22_guvec(A, out):
"""Numba Generalized Ufunc computing A @ A^T across matrix slices.
Parameters
----------
A : ndarray
Input matrix slices of shape `(m, m)`.
out : ndarray
Output matrix slices of shape `(m, m)`.
"""
AAT22_kernel(A, out)
@nb.njit(parallel = True, cache =True, fastmath = FASTMATH)
def _AAT22_pnjit(A, out):
"""Parallel Numba kernel computing A @ A^T over an array of matrices.
Parameters
----------
A : ndarray
Batch of matrices of shape `(N, 2, 2)`.
out : ndarray
Output buffer of shape `(N, 2, 2)`.
"""
for j in nb.prange(A.shape[0]): # pylint: disable=not-an-iterable
AAT22_kernel(A[j], out[j])
# Wrapper Python propre
[docs]
def AAT22_guvec(A: np.ndarray, out=None) -> np.ndarray:
"""Compute A @ A^T for batches of 2x2 matrices via Numba guvectorize.
Parameters
----------
A : np.ndarray
Array of matrices of shape `(..., 2, 2)`.
out : np.ndarray, optional
Destination buffer. Defaults to None.
Returns
-------
np.ndarray
Symmetric matrix products of shape `(..., 2, 2)`.
"""
if out is None:
out = np.empty_like(A)
_AAT22_guvec(A, out=out) # type: ignore
return out
AAT22 = AAT22_guvec
def _cross_ijkl(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor outer product C_{ijkl} = A_{ij} * B_{kl}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
return np.einsum("...ij, ...kl-> ...ijkl", A, B)
cross_ijkl = _cross_ijkl
[docs]
def cross_ikjl(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{ik} * B_{jl}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
return np.einsum("...ik, ...jl-> ...ijkl", A, B)
[docs]
def cross_iklj(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{ik} * B_{lj}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
return np.einsum("...ik, ...lj-> ...ijkl", A, B)
[docs]
def cross_ilkj(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{il} * B_{kj}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
return np.einsum("...il, ...kj-> ...ijkl", A, B)
[docs]
def cross_likj(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{li} * B_{kj}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
AijBkl= cross_ijkl(A,B)
axis = list(range(AijBkl.ndim-4)) + [-3,-1,-2,-4]
AliBkj = AijBkl.transpose(axis)
return AliBkj
[docs]
def cross_iljk(A:FloatArray, B: FloatArray) -> FloatArray:
"""Compute the fourth-order tensor product with index permutation C_{ijkl} = A_{il} * B_{jk}.
Parameters
----------
A : FloatArray
First second-order tensor batch of shape `(..., N, N)`.
B : FloatArray
Second second-order tensor batch of shape `(..., N, N)`.
Returns
-------
FloatArray
Fourth-order tensor batch of shape `(..., N, N, N, N)`.
"""
AijBkl = cross_ijkl(A,B)
axis = list(range(AijBkl.ndim-4)) + [-4,-2,-1,-3]
AilBjk = AijBkl.transpose(axis)
return AilBjk
[docs]
def df_num( f:Callable[[FloatArray], FloatArray|float], A_in:FloatArray|float,
shape_iA:tuple[int,...], shape_ifA:tuple[int,...])-> FloatArray:
"""Compute numerical Jacobian derivatives of tensor functions via finite differences.
Uses `scipy.optimize.approx_fprime` iteratively over batch dimensions.
Parameters
----------
f : Callable[[FloatArray], FloatArray or float]
Tensor function to differentiate.
A_in : FloatArray or float
Evaluation point or batch of evaluation points.
shape_iA : tuple of int
Shape of a single input argument to `f`.
shape_ifA : tuple of int
Shape of a single output evaluation from `f`.
Returns
-------
FloatArray
Numerical derivative tensor of shape `shape_coll + shape_ifA + shape_iA`.
"""
A = np.atleast_1d(A_in)
dim_coll = A.ndim - len(shape_iA)
shape_coll = A.shape[:dim_coll]
shape_idfA = shape_ifA + shape_iA
df_num_res = np.zeros(shape_coll+shape_idfA)
def fflat(iAflat):
fiA = f(iAflat.reshape(shape_iA))
if isinstance(fiA, np.ndarray) :
fiA = fiA.flatten()
return fiA
for k in np.ndindex(A.shape[:dim_coll]):
Ak = A[k]
df_num_res[k] = sp.optimize.approx_fprime(Ak.flatten(), fflat).reshape(shape_idfA)
if not isinstance(A_in, np.ndarray) :
df_num_res = df_num_res.squeeze()
return df_num_res
[docs]
def check_relerr(ref:FloatArray|npfloat, val:FloatArray|npfloat, innershape,
reltol:float = 1.e-4, abstol:float = 1.e-5,
checkname:str="", verbose:bool = True )-> bool:
"""Compare two tensor arrays elementwise using absolute and relative error tolerances.
Parameters
----------
ref : FloatArray or npfloat
Reference ground-truth values.
val : FloatArray or npfloat
Computed values to validate.
innershape : tuple of int
Shape of each tensor element inside the collection.
reltol : float, optional
Maximum allowed relative error. Defaults to 1e-4.
abstol : float, optional
Maximum allowed absolute error. Defaults to 1e-5.
checkname : str, optional
Name or label describing the test. Defaults to "".
verbose : bool, optional
If True, prints detailed diagnostics upon failure. Defaults to True.
Returns
-------
bool
True if all slices pass within error tolerances, False otherwise.
"""
ref = np.asarray(ref).reshape((-1,)+ innershape)
val = np.asarray(val).reshape((-1,)+ innershape)
n = ref.shape[0]
for i in range(n):
err = np.linalg.norm((ref[i]-val[i]))
refnorm = np.linalg.norm(ref)
if err > abstol and err > refnorm*reltol :
if verbose:
print(f"check {checkname} Failed, for abstol = {abstol:.4e}, reltol {reltol:.4e}" )
print(f" index {i}, error norm : {err :.4e}" + \
f" relative error norm : {err/refnorm:.4e}")
return False
return True
[docs]
def check_df(f, df, x, ifshape, ixshape, verbose = False):
"""Verify an analytical derivative against finite-difference approximation.
Parameters
----------
f : Callable
Original function.
df : Callable
Analytical derivative function.
x : FloatArray
Evaluation inputs.
ifshape : tuple of int
Output shape of f for a single input item.
ixshape : tuple of int
Input shape for a single item.
verbose : bool, optional
If True, prints verification mismatch details. Defaults to False.
Returns
-------
bool
True if the analytical and numerical gradients agree, False otherwise.
"""
dfx_exa = df(x)
dfx_num = df_num(f, x, ixshape, ifshape)
return check_relerr(dfx_exa, dfx_num, ifshape+ixshape, verbose = verbose)
# The following block needs mygrad library ...
# def df_vector_auto(Fun, X):
# """ compute the auto derivative of fun(A) at A where A is a vector and fun(A) a vector"""
# if not HAS_MYGRAD:
# frame = inspect.currentframe()
# fun_name = frame.f_code.co_name if frame else "df_vector_auto"
# raise ImportError(f"mygrad library is needed to run {fun_name}")
# else:
#
# dim = X.shape[-1]
# dFundX_shape = X.shape + (dim,)
# dFundX = np.zeros(dFundX_shape)
# for i in range(dim):
# mgX = mgTensor(X)
# fxi = Fun(mgX)[..., i]
# fxi.backward()
# dFundX[...,i,:] = np.array(mgX.grad)
# return dFundX
# def df_tensor_auto(Fun, X, shape_fun=None):
#
# """ compute the auto derivative of fun(A) at A where A is a tensor and fun(A) a tensor"""
# if not HAS_MYGRAD:
# frame = inspect.currentframe()
# fun_name = frame.f_code.co_name if frame else "df_tensor_auto"
# raise ImportError(f"mygrad library is needed to run {fun_name}")
# xdim0 = X.shape[-2]
# xdim1 = X.shape[-1]
# if shape_fun is None:
# shape_fun=(xdim0, xdim1)
# fdim0 = shape_fun[0]
# fdim1 = shape_fun[1]
# shapedfdx = X.shape[:-2] + ( fdim0, fdim1, xdim0, xdim1)
# dFdX = np.empty(shapedfdx)
# mgX = mgTensor(X)
# for i in range(fdim0) :
# for j in range(fdim1):
# FX = Fun(mgX)[..., i, j]
# FX.backward()
# dFdX[..., i,j, :,: ] = mgX.grad
# return dFdX
[docs]
def eigval22sym(A:FloatArray):
"""Compute analytical eigenvalues for a batch of symmetric 2x2 matrices.
Parameters
----------
A : FloatArray
Collection of symmetric 2x2 matrices of shape `(..., 2, 2)`.
Returns
-------
FloatArray
Eigenvalues array of shape `(..., 2)` with entries `[..., 0]` (max) and `[..., 1]` (min).
Raises
------
ValueError
If trailing dimensions are not (2, 2) or matrices have negative discriminant.
"""
if A.shape[-2:] != (2,2) :
raise ValueError('A must be of shape[-2:] == (2,2)')
a = A[..., 0, 0]
b = A[..., 0, 1]
c = A[..., 1, 1]
disc = (a-c)**2 + 4*b**2
if np.any(disc < 0):
raise ValueError('all 2*2 matrix in a must be positive definite')
eigsvals = np.empty(A.shape[:-2]+(2,))
disc = np.sqrt(disc)
eigsvals[..., 0] = (a+c + disc )/2.
eigsvals[..., 1] = (a+c - disc )/2.
return eigsvals
[docs]
def max_eigval22sym(A:FloatArray, outlmax : Optional[FloatArray] = None):
"""Compute the maximum eigenvalue for a batch of symmetric 2x2 matrices.
Parameters
----------
A : FloatArray
Batch of symmetric 2x2 matrices of shape `(..., 2, 2)`.
outlmax : FloatArray, optional
Output buffer receiving eigenvalues. Defaults to None.
Returns
-------
FloatArray
Maximum eigenvalues array of shape `(...)`.
Raises
------
ValueError
If trailing dimensions are not (2, 2) or matrices have negative discriminant.
"""
if A.shape[-2:] != (2,2) :
raise ValueError('A must be of shape[-2:] == (2,2)')
if outlmax is None:
outlmax = np.empty(A.shape[:-2])
a = A[..., 0, 0]
b = A[..., 0, 1]
c = A[..., 1, 1]
disc = (a-c)**2 + 4*b**2
if np.any(disc < 0):
raise ValueError('all 2*2 matrix in a must be positive definite')
disc = np.sqrt(disc)
outlmax[...] = (a+c + disc )/2.
return outlmax
[docs]
@nb.njit(cache=True, fastmath = FASTMATH)
def max_eigval22sym_kernel_njit(A:FloatArray):
"""Numba scalar kernel computing the maximum eigenvalue of a single symmetric 2x2 matrix.
Parameters
----------
A : FloatArray
2x2 symmetric matrix slice.
Returns
-------
float
Maximum eigenvalue.
Raises
------
ValueError
If the discriminant is negative.
"""
a = A[0, 0]
b = A[0, 1]
c = A[1, 1]
disc = (a-c)**2 + 4*b**2
if disc < 0:
raise ValueError('2*2 matrix in a must be positive definite')
return (a+c + np.sqrt(disc ))/2.