Skip to content

Mpi wrapper

DummyMPI

Dummy MPI implementation for single-process mode.

Source code in june/mpi_wrapper.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
class DummyMPI:
    """Dummy MPI implementation for single-process mode."""
    # Define constants that would normally come from MPI
    UINT32_T = 'UINT32_T'  # Just a placeholder
    ANY_SOURCE = -1

    class Status:
        """ """
        def __init__(self):
            pass

        def Get_source(self):
            """ """
            return 0

    class COMM_WORLD:
        """ """
        @staticmethod
        def Get_rank():
            """ """
            return 0

        @staticmethod
        def Get_size():
            """ """
            return 1

        @staticmethod
        def Barrier():
            """ """
            # No-op in single process mode
            pass

        @staticmethod
        def bcast(obj, root=0):
            """

            Args:
                obj: 
                root: (Default value = 0)

            """
            return obj

        @staticmethod
        def allgather(obj):
            """

            Args:
                obj: 

            """
            return [obj]

        @staticmethod
        def alltoall(obj):
            """

            Args:
                obj: 

            """
            # Return same length array
            return [0] * len(obj) if isinstance(obj, (list, tuple, np.ndarray)) else [0]

        @staticmethod
        def Alltoallv(sendbuf, recvbuf):
            """

            Args:
                sendbuf: 
                recvbuf: 

            """
            # In single process mode, this is a no-op
            # We would normally just return the buffer, but we need to handle
            # the complex MPI parameters
            return recvbuf

        @staticmethod
        def iprobe(source=0, tag=0, status=None):
            """

            Args:
                source: (Default value = 0)
                tag: (Default value = 0)
                status: (Default value = None)

            """
            # Always return no messages in non-MPI mode
            return False

        @staticmethod
        def send(obj, dest=0, tag=0):
            """

            Args:
                obj: 
                dest: (Default value = 0)
                tag: (Default value = 0)

            """
            # No-op in single process mode
            pass

        @staticmethod
        def recv(source=0, tag=0):
            """

            Args:
                source: (Default value = 0)
                tag: (Default value = 0)

            """
            # Return None in single process mode
            return None

COMM_WORLD

Source code in june/mpi_wrapper.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
class COMM_WORLD:
    """ """
    @staticmethod
    def Get_rank():
        """ """
        return 0

    @staticmethod
    def Get_size():
        """ """
        return 1

    @staticmethod
    def Barrier():
        """ """
        # No-op in single process mode
        pass

    @staticmethod
    def bcast(obj, root=0):
        """

        Args:
            obj: 
            root: (Default value = 0)

        """
        return obj

    @staticmethod
    def allgather(obj):
        """

        Args:
            obj: 

        """
        return [obj]

    @staticmethod
    def alltoall(obj):
        """

        Args:
            obj: 

        """
        # Return same length array
        return [0] * len(obj) if isinstance(obj, (list, tuple, np.ndarray)) else [0]

    @staticmethod
    def Alltoallv(sendbuf, recvbuf):
        """

        Args:
            sendbuf: 
            recvbuf: 

        """
        # In single process mode, this is a no-op
        # We would normally just return the buffer, but we need to handle
        # the complex MPI parameters
        return recvbuf

    @staticmethod
    def iprobe(source=0, tag=0, status=None):
        """

        Args:
            source: (Default value = 0)
            tag: (Default value = 0)
            status: (Default value = None)

        """
        # Always return no messages in non-MPI mode
        return False

    @staticmethod
    def send(obj, dest=0, tag=0):
        """

        Args:
            obj: 
            dest: (Default value = 0)
            tag: (Default value = 0)

        """
        # No-op in single process mode
        pass

    @staticmethod
    def recv(source=0, tag=0):
        """

        Args:
            source: (Default value = 0)
            tag: (Default value = 0)

        """
        # Return None in single process mode
        return None

Alltoallv(sendbuf, recvbuf) staticmethod

Parameters:

Name Type Description Default
sendbuf
required
recvbuf
required
Source code in june/mpi_wrapper.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@staticmethod
def Alltoallv(sendbuf, recvbuf):
    """

    Args:
        sendbuf: 
        recvbuf: 

    """
    # In single process mode, this is a no-op
    # We would normally just return the buffer, but we need to handle
    # the complex MPI parameters
    return recvbuf

Barrier() staticmethod

Source code in june/mpi_wrapper.py 
54
55
56
57
58
@staticmethod
def Barrier():
    """ """
    # No-op in single process mode
    pass

Get_rank() staticmethod

