Slow ArrayD operation

[ungsik]

I am converting the following python code into rust using ArrayD.

tmp = numpy.zeros_like(h4e)
dim = h4e.shape[0]

for i in range(dim):
    for j in range(dim):
        for k in range(dim):
            for m in range(dim):
                tmp[i, j, k, m, :, :, :, :] = (
                    (h4e[i, j, k, m, :, :, :, :] -
                     h4e[i, j, m, k, :, :, :, :] +
                     h4e[i, m, j, k, :, :, :, :] -
                     h4e[i, m, k, j, :, :, :, :] -
                     h4e[i, k, j, m, :, :, :, :] +
                     h4e[i, k, m, j, :, :, :, :]) -
                    (h4e[j, i, k, m, :, :, :, :] -
                     h4e[j, i, m, k, :, :, :, :] +
                     h4e[j, m, i, k, :, :, :, :] -
                     h4e[j, m, k, i, :, :, :, :] -
                     h4e[j, k, i, m, :, :, :, :] +
                     h4e[j, k, m, i, :, :, :, :]) -
                    (h4e[k, j, i, m, :, :, :, :] -
                     h4e[k, j, m, i, :, :, :, :] +
                     h4e[k, m, j, i, :, :, :, :] -
                     h4e[k, m, i, j, :, :, :, :] -
                     h4e[k, i, j, m, :, :, :, :] +
                     h4e[k, i, m, j, :, :, :, :]) -
                    (h4e[m, j, k, i, :, :, :, :] -
                     h4e[m, j, i, k, :, :, :, :] +
                     h4e[m, i, j, k, :, :, :, :] -
                     h4e[m, i, k, j, :, :, :, :] -
                     h4e[m, k, j, i, :, :, :, :] +
                     h4e[m, k, i, j, :, :, :, :])) / 24.0

When I tried the following code, it takes about 50 seconds when the shape is [6; 8].
Is there more performant and correct way?

let mut tmp = ArrayD::<c64>::zeros(IxDyn(&[dim; 8]));
for (i, j, k, m) in iproduct!(0..dim, 0..dim, 0..dim, 0..dim) {
    for (p, q, r, s) in iproduct!(0..dim, 0..dim, 0..dim, 0..dim) {
        tmp[&[i, j, k, m, p, q, r, s][..]] = (
            (h4e[&[i, j, k, m, p, q, r, s][..]] -
            h4e[&[i, j, m, k, p, q, r, s][..]] +
        h4e[&[i, m, j, k, p, q, r, s][..]] -
            h4e[&[i, m, k, j, p, q, r, s][..]] -
            h4e[&[i, k, j, m, p, q, r, s][..]] +
        h4e[&[i, k, m, j, p, q, r, s][..]]) -
            (h4e[&[j, i, k, m, p, q, r, s][..]] -
            h4e[&[j, i, m, k, p, q, r, s][..]] +
        h4e[&[j, m, i, k, p, q, r, s][..]] -
            h4e[&[j, m, k, i, p, q, r, s][..]] -
            h4e[&[j, k, i, m, p, q, r, s][..]] +
        h4e[&[j, k, m, i, p, q, r, s][..]]) -
            (h4e[&[k, j, i, m, p, q, r, s][..]] -
            h4e[&[k, j, m, i, p, q, r, s][..]] +
        h4e[&[k, m, j, i, p, q, r, s][..]] -
            h4e[&[k, m, i, j, p, q, r, s][..]] -
            h4e[&[k, i, j, m, p, q, r, s][..]] +
        h4e[&[k, i, m, j, p, q, r, s][..]]) -
            (h4e[&[m, j, k, i, p, q, r, s][..]] -
            h4e[&[m, j, i, k, p, q, r, s][..]] +
        h4e[&[m, i, j, k, p, q, r, s][..]] -
            h4e[&[m, i, k, j, p, q, r, s][..]] -
            h4e[&[m, k, j, i, p, q, r, s][..]] +
        h4e[&[m, k, i, j, p, q, r, s][..]])) / 24.0
    }
}

[bluss]

General performance tips are written on this first page: https://docs.rs/ndarray/0.15.3/ndarray/

My additional summary:

• Avoid using iproduct, it is slower than nested for loops
• Avoid using indexing for the inner loop - use higher level traversals
• Avoid using ArrayD for the inner loop (your inner loop spans over 4 axes, so you can easily avoid)

Think of ArrayD as a container of arrays. Doing operations on individual elements is slow.

it would be very slow if you wrote it the same way in Python, as in Rust. Not your fault that these primitives are slow, but they are. We avoid indexing in Rust like you'd like to do in Python.