to_xp: AoS, SoA, PODVector, MF, Array4

Add a `to_xp()` method that returns either NumPy or CuPy
containers, depending on `amr.Config.have_gpu`, to all
major AMReX data containers.

src/amrex/Array4.py

@@ -9,7 +9,7 @@
 
 def array4_to_numpy(self, copy=False, order="F"):
     """
-    Provide a Numpy view into an Array4.
+    Provide a NumPy view into an Array4.
 
     This includes ngrow guard cells of the box.
 
@@ -32,7 +32,7 @@ def array4_to_numpy(self, copy=False, order="F"):
     Returns
     -------
     np.array
-        A numpy n-dimensional array.
+        A NumPy n-dimensional array.
     """
     import numpy as np
 
@@ -52,7 +52,7 @@ def array4_to_numpy(self, copy=False, order="F"):
 
 def array4_to_cupy(self, copy=False, order="F"):
     """
-    Provide a Cupy view into an Array4.
+    Provide a CuPy view into an Array4.
 
     This includes ngrow guard cells of the box.
 
@@ -92,6 +92,44 @@ def array4_to_cupy(self, copy=False, order="F"):
         raise ValueError("The order argument must be F or C.")
 
 
+def array4_to_xp(self, copy=False, order="F"):
+    """
+    Provide a NumPy or CuPy view into an Array4, depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    This includes ngrow guard cells of the box.
+
+    Note on the order of indices:
+    By default, this is as in AMReX in Fortran contiguous order, indexing as
+    x,y,z. This has performance implications for use in external libraries such
+    as cupy.
+    The order="C" option will index as z,y,x and perform better with cupy.
+    https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
+
+    Parameters
+    ----------
+    self : amrex.Array4_*
+        An Array4 class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+    order : string, optional
+        F order (default) or C. C is faster with external libraries.
+
+    Returns
+    -------
+    xp.array
+        A NumPy or CuPy n-dimensional array.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return (
+        self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order)
+    )
+
+
 def register_Array4_extension(amr):
     """Array4 helper methods"""
     import inspect
@@ -106,3 +144,4 @@ def register_Array4_extension(amr):
     ):
         Array4_type.to_numpy = array4_to_numpy
         Array4_type.to_cupy = array4_to_cupy
+        Array4_type.to_xp = array4_to_xp

src/amrex/ArrayOfStructs.py

@@ -9,7 +9,7 @@
 
 def aos_to_numpy(self, copy=False):
     """
-    Provide Numpy views into a ArrayOfStructs.
+    Provide NumPy views into a ArrayOfStructs.
 
     Parameters
     ----------
@@ -22,7 +22,7 @@ def aos_to_numpy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each lists
-        of 1D numpy arrays.
+        of 1D NumPy arrays.
     """
     import numpy as np
 
@@ -42,7 +42,7 @@ def aos_to_numpy(self, copy=False):
 
 def aos_to_cupy(self, copy=False):
     """
-    Provide Cupy views into a ArrayOfStructs.
+    Provide CuPy views into a ArrayOfStructs.
 
     Parameters
     ----------
@@ -55,7 +55,7 @@ def aos_to_cupy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each lists
-        of 1D numpy arrays.
+        of 1D NumPy arrays.
 
     Raises
     ------
@@ -70,6 +70,32 @@ def aos_to_cupy(self, copy=False):
     return cp.array(self, copy=copy)
 
 
+def aos_to_xp(self, copy=False):
+    """
+    Provide NumPy or CuPy views into a ArrayOfStructs, depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    self : amrex.ArrayOfStructs_*
+        An ArrayOfStructs class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    namedtuple
+        A tuple with real and int components that are each lists
+        of 1D NumPy or CuPy arrays.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
+
+
 def register_AoS_extension(amr):
     """ArrayOfStructs helper methods"""
     import inspect
@@ -84,3 +110,4 @@ def register_AoS_extension(amr):
     ):
         AoS_type.to_numpy = aos_to_numpy
         AoS_type.to_cupy = aos_to_cupy
