From 2d5f6c3913bc08ceeaa78f3517f47d70b9ebe348 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 16:58:26 +0100 Subject: [PATCH 01/12] add eigh which supports only first derivatives --- tensortrax/math/_linalg/__init__.py | 2 +- tensortrax/math/_linalg/_linalg_tensor.py | 50 +++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/tensortrax/math/_linalg/__init__.py b/tensortrax/math/_linalg/__init__.py index 92d414c..78bded3 100644 --- a/tensortrax/math/_linalg/__init__.py +++ b/tensortrax/math/_linalg/__init__.py @@ -9,4 +9,4 @@ """ from ._linalg_array import det as _det, inv as _inv -from ._linalg_tensor import det, inv, eigvalsh +from ._linalg_tensor import det, inv, eigvalsh, eigh diff --git a/tensortrax/math/_linalg/_linalg_tensor.py b/tensortrax/math/_linalg/_linalg_tensor.py index 15e89fe..d92f679 100644 --- a/tensortrax/math/_linalg/_linalg_tensor.py +++ b/tensortrax/math/_linalg/_linalg_tensor.py @@ -104,3 +104,53 @@ def eigvalsh(A): Δδx=Δδλ, ntrax=A.ntrax, ) + + +def eigh(A): + "Eigenvalues and -bases of a symmetric Tensor (only first derivative)." + + λ, N = [x.T for x in np.linalg.eigh(f(A).T)] + N = transpose(N) + M = einsum("ai...,aj...->aij...", N, N) + + δλ = einsum("aij...,ij...->a...", M, δ(A)) + Δλ = einsum("aij...,ij...->a...", M, Δ(A)) + + Γ = [(1, 2), (2, 0), (0, 1)] + + δN = [] + for α in range(3): + δNα = [] + for γ in Γ[α]: + Mαγ = einsum("i...,j...->ij...", N[α], N[γ]) + δAαγ = einsum("ij...,ij...->...", Mαγ, δ(A)) + λαγ = λ[α] - λ[γ] + λ_equal = np.isclose(λ[α], λ[γ]) + if np.any(λ_equal): + if len(λαγ.shape) == 0: + λαγ = np.inf + else: + λαγ[λ_equal] = np.inf + δNα.append(1 / λαγ * N[γ] * δAαγ) + δN.append(sum(δNα, axis=0)) + + δM = einsum("ai...,aj...->aij...", δN, N) + einsum("ai...,aj...->aij...", N, δN) + Δδλ = einsum("aij...,ij...->a...", δM, Δ(A)) + einsum( + "aij...,ij...->a...", M, Δδ(A) + ) + + return ( + Tensor( + x=λ, + δx=δλ, + Δx=Δλ, + Δδx=Δδλ, + ntrax=A.ntrax, + ), Tensor( + x=M, + δx=δM, + Δx=δM * np.nan, + Δδx=δM * np.nan, + ntrax=A.ntrax, + ) + ) \ No newline at end of file From ee1b801ea686dfbc0c76011f3208e7080c8d6ab1 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 16:59:23 +0100 Subject: [PATCH 02/12] lint black --- tensortrax/math/_linalg/_linalg_tensor.py | 7 ++++--- tests/test_math.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tensortrax/math/_linalg/_linalg_tensor.py b/tensortrax/math/_linalg/_linalg_tensor.py index d92f679..4447a60 100644 --- a/tensortrax/math/_linalg/_linalg_tensor.py +++ b/tensortrax/math/_linalg/_linalg_tensor.py @@ -146,11 +146,12 @@ def eigh(A): Δx=Δλ, Δδx=Δδλ, ntrax=A.ntrax, - ), Tensor( + ), + Tensor( x=M, δx=δM, Δx=δM * np.nan, Δδx=δM * np.nan, ntrax=A.ntrax, - ) - ) \ No newline at end of file + ), + ) diff --git a/tests/test_math.py b/tests/test_math.py index a4b9ab2..a48b4de 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -104,7 +104,7 @@ def test_math(): tm.reshape(t, (9,)) tm.reshape(t, (3, 3)) - + tm.reshape(x, (3, 3, 100)) From 6c4f08760f77f4fecaaecb69c6c061365194e465 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 17:05:42 +0100 Subject: [PATCH 03/12] test expm, eigh --- tensortrax/math/_linalg/__init__.py | 2 +- tensortrax/math/_linalg/_linalg_tensor.py | 7 ++++++- tests/test_math.py | 6 +++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tensortrax/math/_linalg/__init__.py b/tensortrax/math/_linalg/__init__.py index 78bded3..20b316e 100644 --- a/tensortrax/math/_linalg/__init__.py +++ b/tensortrax/math/_linalg/__init__.py @@ -9,4 +9,4 @@ """ from ._linalg_array import det as _det, inv as _inv -from ._linalg_tensor import det, inv, eigvalsh, eigh +from ._linalg_tensor import det, inv, eigvalsh, eigh, expm diff --git a/tensortrax/math/_linalg/_linalg_tensor.py b/tensortrax/math/_linalg/_linalg_tensor.py index 4447a60..ee1edaa 100644 --- a/tensortrax/math/_linalg/_linalg_tensor.py +++ b/tensortrax/math/_linalg/_linalg_tensor.py @@ -12,7 +12,7 @@ import numpy as np from ..._tensor import Tensor, einsum, matmul, f, δ, Δ, Δδ -from .._math_tensor import transpose, ddot, sum +from .._math_tensor import transpose, ddot, sum, exp from . import _linalg_array as array @@ -155,3 +155,8 @@ def eigh(A): ntrax=A.ntrax, ), ) + + +def expm(A): + "Compute the matrix exponential of a symmetric array." + return einsum("a...,aij...->ij...", *eigh(A)) diff --git a/tests/test_math.py b/tests/test_math.py index a48b4de..df8a0f2 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -63,10 +63,14 @@ def test_math(): ]: assert np.allclose(fun(F), fun(T).x) - for fun in [tm.linalg.det]: + for fun in [tm.linalg.det, tm.linalg.inv]: assert np.allclose(fun(F), fun(T).x) assert tm.linalg.eigvalsh(T).shape == (3,) + assert tm.linalg.eigh(T)[0].shape == (3,) + assert tm.linalg.eigh(T)[1].shape == (3, 3, 3) + + assert tm.linalg.expm(T).shape == (3, 3) assert tm.array.cross(F, F).shape == F.shape assert tm.array.eye(F).shape == F.shape From 181306c995cc5c97147ce82060d1d794d4f18a61 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 19:40:26 +0100 Subject: [PATCH 04/12] Update _linalg_tensor.py --- tensortrax/math/_linalg/_linalg_tensor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensortrax/math/_linalg/_linalg_tensor.py b/tensortrax/math/_linalg/_linalg_tensor.py index ee1edaa..7521be8 100644 --- a/tensortrax/math/_linalg/_linalg_tensor.py +++ b/tensortrax/math/_linalg/_linalg_tensor.py @@ -159,4 +159,5 @@ def eigh(A): def expm(A): "Compute the matrix exponential of a symmetric array." - return einsum("a...,aij...->ij...", *eigh(A)) + λ, M = eigh(A) + return einsum("a...,aij...->ij...", exp(A), M) From 5c7c0d272ac58cb54da946fb3bb518b257aaea20 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 20:13:57 +0100 Subject: [PATCH 05/12] fix expm --- tensortrax/math/_linalg/_linalg_tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensortrax/math/_linalg/_linalg_tensor.py b/tensortrax/math/_linalg/_linalg_tensor.py index 7521be8..68b7549 100644 --- a/tensortrax/math/_linalg/_linalg_tensor.py +++ b/tensortrax/math/_linalg/_linalg_tensor.py @@ -160,4 +160,4 @@ def eigh(A): def expm(A): "Compute the matrix exponential of a symmetric array." λ, M = eigh(A) - return einsum("a...,aij...->ij...", exp(A), M) + return einsum("a...,aij...->ij...", exp(λ), M) From ccc588205dc4c5c1fed7678ba5b46dab2cf6f5a3 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 20:19:40 +0100 Subject: [PATCH 06/12] Update test_math.py --- tests/test_math.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_math.py b/tests/test_math.py index df8a0f2..737a4ef 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -67,6 +67,9 @@ def test_math(): assert np.allclose(fun(F), fun(T).x) assert tm.linalg.eigvalsh(T).shape == (3,) + + F = np.diag([1.2, 1.2, 2.]) + assert tm.linalg.eigh(T)[0].shape == (3,) assert tm.linalg.eigh(T)[1].shape == (3, 3, 3) From 0fdbb507ba64dcf4cad0df5ca5969be3e6370420 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 20:21:41 +0100 Subject: [PATCH 07/12] Update test_math.py --- tests/test_math.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/test_math.py b/tests/test_math.py index 737a4ef..00f5d31 100644 --- a/tests/test_math.py +++ b/tests/test_math.py @@ -67,11 +67,6 @@ def test_math(): assert np.allclose(fun(F), fun(T).x) assert tm.linalg.eigvalsh(T).shape == (3,) - - F = np.diag([1.2, 1.2, 2.]) - - assert tm.linalg.eigh(T)[0].shape == (3,) - assert tm.linalg.eigh(T)[1].shape == (3, 3, 3) assert tm.linalg.expm(T).shape == (3, 3) @@ -114,6 +109,18 @@ def test_math(): tm.reshape(x, (3, 3, 100)) + F = np.diag([1.2, 1.2, 2.0]) + T = tr.Tensor(F) + + assert tm.linalg.eigh(T)[0].shape == (3,) + assert tm.linalg.eigh(T)[1].shape == (3, 3, 3) + + F = np.tile(F.reshape(3, 3, 1), 5) + T = tr.Tensor(F, ntrax=1) + + assert tm.linalg.eigh(T)[0].shape == (3,) + assert tm.linalg.eigh(T)[1].shape == (3, 3, 3) + if __name__ == "__main__": test_math() From 7cf6dfd28a546deb9d1852bfb4b3f8245f4344ce Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 21:59:02 +0100 Subject: [PATCH 08/12] add jacobian --- tensortrax/_evaluate.py | 47 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tensortrax/_evaluate.py b/tensortrax/_evaluate.py index fdb7387..778bdd4 100644 --- a/tensortrax/_evaluate.py +++ b/tensortrax/_evaluate.py @@ -53,6 +53,53 @@ def evaluate_function(*args, **kwargs): return evaluate_function +def jacobian(fun, wrt=0, ntrax=0, parallel=False, full_output=False): + "Evaluate the jacobian of a function." + + def evaluate_jacobian(*args, **kwargs): + + x = arg_to_tensor(args, kwargs, wrt) + t = Tensor(x, ntrax=ntrax) + δx = Δx = np.eye(t.size) + indices = range(t.size) + + args0, kwargs0 = add_tensor(args, kwargs, wrt, δx[0], Δx[0], ntrax) + shape = fun(*args0, **kwargs0).shape + axes = tuple([slice(None)] * len(shape)) + + fx = np.zeros((*shape, *t.trax)) + dfdx = np.zeros((*shape, t.size, *t.trax)) + + def kernel(a, wrt, δx, Δx, ntrax, args, kwargs): + args, kwargs = add_tensor(args, kwargs, wrt, δx[a], Δx[a], ntrax) + func = fun(*args, **kwargs) + fx[axes] = f(func) + dfdx[(*axes, a)] = δ(func) + + if not parallel: + for a in indices: + kernel(a, wrt, δx, Δx, ntrax, args, kwargs) + + else: + threads = [ + Thread(target=kernel, args=(a, wrt, δx, Δx, ntrax, args, kwargs)) + for a in indices + ] + + for th in threads: + th.start() + + for th in threads: + th.join() + + if full_output: + return np.array(dfdx).reshape(*shape, *t.shape, *t.trax), fx[0] + else: + return np.array(dfdx).reshape(*shape, *t.shape, *t.trax) + + return evaluate_jacobian + + def gradient(fun, wrt=0, ntrax=0, parallel=False, full_output=False): "Evaluate the gradient of a scalar-valued function." From 0ec1b0c5522125c59ed8bf0e7fcf461db3830173 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 21:59:18 +0100 Subject: [PATCH 09/12] Update __init__.py --- tensortrax/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensortrax/__init__.py b/tensortrax/__init__.py index ba589bf..e528aea 100644 --- a/tensortrax/__init__.py +++ b/tensortrax/__init__.py @@ -6,6 +6,7 @@ from ._helpers import f, δ, Δ, Δδ from ._evaluate import ( function, + jacobian, gradient, hessian, gradient_vector_product, From 120c78c3788f6eddfb591dcbedb41c8c2d58c364 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 22:06:58 +0100 Subject: [PATCH 10/12] test jacobian --- tensortrax/_evaluate.py | 4 ++-- tests/test_jacobian.py | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) create mode 100644 tests/test_jacobian.py diff --git a/tensortrax/_evaluate.py b/tensortrax/_evaluate.py index 778bdd4..efcd0f7 100644 --- a/tensortrax/_evaluate.py +++ b/tensortrax/_evaluate.py @@ -63,7 +63,7 @@ def evaluate_jacobian(*args, **kwargs): δx = Δx = np.eye(t.size) indices = range(t.size) - args0, kwargs0 = add_tensor(args, kwargs, wrt, δx[0], Δx[0], ntrax) + args0, kwargs0 = add_tensor(args, kwargs, wrt, None, None, ntrax) shape = fun(*args0, **kwargs0).shape axes = tuple([slice(None)] * len(shape)) @@ -93,7 +93,7 @@ def kernel(a, wrt, δx, Δx, ntrax, args, kwargs): th.join() if full_output: - return np.array(dfdx).reshape(*shape, *t.shape, *t.trax), fx[0] + return np.array(dfdx).reshape(*shape, *t.shape, *t.trax), fx else: return np.array(dfdx).reshape(*shape, *t.shape, *t.trax) diff --git a/tests/test_jacobian.py b/tests/test_jacobian.py new file mode 100644 index 0000000..5db5592 --- /dev/null +++ b/tests/test_jacobian.py @@ -0,0 +1,27 @@ +import tensortrax as tr +import tensortrax.math as tm +import numpy as np + + +def right_cauchy_green(F): + return F.T() @ F + + +def test_jacobian(): + + F = (np.eye(3).ravel() + np.arange(9) / 10).reshape(3, 3, 1, 1) + + for parallel in [False, True]: + for fun in [right_cauchy_green]: + c = tr.function(fun, ntrax=2, parallel=parallel)(F) + dCdF, C = tr.jacobian(fun, ntrax=2, parallel=parallel, full_output=True)(F) + + assert c.shape == (3, 3, 1, 1) + assert C.shape == (3, 3, 1, 1) + assert dCdF.shape == (3, 3, 3, 3, 1, 1) + + assert np.allclose(C, c) + + +if __name__ == "__main__": + test_jacobian() From 6a98050a0beef0a66dde927fa7b7e1b649d0ecb7 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 22:09:01 +0100 Subject: [PATCH 11/12] bump version --- README.md | 11 ++++++++++- pyproject.toml | 2 +- setup.cfg | 2 +- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 571d23b..3226804 100644 --- a/README.md +++ b/README.md @@ -154,7 +154,7 @@ P_12 = δ(I1_C) # (= Δ(I1_C)) A_1223 = Δδ(I1_C) ``` -To obtain full gradients and hessians in one function call, `tensortrax` provides helpers (decorators) which handle the multiple function calls. +To obtain full gradients and hessians of scalar-valued functions in one function call, `tensortrax` provides helpers (decorators) which handle the multiple function calls. ```python fun = lambda F: trace(F.T() @ F) @@ -164,6 +164,15 @@ grad = tr.gradient(fun)(x) hess = tr.hessian(fun)(x) ``` +For tensor-valued functions, use `jacobian()` instead of `gradient()`. + +```python +fun = lambda F: F.T() @ F + +func = tr.jacobian(fun)(x) +``` + + Evaluate the gradient- as well as the hessian-vector-product: ```python diff --git a/pyproject.toml b/pyproject.toml index c82d7a8..455fdf1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "tensortrax" -version = "0.2.2" +version = "0.2.3" description = "Math on (Hyper-Dual) Tensors with Trailing Axes" readme = "README.md" requires-python = ">=3.7" diff --git a/setup.cfg b/setup.cfg index 155c770..d40b509 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = tensortrax -version = 0.2.2 +version = 0.2.3 author = Andreas Dutzler author_email = a.dutzler@gmail.com description = Math on (Hyper-Dual) Tensors with Trailing Axes From cbc0ecef2a4fbb99e9c47695181d76f6ee538775 Mon Sep 17 00:00:00 2001 From: Andreas Dutzler Date: Mon, 12 Dec 2022 22:10:58 +0100 Subject: [PATCH 12/12] Update test_jacobian.py --- tests/test_jacobian.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_jacobian.py b/tests/test_jacobian.py index 5db5592..2f169bc 100644 --- a/tests/test_jacobian.py +++ b/tests/test_jacobian.py @@ -14,6 +14,7 @@ def test_jacobian(): for parallel in [False, True]: for fun in [right_cauchy_green]: c = tr.function(fun, ntrax=2, parallel=parallel)(F) + dCdF = tr.jacobian(fun, ntrax=2, parallel=parallel)(F) dCdF, C = tr.jacobian(fun, ntrax=2, parallel=parallel, full_output=True)(F) assert c.shape == (3, 3, 1, 1)