Get rid of cstagger, use python module with numba

[tiagopereira]

This will get rid of all cstagger machinery, cython compilation, etc. and enable simpler functions for the stagger up/down operations.

[M1kol4j]

I've tested it with Bifrost 3d run in Float32 & Float64, both work fine with the same package. it is also super nice to see correct values at the BC in z-direction. The only problem is speed, it is much slower when compared to cython version, but I am mostly using it on the whole box while I should extract a specific slice first.

Speaking of which, it would nice to have cube_slice argument in get_var, I know there are iix,iiy & iiz but what about additional argument cube_slicer = [slice(None)] * snap.dim. Well ok, this would be mostly for my convenience ;-) See example from bifrost_make_movie.py.

        cube_slicer = [slice(None)] * m.ndim

        if cut_dim == 'x':
          cube_slicer[0] = cut_ax_idx
          if (bobj.params['u_l'] > 1.0):
            xaxis_label = 'x-axsis [Mm]'
            yaxis_label = 'z-axsis [Mm]'
          else:
            xaxis_label = 'x-axsis'
            yaxis_label = 'z-axsis'
        elif cut_dim == 'y':
          cube_slicer[1] = cut_ax_idx
          if (bobj.params['u_l'] > 1.0):
            xaxis_label = 'x-axsis [Mm]'
            yaxis_label = 'z-axsis [Mm]'
          else:
            xaxis_label = 'x-axsis'
            yaxis_label = 'z-axsis'
        elif cut_dim == 'z':
          cube_slicer[2] = cut_ax_idx
          if (bobj.params['u_l'] > 1.0):
            xaxis_label = 'x-axsis [Mm]'
            yaxis_label = 'y-axsis [Mm]'
          else:
            xaxis_label = 'x-axsis'
            yaxis_label = 'y-axsis'

        print('Your 3D cubes will be cut in %s-direction at %s[%d] = %0.2f Mm' % (cut_dim, cut_dim, cut_ax_idx, cut_ax_val))

I will however try to adjust my script to iix ... interface.

[tiagopereira]

Interesting you point out about the speed. In my first tests, it was faster.

Here are some benchmarks run on eagle4, for two simulations: one 720x720x1116 and the other 504x504x496:

from helita.sim import bifrost, cstagger, stagger
bb = bifrost.BifrostData("nw072100", snap=525, fdir=".")
bb.nx, bb.ny, bb.nz
(720, 720, 1116)

# Old cstagger:
%timeit cstagger.do(bb.px, 'xup')
1.49 s ± 1.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit  cstagger.do(bb.py, 'yup')
1.59 s ± 64.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit cstagger.do(bb.pz, 'zup')
3.4 s ± 18.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# New stagger:
%timeit stagger.do(bb.px, 'xup')
1.58 s ± 76.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit stagger.do(bb.py, 'yup')
1.54 s ± 66.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit stagger.do(bb.pz, 'zup')
2.4 s ± 228 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

For a 504x504x496 sim:

bb = bifrost.BifrostData("cb24bih", snap=525, fdir=".")
bb.nx, bb.ny, bb.nz
(504, 504, 496)

# Old cstagger:
%timeit -n 5 cstagger.do(bb.px, 'xup')
255 ms ± 69.5 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
%timeit -n 5 cstagger.do(bb.py, 'yup')
181 ms ± 993 µs per loop (mean ± std. dev. of 7 runs, 5 loops each)
%timeit -n 5 cstagger.do(bb.pz, 'zup')
602 ms ± 58.8 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

# New stagger:
%timeit -n 5 stagger.do(bb.px, 'xup')
373 ms ± 17.9 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
%timeit -n 5 stagger.do(bb.py, 'yup')
332 ms ± 18.8 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)
%timeit -n 5 stagger.do(bb.pz, 'zup')
500 ms ± 25.2 ms per loop (mean ± std. dev. of 7 runs, 5 loops each)

For the up operations, the new stagger is a bit slower for the x and y dimensions, but a bit faster for the z direction. However, it seems that cstagger is slower for the down operations, while stagger takes about the same, so they are more comparable there.

[tiagopereira]

In some cases for a smaller (256, 256, 512) simulation the stagger version got speedups of about 3x for some of the operations, but this doesn't seem to happen for all machines. In any case, I am not seeing it much slower, so maybe the fault lies somewhere else?

[M1kol4j]

I am start thinking that my entire python is slow now, but since I am using more and more Julia now I can simply re-install anaconda and check again. Results you have shown looks awesome. You should not wait with pull request, this speed with and fix for double precision are great update to current version and I will switch to it as soon as possible.

…

On May 4 2021, at 4:41 pm, Tiago M. D. Pereira ***@***.***> wrote: In some cases for a smaller (256, 256, 512) simulation the stagger version got speedups of about 3x for some of the operations, but this doesn't seem to happen for all machines. In any case, I am not seeing it much slower, so maybe the fault lies somewhere else? — You are receiving this because your review was requested. Reply to this email directly, view it on GitHub (#31 (comment)), or unsubscribe (https://github.com/notifications/unsubscribe-auth/ACAVX37NYKBRATGNKJ6BLHTTMABTPANCNFSM434CYEVA).

[tiagopereira]

Sorry, this was not ready to be merged yet (need to remove the cstagger part).