Source code in june/mpi_wrapper.py 
44
45
46
47
@staticmethod
def Get_rank():
    """ """
    return 0

Get_size() staticmethod

Source code in june/mpi_wrapper.py 
49
50
51
52
@staticmethod
def Get_size():
    """ """
    return 1

allgather(obj) staticmethod

Parameters:

Name Type Description Default
obj
required
Source code in june/mpi_wrapper.py 
71
72
73
74
75
76
77
78
79
@staticmethod
def allgather(obj):
    """

    Args:
        obj: 

    """
    return [obj]

alltoall(obj) staticmethod

Parameters:

Name Type Description Default
obj
required
Source code in june/mpi_wrapper.py 
81
82
83
84
85
86
87
88
89
90
@staticmethod
def alltoall(obj):
    """

    Args:
        obj: 

    """
    # Return same length array
    return [0] * len(obj) if isinstance(obj, (list, tuple, np.ndarray)) else [0]

bcast(obj, root=0) staticmethod

Parameters:

Name Type Description Default
obj
required
root

(Default value = 0)

 0
Source code in june/mpi_wrapper.py 
60
61
62
63
64
65
66
67
68
69
@staticmethod
def bcast(obj, root=0):
    """

    Args:
        obj: 
        root: (Default value = 0)

    """
    return obj

iprobe(source=0, tag=0, status=None) staticmethod

Parameters:

Name Type Description Default
source

(Default value = 0)

 0
tag

(Default value = 0)

 0
status

(Default value = None)

 None
Source code in june/mpi_wrapper.py 
106
107
108
109
110
111
112
113
114
115
116
117
@staticmethod
def iprobe(source=0, tag=0, status=None):
    """

    Args:
        source: (Default value = 0)
        tag: (Default value = 0)
        status: (Default value = None)

    """
    # Always return no messages in non-MPI mode
    return False

recv(source=0, tag=0) staticmethod

Parameters:

Name Type Description Default
source

(Default value = 0)

 0
tag

(Default value = 0)

 0
Source code in june/mpi_wrapper.py 
132
133
134
135
136
137
138
139
140
141
142
@staticmethod
def recv(source=0, tag=0):
    """

    Args:
        source: (Default value = 0)
        tag: (Default value = 0)

    """
    # Return None in single process mode
    return None

send(obj, dest=0, tag=0) staticmethod

Parameters:

Name Type Description Default
obj
required
dest

(Default value = 0)

 0
tag

(Default value = 0)

 0
Source code in june/mpi_wrapper.py 
119
120
121
122
123
124
125
126
127
128
129
130
@staticmethod
def send(obj, dest=0, tag=0):
    """

    Args:
        obj: 
        dest: (Default value = 0)
        tag: (Default value = 0)

    """
    # No-op in single process mode
    pass

Status

Source code in june/mpi_wrapper.py 
33
34
35
36
37
38
39
40
class Status:
    """ """
    def __init__(self):
        pass

    def Get_source(self):
        """ """
        return 0

Get_source()

Source code in june/mpi_wrapper.py 
38
39
40
def Get_source(self):
    """ """
    return 0

MovablePeople

Class for managing mobile people across domains. This version is MPI-aware but works in both MPI and non-MPI modes.

