Experimental: NumPy Interoperability

[woodtp]

This PR introduces C-extensions that handle NumPy interop with Bunch and Grid3D.

New for Bunch:

import numpy as np
from orbit.core.bunch import Bunch

# construct a new Bunch instance from an existing numpy array
coords = np.random.normal(size=(1000, 6))

bunch = Bunch.from_numpy(coords)

# Get the coordinates stored within a bunch
coords = bunch.to_numpy()

# Update coordinates of an existing Bunch instance with the values from an array
bunch.update_from_numpy(coords)

Similarly for Grid3D:

import numpy as np
from orbit.core.spacecharge import Grid3D

nx = ny = nz = 100

np_grid = np.random.normal(size=(nx, ny, nz))


# Initialize grid
g = Grid3D(nx, ny, nz)

# Populate grid with external data
g.from_numpy(np_grid)

# Maybe there should be something like:
# g = Grid3D.from_numpy(np_grid)

# Extract grid data to numpy
np_grid = g.to_numpy()

To test the build:

meson setup build -DPyORBIT_EXPERIMENTAL_WITH_NUMPY=true                                                                                                                                          
meson compile -C ./build
pip install . --config-settings=builddir=./build

Resolves #36

[austin-hoover]

I was able to compile following those instructions and looks like it's working. I think this would be a useful feature since NumPy arrays are used in so many other packages. My own preference would be bunch.from_numpy(...) for an already existing bunch object rather than bunch = Bunch.from_numpy(...).

[woodtp]

I was able to compile following those instructions and looks like it's working. I think this would be a useful feature since NumPy arrays are used in so many other packages. My own preference would be bunch.from_numpy(...) for an already existing bunch object rather than bunch = Bunch.from_numpy(...).

Yeah, I think that's fine, too. One question I have about that is: what should the behavior be if the bunch is already filled? It could no-op, throw an exception, or update the existing bunch, which is what update_from_numpy method does.

[woodtp]

The other thing I've been thinking about is how much responsibility the user should have for reading/writing numpy data to/from a file when USE_MPI is enabled. The pattern that I've been using for initializing large bunches involves loading the file in memmap mode and calculating offsets into the array in order to distribute, as evenly as possible, across ranks:

mpi_comm = orbit_mpi.mpi_comm.MPI_COMM_WORLD
mpi_rank = orbit_mpi.MPI_Comm_rank(mpi_comm)
mpi_size = orbit_mpi.MPI_Comm_size(mpi_comm)

coords = np.load("bunch_data.npy", mmap_mode="r")

global_size = coords.shape[0]

base = global_size // mpi_size
remainder = global_size % mpi_size

local_size = base + (1 if mpi_rank < remainder else 0)
start_row = rank * base + min(mpi_rank, remainder)
stop_row = start_row + local_size

local_coords = coords[start_row:stop_row]
bunch = Bunch.from_numpy(local_coords)

Similarly, to save a bunch with MPI, you have to create the file on the primary rank and place a barrier to wait on that operation before you can write into it on the other ranks:

output_file = "output_bunch.npy"
dtype = np.float64

if mpi_rank == 0:
    mm = np.lib.format.open_memmap(
        output_file,
        mode="w+",
        dtype=dtype,
        shape=(global_size, 6),
    )
    del mm # flush

orbit_mpi.MPI_Barrier(mpi_comm)

mm = np.lib.format.open_memmap(
    output_file,
    mode="r+",
    dtype=dtype,
    shape=(global_size, 6),
)

mm[start_row:stop_row, :] = bunch.to_numpy()

and if those things aren't done correctly then you'll wind up with incomplete I/O or maybe worse. It might not be a bad idea to wrap those steps into the bindings if you pass a file name instead of a numpy object.

[austin-hoover]

That makes sense. I was thinking that this would be like the 'readBunch' function, which we currently call from an existing bunch, and it overwrites the coordinates + the bunch properties. But I guess it doesn't really matter if it's overwriting anyway.

In the past we talked about using a different file type for 'dumpBunch' and 'readBunch' instead of a plain text file. I think that would be very helpful especially when there are additional properties for each particle, like macro size, tune, etc. It would also be very helpful to hide the MPI stuff from the user when loading bunches like you're doing here.

What do you think @azukov?