mtitze/njet

A lightweight AD package, using forward-mode automatic differentiation, in order to determine the higher-order derivatives of a given function in multiple variables.

njet's Issues

Gradient w.r.t. an Array (Tensor)

Hi @mtitze,

do you think that this is the recommended way to use njet on many 3x3 tensors? I took the strain energy function of a Neo-Hookean solid as an example.

from njet.functions import log
from njet import derive
import numpy as np

# init `n` random 3x3 tensors
n = 2 ** 18
x = (np.eye(3) + np.random.rand(n, 3, 3) / 10).T

def grad(fun):
    "A decorator to evaluate the gradient."
    
    def inner(x, **kwargs):
        # init the output
        out = np.zeros((9, *x.shape[2:]))
        
        # obtain the gradient and loop over items
        for key, value in fun.grad(*flatten(x), **kwargs).items():
            out[key] = value
        
        # reshape the gradient
        return out.reshape(3, 3, -1)
    
    return inner

def hess(fun):
    "A decorator to evaluate the hessian."

    def inner(x, **kwargs):
        # init the output
        out = np.zeros((9, 9, *x.shape[2:]))
        
        # obtain the hessian and loop over items
        for key, value in fun.hess(*flatten(x), **kwargs).items():
            out[key] = value
        
        # reshape the hessian
        return out.reshape(3, 3, 3, 3, -1)
    
    return inner

def trace(x):
    return x[0][0] + x[1][1] + x[2][2]

def det(x):
    a = x[0][0] * x[1][1] * x[2][2]
    b = x[0][1] * x[1][2] * x[2][0]
    c = x[0][2] * x[1][0] * x[2][1]
    d = x[2][0] * x[1][1] * x[0][2]
    e = x[2][1] * x[1][2] * x[0][0]
    f = x[2][2] * x[1][0] * x[0][1]
    return a + b + c - d - e - f

flatten = lambda x: x.reshape(9, -1)
reshape = lambda x: [[x[0], x[1], x[2]], [x[3], x[4], x[5]], [x[6], x[7], x[8]]]
# reshape = lambda x: np.array(x, dtype=object).reshape(3, 3, -1) # this fails

def tensorjet(fun):
    "A decorator to reshape the input."
    def inner(*x, **kwargs):
        return fun(reshape(x), **kwargs)
    return inner

@tensorjet
def fun(C):
    "The Neo-Hookean material formulation (isotropic hyperelasticity)."
    return trace(C) - log(det(C))

dfun = derive(fun, order=2, n_args=9)
dfundx = grad(dfun)(x)
d2fundx2 = hess(dfun)(x)

# %timeit hess(dfun)(x)
# 2.08 s ± 27.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Any chance to use NumPy instead of nested lists, i.e. replace the reshape-method?

 reshape = lambda x: np.array(x, dtype=object).reshape(3, 3, -1) # (this fails)

I couldn't get it to work.

Thanks!

Package contains .ipynb_checkpoints directory

Do git clean -fdx when packaging :).

ps. got some scary moments debugging the import errors in CI pipelines because of naming collision with ours internal package having current version 0.1.1 :-D
pps. congrats with public release!