Source code in june/mpi_wrapper.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
class MovablePeople:
    """Class for managing mobile people across domains.
    This version is MPI-aware but works in both MPI and non-MPI modes.

    """
    def __init__(self):
        self.skinny_out = {}
        self.skinny_in = {}
        self.index = {}

    def count_total_people(self):
        """Calculate the total number of people in skinny_out across all domains.


        Returns:
            int: Total number of people scheduled for transfer

        """
        total_count = 0
        for domain_id in self.skinny_out:
            for group_spec in self.skinny_out[domain_id]:
                for group_id in self.skinny_out[domain_id][group_spec]:
                    for subgroup_type in self.skinny_out[domain_id][group_spec][group_id]:
                        total_count += len(self.skinny_out[domain_id][group_spec][group_id][subgroup_type])
        return total_count

    def count_people_by_domain(self):
        """Calculate the number of people going to each domain.


        Returns:
            dict: Dictionary mapping domain_id to number of people

        """
        counts_by_domain = {}
        for domain_id in self.skinny_out:
            domain_count = 0
            for group_spec in self.skinny_out[domain_id]:
                for group_id in self.skinny_out[domain_id][group_spec]:
                    for subgroup_type in self.skinny_out[domain_id][group_spec][group_id]:
                        domain_count += len(self.skinny_out[domain_id][group_spec][group_id][subgroup_type])
            counts_by_domain[domain_id] = domain_count
        return counts_by_domain

    def add_person(self, person, external_subgroup):
        """Add or update a person to the outward facing group

        Args:
            person: 
            external_subgroup: 

        """
        domain_id = external_subgroup.domain_id

        # Rest of implementation only matters in MPI mode
        if not mpi_available and domain_id != 0:
            return  # In non-MPI mode, we only care about domain 0

        group_spec = external_subgroup.spec
        group_id = external_subgroup.group_id
        subgroup_type = external_subgroup.subgroup_type

        # Initialise domain structures if needed
        if domain_id not in self.skinny_out:
            self.skinny_out[domain_id] = {}
        if group_spec not in self.skinny_out[domain_id]:
            self.skinny_out[domain_id][group_spec] = {}
        if group_id not in self.skinny_out[domain_id][group_spec]:
            self.skinny_out[domain_id][group_spec][group_id] = {}
        if subgroup_type not in self.skinny_out[domain_id][group_spec][group_id]:
            self.skinny_out[domain_id][group_spec][group_id][subgroup_type] = {}

        # Create the view array with person data
        if person.infected:
            view = [
                person.id,
                person.infection.transmission.probability,
                person.infection.infection_id(),
                False,
                np.array([], dtype=np.int64),
                np.array([], dtype=np.float64),
                mpi_rank,
                True,
            ]
        else:
            susceptibility_inf_ids, susceptibility_inf_suscs = person.immunity.serialise()
            view = [
                person.id,
                0.0,
                0,
                True,
                np.array(susceptibility_inf_ids, dtype=np.int64),
                np.array(susceptibility_inf_suscs, dtype=np.float64),
                mpi_rank,
                True,
            ]

        self.skinny_out[domain_id][group_spec][group_id][subgroup_type][person.id] = view

    def delete_person(self, person, external_subgroup):
        """Remove a person from the external subgroup

        Args:
            person: 
            external_subgroup: 

        """
        domain_id = external_subgroup.domain_id

        if not mpi_available and domain_id != 0:
            return 0  # In non-MPI mode, success by default

        group_spec = external_subgroup.spec
        group_id = external_subgroup.group_id
        subgroup_type = external_subgroup.subgroup_type

        try:
            # Remove from skinny_out
            del self.skinny_out[domain_id][group_spec][group_id][subgroup_type][person.id]
            return 0
        except KeyError:
            return 1

    def serialise(self, rank):
        """Serialise person data for MPI communication

        Args:
            rank: 

        """
        # In single-process mode, this is simplified
        if not mpi_available:
            return None, None, 0

        keys, data = [], []
        if rank not in self.skinny_out:
            return None, None, 0

        # Serialise person data
        for group_spec in self.skinny_out[rank]:
            for group_id in self.skinny_out[rank][group_spec]:
                for subgroup_type in self.skinny_out[rank][group_spec][group_id]:
                    keys.append((
                        group_spec,
                        group_id,
                        subgroup_type,
                        len(self.skinny_out[rank][group_spec][group_id][subgroup_type])
                    ))
                    data += [
                        view
                        for pid, view in self.skinny_out[rank][group_spec][group_id][
                            subgroup_type
                        ].items()
                    ]

        outbound = np.array(data, dtype=object)

        return keys, outbound, outbound.shape[0]

    def update(self, rank, keys, rank_data):
        """Update with person registration in the destination rank

        Args:
            rank: 
            keys: 
            rank_data: 

        """
        # In single-process mode, this is simplified
        if not mpi_available:
            return

        index = 0

        for key in keys:
            group_spec, group_id, subgroup_type, n_data = key
            if group_spec not in self.skinny_in:
                self.skinny_in[group_spec] = {}
            if group_id not in self.skinny_in[group_spec]:
                self.skinny_in[group_spec][group_id] = {}
            if subgroup_type not in self.skinny_in[group_spec][group_id]:
                self.skinny_in[group_spec][group_id][subgroup_type] = {}

            data = rank_data[index : index + n_data]
            index += n_data

            try:
                # Update skinny_in with person data
                for k, i, t, s, iids, iis, d, a in data:
                    person_id = int(k)
                    # Register person in Person._persons dictionary
                    from june.demography import Person
                    if person_id not in Person._persons:
                        # Create minimal person instance
                        person = Person(
                            id=person_id,
                        )
                        person._home_rank = d
                        person._current_rank = mpi_rank
                        # This will automatically register in Person._persons
                    else:
                        person = Person._persons[person_id]
                        # Update person's rank since they've moved
                        person._current_rank = mpi_rank

                    self.skinny_in[group_spec][group_id][subgroup_type][person_id] = {
                        "inf_prob": i,
                        "inf_id": t,
                        "susc": s,
                        "immunity_inf_ids": iids,
                        "immunity_suscs": iis,
                        "dom": d,
                        "active": a,
                    }
            except Exception:
                print("failing", rank, "f-done")
                raise

