Package: tensortrax
Math on (Hyper-Dual) Tensors with Trailing Axes.HighlightsDesigned to operate on input arrays with (elementwise-operating) trailing axesEssential vector/tensor Hyper-Dual number math, including limited support foreinsum(restricted to max. three operands)Math is limited but similar to NumPy, try to useimport tensortrax.math as tminstead ofimport numpy as npinside functions to be differentiatedForward Mode Automatic Differentiation (AD) using Hyper-Dual Tensors, up to second order derivativesCreate functions in terms of Hyper-Dual TensorsEvaluate the function, the gradient (jacobian) and the hessian of scalar-valued functions or functionals on given input arraysStraight-forward definition of custom functions in variational-calculus notationStable gradient and hessian of eigenvalues obtained fromeigvalshin case of repeated equal eigenvaluesPlease keep in mind thattensortraxis not imitating a 100% full-featured NumPy, e.g. likeAutograd[1]. No arbitrary-order gradients or gradients-of-gradients are supported. The capability is limited to first- and second order gradients of a given function. Also,tensortraxprovides no support fordtype=complex.MotivationGradient and hessian evaluations of functions or functionals based on tensor-valued input arguments are a fundamental repetitive and (error-prone) task in constitutive hyperelastic material formulations used in continuum mechanics of solid bodies. In the worst case, conceptual ideas are impossible to pursue because the required tensorial derivatives are not readily achievable. The Hyper-Dual number approach enables a generalized and systematic way to overcome this deficiency[2]. Compared to existing Hyper-Dual Number libaries ([3],[4]) which introduce a new (hyper-) dualdtype(treated asdtype=objectin NumPy),tensortraxrelies on its ownTensorclass. This approach involves a re-definition of all essential math operations (and NumPy-functions), whereas thedtype-approach supports most basic math operations out of the box. However, intensortraxNumPy and all its underlying linear algebra functions operate on default data types (e.g.dtype=float). This allows to support functions likenp.einsum(). Beside the differences concerning the underlyingdtype,tensortraxis formulated on easy-to-understand (tensorial) calculus of variation. Hence, gradient- and hessian-vector products are evaluated with very little overhead compared to analytic formulations.UsageLet's define a scalar-valued function which operates on a tensor. The math moduletensortrax.mathprovides some essential NumPy-like functions including linear algebra.importtensortraxastrimporttensortrax.mathastmdeffun(F,mu=1):C=F.T@FI1=tm.trace(C)J=tm.linalg.det(F)returnmu/2*(J**(-2/3)*I1-3)The Hessian of the scalar-valued function w.r.t. the chosen function argument (here,wrt=0orwrt="F") is evaluated by variational calculus (Forward Mode AD implemented as Hyper-Dual Tensors). The function is called once for each component of the hessian (symmetry is taken care of). The function and the gradient are evaluated with no additional computational cost. Optionally, the function, gradient and Hessian calls are executed in parallel (threaded).importnumpyasnp# some random input datanp.random.seed(125161)F=(np.eye(3)+np.random.rand(50,8,3,3)/10).T# W = tr.function(fun, wrt=0, ntrax=2)(F)# dWdF = tr.gradient(fun, wrt=0, ntrax=2)(F)# d2WdF2, dWdF, W = tr.hessian(fun, wrt="F", ntrax=2, full_output=True)(F=F)d2WdF2=tr.hessian(fun,wrt="F",ntrax=2,parallel=False)(F=F)Another possibility is to define and operate on Tensors manually. This enables more flexible coding, which wouldn't be possible with the builtin functions. The Hu-Washizu Three-Field-Variational principle for nearly incompressible hyperelastic solids[5]is used here to obtain mixed partial derivatives. Some random input arrays are generated and a Tensor is created for each variable. After performing some math, the hessian of the resulting tensor object is extracted.# some random input datan=10x=(np.eye(3)+np.random.rand(n,3,3)/10).Ty=np.random.rand(n)z=np.random.rand(n)/10+1# create tensorsF=tr.Tensor(x,ntrax=1)p=tr.Tensor(y,ntrax=1)J=tr.Tensor(z,ntrax=1)defneo_hooke(F,mu=1):"Strain energy function of the Neo-Hookean material formulation."C=F.T@FI3=tm.linalg.det(C)returnmu*(I3**(-1/3)*tm.trace(C)-3)/2defvolumetric(J,bulk=20):"Volumetric strain energy function."returnbulk*(J-1)**2/2defW(F,p,J):"Hu-Washizu (u, p, J) - Three-Field-Variation."detF=tm.linalg.det(F)returnneo_hooke(F)+volumetric(J)+p*(detF-J)# init Tensors to be used with second partial derivativesF.init(hessian=True,δx=False,Δx=False)p.init(hessian=True,δx=True,Δx=False)J.init(hessian=True,δx=False,Δx=True)# evaluate a mixed second partial derivativedWdpdJ=tr.Δδ(W(F,p,J))In a similar way, the gradient may be obtained by initiating a Tensor with the gradient argument.# init Tensors to be used with first partial derivativesF.init(gradient=True,δx=False)p.init(gradient=True,δx=True)J.init(gradient=True,δx=False)# evaluate a partial derivativedWdp=tr.δ(W(F,p,J))PerformanceAbenchmarkfor the gradient and hessian runtimes of an isotropic hyperelastic strain energy function demonstrates the performance of this package compared toAutograd[1]. The hessian is evaluated in about 2.5 seconds for one million input tensors (Intel Core i7-11850H, 32GB RAM). While the runtimes of the gradients evaluated by Tensortrax are similar compared to those of Autograd, the Hessian evaluation in Tensortrax ismuch fasterthan in Autograd (seeTable 1,Table 2andFigure 1). Note that this is only one exemplaric benchmark for a specific kind of function - runtimes may vary.\psi(\boldsymbol{C}) = tr(\boldsymbol{C}) - \ln(\det(\boldsymbol{C}))Table 1:Runtimes of Gradients evaluated bytensortraxandautograd.TensorsGradient (Tensortrax) in sGradient (Autograd) in sSpeedup10.000710.00082x  1.1640.000750.00070x  0.93160.000750.00071x  0.94640.000830.00077x  0.942560.000970.00094x  0.9710240.001600.00161x  1.0140960.005680.00428x  0.75163840.012330.01690x  1.37655360.061030.06756x  1.112621440.279100.26911x  0.9610485761.155611.09512x  0.95Table 2:Runtimes of Hessians evaluated bytensortraxandautograd.TensorsHessian (Tensortrax) in sHessian (Autograd) in sSpeedup10.000680.01348x 19.7740.000730.00723x  9.87160.000760.00718x  9.41640.000830.00792x  9.582560.001110.01044x  9.4310240.002140.02094x  9.7940960.008440.06148x  7.28163840.025440.23778x  9.35655360.116850.95208x  8.152621440.515613.94024x  7.6410485762.1362916.15154x  7.56Figure 1:Runtime vs. Number of input tensors - plot for Gradients and Hessians evaluated bytensortraxandautograd.TheoryThe calculus of variation deals with variations, i.e. small changes in functions and functionals. A small-change in a function is evaluated by applying small changes on the tensor components.\psi = \psi(\boldsymbol{F})\delta \psi = \delta \psi(\boldsymbol{F}, \delta \boldsymbol{F})Let's take the trace of a tensor product as an example. The variation is evaluated as follows:\psi = tr(\boldsymbol{F}^T \boldsymbol{F}) = \boldsymbol{F} : \boldsymbol{F}\delta \psi = \delta \boldsymbol{F} : \boldsymbol{F} + \boldsymbol{F} : \delta \boldsymbol{F} = 2 \ \boldsymbol{F} : \delta \boldsymbol{F}The $P_{ij}$ - component of the jacobian $\boldsymbol{P}$ is now numerically evaluated by setting the respective variational component $\delta F_{ij}$ of the tensor to one and all other components to zero. In total, $i \cdot j$ function calls are necessary to assemble the full jacobian. For example, the $12$ - component is evaluated as follows:\delta \boldsymbol{F}_{(12)} = \begin{bmatrix} 0 & 1 & 0 \\ 0 & 0 & 0 \\ 0 & 0 & 0 \end{bmatrix}\delta_{(12)} \psi = \frac{\partial \psi}{\partial F_{12}} = 2 \ \boldsymbol{F} : \delta \boldsymbol{F}_{(12)} = 2 \ \boldsymbol{F} : \begin{bmatrix} 0 & 1 & 0 \\ 0 & 0 & 0 \\ 0 & 0 & 0 \end{bmatrix}The second order variation, i.e. a variation applied on another variation of a function is evaluated in the same way as a first order variation.\Delta \delta \psi = 2 \ \delta \boldsymbol{F} : \Delta \boldsymbol{F} + 2 \ \boldsymbol{F} : \Delta \delta \boldsymbol{F}Once again, each component $A_{ijkl}$ of the fourth-order hessian is numerically evaluated. In total, $i \cdot j \cdot k \cdot l$ function calls are necessary to assemble the full hessian (without considering symmetry). For example, the $1223$ - component is evaluated by setting $\Delta \delta \boldsymbol{F} = \boldsymbol{0}$ and $\delta \boldsymbol{F}$ and $\Delta \boldsymbol{F}$ as follows:\delta \boldsymbol{F}_{(12)} = \begin{bmatrix} 0 & 1 & 0 \\ 0 & 0 & 0 \\ 0 & 0 & 0 \end{bmatrix}\Delta \boldsymbol{F}_{(23)} = \begin{bmatrix} 0 & 0 & 0 \\ 0 & 0 & 1 \\ 0 & 0 & 0 \end{bmatrix}\Delta \delta \boldsymbol{F} = \begin{bmatrix} 0 & 0 & 0 \\ 0 & 0 & 0 \\ 0 & 0 & 0 \end{bmatrix}\Delta_{(23)} \delta_{(12)} \psi = \Delta_{(12)} \delta_{(23)} \psi = \frac{\partial^2 \psi}{\partial F_{12}\ \partial F_{23}}\Delta_{(23)} \delta_{(12)} \psi = 2 \ \delta \boldsymbol{F}_{(12)} : \Delta \boldsymbol{F}_{(23)} + 2 \ \boldsymbol{F} : \Delta \delta \boldsymbol{F}Numeric calculus of variation intensortraxEach Tensor has four attributes: the (real) tensor array and the (hyper-dual) variational arrays. To obtain the above mentioned $12$ - component of the gradient and the $1223$ - component of the hessian, a tensor has to be created with the appropriate small-changes of the tensor components (dual arrays).importtensortraxastrfromtensortraximportTensor,f,δ,Δ,Δδfromtensortrax.mathimporttraceδF_12=np.array([[0,1,0],[0,0,0],[0,0,0],],dtype=float,)ΔF_23=np.array([[0,0,0],[0,0,1],[0,0,0],],dtype=float,)x=np.eye(3)+np.arange(9).reshape(3,3)/10F=Tensor(x=x,δx=δF_12,Δx=ΔF_23,Δδx=None)I1_C=trace(F.T@F)The function as well as the gradient and hessian components are accessible as:ψ=f(I1_C)P_12=δ(I1_C)# (= Δ(I1_C))A_1223=Δδ(I1_C)To obtain full gradients and hessians of scalar-valued functions in one function call,tensortraxprovides helpers (decorators) which handle the multiple function calls.fun=lambdaF:trace(F.T@F)func=tr.function(fun)(x)grad=tr.gradient(fun)(x)hess=tr.hessian(fun)(x)For tensor-valued functions, usejacobian()instead ofgradient().fun=lambdaF:F.T@Fjac=tr.jacobian(fun)(x)Evaluate the gradient- as well as the hessian-vector(s)-product:gvp=tr.gradient_vector_product(fun)(x,δx=x)hvp=tr.hessian_vector_product(fun)(x,δx=x)hvsp=tr.hessian_vectors_product(fun)(x,δx=x,Δx=x)ExtensionsCustom functions (extensions) are easy to implement intensortrax. Beside the function expression, three additional (dual) variation expressions have to be defined.importnumpyasnpfromtensortraximportTensor,f,δ,Δ,Δδdefsin(A):returnTensor(x=np.sin(f(A)),δx=np.cos(f(A))*δ(A),Δx=np.cos(f(A))*Δ(A),Δδx=-np.sin(f(A))*δ(A)*Δ(A)+np.cos(f(A))*Δδ(A),ntrax=A.ntrax,)x=np.eye(3)y=sin(Tensor(x))Hint:Feel free tocontributemissing math-functions tosrc/tensortrax/math/_math_tensor.py:page_with_curl: :pencil2:.ReferencesD. Maclaurin, D. Duvenaud, M. Johnson and J. Townsend,Autograd. [Online]. Available:https://github.com/HIPS/autograd.J. Fike and J. Alonso,The Development of Hyper-Dual Numbers for Exact Second-Derivative Calculations, 49th AIAA Aerospace Sciences Meeting including the New Horizons Forum and Aerospace Exposition. American Institute of Aeronautics and Astronautics, Jan. 04, 2011,.P. Rehner and G. Bauer,Application of Generalized (Hyper-) Dual Numbers in Equation of State Modeling, Frontiers in Chemical Engineering, vol. 3, 2021. Available:https://github.com/itt-ustutt/num-dual.T. Oberbichler,HyperJet. [Online]. Available:http://github.com/oberbichler/HyperJet.J. Bonet and R. D. Wood,Nonlinear Continuum Mechanics for Finite Element Analysis, 2nd ed. Cambridge: Cambridge University Press, 2008,.