-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathbins_3d.py
175 lines (148 loc) · 6.2 KB
/
bins_3d.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
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
import numpy as np
from flux_accumulator import AccumulatorBase
class Bins3D(AccumulatorBase):
def __init__(self, dims, ranges, ar_existing_data=None, filename=''):
"""
:param dims: the shape of the bins
:param ranges: a 2-by-3 array representing the minimum and maximum ranges of the 3 coordinates
:type dims: np.ndarray
:type ranges np.ndarray
:type ar_existing_data np.ndarray
:type filename str
"""
if ar_existing_data is not None:
expected_shape = (dims[0], dims[1], dims[2], 3)
ravelled_shape = (dims[0] * dims[1] * dims[2] * 3,)
if ar_existing_data.shape == ravelled_shape:
self.ar_data = ar_existing_data.reshape(expected_shape)
else:
assert ar_existing_data.shape == expected_shape, "incompatible shape:{0}".format(
ar_existing_data.shape)
self.ar_data = ar_existing_data
else:
self.ar_data = np.zeros((dims[0], dims[1], dims[2], 3))
self.ar_flux = None
self.ar_weights = None
self.ar_count = None
self.update_array_slices()
self.dims = dims
self.index_type = ''
self.update_index_type()
self.filename = filename
self.ranges = ranges
self.bin_sizes = np.abs(ranges[1] - ranges[0]) / dims
def add_array_with_mask(self, ar_flux, ar_x, ar_y, ar_z, mask, ar_weights):
"""
add flux to x,y bins with weights and a filter mask.
note: the x,y values should be rescaled prior to calling this method.
:type ar_flux: np.multiarray.ndarray
:param ar_x: np.multiarray.ndarray
:param ar_y: np.multiarray.ndarray
:param ar_z: np.multiarray.ndarray
:param mask: np.multiarray.ndarray
:param ar_weights: np.multiarray.ndarray
"""
ar_x_int = ((ar_x - self.ranges[0, 0]) / self.bin_sizes[0]).astype(self.index_type)
ar_y_int = ((ar_y - self.ranges[0, 1]) / self.bin_sizes[1]).astype(self.index_type)
ar_z_int = ((ar_z - self.ranges[0, 2]) / self.bin_sizes[2]).astype(self.index_type)
# restrict the mask to pixels inside the bin range.
m = np.all((0 <= ar_x_int, ar_x_int < self.dims[0],
0 <= ar_y_int, ar_y_int < self.dims[1],
0 <= ar_z_int, ar_z_int < self.dims[2],
mask), axis=0)
ar_flux_masked = ar_flux[m]
ar_weights_masked = ar_weights[m]
ar_indices_x = ar_x_int[m]
ar_indices_y = ar_y_int[m]
ar_indices_z = ar_z_int[m]
# make sure we don't invert x, y and z
# z is the innermost coordinate, x is the outermost.
# represent bins in 1D. this is faster than a 2D numpy histogram
ar_indices_xyz = ar_indices_z + self.dims[2] * (ar_indices_y + (self.dims[1] * ar_indices_x))
# bin data according to x,y values
flux_hist_1d = np.bincount(ar_indices_xyz, weights=ar_flux_masked,
minlength=int(np.prod(self.dims)))
weights_hist_1d = np.bincount(ar_indices_xyz, weights=ar_weights_masked,
minlength=int(np.prod(self.dims)))
count_hist_1d = np.bincount(ar_indices_xyz, weights=None,
minlength=int(np.prod(self.dims)))
# return from 1D to a 2d array
flux_hist = flux_hist_1d.reshape(self.dims)
count_hist = count_hist_1d.reshape(self.dims)
weights_hist = weights_hist_1d.reshape(self.dims)
# accumulate new data
self.ar_flux += flux_hist
self.ar_weights += weights_hist
self.ar_count += count_hist
def merge(self, bins_3d_2):
assert self.ar_data.shape == bins_3d_2.ar_data.shape
assert np.all(self.ranges == bins_3d_2.ranges)
assert np.all(self.dims == bins_3d_2.dims)
self.ar_data += bins_3d_2.ar_data
return self
def merge_array(self, ar_data):
assert self.ar_data.shape == ar_data.shape
self.ar_data += ar_data
return self
def save(self, filename):
self.filename = filename
self.flush()
def from_4d_array(self, stacked_array):
self.ar_data = stacked_array
self.update_array_slices()
self.dims = self.ar_count.shape
self.update_index_type()
def load(self, filename):
# TODO: to static
stacked_array = np.load(filename)
self.from_4d_array(stacked_array)
def update_index_type(self):
# choose integer type according to number of bins
self.index_type = 'int32' if np.prod(self.dims) > 32767 else 'int16'
def update_array_slices(self):
self.ar_flux = self.ar_data[:, :, :, 0]
self.ar_count = self.ar_data[:, :, :, 1]
self.ar_weights = self.ar_data[:, :, :, 2]
def __radd__(self, other):
return self.merge(other)
def __add__(self, other):
return self.merge(other)
@classmethod
def init_as(cls, other):
"""
:type other: Bins3D
"""
new_obj = cls(other.dims, other.ranges, filename=other.filename)
return new_obj
@classmethod
def from_other(cls, other):
new_obj = cls.init_as(other)
new_obj.merge(other)
def set_filename(self, filename):
self.filename = filename
def to_4d_array(self):
return self.ar_data
def flush(self):
np.save(self.filename, self.to_4d_array())
def get_ranges(self):
return self.ranges
def get_dims(self):
return self.dims
def get_pair_count(self):
return self.ar_count.sum()
def get_data_as_array(self):
return self.to_4d_array()
def get_array_shape(self):
return self.ar_data.shape
def get_metadata(self):
return [self.dims, self.index_type,
self.filename,
self.ranges, self.bin_sizes]
@classmethod
def load_from(cls, ar, metadata):
new_bins = cls(dims=(1, 1, 1), ranges=((1, 1, 1), (1, 1, 1)))
(new_bins.dims, new_bins.index_type, new_bins.filename,
new_bins.ranges, new_bins.bin_sizes) = metadata
new_bins.ar_data = ar
new_bins.update_array_slices()
return new_bins