Implement multifab wrapper functionality from WarpX

[dpgrote]

This adds the code from the WarpX fields.py that allows global indexing of MultiFabs. This is done by adding __getitem__ and __setitem__ methods for the MultiFab object.

Note that this incorporates the changes from WarpX PR #4370 that had not been merged in.

This currently only does the get/set without including the ghost cells. What is needed is a way of specifying inclusion of the ghost cells. Perhaps this?

mf = MultiFab()
mf[:,5,6]  # without ghost cells
mf[:,5,6] = 2.  # without ghost cells
mf.with_ghosts([slice(None),5,6])  # with ghost cells (in the first dimension)
mf.with_ghosts([slice(None),5,6], 2.)  # with ghost cells

Maybe this, using imaginary numbers to refer to ghost cells?

mf[-1j,5,6]  # the first ghost cell on the lower boundary
mf[+1j,5,6]  # the first ghost cell on the upper boundary
mf[(),5,6]  # the full first dimension including ghosts
mf[:,5,6]  # the first dimension without ghosts

I prefer this method.

[ax3l]

Hi @dpgrote,

That is fantastic, thanks a lot for porting this!

I think in symmetry with AMReX functions like Box::grow(), the syntax .with_ghosts is clearer. Could also be .grow if we want, but either way is great.

We could even add default arguments to =nGrowVect() for the function so it can be called .with_ghosts() and get the ghosts.

[ax3l]

@dpgrote This is awesome. To add CI and show usage examples, can you please add a few new examples in tests/?

This then makes it very easy to add add a new tab in fields, using the snippet from the test then in the user-facing doc page:

• user-facing: https://pyamrex.readthedocs.io/en/latest/usage/compute.html#fields
• source annotation example:

  pyamrex/tests/test_multifab.py

  Lines 47 to 64 in 1c66690

  	# Manual: Compute Mfab Simple START
  	# finest active MR level, get from a
  	# simulation's AmrMesh object, e.g.:
  	# finest_level = sim.finest_level
  	finest_level = 0 # no MR
  	# iterate over mesh-refinement levels
  	for lev in range(finest_level + 1):
  	# get an existing MultiFab, e.g.,
  	# from a simulation:
  	# mfab = sim.get_field(lev=lev)
  	# Config = sim.extension.Config
  	field_list = mfab.to_xp()
  	for field in field_list:
  	field[()] = 42.0
  	# Manual: Compute Mfab Simple END

• https://github.com/AMReX-Codes/pyamrex/blob/development/docs/source/usage/compute.rst

[dpgrote]

I think in symmetry with AMReX functions like Box::grow(), the syntax .with_ghosts is clearer. Could also be .grow if we want, but either way is great.

We could even add default arguments to =nGrowVect() for the function so it can be called .with_ghosts() and get the ghosts.

How about this, add the method mf.alias_with_ghosts() that returns an alias of the MF with the boxes grown by the number of ghost cells. Then the __getitem__ and __setitem__ can be use directly.

Note that there is a subtlety in __setitem__ regarding ghosts. With include_ghosts False, only valid cells are set. When True, the ghost cells of the individual blocks are set as well (as if FillBoundary had been called). Perhaps the internal ghost cells should always be set?

[ax3l]

We need to put this into a function that has access to amr, so we only attempt an import if:

if amr.Config.have_mpi:

Otherwise, this can cause issues in scenarios, e.g., running serially on Cray machines.

[ax3l]

While we have access to amr, we can replace mpi.COMM_WORLD with amr.ParallelDescriptor.Communicator() later on.

I did not expose that element yet, but we already created an mpi4py compatible MPI_Communicator for the MPMD logic in pyAMReX. This can be a follow-up PR.

[dpgrote]

The more I think about this (and work on implementing something), the more I like the idea of using imaginary numbers for the guard cells. The idea of using negative numbers for the lower ghost cells breaks the standard python usage, since for every sequence type, negative indices are always relative to the end of the sequence. It is asymmetric since there then no nice way of referring to the upper ghost cells. The imaginary indices are a very concise and unambiguous way is specifying the ghosts, both lower and upper. I'll implement this and provide examples to show what it's like.

[ax3l]

So to get a plane with ghost cells, something like this? :)

mf[-2j:+2j, 0, 0]  # include 2 ghost cells in the first plane of the global domain

or is that

mf[-2j:-1+3j, 0, 0]  # include 2 ghost cells in the first plane of the global domain

[dpgrote]

So to get a plane with ghost cells, something like this? :)

mf[-2j:+2j, 0, 0]  # include 2 ghost cells in the first plane of the global domain

or is that

mf[-2j:-1+3j, 0, 0]  # include 2 ghost cells in the first plane of the global domain

It would be

mf[-2j:0, 0, 0]

The -2j refers to the 2nd ghost cell. 0 refers to the first valid cell, remembering python slice range that goes up to but does not include the end of the slice range.

If you want all of the cells, the valid cells plus the two guard cells on each side, then it would be

mf[-2j:3j,0,0]

One thing that is missing is a clean way to specify and open ended range on one side including the ghost cells.

[ax3l]

This is awesome and impressive, thank you! 🚀 ✨