+        AoS_type.to_xp = aos_to_xp

src/amrex/MultiFab.py

@@ -9,9 +9,9 @@
 from .Iterator import next
 
 
-def mf_to_numpy(amr, self, copy=False, order="F"):
+def mf_to_numpy(self, copy=False, order="F"):
     """
-    Provide a Numpy view into a MultiFab.
+    Provide a NumPy view into a MultiFab.
 
     This includes ngrow guard cells of each box.
 
@@ -34,9 +34,13 @@ def mf_to_numpy(amr, self, copy=False, order="F"):
     Returns
     -------
     list of numpy.array
-        A list of numpy n-dimensional arrays, for each local block in the
+        A list of NumPy n-dimensional arrays, for each local block in the
         MultiFab.
     """
+    import inspect
+
+    amr = inspect.getmodule(self)
+
     mf = self
     if copy:
         mf = amr.MultiFab(
@@ -58,7 +62,7 @@ def mf_to_numpy(amr, self, copy=False, order="F"):
 
 def mf_to_cupy(self, copy=False, order="F"):
     """
-    Provide a Cupy view into a MultiFab.
+    Provide a CuPy view into a MultiFab.
 
     This includes ngrow guard cells of each box.
 
@@ -81,7 +85,7 @@ def mf_to_cupy(self, copy=False, order="F"):
     Returns
     -------
     list of cupy.array
-        A list of cupy n-dimensional arrays, for each local block in the
+        A list of CuPy n-dimensional arrays, for each local block in the
         MultiFab.
 
     Raises
@@ -96,6 +100,46 @@ def mf_to_cupy(self, copy=False, order="F"):
     return views
 
 
+def mf_to_xp(self, copy=False, order="F"):
+    """
+    Provide a NumPy or CuPy view into a MultiFab,
+    depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    This includes ngrow guard cells of each box.
+
+    Note on the order of indices:
+    By default, this is as in AMReX in Fortran contiguous order, indexing as
+    x,y,z. This has performance implications for use in external libraries such
+    as cupy.
+    The order="C" option will index as z,y,x and perform better with cupy.
+    https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
+
+    Parameters
+    ----------
+    self : amrex.MultiFab
+        A MultiFab class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+    order : string, optional
+        F order (default) or C. C is faster with external libraries.
+
+    Returns
+    -------
+    list of xp.array
+        A list of NumPy or CuPy n-dimensional arrays, for each local block in the
+        MultiFab.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return (
+        self.to_cupy(copy, order) if amr.Config.have_gpu else self.to_numpy(copy, order)
+    )
+
+
 def copy_multifab(amr, self):
     """
     Create a copy of this MultiFab, using the same Arena.
@@ -141,12 +185,9 @@ def register_MultiFab_extension(amr):
     # register member functions for the MultiFab type
     amr.MultiFab.__iter__ = lambda mfab: amr.MFIter(mfab)
 
-    amr.MultiFab.to_numpy = lambda self, copy=False, order="F": mf_to_numpy(
-        amr, self, copy, order
-    )
-    amr.MultiFab.to_numpy.__doc__ = mf_to_numpy.__doc__
-
+    amr.MultiFab.to_numpy = mf_to_numpy
     amr.MultiFab.to_cupy = mf_to_cupy
+    amr.MultiFab.to_xp = mf_to_xp
 
     amr.MultiFab.copy = lambda self: copy_multifab(amr, self)
     amr.MultiFab.copy.__doc__ = copy_multifab.__doc__

src/amrex/PODVector.py

@@ -9,7 +9,7 @@
 
 def podvector_to_numpy(self, copy=False):
     """
-    Provide a Numpy view into a PODVector (e.g., RealVector, IntVector).
+    Provide a NumPy view into a PODVector (e.g., RealVector, IntVector).
 
     Parameters
     ----------
@@ -21,7 +21,7 @@ def podvector_to_numpy(self, copy=False):
     Returns
     -------
     np.array
-        A 1D numpy array.
+        A 1D NumPy array.
     """
     import numpy as np
 
@@ -41,7 +41,7 @@ def podvector_to_numpy(self, copy=False):
 
 def podvector_to_cupy(self, copy=False):
     """
-    Provide a Cupy view into a PODVector (e.g., RealVector, IntVector).
+    Provide a CuPy view into a PODVector (e.g., RealVector, IntVector).
 
     Parameters
     ----------
@@ -68,6 +68,32 @@ def podvector_to_cupy(self, copy=False):
         raise ValueError("Vector is empty.")
 
 
+def podvector_to_xp(self, copy=False):
+    """
+    Provide a NumPy or CuPy view into a PODVector (e.g., RealVector, IntVector),
+    depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    xp.array
+        A 1D NumPy or CuPy array.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
+
+
 def register_PODVector_extension(amr):
     """PODVector helper methods"""
     import inspect
@@ -82,3 +108,4 @@ def register_PODVector_extension(amr):
     ):
         POD_type.to_numpy = podvector_to_numpy
         POD_type.to_cupy = podvector_to_cupy
+        POD_type.to_xp = podvector_to_xp

src/amrex/StructOfArrays.py

@@ -76,7 +76,7 @@ def soa_int_comps(self, num_comps):
 
 def soa_to_numpy(self, copy=False):
     """
-    Provide Numpy views into a StructOfArrays.
+    Provide NumPy views into a StructOfArrays.
 
     Parameters
     ----------
@@ -89,7 +89,7 @@ def soa_to_numpy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each dicts
-        of 1D numpy arrays. The dictionary key order is the same as
+        of 1D NumPy arrays. The dictionary key order is the same as
         in the C++ component order.
         For pure SoA particle layouts, an additional component idcpu
         with global particle indices is populated.
@@ -129,7 +129,7 @@ def soa_to_numpy(self, copy=False):
 
 def soa_to_cupy(self, copy=False):
     """
-    Provide Cupy views into a StructOfArrays.
+    Provide CuPy views into a StructOfArrays.
 
     Parameters
     ----------
@@ -142,7 +142,7 @@ def soa_to_cupy(self, copy=False):
     -------
     namedtuple
         A tuple with real and int components that are each dicts
-        of 1D numpy arrays. The dictionary key order is the same as
+        of 1D NumPy arrays. The dictionary key order is the same as
         in the C++ component order.
         For pure SoA particle layouts, an additional component idcpu
         with global particle indices is populated.
@@ -183,6 +183,35 @@ def soa_to_cupy(self, copy=False):
     return soa_view
 
 
+def soa_to_xp(self, copy=False):
+    """
+    Provide NumPy or CuPy views into a StructOfArrays, depending on amr.Config.have_gpu .
+
+    This function is similar to CuPy's xp naming suggestion for CPU/GPU agnostic code:
+    https://docs.cupy.dev/en/stable/user_guide/basic.html#how-to-write-cpu-gpu-agnostic-code
+
+    Parameters
+    ----------
+    self : amrex.StructOfArrays_*
+        A StructOfArrays class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    namedtuple
+        A tuple with real and int components that are each dicts
+        of 1D NumPy or CuPy arrays. The dictionary key order is the same as
+        in the C++ component order.
+        For pure SoA particle layouts, an additional component idcpu
+        with global particle indices is populated.
+    """
+    import inspect
+
+    amr = inspect.getmodule(self)
+    return self.to_cupy(copy) if amr.Config.have_gpu else self.to_numpy(copy)
+
+
 def register_SoA_extension(amr):
     """StructOfArrays helper methods"""
     import inspect
@@ -202,3 +231,4 @@ def register_SoA_extension(amr):
         # converters
         SoA_type.to_numpy = soa_to_numpy
         SoA_type.to_cupy = soa_to_cupy
+        SoA_type.to_xp = soa_to_xp