Skip to content

Commit

Permalink
Helpers: to_numpy/cupy
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Sep 13, 2023
1 parent 07b43c8 commit 006e8a3
Show file tree
Hide file tree
Showing 8 changed files with 324 additions and 12 deletions.
4 changes: 4 additions & 0 deletions MANIFEST.in
Original file line number Diff line number Diff line change
Expand Up @@ -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*
98 changes: 98 additions & 0 deletions src/amrex/Array4.py
Original file line number Diff line number Diff line change
@@ -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, 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
76 changes: 76 additions & 0 deletions src/amrex/PODVector.py
Original file line number Diff line number Diff line change
@@ -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, 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
105 changes: 105 additions & 0 deletions src/amrex/StructOfArrays.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
"""
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, 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
9 changes: 9 additions & 0 deletions src/amrex/space1d/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,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)
9 changes: 9 additions & 0 deletions src/amrex/space2d/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,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)
9 changes: 9 additions & 0 deletions src/amrex/space3d/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,3 +41,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)
Loading

0 comments on commit 006e8a3

Please sign in to comment.