translate_to_mpi_operator.py 1.13 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
# -*- coding: utf-8 -*-

import numpy as np

from nifty.keepers import global_configuration as gc,\
                          global_dependency_injector as gdi

MPI = gdi[gc['mpi_module']]

custom_NANMIN = MPI.Op.Create(lambda x, y, datatype:
csongor's avatar
csongor committed
11
                              np.nanmin(np.vstack((x, y)), axis=0))
12
13

custom_NANMAX = MPI.Op.Create(lambda x, y, datatype:
csongor's avatar
csongor committed
14
                              np.nanmax(np.vstack((x, y)), axis=0))
15

16
17
18
custom_UNIQUE = MPI.Op.Create(lambda x, y, datatype:
                              np.unique(np.concatenate([x, y])))

19
20
21
22
23
24
25
26
op_translate_dict = {}

# the value tuple contains the operator and a boolean which specifies
# 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)
csongor's avatar
csongor committed
27
28
op_translate_dict[np.all] = (MPI.BAND, True)
op_translate_dict[np.any] = (MPI.BOR, True)
29
30
op_translate_dict[np.nanmin] = (custom_NANMIN, False)
op_translate_dict[np.nanmax] = (custom_NANMAX, False)
31
op_translate_dict[np.unique] = (custom_UNIQUE, False)