Skip to content

Commit 3e3b1d4

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

File tree

10 files changed

+504
-12
lines changed

10 files changed

+504
-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/ArrayOfStructs.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
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 aos_to_numpy(self, copy=False):
12+
"""
13+
Provide Numpy views into a ArrayOfStructs.
14+
15+
Parameters
16+
----------
17+
self : amrex.ArrayOfStructs_*
18+
An ArrayOfStructs 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+
if self.size() == 0:
31+
raise ValueError("AoS is empty.")
32+
33+
return np.array(self, copy=copy)
34+
35+
36+
def aos_to_cupy(self, copy=False):
37+
"""
38+
Provide Cupy views into a ArrayOfStructs.
39+
40+
Parameters
41+
----------
42+
self : amrex.ArrayOfStructs_*
43+
An ArrayOfStructs class in pyAMReX
44+
copy : bool, optional
45+
Copy the data if true, otherwise create a view (default).
46+
47+
Returns
48+
-------
49+
namedtuple
50+
A tuple with real and int components that are each lists
51+
of 1D numpy arrays.
52+
53+
Raises
54+
------
55+
ImportError
56+
Raises an exception if cupy is not installed
57+
"""
58+
import cupy as cp
59+
60+
SoA_cp = namedtuple(type(self).__name__ + "_cp", ["real", "int"])
61+
62+
soa_view = SoA_cp([], [])
63+
64+
if self.size() == 0:
65+
raise ValueError("AoS is empty.")
66+
67+
return cp.array(self, copy=copy)
68+
69+
70+
def register_AoS_extension(amr):
71+
"""ArrayOfStructs helper methods"""
72+
import inspect
73+
import sys
74+
75+
# register member functions for every ArrayOfStructs_* type
76+
for _, AoS_type in inspect.getmembers(
77+
sys.modules[amr.__name__],
78+
lambda member: inspect.isclass(member)
79+
and member.__module__ == amr.__name__
80+
and member.__name__.startswith("ArrayOfStructs_"),
81+
):
82+
AoS_type.to_numpy = aos_to_numpy
83+
AoS_type.to_cupy = aos_to_cupy

‎src/amrex/MultiFab.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
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 mf_to_numpy(self, copy=False, order="F"):
11+
"""
12+
Provide a Numpy view into a MultiFab.
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.MultiFab
24+
A MultiFab 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+
list of np.array
33+
A list of numpy n-dimensional arrays, for each local block in the
34+
MultiFab.
35+
"""
36+
views = []
37+
for mfi in self:
38+
views.append(self.array(mfi).to_numpy(copy, order))
39+
40+
return views
41+
42+
43+
def mf_to_cupy(self, copy=False, order="F"):
44+
"""
45+
Provide a Cupy view into a MultiFab.
46+
47+
Note on the order of indices:
48+
By default, this is as in AMReX in Fortran contiguous order, indexing as
49+
x,y,z. This has performance implications for use in external libraries such
50+
as cupy.
51+
The order="C" option will index as z,y,x and perform better with cupy.
52+
https://github.com/AMReX-Codes/pyamrex/issues/55#issuecomment-1579610074
53+
54+
Parameters
55+
----------
56+
self : amrex.MultiFab
57+
A MultiFab class in pyAMReX
58+
copy : bool, optional
59+
Copy the data if true, otherwise create a view (default).
60+
order : string, optional
61+
F order (default) or C. C is faster with external libraries.
62+
63+
Returns
64+
-------
65+
list of cupy.array
66+
A list of cupy n-dimensional arrays, for each local block in the
67+
MultiFab.
68+
69+
Raises
70+
------
71+
ImportError
72+
Raises an exception if cupy is not installed
73+
"""
74+
views = []
75+
for mfi in self:
76+
views.append(self.array(mfi).to_cupy(copy, order))
77+
78+
return views
79+
80+
81+
def register_MultiFab_extension(amr):
82+
"""MultiFab helper methods"""
83+
import inspect
84+
import sys
85+
86+
# register member functions for every MultiFab* type
87+
for _, MultiFab_type in inspect.getmembers(
88+
sys.modules[amr.__name__],
89+
lambda member: inspect.isclass(member)
90+
and member.__module__ == amr.__name__
91+
and member.__name__.startswith("MultiFab"),
92+
):
93+
MultiFab_type.to_numpy = mf_to_numpy
94+
MultiFab_type.to_cupy = mf_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

0 commit comments

Comments
 (0)