def makeHeatContent(salt, temp, destMask, thetao, pressure): """ The makeHeatContent() function takes 3D (not temporal) arguments and creates heat content which is then mapped to a destination grid and written to a specified variable Author: Paul J. Durack : [email protected] : @durack1. Created on Tue Nov 24 15:34:30 2015. Inputs: ------ - salt(lev,lat,lon) - 3D array. - temp(lev,lat,lon) - 3D array either in-situ or potential temperature. - destGridArray(str) - 2D array with valid grid and mask. - thetao(bool) - boolean value specifying either in-situ or potential temperature arrays provided. - pressure(bool) - boolean value specifying whether lev-coordinate is pressure (dbar) or depth (m). Usage: ------ >>> from oceanLib import makeHeatContent >>> makeHeatContent(salt,temp,destGridArray,thetao=True,pressure=False) Notes: ----- - PJD 24 Nov 2015 - Migrated into new oceanLib from heatContentLib - TODO: Better deal with insitu vs thetao variables - TODO: """ # Remap variables to short names #print salt.getAxisIds() s = salt(squeeze=1) # Trim off singleton time dimension #print s.getAxisIds() t = temp(squeeze=1) mask = destMask #print mask.getAxisIds() del (salt, temp, destMask) gc.collect() depthInd = 0 # Set depth coordinate index #print 's: ',s.min(),s.max() #print 't: ',t.min(),t.max() # Fix out of bounds values t = mv.where(t < -2.6, -2.6, t) # Fix for NaN values # Calculate pressure - inputs depth & lat # Create z-coordinate from salinity input if not pressure: zCoord = s.getAxis(depthInd) # Assume time,depth,latitude,longitude grid yCoord = s.getAxis(depthInd + 1) yCoord = tile(yCoord, (s.shape[depthInd + 2], 1)).transpose() depthLevels = tile( zCoord.getValue(), (s.shape[depthInd + 2], s.shape[depthInd + 1], 1)).transpose() pressureLevels = sw.pres(np.array(depthLevels), np.array(yCoord)) del (zCoord, yCoord, depthLevels) gc.collect() else: pressureLevels = s.getAxis(depthInd) #print pressureLevels.getValue() pressureLevels = transpose( tile(pressureLevels, (s.shape[depthInd + 2], s.shape[depthInd + 1], 1))) pressureLevels = cdm.createVariable(pressureLevels, id='pressureLevels') pressureLevels.setAxis(0, s.getAxis(depthInd)) pressureLevels.setAxis(1, s.getAxis(depthInd + 1)) pressureLevels.setAxis(2, s.getAxis(depthInd + 2)) pressureLevels.units_long = 'decibar (pressure)' pressureLevels.positive = 'down' pressureLevels.long_name = 'sea_water_pressure' pressureLevels.standard_name = 'sea_water_pressure' pressureLevels.units = 'decibar' pressureLevels.axis = 'Z' #print 'pres: ',pressureLevels.min(),pressureLevels.max() #print pressureLevels.shape #print s.shape #print t.shape #print mask.shape # Calculate temp,rho,cp - inputs temp,salt,pressure if thetao: # Process potential temperature to in-situ temp = sw.temp(np.array(s), np.array(t), np.array(pressureLevels)) # units degrees C rho = sw.dens(np.array(s), np.array(temp), np.array(pressureLevels)) # units kg m-3 cp = sw.cp(np.array(s), np.array(temp), np.array(pressureLevels)) # units J kg-1 C-1 # Correct instances of NaN values and fix masks - applied before cdms variables are created otherwise names/ids/attributes are reset temp = scrubNaNAndMask(temp, s) rho = scrubNaNAndMask(rho, s) cp = scrubNaNAndMask(cp, s) #print 'temp: ',temp.min(),temp.max() #print 'rho: ',rho.min(),rho.max() #print 'cp: ',cp.min(),cp.max() # Calculate heatContent - inputs temp,rho,cp heatContent = np.array(temp) * np.array(rho) * np.array(cp) # units J # Correct instances of NaN values and fix masks - applied before cdms variables are created otherwise names/ids/attributes are reset heatContent = scrubNaNAndMask(heatContent, s) #print 'hc: ',heatContent.min(),heatContent.max() # Interpolate to standard levels - inputs heatContent,levels newDepth = np.array([ 5, 10, 20, 30, 40, 50, 75, 100, 125, 150, 200, 300, 500, 700, 1000, 1500, 1800, 2000 ]).astype('f') newDepth_bounds = np.array([[0, 5], [5, 10], [10, 20], [20, 30], [30, 40], [40, 50], [50, 75], [75, 100], [100, 125], [125, 150], [150, 200], [200, 300], [300, 500], [500, 700], [700, 1000], [1000, 1500], [1500, 1800], [1800, 2000]]).astype('f') # Interpolate to standard levels #print heatContent.shape #print heatContent.getAxisIds() #print pressureLevels.shape #print pressureLevels.getAxisIds() # Reset variable axes heatContent.setAxis(0, s.getAxis(0)) #heatContent.setAxis(1,s.getAxis(1)) #heatContent.setAxis(2,s.getAxis(2)) pdb.set_trace() heatContent.setGrid(s.getGrid()) #print heatContent.shape #print heatContent.getAxisIds() #print pressureLevels.shape #print pressureLevels.getAxisIds() heatContent_depthInterp = cdu.linearInterpolation(heatContent, pressureLevels, levels=newDepth) # Fix bounds newDepth = heatContent_depthInterp.getAxis(0) newDepth.setBounds(newDepth_bounds) del (newDepth_bounds) newDepth.id = 'depth2' newDepth.units_long = 'decibar (pressure)' newDepth.positive = 'down' newDepth.long_name = 'sea_water_pressure' newDepth.standard_name = 'sea_water_pressure' newDepth.units = 'decibar' newDepth.axis = 'Z' #print 'hc_interp:',heatContent_depthInterp.min(),heatContent_depthInterp.max() # Integrate to 700 dbar - inputs heatContent heatContent_depthInteg = cdu.averager( heatContent_depthInterp[0:14, ...], axis=0, weights='weighted', action='sum')(squeeze=1) # Calculate depth-weighted-integrated thetao # Assign all axis info #print heatContent_depthInteg.shape pdb.set_trace() # Interpolate in x,y - inputs heatContent #tmp1 = heatContent_depthInteg.regrid(mask.getGrid(),regridTool='esmf',regridMethod='linear') ; # Use defaults - ,coordSys='deg',diag = {},periodicity=1) tmp1 = heatContent_depthInteg.regrid(mask, regridTool='esmf', regridMethod='linear') # Use defaults - ,coordSys='deg',diag = {},periodicity=1) #print tmp1.shape tmp1 = mv.where(tmp1 < 0, 0, tmp1) # Fix for negative values # Infill - inputs heatContent # Create inputs for interpolation points = np.zeros([(mask.shape[0] * mask.shape[1]), 2]) # Create 25380 vectors of lon/lat latcounter = 0 loncounter = 0 for count, data in enumerate(points): if not np.mod(count, 180) and not count == 0: latcounter = latcounter + 1 loncounter = 0 points[count, 0] = mask.getLatitude().getValue()[latcounter] points[count, 1] = mask.getLongitude().getValue()[loncounter] loncounter = loncounter + 1 del (count, data, latcounter, loncounter) gc.collect() valid = np.logical_not(tmp1.mask) # Get inverted-logic boolean mask from variable if valid.size == 1: print '** No valid mask found, skipping **' return valid = valid.flatten() # Flatten 2D to 1D #maskFilled = mask(tmp,points,valid) interpolant = interpolate.LinearNDInterpolator( points[valid, :], np.array(tmp1.flatten())[valid]) # Create interpolant maskFill = interpolant(points[:, 0].squeeze(), points[:, 1].squeeze()) # Use interpolant to create filled matrix maskFill = np.reshape(maskFill, mask.shape) # Resize to original dimensions # Fix issues with interpolant tmp2 = mv.where(np.isnan(maskFill), 1e+20, maskFill) # Fix for NaN values tmp2 = mv.where(tmp2 > tmp1.max(), 0, tmp2) # Fix for max values tmp2 = mv.where(tmp2 < tmp1.min(), 0, tmp2) # Fix for min values tmp = mv.masked_where(mask.mask, tmp2) #print tmp.shape # Redress variable heatContent = cdm.createVariable([tmp], id='heatContent') depthInt = cdm.createAxis([350], id='depth') depthInt.setBounds(np.array([0, 700])) depthInt.units_long = 'decibar (pressure)' depthInt.positive = 'down' depthInt.long_name = 'sea_water_pressure' depthInt.standard_name = 'sea_water_pressure' depthInt.units = 'decibar' depthInt.axis = 'Z' heatContent.setAxis(0, depthInt) heatContent.setAxis(1, mask.getAxis(0)) heatContent.setAxis(2, mask.getAxis(1)) heatContent.units_long = 'Joules' heatContent.long_name = 'sea_water_heat_content' heatContent.standard_name = 'sea_water_heat_content' heatContent.units = 'J' return heatContent
def __init__(self, x_vector, y): # x_vector is of dimension (n_points, n_dims) super().__init__(x_vector, y) self.NDfunction = interpolate.LinearNDInterpolator(x_vector, y)
nodeDisplacementX = np.empty(nodes.shape[0]) nodeDisplacementY = np.empty(nodes.shape[0]) nodeDisplacementZ = np.empty(nodes.shape[0]) for inode in range(nodes.shape[0]): cord = np.array([nodes[inode][0],nodes[inode][1],nodes[inode][2],1]) new_cord = affineMatrix.dot(cord) nodeDisplacementX[inode] = new_cord[0]-nodes[inode][0] nodeDisplacementY[inode] = new_cord[1]-nodes[inode][1] nodeDisplacementZ[inode] = new_cord[2]-nodes[inode][2] newim3_map = ndimage.map_coordinates(im2_matched, [mapX, mapY, mapZ], order=1) surf = plot.surf_points inter_z = interpolate.LinearNDInterpolator(nodes, nodeDisplacementZ, fill_value=0.0, rescale=True) extrapo_z = inter_z.__call__(surf) plot.plot3dpoint(extrapo_z) ##### plot3 = Visualization.DataVisulization(ndimage.gaussian_filter(newim3_map,5), 110) plot3.contour3d() plot4 = Visualization.DataVisulization(ndimage.gaussian_filter(newim3,5), 110) plot4.contour3d() #src3 = mlab.pipeline.scalar_field(ndimage.gaussian_filter(newim3,3)) #mlab.pipeline.iso_surface(src3, contours=[70], colormap='Oranges') save = SaveLoad.saveDataBase("/Users/junchaowei/Desktop/SpaceRegistration_000_125/") datadict={}
def active_from_xyz(mesh, xyz, grid_reference='CC', method='linear'): """Returns an active cell index array below a surface Get active cells in the `mesh` that are below the surface create by interpolating over the last dimension of the input points. This method will uses scipy's interpolation routine to interpolate between input points. This will use a nearest neighbour interpolation for cell values outside the convex hull of points. For `grid_reference='CC'`, this will check if the center of a given cell in the mesh is below the interpolation surface to determine if that cell is active. For `grid_reference='N'`, this will check if **every** node of a given cell in the mesh is below the interpolation surface. For the Parameters ---------- mesh : TensorMesh or TreeMesh or CylMesh Mesh object, (if CylMesh: mesh must be symmetric). xyz : numpy.ndarray Points coordinates shape (*, mesh.dim). grid_reference : {'CC', 'N'} Use cell coordinates from cells-center 'CC' or nodes 'N'. method : {'linear', 'nearest'} Interpolation method for the xyz points. Returns ------- active : numpy.ndarray 1D mask array of `bool` for the active cells below xyz. Examples -------- .. plot:: :include-source: import matplotlib.pyplot as plt import numpy as np from discretize import TensorMesh from discretize.utils import active_from_xyz mesh = TensorMesh([5, 5]) topo_func = lambda x: -3*(x-0.2)*(x-0.8)+.5 topo_points = np.linspace(0, 1) topo_vals = topo_func(topo_points) active_cc = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='CC') ax = plt.subplot(121) mesh.plotImage(active_cc, ax=ax) mesh.plotGrid(centers=True, ax=ax) ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3') ax.set_title("CC") active_n = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='N') ax = plt.subplot(122) mesh.plotImage(active_n, ax=ax) mesh.plotGrid(nodes=True, ax=ax) ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3') ax.set_title("N") plt.show() """ if isinstance(mesh, discretize.CylMesh) and not mesh.isSymmetric: raise NotImplementedError('Unsymmetric CylMesh is not yet supported') if not isinstance(mesh, (discretize.TensorMesh, discretize.TreeMesh, discretize.CylMesh)): raise TypeError("Mesh must be either TensorMesh, TreeMesh, or Symmetric CylMesh") if grid_reference not in ["N", "CC"]: raise ValueError("Value of grid_reference must be 'N' (nodal) or 'CC' (cell center)") dim = mesh.dim - 1 if mesh.dim == 3: if xyz.shape[1] != 3: raise ValueError("xyz locations of shape (*, 3) required for 3D mesh") if method == 'linear': tri2D = Delaunay(xyz[:, :2]) z_interpolate = interpolate.LinearNDInterpolator(tri2D, xyz[:, 2]) else: z_interpolate = interpolate.NearestNDInterpolator(xyz[:, :2], xyz[:, 2]) elif mesh.dim == 2: if xyz.shape[1] != 2: raise ValueError("xyz locations of shape (*, 2) required for 2D mesh") z_interpolate = interpolate.interp1d( xyz[:, 0], xyz[:, 1], bounds_error=False, fill_value=np.nan, kind=method ) else: if xyz.ndim != 1: raise ValueError("xyz locations of shape (*, ) required for 1D mesh") if grid_reference == 'CC': locations = mesh.gridCC if mesh.dim == 1: active = np.zeros(mesh.nC, dtype='bool') active[np.searchsorted(mesh.vectorCCx, xyz).max():] = True return active elif grid_reference == 'N': if mesh.dim == 3: locations = np.vstack([ mesh.gridCC + (np.c_[-1, 1, 1][:, None] * mesh.h_gridded / 2.).squeeze(), mesh.gridCC + (np.c_[-1, -1, 1][:, None] * mesh.h_gridded / 2.).squeeze(), mesh.gridCC + (np.c_[1, 1, 1][:, None] * mesh.h_gridded / 2.).squeeze(), mesh.gridCC + (np.c_[1, -1, 1][:, None] * mesh.h_gridded / 2.).squeeze() ]) elif mesh.dim == 2: locations = np.vstack([ mesh.gridCC + (np.c_[-1, 1][:, None] * mesh.h_gridded / 2.).squeeze(), mesh.gridCC + (np.c_[1, 1][:, None] * mesh.h_gridded / 2.).squeeze(), ]) else: active = np.zeros(mesh.nC, dtype='bool') active[np.searchsorted(mesh.vectorNx, xyz).max():] = True return active # Interpolate z values on CC or N z_xyz = z_interpolate(locations[:, :-1]).squeeze() # Apply nearest neighbour if in extrapolation ind_nan = np.isnan(z_xyz) if any(ind_nan): tree = cKDTree(xyz) _, ind = tree.query(locations[ind_nan, :]) z_xyz[ind_nan] = xyz[ind, dim] # Create an active bool of all True active = np.all( (locations[:, dim] < z_xyz).reshape((mesh.nC, -1), order='F'), axis=1 ) return active.ravel()
#dispz = np.zeros(mesh2.points.T.shape[1]) colors = (dispz) * 0.0495 #colors = np.zeros(mesh2.points.T.shape[1]) #sio.savemat('/Users/junchaowei/Desktop/0_90_z.mat', {'p':colors}) #sio.savemat('/Users/junchaowei/Desktop/0_90_element.mat',{'p':mesh2.ntri}) @mlab.show def main(): mesh2.view(mesh2.polydata(mesh2.points.T, colors)) main() mayavi.mlab.colorbar(object=None, title="U(Z)(mm)", orientation="vertical", nb_labels=20, nb_colors=None, label_fmt=None) "colorbar reference: http://docs.enthought.com/mayavi/mayavi/auto/mlab_decorations.html" plot = Visualization.DataVisulization(image1, threshold1 + 5000) plot.contour3d() surf = plot.surf_points inter_z = interpolate.LinearNDInterpolator(mesh2.points, colors, fill_value=0.0, rescale=True) extrapo_z = inter_z.__call__(surf) plot.plot3dpoint(extrapo_z)
def _ip(self): """A `scipy.interpolate.LinearNDInterpolator` instance containing the learner's data.""" # XXX: take our own triangulation into account when generating the _ip return interpolate.LinearNDInterpolator(self.points, self.values)
def ax_scalp( v, channels, ax=None, annotate=False, vmin=None, vmax=None, cmap=cm.coolwarm, scalp_line_width=1, scalp_line_style="solid", chan_pos_list=CHANNEL_10_20_APPROX, interpolation="bilinear", fontsize=8, ): """Draw a scalp plot. Draws a scalp plot on an existing axes. The method takes an array of values and an array of the corresponding channel names. It matches the channel names with an channel position list to project them correctly on the scalp. Parameters ---------- v : 1d-array of floats The values for the channels channels : 1d array of strings The corresponding channel names for the values in ``v`` ax : Axes, optional The axes to draw the scalp plot on. If not provided, the currently activated axes (i.e. ``gca()``) will be taken annotate : Boolean, optional Draw the channel names next to the channel markers. vmin, vmax : float, optional The display limits for the values in ``v``. If the data in ``v`` contains values between -3..3 and ``vmin`` and ``vmax`` are set to -1 and 1, all values smaller than -1 and bigger than 1 will appear the same as -1 and 1. If not set, the maximum absolute value in ``v`` is taken to calculate both values. cmap : matplotlib.colors.colormap, optional A colormap to define the color transitions. scalp_line_width: float Line width for outline of scalp scalp_line_style: float Line style for outline of scalp chan_pos_list: iterable of tuples First entry should be 'angle' or 'cartesian', remaining entries 2-tuples of x and y. interpolation: str Returns ------- ax : Axes the axes on which the plot was drawn Notes ----- Code adapted from Wyrm [1]_ toolbox https://github.com/bbci/wyrm. References ---------- .. [1] Schirrmeister, R. T., Springenberg, J. T., Fiederer, L. D. J., Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F. & Ball, T. (2017). Deep learning with convolutional neural networks for EEG decoding and visualization. Human Brain Mapping , Aug. 2017. Online: http://dx.doi.org/10.1002/hbm.23730 """ if ax is None: ax = plt.gca() assert len(v) == len(channels), "Should be as many values as channels" assert interpolation == "bilinear" or interpolation == "nearest" if vmin is None: # added by me ([email protected]) assert vmax is None vmin, vmax = -np.max(np.abs(v)), np.max(np.abs(v)) # what if we have an unknown channel? points = [get_channelpos(c, chan_pos_list) for c in channels] for c in channels: assert get_channelpos( c, chan_pos_list) is not None, ("Expect " + c + " to exist in positions") z = [v[i] for i in range(len(points))] # calculate the interpolation x = [i[0] for i in points] y = [i[1] for i in points] # interpolate the in-between values xx = np.linspace(min(x), max(x), 500) yy = np.linspace(min(y), max(y), 500) if interpolation == "bilinear": xx_grid, yy_grid = np.meshgrid(xx, yy) f = interpolate.LinearNDInterpolator(list(zip(x, y)), z) zz = f(xx_grid, yy_grid) else: assert interpolation == "nearest" f = interpolate.NearestNDInterpolator(list(zip(x, y)), z) assert len(xx) == len(yy) zz = np.ones((len(xx), len(yy))) for i_x in range(len(xx)): for i_y in range(len(yy)): # somehow this is correct. don't know why :( zz[i_y, i_x] = f(xx[i_x], yy[i_y]) # zz[i_x,i_y] = f(xx[i_x], yy[i_y]) assert not np.any(np.isnan(zz)) # plot map image = ax.imshow( zz, vmin=vmin, vmax=vmax, cmap=cmap, extent=[min(x), max(x), min(y), max(y)], origin="lower", interpolation=interpolation, ) if scalp_line_width > 0: # paint the head ax.add_artist( plt.Circle( (0, 0), 1, linestyle=scalp_line_style, linewidth=scalp_line_width, fill=False, )) # add a nose ax.plot( [-0.1, 0, 0.1], [1, 1.1, 1], color="black", linewidth=scalp_line_width, linestyle=scalp_line_style, ) # add ears _add_ears(ax, scalp_line_width, scalp_line_style) # add markers at channels positions # set the axes limits, so the figure is centered on the scalp ax.set_ylim([-1.05, 1.15]) ax.set_xlim([-1.15, 1.15]) # hide the frame and ticks ax.set_frame_on(False) ax.set_xticks([]) ax.set_yticks([]) # draw the channel names if annotate: for i in zip(channels, list(zip(x, y))): ax.annotate( " " + i[0], i[1], horizontalalignment="center", verticalalignment="center", fontsize=fontsize, ) ax.set_aspect(1) return image
def create_model_maps( altitude, variable=None, model=None, file=None, numContours=25, fill=False, ct='viridis', # https://matplotlib.org/examples/color/colormaps_reference.html transparency=1, nearest=False, linear=True, saveFig=True): import matplotlib #matplotlib.use('tkagg') import matplotlib.pyplot as plt if model == None and file == None: print("Please input either a model or the file path/name to a model.") return if file != None: model = read_model_results(file) if (variable == None) or (variable not in model.keys()): print("Variable not entered or not found, Please select one from: ") index = 0 name_index_dict = {} for name in model: if name.lower() == 'dim': continue if name.lower() == 'meta': continue print(index, ": ", name) name_index_dict[index] = name index += 1 i_choice = int(input("Enter Selection: ")) dataname = name_index_dict[i_choice].lower() print(dataname) else: dataname = variable.lower() mars_radius = model['meta']['mars_radius'] lats = np.arange(181) - 90 lons = np.arange(361) - 180 sc_lat_mso = np.repeat(lats, len(lons)) sc_lon_mso = np.tile(lons, len(lats)) r = np.full(len(sc_lon_mso), altitude + mars_radius) sc_alt_array = np.full(len(sc_lon_mso), altitude) sc_mso_x = r * np.sin( (90 - sc_lat_mso) * (np.pi / 180)) * np.cos(sc_lon_mso * (np.pi / 180)) sc_mso_y = r * np.sin( (90 - sc_lat_mso) * (np.pi / 180)) * np.sin(sc_lon_mso * (np.pi / 180)) sc_mso_z = r * np.cos((90 - sc_lat_mso) * (np.pi / 180)) sc_path = np.array([sc_mso_x, sc_mso_y, sc_mso_z]).T if nearest: interp_method = 'nearest' else: interp_method = 'linear' if model == None and file == None: print( "Please input either a model dictionary from read_model_results, or a model file." ) return if 'lon' in model['dim']: if 'mso' == model['meta']['coord_sys'].lower(): #Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows lat_mso_model = model['dim']['lat'] lon_mso_model = model['dim']['lon'] alt_mso_model = model['dim']['alt'] lat_array = np.repeat(lat_mso_model, len(lon_mso_model)) lon_array = np.tile(lon_mso_model, len(lat_mso_model)) data_points = np.transpose(np.array([lon_array, lat_array])) index = 0 for point in sc_path: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) alt_mso = altitude lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180) sc_path[index] = np.array([lon_mso, lat_mso, alt_mso]) index += 1 latlon_triangulation = spatial.Delaunay(data_points) for var in model: if var.lower() != dataname: continue print("Interpolating variable " + var) if var == 'dim' or var == 'meta': continue #Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'longitude': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'latitude': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'altitude': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) #Build an array of values that correspond to the points in data point values = np.empty([ len(lat_mso_model) * len(lon_mso_model), len(alt_mso_model) ]) index = 0 for alt in range(0, len(alt_mso_model)): for lat in range(0, len(lat_mso_model)): for lon in range(0, len(lon_mso_model)): values[index, alt] = data_new[lon][lat][alt] index += 1 index = 0 x = np.empty(len(sc_path)) index = 0 for sc_pos in sc_path: if sc_pos[2] > np.max(alt_mso_model): x[index] = np.NaN index += 1 continue if sc_pos[2] < np.min(alt_mso_model): x[index] = np.NaN index += 1 continue sorted_x_distance = np.argsort( np.abs(alt_mso_model - sc_pos[2])) alti1 = sorted_x_distance[0] if nearest: x[index] = interpolate.griddata(data_points, values[:, alti1], [sc_pos[0], sc_pos[1]], method='nearest') else: if alt_mso_model[alti1] < sc_pos[2]: alti2 = alti1 + 1 else: temp = alti1 - 1 alti2 = alti1 alti1 = temp #Interpolate through space first_val_calc = interpolate.LinearNDInterpolator( latlon_triangulation, values[:, alti1]) second_val_calc = interpolate.LinearNDInterpolator( latlon_triangulation, values[:, alti2]) first_val = first_val_calc([sc_pos[0], sc_pos[1]]) second_val = second_val_calc([sc_pos[0], sc_pos[1]]) delta_1 = sc_pos[2] - alt_mso_model[alti1] delta_2 = alt_mso_model[alti2] - sc_pos[2] delta_tot = alt_mso_model[alti2] - alt_mso_model[alti1] x[index] = ((first_val * delta_2) + (second_val * delta_1)) / (delta_tot) index += 1 tracer = np.array(x) if 'geo' == model['meta']['coord_sys'].lower(): #Build the Matrix that transforms GEO to MSO coordinates ls_rad = model['meta']['ls'] * np.pi / 180 rads_tilted_y = 25.19 * np.sin(ls_rad) * np.pi / 180 rads_tilted_x = -25.19 * np.cos(ls_rad) * np.pi / 180 lonsubsol_rad = -model['meta']['longsubsol'] * np.pi / 180 z_rotation = np.matrix( [[np.cos(lonsubsol_rad), -np.sin(lonsubsol_rad), 0], [np.sin(lonsubsol_rad), np.cos(lonsubsol_rad), 0], [0, 0, 1]]) y_rotation = np.matrix( [[np.cos(rads_tilted_y), 0, np.sin(rads_tilted_y)], [0, 1, 0], [-np.sin(rads_tilted_y), 0, np.cos(rads_tilted_y)]]) x_rotation = np.matrix( [[1, 0, 0], [0, np.cos(rads_tilted_x), -np.sin(rads_tilted_x)], [0, np.sin(rads_tilted_x), np.cos(rads_tilted_x)]]) geo_to_mso_matrix = np.dot(x_rotation, np.dot(y_rotation, z_rotation)) #Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows lat_geo_model = model['dim']['lat'] lon_geo_model = model['dim']['lon'] alt_geo_model = model['dim']['alt'] alt_array = np.repeat(alt_geo_model, len(lon_geo_model) * len(lat_geo_model)) lat_array = np.tile(np.repeat(lat_geo_model, len(lon_geo_model)), len(alt_geo_model)) lon_array = np.tile(lon_geo_model, len(lat_geo_model) * len(alt_geo_model)) data_points = np.transpose( np.array([lon_array, lat_array, alt_array])) #Convert to GEO coordinates, then to MSO index = 0 for point in data_points: r = point[2] + mars_radius x = r * np.sin((90 - point[1]) * np.pi / 180) * np.cos( point[0] * np.pi / 180) y = r * np.sin((90 - point[1]) * np.pi / 180) * np.sin( point[0] * np.pi / 180) z = r * np.cos((90 - point[1]) * np.pi / 180) data_points[index] = np.dot(geo_to_mso_matrix, np.array([x, y, z])) index += 1 #Convert to MSO Lon/Lat/Alt in order to weight the interpolation better lat_mso = np.empty(len(lon_geo_model) * len(lat_geo_model)) lon_mso = np.empty(len(lon_geo_model) * len(lat_geo_model)) index = 0 for point in data_points: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) lat_mso[index] = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso[index] = np.arctan2(point[1], point[0]) / (np.pi / 180) index += 1 if index >= len(lon_geo_model) * len(lat_geo_model): break latlon_points = np.transpose(np.array([lon_mso, lat_mso])) index = 0 for point in sc_path: r_mso = np.sqrt(point[0]**2 + point[1]**2 + point[2]**2) alt_mso = altitude lat_mso = 90 - (np.arccos(point[2] / r_mso) / (np.pi / 180)) lon_mso = np.arctan2(point[1], point[0]) / (np.pi / 180) sc_path[index] = np.array([lon_mso, lat_mso, alt_mso]) index += 1 latlon_triangulation = spatial.Delaunay(latlon_points) #Loop through the variables in the model for var in model: if var.lower() != dataname: continue print("Interpolating variable " + var) #Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'longitude': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'latitude': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'altitude': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) #Build an array of values that correspond to the points in data point values = np.empty([ len(lat_geo_model) * len(lon_geo_model), len(alt_geo_model) ]) index = 0 for alt in range(0, len(alt_geo_model)): for lat in range(0, len(lat_geo_model)): for lon in range(0, len(lon_geo_model)): values[index, alt] = data_new[lon][lat][alt] index += 1 index = 0 x = np.empty(len(sc_path)) index = 0 for sc_pos in sc_path: if sc_pos[2] > np.max(alt_geo_model): x[index] = np.NaN index += 1 continue if sc_pos[2] < np.min(alt_geo_model): x[index] = np.NaN index += 1 continue sorted_x_distance = np.argsort( np.abs(alt_geo_model - sc_pos[2])) alti1 = sorted_x_distance[0] if nearest: x[index] = interpolate.griddata(latlon_points, values[:, alti1], [sc_pos[0], sc_pos[1]], method='nearest') else: if alt_geo_model[alti1] < sc_pos[2]: alti2 = alti1 + 1 else: temp = alti1 - 1 alti2 = alti1 alti1 = temp #Interpolate through space first_val_calc = interpolate.LinearNDInterpolator( latlon_triangulation, values[:, alti1]) second_val_calc = interpolate.LinearNDInterpolator( latlon_triangulation, values[:, alti2]) first_val = first_val_calc([sc_pos[0], sc_pos[1]]) second_val = second_val_calc([sc_pos[0], sc_pos[1]]) delta_1 = sc_pos[2] - alt_geo_model[alti1] delta_2 = alt_geo_model[alti2] - sc_pos[2] delta_tot = alt_geo_model[alti2] - alt_geo_model[alti1] x[index] = ((first_val * delta_2) + (second_val * delta_1)) / (delta_tot) index += 1 tracer = np.array(x) else: if 'mso' == model['meta']['coord_sys'].lower(): #Build a big matrix with dimensions 3 columns by (num lat * num lon * num alt) rows x_mso_model = model['dim']['x'] y_mso_model = model['dim']['y'] z_mso_model = model['dim']['z'] #Loop through the variables in the model for var in model: if var.lower() != dataname: continue #Rearrange the data to lon/lat/alt data = model[var]['data'] dim_order_array = [0, 1, 2] for j in [0, 1, 2]: if model[var]['dim_order'][j] == 'x': dim_order_array[0] = j elif model[var]['dim_order'][j] == 'y': dim_order_array[1] = j elif model[var]['dim_order'][j] == 'z': dim_order_array[2] = j data_new = np.transpose(data, dim_order_array) #Interpolate through space tracer = mvn_kp_sc_traj_xyz(x_mso_model, y_mso_model, z_mso_model, data_new, sc_mso_x, sc_mso_y, sc_mso_z, nn=interp_method) xi, yi = np.linspace(sc_lon_mso.min(), sc_lon_mso.max(), 300), np.linspace(sc_lat_mso.min(), sc_lat_mso.max(), 300) xi, yi = np.meshgrid(xi, yi) zi = scipy.interpolate.griddata((sc_lon_mso, sc_lat_mso), tracer, (xi, yi), method=interp_method) if saveFig: fig = plt.figure() ax = fig.add_subplot(1, 1, 1) if fill: plt.contourf(xi, yi, zi, numContours, alpha=transparency, cmap=ct, extent=(-180, -90, 180, 90)) else: CS = plt.contour(xi, yi, zi, numContours, alpha=transparency, cmap=ct) plt.clabel(CS, inline=1, fontsize=7, fmt='%1.0f') extent = ax.get_window_extent().transformed( fig.dpi_scale_trans.inverted()) plt.axis('off') save_name = "ModelData_" + dataname + "_" + str(altitude) + "km" if fill: save_name = save_name + "_filled" print("Saving file as:" + os.path.join(os.path.dirname(file), save_name + ".png")) plt.savefig(os.path.join(os.path.dirname(file), save_name + ".png"), transparent=False, bbox_inches=extent, pad_inches=0, dpi=150) plt.show(block=True) else: # convert back to east longitude (as in model file) return dict(lon=xi[0, :], elon=((xi[0, :] + 360) % 360), longsubsol=model['meta']['longsubsol'], dec=model['meta']['dec'], Ls=model['meta']['ls'], lat=yi[:, 0], param=zi)
from scipy import interpolate import matplotlib.pyplot as plt from scipy.integrate import solve_ivp # http://www.ece.uci.edu/docs/hspice/hspice_2001_2-87.html # i guess we can do optimization ? try: import cPickle as pickle except ImportError: import pickle print("loading data") vd = np.genfromtxt('backward_vd.csv', delimiter=',').reshape(-1, 1) vg = np.genfromtxt('backward_vg.csv', delimiter=',').reshape(-1, 1) vc = np.genfromtxt('backward_vc.csv', delimiter=',').reshape(-1, 1) i = np.genfromtxt('backward_i.csv', delimiter=',').reshape(-1, 1) print("building interpolator") # points = np.concatenate((vd, vg), axis=1) points = np.concatenate((vc, vg), axis=1) x = points y = i fit = interpolate.LinearNDInterpolator(x, y, fill_value=0.0, rescale=True) with open('backward.pkl', 'wb') as f: pickle.dump(fit, f)
""" this module creates image transformation or project for two images' registration """ import SaveLoad import Visulization from scipy import ndimage from scipy import interpolate import numpy as np import scipy.ndimage as ndimage inputDVC = SaveLoad.DVCdata("/Users/junchaowei/Desktop/MultiStrainGage11182016/") # import the DVC database ################### save the displacement, and create interpolation ############## inter_x = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementX()[0], fill_value=0.0, rescale=True) inter_y = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementY()[0], fill_value=0.0, rescale=True) inter_z = interpolate.LinearNDInterpolator(inputDVC.getPoints1(), inputDVC.getDisplacementZ()[0], fill_value=0.0, rescale=True) image1 = inputDVC.getImage1() image2 = inputDVC.getImage2() sx, sy = image1[:,:,0].shape # get dimensionality of image mx, my = np.meshgrid(np.linspace(1,sx,sx),np.linspace(1,sy,sy)) # meshgrid change the orientation of x,y: x becomes the horizontal axis, y becomes the vertical # this change would affect the return statement of reshape. Cord_xy = np.vstack([mx.flatten(), my.flatten()]) # column major vector displacementDataX = np.zeros([sy,sx, image1.shape[2]]) displacementDataY = np.zeros([sy,sx, image1.shape[2]]) displacementDataZ = np.zeros([sy,sx, image1.shape[2]])
def trace_mline(self, init_state, time, direction='FWD', \ length_line=None, stp=None, ripple=True): ''' Traces the field line given a starting point. Integration step defined by stp and maximum length of the field line defined by s. Collision with the wall stops integration. input: - init_state [[r_1, r_2, ...], [phi_1, phi_2, ...], [z_1, z_2, ...]] : the coordinates of the starting points (list or numpy array) OR list of lists OR numpy 2D array - time : array of times requested (list or numpy array) - direction : direction of integration 'FWD' forward or 'REV' reverse (string) - ripple : take into account magnetic ripple or not (boolean) output: returns a dictionary containing: - r : radial coordinate (np.array) - p : toridal coordinate (np.array) - z : vertical coordinate (np.array) - x : x coordinate (np.array) - y : y coordinate (np.array) - cp : collision point with the wall (list) ''' # Step for the integration if (stp is None): stp = 0.001 # Length of the field line if (length_line is None): s = 100 else: s = length_line ds = np.linspace(0, s, int(s / stp)) init_state = np.squeeze(np.asarray(init_state)) if (init_state.ndim < 2): init_state = init_state[:, np.newaxis] self.ar_time = np.atleast_1d(np.squeeze(np.asarray([time]))) br_intp_t = np.atleast_2d(np.squeeze(self.f_intp_br(self.ar_time))) bt_intp_t = np.atleast_2d(np.squeeze(self.f_intp_bt(self.ar_time))) bz_intp_t = np.atleast_2d(np.squeeze(self.f_intp_bz(self.ar_time))) # !!!!!!!!!!!!!!!!!!!!! # HARD CODED CORRECTION if (np.nanmean(bt_intp_t) > 0): bt_intp_t *= -1 # !!!!!!!!!!!!!!!!!!!!! # Interpolate current itor_intp_t_vect = np.interp(self.ar_time, self.t_itor[:, 0], \ self.itor[:, 0]) b0_intp_t_vect = np.interp(self.ar_time, self.equi.time[self.mask], \ self.equi.vacuum_toroidal_field.b0[self.mask]) outMagLine = [] for ii in range(self.ar_time.size): self.br_lin_intp = \ interpolate.LinearNDInterpolator(self.delaunay, br_intp_t[ii]) self.bt_lin_intp = \ interpolate.LinearNDInterpolator(self.delaunay, bt_intp_t[ii]) self.bz_lin_intp = \ interpolate.LinearNDInterpolator(self.delaunay, bz_intp_t[ii]) # Interpolated current and b0 self.itor_intp_t = itor_intp_t_vect[ii] self.b0_intp_t = b0_intp_t_vect[ii] out = [] for jj in range(init_state.shape[1]): out_prime = self.integrate_solve_ivp(init_state[:, jj], \ direction, s, ds, ripple) out_prime['init_point'] = init_state[:, jj] out_prime['time'] = self.ar_time[ii] # Return list of dict out.append(out_prime) outMagLine.append(out) return outMagLine
def interp_quantity(quantity, interp_points, points, time_points, equi, \ ar_time, ar_R, ar_Phi, ar_Z, mask, mask_time, \ firstSpaceInterp, itor, t_itor, t_ignitron, \ no_ripple, ind_mid, equiDict): value_interpolated = np.full((ar_time.size, ar_R.size), np.nan) if (firstSpaceInterp): value_interpSpace = np.full((time_points.size, ar_R.size), np.nan) # Computation of requested quantities if (quantity == 'b_field_norm'): if (no_ripple): b_field_norm = np.sqrt(equiDict['b_field_r']**2. \ + equiDict['b_field_z']**2. \ + equiDict['b_field_tor']**2.) if (firstSpaceInterp): # Space interpolation for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ b_field_norm[mask, :][mask_time][ii, :]) value_interpSpace[ii, :] = lin_intp.__call__(interp_points) else: # Time interpolation f_intp = interpolate.interp1d(equi.time[mask], \ b_field_norm[mask, :], axis=0, \ bounds_error=False) else: # B_norm with ripple calculation # Declaration arrays br_intp = np.full((ar_time.size, ar_R.size), np.nan) bt_intp = np.full((ar_time.size, ar_R.size), np.nan) bz_intp = np.full((ar_time.size, ar_R.size), np.nan) if (firstSpaceInterp): br_Sintp = np.full((time_points.size, ar_R.size), np.nan) bt_Sintp = np.full((time_points.size, ar_R.size), np.nan) bz_Sintp = np.full((time_points.size, ar_R.size), np.nan) # Space interpolation for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ equiDict['b_field_r'][mask, :][mask_time][ii, :]) br_Sintp[ii, :] = lin_intp.__call__(interp_points) lin_intp = interpolate.LinearNDInterpolator(points, \ equiDict['b_field_tor'][mask, :][mask_time][ii, :]) bt_Sintp[ii, :] = lin_intp.__call__(interp_points) lin_intp = interpolate.LinearNDInterpolator(points, \ equiDict['b_field_z'][mask, :][mask_time][ii, :]) bz_Sintp[ii, :] = lin_intp.__call__(interp_points) # Time interpolation f_intp = interpolate.interp1d(time_points, \ br_Sintp, axis=0, \ bounds_error=False) br_intp = np.atleast_2d(np.squeeze(f_intp(ar_time))) f_intp = interpolate.interp1d(time_points, \ bt_Sintp, axis=0, \ bounds_error=False) bt_intp = np.atleast_2d(np.squeeze(f_intp(ar_time))) f_intp = interpolate.interp1d(time_points, \ bz_Sintp, axis=0, \ bounds_error=False) bz_intp = np.atleast_2d(np.squeeze(f_intp(ar_time))) else: # Time interpolation f_intp_br = interpolate.interp1d(equi.time[mask], \ equiDict['b_field_r'][mask, :], axis=0, \ bounds_error=False) f_intp_bt = interpolate.interp1d(equi.time[mask], \ equiDict['b_field_tor'][mask, :], axis=0, \ bounds_error=False) f_intp_bz = interpolate.interp1d(equi.time[mask], \ equiDict['b_field_z'][mask, :], axis=0, \ bounds_error=False) br_intp_t = np.atleast_2d(np.squeeze(f_intp_br(ar_time))) bt_intp_t = np.atleast_2d(np.squeeze(f_intp_bt(ar_time))) bz_intp_t = np.atleast_2d(np.squeeze(f_intp_bz(ar_time))) # Space interpolation for ii in range(ar_time.size): lin_intp = interpolate.LinearNDInterpolator( points, br_intp_t[ii]) br_intp[ii, :] = lin_intp.__call__(interp_points) lin_intp = interpolate.LinearNDInterpolator( points, bt_intp_t[ii]) bt_intp[ii, :] = lin_intp.__call__(interp_points) lin_intp = interpolate.LinearNDInterpolator( points, bz_intp_t[ii]) bz_intp[ii, :] = lin_intp.__call__(interp_points) # Interpolate current itor_intp_t = np.interp(ar_time, t_itor[:, 0], itor[:, 0]) b0_intp_t = np.interp(ar_time, equi.time[mask], \ equi.vacuum_toroidal_field.b0[mask]) # Compute reference vaccuum magnetic field bt_vac = equi.vacuum_toroidal_field.r0*b0_intp_t[:, np.newaxis] \ / ar_R[np.newaxis, :] # Compute magnetic field for given Phi br_ripple, bt_ripple, bz_ripple = mr.mag_ripple(ar_R, ar_Phi, \ ar_Z, itor_intp_t) # Check and correct Br if needed r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r z_mid_ax = \ equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_mid_ax)**2).argmin() if (equiDict['b_field_r'][ind_mid, ind_p] > 0): br_intp *= -1 warnings.warn( 'Correcting b_field_r in b_field_norm, for negative toroidal current COCOS 11' ) # Check and correct Bz if needed r_mid_ax = \ equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_z'][ind_mid, ind_p] < 0): bz_intp *= -1 warnings.warn( 'Correcting b_field_z in b_field_norm, for negative toroidal current COCOS 11' ) # Check and correct Btor if needed ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_tor'][ind_mid, ind_p] > 0): bt_intp *= -1 warnings.warn( 'Correcting b_field_tor in b_field_norm, for negative toroidal field COCOS 11' ) # Value interpolated value_interpolated = \ np.sqrt((br_intp - br_ripple)**2. \ + (np.abs(bt_intp - bt_ripple) - np.abs(bt_vac))**2. \ + (bz_intp - bz_ripple)**2.) return np.squeeze(value_interpolated) elif (not no_ripple and re.match('b_field_*', quantity)): # Declaration arrays br_ripple = np.full((ar_time.size, ar_R.size), np.nan) bt_ripple = np.full((ar_time.size, ar_R.size), np.nan) bz_ripple = np.full((ar_time.size, ar_R.size), np.nan) if (firstSpaceInterp): # Space interpolation quant_mask = eval('equiDict["' + quantity + '"][mask, :][mask_time]') for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ quant_mask[ii, :]) value_interpSpace[ii, :] = lin_intp.__call__(interp_points) # Time interpolation f_intp = interpolate.interp1d(time_points, \ value_interpSpace, axis=0, \ bounds_error=False) value_interpolated = np.atleast_2d(np.squeeze(f_intp(ar_time))) else: # Time interpolation f_intp = interpolate.interp1d(equi.time[mask], \ eval('equiDict["'+quantity+'"][mask, :]'), axis=0, \ bounds_error=False) b_intp_t = np.atleast_2d(np.squeeze(f_intp(ar_time))) # Space interpolation for ii in range(ar_time.size): lin_intp = interpolate.LinearNDInterpolator( points, b_intp_t[ii]) value_interpolated[ii, :] = lin_intp.__call__(interp_points) # Interpolate current itor_intp_t = np.interp(ar_time, t_itor[:, 0], itor[:, 0]) # Compute magnetic field for given Phi br_ripple, bt_ripple, bz_ripple = mr.mag_ripple(ar_R, ar_Phi, \ ar_Z, itor_intp_t) if (quantity == 'b_field_r'): r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r z_mid_ax = \ equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_mid_ax)**2).argmin() if (equiDict['b_field_r'][ind_mid, ind_p] > 0): value_interpolated *= -1 value_interpolated -= br_ripple warnings.warn( 'Correcting b_field_r, for negative toroidal current COCOS 11' ) else: value_interpolated -= br_ripple elif (quantity == 'b_field_z'): r_mid_ax = \ equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_z'][ind_mid, ind_p] < 0): value_interpolated *= -1 value_interpolated -= bz_ripple warnings.warn( 'Correcting b_field_z, for negative toroidal current COCOS 11' ) else: value_interpolated -= bz_ripple elif (quantity == 'b_field_tor'): b0_intp_t = np.interp(ar_time, equi.time[mask], \ equi.vacuum_toroidal_field.b0[mask]) # Compute reference vaccuum magnetic field bt_vac = equi.vacuum_toroidal_field.r0*b0_intp_t[:, np.newaxis] \ / ar_R[np.newaxis, :] r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_tor'][ind_mid, ind_p] > 0): value_interpolated *= -1 value_interpolated -= (bt_ripple - bt_vac) warnings.warn( 'Correcting b_field_tor, for negative toroidal field COCOS 11' ) else: value_interpolated -= (bt_ripple - bt_vac) else: print() print('ERROR: not valid quantity input:', quantity) print() raise SyntaxError return np.squeeze(value_interpolated) elif (quantity == 'rho_pol_norm'): rho_pol_norm = np.sqrt((equiDict['psi'] \ - equiDict['prof_1d_psi'][:, 0, np.newaxis]) \ / (equiDict['prof_1d_psi'][:, -1, np.newaxis] \ - equiDict['prof_1d_psi'][:, 0, np.newaxis])) if (firstSpaceInterp): # Space interpolation for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ rho_pol_norm[mask, :][mask_time][ii, :]) value_interpSpace[ii, :] = lin_intp.__call__(interp_points) else: # Time interpolation f_intp = interpolate.interp1d(equi.time[mask], rho_pol_norm[mask, :], \ axis=0, bounds_error=False) elif (quantity == 'rho_tor_norm' or quantity == 'rho_tor'): if (firstSpaceInterp): # Space interpolation for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ equiDict['psi'][mask, :][mask_time][ii, :]) value_interpSpace[ii, :] = lin_intp.__call__(interp_points) else: # Time interpolation f_intp = interpolate.interp1d(equi.time[mask], equiDict['psi'][mask, :], \ axis=0, bounds_error=False) elif (quantity == 'theta'): mag_ax_r = np.atleast_1d(np.squeeze(np.interp(ar_time, \ equi.time[mask], equiDict['mag_axis_r'][mask]))) mag_ax_z = np.atleast_1d(np.squeeze(np.interp(ar_time, \ equi.time[mask], equiDict['mag_axis_z'][mask]))) for ii in range(ar_time.size): delta_R = ar_R - mag_ax_r[ii, np.newaxis] delta_Z = ar_Z - mag_ax_z[ii, np.newaxis] val_arctan = np.arctan(delta_Z / delta_R) mask_theta = (delta_R >= 0) & (delta_Z >= 0) value_interpolated[ii, mask_theta] = val_arctan[mask_theta] mask_theta = (delta_R >= 0) & (delta_Z < 0) value_interpolated[ii, mask_theta] = 2 * np.pi + val_arctan[mask_theta] mask_theta = (delta_R < 0) value_interpolated[ii, mask_theta] = np.pi + val_arctan[mask_theta] return np.squeeze(2. * np.pi - value_interpolated) else: if (firstSpaceInterp): # Space interpolation for ii in range(time_points.size): lin_intp = interpolate.LinearNDInterpolator(points, \ eval('equiDict["'+quantity+'"][mask, :][mask_time][ii, :]')) value_interpSpace[ii, :] = lin_intp.__call__(interp_points) else: # Time interpolation f_intp = interpolate.interp1d(equi.time[mask], \ eval('equiDict["'+quantity+'"][mask, :]'), \ axis=0, bounds_error=False) if (firstSpaceInterp): # Time interpolation f_intp = interpolate.interp1d(time_points, \ value_interpSpace, axis=0, \ bounds_error=False) value_interpolated = np.atleast_2d(np.squeeze(f_intp(ar_time))) else: # Time interpolation out_time_interp = np.atleast_2d(np.squeeze(f_intp(ar_time))) # Space interpolation for ii in range(ar_time.size): lin_intp = interpolate.LinearNDInterpolator( points, out_time_interp[ii]) value_interpolated[ii, :] = lin_intp.__call__(interp_points) # Extra calculations for rho_tor, rho_tor_norm, b_field_r and b_field_z if (quantity == 'rho_tor_norm' or quantity == 'rho_tor'): value_interp_rho_tor = np.full(value_interpolated.shape, np.nan) if (quantity == 'rho_tor_norm'): try: rho_tor_norm = equiDict['prof_1d_rho_tor'] \ / equiDict['prof_1d_rho_tor'][:, -1, np.newaxis] except ZeroDivisionError as err: print('Division by zero for rho_tor_norm interpolation:', err) raise f_intp_rho_tor = interpolate.interp1d(equi.time[mask], \ rho_tor_norm[mask, :], axis=0, bounds_error=False) elif (quantity == 'rho_tor'): f_intp_rho_tor = interpolate.interp1d(equi.time[mask], \ equiDict['prof_1d_rho_tor'][mask, :], \ axis=0, bounds_error=False) rho_tor1D = np.atleast_2d(np.squeeze(f_intp_rho_tor(ar_time))) f_intp_psi = interpolate.interp1d(equi.time[mask], \ equiDict['prof_1d_psi'][mask, :], axis=0, bounds_error=False) psi1D = np.atleast_2d(np.squeeze(f_intp_psi(ar_time))) for ii in range(ar_time.size): f_intp_psi_rho_tor = interpolate.interp1d(psi1D[ii], rho_tor1D[ii], \ bounds_error=False) value_interp_rho_tor[ii] = f_intp_psi_rho_tor( value_interpolated[ii]) return np.squeeze(value_interp_rho_tor) elif (quantity == 'b_field_r'): r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r z_mid_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_mid_ax)**2).argmin() if (equiDict['b_field_r'][ind_mid, ind_p] > 0): value_interpolated *= -1 warnings.warn('Correcting b_field_r, if Ip negative COCOS 11') return np.squeeze(value_interpolated) elif (quantity == 'b_field_z'): r_mid_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r \ + 0.5*equi.time_slice[ind_mid].boundary.minor_radius z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z ind_p = np.abs((equiDict['r'] - r_mid_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_z'][ind_mid, ind_p] < 0): value_interpolated *= -1 warnings.warn('Correcting b_field_z, if Ip negative COCOS 11') return np.squeeze(value_interpolated) elif (quantity == 'b_field_tor'): r_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.r z_ax = equi.time_slice[ind_mid].global_quantities.magnetic_axis.z ind_p = np.abs((equiDict['r'] - r_ax)**2. \ + (equiDict['z'] - z_ax)**2).argmin() if (equiDict['b_field_tor'][ind_mid, ind_p] > 0): value_interpolated *= -1 warnings.warn( 'Correcting b_field_tor, if b_field_tor negative COCOS 11') return np.squeeze(value_interpolated) else: return np.squeeze(value_interpolated)
nfiles = len(filenames) beam_E = np.zeros((hp.nside2npix(nside), len(freqs))) thetai, phii = hp.pix2ang(nside, np.arange(hp.nside2npix(nside))) for fi, f in enumerate(filenames): data = np.loadtxt(f, skiprows=2) if opts.linear: lat = data[:, 0] * np.pi / 180.0 nlat = len(lat) lon = data[:, 1] * np.pi / 180.0 nlon = len(lon) coords = np.stack((lat, lon)).T gain = data[:, 2] lut = interpolate.LinearNDInterpolator(coords, gain) else: lat = np.unique(data[:, 0]) * np.pi / 180.0 nlat = len(lat) lon = np.unique(data[:, 1]) * np.pi / 180.0 nlon = len(lon) gain = data[:, 2].reshape(nlon, nlat).transpose() lut = interpolate.RectBivariateSpline(lat, lon, gain) for i in np.arange(hp.nside2npix(nside)): beam_E[i, fi] = lut(thetai[i], phii[i]) # Write to .fits file if opts.linear: filename = 'healpix_beam_%.0f-%.0fMHz_nside-%d_linear-interp.fits' %(freqs[0], freqs[-1], nside) else: filename = 'healpix_beam_%.0f-%.0fMHz_nside-%d.fits' %(freqs[0], freqs[-1], nside)
def compute_disps_on_grid(E_grid_data, disp_sensor_coords, L=1., w=0.2, Nx=200, Ny=40, nu=0.3, traction=0.1, output_paraview=False): ''' E_grid_data: #grid pts x 3 array with {x_pts, y_pts, E_vals} disp_sensor_coords: # sensors x 2 array with x/y coordinates of u sensor loc L: length w: width Nx: number of elements in x direction Ny: ... nu: poisson ratio traction: applied traction value on right side returns u_x, u_y - arrays with x/y components of displacement evaluated at the 2d grid defined by disp_x/y_sensors ''' mesh = RectangleMesh(Point(0., 0.), Point(L, w), Nx, Ny, "crossed") interp = scipy_interp.LinearNDInterpolator(E_grid_data[:, 0:2], E_grid_data[:, 2]) E = StiffnessInterp(interp) D = FunctionSpace(mesh, "DG", 0) E_field = interpolate(E, D) nu = Constant(nu) def eps(v): return sym(grad(v)) # https://en.wikipedia.org/wiki/Hooke%27s_law#Plane_stress: def sigma(v): return E_field / (1. - nu**2) * ( (1. - nu) * eps(v) + nu * tr(eps(v)) * Identity(2)) V = VectorFunctionSpace(mesh, 'Lagrange', degree=2) du = TrialFunction(V) u_ = TestFunction(V) a = inner(sigma(du), eps(u_)) * dx # Handle boundary conditions # Dirichlet displacement BC on the left side of beam: def left(x, on_boundary): return near(x[0], 0.) def bottomleftcorner(x, on_boundary): return (near(x[0], 0.) and near(x[1], 0.)) bc1 = DirichletBC(V.sub(0), Constant(0.), left) bc2 = DirichletBC(V, Constant((0., 0.)), bottomleftcorner, method='pointwise') bcs = [bc1, bc2] # Traction boundary condition on the right side of beam: class Right(SubDomain): def inside(self, x, on_boundary): return near(x[0], L) and on_boundary # Define traction: T = Constant((traction, 0.0)) facets = MeshFunction("size_t", mesh, 1) facets.set_all(0) Right().mark(facets, 1) ds = Measure('ds')(subdomain_data=facets) # Define right hand side of variational form l = dot(T, u_) * ds(1) # Solve problem u = Function(V, name="Displacement") set_log_level(50) solve(a == l, u, bcs) if output_paraview: file_results = XDMFFile("elasticity_results.xdmf") file_results.parameters["flush_output"] = True file_results.parameters["functions_share_mesh"] = True file_results.write(u, 0.) u_x = np.zeros(len(disp_sensor_coords)) u_y = np.zeros(len(disp_sensor_coords)) for i, coord in enumerate(disp_sensor_coords): p = Point(coord[0], coord[1]) disp = u(p) u_x[i] = disp[0] u_y[i] = disp[1] return u_x, u_y
def get_interpolation(cls, geo_data, **kwargs): """读等值线数据文件,并返回插值对象""" x = geo_data.x y = geo_data.y z = geo_data.z return interpolate.LinearNDInterpolator(np.array([x, y]).T, z)
def interpolate_heatmap(x, y, z, n: int=None, interp_method:str='linear'): """ Args: x (array): x data points y (array): y data points z (array): z data points n (int): number of points for each dimension on the interpolated grid interp_method {"linear", "nearest", "deg"} determines what interpolation method is used. Returns: x_grid : N*1 array of x-values of the interpolated grid y_grid : N*1 array of x-values of the interpolated grid z_grid : N*N array of z-values that form a grid. The output of this method can directly be used for plt.imshow(z_grid, extent=extent, aspect='auto') where the extent is determined by the min and max of the x_grid and y_grid. The output can also be used as input for ax.pcolormesh(x, y, Z,**kw) """ points = list(zip(x, y)) lbrt = np.min(points, axis=0), np.max(points, axis=0) lbrt = lbrt[0][0], lbrt[0][1], lbrt[1][0], lbrt[1][1] xy_mean = np.mean([lbrt[0], lbrt[2]]), np.mean([lbrt[1], lbrt[3]]) xy_scale = np.ptp([lbrt[0], lbrt[2]]), np.ptp([lbrt[1], lbrt[3]]) # interpolation needs to happen on a rescaled grid, this is somewhat akin to an # assumption in the interpolation that the scale of the experiment is chosen sensibly. # N.B. even if interp_method == "nearest" the linear interpolation is used # to determine the amount of grid points. Could be improved. ip = interpolate.LinearNDInterpolator( scale(points, xy_mean=xy_mean, xy_scale=xy_scale), z) if n is None: # Calculate how many grid points are needed. # factor from A=√3/4 * a² (equilateral triangle) # N.B. a factor 4 was added as there were to few points for uniform # grid otherwise. n = int(0.658 / np.sqrt(areas(ip).min()))*4 n = max(n, 10) if n > 500: logging.warning('n: {} larger than 500'.format(n)) n=500 x_lin = y_lin = np.linspace(-0.5, 0.5, n) if interp_method == 'linear': z_grid = ip(x_lin[:, None], y_lin[None, :]).squeeze() elif interp_method == "nearest": ip = interpolate.NearestNDInterpolator( scale(points, xy_mean=xy_mean, xy_scale=xy_scale), z) z_grid = ip(x_lin[:, None], y_lin[None, :]).squeeze() elif interp_method == "deg": # Circular interpolation in deg units phases=np.deg2rad(z) newdata_cos=np.cos(phases) newdata_sin=np.sin(phases) ip_cos = interpolate.LinearNDInterpolator( scale(points, xy_mean=xy_mean, xy_scale=xy_scale), newdata_cos) newdata_cos = ip_cos(x_lin[:, None], y_lin[None, :]).squeeze() ip_sin = interpolate.LinearNDInterpolator( scale(points, xy_mean=xy_mean, xy_scale=xy_scale), newdata_sin) newdata_sin = ip_sin(x_lin[:, None], y_lin[None, :]).squeeze() z_grid = (np.rad2deg(np.arctan2(newdata_sin, newdata_cos)) % 360).squeeze() # x and y grid points need to be rescaled from the linearly chosen points points_grid = unscale(list(zip(x_lin, y_lin)), xy_mean=xy_mean, xy_scale=xy_scale) x_grid = points_grid[:, 0] y_grid = points_grid[:, 1] return x_grid, y_grid, (z_grid).T
S = [] sig = [] #This here makes a grid for a given Re (alfa,S) and for each point sigma for case in caseList: # print case.alfa, case.Re for p in zip(case.S, case.sigma): alfa.append(float(case.alfa)) S.append(p[0]) sig.append(p[1]) p = [] #this grid (alfa,S) is stored in p, and used to find an interpolant for pair in zip(alfa, S): p.append([pair[0], pair[1]]) if InterType == 1: interp = interpolate.LinearNDInterpolator(p, sig) elif InterType == 2: interp = interpolate.CloughTocher2DInterpolator(p, sig) else: raise "Unknown Interpolation type" if with3D: Anew = np.linspace(min(alfa), max(alfa), num=200, endpoint=True) Snew = np.linspace(min(S), max(S), num=200, endpoint=True) Anew, Snew = np.meshgrid(Anew, Snew) zz = interp(Anew, Snew) fig = plt.figure() ax = fig.gca(projection='3d') ax.set_xlabel('alpha') ax.set_ylabel('S') ax.set_zlabel('Sigma')
# lat_values_1d = lat_values.flatten() # data_values_1d = data_values.flatten() # data_interp_values = interpolate.griddata((lon_values_1d, lat_values_1d), data_values_1d, (lon0_values, lat0_values), method='nearest') ######## ### For consistency I have to bring Vector of lats and cords to ### scipy.interpolate.LinearNDInterpolator if lon_values.ndim == 1: lon_values, lat_values = np.meshgrid(lon_values, lat_values) ### lon_values_1d = lon_values.flatten() lon_values_1d[lon_values_1d < 0] = lon_values_1d[lon_values_1d < 0] + 360 lat_values_1d = lat_values.flatten() data_values_1d = data_values.flatten() points = list(zip(lon_values_1d, lat_values_1d)) interp_surface = interpolate.LinearNDInterpolator(points, data_values_1d) data_interp_values = interp_surface(lon0_grid, lat0_grid) # data_interp_values = scipy.interpolate.interpn((lat, lon), data_values, (lat0, lon0), method='nearest') # xr.concat([data0, data_interp], 'ensembles') # stopped using xarray data_ensembles = np.dstack((data_ensembles, data_interp_values)) print('Finished #' + str(iexperiment) + ' From ' + str(len(cmip_list))) fig, ax = plt.subplots(3) # data.plot(ax=ax[0], vmin=5, vmax=30) ax[0].pcolormesh(data, cmap='RdBu_r', vmin=5, vmax=30) ax[0].set_title('Data : ' + str(ds.source_id)) ax[1].pcolormesh(lon0_values, lat0_values, data0_values, vmin=5,
dis_map_h = np.load('../data/distortion/194-437-xa_high-centroids.npy') #dis_map_l = np.load('../data/distortion/5-185-xa_low-centroids.npy') #dis_map_h = np.load('../data/distortion/5-185-xa_high-centroids.npy') dis_map_l = dis_map_l.reshape( 50 * 50, 2) / detector_size[0] * 36000 * 800 * 0.001666 dis_map_h = dis_map_h.reshape( 50 * 50, 2) / detector_size[0] * 36000 * 800 * 0.001666 cen_x = np.arange(detector_size[0]) + 0.5 cen_y = np.arange(detector_size[1]) + 0.5 xi = np.repeat(detector2gon(cen_x, detector_size[0]), detector_size[1]) eta = np.tile(detector2gon(cen_y, detector_size[1]), detector_size[0]) points = np.array([xi, eta]).T fxi_l = interpolate.LinearNDInterpolator(points, dis_map_l[:, 0]) feta_l = interpolate.LinearNDInterpolator(points, dis_map_l[:, 1]) fxi_h = interpolate.LinearNDInterpolator(points, dis_map_h[:, 0]) feta_h = interpolate.LinearNDInterpolator(points, dis_map_h[:, 1]) dis_map = { 'xi_l': fxi_l, 'xi_h': fxi_h, 'eta_l': feta_l, 'eta_h': feta_h } manager = Manager() return_dict = manager.dict() p_list = [] pid = 0
def _get_t_isochrone(self, age, metal=None, FeH=None, masses=None, *args, **kwargs): """ Retrieve isochrone from the original source internal use to adapt any library """ # make sure unit is in years and then only give the value (no units) _age = int(units.Quantity(age, units.year).value) # if hasUnit(age): # _age = int(age.to('yr').magnitude) # else: # _age = int(age * inputUnit.to('yr').magnitude) _logA = np.log10(_age) assert ((metal is not None) | (FeH is not None)),\ "Need a chemical par. value." if (metal is not None) & (FeH is not None): print("Warning: both Z & [Fe/H] provided, ignoring [Fe/H].") if metal is None: metal = self.FeHtometal(FeH) if self.interpolation(): # Do the actual nd interpolation # Maybe already exists? if (metal in self.Z) & (_age in self.ages): t = self.selectWhere( '*', '(round(Z, 6) == {0}) & (round(logA, 6) == {1})'.format( metal, _logA)) if t.nrows > 0: return t # apparently not # find 2 closest metal values ca1 = (self.ages <= _age) ca2 = (self.ages > _age) cz1 = (self.Z <= metal) cz2 = (self.Z > metal) if (metal in self.Z): # perfect match in metal, need to find ages if (_age in self.ages): return self.selectWhere( '*', '(round(Z, 6) == {0}) & (round(logA, 6) == {1})'. format(metal, _logA)) elif (True in ca1) & (True in ca2): # bracket on _age: closest values a1, a2 = (np.log10(max(self.ages[ca1])), np.log10(min(self.ages[ca2]))) iso = self.selectWhere( '*', '(Z == 0.02) & ( (abs(logA - {0}) < 1e-4) | (abs(logA - {1}) < 1e-4 ) )' .format(a1, a2)) if masses is None: _logM = np.unique(iso['logM']) else: _logM = masses # define interpolator points = np.array([self[k] for k in 'logA logM Z'.split()]).T values = np.array( [self[k] for k in list(self.data.keys())]).T _ifunc = interpolate.LinearNDInterpolator(points, values) pts = np.array([(_logA, logMk, metal) for logMk in _logM]) r = _ifunc(pts) return Table(r) else: raise Exception('Age not covered by the isochrones') elif (True in cz1) & (True in cz2): # need to find closest Z pass return else: # find the closest match _Z = self.Z[((metal - self.Z)**2).argmin()] # _logA = np.log10(self.ages[((_age - self.ages) ** 2).argmin()]) _logA = self.logages[((np.log10(_age) - self.logages)**2).argmin()] tab = self.data.selectWhere( '*', "(round(Z, 6) == {0}) & (round(logA,6) == {1})".format( _Z, _logA)) # mass selection if masses is not None: # masses are expected in logM for interpolation # if masses.max() > 2.3: # _m = np.log10(masses) # else: _m = masses data_logM = tab['logM'][:] # refuse extrapolation! # ind = np.where(_m <= max(data_logM)) data = {} for kn in list(tab.keys()): data[kn] = interp(_m, data_logM, tab[kn], left=np.nan, right=np.nan) return Table(data)
def refine_tree_xyz( mesh, xyz, method="radial", octree_levels=[1, 1, 1], octree_levels_padding=None, finalize=False, min_level=0, max_distance=np.inf ): """ Refine a TreeMesh based on xyz point locations Parameters ---------- mesh: TreeMesh The TreeMesh object to be refined xyz: numpy.ndarray 2D array of points method: str Method used to refine the mesh based on xyz locations - "radial": Based on radial distance xyz and cell centers - "surface": Along triangulated surface repeated vertically - "box": Inside limits defined by outer xyz locations octree_levels: list Minimum number of cells around points in each k octree level [N(k), N(k-1), ...] octree_levels_padding: list Padding cells added to the outer limits of the data each octree levels used for method= "surface" and "box" [N(k), N(k-1), ...]. finalize: bool True | [False] Finalize the TreeMesh max_distance: float Maximum refinement distance from xyz locations. Used for method="surface" to reduce interpolation distance. Returns -------- discretize.TreeMesh mesh """ if octree_levels_padding is not None: if len(octree_levels_padding) != len(octree_levels): raise ValueError( "'octree_levels_padding' must be the length %i" % len(octree_levels) ) else: octree_levels_padding = np.zeros_like(octree_levels) # Prime the refinement against large cells mesh.insert_cells( xyz, [mesh.max_level - np.nonzero(octree_levels)[0][0]]*xyz.shape[0], finalize=False ) # Trigger different refine methods if method.lower() == "radial": # Build a cKDTree for fast nearest lookup tree = cKDTree(xyz) # Compute the outer limits of each octree level rMax = np.cumsum( mesh.hx.min() * np.asarray(octree_levels) * 2**np.arange(len(octree_levels)) ) # Radial function def inBall(cell): xyz = cell.center r, ind = tree.query(xyz) for ii, nC in enumerate(octree_levels): if r < rMax[ii]: return mesh.max_level-ii return min_level mesh.refine(inBall, finalize=finalize) elif method.lower() == 'surface': # Compute centroid centroid = np.mean(xyz, axis=0) if mesh.dim == 2: rOut = np.abs(centroid[0]-xyz).max() hz = mesh.hy.min() else: # Largest outer point distance rOut = np.linalg.norm(np.r_[ np.abs(centroid[0]-xyz[:, 0]).max(), np.abs(centroid[1]-xyz[:, 1]).max() ] ) hz = mesh.hz.min() # Compute maximum depth of refinement zmax = np.cumsum( hz * np.asarray(octree_levels) * 2**np.arange(len(octree_levels)) ) # Compute maximum horizontal padding offset padWidth = np.cumsum( mesh.hx.min() * np.asarray(octree_levels_padding) * 2**np.arange(len(octree_levels_padding)) ) # Increment the vertical offset zOffset = 0 xyPad = -1 depth = zmax[-1] # Cycle through the Tree levels backward for ii in range(len(octree_levels)-1, -1, -1): dx = mesh.hx.min() * 2**ii if mesh.dim == 3: dy = mesh.hy.min() * 2**ii dz = mesh.hz.min() * 2**ii else: dz = mesh.hy.min() * 2**ii # Increase the horizontal extent of the surface if xyPad != padWidth[ii]: xyPad = padWidth[ii] # Calculate expansion for padding XY cells expansion_factor = (rOut + xyPad) / rOut xLoc = (xyz - centroid)*expansion_factor + centroid if mesh.dim == 3: # Create a new triangulated surface tri2D = Delaunay(xLoc[:, :2]) F = interpolate.LinearNDInterpolator(tri2D, xLoc[:, 2]) else: F = interpolate.interp1d(xLoc[:, 0], xLoc[:, 1], fill_value='extrapolate') limx = np.r_[xLoc[:, 0].max(), xLoc[:, 0].min()] nCx = int(np.ceil((limx[0]-limx[1]) / dx)) if mesh.dim == 3: limy = np.r_[xLoc[:, 1].max(), xLoc[:, 1].min()] nCy = int(np.ceil((limy[0]-limy[1]) / dy)) # Create a grid at the octree level in xy CCx, CCy = np.meshgrid( np.linspace( limx[1], limx[0], nCx ), np.linspace( limy[1], limy[0], nCy ) ) xy = np.c_[CCx.reshape(-1), CCy.reshape(-1)] # Only keep points within triangulation indexTri = tri2D.find_simplex(xy) else: xy = np.linspace( limx[1], limx[0], nCx ) indexTri = np.ones_like(xy, dtype='bool') # Interpolate the elevation linearly z = F(xy[indexTri != -1]) newLoc = np.c_[xy[indexTri != -1], z] # Only keep points within max_distance tree = cKDTree(xyz) r, ind = tree.query(newLoc) # Apply vertical padding for current octree level dim = mesh.dim - 1 zOffset = 0 while zOffset < depth: indIn = r < (max_distance + padWidth[ii]) nnz = int(np.sum(indIn)) if nnz > 0: mesh.insert_cells( np.c_[ newLoc[indIn, :dim], newLoc[indIn, -1]-zOffset], np.ones(nnz)*mesh.max_level-ii, finalize=False ) zOffset += dz depth -= dz * octree_levels[ii] if finalize: mesh.finalize() elif method.lower() == 'box': # Define the data extend [bottom SW, top NE] bsw = np.min(xyz, axis=0) tne = np.max(xyz, axis=0) hx = mesh.hx.min() if mesh.dim == 2: hz = mesh.hy.min() else: hz = mesh.hz.min() # Pre-calculate max depth of each level zmax = np.cumsum( hz * np.asarray(octree_levels) * 2**np.arange(len(octree_levels)) ) if mesh.dim == 2: # Pre-calculate outer extent of each level padWidth = np.cumsum( mesh.hx.min() * np.asarray(octree_levels_padding) * 2**np.arange(len(octree_levels_padding)) ) # Make a list of outer limits BSW = [ bsw - np.r_[padWidth[ii], zmax[ii]] for ii, (octZ, octXY) in enumerate( zip(octree_levels, octree_levels_padding) ) ] TNE = [ tne + np.r_[padWidth[ii], zmax[ii]] for ii, (octZ, octXY) in enumerate( zip(octree_levels, octree_levels_padding) ) ] else: hy = mesh.hy.min() # Pre-calculate outer X extent of each level padWidth_x = np.cumsum( hx * np.asarray(octree_levels_padding) * 2**np.arange(len(octree_levels_padding)) ) # Pre-calculate outer Y extent of each level padWidth_y = np.cumsum( hy * np.asarray(octree_levels_padding) * 2**np.arange(len(octree_levels_padding)) ) # Make a list of outer limits BSW = [ bsw - np.r_[padWidth_x[ii], padWidth_y[ii], zmax[ii]] for ii, (octZ, octXY) in enumerate( zip(octree_levels, octree_levels_padding) ) ] TNE = [ tne + np.r_[padWidth_x[ii], padWidth_y[ii], zmax[ii]] for ii, (octZ, octXY) in enumerate( zip(octree_levels, octree_levels_padding) ) ] def inBox(cell): xyz = cell.center for ii, (nC, bsw, tne) in enumerate(zip(octree_levels, BSW, TNE)): if np.all([xyz > bsw, xyz < tne]): return mesh.max_level-ii return cell._level mesh.refine(inBox, finalize=finalize) else: raise NotImplementedError( "Only method= 'radial', 'surface'" " or 'box' have been implemented" ) return mesh
def write2netCDF(filename, lon, lat, z, increments=None, nSamples=None, title='CSI product', name='z', scale=1.0, offset=0.0, mask=None, xyunits=['Lon', 'Lat'], units='None', interpolation=True, verbose=True, noValues=np.nan): ''' Creates a netCDF file with the arrays in Z. Z can be list of array or an array, the size of lon. .. Args: * filename -> Output file name * lon -> 1D Array of lon values * lat -> 1D Array of lat values * z -> 2D slice to be saved * mask -> if not None, must be a 2d-array of a polynome to mask what is outside of it. This option is really long, so I don't use it... .. Kwargs: * title -> Title for the grd file * name -> Name of the field in the grd file * scale -> Scale value in the grd file * offset -> Offset value in the grd file .. Returns: * None ''' if interpolation: # Check if nSamples is not None: if type(nSamples) is int: nSamples = [nSamples, nSamples] dlon = (lon.max() - lon.min()) / nSamples[0] dlat = (lat.max() - lat.min()) / nSamples[1] if increments is not None: dlon, dlat = increments # Resample on a regular grid olon, olat = np.meshgrid(np.arange(lon.min(), lon.max(), dlon), np.arange(lat.min(), lat.max(), dlat)) else: # Get lon lat olon = lon olat = lat if increments is not None: dlon, dlat = increments else: dlon = olon[0, 1] - olon[0, 0] dlat = olat[1, 0] - olat[0, 0] # Create a file fid = netcdf(filename, 'w') # Create a dimension variable fid.createDimension('side', 2) if verbose: print('Create dimension xysize with size {}'.format(np.prod( olon.shape))) fid.createDimension('xysize', np.prod(olon.shape)) # Range variables fid.createVariable('x_range', 'd', ('side', )) fid.variables['x_range'].units = xyunits[0] fid.createVariable('y_range', 'd', ('side', )) fid.variables['y_range'].units = xyunits[1] # Spacing fid.createVariable('spacing', 'd', ('side', )) fid.createVariable('dimension', 'i4', ('side', )) # Informations if title is not None: fid.title = title fid.source = 'CSI.utils.write2netCDF' # Filing rnage and spacing if verbose: print('x_range from {} to {} with spacing {}'.format( olon[0, 0], olon[0, -1], dlon)) fid.variables['x_range'][0] = olon[0, 0] fid.variables['x_range'][1] = olon[0, -1] fid.variables['spacing'][0] = dlon if verbose: print('y_range from {} to {} with spacing {}'.format( olat[0, 0], olat[-1, 0], dlat)) fid.variables['y_range'][0] = olat[0, 0] fid.variables['y_range'][1] = olat[-1, 0] fid.variables['spacing'][1] = dlat if interpolation: # Interpolate interpZ = sciint.LinearNDInterpolator(np.vstack((lon, lat)).T, z, fill_value=noValues) oZ = interpZ(olon, olat) else: # Get values oZ = z # Masking? if mask is not None: # Import matplotlib.path import matplotlib.path as path # Create the path poly = path.Path([[lo, la] for lo, la in zip(mask[:, 0], mask[:, 1])], closed=False) # Create the list of points xy = np.vstack((olon.flatten(), olat.flatten())).T # Findthose outside bol = poly.contains_points(xy) # Mask those out oZ = oZ.flatten() oZ[bol] = np.nan oZ = oZ.reshape(olon.shape) # Range zmin = np.nanmin(oZ) zmax = np.nanmax(oZ) fid.createVariable('{}_range'.format(name), 'd', ('side', )) fid.variables['{}_range'.format(name)].units = units fid.variables['{}_range'.format(name)][0] = zmin fid.variables['{}_range'.format(name)][1] = zmax # Create Variable fid.createVariable(name, 'd', ('xysize', )) fid.variables[name].long_name = name fid.variables[name].scale_factor = scale fid.variables[name].add_offset = offset fid.variables[name].node_offset = 0 # Fill it fid.variables[name][:] = np.flipud(oZ).flatten() # Set dimension fid.variables['dimension'][:] = oZ.shape[::-1] # Synchronize and close fid.sync() fid.close() # All done return
def ax_scalp(v, channels, ax=None, annotate=False, vmin=None, vmax=None, colormap=None): """Draw a scalp plot. Draws a scalp plot on an existing axes. The method takes an array of values and an array of the corresponding channel names. It matches the channel names with an internal list of known channels and their positions to project them correctly on the scalp. .. warning:: The behaviour for unkown channels is undefined. Parameters ---------- v : 1d-array of floats The values for the channels channels : 1d array of strings The corresponding channel names for the values in ``v`` ax : Axes, optional The axes to draw the scalp plot on. If not provided, the currently activated axes (i.e. ``gca()``) will be taken annotate : Boolean, optional Draw the channel names next to the channel markers. vmin, vmax : float, optional The display limits for the values in ``v``. If the data in ``v`` contains values between -3..3 and ``vmin`` and ``vmax`` are set to -1 and 1, all values smaller than -1 and bigger than 1 will appear the same as -1 and 1. If not set, the maximum absolute value in ``v`` is taken to calculate both values. colormap : matplotlib.colors.colormap, optional A colormap to define the color transitions. Returns ------- ax : Axes the axes on which the plot was drawn See Also -------- ax_colorbar """ if ax is None: ax = plt.gca() # what if we have an unknown channel? points = [get_channelpos(c) for c in channels] # calculate the interpolation x = [i[0] for i in points] y = [i[1] for i in points] z = v # interplolate the in-between values xx = np.linspace(min(x), max(x), 500) yy = np.linspace(min(y), max(y), 500) xx, yy = np.meshgrid(xx, yy) f = interpolate.LinearNDInterpolator(list(zip(x, y)), z) zz = f(xx, yy) # draw the contour map ctr = ax.contourf(xx, yy, zz, 20, vmin=vmin, vmax=vmax, cmap=colormap) ax.contour(xx, yy, zz, 5, colors="k", vmin=vmin, vmax=vmax, linewidths=.1) # paint the head ax.add_artist( plt.Circle((0, 0), 1, linestyle='solid', linewidth=2, fill=False)) # add a nose ax.plot([-0.1, 0, 0.1], [1, 1.1, 1], 'k-') # add markers at channels positions ax.plot(x, y, 'k+') # set the axes limits, so the figure is centered on the scalp ax.set_ylim([-1.05, 1.15]) ax.set_xlim([-1.15, 1.15]) # hide the frame and axes # hiding the axes might be too much, as this will also hide the x/y # labels :/ ax.set_frame_on(False) ax.get_xaxis().set_visible(False) ax.get_yaxis().set_visible(False) # draw the channel names if annotate: for i in zip(channels, list(zip(x, y))): ax.annotate(" " + i[0], i[1]) ax.set_aspect(1) plt.sci(ctr) return ax
def refine_tree_xyz( mesh, xyz, method="radial", octree_levels=[1, 1, 1], octree_levels_padding=None, finalize=False, min_level=0, max_distance=np.inf, ): """Refine region within a :class:`~discretize.TreeMesh` This function refines the specified region of a tree mesh using one of several methods. These are summarized below: **radial:** refines based on radial distances from a set of xy[z] locations. Consider a tree mesh whose smallest cell size has a width of *h* . And *octree_levels = [nc1, nc2, nc3, ...]* . Within a distance of *nc1 x h* from any of the points supplied, the smallest cell size is used. Within a distance of *nc2 x (2h)* , the cells will have a width of *2h* . Within a distance of *nc3 x (4h)* , the cells will have a width of *4h* . Etc... **surface:** refines downward from a triangulated surface. Consider a tree mesh whose smallest cell size has a width of *h*. And *octree_levels = [nc1, nc2, nc3, ...]* . Within a downward distance of *nc1 x h* from the topography (*xy[z]* ) supplied, the smallest cell size is used. The topography is triangulated if the points supplied are coarser than the cell size. No refinement is done above the topography. Within a vertical distance of *nc2 x (2h)* , the cells will have a width of *2h* . Within a vertical distance of *nc3 x (4h)* , the cells will have a width of *4h* . Etc... **box:** refines inside the convex hull defined by the xy[z] locations. Consider a tree mesh whose smallest cell size has a width of *h*. And *octree_levels = [nc1, nc2, nc3, ...]* . Within the convex hull defined by *xyz* , the smallest cell size is used. Within a distance of *nc2 x (2h)* from that convex hull, the cells will have a width of *2h* . Within a distance of *nc3 x (4h)* , the cells will have a width of *4h* . Etc... Parameters ---------- mesh : discretize.TreeMesh The tree mesh object to be refined xyz : numpy.ndarray 2D array of points (n, dim) method : {'radial', 'surface', 'box'} Method used to refine the mesh based on xyz locations. - *radial:* Based on radial distance xy[z] and cell centers - *surface:* Refines downward from a triangulated surface - *box:* Inside limits defined by outer xy[z] locations octree_levels : list of int, optional Minimum number of cells around points in each *k* octree level starting from the smallest cells size; i.e. *[nc(k), nc(k-1), ...]* . Note that you *can* set entries to 0; e.g. you don't want to discretize using the smallest cell size. octree_levels_padding : list of int, optional Padding cells added to extend the region of refinement at each level. Used for *method = surface* and *box*. Has the form *[nc(k), nc(k-1), ...]* finalize : bool, optional Finalize the tree mesh. min_level : int, optional Sets the largest cell size allowed in the mesh. The default (*0*), allows the largest cell size to be used. max_distance : float Maximum refinement distance from xy[z] locations. Used if *method* = "surface" to reduce interpolation distance Returns ------- discretize.TreeMesh The refined tree mesh Examples -------- Here we use the **refine_tree_xyz** function refine a tree mesh based on topography as well as a cluster of points. >>> from discretize import TreeMesh >>> from discretize.utils import mkvc, refine_tree_xyz >>> import matplotlib.pyplot as plt >>> import numpy as np >>> dx = 5 # minimum cell width (base mesh cell width) in x >>> dy = 5 # minimum cell width (base mesh cell width) in y >>> x_length = 300.0 # domain width in x >>> y_length = 300.0 # domain width in y Compute number of base mesh cells required in x and y >>> nbcx = 2 ** int(np.round(np.log(x_length / dx) / np.log(2.0))) >>> nbcy = 2 ** int(np.round(np.log(y_length / dy) / np.log(2.0))) Define the base mesh >>> hx = [(dx, nbcx)] >>> hy = [(dy, nbcy)] >>> mesh = TreeMesh([hx, hy], x0="CC") Refine surface topography >>> xx = mesh.vectorNx >>> yy = -3 * np.exp((xx ** 2) / 100 ** 2) + 50.0 >>> pts = np.c_[mkvc(xx), mkvc(yy)] >>> mesh = refine_tree_xyz( ... mesh, pts, octree_levels=[2, 4], method="surface", finalize=False ... ) Refine mesh near points >>> xx = np.array([-10.0, 10.0, 10.0, -10.0]) >>> yy = np.array([-40.0, -40.0, -60.0, -60.0]) >>> pts = np.c_[mkvc(xx), mkvc(yy)] >>> mesh = refine_tree_xyz( ... mesh, pts, octree_levels=[4, 2], method="radial", finalize=True ... ) Plot the mesh >>> fig = plt.figure(figsize=(6, 6)) >>> ax = fig.add_subplot(111) >>> mesh.plotGrid(ax=ax) >>> ax.set_xbound(mesh.x0[0], mesh.x0[0] + np.sum(mesh.hx)) >>> ax.set_ybound(mesh.x0[1], mesh.x0[1] + np.sum(mesh.hy)) >>> ax.set_title("QuadTree Mesh") >>> plt.show() """ if octree_levels_padding is not None: if len(octree_levels_padding) != len(octree_levels): raise ValueError("'octree_levels_padding' must be the length %i" % len(octree_levels)) else: octree_levels_padding = np.zeros_like(octree_levels) octree_levels = np.asarray(octree_levels) octree_levels_padding = np.asarray(octree_levels_padding) # Trigger different refine methods if method.lower() == "radial": # Compute the outer limits of each octree level rMax = np.cumsum(mesh.h[0].min() * octree_levels * 2**np.arange(len(octree_levels))) rs = np.ones(xyz.shape[0]) level = np.ones(xyz.shape[0], dtype=np.int32) for ii, nC in enumerate(octree_levels): # skip "zero" sized balls if rMax[ii] > 0: mesh.refine_ball(xyz, rs * rMax[ii], level * (mesh.max_level - ii), finalize=False) if finalize: mesh.finalize() elif method.lower() == "surface": # Compute centroid centroid = np.mean(xyz, axis=0) if mesh.dim == 2: rOut = np.abs(centroid[0] - xyz).max() hz = mesh.h[1].min() else: # Largest outer point distance rOut = np.linalg.norm(np.r_[np.abs(centroid[0] - xyz[:, 0]).max(), np.abs(centroid[1] - xyz[:, 1]).max(), ]) hz = mesh.h[2].min() # Compute maximum depth of refinement zmax = np.cumsum(hz * octree_levels * 2**np.arange(len(octree_levels))) # Compute maximum horizontal padding offset padWidth = np.cumsum(mesh.h[0].min() * octree_levels_padding * 2**np.arange(len(octree_levels_padding))) # Increment the vertical offset zOffset = 0 xyPad = -1 depth = zmax[-1] # Cycle through the Tree levels backward for ii in range(len(octree_levels) - 1, -1, -1): dx = mesh.h[0].min() * 2**ii if mesh.dim == 3: dy = mesh.h[1].min() * 2**ii dz = mesh.h[2].min() * 2**ii else: dz = mesh.h[1].min() * 2**ii # Increase the horizontal extent of the surface if xyPad != padWidth[ii]: xyPad = padWidth[ii] # Calculate expansion for padding XY cells expansion_factor = (rOut + xyPad) / rOut xLoc = (xyz - centroid) * expansion_factor + centroid if mesh.dim == 3: # Create a new triangulated surface tri2D = Delaunay(xLoc[:, :2]) F = interpolate.LinearNDInterpolator(tri2D, xLoc[:, 2]) else: F = interpolate.interp1d(xLoc[:, 0], xLoc[:, 1], fill_value="extrapolate") limx = np.r_[xLoc[:, 0].max(), xLoc[:, 0].min()] nCx = int(np.ceil((limx[0] - limx[1]) / dx)) if mesh.dim == 3: limy = np.r_[xLoc[:, 1].max(), xLoc[:, 1].min()] nCy = int(np.ceil((limy[0] - limy[1]) / dy)) # Create a grid at the octree level in xy CCx, CCy = np.meshgrid( np.linspace(limx[1], limx[0], nCx), np.linspace(limy[1], limy[0], nCy), ) xy = np.c_[CCx.reshape(-1), CCy.reshape(-1)] # Only keep points within triangulation indexTri = tri2D.find_simplex(xy) else: xy = np.linspace(limx[1], limx[0], nCx) indexTri = np.ones_like(xy, dtype="bool") # Interpolate the elevation linearly z = F(xy[indexTri != -1]) newLoc = np.c_[xy[indexTri != -1], z] # Only keep points within max_distance tree = cKDTree(xyz) r, ind = tree.query(newLoc) # Apply vertical padding for current octree level dim = mesh.dim - 1 zOffset = 0 while zOffset < depth: indIn = r < (max_distance + padWidth[ii]) nnz = int(np.sum(indIn)) if nnz > 0: mesh.insert_cells( np.c_[newLoc[indIn, :dim], newLoc[indIn, -1] - zOffset], np.ones(nnz) * mesh.max_level - ii, finalize=False, ) zOffset += dz depth -= dz * octree_levels[ii] if finalize: mesh.finalize() elif method.lower() == "box": # Define the data extent [bottom SW, top NE] bsw = np.min(xyz, axis=0) tne = np.max(xyz, axis=0) hs = np.asarray([h.min() for h in mesh.h]) hx = hs[0] hz = hs[-1] # Pre-calculate outer extent of each level # x_pad padWidth = np.cumsum(hx * octree_levels_padding * 2**np.arange(len(octree_levels))) if mesh.dim == 3: # y_pad hy = hs[1] padWidth = np.c_[padWidth, np.cumsum(hy * octree_levels_padding * 2**np.arange(len(octree_levels))), ] # Pre-calculate max depth of each level padWidth = np.c_[padWidth, np.cumsum(hz * np.maximum(octree_levels - 1, 0) * 2**np.arange(len(octree_levels))), ] levels = [] BSW = [] TNE = [] for ii, octZ in enumerate(octree_levels): if octZ > 0: levels.append(mesh.max_level - ii) BSW.append(bsw - padWidth[ii]) TNE.append(tne + padWidth[ii]) mesh.refine_box(BSW, TNE, levels, finalize=finalize) else: raise NotImplementedError("Only method= 'radial', 'surface'" " or 'box' have been implemented") return mesh
matData = sio.loadmat(sys.argv[1]) cellWidth = matData['cellWidth'] LonPos = matData['lon'].T * dtor LatPos = matData['lat'].T * dtor minCellWidth = cellWidth.min() meshDensityVsLatLon = (minCellWidth / cellWidth)**4 print 'minimum cell width in grid definition:', minCellWidth print 'maximum cell width in grid definition:', cellWidth.max() print 'cellWidth, south to north:' print cellWidth print 'meshDensityVsLatLon, south to north, smallest cell gets a 1:' print meshDensityVsLatLon LON, LAT = np.meshgrid(LonPos, LatPos) ds = nc4.Dataset(sys.argv[2], 'r+') meshDensity = ds.variables["meshDensity"][:] print "Preparing interpolation of meshDensity from lat/lon to mesh..." meshDensityInterp = interpolate.LinearNDInterpolator( np.vstack((LAT.ravel(), LON.ravel())).T, meshDensityVsLatLon.ravel()) print "Interpolating and writing meshDensity..." ds.variables['meshDensity'][:] = meshDensityInterp( np.vstack( (ds.variables['latCell'][:], np.mod(ds.variables['lonCell'][:] + np.pi, 2 * np.pi) - np.pi)).T) ds.close()
def active_from_xyz(mesh, xyz, grid_reference="CC", method="linear"): """Return boolean array indicating which cells are below surface For a set of locations defining a surface, **active_from_xyz** outputs a boolean array indicating which mesh cells like below the surface points. This method uses SciPy's interpolation routine to interpolate between location points defining the surface. Nearest neighbour interpolation is used for cells outside the convex hull of the surface points. Parameters ---------- mesh : discretize.TensorMesh or discretize.TreeMesh or discretize.CylindricalMesh Mesh object. If *mesh* is a cylindrical mesh, it must be symmetric xyz : (N, dim) numpy.ndarray Points defining the surface topography. grid_reference : {'CC', 'N'} Define where the cell is defined relative to surface. Choose between {'CC','N'} - If 'CC' is used, cells are active if their centers are below the surface. - If 'N' is used, cells are active if they lie entirely below the surface. method : {'linear', 'nearest'} Interpolation method for locations between the xyz points. Returns ------- (n_cells) numpy.ndarray of bool 1D mask array of *bool* for the active cells below xyz. Examples -------- Here we define the active cells below a parabola. We demonstrate the differences that appear when using the 'CC' and 'N' options for *reference_grid*. >>> import matplotlib.pyplot as plt >>> import numpy as np >>> from discretize import TensorMesh >>> from discretize.utils import active_from_xyz Determine active cells for a given mesh and topography >>> mesh = TensorMesh([5, 5]) >>> topo_func = lambda x: -3*(x-0.2)*(x-0.8)+.5 >>> topo_points = np.linspace(0, 1) >>> topo_vals = topo_func(topo_points) >>> active_cc = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='CC') >>> active_n = active_from_xyz(mesh, np.c_[topo_points, topo_vals], grid_reference='N') Plot visual representation .. collapse:: Expand to show scripting for plot >>> ax = plt.subplot(121) >>> mesh.plot_image(active_cc, ax=ax) >>> mesh.plot_grid(centers=True, ax=ax) >>> ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3') >>> ax.set_title("CC") >>> ax = plt.subplot(122) >>> mesh.plot_image(active_n, ax=ax) >>> mesh.plot_grid(nodes=True, ax=ax) >>> ax.plot(np.linspace(0,1), topo_func(np.linspace(0,1)), color='C3') >>> ax.set_title("N") >>> plt.show() """ try: if not mesh.is_symmetric: raise NotImplementedError( "Unsymmetric CylindricalMesh is not yet supported") except AttributeError: pass if grid_reference not in ["N", "CC"]: raise ValueError( "Value of grid_reference must be 'N' (nodal) or 'CC' (cell center)" ) dim = mesh.dim - 1 if mesh.dim == 3: if xyz.shape[1] != 3: raise ValueError( "xyz locations of shape (*, 3) required for 3D mesh") if method == "linear": tri2D = Delaunay(xyz[:, :2]) z_interpolate = interpolate.LinearNDInterpolator(tri2D, xyz[:, 2]) else: z_interpolate = interpolate.NearestNDInterpolator( xyz[:, :2], xyz[:, 2]) elif mesh.dim == 2: if xyz.shape[1] != 2: raise ValueError( "xyz locations of shape (*, 2) required for 2D mesh") z_interpolate = interpolate.interp1d(xyz[:, 0], xyz[:, 1], bounds_error=False, fill_value=np.nan, kind=method) else: if xyz.ndim != 1: raise ValueError( "xyz locations of shape (*, ) required for 1D mesh") if grid_reference == "CC": # this should work for all 4 mesh types... locations = mesh.cell_centers if mesh.dim == 1: active = np.zeros(mesh.nC, dtype="bool") active[np.searchsorted(mesh.cell_centers_x, xyz).max():] = True return active elif grid_reference == "N": try: # try for Cyl, Tensor, and Tree operations if mesh.dim == 3: locations = np.vstack([ mesh.cell_centers + (np.c_[-1, 1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), mesh.cell_centers + (np.c_[-1, -1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), mesh.cell_centers + (np.c_[1, 1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), mesh.cell_centers + (np.c_[1, -1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), ]) elif mesh.dim == 2: locations = np.vstack([ mesh.cell_centers + (np.c_[-1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), mesh.cell_centers + (np.c_[1, 1][:, None] * mesh.h_gridded / 2.0).squeeze(), ]) else: active = np.zeros(mesh.nC, dtype="bool") active[np.searchsorted(mesh.nodes_x, xyz).max():] = True return active except AttributeError: # Try for Curvilinear Mesh gridN = mesh.gridN.reshape((*mesh.vnN, mesh.dim), order="F") if mesh.dim == 3: locations = np.vstack([ gridN[:-1, 1:, 1:].reshape((-1, mesh.dim), order="F"), gridN[:-1, :-1, 1:].reshape((-1, mesh.dim), order="F"), gridN[1:, 1:, 1:].reshape((-1, mesh.dim), order="F"), gridN[1:, :-1, 1:].reshape((-1, mesh.dim), order="F"), ]) elif mesh.dim == 2: locations = np.vstack([ gridN[:-1, 1:].reshape((-1, mesh.dim), order="F"), gridN[1:, 1:].reshape((-1, mesh.dim), order="F"), ]) # Interpolate z values on CC or N z_xyz = z_interpolate(locations[:, :-1]).squeeze() # Apply nearest neighbour if in extrapolation ind_nan = np.isnan(z_xyz) if any(ind_nan): tree = cKDTree(xyz) _, ind = tree.query(locations[ind_nan, :]) z_xyz[ind_nan] = xyz[ind, dim] # Create an active bool of all True active = np.all((locations[:, dim] < z_xyz).reshape((mesh.nC, -1), order="F"), axis=1) return active.ravel()
def ax_scalp(v, channels, ax=None, annotate=True, vmin=None, vmax=None, cmap=cm.coolwarm, scalp_line_width=1, scalp_line_style='solid', chan_pos_list=CHANNEL_9_APPROX, interpolation='bilinear', fontsize=8): """ :param v: 输入通道对应的相关性值 :param channels: 通道的坐标位置 :param ax: 最后的存到ax这个框架里 :param annotate: :param vmin: :param vmax: :param cmap: :param scalp_line_width:头皮轮廓的线宽 :param scalp_line_style:头皮轮廓的线的形状 :param chan_pos_list:整个电极通道的位置 :param interpolation:双线性插值 :param fontsize:字体大小 :return: 一张头皮图 """ if ax is None: ax = plt.gca() assert len(v) == len(channels), "Should be as many values as channels" assert interpolation == 'bilinear' or interpolation == 'nearest' if vmin is None: # added by me ([email protected]) assert vmax is None vmin, vmax = -np.max(np.abs(v)), np.max(np.abs(v)) # what if we have an unknown channel?得到对应电极通道在列表中正确的位置 points = [get_channelpos(c, chan_pos_list) for c in channels] for c in channels: assert get_channelpos( c, chan_pos_list) is not None, ("Expect " + c + " to exist in positions") z = [v[i] for i in range(len(points))] # calculate the interpolation x = [i[0] for i in points] y = [i[1] for i in points] # interpolate the in-between values插值填充 xx = np.linspace(min(x), max(x), 1000) yy = np.linspace(min(y), max(y), 1000) if interpolation == 'bilinear': xx_grid, yy_grid = np.meshgrid(xx, yy) #生成网格点坐标矩阵 #对输入数据进行三角测量,并在每个执行线性重心插值的三角形上构造插值 f = interpolate.LinearNDInterpolator(list(zip(x, y)), z) # zz = f(xx_grid, yy_grid) else: assert interpolation == 'nearest' f = interpolate.NearestNDInterpolator(list(zip(x, y)), z) assert len(xx) == len(yy) zz = np.ones((len(xx), len(yy))) for i_x in xrange(len(xx)): for i_y in xrange(len(yy)): # somehow this is correct. don't know why :( zz[i_y, i_x] = f(xx[i_x], yy[i_y]) # zz[i_x,i_y] = f(xx[i_x], yy[i_y]) assert not np.any(np.isnan(zz)) # plot map开始画头皮图 image = ax.imshow(zz, vmin=vmin, vmax=vmax, cmap=cmap, extent=[min(x), max(x), min(y), max(y)], origin='lower', interpolation=interpolation) if scalp_line_width > 0: # paint the head ax.add_artist( plt.Circle((0, 0), 1, linestyle=scalp_line_style, linewidth=scalp_line_width, fill=False)) # add a nose ax.plot([-0.1, 0, 0.1], [1, 1.1, 1], color='black', linewidth=scalp_line_width, linestyle=scalp_line_style) # add ears _add_ears(ax, scalp_line_width, scalp_line_style) # add markers at channels positions # set the axes limits, so the figure is centered on the scalp ax.set_ylim([-1.05, 1.15]) ax.set_xlim([-1.15, 1.15]) # hide the frame and ticks ax.set_frame_on(False) ax.set_xticks([]) ax.set_yticks([]) # draw the channel names if annotate: for i in zip(channels, list(zip(x, y))): ax.annotate(" " + i[0], i[1], horizontalalignment="center", verticalalignment='center', fontsize=fontsize) ax.set_aspect(1) plt.show() return image
for m in trange(m_+1,desc='Day'): tau = np.busday_count(t_m[m], t_end)/252 if tau < tau_implvol[0]: tau = tau_implvol[0] for j in range(j_): # compute shadow yield x_y = ytm_shadowrates(np.array([y])) x_y = np.atleast_1d(x_y) # compute call option value v_call_thor[j, m] = \ call_option_value(x_proj[j, m, 0], x_y, tau, x_proj[j, m, 1:], m_moneyness, tau_implvol, k_strk, t_end, t_m[m]) # compute log-implied volatility at the moneyness log_sigma_atm[j, m] = \ interpolate.LinearNDInterpolator(points, x_proj[j, m, 1:])(*np.r_[tau, 0]) # - # ## [Step 2](https://www.arpm.co/lab/redirect.php?permalink=s_pricing_calloption-implementation-step02): Scenario-probability expectations and standard deviations # + mu_v = np.zeros(m_+1) sig_v = np.zeros(m_+1) for m in range(len(t_m)): mu_v[m], sig1 = meancov_sp(v_call_thor[:, m].reshape(-1, 1)) sig_v[m] = np.sqrt(sig1) # - # ## [Step 3](https://www.arpm.co/lab/redirect.php?permalink=s_pricing_calloption-implementation-step03): Save databases
def interpf(x, z, method=None, extrap=False, arrays=False): '''interpolate to find f, where z = f(*x) Input: x: an array of shape (npts, dim) or (npts,) z: an array of shape (npts,) method: string for kind of interpolator extrap: if True, extrapolate a bounding box (can reduce # of nans) arrays: if True, z = f(*x) is a numpy array; otherwise don't use arrays Output: interpolated function f, where z = f(*x) NOTE: if scipy is not installed, will use np.interp for 1D (non-rbf), or mystic's rbf otherwise. default method is 'nearest' for 1D and 'linear' otherwise. method can be one of ('rbf','linear','cubic', 'nearest','inverse','gaussian','multiquadric','quintic','thin_plate'). NOTE: if extrap is True, extrapolate using interpf with method='thin_plate' (or 'rbf' if scipy is not found). Alternately, any one of ('rbf', 'linear','cubic','nearest','inverse','gaussian','multiquadric', 'quintic','thin_plate') can be used. If extrap is a cost function z = f(x), then directly use it in the extrapolation. ''' #XXX: return f(*x) or f(x)? if arrays: _f, _fx = _to_array, _array else: _f, _fx = _to_nonarray, _nonarray if len(x) > len(z): # if len(x) < len(z), points are dropped x = x[:len(z)] # drop points x, z = extrapolate(x, z, method=extrap) x, z = _unique(x, z, sort=True) # avoid nan as first value #XXX: better choice than 'nearest'? if method is None: method = 'nearest' if x.ndim == 1 else 'linear' methods = dict(rbf=0, linear=1, nearest=2, cubic=3) functions = {0: 'multiquadric', 1: 'linear', 2: 'nearest', 3: 'cubic'} # also: ['thin_plate','inverse','gaussian','quintic'] kind = methods[method] if method in methods else None function = functions[kind] if (not kind is None) else method if kind is None: kind = 0 #XXX: or raise KeyError ? try: import scipy.interpolate as si except ImportError: if not kind == 0: # non-rbf if x.ndim == 1: # is 1D, so use np.interp import numpy as np return lambda xn: _fx(np.interp(xn, xp=x, fp=z)) kind = 0 # otherwise, utilize mystic's rbf si = _rbf if kind == 0: # 'rbf' import numpy as np return _f(si.Rbf(*np.vstack((x.T, z)), function=function, smooth=0)) elif x.ndim == 1: return _f( si.interp1d(x, z, fill_value='extrapolate', bounds_error=False, kind=method)) elif kind == 1: # 'linear' return _f(si.LinearNDInterpolator(x, z, rescale=False)) elif kind == 2: # 'nearest' return _f(si.NearestNDInterpolator(x, z, rescale=False)) #elif x.ndim == 1: # 'cubic' # return lambda xn: _fx(si.spline(x, z, xn)) return _f(si.CloughTocher2DInterpolator(x, z, rescale=False))
def __init__(self, x_values, y_values, z_values, classification, scalar_ans=None, x_vector_ans=None, y_vector_ans=None, z_vector_ans=None): """A class representing either a scalar or vector table. If a lookup table is to be used instead of a function for the model observing method, it is required to standardize the information given by the user's lookup table, as is the purpose of this class. Arguments --------- x_values : array_like The values of the x points for the table. Array must be parallel with y_values and z_values along with the scalar/vector answer. y_values : array_like The values of the x points for the table. Array must be parallel with x_values and z_values along with the scalar/vector answer. z_values : array_like The values of the x points for the table. Array must be parallel with x_values and y_values along with the scalar/vector answer. classification : string The classification of this table, either as a scalar lookup table or a vector lookup table. Should be one of: - 'scalar' A scalar based lookup table. - 'vector' A vector based lookup table. scalar_ans : array_like, {for | classification == 'scalar'} The scalar answers to the (x,y,z) point given by the input values. Must be parallel with x_values, y_values, and z_values. Ignored if classification == 'vector'. x_vector_ans : array_like, {for | classification == 'vector'} The x component of the answer vector that exists at the point (x,y,z) given by the input values. Must be parallel with x_values, y_values, and z_values along with other components. Ignored if classification == 'scalar' y_vector_ans : array_like, {for | classification == 'vector'} The y component of the answer vector that exists at the point (x,y,z) given by the input values. Must be parallel with x_values, y_values, and z_values along with other components. Ignored if classification == 'scalar' z_vector_ans : array_like, {for | classification == 'vector'} The z component of the answer vector that exists at the point (x,y,z) given by the input values. Must be parallel with x_values, y_values, and z_values along with other components. Ignored if classification == 'scalar' """ # Type check. x_values = Robust.valid.validate_float_array(x_values) y_values = Robust.valid.validate_float_array(y_values) z_values = Robust.valid.validate_float_array(z_values) classification = Robust.valid.validate_string(classification).lower() # Decide on the type before type checking. if (classification == 'scalar'): if (scalar_ans is not None): scalar_ans = Robust.valid.validate_float_array(scalar_ans) else: raise TypeError('Scalar answer array must be provided if ' 'table classification is set to scalar.' ' --Kyubey') elif (classification == 'vector'): if (x_vector_ans is not None): x_vector_ans = Robust.valid.validate_float_array(x_vector_ans) else: raise TypeError('The x component of the vector answer array ' 'must be provided if table classification is ' 'set to vector.' ' --Kyubey') if (y_vector_ans is not None): y_vector_ans = Robust.valid.validate_float_array(y_vector_ans) else: raise TypeError('The y component of the vector answer array ' 'must be provided if table classification is ' 'set to vector.' ' --Kyubey') if (z_vector_ans is not None): z_vector_ans = Robust.valid.validate_float_array(z_vector_ans) else: raise TypeError('The z component of the vector answer array ' 'must be provided if table classification is ' 'set to vector.' ' --Kyubey') else: raise Robust.InputError( 'Table classification must be one of the ' 'following: \n' 'scalar, vector \n' 'It is currently: < {table_cls} >' ' --Kyubey'.format(table_cls=classification)) # Precompute the Delaunay triangulation, this is done under the # assumption that the table should not be changed after data is # put into it. # pylint: disable=maybe-no-member try: Delanuay_tri = sp_spt.Delaunay( np.array([x_values, y_values, z_values]).T) except (TypeError, ValueError): raise except Exception: # If there is a Qhull error, we don't want to deal with it. Does # not currently know how to specify the QHull error on its own. Delanuay_tri = None # pylint: enable=maybe-no-member # Attempt to make the linear interpolators. if (Delanuay_tri is not None): if (classification == 'scalar'): linear_interp_scalar = \ sp_inter.LinearNDInterpolator(Delanuay_tri, scalar_ans,fill_value=0) # Just for safety reasons. linear_interp_x_axis = None linear_interp_y_axis = None linear_interp_z_axis = None elif (classification == 'vector'): linear_interp_x_axis = \ sp_inter.LinearNDInterpolator(Delanuay_tri, x_vector_ans, fill_value=0) linear_interp_y_axis = \ sp_inter.LinearNDInterpolator(Delanuay_tri, y_vector_ans, fill_value=0) linear_interp_z_axis = \ sp_inter.LinearNDInterpolator(Delanuay_tri, z_vector_ans, fill_value=0) # For safety reasons. linear_interp_scalar = None else: linear_interp_scalar = None linear_interp_x_axis = None linear_interp_y_axis = None linear_interp_z_axis = None # Assign variables. Depending on the actual self.x_values = x_values self.y_values = y_values self.z_values = z_values self.classification = classification self.scalar_ans = scalar_ans self.x_vector_ans = x_vector_ans self.y_vector_ans = y_vector_ans self.z_vector_ans = z_vector_ans self._Delanuay_triangulation = Delanuay_tri self._linear_interp_scalar = linear_interp_scalar self._linear_interp_x_axis = linear_interp_x_axis self._linear_interp_y_axis = linear_interp_y_axis self._linear_interp_z_axis = linear_interp_z_axis