add_person(person, external_subgroup)

Add or update a person to the outward facing group

Parameters:

Name Type Description Default
person
required
external_subgroup
required
Source code in june/mpi_wrapper.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def add_person(self, person, external_subgroup):
    """Add or update a person to the outward facing group

    Args:
        person: 
        external_subgroup: 

    """
    domain_id = external_subgroup.domain_id

    # Rest of implementation only matters in MPI mode
    if not mpi_available and domain_id != 0:
        return  # In non-MPI mode, we only care about domain 0

    group_spec = external_subgroup.spec
    group_id = external_subgroup.group_id
    subgroup_type = external_subgroup.subgroup_type

    # Initialise domain structures if needed
    if domain_id not in self.skinny_out:
        self.skinny_out[domain_id] = {}
    if group_spec not in self.skinny_out[domain_id]:
        self.skinny_out[domain_id][group_spec] = {}
    if group_id not in self.skinny_out[domain_id][group_spec]:
        self.skinny_out[domain_id][group_spec][group_id] = {}
    if subgroup_type not in self.skinny_out[domain_id][group_spec][group_id]:
        self.skinny_out[domain_id][group_spec][group_id][subgroup_type] = {}

    # Create the view array with person data
    if person.infected:
        view = [
            person.id,
            person.infection.transmission.probability,
            person.infection.infection_id(),
            False,
            np.array([], dtype=np.int64),
            np.array([], dtype=np.float64),
            mpi_rank,
            True,
        ]
    else:
        susceptibility_inf_ids, susceptibility_inf_suscs = person.immunity.serialise()
        view = [
            person.id,
            0.0,
            0,
            True,
            np.array(susceptibility_inf_ids, dtype=np.int64),
            np.array(susceptibility_inf_suscs, dtype=np.float64),
            mpi_rank,
            True,
        ]

    self.skinny_out[domain_id][group_spec][group_id][subgroup_type][person.id] = view

count_people_by_domain()

Calculate the number of people going to each domain.

Returns:

Name Type Description
dict

Dictionary mapping domain_id to number of people
Source code in june/mpi_wrapper.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
def count_people_by_domain(self):
    """Calculate the number of people going to each domain.


    Returns:
        dict: Dictionary mapping domain_id to number of people

    """
    counts_by_domain = {}
    for domain_id in self.skinny_out:
        domain_count = 0
        for group_spec in self.skinny_out[domain_id]:
            for group_id in self.skinny_out[domain_id][group_spec]:
                for subgroup_type in self.skinny_out[domain_id][group_spec][group_id]:
                    domain_count += len(self.skinny_out[domain_id][group_spec][group_id][subgroup_type])
        counts_by_domain[domain_id] = domain_count
    return counts_by_domain

count_total_people()

Calculate the total number of people in skinny_out across all domains.

Returns:

Name Type Description
int

Total number of people scheduled for transfer
Source code in june/mpi_wrapper.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def count_total_people(self):
    """Calculate the total number of people in skinny_out across all domains.


    Returns:
        int: Total number of people scheduled for transfer

    """
    total_count = 0
    for domain_id in self.skinny_out:
        for group_spec in self.skinny_out[domain_id]:
            for group_id in self.skinny_out[domain_id][group_spec]:
                for subgroup_type in self.skinny_out[domain_id][group_spec][group_id]:
                    total_count += len(self.skinny_out[domain_id][group_spec][group_id][subgroup_type])
    return total_count

delete_person(person, external_subgroup)

Remove a person from the external subgroup

Parameters:

Name Type Description Default
person
required
external_subgroup
required
Source code in june/mpi_wrapper.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
def delete_person(self, person, external_subgroup):
    """Remove a person from the external subgroup

    Args:
        person: 
        external_subgroup: 

    """
    domain_id = external_subgroup.domain_id

    if not mpi_available and domain_id != 0:
        return 0  # In non-MPI mode, success by default

    group_spec = external_subgroup.spec
    group_id = external_subgroup.group_id
    subgroup_type = external_subgroup.subgroup_type

    try:
        # Remove from skinny_out
        del self.skinny_out[domain_id][group_spec][group_id][subgroup_type][person.id]
        return 0
    except KeyError:
        return 1

