forked from tritemio/PyBroMo
-
Notifications
You must be signed in to change notification settings - Fork 3
/
iter_chunks.py
66 lines (54 loc) · 2.16 KB
/
iter_chunks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
#
# PyBroMo - A single molecule diffusion simulator in confocal geometry.
#
# Copyright (C) 2013-2015 Antonino Ingargiola tritemio@gmail.com
#
"""
PyBroMo - A single molecule diffusion simulator in confocal geometry.
Copyright (C) 2013-2014 Antonino Ingargiola tritemio@gmail.com
This module implements iterator functions to loop over arrays in chunks.
"""
from __future__ import print_function, absolute_import, division
from builtins import range, zip
import numpy as np
def iter_chunksize(num_samples, chunksize):
"""Iterator used to iterate in chunks over an array of size `num_samples`.
At each iteration returns `chunksize` except for the last iteration.
"""
last_chunksize = np.mod(num_samples, chunksize)
for i in range(int(num_samples / chunksize)):
yield chunksize
if last_chunksize > 0:
yield last_chunksize
def iter_chunk_slice(num_samples, chunksize):
"""Iterator used to iterate in chunks over an array of size `num_samples`.
At each iteration returns a slice of size `chunksize`. In the last
iteration the slice may be smaller.
"""
i = 0
for c_size in iter_chunksize(num_samples, chunksize):
yield slice(i, i + c_size)
i += c_size
def iter_chunk_index(num_samples, chunksize):
"""Iterator used to iterate in chunks over an array of size `num_samples`.
At each iteration returns a start and stop index for a slice of size
`chunksize`. In the last iteration the slice may be smaller.
"""
i = 0
for c_size in iter_chunksize(num_samples, chunksize):
yield i, i + c_size
i += c_size
def reduce_chunk(func, array):
"""Reduce with `func`, chunk by chunk, the passed pytable `array`.
"""
res = []
for slice in iter_chunk_slice(array.shape[-1], array.chunkshape[-1]):
res.append(func(array[..., slice]))
return func(res)
def map_chunk(func, array, out_array):
"""Map with `func`, chunk by chunk, the input pytable `array`.
The result is stored in the output pytable array `out_array`.
"""
for slice in iter_chunk_slice(array.shape[-1], array.chunkshape[-1]):
out_array.append(func(array[..., slice]))
return out_array