| 1 | import sys |
|---|
| 2 | import numpy |
|---|
| 3 | |
|---|
| 4 | |
|---|
| 5 | def build_matrix(data, qx_data, qy_data): |
|---|
| 6 | """ |
|---|
| 7 | Build a matrix for 2d plot from a vector |
|---|
| 8 | Returns a matrix (image) with ~ square binning |
|---|
| 9 | Requirement: need 1d array formats of |
|---|
| 10 | data, qx_data, and qy_data |
|---|
| 11 | where each one corresponds to z, x, or y axis values |
|---|
| 12 | |
|---|
| 13 | """ |
|---|
| 14 | # No qx or qy given in a vector format |
|---|
| 15 | if qx_data == None or qy_data == None \ |
|---|
| 16 | or qx_data.ndim != 1 or qy_data.ndim != 1: |
|---|
| 17 | return data |
|---|
| 18 | |
|---|
| 19 | # maximum # of loops to fillup_pixels |
|---|
| 20 | # otherwise, loop could never stop depending on data |
|---|
| 21 | max_loop = 1 |
|---|
| 22 | # get the x and y_bin arrays. |
|---|
| 23 | x_bins, y_bins = get_bins(qx_data, qy_data) |
|---|
| 24 | # set zero to None |
|---|
| 25 | |
|---|
| 26 | #Note: Can not use scipy.interpolate.Rbf: |
|---|
| 27 | # 'cause too many data points (>10000)<=JHC. |
|---|
| 28 | # 1d array to use for weighting the data point averaging |
|---|
| 29 | #when they fall into a same bin. |
|---|
| 30 | weights_data = numpy.ones([data.size]) |
|---|
| 31 | # get histogram of ones w/len(data); this will provide |
|---|
| 32 | #the weights of data on each bins |
|---|
| 33 | weights, xedges, yedges = numpy.histogram2d(x=qy_data, |
|---|
| 34 | y=qx_data, |
|---|
| 35 | bins=[y_bins, x_bins], |
|---|
| 36 | weights=weights_data) |
|---|
| 37 | # get histogram of data, all points into a bin in a way of summing |
|---|
| 38 | image, xedges, yedges = numpy.histogram2d(x=qy_data, |
|---|
| 39 | y=qx_data, |
|---|
| 40 | bins=[y_bins, x_bins], |
|---|
| 41 | weights=data) |
|---|
| 42 | # Now, normalize the image by weights only for weights>1: |
|---|
| 43 | # If weight == 1, there is only one data point in the bin so |
|---|
| 44 | # that no normalization is required. |
|---|
| 45 | image[weights > 1] = image[weights > 1] / weights[weights > 1] |
|---|
| 46 | # Set image bins w/o a data point (weight==0) as None (was set to zero |
|---|
| 47 | # by histogram2d.) |
|---|
| 48 | image[weights == 0] = None |
|---|
| 49 | |
|---|
| 50 | # Fill empty bins with 8 nearest neighbors only when at least |
|---|
| 51 | #one None point exists |
|---|
| 52 | loop = 0 |
|---|
| 53 | |
|---|
| 54 | # do while loop until all vacant bins are filled up up |
|---|
| 55 | #to loop = max_loop |
|---|
| 56 | while not(numpy.isfinite(image[weights == 0])).all(): |
|---|
| 57 | if loop >= max_loop: # this protects never-ending loop |
|---|
| 58 | break |
|---|
| 59 | image = fillup_pixels(image=image, weights=weights) |
|---|
| 60 | loop += 1 |
|---|
| 61 | |
|---|
| 62 | return image |
|---|
| 63 | |
|---|
| 64 | def get_bins(qx_data, qy_data): |
|---|
| 65 | """ |
|---|
| 66 | get bins |
|---|
| 67 | return x_bins and y_bins: 1d arrays of the index with |
|---|
| 68 | ~ square binning |
|---|
| 69 | Requirement: need 1d array formats of |
|---|
| 70 | qx_data, and qy_data |
|---|
| 71 | where each one corresponds to x, or y axis values |
|---|
| 72 | """ |
|---|
| 73 | # No qx or qy given in a vector format |
|---|
| 74 | if qx_data == None or qy_data == None \ |
|---|
| 75 | or qx_data.ndim != 1 or qy_data.ndim != 1: |
|---|
| 76 | return data |
|---|
| 77 | |
|---|
| 78 | # find max and min values of qx and qy |
|---|
| 79 | xmax = qx_data.max() |
|---|
| 80 | xmin = qx_data.min() |
|---|
| 81 | ymax = qy_data.max() |
|---|
| 82 | ymin = qy_data.min() |
|---|
| 83 | |
|---|
| 84 | # calculate the range of qx and qy: this way, it is a little |
|---|
| 85 | # more independent |
|---|
| 86 | x_size = xmax - xmin |
|---|
| 87 | y_size = ymax - ymin |
|---|
| 88 | |
|---|
| 89 | # estimate the # of pixels on each axes |
|---|
| 90 | npix_y = int(numpy.floor(numpy.sqrt(len(qy_data)))) |
|---|
| 91 | npix_x = int(numpy.floor(len(qy_data) / npix_y)) |
|---|
| 92 | |
|---|
| 93 | # bin size: x- & y-directions |
|---|
| 94 | xstep = x_size / (npix_x - 1) |
|---|
| 95 | ystep = y_size / (npix_y - 1) |
|---|
| 96 | |
|---|
| 97 | # max and min taking account of the bin sizes |
|---|
| 98 | xmax = xmax + xstep / 2.0 |
|---|
| 99 | xmin = xmin - xstep / 2.0 |
|---|
| 100 | ymax = ymax + ystep / 2.0 |
|---|
| 101 | ymin = ymin - ystep / 2.0 |
|---|
| 102 | |
|---|
| 103 | # store x and y bin centers in q space |
|---|
| 104 | x_bins = numpy.linspace(xmin, xmax, npix_x) |
|---|
| 105 | y_bins = numpy.linspace(ymin, ymax, npix_y) |
|---|
| 106 | |
|---|
| 107 | #set x_bins and y_bins |
|---|
| 108 | return x_bins, y_bins |
|---|
| 109 | |
|---|
| 110 | def fillup_pixels(image=None, weights=None): |
|---|
| 111 | """ |
|---|
| 112 | Fill z values of the empty cells of 2d image matrix |
|---|
| 113 | with the average over up-to next nearest neighbor points |
|---|
| 114 | |
|---|
| 115 | :param image: (2d matrix with some zi = None) |
|---|
| 116 | |
|---|
| 117 | :return: image (2d array ) |
|---|
| 118 | |
|---|
| 119 | :TODO: Find better way to do for-loop below |
|---|
| 120 | |
|---|
| 121 | """ |
|---|
| 122 | # No image matrix given |
|---|
| 123 | if image == None or numpy.ndim(image) != 2 \ |
|---|
| 124 | or numpy.isfinite(image).all() \ |
|---|
| 125 | or weights == None: |
|---|
| 126 | return image |
|---|
| 127 | # Get bin size in y and x directions |
|---|
| 128 | len_y = len(image) |
|---|
| 129 | len_x = len(image[1]) |
|---|
| 130 | temp_image = numpy.zeros([len_y, len_x]) |
|---|
| 131 | weit = numpy.zeros([len_y, len_x]) |
|---|
| 132 | # do for-loop for all pixels |
|---|
| 133 | for n_y in range(len(image)): |
|---|
| 134 | for n_x in range(len(image[1])): |
|---|
| 135 | # find only null pixels |
|---|
| 136 | if weights[n_y][n_x] > 0 or numpy.isfinite(image[n_y][n_x]): |
|---|
| 137 | continue |
|---|
| 138 | else: |
|---|
| 139 | # find 4 nearest neighbors |
|---|
| 140 | # check where or not it is at the corner |
|---|
| 141 | if n_y != 0 and numpy.isfinite(image[n_y - 1][n_x]): |
|---|
| 142 | temp_image[n_y][n_x] += image[n_y - 1][n_x] |
|---|
| 143 | weit[n_y][n_x] += 1 |
|---|
| 144 | if n_x != 0 and numpy.isfinite(image[n_y][n_x - 1]): |
|---|
| 145 | temp_image[n_y][n_x] += image[n_y][n_x - 1] |
|---|
| 146 | weit[n_y][n_x] += 1 |
|---|
| 147 | if n_y != len_y - 1 and numpy.isfinite(image[n_y + 1][n_x]): |
|---|
| 148 | temp_image[n_y][n_x] += image[n_y + 1][n_x] |
|---|
| 149 | weit[n_y][n_x] += 1 |
|---|
| 150 | if n_x != len_x - 1 and numpy.isfinite(image[n_y][n_x + 1]): |
|---|
| 151 | temp_image[n_y][n_x] += image[n_y][n_x + 1] |
|---|
| 152 | weit[n_y][n_x] += 1 |
|---|
| 153 | # go 4 next nearest neighbors when no non-zero |
|---|
| 154 | # neighbor exists |
|---|
| 155 | if n_y != 0 and n_x != 0 and\ |
|---|
| 156 | numpy.isfinite(image[n_y - 1][n_x - 1]): |
|---|
| 157 | temp_image[n_y][n_x] += image[n_y - 1][n_x - 1] |
|---|
| 158 | weit[n_y][n_x] += 1 |
|---|
| 159 | if n_y != len_y - 1 and n_x != 0 and \ |
|---|
| 160 | numpy.isfinite(image[n_y + 1][n_x - 1]): |
|---|
| 161 | temp_image[n_y][n_x] += image[n_y + 1][n_x - 1] |
|---|
| 162 | weit[n_y][n_x] += 1 |
|---|
| 163 | if n_y != len_y and n_x != len_x - 1 and \ |
|---|
| 164 | numpy.isfinite(image[n_y - 1][n_x + 1]): |
|---|
| 165 | temp_image[n_y][n_x] += image[n_y - 1][n_x + 1] |
|---|
| 166 | weit[n_y][n_x] += 1 |
|---|
| 167 | if n_y != len_y - 1 and n_x != len_x - 1 and \ |
|---|
| 168 | numpy.isfinite(image[n_y + 1][n_x + 1]): |
|---|
| 169 | temp_image[n_y][n_x] += image[n_y + 1][n_x + 1] |
|---|
| 170 | weit[n_y][n_x] += 1 |
|---|
| 171 | |
|---|
| 172 | # get it normalized |
|---|
| 173 | ind = (weit > 0) |
|---|
| 174 | image[ind] = temp_image[ind] / weit[ind] |
|---|
| 175 | |
|---|
| 176 | return image |
|---|