serialise(rank)

Serialise person data for MPI communication

Parameters:

Name Type Description Default
rank
required
Source code in june/mpi_wrapper.py 
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def serialise(self, rank):
    """Serialise person data for MPI communication

    Args:
        rank: 

    """
    # In single-process mode, this is simplified
    if not mpi_available:
        return None, None, 0

    keys, data = [], []
    if rank not in self.skinny_out:
        return None, None, 0

    # Serialise person data
    for group_spec in self.skinny_out[rank]:
        for group_id in self.skinny_out[rank][group_spec]:
            for subgroup_type in self.skinny_out[rank][group_spec][group_id]:
                keys.append((
                    group_spec,
                    group_id,
                    subgroup_type,
                    len(self.skinny_out[rank][group_spec][group_id][subgroup_type])
                ))
                data += [
                    view
                    for pid, view in self.skinny_out[rank][group_spec][group_id][
                        subgroup_type
                    ].items()
                ]

    outbound = np.array(data, dtype=object)

    return keys, outbound, outbound.shape[0]

update(rank, keys, rank_data)

Update with person registration in the destination rank

Parameters:

Name Type Description Default
rank
required
keys
required
rank_data
required
Source code in june/mpi_wrapper.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
def update(self, rank, keys, rank_data):
    """Update with person registration in the destination rank

    Args:
        rank: 
        keys: 
        rank_data: 

    """
    # In single-process mode, this is simplified
    if not mpi_available:
        return

    index = 0

    for key in keys:
        group_spec, group_id, subgroup_type, n_data = key
        if group_spec not in self.skinny_in:
            self.skinny_in[group_spec] = {}
        if group_id not in self.skinny_in[group_spec]:
            self.skinny_in[group_spec][group_id] = {}
        if subgroup_type not in self.skinny_in[group_spec][group_id]:
            self.skinny_in[group_spec][group_id][subgroup_type] = {}

        data = rank_data[index : index + n_data]
        index += n_data

        try:
            # Update skinny_in with person data
            for k, i, t, s, iids, iis, d, a in data:
                person_id = int(k)
                # Register person in Person._persons dictionary
                from june.demography import Person
                if person_id not in Person._persons:
                    # Create minimal person instance
                    person = Person(
                        id=person_id,
                    )
                    person._home_rank = d
                    person._current_rank = mpi_rank
                    # This will automatically register in Person._persons
                else:
                    person = Person._persons[person_id]
                    # Update person's rank since they've moved
                    person._current_rank = mpi_rank

                self.skinny_in[group_spec][group_id][subgroup_type][person_id] = {
                    "inf_prob": i,
                    "inf_id": t,
                    "susc": s,
                    "immunity_inf_ids": iids,
                    "immunity_suscs": iis,
                    "dom": d,
                    "active": a,
                }
        except Exception:
            print("failing", rank, "f-done")
            raise

move_info(info2move)

Move information between processes in MPI mode, or simply return the information as-is in non-MPI mode.

Parameters:

Name Type Description Default
info2move list

Information to move between processes

 required

Returns:

Name Type Description
tuple

(buffer, n_sending, n_receiving)
Source code in june/mpi_wrapper.py 
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
def move_info(info2move):
    """Move information between processes in MPI mode, or simply return the
    information as-is in non-MPI mode.

    Args:
        info2move (list): Information to move between processes

    Returns:
        tuple: (buffer, n_sending, n_receiving)

    """
    if not mpi_available:
        # In non-MPI mode, just flatten the list and return it
        buffer = np.concatenate(info2move)
        return buffer, len(buffer), len(buffer)

    # In MPI mode, do the actual movement
    assert len(info2move) == mpi_size
    buffer = np.concatenate(info2move)
    assert buffer.dtype == np.uint32

    n_sending = len(buffer)
    count = np.array([len(x) for x in info2move])
    displ = np.array([sum(count[:p]) for p in range(len(info2move))])

    values = mpi_comm.alltoall(count)
    n_receiving = sum(values)
    r_buffer = np.zeros(n_receiving, dtype=np.uint32)
    rdisp = np.array([sum(values[:p]) for p in range(len(values))])

    mpi_comm.Alltoallv(
        [buffer, count, displ, MPI.UINT32_T],
        [r_buffer, values, rdisp, MPI.UINT32_T]
    )

    return r_buffer, n_sending, n_receiving