Skip to content

Commit d8da7df

Browse files
ax3lax3l
committed
Helpers: to_numpy/cupy
1 parent 0aa7fe9 commit d8da7df

File tree

8 files changed

+321
-12
lines changed

8 files changed

+321
-12
lines changed

‎MANIFEST.in

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@ recursive-include cmake *
66
recursive-include src *
77
recursive-include tests *
88

9+
# avoid accidentially copying compiled Python files
10+
global-exclude */__pycache__/*
11+
global-exclude *.pyc
12+
913
# see .gitignore
1014
prune cmake-build*
1115
prune .spack-env*

‎src/amrex/Array4.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
9+
10+
def array4_to_numpy(self, copy=False, order="F"):
11+
"""
12+
Provide a Numpy view into an Array4.
13+
14+
Note on the order of indices:
15+
By default, this is as in AMReX in Fortran contiguous order, indexing as
16+
x,y,z. This has performance implications for use in external libraries such
17+
as cupy.
18+
The order="C" option will index as z,y,x and perform better with cupy.
19+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
20+
21+
Parameters
22+
----------
23+
self : amrex.Array4_*
24+
An Array4 class in pyAMReX
25+
copy : bool, optional
26+
Copy the data if true, otherwise create a view (default).
27+
order : string, optional
28+
F order (default) or C. C is faster with external libraries.
29+
30+
Returns
31+
-------
32+
np.array
33+
A numpy n-dimensional array.
34+
"""
35+
import numpy as np
36+
37+
if order == "F":
38+
return np.array(self, copy=copy).T
39+
elif order == "C":
40+
return np.array(self, copy=copy)
41+
else:
42+
raise ValueError("The order argument must be F or C.")
43+
44+
45+
def array4_to_cupy(self, copy=False, order="F"):
46+
"""
47+
Provide a Cupy view into an Array4.
48+
49+
Note on the order of indices:
50+
By default, this is as in AMReX in Fortran contiguous order, indexing as
51+
x,y,z. This has performance implications for use in external libraries such
52+
as cupy.
53+
The order="C" option will index as z,y,x and perform better with cupy.
54+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
55+
56+
Parameters
57+
----------
58+
self : amrex.Array4_*
59+
An Array4 class in pyAMReX
60+
copy : bool, optional
61+
Copy the data if true, otherwise create a view (default).
62+
order : string, optional
63+
F order (default) or C. C is faster with external libraries.
64+
65+
Returns
66+
-------
67+
cupy.array
68+
A cupy n-dimensional array.
69+
70+
Raises
71+
------
72+
ImportError
73+
Raises an exception if cupy is not installed
74+
"""
75+
import cupy as cp
76+
77+
if order == "F":
78+
return cp.array(self, copy=copy).T
79+
elif order == "C":
80+
return cp.array(self, copy=copy)
81+
else:
82+
raise ValueError("The order argument must be F or C.")
83+
84+
85+
def register_Array4_extension(amr):
86+
"""Array4 helper methods"""
87+
import inspect
88+
import sys
89+
90+
# register member functions for every Array4_* type
91+
for _, Array4_type in inspect.getmembers(
92+
sys.modules[amr.__name__],
93+
lambda member: inspect.isclass(member)
94+
and member.__module__ == amr.__name__
95+
and member.__name__.startswith("Array4_"),
96+
):
97+
Array4_type.to_numpy = array4_to_numpy
98+
Array4_type.to_cupy = array4_to_cupy

‎src/amrex/PODVector.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
9+
10+
def podvector_to_numpy(self, copy=False):
11+
"""
12+
Provide a Numpy view into a PODVector (e.g., RealVector, IntVector).
13+
14+
Parameters
15+
----------
16+
self : amrex.PODVector_*
17+
A PODVector class in pyAMReX
18+
copy : bool, optional
19+
Copy the data if true, otherwise create a view (default).
20+
21+
Returns
22+
-------
23+
np.array
24+
A 1D numpy array.
25+
"""
26+
import numpy as np
27+
28+
if self.size() > 0:
29+
return np.array(self, copy=copy)
30+
else:
31+
raise ValueError("Vector is empty.")
32+
33+
34+
def podvector_to_cupy(self, copy=False):
35+
"""
36+
Provide a Cupy view into a PODVector (e.g., RealVector, IntVector).
37+
38+
Parameters
39+
----------
40+
self : amrex.PODVector_*
41+
A PODVector class in pyAMReX
42+
copy : bool, optional
43+
Copy the data if true, otherwise create a view (default).
44+
45+
Returns
46+
-------
47+
cupy.array
48+
A 1D cupy array.
49+
50+
Raises
51+
------
52+
ImportError
53+
Raises an exception if cupy is not installed
54+
"""
55+
import cupy as cp
56+
57+
if self.size() > 0:
58+
return cp.array(self, copy=copy)
59+
else:
60+
raise ValueError("Vector is empty.")
61+
62+
63+
def register_PODVector_extension(amr):
64+
"""PODVector helper methods"""
65+
import inspect
66+
import sys
67+
68+
# register member functions for every PODVector_* type
69+
for _, POD_type in inspect.getmembers(
70+
sys.modules[amr.__name__],
71+
lambda member: inspect.isclass(member)
72+
and member.__module__ == amr.__name__
73+
and member.__name__.startswith("PODVector_"),
74+
):
75+
POD_type.to_numpy = podvector_to_numpy
76+
POD_type.to_cupy = podvector_to_cupy

‎src/amrex/StructOfArrays.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
"""
2+
This file is part of pyAMReX
3+
4+
Copyright 2023 AMReX community
5+
Authors: Axel Huebl
6+
License: BSD-3-Clause-LBNL
7+
"""
8+
from collections import namedtuple
9+
10+
11+
def soa_to_numpy(self, copy=False):
12+
"""
13+
Provide Numpy views into a StructOfArrays.
14+
15+
Parameters
16+
----------
17+
self : amrex.StructOfArrays_*
18+
A StructOfArrays class in pyAMReX
19+
copy : bool, optional
20+
Copy the data if true, otherwise create a view (default).
21+
22+
Returns
23+
-------
24+
namedtuple
25+
A tuple with real and int components that are each lists
26+
of 1D numpy arrays.
27+
"""
28+
import numpy as np
29+
30+
SoA_np = namedtuple(type(self).__name__ + "_np", ["real", "int"])
31+
32+
soa_view = SoA_np([], [])
33+
34+
if self.size() == 0:
35+
raise ValueError("SoA is empty.")
36+
37+
for idx_real in range(self.NumRealComps()):
38+
soa_view.real.append(np.array(self.GetRealData(idx_real), copy=copy))
39+
40+
for idx_int in range(self.NumIntComps()):
41+
soa_view.int.append(np.array(self.GetIntData(idx_int), copy=copy))
42+
43+
return soa_view
44+
45+
46+
def soa_to_cupy(self, copy=False):
47+
"""
48+
Provide Cupy views into a StructOfArrays.
49+
50+
Parameters
51+
----------
52+
self : amrex.StructOfArrays_*
53+
A StructOfArrays class in pyAMReX
54+
copy : bool, optional
55+
Copy the data if true, otherwise create a view (default).
56+
57+
Returns
58+
-------
59+
namedtuple
60+
A tuple with real and int components that are each lists
61+
of 1D numpy arrays.
62+
63+
Raises
64+
------
65+
ImportError
66+
Raises an exception if cupy is not installed
67+
"""
68+
import cupy as cp
69+
70+
SoA_cp = namedtuple(type(self).__name__ + "_cp", ["real", "int"])
71+
72+
soa_view = SoA_cp([], [])
73+
74+
if self.size() == 0:
75+
raise ValueError("SoA is empty.")
76+
77+
for idx_real in range(self.NumRealComps()):
78+
soa_view.real.append(cp.array(self.GetRealData(idx_real), copy=copy))
79+
80+
for idx_int in range(self.NumIntComps()):
81+
soa_view.int.append(cp.array(self.GetIntData(idx_int), copy=copy))
82+
83+
return soa_view
84+
85+
86+
def register_SoA_extension(amr):
87+
"""StructOfArrays helper methods"""
88+
import inspect
89+
import sys
90+
91+
# register member functions for every StructOfArrays_* type
92+
for _, SoA_type in inspect.getmembers(
93+
sys.modules[amr.__name__],
94+
lambda member: inspect.isclass(member)
95+
and member.__module__ == amr.__name__
96+
and member.__name__.startswith("StructOfArrays_"),
97+
):
98+
SoA_type.to_numpy = soa_to_numpy
99+
SoA_type.to_cupy = soa_to_cupy

‎src/amrex/space1d/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
# in pure Python or add some other Python logic
3232
#
3333
def d_decl(x, y, z):
34+
"""Return a tuple of the first passed element"""
3435
return (x,)
3536

3637

@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
4142
print(*args, **kwargs)
4243
elif ParallelDescriptor.IOProcessor():
4344
print(*args, **kwargs)
45+
46+
47+
from ..Array4 import register_Array4_extension
48+
from ..PODVector import register_PODVector_extension
49+
from ..StructOfArrays import register_SoA_extension
50+
51+
register_Array4_extension(amrex_1d_pybind)
52+
register_PODVector_extension(amrex_1d_pybind)
53+
register_SoA_extension(amrex_1d_pybind)

‎src/amrex/space2d/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
# in pure Python or add some other Python logic
3232
#
3333
def d_decl(x, y, z):
34+
"""Return a tuple of the first two passed elements"""
3435
return (x, y)
3536

3637

@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
4142
print(*args, **kwargs)
4243
elif ParallelDescriptor.IOProcessor():
4344
print(*args, **kwargs)
45+
46+
47+
from ..Array4 import register_Array4_extension
48+
from ..PODVector import register_PODVector_extension
49+
from ..StructOfArrays import register_SoA_extension
50+
51+
register_Array4_extension(amrex_2d_pybind)
52+
register_PODVector_extension(amrex_2d_pybind)
53+
register_SoA_extension(amrex_2d_pybind)

‎src/amrex/space3d/__init__.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
# in pure Python or add some other Python logic
3232
#
3333
def d_decl(x, y, z):
34+
"""Return a tuple of the three passed elements"""
3435
return (x, y, z)
3536

3637

@@ -41,3 +42,12 @@ def Print(*args, **kwargs):
4142
print(*args, **kwargs)
4243
elif ParallelDescriptor.IOProcessor():
4344
print(*args, **kwargs)
45+
46+
47+
from ..Array4 import register_Array4_extension
48+
from ..PODVector import register_PODVector_extension
49+
from ..StructOfArrays import register_SoA_extension
50+
51+
register_Array4_extension(amrex_3d_pybind)
52+
register_PODVector_extension(amrex_3d_pybind)
53+
register_SoA_extension(amrex_3d_pybind)

0 commit comments

Comments
 (0)