test.py

import numpy as np
import timeit
from nutils import core, log, function as fn, plot, _, mesh


def one(domain, basis, z):
    mx = basis[:,_,_] * basis[_,:,_] * basis[_,_,:]
    mx = (mx[:,_,:,_,:,_] * z[:,:,_,_,_,_] * z[_,_,:,:,_,_] * z[_,_,_,_,:,:]).sum([0, 2, 4])
    return domain.integrate(mx, ischeme='gauss9')


def two(domain, basis, z):
    basis = fn.matmat(basis, z)
    mx = basis[:,_,_] * basis[_,:,_] * basis[_,_,:]
    return domain.integrate(mx, ischeme='gauss9')


def three(domain, basis, z):
    mx = fn.ravel(basis[:,_,_] * basis[_,:,_] * basis[_,_,:], 1)
    mx = domain.integrate(mx, ischeme='gauss9').core

    nb, nr = len(basis), z.shape[1]

    mx = mx.T.dot(z).T.reshape((nr, nb, nb))
    mx = mx.dot(z)
    return mx.transpose((0,2,1)).dot(z).transpose((0,1,2))


if __name__ == '__main__':
    pts = np.linspace(0, 1, 31)
    domain, geom = mesh.rectilinear([pts,pts])
    basis = domain.basis('spline', degree=1)
    z = np.random.random((len(basis), 15))

    log.user(timeit.timeit(
        'one(domain,basis,z)', number=1,
        globals={'one': one, 'domain': domain, 'geom': geom, 'basis': basis, 'z': z}
    ))

    log.user(timeit.timeit(
        'two(domain,basis,z)', number=1,
        globals={'two': two, 'domain': domain, 'geom': geom, 'basis': basis, 'z': z}
    ))

    log.user(timeit.timeit(
        'three(domain,basis,z)', number=1,
        globals={'three': three, 'domain': domain, 'geom': geom, 'basis': basis, 'z': z}
    ))