Skip to content

Commit

Permalink
Merge pull request #58 from adtzlr/remove-for-loops
Browse files Browse the repository at this point in the history 
Remove For-loops
  • Loading branch information
@adtzlr
adtzlr authored Feb 4, 2023
2 parents 57dc11a + de8e191 commit a011ee6
Show file tree
Hide file tree
Showing 6 changed files with 182 additions and 195 deletions.
30 changes: 15 additions & 15 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def fun(F, mu=1):
return mu / 2 * (J ** (-2 / 3) * I1 - 3)
```

The hessian of the scalar-valued function w.r.t. the chosen function argument (here, `wrt=0` or `wrt="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 calls are executed in parallel (threaded).
The hessian of the scalar-valued function w.r.t. the chosen function argument (here, `wrt=0` or `wrt="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 calls are executed in parallel (threaded)~~ (currently disabled, will be re-enabled soon).

```python
import numpy as np
Expand All @@ -64,27 +64,27 @@ d2WdF2 = tr.hessian(fun, wrt="F", ntrax=2, parallel=False)(F=F)
```

# Performance
A [benchmark](https://github.com/adtzlr/tensortrax/blob/main/docs/benchmark/benchmark.py) for the gradient and hessian runtimes of an isotropic hyperelastic strain energy function demonstrates the performance of this package. The hessian is evaluated in about five seconds for one million input tensors (Intel Core i7-11850H, 32GB RAM).
A [benchmark](https://github.com/adtzlr/tensortrax/blob/main/docs/benchmark/benchmark.py) for the gradient and hessian runtimes of an isotropic hyperelastic strain energy function demonstrates the performance of this package. The hessian is evaluated in about two seconds for one million input tensors (Intel Core i7-11850H, 32GB RAM).

```math
\psi(\boldsymbol{C}) = tr(\boldsymbol{C}) - \ln(\det(\boldsymbol{C}))
```

| Tensors | Gradient in s | Hessian in s |
| ------- | ------------- | ------------ |
| 2 | 0.00552 | 0.01474 |
| 8 | 0.00429 | 0.01420 |
| 32 | 0.00415 | 0.01364 |
| 128 | 0.00418 | 0.01453 |
| 512 | 0.00697 | 0.02465 |
| 2048 | 0.00831 | 0.03134 |
| 8192 | 0.01289 | 0.05174 |
| 32768 | 0.02837 | 0.11737 |
| 131072 | 0.16327 | 0.58191 |
| 524288 | 0.86078 | 2.65141 |
| 2097152 | 2.97900 | 10.95087 |

![benchmark](https://user-images.githubusercontent.com/5793153/214539409-63d9418e-9cb6-4e38-9a86-572665da30fe.svg)
| 2 | 0.00077 | 0.00058 |
| 8 | 0.00070 | 0.00057 |
| 32 | 0.00053 | 0.00063 |
| 128 | 0.00061 | 0.00068 |
| 512 | 0.00095 | 0.00126 |
| 2048 | 0.00261 | 0.00338 |
| 8192 | 0.00604 | 0.01298 |
| 32768 | 0.02806 | 0.05353 |
| 131072 | 0.13473 | 0.25056 |
| 524288 | 0.56922 | 1.03252 |
| 2097152 | 2.42609 | 4.59884 |

![benchmark](https://user-images.githubusercontent.com/5793153/216739312-c0a199fb-c905-43fc-9f2a-5550ce045b32.svg)

# Theory
The 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.
Expand Down
2 changes: 1 addition & 1 deletion src/tensortrax/__about__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@
tensorTRAX: Math on (Hyper-Dual) Tensors with Trailing Axes.
"""

__version__ = "0.5.1"
__version__ = "0.6.0"
253 changes: 77 additions & 176 deletions src/tensortrax/_evaluate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,250 +28,151 @@ def evaluate(*args, **kwargs):
return evaluate


def add_tensor(args, kwargs, wrt, δx, Δx, ntrax, sym):
def add_tensor(
args, kwargs, wrt, ntrax, sym=False, gradient=False, hessian=False, δx=None, Δx=None
):
"Modify the arguments and replace the w.r.t.-argument by a tensor."

kwargs_out = copy(kwargs)
args_out = list(args)

if isinstance(wrt, str):
if sym:
kwargs_out[wrt] = from_triu_1d(
Tensor(x=triu_1d(kwargs[wrt]), δx=δx, Δx=Δx, ntrax=ntrax)
)
else:
kwargs_out[wrt] = Tensor(x=kwargs[wrt], δx=δx, Δx=Δx, ntrax=ntrax)

args_old = kwargs
args_new = kwargs_out
elif isinstance(wrt, int):
if sym:
args_out[wrt] = from_triu_1d(
Tensor(x=triu_1d(args[wrt]), δx=δx, Δx=Δx, ntrax=ntrax)
)
else:
args_out[wrt] = Tensor(x=args[wrt], δx=δx, Δx=Δx, ntrax=ntrax)
args_old = args
args_new = args_out
else:
raise TypeError(
f"Type of wrt not supported. type(wrt) is {type(wrt)} (must be str or int)."
)

return args_out, kwargs_out
x = args_old[wrt]

if sym:
x = x = triu_1d(x)

def arg_to_tensor(args, kwargs, wrt, sym):
"Return the argument which will be replaced by a tensor."
tensor = Tensor(x=x, ntrax=ntrax)
trax = tensor.trax

if isinstance(wrt, str):
x = kwargs[wrt] if not sym else triu_1d(kwargs[wrt])
elif isinstance(wrt, int):
x = args[wrt] if not sym else triu_1d(args[wrt])
else:
raise TypeError(f"w.r.t. {wrt} not supported.")
tensor.init(gradient=gradient, hessian=hessian, sym=sym, δx=δx, Δx=Δx)

if sym:
tensor = from_triu_1d(tensor)

args_new[wrt] = tensor

return x
return args_out, kwargs_out, tensor.shape, trax


def function(fun, wrt=0, ntrax=0, parallel=False):
"Evaluate a scalar-valued function."

@wraps(fun)
def evaluate_function(*args, **kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, None, None, ntrax, False)
return fun(*args, **kwargs).x

return evaluate_function


def jacobian(fun, wrt=0, ntrax=0, parallel=False, full_output=False):
"Evaluate the jacobian of a function."

@wraps(fun)
def evaluate_jacobian(*args, **kwargs):

x = arg_to_tensor(args, kwargs, wrt, False)
t = Tensor(x, ntrax=ntrax)
δx = Δx = np.eye(t.size)
indices = range(t.size)

args0, kwargs0 = add_tensor(args, kwargs, wrt, None, None, ntrax, False)
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, False)
func = fun(*args, **kwargs)
fx[axes] = f(func)
dfdx[(*axes, a)] = δ(func)

run(target=kernel, parallel=parallel)(
(a, wrt, δx, Δx, ntrax, args, kwargs) for a in indices
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, False, False, False
)
func = fun(*args, **kwargs)
return f(func)

if full_output:
return np.array(dfdx).reshape(*shape, *t.shape, *t.trax), fx
else:
return np.array(dfdx).reshape(*shape, *t.shape, *t.trax)

return evaluate_jacobian
return evaluate_function


def gradient(fun, wrt=0, ntrax=0, parallel=False, full_output=False, sym=False):
"Evaluate the gradient of a scalar-valued function."
"Evaluate a scalar-valued function."

@wraps(fun)
def evaluate_gradient(*args, **kwargs):

x = arg_to_tensor(args, kwargs, wrt, sym)
t = Tensor(x, ntrax=ntrax)
indices = range(t.size)

fx = np.zeros((1, *t.trax))
dfdx = np.zeros((t.size, *t.trax))
δx = Δx = np.eye(t.size)

if sym:
idx_off_diag = {1: None, 3: [1], 6: [1, 2, 4]}[t.size]
δx[idx_off_diag] /= 2

def kernel(a, wrt, δx, Δx, ntrax, sym, args, kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, δx[a], Δx[a], ntrax, sym)
func = fun(*args, **kwargs)
fx[:] = f(func)
dfdx[a] = δ(func)

run(target=kernel, parallel=parallel)(
(a, wrt, δx, Δx, ntrax, sym, args, kwargs) for a in indices
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, sym, True, False
)

if sym:
dfdx = from_triu_1d(dfdx)
shape = dfdx.shape[:2]
else:
shape = t.shape
func = fun(*args, **kwargs)
grad = δ(func) if sym is False else from_triu_1d(δ(func))

if full_output:
return np.array(dfdx).reshape(*shape, *t.trax), fx[0]
trax = (1,) if len(trax) == 0 else trax
return grad, f(func).reshape(*trax)
else:
return np.array(dfdx).reshape(*shape, *t.trax)
return grad

return evaluate_gradient


def hessian(fun, wrt=0, ntrax=0, parallel=False, full_output=False, sym=False):
"Evaluate the hessian of a scalar-valued function."
"Evaluate a scalar-valued function."

@wraps(fun)
def evaluate_hessian(*args, **kwargs):
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, sym, False, True
)
func = fun(*args, **kwargs)
hess = Δδ(func) if sym is False else from_triu_2d(Δδ(func))

x = arg_to_tensor(args, kwargs, wrt, sym)
t = Tensor(x, ntrax=ntrax)
indices = np.array(np.triu_indices(t.size)).T
if full_output:
grad = δ(func) if sym is False else from_triu_1d(δ(func))
grad = grad.reshape(*shape, *trax)
trax = (1,) if len(trax) == 0 else trax
return hess, grad, f(func).reshape(*trax)
else:
return hess

fx = np.zeros((1, *t.trax))
dfdx = np.zeros((t.size, *t.trax))
d2fdx2 = np.zeros((t.size, t.size, *t.trax))
δx = Δx = np.eye(t.size)
return evaluate_hessian

if sym:
idx_off_diag = {1: None, 3: [1], 6: [1, 2, 4]}[t.size]
δx[idx_off_diag] /= 2

def kernel(a, b, wrt, δx, Δx, ntrax, sym, args, kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, δx[a], Δx[b], ntrax, sym)
func = fun(*args, **kwargs)
fx[:] = f(func)
dfdx[a] = δ(func)
d2fdx2[a, b] = d2fdx2[b, a] = Δδ(func)
def jacobian(fun, wrt=0, ntrax=0, parallel=False, full_output=False):
"Evaluate a scalar-valued function."

run(target=kernel, parallel=parallel)(
(a, b, wrt, δx, Δx, ntrax, sym, args, kwargs) for a, b in indices
@wraps(fun)
def evaluate_jacobian(*args, **kwargs):
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, False, True, False
)

if sym:
dfdx = from_triu_1d(dfdx)
d2fdx2 = from_triu_2d(d2fdx2)
shape = dfdx.shape[:2]
else:
shape = t.shape
func = fun(*args, **kwargs)

if full_output:
return (
np.array(d2fdx2).reshape(*shape, *shape, *t.trax),
np.array(dfdx).reshape(*shape, *t.trax),
fx[0],
)
return δ(func), f(func).reshape(*func.shape, *trax)
else:
return np.array(d2fdx2).reshape(*shape, *shape, *t.trax)
return δ(func)

return evaluate_hessian
return evaluate_jacobian


def gradient_vector_product(fun, wrt=0, ntrax=0, parallel=False):
"Evaluate the gradient-vector-product of a function."

@wraps(fun)
def evaluate_gradient_vector_product(*args, δx, **kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, δx, None, ntrax, False)
return fun(*args, **kwargs).δx
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, False, gradient=True, δx=δx
)
return δ(fun(*args, **kwargs)).reshape(*trax)

return evaluate_gradient_vector_product


def hessian_vectors_product(fun, wrt=0, ntrax=0, parallel=False):
"Evaluate the hessian-vectors-product of a function."

@wraps(fun)
def evaluate_hessian_vectors_product(*args, δx, Δx, **kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, δx, Δx, ntrax, False)
return fun(*args, **kwargs).Δδx

return evaluate_hessian_vectors_product


def hessian_vector_product(fun, wrt=0, ntrax=0, parallel=False, full_output=False):
def hessian_vector_product(fun, wrt=0, ntrax=0, parallel=False):
"Evaluate the hessian-vector-product of a function."

@wraps(fun)
def evaluate_hessian_vector_product(*args, δx, **kwargs):

x = arg_to_tensor(args, kwargs, wrt, False)
t = Tensor(x, ntrax=ntrax)
indices = range(t.size)

fx = np.zeros((1, *t.trax))
hvp = np.zeros((t.size, *t.trax))
Δx = np.eye(t.size)

def kernel(a, wrt, δx, Δx, ntrax, args, kwargs):
args, kwargs = add_tensor(args, kwargs, wrt, δx, Δx[a], ntrax, False)
func = fun(*args, **kwargs)
fx[:] = f(func)
hvp[a] = Δδ(func)

run(target=kernel, parallel=parallel)(
(a, wrt, δx, Δx, ntrax, args, kwargs) for a in indices
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, False, hessian=True, δx=δx
)

return np.array(hvp).reshape(*t.shape, *t.trax)
return Δδ(fun(*args, **kwargs)).reshape(*shape, *trax)

return evaluate_hessian_vector_product


def run(target, parallel=False):
"Serial or threaded execution of a callable target."

@wraps(target)
def run(args=()):
"Run the callable target with iterable args (one item per thread if parallel)."

if not parallel:
[target(*arg) for arg in args]

else:
threads = [Thread(target=target, args=arg) for arg in args]

for th in threads:
th.start()
def hessian_vectors_product(fun, wrt=0, ntrax=0, parallel=False):
"Evaluate the hessian-vectors-product of a function."

for th in threads:
th.join()
@wraps(fun)
def evaluate_hessian_vectors_product(*args, δx, Δx, **kwargs):
args, kwargs, shape, trax = add_tensor(
args, kwargs, wrt, ntrax, False, hessian=True, δx=δx, Δx=Δx
)
return Δδ(fun(*args, **kwargs)).reshape(*trax)

return run
return evaluate_hessian_vectors_product
Loading

0 comments on commit a011ee6

Please sign in to comment.