Helpers: to_numpy/cupy

Author: ax3l

src/amrex/Array4.py

@@ -0,0 +1,81 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2022 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+def array4_to_numpy(self, copy=False, order="F"):
+    """
+    Provide a Numpy view into an Array4.
+
+    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
+    -------
+    np.array
+        A numpy n-dimensional array.
+    """
+    import numpy as np
+
+    if order == "F":
+        return np.array(self, copy=copy).T
+    elif order == "C":
+        return np.array(self, copy=copy)
+    else:
+        raise ValueError("The order argument must be F or C.")
+
+
+def array4_to_cupy(self, copy=False, order="F"):
+    """
+    Provide a Cupy view into an Array4.
+
+    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
+    -------
+    cupy.array
+        A numpy n-dimensional array.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    if order == "F":
+        return cp.array(self, copy=copy).T
+    elif order == "C":
+        return cp.array(self, copy=copy)
+    else:
+        raise ValueError("The order argument must be F or C.")

src/amrex/__init__.py

@@ -5,3 +5,48 @@
     "Axel Huebl, Ryan Sandberg, Shreyas Ananthan, Remi Lehe, " "Weiqun Zhang, et al."
 )
 __license__ = "BSD-3-Clause-LBNL"
+
+import numpy as np
+
+from . import amrex_pybind
+from .Array4 import array4_to_cupy, array4_to_numpy
+from .amrex_pybind import *  # noqa
+
+__version__ = amrex_pybind.__version__
+__doc__ = amrex_pybind.__doc__
+__license__ = amrex_pybind.__license__
+__author__ = amrex_pybind.__author__
+
+# at this place we can enhance Python classes with additional methods written
+# in pure Python or add some other Python logic
+#
+
+# Array4 helper methods
+#
+Array4_float.to_numpy = array4_to_numpy
+Array4_double.to_numpy = array4_to_numpy
+Array4_longdouble.to_numpy = array4_to_numpy
+
+Array4_short.to_numpy = array4_to_numpy
+Array4_int.to_numpy = array4_to_numpy
+Array4_long.to_numpy = array4_to_numpy
+Array4_longlong.to_numpy = array4_to_numpy
+
+Array4_ushort.to_numpy = array4_to_numpy
+Array4_uint.to_numpy = array4_to_numpy
+Array4_ulong.to_numpy = array4_to_numpy
+Array4_ulonglong.to_numpy = array4_to_numpy
+
+Array4_float.to_cupy = array4_to_cupy
+Array4_double.to_cupy = array4_to_cupy
+Array4_longdouble.to_cupy = array4_to_cupy
+
+Array4_short.to_cupy = array4_to_cupy
+Array4_int.to_cupy = array4_to_cupy
+Array4_long.to_cupy = array4_to_cupy
+Array4_longlong.to_cupy = array4_to_cupy
+
+Array4_ushort.to_cupy = array4_to_cupy
+Array4_uint.to_cupy = array4_to_cupy
+Array4_ulong.to_cupy = array4_to_cupy
+Array4_ulonglong.to_cupy = array4_to_cupy

tests/test_multifab.py

@@ -46,25 +46,25 @@ def test_mfab_loop(make_mfab):
         # numpy representation: non-copying view, including the
         # guard/ghost region
         #   note: in numpy, indices are in C-order!
-        marr_np = np.array(marr, copy=False)
+        marr_np = marr.to_numpy()
 
         # check the values at start/end are the same: first component
         assert marr_np[0, 0, 0, 0] == marr[bx.small_end]
-        assert marr_np[0, -1, -1, -1] == marr[bx.big_end]
+        assert marr_np[-1, -1, -1, 0] == marr[bx.big_end]
         # same check, but for all components
         for n in range(mfab.num_comp):
             small_end_comp = list(bx.small_end) + [n]
             big_end_comp = list(bx.big_end) + [n]
-            assert marr_np[n, 0, 0, 0] == marr[small_end_comp]
-            assert marr_np[n, -1, -1, -1] == marr[big_end_comp]
+            assert marr_np[0, 0, 0, n] == marr[small_end_comp]
+            assert marr_np[-1, -1, -1, n] == marr[big_end_comp]
 
         # now we do some faster assignments, using range based access
         #   this should fail as out-of-bounds, but does not
         #     does Numpy not check array access for non-owned views?
         # marr_np[24:200, :, :, :] = 42.
 
         #   all components and all indices set at once to 42
-        marr_np[:, :, :, :] = 42.0
+        marr_np[()] = 42.0
 
         # values in start & end still match?
         assert marr_np[0, 0, 0, 0] == marr[bx.small_end]
@@ -210,8 +210,7 @@ def test_mfab_ops_cuda_cupy(make_mfab_device):
     with cupy.profiler.time_range("assign 3 [()]", color_id=0):
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="C")
             # print(marr_cupy.shape)  # 1, 32, 32, 32
             # print(marr_cupy.dtype)  # float64
 
@@ -244,8 +243,7 @@ def set_to_five(mm):
 
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="C")
 
             # write and read into the marr_cupy
             fives_cp = set_to_five(marr_cupy)
@@ -266,8 +264,7 @@ def set_to_seven(x):
 
         for mfi in mfab_device:
             bx = mfi.tilebox().grow(ngv)
-            marr = mfab_device.array(mfi)
-            marr_cupy = cp.array(marr, copy=False)
+            marr_cupy = mfab_device.array(mfi).to_cupy(order="C")
 
             # write and read into the marr_cupy
             set_to_seven(marr_cupy)