Commit c1cdc4fb authored by theos's avatar theos
Browse files

Added a custom MPI Operator for amin and amax in order to be able to deal with numpy arrays.

parent 6ccfcce7
Pipeline #2203 skipped
......@@ -7,6 +7,16 @@ from nifty.keepers import global_configuration as gc,\
MPI = gdi[gc['mpi_module']]
custom_MIN = MPI.Op.Create(lambda x, y, datatype:
np.amin(np.vstack((x, y)), axis=0)
if isinstance(x, np.ndarray) else
lambda x, y, d: MPI.MIN(x, y))
custom_MAX = MPI.Op.Create(lambda x, y, datatype:
np.amax(np.vstack((x, y)), axis=0)
if isinstance(x, np.ndarray) else
lambda x, y, d: MPI.MAX(x, y))
custom_NANMIN = MPI.Op.Create(lambda x, y, datatype:
np.nanmin(np.vstack((x, y)), axis=0))
......@@ -22,8 +32,8 @@ op_translate_dict = {}
# if the operator is compatible to buffers (for Allreduce instead of allreduce)
op_translate_dict[np.sum] = (MPI.SUM, True)
op_translate_dict[np.prod] = (MPI.PROD, True)
op_translate_dict[np.amin] = (MPI.MIN, True)
op_translate_dict[np.amax] = (MPI.MAX, True)
op_translate_dict[np.amin] = (custom_MIN, False)
op_translate_dict[np.amax] = (custom_MAX, False)
op_translate_dict[np.all] = (MPI.BAND, True)
op_translate_dict[np.any] = (MPI.BOR, True)
op_translate_dict[np.nanmin] = (custom_NANMIN, False)
......
......@@ -1777,8 +1777,8 @@ class Test_axis(unittest.TestCase):
itertools.product(['max', 'sum', 'prod', 'mean', 'var', 'std', 'all', 'any',
'min', 'amin', 'nanmin', 'argmin', 'max', 'amax',
'nanmax', 'argmax'],
all_datatypes[1:],
[(1,)],# (6, 6)],
all_datatypes[1:],
[(1,), (6, 6)],
all_distribution_strategies,
[None, 0, (1, ), (0, 1)]),
testcase_func_name=custom_name_func)
......@@ -1805,10 +1805,10 @@ class Test_axis(unittest.TestCase):
itertools.product(['sum', 'prod', 'mean', 'var', 'std', 'all', 'any',
'min', 'amin', 'nanmin', 'argmin', 'max', 'amax',
'nanmax', 'argmax'],
all_datatypes[1:],
all_datatypes[1:],
[(4, 2, 3)],
all_distribution_strategies,
[(0, 1), (1, 2), (0, 1, 2)]),
[(0, 1)]),#, (1, 2), (0, 1, 2)]),
testcase_func_name=custom_name_func)
def test_axis_with_operations_many_dimentions(self, function, dtype,
global_shape,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment