11"""Useful array functions.
22
33"""
4- from numba import jit , types as nbtypes , c16 , f8 , i8 , float64 , int64 , int32 , int8
4+ # import numba
55import numpy
6- import warnings
7- from numba import NumbaPerformanceWarning
8- warnings .simplefilter ('ignore' , NumbaPerformanceWarning )
96
107
118# @numba.jit([numba.types.Tuple((numba.float64[:], numba.float64[:]))
@@ -49,40 +46,6 @@ def average_chunks_jit(arr, wts, chunksize):
4946 return chunks , weights
5047
5148
52- @jit ([
53- nbtypes .Tuple ((c16 [:], f8 [:]))(c16 [:], f8 [:], i8 ),
54- ], nopython = True )
55- def average_chunks_complex (arr , wts , chunksize ):
56- """ Average the array arr with weights by chunks
57-
58- Array len does not have to be multiple of chunksize
59- The function is used for processing arr with complex data type.
60-
61- :param arr: 1D array of values
62- :param wts: 1D array of weights
63- :param chunksize: averaging size
64- :return: 1D array of averaged data_models, 1d array of weights
65- """
66- if chunksize <= 1 :
67- return arr , wts
68-
69- mask = numpy .zeros (((len (arr ) - 1 ) // chunksize + 1 , arr .shape [0 ]), dtype = arr .dtype )
70- maskwts = numpy .zeros_like (mask , dtype = wts .dtype )
71- for enumerate_id , i in enumerate (range (0 , len (arr ), chunksize )):
72- mask [enumerate_id , i :i + chunksize ] = 1 + 1j
73- maskwts [enumerate_id , i :i + chunksize ] = 1
74- tmp = wts * arr
75- chunks = mask .dot (tmp )
76- weights = maskwts .dot (wts )
77- chunks [weights > 0.0 ] = chunks [weights > 0.0 ] / weights [weights > 0.0 ]
78-
79- return chunks , weights
80-
81- @jit ([nbtypes .Tuple ((float64 [:], float64 [:]))(float64 [:], float64 [:], int64 ),
82- nbtypes .Tuple ((f8 [:],float64 [:]))(f8 [:], float64 [:], int64 ),
83- nbtypes .Tuple ((f8 [:],float64 [:]))(f8 [:], float64 [:], int32 ),
84- nbtypes .Tuple ((f8 [:],float64 [:]))(f8 [:], float64 [:], int8 )
85- ], nopython = True )
8649def average_chunks (arr , wts , chunksize ):
8750 """ Average the array arr with weights by chunks
8851
@@ -108,24 +71,13 @@ def average_chunks(arr, wts, chunksize):
10871
10972 # Codes optimized
11073
111- # mask = numpy.zeros(((len(arr)-1)//chunksize + 1, arr.shape[0]), dtype=bool)
112- # for enumerate_id, i in enumerate(range(0, len(arr), chunksize)):
113- # mask[enumerate_id,i:i+chunksize]= 1
114- # chunks = mask.dot(wts*arr)
115- # weights = mask.dot(wts)
116- # # chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
117- # numpy.putmask(chunks, weights>0.0, chunks/weights)
118-
119- # Numba
120- mask = numpy .zeros (((len (arr )- 1 )// chunksize + 1 , arr .shape [0 ]), dtype = arr .dtype )
121- maskwts = numpy .zeros_like (mask ,dtype = wts .dtype )
122- for enumerate_id , i in enumerate (range (0 , len (arr ), chunksize )):
123- mask [enumerate_id ,i :i + chunksize ]= 1
124- maskwts [enumerate_id ,i :i + chunksize ]= 1
125- tmp = wts * arr
126- chunks = mask .dot (tmp )
127- weights = maskwts .dot (wts )
128- chunks [weights > 0.0 ] = chunks [weights > 0.0 ] / weights [weights > 0.0 ]
74+ mask = numpy .zeros (((len (arr )- 1 )// chunksize + 1 , arr .shape [0 ]), dtype = bool )
75+ for enumerate_id ,i in enumerate (range (0 , len (arr ), chunksize )):
76+ mask [enumerate_id ,i :i + chunksize ]= 1
77+ chunks = mask .dot (wts * arr )
78+ weights = mask .dot (wts )
79+ # chunks[weights > 0.0] = chunks[weights > 0.0] / weights[weights > 0.0]
80+ numpy .putmask (chunks , weights > 0.0 , chunks / weights )
12981
13082 return chunks , weights
13183
@@ -144,66 +96,25 @@ def average_chunks2(arr, wts, chunksize):
14496 # assert arr.shape == wts.shape, "Shapes of arrays must be the same"
14597 # It is possible that there is a dangling null axis on wts
14698 wts = wts .reshape (arr .shape )
147- #
148- # For numba
149- if arr .dtype == 'c16' :
150- l0 = len (average_chunks_complex (arr [:, 0 ].flatten (), wts [:, 0 ].flatten (), chunksize [0 ])[0 ])
151- l1 = len (average_chunks_complex (arr [0 , :].flatten (), wts [0 , :].flatten (), chunksize [1 ])[0 ])
152-
153- tempchunks = numpy .zeros ([arr .shape [0 ], l1 ], dtype = arr .dtype )
154- tempwt = numpy .zeros ([arr .shape [0 ], l1 ])
155-
156- tempchunks *= tempwt
157- for i in range (arr .shape [0 ]):
158- result = average_chunks_complex (arr [i , :], wts [i , :], chunksize [1 ])
159- tempchunks [i , :], tempwt [i , :] = result [0 ].flatten (), result [1 ].flatten ()
160-
161- chunks = numpy .zeros ([l0 , l1 ], dtype = arr .dtype )
162- weights = numpy .zeros ([l0 , l1 ])
163-
164- for i in range (l1 ):
165- result = average_chunks_complex (tempchunks [:, i ], tempwt [:, i ], chunksize [0 ])
166- chunks [:, i ], weights [:, i ] = result [0 ].flatten (), result [1 ].flatten ()
167-
168- else :
169- l0 = len (average_chunks (arr [:, 0 ].flatten (), wts [:, 0 ].flatten (), chunksize [0 ])[0 ])
170- l1 = len (average_chunks (arr [0 , :].flatten (), wts [0 , :].flatten (), chunksize [1 ])[0 ])
171-
172- tempchunks = numpy .zeros ([arr .shape [0 ], l1 ], dtype = arr .dtype )
173- tempwt = numpy .zeros ([arr .shape [0 ], l1 ])
174-
175- tempchunks *= tempwt
176- for i in range (arr .shape [0 ]):
177- result = average_chunks (arr [i , :], wts [i , :], chunksize [1 ])
178- tempchunks [i , :], tempwt [i , :] = result [0 ].flatten (), result [1 ].flatten ()
179-
180- chunks = numpy .zeros ([l0 , l1 ], dtype = arr .dtype )
181- weights = numpy .zeros ([l0 , l1 ])
182-
183- for i in range (l1 ):
184- result = average_chunks (tempchunks [:, i ], tempwt [:, i ], chunksize [0 ])
185- chunks [:, i ], weights [:, i ] = result [0 ].flatten (), result [1 ].flatten ()
186-
187-
188- # Original
189- # l0 = len(average_chunks(arr[:, 0].flatten(), wts[:, 0].flatten(), chunksize[0])[0])
190- # l1 = len(average_chunks(arr[0, :].flatten(), wts[0, :].flatten(), chunksize[1])[0])
191- #
192- # tempchunks = numpy.zeros([arr.shape[0], l1], dtype=arr.dtype)
193- # tempwt = numpy.zeros([arr.shape[0], l1])
194- #
195- # tempchunks *= tempwt
196- # for i in range(arr.shape[0]):
197- # result = average_chunks(arr[i, :], wts[i, :], chunksize[1])
198- # tempchunks[i, :], tempwt[i, :] = result[0].flatten(), result[1].flatten()
199- #
200- # chunks = numpy.zeros([l0, l1], dtype=arr.dtype)
201- # weights = numpy.zeros([l0, l1])
202- #
203- # for i in range(l1):
204- # result = average_chunks(tempchunks[:, i], tempwt[:, i], chunksize[0])
205- # chunks[:, i], weights[:, i] = result[0].flatten(), result[1].flatten()
206-
99+
100+ l0 = len (average_chunks (arr [:, 0 ], wts [:, 0 ], chunksize [0 ])[0 ])
101+ l1 = len (average_chunks (arr [0 , :], wts [0 , :], chunksize [1 ])[0 ])
102+
103+ tempchunks = numpy .zeros ([arr .shape [0 ], l1 ], dtype = arr .dtype )
104+ tempwt = numpy .zeros ([arr .shape [0 ], l1 ])
105+
106+ tempchunks *= tempwt
107+ for i in range (arr .shape [0 ]):
108+ result = average_chunks (arr [i , :], wts [i , :], chunksize [1 ])
109+ tempchunks [i , :], tempwt [i , :] = result [0 ].flatten (), result [1 ].flatten ()
110+
111+ chunks = numpy .zeros ([l0 , l1 ], dtype = arr .dtype )
112+ weights = numpy .zeros ([l0 , l1 ])
113+
114+ for i in range (l1 ):
115+ result = average_chunks (tempchunks [:, i ], tempwt [:, i ], chunksize [0 ])
116+ chunks [:, i ], weights [:, i ] = result [0 ].flatten (), result [1 ].flatten ()
117+
207118 return chunks , weights
208119
209120
0 commit comments