Helpers: to_numpy/cupy

MANIFEST.in

@@ -6,6 +6,10 @@ recursive-include cmake *
 recursive-include src *
 recursive-include tests *
 
+# avoid accidentially copying compiled Python files
+global-exclude */__pycache__/*
+global-exclude *.pyc
+
 # see .gitignore
 prune cmake-build*
 prune .spack-env*

src/amrex/Array4.py

@@ -0,0 +1,98 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 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 cupy 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.")
+
+
+def register_Array4_extension(amr):
+    """Array4 helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every Array4_* type
+    for _, Array4_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("Array4_"),
+    ):
+        Array4_type.to_numpy = array4_to_numpy
+        Array4_type.to_cupy = array4_to_cupy

src/amrex/PODVector.py

@@ -0,0 +1,76 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def podvector_to_numpy(self, copy=False):
+    """
+    Provide a Numpy view into a PODVector (e.g., RealVector, IntVector).
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    np.array
+        A 1D numpy array.
+    """
+    import numpy as np
+
+    if self.size() > 0:
+        return np.array(self, copy=copy)
+    else:
+        raise ValueError("Vector is empty.")
+
+
+def podvector_to_cupy(self, copy=False):
+    """
+    Provide a Cupy view into a PODVector (e.g., RealVector, IntVector).
+
+    Parameters
+    ----------
+    self : amrex.PODVector_*
+        A PODVector class in pyAMReX
+    copy : bool, optional
+        Copy the data if true, otherwise create a view (default).
+
+    Returns
+    -------
+    cupy.array
+        A 1D cupy array.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    if self.size() > 0:
+        return cp.array(self, copy=copy)
+    else:
+        raise ValueError("Vector is empty.")
+
+
+def register_PODVector_extension(amr):
+    """PODVector helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every PODVector_* type
+    for _, POD_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("PODVector_"),
+    ):
+        POD_type.to_numpy = podvector_to_numpy
+        POD_type.to_cupy = podvector_to_cupy

src/amrex/StructOfArrays.py

@@ -0,0 +1,99 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+from collections import namedtuple
+
+
+def soa_to_numpy(self, copy=False):
+    """
+    Provide Numpy views into a StructOfArrays.
+
+    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 lists
+        of 1D numpy arrays.
+    """
+    import numpy as np
+
+    SoA_np = namedtuple(type(self).__name__ + "_np", ["real", "int"])
+
+    soa_view = SoA_np([], [])
+
+    if self.size() == 0:
+        raise ValueError("SoA is empty.")
+
+    for idx_real in range(self.NumRealComps()):
+        soa_view.real.append(np.array(self.GetRealData(idx_real), copy=copy))
+
+    for idx_int in range(self.NumIntComps()):
+        soa_view.int.append(np.array(self.GetIntData(idx_int), copy=copy))
+
+    return soa_view
+
+
+def soa_to_cupy(self, copy=False):
+    """
+    Provide Cupy views into a StructOfArrays.
+
+    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 lists
+        of 1D numpy arrays.
+
+    Raises
+    ------
+    ImportError
+        Raises an exception if cupy is not installed
+    """
+    import cupy as cp
+
+    SoA_cp = namedtuple(type(self).__name__ + "_cp", ["real", "int"])
+
+    soa_view = SoA_cp([], [])
+
+    if self.size() == 0:
+        raise ValueError("SoA is empty.")
+
+    for idx_real in range(self.NumRealComps()):
+        soa_view.real.append(cp.array(self.GetRealData(idx_real), copy=copy))
+
+    for idx_int in range(self.NumIntComps()):
+        soa_view.int.append(cp.array(self.GetIntData(idx_int), copy=copy))
+
+    return soa_view
+
+
+def register_SoA_extension(amr):
+    """StructOfArrays helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every StructOfArrays_* type
+    for _, SoA_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("StructOfArrays_"),
+    ):
+        SoA_type.to_numpy = soa_to_numpy
+        SoA_type.to_cupy = soa_to_cupy

src/amrex/space1d/__init__.py

@@ -31,6 +31,7 @@
 # in pure Python or add some other Python logic
 #
 def d_decl(x, y, z):
+    """Return a tuple of the first passed element"""
     return (x,)
 
 
@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_1d_pybind)
+register_PODVector_extension(amrex_1d_pybind)
+register_SoA_extension(amrex_1d_pybind)

src/amrex/space2d/__init__.py

@@ -31,6 +31,7 @@
 # in pure Python or add some other Python logic
 #
 def d_decl(x, y, z):
+    """Return a tuple of the first two passed elements"""
     return (x, y)
 
 
@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_2d_pybind)
+register_PODVector_extension(amrex_2d_pybind)
+register_SoA_extension(amrex_2d_pybind)

src/amrex/space3d/__init__.py

@@ -31,6 +31,7 @@
 # in pure Python or add some other Python logic
 #
 def d_decl(x, y, z):
+    """Return a tuple of the three passed elements"""
     return (x, y, z)
 
 
@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
         print(*args, **kwargs)
     elif ParallelDescriptor.IOProcessor():
         print(*args, **kwargs)
+
+
+from ..Array4 import register_Array4_extension
+from ..PODVector import register_PODVector_extension
+from ..StructOfArrays import register_SoA_extension
+
+register_Array4_extension(amrex_3d_pybind)
+register_PODVector_extension(amrex_3d_pybind)
+register_SoA_extension(amrex_3d_pybind)