| 1 | import sys |
|---|
| 2 | import numpy |
|---|
| 3 | import string |
|---|
| 4 | |
|---|
| 5 | from collections import OrderedDict |
|---|
| 6 | |
|---|
| 7 | # MPL shapes dictionary with some extra styles rendered internally. |
|---|
| 8 | # Ordered for consistent display in combo boxes |
|---|
| 9 | SHAPES = OrderedDict([ |
|---|
| 10 | ('Circle' , 'o'), |
|---|
| 11 | ('Point' , '.'), |
|---|
| 12 | ('Pixel' , ','), |
|---|
| 13 | ('Triangle Down' , 'v'), |
|---|
| 14 | ('Triangle Up' , '^'), |
|---|
| 15 | ('Triangle Left' , '<'), |
|---|
| 16 | ('Triangle Right' , '>'), |
|---|
| 17 | ('Octagon' , '8'), |
|---|
| 18 | ('Square' , 's'), |
|---|
| 19 | ('Pentagon' , 'p'), |
|---|
| 20 | ('Star' , '*'), |
|---|
| 21 | ('Hexagon1' , 'h'), |
|---|
| 22 | ('Hexagon2' , 'H'), |
|---|
| 23 | ('Cross +' , 'p'), |
|---|
| 24 | ('Cross X ' , 'x'), |
|---|
| 25 | ('Diamond' , 'D'), |
|---|
| 26 | ('Thin Diamond' , 'd'), |
|---|
| 27 | ('Line' , '-'), |
|---|
| 28 | ('Dash' , '--'), |
|---|
| 29 | ('Vline' , 'vline'), |
|---|
| 30 | ('Step' , 'step'), |
|---|
| 31 | ]) |
|---|
| 32 | |
|---|
| 33 | # MPL Colors dictionary. Ordered for consistent display |
|---|
| 34 | COLORS = OrderedDict([ |
|---|
| 35 | ('Blue', 'b'), |
|---|
| 36 | ('Green', 'g'), |
|---|
| 37 | ('Red', 'r'), |
|---|
| 38 | ('Cyan', 'c'), |
|---|
| 39 | ('Magenta', 'm'), |
|---|
| 40 | ('Yellow', 'y'), |
|---|
| 41 | ('Black', 'k'), |
|---|
| 42 | ('Custom', 'x'), |
|---|
| 43 | ]) |
|---|
| 44 | |
|---|
| 45 | def build_matrix(data, qx_data, qy_data): |
|---|
| 46 | """ |
|---|
| 47 | Build a matrix for 2d plot from a vector |
|---|
| 48 | Returns a matrix (image) with ~ square binning |
|---|
| 49 | Requirement: need 1d array formats of |
|---|
| 50 | data, qx_data, and qy_data |
|---|
| 51 | where each one corresponds to z, x, or y axis values |
|---|
| 52 | |
|---|
| 53 | """ |
|---|
| 54 | # No qx or qy given in a vector format |
|---|
| 55 | if qx_data is None or qy_data is None \ |
|---|
| 56 | or qx_data.ndim != 1 or qy_data.ndim != 1: |
|---|
| 57 | return data |
|---|
| 58 | |
|---|
| 59 | # maximum # of loops to fillup_pixels |
|---|
| 60 | # otherwise, loop could never stop depending on data |
|---|
| 61 | max_loop = 1 |
|---|
| 62 | # get the x and y_bin arrays. |
|---|
| 63 | x_bins, y_bins = get_bins(qx_data, qy_data) |
|---|
| 64 | # set zero to None |
|---|
| 65 | |
|---|
| 66 | #Note: Can not use scipy.interpolate.Rbf: |
|---|
| 67 | # 'cause too many data points (>10000)<=JHC. |
|---|
| 68 | # 1d array to use for weighting the data point averaging |
|---|
| 69 | #when they fall into a same bin. |
|---|
| 70 | weights_data = numpy.ones([data.size]) |
|---|
| 71 | # get histogram of ones w/len(data); this will provide |
|---|
| 72 | #the weights of data on each bins |
|---|
| 73 | weights, xedges, yedges = numpy.histogram2d(x=qy_data, |
|---|
| 74 | y=qx_data, |
|---|
| 75 | bins=[y_bins, x_bins], |
|---|
| 76 | weights=weights_data) |
|---|
| 77 | # get histogram of data, all points into a bin in a way of summing |
|---|
| 78 | image, xedges, yedges = numpy.histogram2d(x=qy_data, |
|---|
| 79 | y=qx_data, |
|---|
| 80 | bins=[y_bins, x_bins], |
|---|
| 81 | weights=data) |
|---|
| 82 | # Now, normalize the image by weights only for weights>1: |
|---|
| 83 | # If weight == 1, there is only one data point in the bin so |
|---|
| 84 | # that no normalization is required. |
|---|
| 85 | image[weights > 1] = image[weights > 1] / weights[weights > 1] |
|---|
| 86 | # Set image bins w/o a data point (weight==0) as None (was set to zero |
|---|
| 87 | # by histogram2d.) |
|---|
| 88 | image[weights == 0] = None |
|---|
| 89 | |
|---|
| 90 | # Fill empty bins with 8 nearest neighbors only when at least |
|---|
| 91 | #one None point exists |
|---|
| 92 | loop = 0 |
|---|
| 93 | |
|---|
| 94 | # do while loop until all vacant bins are filled up up |
|---|
| 95 | #to loop = max_loop |
|---|
| 96 | while not(numpy.isfinite(image[weights == 0])).all(): |
|---|
| 97 | if loop >= max_loop: # this protects never-ending loop |
|---|
| 98 | break |
|---|
| 99 | image = fillupPixels(image=image, weights=weights) |
|---|
| 100 | loop += 1 |
|---|
| 101 | |
|---|
| 102 | return image |
|---|
| 103 | |
|---|
| 104 | def get_bins(qx_data, qy_data): |
|---|
| 105 | """ |
|---|
| 106 | get bins |
|---|
| 107 | return x_bins and y_bins: 1d arrays of the index with |
|---|
| 108 | ~ square binning |
|---|
| 109 | Requirement: need 1d array formats of |
|---|
| 110 | qx_data, and qy_data |
|---|
| 111 | where each one corresponds to x, or y axis values |
|---|
| 112 | """ |
|---|
| 113 | # No qx or qy given in a vector format |
|---|
| 114 | if qx_data is None or qy_data is None \ |
|---|
| 115 | or qx_data.ndim != 1 or qy_data.ndim != 1: |
|---|
| 116 | return data |
|---|
| 117 | |
|---|
| 118 | # find max and min values of qx and qy |
|---|
| 119 | xmax = qx_data.max() |
|---|
| 120 | xmin = qx_data.min() |
|---|
| 121 | ymax = qy_data.max() |
|---|
| 122 | ymin = qy_data.min() |
|---|
| 123 | |
|---|
| 124 | # calculate the range of qx and qy: this way, it is a little |
|---|
| 125 | # more independent |
|---|
| 126 | x_size = xmax - xmin |
|---|
| 127 | y_size = ymax - ymin |
|---|
| 128 | |
|---|
| 129 | # estimate the # of pixels on each axes |
|---|
| 130 | npix_y = int(numpy.floor(numpy.sqrt(len(qy_data)))) |
|---|
| 131 | npix_x = int(numpy.floor(len(qy_data) / npix_y)) |
|---|
| 132 | |
|---|
| 133 | # bin size: x- & y-directions |
|---|
| 134 | xstep = x_size / (npix_x - 1) |
|---|
| 135 | ystep = y_size / (npix_y - 1) |
|---|
| 136 | |
|---|
| 137 | # max and min taking account of the bin sizes |
|---|
| 138 | xmax = xmax + xstep / 2.0 |
|---|
| 139 | xmin = xmin - xstep / 2.0 |
|---|
| 140 | ymax = ymax + ystep / 2.0 |
|---|
| 141 | ymin = ymin - ystep / 2.0 |
|---|
| 142 | |
|---|
| 143 | # store x and y bin centers in q space |
|---|
| 144 | x_bins = numpy.linspace(xmin, xmax, npix_x) |
|---|
| 145 | y_bins = numpy.linspace(ymin, ymax, npix_y) |
|---|
| 146 | |
|---|
| 147 | #set x_bins and y_bins |
|---|
| 148 | return x_bins, y_bins |
|---|
| 149 | |
|---|
| 150 | def fillupPixels(image=None, weights=None): |
|---|
| 151 | """ |
|---|
| 152 | Fill z values of the empty cells of 2d image matrix |
|---|
| 153 | with the average over up-to next nearest neighbor points |
|---|
| 154 | |
|---|
| 155 | :param image: (2d matrix with some zi = None) |
|---|
| 156 | |
|---|
| 157 | :return: image (2d array ) |
|---|
| 158 | |
|---|
| 159 | :TODO: Find better way to do for-loop below |
|---|
| 160 | |
|---|
| 161 | """ |
|---|
| 162 | # No image matrix given |
|---|
| 163 | if image == None or numpy.ndim(image) != 2 \ |
|---|
| 164 | or numpy.isfinite(image).all() \ |
|---|
| 165 | or weights == None: |
|---|
| 166 | return image |
|---|
| 167 | # Get bin size in y and x directions |
|---|
| 168 | len_y = len(image) |
|---|
| 169 | len_x = len(image[1]) |
|---|
| 170 | temp_image = numpy.zeros([len_y, len_x]) |
|---|
| 171 | weit = numpy.zeros([len_y, len_x]) |
|---|
| 172 | # do for-loop for all pixels |
|---|
| 173 | for n_y in range(len(image)): |
|---|
| 174 | for n_x in range(len(image[1])): |
|---|
| 175 | # find only null pixels |
|---|
| 176 | if weights[n_y][n_x] > 0 or numpy.isfinite(image[n_y][n_x]): |
|---|
| 177 | continue |
|---|
| 178 | else: |
|---|
| 179 | # find 4 nearest neighbors |
|---|
| 180 | # check where or not it is at the corner |
|---|
| 181 | if n_y != 0 and numpy.isfinite(image[n_y - 1][n_x]): |
|---|
| 182 | temp_image[n_y][n_x] += image[n_y - 1][n_x] |
|---|
| 183 | weit[n_y][n_x] += 1 |
|---|
| 184 | if n_x != 0 and numpy.isfinite(image[n_y][n_x - 1]): |
|---|
| 185 | temp_image[n_y][n_x] += image[n_y][n_x - 1] |
|---|
| 186 | weit[n_y][n_x] += 1 |
|---|
| 187 | if n_y != len_y - 1 and numpy.isfinite(image[n_y + 1][n_x]): |
|---|
| 188 | temp_image[n_y][n_x] += image[n_y + 1][n_x] |
|---|
| 189 | weit[n_y][n_x] += 1 |
|---|
| 190 | if n_x != len_x - 1 and numpy.isfinite(image[n_y][n_x + 1]): |
|---|
| 191 | temp_image[n_y][n_x] += image[n_y][n_x + 1] |
|---|
| 192 | weit[n_y][n_x] += 1 |
|---|
| 193 | # go 4 next nearest neighbors when no non-zero |
|---|
| 194 | # neighbor exists |
|---|
| 195 | if n_y != 0 and n_x != 0 and\ |
|---|
| 196 | numpy.isfinite(image[n_y - 1][n_x - 1]): |
|---|
| 197 | temp_image[n_y][n_x] += image[n_y - 1][n_x - 1] |
|---|
| 198 | weit[n_y][n_x] += 1 |
|---|
| 199 | if n_y != len_y - 1 and n_x != 0 and \ |
|---|
| 200 | numpy.isfinite(image[n_y + 1][n_x - 1]): |
|---|
| 201 | temp_image[n_y][n_x] += image[n_y + 1][n_x - 1] |
|---|
| 202 | weit[n_y][n_x] += 1 |
|---|
| 203 | if n_y != len_y and n_x != len_x - 1 and \ |
|---|
| 204 | numpy.isfinite(image[n_y - 1][n_x + 1]): |
|---|
| 205 | temp_image[n_y][n_x] += image[n_y - 1][n_x + 1] |
|---|
| 206 | weit[n_y][n_x] += 1 |
|---|
| 207 | if n_y != len_y - 1 and n_x != len_x - 1 and \ |
|---|
| 208 | numpy.isfinite(image[n_y + 1][n_x + 1]): |
|---|
| 209 | temp_image[n_y][n_x] += image[n_y + 1][n_x + 1] |
|---|
| 210 | weit[n_y][n_x] += 1 |
|---|
| 211 | |
|---|
| 212 | # get it normalized |
|---|
| 213 | ind = (weit > 0) |
|---|
| 214 | image[ind] = temp_image[ind] / weit[ind] |
|---|
| 215 | |
|---|
| 216 | return image |
|---|
| 217 | |
|---|
| 218 | def rescale(lo, hi, step, pt=None, bal=None, scale='linear'): |
|---|
| 219 | """ |
|---|
| 220 | Rescale (lo,hi) by step, returning the new (lo,hi) |
|---|
| 221 | The scaling is centered on pt, with positive values of step |
|---|
| 222 | driving lo/hi away from pt and negative values pulling them in. |
|---|
| 223 | If bal is given instead of point, it is already in [0,1] coordinates. |
|---|
| 224 | |
|---|
| 225 | This is a helper function for step-based zooming. |
|---|
| 226 | """ |
|---|
| 227 | # Convert values into the correct scale for a linear transformation |
|---|
| 228 | # TODO: use proper scale transformers |
|---|
| 229 | loprev = lo |
|---|
| 230 | hiprev = hi |
|---|
| 231 | if scale == 'log': |
|---|
| 232 | assert lo > 0 |
|---|
| 233 | if lo > 0: |
|---|
| 234 | lo = numpy.log10(lo) |
|---|
| 235 | if hi > 0: |
|---|
| 236 | hi = numpy.log10(hi) |
|---|
| 237 | if pt is not None: |
|---|
| 238 | pt = numpy.log10(pt) |
|---|
| 239 | |
|---|
| 240 | # Compute delta from axis range * %, or 1-% if persent is negative |
|---|
| 241 | if step > 0: |
|---|
| 242 | delta = float(hi - lo) * step / 100 |
|---|
| 243 | else: |
|---|
| 244 | delta = float(hi - lo) * step / (100 - step) |
|---|
| 245 | |
|---|
| 246 | # Add scale factor proportionally to the lo and hi values, |
|---|
| 247 | # preserving the |
|---|
| 248 | # point under the mouse |
|---|
| 249 | if bal is None: |
|---|
| 250 | bal = float(pt - lo) / (hi - lo) |
|---|
| 251 | lo = lo - (bal * delta) |
|---|
| 252 | hi = hi + (1 - bal) * delta |
|---|
| 253 | |
|---|
| 254 | # Convert transformed values back to the original scale |
|---|
| 255 | if scale == 'log': |
|---|
| 256 | if (lo <= -250) or (hi >= 250): |
|---|
| 257 | lo = loprev |
|---|
| 258 | hi = hiprev |
|---|
| 259 | else: |
|---|
| 260 | lo, hi = numpy.power(10., lo), numpy.power(10., hi) |
|---|
| 261 | return (lo, hi) |
|---|
| 262 | |
|---|
| 263 | def getValidColor(color): |
|---|
| 264 | ''' |
|---|
| 265 | Returns a valid matplotlib color |
|---|
| 266 | ''' |
|---|
| 267 | |
|---|
| 268 | if color is None: |
|---|
| 269 | return color |
|---|
| 270 | |
|---|
| 271 | # Check if it's an int |
|---|
| 272 | if isinstance(color, int): |
|---|
| 273 | # Check if it's within the range |
|---|
| 274 | if 0 <= color <=6: |
|---|
| 275 | color = list(COLORS.values())[color] |
|---|
| 276 | # Check if it's an RGB string |
|---|
| 277 | elif isinstance(color, str): |
|---|
| 278 | # Assure the correctnes of the string |
|---|
| 279 | assert(color[0]=="#" and len(color) == 7) |
|---|
| 280 | assert(all(c in string.hexdigits for c in color[1:])) |
|---|
| 281 | else: |
|---|
| 282 | raise AttributeError |
|---|
| 283 | |
|---|
| 284 | return color |
|---|
