def overlapping_bricks(self, candidates, map_petals=False): """Convert a list of potentially overlapping bricks into actual overlaps. Parameters ---------- candidates : :class:`list` A list of candidate bricks. map_petals : bool, optional If ``True`` a map of petal number to a list of overlapping bricks is returned. Returns ------- :class:`list` A list of Polygon objects. """ petals = self.petals() petal2brick = dict() bricks = list() for b in candidates: b_ra1, b_ra2 = self.brick_offset(b) brick_corners = np.array([[b_ra1, b.dec1], [b_ra2, b.dec1], [b_ra2, b.dec2], [b_ra1, b.dec2]]) brick = Polygon(brick_corners, closed=True, facecolor='r') for i, p in enumerate(petals): if brick.get_path().intersects_path(p.get_path()): brick.set_facecolor('g') if i in petal2brick: petal2brick[i].append(b.id) else: petal2brick[i] = [b.id] bricks.append(brick) if map_petals: return petal2brick return bricks
def select_vert(img): """ User-friendly approach which lets the user manually select the area corresponding to the verticle line. If not correctly selected the first time, the user can re-do this procedure. Close plots manually to continue in the procedure. Type Y for 'yes' or N for 'no' when deciding if the area is good. """ # Local variable which breaks loop if area of interest is selected well OK = False # Main while-loop while OK == False: # Plot image fig, ax = plt.subplots(figsize=(10, 10)) ax.imshow(img, cmap="gray") # Let user specify points coord = np.asarray(plt.ginput(4, show_clicks=True)) p = Polygon(coord, linewidth=1, edgecolor='r', facecolor='none') plt.gca().add_artist(p) # Include area of interest in plot plt.draw() plt.show() # Ask user to accept or reject the proposed area of interest val = input("Is the region correct ([Y]/n)?\n") # Break if OK, re-do if not if val == "Y" or val == "": OK = True """ Creates a mask which marks the vertical line based on the coordinates given by the user. """ x, y = np.meshgrid(np.arange(img.shape[0]), np.arange(img.shape[1]), indexing='xy') x, y = x.flatten(), y.flatten() pts = np.vstack((x, y)).T pts_t = tuple(map(tuple, pts)) mask = np.ones((img.shape[0], img.shape[1])) for (x, y) in pts_t: if p.get_path().contains_point((x, y)): mask[y][x] = 0 # Return mask which is the area of interest with value 1, 0 else return mask
def getCatalog(size=10000, surveyname=None): # dummy catalog: uniform on sphere # Marsaglia (1972) xyz = np.random.normal(size=(size, 3)) r = np.sqrt((xyz**2).sum(axis=1)) dec = np.arccos(xyz[:,2]/r) / skm.DEG2RAD - 90 ra = - np.arctan2(xyz[:,0], xyz[:,1]) / skm.DEG2RAD if surveyname is not None: from matplotlib.patches import Polygon # construct survey polygon ra_fp, dec_fp = skm.survey_register[surveyname].getFootprint() poly = Polygon(np.dstack((ra_fp,dec_fp))[0], closed=True) inside = [poly.get_path().contains_point(Point(ra_,dec_)) for (ra_,dec_) in zip(ra,dec)] ra = ra[inside] dec = dec[inside] return ra, dec
def overlapping_bricks(self, session, map_petals=False): """Perform a geometric calculation to find bricks that overlap a tile. Parameters ---------- session : :class:`sqlalchemy.orm.session.Session` Database connection. map_petals : bool, optional If ``True`` a map of petal number to a list of overlapping bricks is returned. Returns ------- :class:`list` If `map_petals` is ``False``, a list of :class:`~matplotlib.patches.Polygon` objects. Otherwise, a :class:`dict` mapping petal number to the :class:`~lvmspec.database.metadata.Brick` objects that overlap that petal. """ if self._brick_polygons is None and self._petal2brick is None: candidates = self._coarse_overlapping_bricks(session) petals = self.petals() self._petal2brick = dict() self._brick_polygons = list() for b in candidates: b_ra1, b_ra2 = self.brick_offset(b) brick_corners = np.array([[b_ra1, b.dec1], [b_ra2, b.dec1], [b_ra2, b.dec2], [b_ra1, b.dec2]]) brick_poly = Polygon(brick_corners, closed=True, facecolor='r') for i, p in enumerate(petals): if brick_poly.get_path().intersects_path(p.get_path()): brick_poly.set_facecolor('g') if i in self._petal2brick: self._petal2brick[i].append(b) else: self._petal2brick[i] = [b] self._brick_polygons.append(brick_poly) if map_petals: return self._petal2brick return self._brick_polygons
def plot_sat(sa_obj, **kwargs): import matplotlib.cm as mplcm import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import cartopy.crs as ccrs import cartopy.feature as cfeature import cmocean from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER import matplotlib.ticker as mticker stdvarname = sa_obj.stdvarname # sort out data/coordinates for plotting slons, slats = sa_obj.vars['longitude'], sa_obj.vars['latitude'] if 'col_obj' in kwargs.keys(): model_lats = kwargs['col_obj'].vars['model_lats'] model_lons = kwargs['col_obj'].vars['model_lons'] model_vals = kwargs['col_obj'].vars['model_values'] if (sa_obj.region in model_dict and 'mc_obj' in kwargs.keys()): grid_date = model_dict[sa_obj.region]['grid_date'] model_var_dict = kwargs['mc_obj'].vars # check region and determine projection if (sa_obj.region == 'global' or (sa_obj.region in region_dict['rect'] and 'boundinglat' in region_dict['rect'][sa_obj.region].keys())): # Polar Stereographic Projection polarproj = ccrs.NorthPolarStereo(central_longitude=0.0, true_scale_latitude=66, globe=None) projection = polarproj land = cfeature.GSHHSFeature(scale='i', levels=[1], facecolor=cfeature.COLORS['land']) else: if sa_obj.region in region_dict['rect']: latmin = region_dict['rect'][sa_obj.region]['llcrnrlat'] latmax = region_dict['rect'][sa_obj.region]['urcrnrlat'] lonmin = region_dict['rect'][sa_obj.region]['llcrnrlon'] lonmax = region_dict['rect'][sa_obj.region]['urcrnrlon'] elif sa_obj.region in region_dict['geojson']: latmin = np.min(sa_obj.vars['latitude']) - .5 latmax = np.max(sa_obj.vars['latitude']) + .5 lonmin = np.min(sa_obj.vars['longitude']) - .5 lonmax = np.max(sa_obj.vars['longitude']) + .5 elif sa_obj.region in region_dict['poly']: latmin = np.min(region_dict['poly'][sa_obj.region]['lats']) - .5 latmax = np.max(region_dict['poly'][sa_obj.region]['lats']) + .5 lonmin = np.min(region_dict['poly'][sa_obj.region]['lons']) - .5 lonmax = np.max(region_dict['poly'][sa_obj.region]['lons']) + .5 elif sa_obj.region in model_dict: # model bounds latmin = np.min(model_var_dict['latitude']) latmax = np.max(model_var_dict['latitude']) lonmin = np.min(model_var_dict['longitude']) lonmax = np.max(model_var_dict['longitude']) else: print("Error: Region not defined!") projection = ccrs.Mercator(central_longitude=(lonmin + lonmax) / 2., min_latitude=latmin, max_latitude=latmax, globe=None, latitude_true_scale=(latmin + latmax) / 2., false_easting=0.0, false_northing=0.0, scale_factor=None) land = cfeature.GSHHSFeature(scale='i', levels=[1], facecolor=cfeature.COLORS['land']) # make figure fig, ax = plt.subplots(nrows=1, ncols=1, subplot_kw=dict(projection=projection), figsize=(9, 9)) # plot domain extent if 'polarproj' not in locals(): ax.set_extent([lonmin, lonmax, latmin, latmax], crs=ccrs.PlateCarree()) else: ax.set_extent([-180, 180, 40, 90], crs=ccrs.PlateCarree()) # plot model domain if region is a model domain if (sa_obj.region in model_dict and len(model_var_dict['latitude'].shape) == 1): lenlons = len(model_var_dict['longitude'][:]) lenlats = len(model_var_dict['latitude'][:]) ax.plot([model_var_dict['longitude'][0]] * lenlats, model_var_dict['latitude'][:], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(model_var_dict['longitude'][:], [model_var_dict['latitude'][-1]] * lenlons, '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot([model_var_dict['longitude'][-1]] * lenlats, model_var_dict['latitude'][::-1], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(model_var_dict['longitude'][::-1], [model_var_dict['latitude'][0]] * lenlons, '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) if (sa_obj.region in model_dict and len(model_var_dict['latitude'].shape) == 2): ax.plot(model_var_dict['longitude'][0, :], model_var_dict['latitude'][0, :], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(model_var_dict['longitude'][-1, :], model_var_dict['latitude'][-1, :], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(model_var_dict['longitude'][:, 0], model_var_dict['latitude'][:, 0], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(model_var_dict['longitude'][:, -1], model_var_dict['latitude'][:, -1], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) # plot polygon if defined if sa_obj.region in region_dict['poly']: ax.plot(region_dict['poly'][sa_obj.region]['lons'], region_dict['poly'][sa_obj.region]['lats'], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) # colors cmap = cmocean.cm.amp if stdvarname == 'sea_surface_wave_significant_height': cmap = cmocean.cm.amp levels = [ 0, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3, 3.25, 3.5, 3.75, 4, 4.5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18 ] elif stdvarname == 'wind_speed': cmap = cmocean.cm.amp levels = [ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 36, 40, 44, 48 ] if 'cmap' in kwargs.keys(): cmap = kwargs['cmap'] if 'levels' in kwargs.keys(): levels = kwargs['levels'] # draw figure features mpl.rcParams['contour.negative_linestyle'] = 'solid' fs = 11 if sa_obj.region in quicklook_dict: if ('cm_levels' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['cm_levels'] is not None): levels = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['cm_levels'] if ('scl' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['scl'] is not None): scl = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['scl'] if ('icl' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['icl'] is not None): scl = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['icl'] # plot lats/lons gridcolor = 'gray' gl = ax.gridlines(draw_labels=True, crs=ccrs.PlateCarree(), linewidth=1, color=gridcolor, alpha=0.4, linestyle='-') gl.top_labels = False gl.right_labels = False gl.xformatter = LONGITUDE_FORMATTER gl.yformatter = LATITUDE_FORMATTER gl.xlabel_style = {'size': fs, 'color': gridcolor} gl.ylabel_style = {'size': fs, 'color': gridcolor} # - add coastline if 'polarproj' not in locals(): ax.add_geometries(land.intersecting_geometries( [lonmin, lonmax, latmin, latmax]), ccrs.PlateCarree(), facecolor=cfeature.COLORS['land'], edgecolor='black', linewidth=1) else: ax.add_geometries(land.intersecting_geometries([-180, 180, 0, 90]), ccrs.PlateCarree(), facecolor=cfeature.COLORS['land'], edgecolor='black', linewidth=1) # - add land color ax.add_feature(land, facecolor='burlywood', alpha=0.5) # scale colormap norm = mpl.colors.BoundaryNorm(levels, cmap.N) extend = 'neither' if 'extend' in kwargs.keys(): extend = kwargs['extend'] # - add satellite sc = ax.scatter( slons, slats, s=10, #c='k',#sa_obj.vars[sa_obj.stdvarname], c=sa_obj.vars[sa_obj.stdvarname], marker='o', edgecolor='face', cmap=cmocean.cm.amp, norm=norm, transform=ccrs.PlateCarree()) if 'col_obj' in kwargs.keys(): # - add satellite sc2 = ax.scatter( model_lons, model_lats, s=10, #c='r',#model_vals, c=model_vals, marker='o', edgecolor='face', cmap=cmocean.cm.amp, norm=norm, transform=ccrs.PlateCarree()) # - point of interests depending on region if ('quicklook_dict' in globals() and sa_obj.region in quicklook_dict and 'poi' in quicklook_dict[sa_obj.region]): for poi in quicklook_dict[sa_obj.region]['poi']: pname = quicklook_dict[sa_obj.region]['poi'][poi]['name'] plat = quicklook_dict[sa_obj.region]['poi'][poi]['lat'] plon = quicklook_dict[sa_obj.region]['poi'][poi]['lon'] scp = ax.scatter( plon, plat, s=20, c=quicklook_dict[sa_obj.region]['poi'][poi].get('color', 'b'), marker=quicklook_dict[sa_obj.region]['poi'][poi]['marker'], transform=ccrs.PlateCarree()) ax.text(plon, plat, pname, transform=ccrs.PlateCarree()) # - plot polygon if sa_obj.region in region_dict['poly']: ax.plot(region_dict['poly'][sa_obj.region]['lons'], region_dict['poly'][sa_obj.region]['lats'], 'k:', transform=ccrs.PlateCarree()) # plot polygon from geojson if defined if sa_obj.region in region_dict['geojson']: import geojson import matplotlib.patches as patches from matplotlib.patches import Polygon fstr = region_dict['geojson'][sa_obj.region]['fstr'] with open(fstr) as f: gj = geojson.load(f) fidx = region_dict['geojson'][sa_obj.region].get('fidx') if fidx is not None: geo = {'type': 'Polygon', 'coordinates':\ gj['features'][fidx]['geometry']['coordinates'][0]} poly = Polygon([tuple(l) for l in geo['coordinates'][0]], closed=True) ax.add_patch(\ patches.Polygon(\ poly.get_path().to_polygons()[0], transform=ccrs.PlateCarree(), facecolor='None', edgecolor='red', lw = 3, alpha=0.2)) else: for i in range(len(gj['features'])): geo = {'type': 'Polygon', 'coordinates':\ gj['features'][i]['geometry']['coordinates'][0]} poly = Polygon([tuple(l) for l in geo['coordinates'][0]], closed=True) ax.add_patch(\ patches.Polygon(\ poly.get_path().to_polygons()[0], transform=ccrs.PlateCarree(), facecolor='None', edgecolor='red', lw = 3, alpha=0.2)) # - colorbar axins = inset_axes( ax, width="5%", # width = 5% of parent_bbox width height="100%", # height : 50% loc='lower left', bbox_to_anchor=(1.01, 0., 1, 1), bbox_transform=ax.transAxes, borderpad=0, ) cbar = fig.colorbar(sc, cax=axins) cbar.ax.set_ylabel(sa_obj.stdvarname + ' [' + variable_info[sa_obj.varalias]['units'] + ']') cbar.ax.tick_params(labelsize=fs) plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9) ax.set_title(sa_obj.mission + ' with ' + str(len(sa_obj.vars[sa_obj.stdvarname])) + ' footprints: ' + '\n' + sa_obj.sdate.strftime("%Y-%m-%d %H:%M:%S UTC") + ' to ' + sa_obj.edate.strftime("%Y-%m-%d %H:%M:%S UTC"), fontsize=fs) # - save figure if ('savepath' in kwargs.keys() and kwargs['savepath'] != None): plt.savefig(kwargs['savepath'] + '/' + 'satellite_coverage_for_' + sa_obj.region + '_from_' + sa_obj.sdate.strftime("%Y%m%d") + 'T' + sa_obj.sdate.strftime("%H") + 'Z' + '_to_' + sa_obj.edate.strftime("%Y%m%d") + 'T' + sa_obj.edate.strftime("%H") + 'Z' + '.png', format='png', dpi=300) # - show figure if ('showfig' in kwargs.keys() and kwargs['showfig'] == True): plt.show()
def comp_fig(sa_obj=None, mc_obj=None, coll_obj=None, **kwargs): import matplotlib.cm as mplcm import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import cartopy.crs as ccrs import cartopy.feature as cfeature import cmocean from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER import matplotlib.ticker as mticker from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter from cartopy.mpl.ticker import LatitudeLocator, LongitudeLocator # sort out data/coordinates for plotting sat = "NA" model = "NA" """ If sa_obj is not None get satellite_altimetry data for plotting """ if sa_obj is not None: slons, slats = sa_obj.vars['longitude'], sa_obj.vars['latitude'] svar = sa_obj.vars[sa_obj.stdvarname] stdvarname = sa_obj.stdvarname sat = sa_obj.mission """ If mc_obj is not None get model data for plotting """ if mc_obj is not None: mlons = mc_obj.vars['longitude'] mlats = mc_obj.vars['latitude'] mvar = mc_obj.vars[mc_obj.stdvarname] # inflate coords if regular lat/lon grid if (len(mlons.shape) == 1): mlons, mlats = np.meshgrid(mlons, mlats) stdvarname = mc_obj.stdvarname model = mc_obj.model """ If sa_obj is not None get satellite_altimetry data for plotting """ if sa_obj is None: slons, slats = coll_obj.vars['obs_lons'], coll_obj.vars['obs_lats'] svar = coll_obj.vars['obs_values'] """ Get all misc in **kwargs """ """ Prepare plotting """ # determine region bounds if sa_obj is not None: if sa_obj.region in region_dict['rect']: latmin = region_dict['rect'][sa_obj.region]['llcrnrlat'] latmax = region_dict['rect'][sa_obj.region]['urcrnrlat'] lonmin = region_dict['rect'][sa_obj.region]['llcrnrlon'] lonmax = region_dict['rect'][sa_obj.region]['urcrnrlon'] elif sa_obj.region in region_dict['geojson']: latmin = np.min(sa_obj.vars['latitude']) - .5 latmax = np.max(sa_obj.vars['latitude']) + .5 lonmin = np.min(sa_obj.vars['longitude']) - .5 lonmax = np.max(sa_obj.vars['longitude']) + .5 elif sa_obj.region in region_dict['poly']: latmin = np.min(region_dict['poly'][sa_obj.region]['lats']) - .5 latmax = np.max(region_dict['poly'][sa_obj.region]['lats']) + .5 lonmin = np.min(region_dict['poly'][sa_obj.region]['lons']) - .5 lonmax = np.max(region_dict['poly'][sa_obj.region]['lons']) + .5 elif sa_obj.region in model_dict: # model bounds latmin = np.min(mlats) latmax = np.max(mlats) lonmin = np.min(mlons) lonmax = np.max(mlons) else: print("Error: Region not defined!") elif (sa_obj is None and mc_obj is not None): # model bounds latmin = np.min(mlats) latmax = np.max(mlats) lonmin = np.min(mlons) lonmax = np.max(mlons) # determine projection """ Here, a routine is needed to determine a suitable projection. As for now, there is Mercator as default. """ projection_default = ccrs.Mercator( central_longitude=(lonmin + lonmax) / 2., min_latitude=latmin, max_latitude=latmax, globe=None, latitude_true_scale=(latmin + latmax) / 2., false_easting=0.0, false_northing=0.0, scale_factor=None) projection = kwargs.get('projection', projection_default) land = cfeature.GSHHSFeature(scale='i', levels=[1], facecolor=cfeature.COLORS['land']) # make figure fig, ax = plt.subplots(nrows=1, ncols=1, subplot_kw=dict(projection=projection), figsize=(9, 9)) # plot domain extent ax.set_extent([lonmin, lonmax, latmin, latmax], crs=ccrs.PlateCarree()) # plot model domain if model is available if mc_obj is not None: ax.plot(mlons[0, :], mlats[0, :], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(mlons[-1, :], mlats[-1, :], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(mlons[:, 0], mlats[:, 0], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) ax.plot(mlons[:, -1], mlats[:, -1], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) # plot polygon if defined if sa_obj.region in region_dict['poly']: ax.plot(region_dict['poly'][sa_obj.region]['lons'], region_dict['poly'][sa_obj.region]['lats'], '-', transform=ccrs.PlateCarree(), color='gray', linewidth=2) # plot polygon from geojson if defined if sa_obj.region in region_dict['geojson']: import geojson import matplotlib.patches as patches from matplotlib.patches import Polygon fstr = region_dict['geojson'][sa_obj.region]['fstr'] with open(fstr) as f: gj = geojson.load(f) fidx = region_dict['geojson'][sa_obj.region].get('fidx') if fidx is not None: geo = {'type': 'Polygon', 'coordinates':\ gj['features'][fidx]['geometry']['coordinates'][0]} poly = Polygon([tuple(l) for l in geo['coordinates'][0]], closed=True) ax.add_patch(\ patches.Polygon(\ poly.get_path().to_polygons()[0], transform=ccrs.PlateCarree(), facecolor='None', edgecolor='red', lw = 3, alpha=0.2)) else: for i in range(len(gj['features'])): geo = {'type': 'Polygon', 'coordinates':\ gj['features'][i]['geometry']['coordinates'][0]} poly = Polygon([tuple(l) for l in geo['coordinates'][0]], closed=True) ax.add_patch(\ patches.Polygon(\ poly.get_path().to_polygons()[0], transform=ccrs.PlateCarree(), facecolor='None', edgecolor='red', lw = 3, alpha=0.2)) # colors if stdvarname == 'sea_surface_wave_significant_height': cmap = cmocean.cm.amp levels = [ 0, 0.25, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3, 3.25, 3.5, 3.75, 4, 4.5, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18 ] elif stdvarname == 'wind_speed': cmap = cmocean.cm.amp levels = [ 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 36, 40, 44, 48 ] if 'cmap' in kwargs.keys(): cmap = kwargs['cmap'] if 'levels' in kwargs.keys(): levels = kwargs['levels'] if 'scl' in kwargs.keys(): scl = kwargs['scl'] else: scl = 18 if 'icl' in kwargs.keys(): icl = kwargs['icl'] else: icl = 1 norm = mpl.colors.BoundaryNorm(levels, cmap.N) extend = 'neither' if 'extend' in kwargs.keys(): extend = kwargs['extend'] if sa_obj.region in quicklook_dict: if ('cm_levels' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['cm_levels'] is not None): levels = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['cm_levels'] if ('scl' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['scl'] is not None): scl = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['scl'] if ('icl' in \ quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]\ and quicklook_dict[sa_obj.region]['varspec'][sa_obj.varalias]['icl'] is not None): scl = quicklook_dict[sa_obj.region]['varspec']\ [sa_obj.varalias]['icl'] # draw figure features mpl.rcParams['contour.negative_linestyle'] = 'solid' fs = 11 # plot lats/lons gridcolor = 'gray' gl = ax.gridlines(draw_labels=True, crs=ccrs.PlateCarree(), linewidth=1, color=gridcolor, alpha=0.4, linestyle='-') gl.top_labels = False gl.right_labels = False gl.xlabel_style = {'size': fs, 'color': gridcolor} gl.ylabel_style = {'size': fs, 'color': gridcolor} #gl.xlocator = mticker.FixedLocator([-180, -45, 0, 45, 180]) gl.xlocator = LongitudeLocator() gl.ylocator = LatitudeLocator() gl.xformatter = LongitudeFormatter() gl.yformatter = LatitudeFormatter() # - model contours im = ax.contourf(mlons, mlats, mvar, levels=levels, transform=ccrs.PlateCarree(), cmap=cmocean.cm.amp, norm=norm, extend=extend) imc = ax.contour(mlons, mlats, mvar, levels=levels[scl::icl], transform=ccrs.PlateCarree(), colors='w', linewidths=0.3) ax.clabel(imc, fmt='%2d', colors='w', fontsize=fs) # - add coastline ax.add_geometries(land.intersecting_geometries( [lonmin, lonmax, latmin, latmax]), ccrs.PlateCarree(), facecolor=cfeature.COLORS['land'], edgecolor='black', linewidth=1) # - add land color ax.add_feature(land, facecolor='burlywood', alpha=0.5) # - add satellite if len(slats) > 0: sc = ax.scatter(slons, slats, s=10, c=svar, marker='o', edgecolor='face', cmap=cmocean.cm.amp, norm=norm, transform=ccrs.PlateCarree()) # - point of interests depending on region if ('quicklook_dict' in globals() and sa_obj.region in quicklook_dict and 'poi' in quicklook_dict[sa_obj.region]): for poi in quicklook_dict[sa_obj.region]['poi']: pname = quicklook_dict[sa_obj.region]['poi'][poi]['name'] plat = quicklook_dict[sa_obj.region]['poi'][poi]['lat'] plon = quicklook_dict[sa_obj.region]['poi'][poi]['lon'] scp = ax.scatter( plon, plat, s=20, c='b', marker=quicklook_dict[sa_obj.region]['poi'][poi]['marker'], transform=ccrs.PlateCarree()) ax.text(plon, plat, pname, transform=ccrs.PlateCarree()) # - colorbar axins = inset_axes( ax, width="5%", # width = 5% of parent_bbox width height="100%", # height : 50% loc='lower left', bbox_to_anchor=(1.01, 0., 1, 1), bbox_transform=ax.transAxes, borderpad=0, ) cbar = fig.colorbar(im, cax=axins) cbar.ax.set_ylabel(stdvarname + ' [' + variable_info[sa_obj.varalias]['units'] + ']', size=fs) cbar.ax.tick_params(labelsize=fs) plt.subplots_adjust(bottom=0.1, right=0.8, top=0.9) # - title ax.set_title( model + ' model time step: ' + coll_obj.vars['valid_date'][0].strftime("%Y-%m-%d %H%M:%S UTC") + '\n' + sat + ' coverage \n from ' + coll_obj.vars['datetime'][0].strftime("%Y-%m-%d %H:%M:%S UTC") + ' to ' + coll_obj.vars['datetime'][-1].strftime("%Y-%m-%d %H:%M:%S UTC"), fontsize=fs) # - save figure if ('savepath' in kwargs.keys() and kwargs['savepath'] != None): plt.savefig(kwargs['savepath'] + '/' + model + '_vs_satellite_' + coll_obj.vars['valid_date'][0].strftime("%Y%m%d") + 'T' + coll_obj.vars['valid_date'][0].strftime("%H") + 'Z.png', format='png', dpi=200) # - show figure if ('showfig' in kwargs.keys() and kwargs['showfig'] == True): plt.show()
def parameter_pdf(parameter_space, fig_comment='', mmi_obs=None, limits_filename=None, bbox=None, localities_file=None, plot_additions=None): """Calculate a pdf for parameter values based on the uncertainty model """ xlabel_dict = { 'mag': 'Magnitude ($M_w$)', 'longitude': 'Longitude', 'latitude': 'Latitude', 'depth': 'Depth (km)', 'strike': 'Strike ($^\circ$)', 'dip': 'Dip ($^\circ$)' } parameter_dict = { 'mag': 0, 'longitude': 1, 'latitude': 2, 'depth': 3, 'strike': 4, 'dip': 5 } parameter_pdf_sums = {} parameter_pdf_values = {} plt.clf() fig = plt.figure(figsize=(16, 8)) #, tight_layout=True) gs = plt.GridSpec(2, 4) for key, value in parameter_dict.iteritems(): unique_vals = np.unique(parameter_space[value]) pdf_sums = [] bin = False # Determine if we need to bin values of strike and dip (i.e. # for non-area sources if key == 'strike': if len(unique_vals) > 24: bin = True if megathrust or slab: bins = np.arange(0, 360, 5.) else: bins = np.arange(0, 360, 15.) if key == 'dip' or key == 'depth': if len(unique_vals) > 4: bin = True if key == 'dip': if (max(parameter_space[value]) - min(parameter_space[value])) > 10.0: bins = np.arange(min(parameter_space[value]), max(parameter_space[value]) + 5, 5.) else: bins = np.arange(min(parameter_space[value]), max(parameter_space[value]) + 1, 1.) elif key == 'depth': if max(parameter_space[value]) > 80.0: bins = np.arange(min(parameter_space[value]), max(parameter_space[value]) + 20, 20.) else: bins = np.arange(0.0, max(parameter_space[value]) + 5.0, 5.) if bin: # Calculate as histogram hist, bins = np.histogram(parameter_space[value], bins) # align to bin centre for plotting #bin_width = bins[1] - bins[0] unique_vals = [] for i, edge in enumerate(bins): try: ind = np.intersect1d( np.where(parameter_space[value] >= edge), np.where(parameter_space[value] < bins[i + 1])) except IndexError: ind = np.where(parameter_space[value] >= edge)[ 0] #[0] to get array from tuple pdf_sum = 0 for index in ind: # likelihood = np.power((1/(self.sigma*np.sqrt(2*np.pi))), len(self.mmi_obs)) * \ # np.exp((-1/2)*((self.sum_squares_list[index]/self.sigma**2))) posterior_prob = parameter_space[7][index] pdf_sum += posterior_prob #pdf_sum = np.sum(self.uncert_fun.pdf(self.rmse[ind])) pdf_sums.append(pdf_sum) unique_vals.append(edge) # + bin_width) else: # Use raw values for val in unique_vals: # Just unique values, as pdf sums are repeated ind = np.argwhere(parameter_space[value] == val) pdf_sum = 0 for index in ind: # likelihood = np.power((1/(self.sigma*np.sqrt(2*np.pi))), len(self.mmi_obs)) * \ # np.exp((-1/2)*((self.sum_squares_list[index]/self.sigma**2))) posterior_prob = parameter_space[7][index] pdf_sum += posterior_prob #pdf_sum = np.sum(self.uncert_fun.pdf(self.rmse[ind])) pdf_sums.append(pdf_sum) pdf_sums = np.array(pdf_sums) # print 'pdf_sums', pdf_sums # print 'sum', sum(pdf_sums) parameter_pdf_sums[key] = pdf_sums parameter_pdf_values[key] = unique_vals # Get the best-fit value for plotting on top index = np.argmin(parameter_space[6]) best_fit_x = parameter_space[value][index] index_posterior = np.argmax(parameter_space[7]) best_fit_x_posterior = parameter_space[value][index_posterior] #y_index = np.where(unique_vals == best_fit_x)[0] try: y_index = (np.abs(unique_vals - best_fit_x)).argmin()[0] except IndexError: y_index = (np.abs(unique_vals - best_fit_x)).argmin() best_fit_y = pdf_sums[y_index] try: y_index_posterior = (np.abs(unique_vals - best_fit_x_posterior)).argmin()[0] except IndexError: y_index_posterior = (np.abs(unique_vals - best_fit_x_posterior)).argmin() best_fit_y_posterior = pdf_sums[y_index_posterior] # Now calculate the range that contains 95% of the distribution # Sort the pdf values, descending pdf_sums_flat = pdf_sums.flatten() try: unique_vals_flat = unique_vals.flatten() except AttributeError: unique_vals_flat = np.array(unique_vals).flatten() sorted_probs_args = np.argsort(pdf_sums_flat)[::-1] #.argsort() sorted_probs = pdf_sums_flat[sorted_probs_args] sorted_values = unique_vals_flat[sorted_probs_args] sum_probs = sum(sorted_probs) prob_limit = 0.95 * sum_probs print 'Sum probs, should be 1', sum_probs print prob_limit prob_sum = 0.0 for pi, prob_val in enumerate(sorted_probs): if prob_sum > prob_limit: prob_index = pi break else: prob_sum += prob_val values_in_bounds = sorted_values[0:prob_index] min_bound = min(values_in_bounds) max_bound = max(values_in_bounds) print 'min_bound', min_bound print 'max_bound', max_bound # Get plot additions to plot on top # if plot_additions is not None: # x_addition = plot_additions[key] # try: # y_index = np.where(unique_(np.abs(unique_vals - best_fit_x)).argmin()[0] # except IndexError: # y_index = (np.abs(unique_vals - best_fit_x)).argmin() # best_fit_y = pdf_sums[y_index] # Now plot the results try: width = unique_vals[1] - unique_vals[ 0] # Scale width by discretisation except IndexError: # if only one value width = 1.0 if key == 'strike': # Plot as rose diagram ax = fig.add_subplot(gs[0, 3], projection='polar') ax.bar(np.deg2rad(unique_vals), pdf_sums, width=np.deg2rad(width), bottom=0.0, align='center', color='0.5', edgecolor='k') ax.scatter(np.deg2rad(best_fit_x), best_fit_y, marker='*', c='#696969', edgecolor='k', s=100, zorder=10) ax.scatter(np.deg2rad(best_fit_x_posterior), best_fit_y_posterior, marker='*', c='w', edgecolor='k', s=500, zorder=9) ax.set_theta_zero_location('N') ax.set_theta_direction(-1) ax.set_thetagrids(np.arange(0, 360, 15)) # Define grids intervals for radial axis if max(pdf_sums) < 0.11: r_int = 0.02 else: r_int = 0.2 ax.set_rgrids(np.arange(r_int, max(pdf_sums) + 0.01, r_int), angle=np.deg2rad(7.5)) #, weight= 'black') ax.set_xlabel(xlabel_dict[key]) ax.text(-0.07, 1.02, 'c)', transform=ax.transAxes, fontsize=14) elif key == 'mag': ax = fig.add_subplot(gs[0, 2]) ymax = max(pdf_sums) * 1.1 # lbx = [self.min_mag, self.min_mag] # lby = [0, ymax] # ubx = [self.max_mag, self.max_mag] # uby = [0, ymax] elif key == 'depth': ax = fig.add_subplot(gs[1, 3]) ymax = max(pdf_sums) * 1.1 # lbx = [self.min_depth, self.min_depth] # lby = [0, ymax] # ubx = [self.max_depth, self.max_depth] # uby = [0, ymax] elif key == 'dip': ax = fig.add_subplot(gs[1, 2]) ymax = max(pdf_sums) * 1.1 # lbx = [self.min_dip, self.min_dip] # lby = [0, ymax] # ubx = [self.max_dip, self.max_dip] # uby = [0, ymax] if key == 'mag' or key == 'dip' or key == 'depth': ax.bar(unique_vals, pdf_sums, width, align='center', color='0.5') ax.scatter(best_fit_x, best_fit_y, marker='*', c='#696969', edgecolor='k', s=100, zorder=11) ax.scatter(best_fit_x_posterior, best_fit_y_posterior, marker='*', c='w', edgecolor='k', s=500, zorder=9) if min_bound != max_bound: ax.plot([min_bound, min_bound], [0.0, ymax], linewidth=0.9, linestyle='--', c='k') ax.plot([max_bound, max_bound], [0.0, ymax], linewidth=0.9, linestyle='--', c='k') if plot_additions is not None: try: x_addition = plot_additions[key] y_addition = best_fit_y_posterior * 1.03 except KeyError: x_addition = None if x_addition is not None: ax.scatter(x_addition, y_addition, marker='*', c='b', edgecolor='k', s=200, zorder=10) #if key != 'latitude' and key != 'longitude': # ax.plot(lbx, lby, color='k') # ax.plot(ubx, uby, color='k') ax.set_ylim(0, ymax) ax.set_ylabel('Probability') ax.set_xlabel(xlabel_dict[key]) if key == 'mag': ax.text(-0.07, 1.02, 'b)', transform=ax.transAxes, fontsize=14) ax.set_xlim((min(unique_vals) - 0.4), (max(unique_vals) + 0.2)) if key == 'dip': ax.text(-0.07, 1.02, 'd)', transform=ax.transAxes, fontsize=14) if key == 'depth': ax.text(-0.07, 1.02, 'e)', transform=ax.transAxes, fontsize=14) # Now plot a map of location uncertainty # First get 2D pdf pdf_sums = [] all_lons = [] all_lats = [] for i, lon in enumerate(parameter_space[parameter_dict['longitude']]): if lon in all_lons: continue else: lat = parameter_space[parameter_dict['latitude']][i] # lat = self.parameter_pdf_values['latitude'][i] index = np.intersect1d(np.argwhere(parameter_space[parameter_dict['longitude']]==lon), \ np.argwhere(parameter_space[parameter_dict['latitude']]==lat)) pdf_sum = np.sum( parameter_space[7][index]) #uncert_fun.pdf(rmse[index])) pdf_sums.append(pdf_sum) all_lons.append(lon) all_lats.append(lat) # Normalise pdf sums pdf_sums = np.array(pdf_sums) # pdf_sums = pdf_sums/np.sum(pdf_sums) parameter_pdf_sums['lon_lat'] = pdf_sums all_lons = np.array(all_lons) all_lats = np.array(all_lats) # Get best fit value index = np.argmin(parameter_space[6]) best_fit_lon = parameter_space[parameter_dict['longitude']][index] best_fit_lat = parameter_space[parameter_dict['latitude']][index] index_posterior = np.argmax(parameter_space[7]) best_fit_lon_posterior = parameter_space[ parameter_dict['longitude']][index_posterior] best_fit_lat_posterior = parameter_space[ parameter_dict['latitude']][index_posterior] ax = fig.add_subplot(gs[:, 0:2]) bbox = bbox.split('/') if bbox is not None: minlon = float(bbox[0]) maxlon = float(bbox[1]) minlat = float(bbox[2]) maxlat = float(bbox[3]) else: minlon = min(parameter_pdf_values['longitude']) maxlon = max(parameter_pdf_values['longitude']) minlat = min(parameter_pdf_values['latitude']) maxlat = max(parameter_pdf_values['latitude']) lat_0 = minlat + (maxlat - minlat) / 2. lon_0 = minlon + (maxlon - minlon) / 2. m = Basemap(projection='tmerc', lat_0=lat_0, lon_0=lon_0, llcrnrlon=minlon, llcrnrlat=minlat, urcrnrlon=maxlon, urcrnrlat=maxlat, resolution='i') m.drawcoastlines(linewidth=0.5, color='k') m.drawcountries(color='0.2') m.drawstates(color='0.2') if maxlon - minlon < 2: gridspace = 0.5 elif maxlon - minlon < 3: gridspace = 1.0 elif maxlon - minlon < 7: gridspace = 2.0 else: gridspace = 5.0 m.drawparallels(np.arange(-90., 90., gridspace), labels=[1, 0, 0, 0], fontsize=10, dashes=[2, 2], color='0.5', linewidth=0.5) m.drawmeridians(np.arange(0., 360., gridspace), labels=[0, 0, 1, 0], fontsize=10, dashes=[2, 2], color='0.5', linewidth=0.5) max_val = max(pdf_sums) * 1.1 print 'pdf_sums', pdf_sums clevs = np.arange(0.0, max_val, (max_val / 50)) # clevs = np.arange(0.0,max(pdf_sums),(max_val/50)) cmap = plt.get_cmap('gray_r') # Adjust resolution to avoid memory intense interpolations res = max((maxlon - minlon) / 50., (maxlat - minlat) / 50.) #75 #30 #50 xy = np.mgrid[minlon:maxlon:res, minlat:maxlat:res] xx, yy = np.meshgrid(xy[0, :, 0], xy[1][0]) griddata = interpolate.griddata((all_lons, all_lats), pdf_sums, (xx, yy), method='nearest') # nearest # linears # now plot filled contours of pdf cs = m.contourf(xx, yy, griddata, clevs, cmap=cmap, vmax=max_val, vmin=0.0, latlon=True) for c in cs.collections: # Fix white space on contour levels for pdf images c.set_edgecolor("face") # Mask areas outside of source model if limits_filename is not None: limits_data = np.genfromtxt(limits_filename, delimiter=',') limits_x = limits_data[:, 0] limits_y = limits_data[:, 1] limits_x, limits_y = m(limits_x, limits_y) # Convert to map coordinates poly = Polygon(np.c_[limits_x, limits_y], closed=True) clippath = poly.get_path() ax = plt.gca() patch = PathPatch(clippath, transform=ax.transData, facecolor='none', linewidth=0.4, linestyle='--') print 'Adding patch' ax.add_patch(patch) for contour in cs.collections: contour.set_clip_path(patch) # Now add some locations if localities_file is not None: f_in = open(localities_file) loc_lon = [] loc_lat = [] loc_name = [] for line in f_in.readlines(): row = line.split() loc_lon.append(float(row[0])) loc_lat.append(float(row[1])) loc_name.append(row[2]) loc_sp = m.scatter(loc_lon, loc_lat, c='k', s=40, marker='s', latlon=True) texts = [] for label, x, y in zip(loc_name, loc_lon, loc_lat): x, y = m(x, y) texts.append( plt.text(x, y, label, fontsize=14, color='0.4', zorder=20)) # adjust_text(texts, only_move='xy', # arrowprops=dict(arrowstyle="-", # color='k', lw=0.5)) # Now add historical points on top if mmi_obs is not None: clevs = np.arange(0.5, 9.5, 1.0) cmap = plt.get_cmap('YlOrRd', 7) mmi_labels = [] for obs in mmi_obs[:, 2]: mmi_labels.append(write_roman(int(obs))) sp = m.scatter(mmi_obs[:, 0], mmi_obs[:, 1], c=mmi_obs[:, 2], cmap=cmap, vmin=1.5, vmax=8.5, s=30, latlon=True) sp_ticks = np.arange(1, 9, 1) # Label only if there aren't too many to avoid plots being too busy if len(mmi_obs[:, 2]) < 20: # texts = [] for label, x, y in zip(mmi_labels, mmi_obs[:, 0], mmi_obs[:, 1]): x, y = m(x, y) texts.append(plt.text(x, y, label, fontsize=10)) if len(texts) > 0: adjust_text(texts, only_move='xy', arrowprops=dict(arrowstyle="-", color='k', lw=0.5)) # Divide the axes to make the colorbar locatable to right of maps #divider = make_axes_locatable(ax) #cax = divider.append_axes("right", size="5%", pad=0.05) # plt.colorbar(im, cax=cax) #fig.add_axes(cax) cbar1 = m.colorbar(sp, ticks=sp_ticks, location='right', pad=0.2) cbar1.ax.set_ylabel('MMI') # Now add best-fit location on top m.scatter( best_fit_lon, best_fit_lat, marker='*', facecolor='none', #c='#696969', edgecolor='k', s=100, zorder=10, latlon=True) m.scatter(best_fit_lon_posterior, best_fit_lat_posterior, marker='*', facecolor='none', edgecolor='k', s=500, zorder=9, latlon=True) #m.text(0.05, 0.95, 'c)', transform=ax.transAxes, fontsize=14) if plot_additions is not None: try: x_addition = plot_additions['longitude'] y_addition = plot_additions['latitude'] except KeyError: x_addition = None if x_addition is not None: m.scatter(x_addition, y_addition, marker='*', c='b', edgecolor='k', s=200, zorder=10, latlon=True) plt.annotate('a)', xy=(-0.01, 1.01), xycoords='axes fraction', fontsize=14) print 'max_val', max_val if max_val < 0.0000001: loc_int = 0.00000001 elif max_val < 0.000001: loc_int = 0.0000001 elif max_val < 0.00001: loc_int = 0.000001 elif max_val < 0.0001: loc_int = 0.00001 elif max_val < 0.001: loc_int = 0.0001 elif max_val < 0.01: loc_int = 0.001 elif max_val < 0.1: loc_int = 0.01 else: loc_int = 0.1 ticks = np.arange(0.0, max_val * 1.1, loc_int) cbar = m.colorbar(cs, ticks=ticks, location='bottom') #orientation='horizontal') cbar.ax.set_xlabel('Probability') figname = '%s_all_parameter_pdf.png' % (fig_comment) figname = figname.replace('()', '') plt.tight_layout() plt.savefig(figname, dpi=600, format='png', bbox_inches='tight')
def map_xpt(self, xpt, magx, xpt_ID='xpt1', visual=False, verbose=False): if xpt_ID == 'xpt1': self.analyze_saddle(xpt, xpt_ID) self._set_function('rho', 'cw') NS_buffer = self.NSEW_lookup[xpt_ID]['coor'] N_sol = solve_ivp(self.function, (0, self.dt), NS_buffer['N'], method='LSODA', first_step=self.first_step, max_step=self.max_step, rtol=1e-13, atol=1e-12).y S_sol = solve_ivp(self.function, (0, self.dt), NS_buffer['S'], method='LSODA', first_step=self.first_step, max_step=self.max_step, rtol=1e-13, atol=1e-12).y N_guess = (N_sol[0][-1], N_sol[1][-1]) S_guess = (S_sol[0][-1], S_sol[1][-1]) r_bounds = (self.grid.rmin, self.grid.rmax) z_bounds = (self.grid.zmin, self.grid.zmax) N_minimizer = minimize(self.PsiCostFunc, N_guess, method='L-BFGS-B', jac=self.grid.Gradient, bounds=[r_bounds, z_bounds]).x S_minimizer = minimize(self.PsiCostFunc, S_guess, method='L-BFGS-B', jac=self.grid.Gradient, bounds=[r_bounds, z_bounds]).x if (np.linalg.norm(N_minimizer - magx) >= self.eps): self.flip_NSEW_lookup(xpt_ID) self.config = 'LSN' if self.NSEW_lookup['xpt1']['coor']['N'][ 1] > xpt[1] else 'USN' elif xpt_ID == 'xpt2': from matplotlib.patches import Polygon def cb_region_check(xk): if core_polygon.get_path().contains_point(xk): raise RegionEntered(message='# Entered Core...', region='Core') elif pf_polygon.get_path().contains_point(xk): raise RegionEntered(message='# Entered PF...', region='PF') def reorder_limiter(new_start, limiter): start_index, = find_split_index(new_start, limiter) return limiter def limiter_split(start, end, limiter): start_index, sls = find_split_index(start, limiter) end_index, sls = find_split_index(end, limiter) if end_index <= start_index: limiter.p = limiter.p[ start_index:] + limiter.p[:start_index + 1] return limiter try: visual = self.grid.parent.settings['DEBUG']['visual'][ 'SF_analysis'] except KeyError: visual = False try: verbose = self.grid.parent.settings['DEBUG']['verbose'][ 'SF_analysis'] except KeyError: verbose = False WestPlate = self.grid.parent.PlateData['plate_W1'] EastPlate = self.grid.parent.PlateData['plate_E1'] limiter = self.grid.parent.LimiterData xpt1 = self.NSEW_lookup['xpt1']['coor'] magx = np.array([ self.grid.parent.settings['grid_settings']['rmagx'], self.grid.parent.settings['grid_settings']['zmagx'] ]) psi_max = self.grid.parent.settings['grid_settings']['psi_1'] psi_core = self.grid.parent.settings['grid_settings']['psi_core'] psi_pf_1 = self.grid.parent.settings['grid_settings']['psi_pf_1'] # Create mid-line LHS_Point = Point(magx[0] - 1e6, magx[1]) RHS_Point = Point(magx[0] + 1e6, magx[1]) midline = Line([LHS_Point, RHS_Point]) # Generate Vertical Mid-Plane line Lower_Point = Point(magx[0], magx[1] - 1e6) Upper_Point = Point(magx[0], magx[1] + 1e6) topline = Line([Lower_Point, Upper_Point]) # Drawing Separatrix print('# Tracing core region...') xptNW_midLine = self.draw_line(xpt1['NW'], {'line': midline}, option='theta', direction='cw', show_plot=visual, text=verbose) xptNE_midLine = self.draw_line(xpt1['NE'], {'line': midline}, option='theta', direction='ccw', show_plot=visual, text=verbose) midline_topline_west = self.draw_line(xptNW_midLine.p[-1], {'line': topline}, option='theta', direction='cw', show_plot=visual, text=verbose) midline_topline_east = self.draw_line(xptNE_midLine.p[-1], {'line': topline}, option='theta', direction='ccw', show_plot=visual, text=verbose) print('# Tracing PF region...') xptSW_limiter = self.draw_line(xpt1['SW'], { 'line': limiter }, option='theta', direction='ccw', show_plot=visual, text=verbose).reverse_copy() xptSE_limiter = self.draw_line(xpt1['SE'], {'line': limiter}, option='theta', direction='cw', show_plot=visual, text=verbose) core_boundary = Line((xptNW_midLine.p + midline_topline_west.p + midline_topline_east.reverse_copy().p + xptNE_midLine.reverse_copy().p)) core_polygon = Polygon(np.column_stack(core_boundary.points()).T, fill=True, closed=True, color='violet', label='Core') pf_boundary = Line( (xptSW_limiter.p + xptSE_limiter.p + (limiter_split( xptSE_limiter.p[-1], xptSW_limiter.p[0], limiter).split( xptSE_limiter.p[-1])[1]).split( xptSW_limiter.p[0], add_split_point=True)[0].p)) pf_polygon = Polygon(np.column_stack(pf_boundary.points()).T, fill=True, closed=True, color='dodgerblue', label='PF_1') self.RegionPolygon = {'Core': core_polygon, 'PF_1': pf_polygon} if (pf_polygon.get_path().contains_point(xpt)): print("# Snowflake-plus...") self.analyze_saddle(xpt, xpt_ID='xpt2') # Rotate NSEW so that 'W' is set to 'N' and 'E' to 'S' self.rotate_NSEW_lookup(xpt_ID=xpt_ID, turns=6) self._set_function('rho', 'ccw') xpt2 = self.NSEW_lookup['xpt2']['coor'] N_sol = solve_ivp(self.function, (0, self.dt), xpt2['N'], method='LSODA', first_step=self.first_step, max_step=self.max_step, rtol=1e-13, atol=1e-12).y N_guess = (N_sol[0][-1], N_sol[1][-1]) self.RegionLineCut = self.draw_line( xpt2['N'], {'line_group': [xptSW_limiter, xptSE_limiter, limiter]}, option='rho', direction='ccw', show_plot=visual, text=verbose) if self.line_group_intersect == limiter.points(): self.flip_NSEW_lookup(xpt_ID) xpt2 = self.NSEW_lookup['xpt2']['coor'] self.RegionLineCut = self.draw_line(xpt2['N'], { 'line_group': [xptSW_limiter, xptSE_limiter, limiter] }, option='rho', direction='ccw', show_plot=visual, text=verbose) if self.line_group_intersect == xptSE_limiter.points(): self.config = 'SF75' elif self.line_group_intersect == xptSW_limiter.points(): self.config = 'SF105' else: print("# Snowflake-minus...") # Obtain candidate NSEW directions self.analyze_saddle(xpt, xpt_ID='xpt2') # Prepare minimizer for identification of SF- type self._set_function('rho', 'cw') xpt2 = self.NSEW_lookup['xpt2']['coor'] N_sol = solve_ivp(self.function, (0, self.dt), xpt2['N'], method='LSODA', first_step=self.first_step, max_step=self.max_step, rtol=1e-13, atol=1e-12).y S_sol = solve_ivp(self.function, (0, self.dt), xpt2['S'], method='LSODA', first_step=self.first_step, max_step=self.max_step, rtol=1e-13, atol=1e-12).y N_guess = (N_sol[0][-1], N_sol[1][-1]) S_guess = (S_sol[0][-1], S_sol[1][-1]) for guess in [S_guess, N_guess]: try: minimize(self.PsiCostFunc, guess, method='trust-ncg', jac=self.grid.parent.PsiNorm.Gradient, hess=self.grid.parent.PsiNorm.Hessian, options={ 'initial_trust_radius': self.eps, 'max_trust_radius': self.dt }, callback=cb_region_check) except RegionEntered as e: region = e.region if region == 'Core': # True south should land in region of interest if guess is S_guess: self.flip_NSEW_lookup(xpt_ID) xpt2 = self.NSEW_lookup['xpt2']['coor'] # Determine whether SF15 or SF165 based off of intersection with core boundary self.RegionLineCut = self.draw_line( xpt2['N'], { 'line_group': [ xptNE_midLine, xptNW_midLine, midline_topline_west, midline_topline_east, limiter ] }, option='rho', direction='cw', show_plot=visual, text=verbose) if self.line_group_intersect == xptNE_midLine.points( ): self.config = 'SF15' elif self.line_group_intersect == xptNW_midLine.points( ): self.config = 'SF165' elif self.line_group_intersect == midline_topline_west.points( ): self.config = 'UDN' elif self.line_group_intersect == midline_topline_east.points( ): self.config = 'UDN' break elif region == 'PF': # True south should land in region of interest if guess is S_guess: self.flip_NSEW_lookup(xpt_ID) xpt2 = self.NSEW_lookup['xpt2']['coor'] # Determine whether SF15 or SF165 based off of intersection with core boundary self.RegionLineCut = self.draw_line( xpt2['N'], {'line_group': [xptSW_limiter, xptSE_limiter]}, option='rho', direction='cw', show_plot=visual, text=verbose) if self.line_group_intersect == xptSE_limiter.points( ): self.config = 'SF45' elif self.line_group_intersect == xptSW_limiter.points( ): self.config = 'SF135' break else: raise ValueError( f'# Invalid xpt_ID "{xpt_ID}"in function map_xpt(). Valid IDs are "xpt1" and "xpt2".' )
class Domain: def __init__(self, vertices, subdomains=[]): ''' This class describes a 2D polygonal domain defined by n vertices following a closed path: vertices = np.array([[x0,y0], [x1,y1], [x2,y2], ... [xn,yn]]) ''' self.vertices = vertices self.polygon = Polygon(vertices) self.path = self.polygon.get_path() self.subdomains = subdomains def __add__(self, other): if isinstance(other, Domain): return Domain(self.vertices, self.subdomains.append(other)) else: return NotImplemented def __sub__(self, other): return None def __len__(self): return len(self.subdomains) + 1 @property def area(self): ''' Implementation of Shoelace formula using Numpy https://en.wikipedia.org/wiki/Shoelace_formula ''' x = self.vertices[:, 0] y = self.vertices[:, 1] return 0.5 * np.abs( np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1))) def contains(self, points): """ Return True for points inside Domain """ return np.array(self.path.contains_points(points)) def contains_all(self, points): """ Return True if all points are inside Domain """ return self.contains(points).all() def bounds(self): ''' Returns a (minx, miny, maxx, maxy) tuple (float values) that bounds the object. ''' return np.array( [np.min(self.vertices, axis=0), np.max(self.vertices, axis=0)]) def addMaterial(self, Subdomain, Value): if self.contains_all(Subdomain.vertices): Subdomain.Value = Value self.subdomains.append(Subdomain) else: raise TypeError def plot(self, show=True): from matplotlib import pylab from matplotlib.lines import Line2D color = ['b', 'g', 'r', 'c', 'm', 'y', 'k', 'w'] fig, ax = pylab.subplots() self.polygon.set_alpha(0.3) x, y = zip(*self.polygon.xy) line = Line2D(x, y, linestyle='-', marker='o', markerfacecolor='k') ax.add_line(line) ax.set_title('Computational Domain') ax.set_xlim((min(x) - 0.1 * max(x), 1.1 * max(x))) ax.set_ylim((min(y) - 0.1 * max(y), 1.1 * max(y))) line.set_color('k') ax.add_patch(self.polygon) for i, sub in enumerate(self.subdomains): x, y = zip(*sub.polygon.xy) line = Line2D(x, y, linestyle='--', color='k') ax.add_line(line) sub.polygon.set_color(color[1]) sub.polygon.set_alpha(0.3) ax.add_patch(sub.polygon) if show: pylab.show()
def get_polygon(data, no_of_vert=4, xlabel=None, xticks=None, ylabel=None, yticks=None, fs=25): """ Interactive function to pick a polygon out of a figure and receive the vertices of it. :param data: :type: :param no_of_vert: number of vertices, default 4, :type no_of_vert: int """ from bowpy.util.polygon_interactor import PolygonInteractor from matplotlib.patches import Polygon no_of_vert = int(no_of_vert) # Define shape of polygon. try: x, y = xticks.max(), yticks.max() xmin= -x/10. xmax= x/10. ymin= y - 3.*y/2. ymax= y - 3.*y/4. except AttributeError: y,x = data.shape xmin= -x/10. xmax= x/10. ymin= y - 3.*y/2. ymax= y - 3.*y/4. xs = [] for i in range(no_of_vert): if i >= no_of_vert/2: xs.append(xmax) else: xs.append(xmin) ys = np.linspace(ymin, ymax, no_of_vert/2) ys = np.append(ys,ys[::-1]).tolist() poly = Polygon(list(zip(xs, ys)), animated=True, closed=False, fill=False) # Add polygon to figure. fig, ax = plt.subplots() ax.add_patch(poly) p = PolygonInteractor(ax, poly) plt.title("Pick polygon, close figure to save vertices") plt.xlabel(xlabel, fontsize=fs) plt.ylabel(ylabel, fontsize=fs) try: im = ax.imshow(abs(data), aspect='auto', extent=(xticks.min(), xticks.max(), 0, yticks.max()), interpolation='none') except AttributeError: im = ax.imshow(abs(data), aspect='auto', interpolation='none') cbar = fig.colorbar(im) cbar.ax.tick_params(labelsize=fs) cbar.ax.set_ylabel('R', fontsize=fs) mpl.rcParams['xtick.labelsize'] = fs mpl.rcParams['ytick.labelsize'] = fs ax.tick_params(axis='both', which='major', labelsize=fs) plt.show() print("Calculate area inside chosen polygon\n") try: vertices = (poly.get_path().vertices) vert_tmp = [] xticks.sort() yticks.sort() for fkvertex in vertices: vert_tmp.append([np.abs(xticks-fkvertex[0]).argmin(), np.abs(yticks[::-1]-fkvertex[1]).argmin()]) vertices = np.array(vert_tmp) except AttributeError: vertices = (poly.get_path().vertices).astype('int') indicies = convert_polygon_to_flat_index(data, vertices) return indicies
class dispASDF(pyasdf.ASDFDataSet): """An class for analyzing dispersion curves based on ASDF database """ def set_poly(self,lst,minlon,minlat,maxlon,maxlat): """Define the polygon for study region and max(min) longitude/latitude for map view Parameters: lst -- list of (lon, lat) defining the polygon for study region minlon, minlat, maxlon, maxlat -- minimum and miximum coordinates for the input age model """ self.poly = Polygon(lst) self.perimeter = self.poly.get_path() self.minlon = minlon; self.minlat = minlat; self.maxlon = maxlon; self.maxlat = maxlat return def point_in(self): """test if a point of given longitude and latitude is in the polygon of study region """ return def path_in(self): """test if the line connecting 2 given points is in the polygon """ return def read_age_mdl(self): """read in crustal age model for the oceans """ # dset = Dataset('/projects/howa1663/Code/ToolKit/Models/Age_Ocean_Crust/age.3.2.nc','r') dset = Dataset('./age.3.2.nc','r') # 2-minute resolution longitude = dset.variables['x'][:] longitude[longitude<0] += 360. latitude = dset.variables['y'][:] z = dset.variables['z'][:] # masked array mask = dset.variables['z'][:].mask data = dset.variables['z'][:].data / 100. data[mask] = 9999. data = data[(latitude >= self.minlat)*(latitude <= self.maxlat),:] data = data[:,(longitude >= self.minlon)*(longitude <= self.maxlon)] longitude = longitude[(longitude >= self.minlon)*(longitude <= self.maxlon)] latitude = latitude[(latitude >= self.minlat)*(latitude <= self.maxlat)] self.age_data = data; self.age_lon = longitude; self.age_lat = latitude return def read_topo_mdl(self): """Read in topography model, return the function for calculating topography at any given point. """ # etopo1 = Dataset('/projects/howa1663/Code/ToolKit/Models/ETOPO1/ETOPO1_Ice_g_gmt4.grd', 'r') etopo1 = Dataset('/work2/wang/Code/ToolKit/ETOPO1_Ice_g_gmt4.grd','r') lons = etopo1.variables["x"][:] west = lons<0 # mask array with negetive longitudes west = 360.*west*np.ones(len(lons)) lons = lons+west lats = etopo1.variables["y"][:] z = etopo1.variables["z"][:] etopoz = z[(lats>=(self.minlat))*(lats<=(self.maxlat)), :] etopoz = etopoz[:, (lons>=self.minlon)*(lons<=self.maxlon)] lats = lats[(lats>=self.minlat)*(lats<=self.maxlat)] lons = lons[(lons>=self.minlon)*(lons<=self.maxlon)] etopox, etopoy = np.meshgrid(lats, lons, indexing='ij') etopox = etopox.flatten() etopoy = etopoy.flatten() points = np.vstack((etopox,etopoy)).T etopoz = etopoz.flatten() f = LinearNDInterpolator(points,etopoz) return f def read_stations(self,stafile,source='CIEI',chans=['BHZ', 'BHE', 'BHN']): """Read in stations from station list file """ with open(stafile, 'r') as f: Sta = [] site = obspy.core.inventory.util.Site(name='01') creation_date = obspy.core.utcdatetime.UTCDateTime(0) inv = obspy.core.inventory.inventory.Inventory(networks=[], source=source) total_number_of_channels = len(chans) for lines in f.readlines(): lines = lines.split() stacode = lines[0] lon = float(lines[1]) lat = float(lines[2]) if not self.perimeter.contains_point((lon,lat)): continue; netcode = lines[3] netsta = netcode+'.'+stacode if Sta.__contains__(netsta): index = Sta.index(netsta) if abs(self[index].lon-lon) >0.01 and abs(self[index].lat-lat) >0.01: raise ValueError('Incompatible Station Location:' + netsta+' in Station List!') else: print('Warning: Repeated Station:' +netsta+' in Station List!') continue Sta.append(netsta) channels = [] if lon > 180.: lon -= 360. for chan in chans: channel = obspy.core.inventory.channel.Channel(code=chan, location_code='01', latitude=lat, longitude=lon, elevation=0.0, depth=0.0) channels.append(channel) station = obspy.core.inventory.station.Station(code=stacode, latitude=lat, longitude=lon, elevation=0.0, site=site, channels=channels, total_number_of_channels = total_number_of_channels, creation_date = creation_date) network = obspy.core.inventory.network.Network(code=netcode, stations=[station]) networks = [network] inv += obspy.core.inventory.inventory.Inventory(networks=networks, source=source) print('Writing obspy inventory to ASDF dataset') self.add_stationxml(inv) print('End writing obspy inventory to ASDF dataset') return def get_ages(self, lons,lats): """ calculate ages for a give list of points on the ocean floor Parameters: lons -- longitude vector lats -- latitude vector """ xx, yy = np.meshgrid(self.age_lon, self.age_lat) # xx for longitude, yy for latitude xx = xx.reshape(xx.size) #nearest yy = yy.reshape(yy.size) x = np.column_stack((xx,yy)) y = self.age_data.reshape(self.age_data.size) f = NearestNDInterpolator(x,y,rescale=False) ages = f(np.column_stack((lons,lats))) # mask = self.age_data > 180 #f = interp2d(xx[~mask],yy[~mask],self.age_data[~mask],kind='cubic') # cubic spline interpolation # ages = f(lons,lats) return ages def read_paths(self,disp_dir='/work3/wang/JdF/FTAN',res=3.5): """read in dispersion curve measurements from files, calculate ages along great circle path, read waveforms Paramenters: disp_dir -- directory for the dispersion measurement results res -- resolution for the great circle sampling points """ yy, xx = np.meshgrid(self.age_lon, self.age_lat) xx = xx.reshape(xx.size) yy = yy.reshape(yy.size) x = np.vstack((xx,yy)).T y = self.age_data.reshape(self.age_data.size) f = NearestNDInterpolator(x,y,rescale=False) # nearest-neighbour interpolation f_topo = self.read_topo_mdl() staLst = self.waveforms.list() print('Reading dispersion curves & averaging crustal age along paths') for staid1 in staLst: lat1 = self.waveforms[staid1].coordinates['latitude'] lon1 = self.waveforms[staid1].coordinates['longitude'] netcode1, stacode1 = staid1.split('.') if lon1 < 0.: lon1 += 360. for staid2 in staLst: if staid1 >= staid2: continue; lat2 = self.waveforms[staid2].coordinates['latitude'] lon2 = self.waveforms[staid2].coordinates['longitude'] if lon2 < 0.: lon2 += 360. gc_path, dist = get_gc_path(lon1,lat1,lon2,lat2,res) # the great circle path travels beyond the study region if not self.perimeter.contains_path(Path(gc_path)): print("Station "+staid1+" not in study region! Will be discarded!") del self.waveforms[staid1] continue; # the great circle path travels beyond the study region netcode2, stacode2 = staid2.split('.') gc_points = np.array(gc_path) ages = f(gc_points[:,::-1]) depths = f_topo(gc_points[:,::-1]) if ages.max() > 300: continue; # the paths went out the model bound #age_avg = 1/((1./ages).mean()) age_avg = ages.mean() depth_avg = depths.mean() disp_file = disp_dir+'/'+stacode1+'/'+'COR_'+stacode1+'_'+stacode2+'.SAC_2_DISP.1' snr_file = disp_dir+'/'+stacode1+'/'+'COR_'+stacode1+'_'+stacode2+'.SAC_2_amp_snr' if not os.path.isfile(disp_file): continue; arr = np.loadtxt(disp_file) arr_snr = np.loadtxt(snr_file) snrp_vec = arr_snr[:,2] # snr for positive lag signal snrn_vec = arr_snr[:,4] # snr for negative lag signal transp = np.vstack((snrp_vec,snrn_vec)).max(axis=0) > 3 # mean snr > 5 # transp = np.vstack((snrp_vec,snrn_vec)).min(axis=0) > 5 # both positive & negative lag have snr > 5 if (transp*1).max() == 0: continue per_vec = arr[transp,2] grv_vec = arr[transp,3] phv_vec = arr[transp,4] d_g = grv_vec[1:]-grv_vec[:-1] try: d_g[per_vec[1:]<6.] = 1. # don't consider period less than 5. sec index = np.where(d_g < -0.2)[0][0] # find out where the group velocity starts to drop larger than 0.2 per_vec = per_vec[:index+1] # only keep results before the group velocity starts to drop grv_vec = grv_vec[:index+1] phv_vec = phv_vec[:index+1] snrp_vec = snrp_vec[:index+1] snrn_vec = snrn_vec[:index+1] except: pass mask = per_vec*phv_vec > dist # interstation distance larger than one wavelength try: ind1 = np.where(mask)[0][0] # 1st True if ind1 == 0: continue except: ind1 = per_vec.size disp_arr = np.vstack((per_vec[:ind1],grv_vec[:ind1],phv_vec[:ind1],snrp_vec[:ind1],snrn_vec[:ind1])) # xcorr_file = disp_dir+'/'+stacode1+'/'+'COR_'+stacode1+'_'+stacode2+'.SAC' # if not os.path.isfile(xcorr_file): continue; # tr = obspy.core.read(xcorr_file)[0] # xcorr_header = {'stacode1': '', 'stacode2': '', 'npts': 12345, 'b': 12345, 'e': 12345, \ # 'delta': 12345, 'dist': 12345, 'stackday': 0} # xcorr_header['b'] = tr.stats.sac.b # xcorr_header['e'] = tr.stats.sac.e # xcorr_header['stacode1'] = stacode1 # xcorr_header['stacode2'] = stacode2 # xcorr_header['npts'] = tr.stats.npts # xcorr_header['delta'] = tr.stats.delta # xcorr_header['stackday'] = tr.stats.sac.user0 # try: # xcorr_header['dist'] = tr.stats.sac.dist # except AttributeError: # xcorr_header['dist'] = dist staid_aux = netcode1+'/'+stacode1+'/'+netcode2+'/'+stacode2 parameters1 = {'age_avg': age_avg, 'depth_avg': depth_avg,'dist':dist, 'L_gc':len(gc_path)} ages_path = np.concatenate((gc_points,ages.reshape(-1,1)),axis=1) self.add_auxiliary_data(data=ages_path, data_type='AgeGc', path=staid_aux, parameters=parameters1) parameters2 = {'T': 0, 'grV': 1, 'phV': 2, 'snr_p': 3, 'snr_n': 4,'dist': dist, 'Np': ind1} self.add_auxiliary_data(data=disp_arr, data_type='DISPArray', path=staid_aux, parameters=parameters2) #self.add_auxiliary_data(data=tr.data, data_type='NoiseXcorr', path=staid_aux, parameters=xcorr_header) print('End of reading dispersion curves') return def intp_disp(self,pers,verbose=False): """interpolate the dispersion curves to a given period band, QC was applied during this process Parameter: pers -- period array """ if pers.size == 0: pers = np.append( np.arange(6.)*2.+6., np.arange(4.)*3.+18.) self.pers = pers staLst = self.waveforms.list() for staid1 in staLst: netcode1, stacode1 = staid1.split('.') for staid2 in staLst: netcode2, stacode2 = staid2.split('.') if staid1 >= staid2: continue try: subdset = self.auxiliary_data['DISPArray'][netcode1][stacode1][netcode2][stacode2] except: continue data = subdset.data.value index = subdset.parameters dist = index['dist'] if verbose: print('Interpolating dispersion curve for '+ netcode1+'.'+stacode1+'_'+netcode2+'.'+stacode2) outindex = { 'T': 0, 'grV': 1, 'phV': 2, 'snr_p': 3, 'snr_n': 4, 'dist': dist, 'Np': pers.size } Np = int(index['Np']) if Np < 5: warnings.warn('Not enough datapoints for: '+ netcode1+'.'+stacode1+'_'+netcode2+'.'+stacode2, UserWarning, stacklevel=1) continue obsT = data[index['T']][:Np] grV = np.interp(pers, obsT, data[index['grV']][:Np] ) phV = np.interp(pers, obsT, data[index['phV']][:Np] ) inbound = (pers > obsT[0])*(pers < obsT[-1]) if grV[inbound].size == pers[inbound].size and phV[inbound].size == pers[inbound].size: interpdata = np.append(pers[inbound], grV[inbound]) interpdata = np.append(interpdata, phV[inbound]) else: continue snr_p = np.interp(pers, obsT, data[index['snr_p']][:Np] ) snr_n = np.interp(pers, obsT, data[index['snr_n']][:Np] ) interpdata = np.append(interpdata, snr_p[inbound]) interpdata = np.append(interpdata, snr_n[inbound]) interpdata = interpdata.reshape(5, pers[inbound].size) staid_aux = netcode1+'/'+stacode1+'/'+netcode2+'/'+stacode2 self.add_auxiliary_data(data=interpdata, data_type='DISPinterp', path=staid_aux, parameters=outindex) return def get_basemap(self,model='age'): """get basemap from given model, use ocean crustal age model or etopo model Parameters: model -- use which type of data for basemap, 'age' or 'etopo' """ if model == 'age': distEW, az, baz=obspy.geodetics.gps2dist_azimuth(self.minlat, self.minlon, self.minlat, self.maxlon) # distance is in m distNS, az, baz=obspy.geodetics.gps2dist_azimuth(self.minlat, self.minlon, self.maxlat+2., self.minlon) # distance is in m m = Basemap(width=distEW, height=distNS, rsphere=(6378137.00,6356752.3142), resolution='i', projection='lcc',\ lat_1=self.minlat-1, lat_2=self.maxlat+1, lon_0=(self.minlon+self.maxlon)/2, lat_0=(self.minlat+self.maxlat)/2) m.drawparallels(np.arange(-80.0,80.0,5.0), linewidth=1, dashes=[2,2], labels=[1,0,0,0], fontsize=15) m.drawmeridians(np.arange(-170.0,170.0,5.0), linewidth=1, dashes=[2,2], labels=[0,0,1,0], fontsize=15) m.drawcoastlines(linewidth=1.0) m.drawcountries(linewidth=1.0) m.drawstates(linewidth=1.0) m.readshapefile('./Plates/PB2002_boundaries', name='PB2002_boundaries', drawbounds=True, \ linewidth=1, color='orange') # draw plate boundary on basemap x, y = m(*np.meshgrid(self.age_lon,self.age_lat)) data = np.ma.masked_array(self.age_data, self.age_data > 9000) img = m.pcolormesh(x, y, data, shading='gouraud', cmap='jet_r', vmin=0, vmax=11,alpha=0.5) # cmap possible choices: "jet","Spectral" m.drawmapboundary(fill_color="white") cbar = m.colorbar(img,location='bottom',pad="3%") cbar.set_alpha(1) cbar.draw_all() cbar.set_label('Age (Ma)') # plt.show() elif model == 'etopo': code else: raise ValueError('Only age or etopo model can be used for constructing basemap') return m def plot_stations(self,ppoly=False): """Plot stations on basemap basemap -- 'age' or 'etopo' ppoly -- flag for if showing the boundary polygon """ staLst = self.waveforms.list() m = self.get_basemap() for staid in staLst: lat = self.waveforms[staid].coordinates['latitude'] lon = self.waveforms[staid].coordinates['longitude'] lat_x, lat_y = m(lon,lat) m.plot(lat_x,lat_y,'^',color='olive') if ppoly: point_arr = self.poly.get_xy() xx, yy = m(point_arr[:,0],point_arr[:,1]) cord_arr = np.vstack((xx,yy)).T poly = Polygon(cord_arr,facecolor='none', edgecolor='k') plt.gca().add_patch(poly) plt.title('Station map',fontsize=16) plt.show() return def get_vel_age(self, c0,c1,c2,vmin,vmax): """ Calculate 2-D velocity map based on the three coefficients """ age_data = self.age_data age_lon = self.age_lon age_lat = self.age_lat llon, llat = np.meshgrid(age_lon, age_lat,indexing='ij') ccords = np.dstack((llon,llat)).reshape(llon.size,2) ind_arr = self.perimeter.contains_points(ccords).reshape(llon.shape[0],llon.shape[1]) # select the data points inside the pre-set polygon vel_arr = (c0+ c1*np.sqrt(age_data) + c2*age_data)*np.transpose(ind_arr.astype(float)) distEW, az, baz=obspy.geodetics.gps2dist_azimuth(self.minlat, self.minlon, self.minlat, self.maxlon) # distance is in m distNS, az, baz=obspy.geodetics.gps2dist_azimuth(self.minlat, self.minlon, self.maxlat+2., self.minlon) # distance is in m m = Basemap(width=distEW, height=distNS, rsphere=(6378137.00,6356752.3142), resolution='i', projection='lcc',\ lat_1=self.minlat-1, lat_2=self.maxlat+1, lon_0=(self.minlon+self.maxlon)/2, lat_0=(self.minlat+self.maxlat)/2) m.drawparallels(np.arange(-80.0,80.0,2.0), linewidth=1, dashes=[2,2], labels=[1,0,0,0], fontsize=15) m.drawmeridians(np.arange(-170.0,170.0,2.0), linewidth=1, dashes=[2,2], labels=[0,0,1,0], fontsize=15) m.drawcoastlines(linewidth=1.0) m.drawcountries(linewidth=1.0) m.drawstates(linewidth=1.0) m.readshapefile('./Plates/PB2002_boundaries', name='PB2002_boundaries', drawbounds=True, \ linewidth=1, color='orange') # draw plate boundary on basemap x, y = m(*np.meshgrid(age_lon,age_lat)) data = np.ma.masked_array(vel_arr, vel_arr < 0.1) img = m.pcolormesh(x, y, data, shading='gouraud', cmap='jet_r', vmin=vmin, vmax=vmax) # cmap possible choices: "jet","Spectral" m.drawmapboundary(fill_color="white") cbar = m.colorbar(img,location='bottom',pad="3%") #cbar.solids.set_eddgecolor("face") cbar.set_label('Vel (km/s)') # plt.title('Age-dependent velocity map',fontsize=16) plt.show() return def count_path(self): """get number of paths at all periods for all stations """ return def fit_Harmon(self,period,vel_type='phase'): """find the 3 coefficients which best fit the Harmon velocity-age relationship. [Ye et al, GGG, 2013, eq(1)] Parameters: period -- the period of the group or phase velocity to analyse for vel_type -- type of velocity to invert for. 'group' or 'phase' """ try: subset = self.auxiliary_data.FitResult[str_per][vel_type] warnings.warn("Funciton fit_Harmon has been run for this period and velocity type!") return except: pass staLst = self.waveforms.list() dist_lst = [] int_dis_lst = [] ages_lst = [] age_avg_lst = [] V_lst = [] netcode1_lst = []# used to save which stations were left after final fitting model stacode1_lst = [] netcode2_lst = [] stacode2_lst = [] # used to save which stations were left after final fitting model for staid1 in staLst: netcode1, stacode1 = staid1.split('.') for staid2 in staLst: if staid1 >= staid2: continue; netcode2, stacode2 = staid2.split('.') try: T_V = self.auxiliary_data['DISPinterp'][netcode1][stacode1][netcode2][stacode2].data.value # T, grV, phV except: continue ind_T = np.where(T_V[0,:]==period)[0] if ind_T.size != 1: continue; if vel_type == 'group': ind_V = 1 elif vel_type == 'phase': ind_V = 2 else: raise AttributeError('velocity type can only be group or phase') try: subset_age = self.auxiliary_data['AgeGc'][netcode1][stacode1][netcode2][stacode2] except: print((stacode1+'_'+stacode2+'pair has interpolated dispersion curve but dont have age along path')) continue if T_V[3,ind_T] < 5. or T_V[4,ind_T] < 5.: # snr_p or snr_n smaller than 5. continue d_T = min(period/3, 2.) # time_delay = get_ph_misfit(period,1./(period+d_T),1./(period-d_T),stacode1,stacode2,T_V[1,ind_T]) # if np.abs(time_delay) > 1. or np.abs(time_delay / period) > 0.2: # continue if T_V[2,ind_T]<T_V[1,ind_T]: #phase velocity smaller than group velocity continue V_lst.append(T_V[ind_V,ind_T]) dist = subset_age.parameters['dist'] dist_lst.append(dist) inter_dist = dist / (subset_age.parameters['L_gc']-1) int_dis_lst.append(inter_dist) ages = subset_age.data.value[:,2] age_avg_lst.append(subset_age.parameters['age_avg']) ages_lst.append(ages) netcode1_lst.append(netcode1) stacode1_lst.append(stacode1) netcode2_lst.append(netcode2) stacode2_lst.append(stacode2) if not len(ages_lst)==len(V_lst) & len(V_lst)==len(int_dis_lst): raise AttributeError('The number of inter-station paths are incompatible for inverting c0, c1 & c2') dist = np.array(dist_lst) d_dist = np.array(int_dis_lst) age_avgs = np.array(age_avg_lst) V = np.array(V_lst).reshape(len(V_lst)) fits = np.polyfit(np.sqrt(age_avgs),V,2) p = np.poly1d(fits) predict_V = p(np.sqrt(age_avgs)) diffs = np.abs(predict_V-V) diffs_mean = np.mean(diffs) diffs_std = np.std(diffs) ind = np.abs(diffs-diffs_mean) < 3*diffs_std # discard those datapoints which are far away from a crude model. ages_lst_final = list(compress(ages_lst,ind)) params = (dist[ind], d_dist[ind], V[ind], ages_lst_final) # params = (dist, d_dist, V, ages_lst) cranges = (slice(0.5,4,0.1), slice(-0.5,0.5,0.05), slice(-1.,1,0.05)) resbrute = optimize.brute(sq_misfit,cranges,args=params,full_output=True,finish=None) print("Age-depedent coefficients found for period {} sec.".format(period)) c0, c1, c2 = resbrute[0] data_out = np.vstack((age_avgs[ind],V[ind])).T # [-1,2] array, storing average age and velocity netcode1_lst_final = list(compress(netcode1_lst,ind)) netcode2_lst_final = list(compress(netcode2_lst,ind)) stacode1_lst_final = list(compress(stacode1_lst,ind)) stacode2_lst_final = list(compress(stacode2_lst,ind)) netcode1_arr = np.chararray(len(netcode1_lst_final),itemsize=2); netcode2_arr = np.chararray(len(netcode2_lst_final),itemsize=2) stacode1_arr = np.chararray(len(stacode1_lst_final),itemsize=5); stacode2_arr = np.chararray(len(stacode2_lst_final),itemsize=5) netcode1_arr[:] = netcode1_lst_final[:] stacode1_arr[:] = stacode1_lst_final[:] netcode2_arr[:] = netcode2_lst_final[:] stacode2_arr[:] = stacode2_lst_final[:] # data_out = np.vstack((age_avgs,V)).T # [-1,2] array, storing average age and velocity # netcode1_arr = np.chararray(len(netcode1_lst),itemsize=2); netcode2_arr = np.chararray(len(netcode2_lst),itemsize=2) # stacode1_arr = np.chararray(len(stacode1_lst),itemsize=5); stacode2_arr = np.chararray(len(stacode2_lst),itemsize=5) # netcode1_arr[:] = netcode1_lst[:] # stacode1_arr[:] = stacode1_lst[:] # netcode2_arr[:] = netcode2_lst[:] # stacode2_arr[:] = stacode2_lst[:] stas_arr_final = np.vstack((netcode1_arr,stacode1_arr,netcode2_arr,stacode2_arr)).T # plt.plot(age_avgs[ind],V[ind],'r.') # t0 = np.linspace(0,10,100) # plt.plot(t0,p(np.sqrt(t0)), 'b-') # plt.plot(t0,c0+c1*np.sqrt(t0)+c2*t0, 'g-') # plt.show() para_aux = str(period).zfill(2)+'/'+vel_type parameters = {'c0':c0, 'c1':c1, 'c2':c2} self.add_auxiliary_data(data=data_out, data_type='FitResult', path=para_aux, parameters=parameters) out_index = {'netcode1':0, 'stacode1':1, 'netcode2':2, 'stacode2':3} self.add_auxiliary_data(data=stas_arr_final, data_type='FinalStas', path=para_aux, parameters=out_index) return resbrute def plot_vel_age(self,period,vel_type='phase'): """Plot velocity vs. oceanic crust age, both from model and measurement """ str_per = str(period).zfill(2) age_vel = self.auxiliary_data.FitResult[str_per][vel_type].data.value codes = self.auxiliary_data.FinalStas[str_per][vel_type].data.value c0 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c0'] c1 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c1'] c2 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c2'] plt.plot(age_vel[:,0], age_vel[:,1], 'r.') t0 = np.linspace(0,10,50) plt.plot(t0,c0+c1*np.sqrt(t0)+c2*t0, 'b-') # for i in range(len(codes)): # plt.text(age_vel[i,0],age_vel[i,1],codes[i,1]+'_'+codes[i,3]) plt.title(str(period)+' sec '+vel_type+' velocity vs. oceanic age') plt.xlim(xmin=0.) plt.xlabel('age (ma)') plt.ylabel('km/s') plt.show() return def plot_age_topo(self,period,vel_type='phase'): """Plot age vs. ocean depth averaged along path. """ str_per = str(period).zfill(2) codes = self.auxiliary_data.FinalStas[str_per][vel_type].data.value ages = np.array([]) depths = np.array([]) for code in codes: age_avg = self.auxiliary_data['AgeGc'][code[0]][code[1]][code[2]][code[3]].parameters['age_avg'] depth_avg = self.auxiliary_data['AgeGc'][code[0]][code[1]][code[2]][code[3]].parameters['depth_avg'] ages = np.append(ages,age_avg) depths = np.append(depths,depth_avg) plt.plot(ages, depths, 'r.') plt.title('Oceanic depth vs. age for '+str(period)+' sec paths') plt.ylabel('Depth (m)') plt.xlabel('Age (ma)') plt.show() return def plot_vel_topo(self, period, vel_type='phase'): """ Plot velocity vs. ocean depth averaged along paths. """ str_per = str(period).zfill(2) codes = self.auxiliary_data.FinalStas[str_per][vel_type].data.value depths = np.array([]) for code in codes: depth_avg = self.auxiliary_data['AgeGc'][code[0]][code[1]][code[2]][code[3]].parameters['depth_avg'] depths = np.append(depths,depth_avg) age_vel = self.auxiliary_data.FitResult[str_per][vel_type].data.value plt.plot(depths, age_vel[:,1], 'r.') plt.title(str(period)+' sec '+vel_type+' velocity vs. oceanic depth') plt.xlabel('Depth (m)') plt.ylabel('km/s') plt.show() return def plot_all_vel(self,pers=np.array([6,8,10,14,18,24,27]),vel_type='phase'): """ Plot the interpolated dispersion result in the same plot for a certain type of velocity """ colors = ['red','green','wheat','blue','orange','black','cyan'] i = 0 for period in pers: str_per = str(period).zfill(2) age_vel = self.auxiliary_data.FitResult[str_per][vel_type].data.value c0 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c0'] c1 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c1'] c2 = self.auxiliary_data.FitResult[str_per][vel_type].parameters['c2'] plt.plot(age_vel[:,0], age_vel[:,1], '.',color=colors[i],label=str(period)+' sec') t0 = np.linspace(0,np.max(age_vel[:,0]),100) plt.plot(t0,c0+c1*np.sqrt(t0)+c2*t0,color=colors[i]) i += 1 plt.legend(loc='best',fontsize=16) plt.title(vel_type+' velocity vs. oceanic age',fontsize=16) plt.xlabel('age (ma)',fontsize=16) plt.ylabel('km/s',fontsize=16) plt.xlim(xmin=0.) plt.xticks(fontsize=18) plt.yticks(fontsize=18) plt.show() return def plot_paths(self,period,vel_type='phase'): """ Plot all the paths for dispersion curves used for a given period. The paths are great circle paths """ str_per = str(period).zfill(2) codes = self.auxiliary_data.FinalStas[str_per][vel_type].data.value m = self.get_basemap() for code in codes: lat1 = self.waveforms[code[0]+'.'+code[1]].coordinates['latitude'] lon1 = self.waveforms[code[0]+'.'+code[1]].coordinates['longitude'] lat2 = self.waveforms[code[2]+'.'+code[3]].coordinates['latitude'] lon2 = self.waveforms[code[2]+'.'+code[3]].coordinates['longitude'] gc_path, _ = get_gc_path(lon1,lat1,lon2,lat2,3) gc_points = np.array(gc_path) path_x, path_y = m(gc_points[:,0], gc_points[:,1]) #m.plot(path_x,path_y,color='black',linewidth=0.5) m.plot(path_x[0],path_y[0],'^',color='olive') m.plot(path_x[-1],path_y[-1],'^',color='olive') plt.title(str(period)+' sec') plt.show() return def plot_age_model(lons, lats, resolution='i', cpt='/projects/howa1663/Code/ToolKit/Models/Age_Ocean_Crust/age1.cpt'): """ Not finished """ #mycm=pycpt.load.gmtColormap(cpt) try: etopo1 = Dataset('/work2/wang/Code/ToolKit/ETOPO1_Ice_g_gmt4.grd','r') # read in the etopo1 file which was used as the basemap llons = etopo1.variables["x"][:] west = llons<0 # mask array with negetive longitudes west = 360.*west*np.ones(len(llons)) llons = llons+west llats = etopo1.variables["y"][:] zz = etopo1.variables["z"][:] etopoz = zz[(llats>(lats[0]-2))*(llats<(lats[1]+2)), :] etopoz = etopoz[:, (llons>(lons[0]-2))*(llons<(lons[1]+2))] llats = llats[(llats>(lats[0]-2))*(llats<(lats[1]+2))] llons = llons[(llons>(lons[0]-2))*(llons<(lons[1]+2))] etopoZ = m.transform_scalar(etopoz, llons-360*(llons>180)*np.ones(len(llons)), llats, etopoz.shape[0], etopoz.shape[1]) # tranform the altitude grid into the projected coordinate ls = LightSource(azdeg=315, altdeg=45) m.imshow(ls.hillshade(etopoZ, vert_exag=0.05),cmap='gray') except IOError: print("Couldn't read etopo data or color map file! Check file directory!") return def write_2_sta_in(self, outdir="/work3/wang/JdF/Input_4_Ray", channel='ZZ', pers = np.array([]), outpfx="2_sta_in_",cut=0.8): """ Write the FTAN measurements as input files for 2-station tomography Parameters: outdir -- directory for the generated input files outpfx -- prefix for the name of generated files cut -- cut those lines in input file whose velocity differ more than this threshold from the mean or median """ if not os.path.isdir(outdir): os.makedirs(outdir) if pers.size == 0: pers=np.append( np.arange(18.)*2.+6.) staLst = self.waveforms.list() ph_f_lst = [] gr_f_lst = [] for prd in pers: ph_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_ph.lst_tmp" gr_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_gr.lst_tmp" ph_f = open(ph_name, 'w') gr_f = open(gr_name, 'w') ph_f_lst.append(ph_f) gr_f_lst.append(gr_f) # to be written on the input file # N(id), lat, lon1, lat2, lon2, pvel(gvel), 1.(weight), (strings doesn't really affect to result)staid1, staid2, 1, 1 i = -1 for staid1 in staLst: netcode1, stacode1 = staid1.split('.') lat1 = self.waveforms[staid1].coordinates['latitude'] lon1 = self.waveforms[staid1].coordinates['longitude'] if lon1<0: lon1 += 360. for staid2 in staLst: if staid1 >= staid2: continue; netcode2, stacode2 = staid2.split('.') i = i + 1 try: itp_arr = self.auxiliary_data['DISPinterp'][netcode1][stacode1][netcode2][stacode2].data.value except: continue if itp_arr.shape[1] == 0: continue lat2 = self.waveforms[staid2].coordinates['latitude'] lon2 = self.waveforms[staid2].coordinates['longitude'] dist, az, baz = obspy.geodetics.gps2dist_azimuth(lat1, lon1, lat2, lon2) # distance is in m dist = dist/1000. if lon2 < 0: lon2 += 360. for iper in range(pers.size): per = pers[iper] ind_per = np.where(itp_arr[0,:] == per)[0] if not ind_per.size == 1: continue pvel = itp_arr[2,ind_per][0] gvel = itp_arr[1,ind_per][0] # if dist < 2.*per*pvel: continue # inbound=data[index['inbound']][ind_per] # quality control if pvel < 0 or gvel < 0 or pvel > 5 or gvel > 5: continue if max(np.isnan([pvel, gvel])) != False: continue # skip if parameters in dispersion curve is nan # if inbound != 1.: continue fph = ph_f_lst[iper] fgr = gr_f_lst[iper] fph.writelines("%d %g %g %g %g %g 1. %s %s 1 1 \n" %(i, lat1, lon1, lat2, lon2, pvel, staid1, staid2)) fgr.writelines("%d %g %g %g %g %g 1. %s %s 1 1 \n" %(i, lat1, lon1, lat2, lon2, gvel, staid1, staid2)) for iper in range(pers.size): fph = ph_f_lst[iper] fgr = gr_f_lst[iper] fph.close() fgr.close() from itertools import compress for prd in pers: in_ph_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_ph.lst_tmp" in_gr_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_gr.lst_tmp" fn_ph_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_ph.lst" fn_gr_name = outdir+"/"+outpfx+"%g"%(prd)+"_"+channel+"_gr.lst" with open(in_ph_name) as ph_f_tmp: ph_V_arr = np.array(ph_f_tmp.read().split()[5::11],dtype='f') ph_ind = np.logical_and(abs(ph_V_arr-ph_V_arr.mean()) < cut, abs(ph_V_arr-np.median(ph_V_arr)) < cut) ph_lines_in = [line.rstrip('\n') for line in open(in_ph_name)] ph_lines_fn = list(compress(ph_lines_in, ph_ind)) ph_fn = open (fn_ph_name,'w') ph_fn.writelines('%s\n'%item for item in ph_lines_fn) ph_fn.close() with open(in_gr_name) as gr_f_tmp: gr_V_arr = np.array(gr_f_tmp.read().split()[5::11],dtype='f') gr_ind = np.logical_and(abs(gr_V_arr-gr_V_arr.mean()) < cut, abs(gr_V_arr-np.median(gr_V_arr)) < cut) gr_lines_in = [line.rstrip('\n') for line in open(in_gr_name)] gr_lines_fn = list(compress(gr_lines_in, gr_ind)) gr_fn = open (fn_gr_name,'w') gr_fn.writelines('%s\n'%item for item in gr_lines_fn) gr_fn.close() print('End of Generating Misha Tomography Input File!') return
lat_3 = xlat[row+1][col] data = Conc_var[row][col] if data > 0: if row % 10 == 0 and col % 10 == 0: print row,col cell_vertices = np.array([m(lon_0,lat_0),m(lon_1,lat_1),m(lon_2,lat_2),m(lon_3,lat_3)]) cell_poly = _geoslib.Polygon(cell_vertices) #tundra fraction = 0 for info, shape in zip(m.surface_cover_info, m.surface_cover): #print info if info['DN'] == 1 and info['RINGNUM'] == 1: patch = Polygon(np.array(shape), True) poly = _geoslib.Polygon(patch.get_path().vertices) fraction, intersection_poly = INP.calcPolygonOverlapArea(cell_poly, poly,fraction) if fraction >= 1: break land_grid[row][col] = max(0,fraction) with open(os.path.join(output_path, 'gridded_fraction_tundra-'+sim_start_datetime.strftime("%Y%m%d-%H%M") + '_AmundsenINP.pckl'), 'wb') as file: pickle.dump(land_grid, file) print datetime.now() - start
def get_polygon(data, no_of_vert=4, xlabel=None, xticks=None, ylabel=None, yticks=None, fs=25): """ Interactive function to pick a polygon out of a figure and receive the vertices of it. :param data: :type: :param no_of_vert: number of vertices, default 4, :type no_of_vert: int """ from sipy.util.polygon_interactor import PolygonInteractor from matplotlib.patches import Polygon no_of_vert = int(no_of_vert) # Define shape of polygon. try: x, y = xticks.max(), yticks.max() xmin= -x/10. xmax= x/10. ymin= y - 3.*y/2. ymax= y - 3.*y/4. except AttributeError: y,x = data.shape xmin= -x/10. xmax= x/10. ymin= y - 3.*y/2. ymax= y - 3.*y/4. xs = [] for i in range(no_of_vert): if i >= no_of_vert/2: xs.append(xmax) else: xs.append(xmin) ys = np.linspace(ymin, ymax, no_of_vert/2) ys = np.append(ys,ys[::-1]).tolist() poly = Polygon(list(zip(xs, ys)), animated=True, closed=False, fill=False) # Add polygon to figure. fig, ax = plt.subplots() ax.add_patch(poly) p = PolygonInteractor(ax, poly) plt.title("Pick polygon, close figure to save vertices") plt.xlabel(xlabel, fontsize=fs) plt.ylabel(ylabel, fontsize=fs) try: im = ax.imshow(abs(data), aspect='auto', extent=(xticks.min(), xticks.max(), 0, yticks.max()), interpolation='none') except AttributeError: im = ax.imshow(abs(data), aspect='auto', interpolation='none') cbar = fig.colorbar(im) cbar.ax.tick_params(labelsize=fs) cbar.ax.set_ylabel('R', fontsize=fs) mpl.rcParams['xtick.labelsize'] = fs mpl.rcParams['ytick.labelsize'] = fs ax.tick_params(axis='both', which='major', labelsize=fs) plt.show() print("Calculate area inside chosen polygon\n") try: vertices = (poly.get_path().vertices) vert_tmp = [] xticks.sort() yticks.sort() for fkvertex in vertices: vert_tmp.append([np.abs(xticks-fkvertex[0]).argmin(), np.abs(yticks[::-1]-fkvertex[1]).argmin()]) vertices = np.array(vert_tmp) except AttributeError: vertices = (poly.get_path().vertices).astype('int') indicies = convert_polygon_to_flat_index(data, vertices) return indicies
class MaskDrawer(object): """An interactive polygon mask drawer on an image. Parameters ---------- ax : matplotlib plot axis Inpimg : 2d numpy array Input image to overlay for drawing mask Mask : Boolean numpy array same size of inpimg A Mask which will be used as initial mask and updated upon confirmation max_ds : float Max pixel distance to count as a vertex hit. PolyAtStart : List of vertices A list of square vertices to draw the initial polygon Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices 'n' : Invert the region selected by polynomial for masking 'c' : Confirm the polygon and update the mask """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit def __init__(self, ax, Inpimg, Mask, max_ds=10, PolyAtStart=[(50, 50), (100, 50), (100, 100), (50, 100)]): self.showverts = True self.max_ds = max_ds self.Mask = Mask self.img = Inpimg self.maskinvert = False # imshow the image self.imgplot = ax.imshow(np.ma.filled( np.ma.array(self.img, mask=self.Mask, fill_value=np.nan)), cmap=cm.gray) self.poly = Polygon(PolyAtStart, animated=True, fc='y', ec='none', alpha=0.5) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title( "Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "'n' to invert the region for masking, 'c' to confirm & apply the mask." ) self.ax = ax x, y = zip(*self.poly.xy) self.line = Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.7, animated=True) # self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt - event.x)**2 + (yt - event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind] >= self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes == None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key == 't': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key == 'd': ind = self.get_ind_under_point(event) if ind is not None: self.poly.xy = [ tup for i, tup in enumerate(self.poly.xy) if i != ind ] self.line.set_data(zip(*self.poly.xy)) elif event.key == 'i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # display coords for i in range(len(xys) - 1): s0 = xys[i] s1 = xys[i + 1] d = dist_point_to_segment(p, s0, s1) if d <= self.epsilon: self.poly.xy = np.array( list(self.poly.xy[:i]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i:])) self.line.set_data(zip(*self.poly.xy)) break elif event.key == 'n': """ Flips the region inside out of the polygon to be masked """ print('Inverting the mask region') self.maskinvert = not self.maskinvert elif event.key == 'c': """ Confirm the drawn polynomial shape and add update the mask """ self.UpdateMask() #Update the imshowed image with new mask self.imgplot.set_data( np.ma.filled( np.ma.array(self.img, mask=self.Mask, fill_value=np.nan))) self.imgplot.figure.canvas.draw() self.canvas.draw() def UpdateMask(self): """ Updates the maks array with points insied the polygon """ print('Updating the original Mask..') Path = self.poly.get_path() h, w = self.Mask.shape y, x = np.mgrid[:h, :w] XYpoints = np.transpose((x.ravel(), y.ravel())) NewMask = Path.contains_points(XYpoints) if self.maskinvert: NewMask = ~NewMask # Combine the two mask by taing an elemet wise or self.Mask = self.Mask | NewMask.reshape((h, w)) def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x, y = event.xdata, event.ydata self.poly.xy[self._ind] = x, y self.line.set_data(zip(*self.poly.xy)) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox)
def parameter_pdf(self, fig_comment='', limits_filename=None): """Calculate a pdf for parameter values based on the uncertainty model """ try: self.parameter_space except AttributeError: self.rupture_params_2_array() xlabel_dict = { 'mag': 'Magnitude ($M_w$)', 'longitude': 'Longitude', 'latitude': 'Latitude', 'depth': 'Depth (km)', 'strike': 'Strike ($^\circ$)', 'dip': 'Dip ($^\circ$)' } #self.parameter_pdfs = self.parameter_space*self.uncert_fun.pdf(self.rmse) # Now sum equal values to build pdfs self.parameter_pdf_sums = {} self.parameter_pdf_values = {} plt.clf() fig = plt.figure(figsize=(16, 8)) #, tight_layout=True) gs = plt.GridSpec(2, 4) for key, value in self.parameter_dict.iteritems(): unique_vals = np.unique(self.parameter_space[value]) pdf_sums = [] bin = False # Determine if we need to bin values of strike and dip (i.e. # for non-area sources if key == 'strike': if len(unique_vals) > 24: bin = True bins = np.arange(0, 360, 15.) if key == 'dip' or key == 'depth': if len(unique_vals) > 4: bin = True if key == 'dip': bins = np.arange(min(self.parameter_space[value]), max(self.parameter_space[value]) + 5, 5.) elif key == 'depth': bins = np.arange(min(self.parameter_space[value]), max(self.parameter_space[value]) + 20, 20.) if bin: # Calculate as histogram hist, bins = np.histogram(self.parameter_space[value], bins) # align to bin centre for plotting #bin_width = bins[1] - bins[0] unique_vals = [] for i, edge in enumerate(bins): try: ind = np.intersect1d( np.where(self.parameter_space[value] >= edge), np.where( self.parameter_space[value] < bins[i + 1])) except IndexError: ind = np.where(self.parameter_space[value] >= edge) pdf_sum = 0 if len(ind) < 1: print ind else: for index in ind: # print index try: likelihood = np.power((1/(self.sigma*np.sqrt(2*np.pi))), len(self.mmi_obs)) * \ np.exp((-1/2)*((self.sum_squares_list[index]/self.sigma**2))) pdf_sum += likelihood except TypeError: print 'ind', ind print 'index', index pdf_sums.append(pdf_sum) unique_vals.append(edge) # + bin_width) else: # Use raw values for val in unique_vals: # Just unique values, as pdf sums are repeated ind = np.argwhere(self.parameter_space[value] == val) pdf_sum = 0 # print ind for index in ind: # print index likelihood = np.power((1/(self.sigma*np.sqrt(2*np.pi))), len(self.mmi_obs)) * \ np.exp((-1/2)*((self.sum_squares_list[index]/self.sigma**2))) pdf_sum += likelihood #pdf_sum = np.sum(self.uncert_fun.pdf(self.rmse[ind])) pdf_sums.append(pdf_sum) # Normalise pdf sums pdf_sums = np.array(pdf_sums) pdf_sums = pdf_sums / np.sum(pdf_sums) print 'pdf_sums', pdf_sums self.parameter_pdf_sums[key] = pdf_sums self.parameter_pdf_values[key] = unique_vals # Get the best-fit value for plotting on top index = np.argmin(self.rmse) best_fit_x = self.parameter_space[value][index] #y_index = np.where(unique_vals == best_fit_x)[0] try: y_index = (np.abs(unique_vals - best_fit_x)).argmin()[0] except IndexError: y_index = (np.abs(unique_vals - best_fit_x)).argmin() best_fit_y = pdf_sums[y_index] # Now plot the results try: width = unique_vals[1] - unique_vals[ 0] # Scale width by discretisation except IndexError: # if only one value width = 1.0 if key == 'strike': # Plot as rose diagram ax = fig.add_subplot(gs[0, 2], projection='polar') ax.bar(np.deg2rad(unique_vals), pdf_sums, width=np.deg2rad(width), bottom=0.0, align='center', color='0.5', edgecolor='k') ax.scatter(np.deg2rad(best_fit_x), best_fit_y, marker='*', color='0.5', edgecolor='k', s=200, zorder=10) ax.set_theta_zero_location('N') ax.set_theta_direction(-1) ax.set_thetagrids(np.arange(0, 360, 15)) # Define grids intervals for radial axis if max(pdf_sums) < 0.11: r_int = 0.02 else: r_int = 0.1 ax.set_rgrids(np.arange(0.01, max(pdf_sums) + 0.01, r_int), angle=np.deg2rad(7.5), weight='black') ax.set_xlabel(xlabel_dict[key]) ax.text(-0.07, 1.02, 'b)', transform=ax.transAxes, fontsize=14) #ax.set_title('Rose Diagram of Strike PDF"', y=1.10, fontsize=15) #plt.savefig('%s_%s_%s_parameter_pdf_rose_diagram.png' % (fig_comment,self.gsim, key), \ # dpi=300, format='png', bbox_inches='tight') elif key == 'mag': ax = fig.add_subplot(gs[:, 0:2]) ymax = max(pdf_sums) * 1.1 lbx = [self.min_mag, self.min_mag] lby = [0, ymax] ubx = [self.max_mag, self.max_mag] uby = [0, ymax] elif key == 'depth': ax = fig.add_subplot(gs[1, 3]) ymax = max(pdf_sums) * 1.1 lbx = [self.min_depth, self.min_depth] lby = [0, ymax] ubx = [self.max_depth, self.max_depth] uby = [0, ymax] elif key == 'dip': ax = fig.add_subplot(gs[1, 2]) ymax = max(pdf_sums) * 1.1 lbx = [self.min_dip, self.min_dip] lby = [0, ymax] ubx = [self.max_dip, self.max_dip] uby = [0, ymax] if key == 'mag' or key == 'dip' or key == 'depth': ax.bar(unique_vals, pdf_sums, width, align='center', color='0.5') ax.scatter(best_fit_x, best_fit_y, marker='*', color='0.5', edgecolor='k', s=200, zorder=10) if key != 'latitude' and key != 'longitude': ax.plot(lbx, lby, color='k') ax.plot(ubx, uby, color='k') ax.set_ylim(0, ymax) ax.set_ylabel('Probability') ax.set_xlabel(xlabel_dict[key]) if key == 'mag': ax.text(0.05, 0.95, 'a)', transform=ax.transAxes, fontsize=14) if key == 'dip': ax.text(0.05, 0.92, 'd)', transform=ax.transAxes, fontsize=14) if key == 'depth': ax.text(0.05, 0.92, 'e)', transform=ax.transAxes, fontsize=14) # Now plot a map of location uncertainty # First get 2D pdf pdf_sums = [] all_lons = [] all_lats = [] for i, lon in enumerate( self.parameter_space[self.parameter_dict['longitude']]): if lon in all_lons: continue else: lat = self.parameter_space[self.parameter_dict['latitude']][i] # lat = self.parameter_pdf_values['latitude'][i] index = np.intersect1d(np.argwhere(self.parameter_space[self.parameter_dict['longitude']]==lon), \ np.argwhere(self.parameter_space[self.parameter_dict['latitude']]==lat)) pdf_sum = np.sum(self.uncert_fun.pdf(self.rmse[index])) pdf_sums.append(pdf_sum) all_lons.append(lon) all_lats.append(lat) # Normalise pdf sums pdf_sums = np.array(pdf_sums) pdf_sums = pdf_sums / np.sum(pdf_sums) self.parameter_pdf_sums['lon_lat'] = pdf_sums all_lons = np.array(all_lons) all_lats = np.array(all_lats) # Get best fit value index = np.argmin(self.rmse) best_fit_lon = self.parameter_space[ self.parameter_dict['longitude']][index] best_fit_lat = self.parameter_space[ self.parameter_dict['latitude']][index] ax = fig.add_subplot(gs[0, 3]) minlon = min(self.parameter_pdf_values['longitude']) maxlon = max(self.parameter_pdf_values['longitude']) minlat = min(self.parameter_pdf_values['latitude']) maxlat = max(self.parameter_pdf_values['latitude']) lat_0 = minlat + (maxlat - minlat) / 2. lon_0 = minlon + (maxlon - minlon) / 2. m = Basemap(projection='tmerc', lat_0=lat_0, lon_0=lon_0, llcrnrlon=minlon, llcrnrlat=minlat, urcrnrlon=maxlon, urcrnrlat=maxlat, resolution='i') m.drawcoastlines(linewidth=0.5, color='k') m.drawcountries(color='0.2') m.drawstates(color='0.2') if maxlon - minlon < 2: gridspace = 0.5 elif maxlon - minlon < 3: gridspace = 1.0 elif maxlon - minlon < 7: gridspace = 2.0 else: gridspace = 5.0 m.drawparallels(np.arange(-90., 90., gridspace), labels=[1, 0, 0, 0], fontsize=10, dashes=[2, 2], color='0.5', linewidth=0.5) m.drawmeridians(np.arange(0., 360., gridspace), labels=[0, 0, 0, 1], fontsize=10, dashes=[2, 2], color='0.5', linewidth=0.5) max_val = max(pdf_sums) * 1.1 print 'pdf_sums', pdf_sums clevs = np.arange(0.0, max_val, (max_val / 50)) cmap = plt.get_cmap('gray_r') # Adjust resolution to avoid memory intense interpolations res = max((maxlon - minlon) / 50., (maxlat - minlat) / 50.) xy = np.mgrid[minlon:maxlon:res, minlat:maxlat:res] xx, yy = np.meshgrid(xy[0, :, 0], xy[1][0]) griddata = interpolate.griddata((all_lons, all_lats), pdf_sums, (xx, yy), method='nearest') # now plot filled contours of pdf cs = m.contourf(xx, yy, griddata, clevs, cmap=cmap, vmax=max(pdf_sums), vmin=min(pdf_sums), latlon=True) # Mask areas outside of source model if limits_filename is not None: limits_data = np.genfromtxt(limits_filename, delimiter=',') limits_x = limits_data[:, 0] limits_y = limits_data[:, 1] limits_x, limits_y = m(limits_x, limits_y) # Convert to map coordinates poly = Polygon(np.c_[limits_x, limits_y], closed=True) clippath = poly.get_path() ax = plt.gca() patch = PathPatch(clippath, transform=ax.transData, facecolor='none', linewidth=0.4) print 'Adding patch' ax.add_patch(patch) for contour in cs.collections: contour.set_clip_path(patch) # Now add best-fit location on top m.scatter(best_fit_lon, best_fit_lat, marker='*', color='w', edgecolor='k', s=200, zorder=10, latlon=True) #m.text(0.05, 0.95, 'c)', transform=ax.transAxes, fontsize=14) plt.annotate('c)', xy=(0.05, 0.9), xycoords='axes fraction', fontsize=14) if max_val < 0.001: loc_int = 0.0005 elif max_val < 0.01: loc_int = 0.005 else: loc_int = 0.05 ticks = np.arange(0.0, max_val * 1.1, loc_int) cbar = m.colorbar(cs, ticks=ticks) cbar.ax.set_ylabel('Probability') figname = '%s_%s_all_parameter_pdf.png' % (fig_comment, self.gsim) figname = figname.replace('()', '') plt.tight_layout() plt.savefig(figname, dpi=300, format='png', bbox_inches='tight')
class MaskDrawer(object): """An interactive polygon mask drawer on an image. Parameters ---------- ax : matplotlib plot axis Inpimg : 2d numpy array Input image to overlay for drawing mask Mask : Boolean numpy array same size of inpimg A Mask which will be used as initial mask and updated upon confirmation max_ds : float Max pixel distance to count as a vertex hit. PolyAtStart : List of vertices A list of square vertices to draw the initial polygon Key-bindings ------------ 't' : toggle vertex markers on and off. When vertex markers are on, you can move them, delete them 'd' : delete the vertex under point 'i' : insert a vertex at point. You must be within max_ds of the line connecting two existing vertices 'n' : Invert the region selected by polynomial for masking 'c' : Confirm the polygon and update the mask """ showverts = True epsilon = 5 # max pixel distance to count as a vertex hit def __init__(self, ax, Inpimg, Mask, max_ds=10,PolyAtStart = [(50,50),(100,50),(100,100),(50,100)]): self.showverts = True self.max_ds = max_ds self.Mask = Mask self.img = Inpimg self.maskinvert = False # imshow the image self.imgplot = ax.imshow(np.ma.filled(np.ma.array(self.img,mask=self.Mask,fill_value=np.nan)), cmap=cm.gray) self.poly = Polygon(PolyAtStart, animated=True, fc='y', ec='none', alpha=0.5) ax.add_patch(self.poly) ax.set_clip_on(False) ax.set_title("Click and drag a point to move it; " "'i' to insert; 'd' to delete.\n" "'n' to invert the region for masking, 'c' to confirm & apply the mask.") self.ax = ax x, y = zip(*self.poly.xy) self.line = Line2D(x, y, color='none', marker='o', mfc='r', alpha=0.7, animated=True) # self._update_line() self.ax.add_line(self.line) self.poly.add_callback(self.poly_changed) self._ind = None # the active vert canvas = self.poly.figure.canvas canvas.mpl_connect('draw_event', self.draw_callback) canvas.mpl_connect('button_press_event', self.button_press_callback) canvas.mpl_connect('button_release_event', self.button_release_callback) canvas.mpl_connect('key_press_event', self.key_press_callback) canvas.mpl_connect('motion_notify_event', self.motion_notify_callback) self.canvas = canvas def draw_callback(self, event): self.background = self.canvas.copy_from_bbox(self.ax.bbox) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox) def poly_changed(self, poly): 'this method is called whenever the polygon object is called' # only copy the artist props to the line (except visibility) vis = self.line.get_visible() Artist.update_from(self.line, poly) self.line.set_visible(vis) # don't use the poly visibility state def get_ind_under_point(self, event): 'get the index of the vertex under point if within epsilon tolerance' # display coords xy = np.asarray(self.poly.xy) xyt = self.poly.get_transform().transform(xy) xt, yt = xyt[:, 0], xyt[:, 1] d = np.sqrt((xt-event.x)**2 + (yt-event.y)**2) indseq = np.nonzero(np.equal(d, np.amin(d)))[0] ind = indseq[0] if d[ind]>=self.epsilon: ind = None return ind def button_press_callback(self, event): 'whenever a mouse button is pressed' if not self.showverts: return if event.inaxes==None: return if event.button != 1: return self._ind = self.get_ind_under_point(event) def button_release_callback(self, event): 'whenever a mouse button is released' if not self.showverts: return if event.button != 1: return self._ind = None def key_press_callback(self, event): 'whenever a key is pressed' if not event.inaxes: return if event.key=='t': self.showverts = not self.showverts self.line.set_visible(self.showverts) if not self.showverts: self._ind = None elif event.key=='d': ind = self.get_ind_under_point(event) if ind is not None: self.poly.xy = [tup for i,tup in enumerate(self.poly.xy) if i!=ind] self.line.set_data(zip(*self.poly.xy)) elif event.key=='i': xys = self.poly.get_transform().transform(self.poly.xy) p = event.x, event.y # display coords for i in range(len(xys)-1): s0 = xys[i] s1 = xys[i+1] d = dist_point_to_segment(p, s0, s1) if d<=self.epsilon: self.poly.xy = np.array( list(self.poly.xy[:i]) + [(event.xdata, event.ydata)] + list(self.poly.xy[i:])) self.line.set_data(zip(*self.poly.xy)) break elif event.key=='n': """ Flips the region inside out of the polygon to be masked """ print('Inverting the mask region') self.maskinvert = not self.maskinvert elif event.key=='c': """ Confirm the drawn polynomial shape and add update the mask """ self.UpdateMask() #Update the imshowed image with new mask self.imgplot.set_data(np.ma.filled(np.ma.array(self.img,mask=self.Mask,fill_value=np.nan))) self.imgplot.figure.canvas.draw() self.canvas.draw() def UpdateMask(self): """ Updates the maks array with points insied the polygon """ print('Updating the original Mask..') Path = self.poly.get_path() h, w = self.Mask.shape y, x = np.mgrid[:h,:w] XYpoints = np.transpose((x.ravel(), y.ravel())) NewMask = Path.contains_points(XYpoints) if self.maskinvert : NewMask = ~NewMask # Combine the two mask by taing an elemet wise or self.Mask = self.Mask | NewMask.reshape((h,w)) def motion_notify_callback(self, event): 'on mouse movement' if not self.showverts: return if self._ind is None: return if event.inaxes is None: return if event.button != 1: return x,y = event.xdata, event.ydata self.poly.xy[self._ind] = x,y self.line.set_data(zip(*self.poly.xy)) self.canvas.restore_region(self.background) self.ax.draw_artist(self.poly) self.ax.draw_artist(self.line) self.canvas.blit(self.ax.bbox)
for i in range(len(internalEdges)): reqEdge = internalEdges[i] figs = analyzeFig(figs, reqEdge) # check internals, one fig inside if fMap[inputFace].internal != None and not isinstance( fMap[inputFace].internal, list): reqEdge = fMap[inputFace].internal figs = analyzeFig(figs, reqEdge) fig = plt.figure() fig.add_subplot() ax1 = plt.gca() for f in figs: xp = [p[0] for p in f] yp = [p[1] for p in f] ax1.scatter(xp, yp, s=100, marker="o", zorder=10, color='black') p = Polygon(np.array(f), facecolor='powderblue') ax1.add_patch(p) path = p.get_path() patch = PathPatch(path, facecolor='powderblue', lw=2) ax1.add_patch(patch) plt.show() #printMap(objMap) #verts = [objMap[key] for key in objMap.keys() if "p" in key] #edges = [objMap[key] for key in objMap.keys() if "s" in key] #faces = [objMap[key] for key in objMap.keys() if "f" in key] #print(verts, edges, faces)
except: for coords in s.find_all('coordinates'): space_splits = coords.text.split(' ') lat = [] lng = [] #tab_split=[] space_splits[0] = space_splits[0].split('\t')[-1] del space_splits[-1] #for t in space_splits: # tab_split.append(t.split('\t')[5]) for split in space_splits: comma_split = split.split(',') lat.append(comma_split[1]) # lat lng.append(comma_split[0]) # lng lat = [float(i) for i in lat] lng = [float(i) for i in lng] #print(lat) coords = [[i, j] for i, j in zip(lat, lng)] coords2 = np.array(coords) poly = Polygon(coords2, closed=True) if poly.get_path().contains_point(a): found = True block.append(j) if found == False: block.append(np.nan) shutil.rmtree(r'C:/Users/koku8001/Desktop/renaming/unzipped') zip_ref.close() output_final_file['Block'] = block
class Map(): def __init__(self, proj, ax=None, interactive=True, **kwargs): """Create Map with a given projection. A `skymapper.Map` holds a predefined projection and `matplotlib` axes and figures to enable plotting on the sphere with proper labeling and inter/extrapolations. It also allows for interactive and exploratory work by updateing the maps after pan/zoom events. Most of the methods are wrappers of `matplotlib` functions by the same names, so that one can mostly interact with a `Map` instance as one would do with a `matplotlib.axes`. For plotting purposes, it is recommended to switch `interactive` off. Args: proj: `skymapper.Projection` instance ax: `matplotlib.axes` instance, will be created otherwise interactive: if pan/zoom is enabled for map updates **kwargs: styling of the `matplotlib.patches.Polygon` that shows the outline of the map. """ # store arguments to regenerate the map self._config = {'__init__': _parseArgs(locals())} self.proj = proj self._setFigureAx(ax, interactive=interactive) self._resolution = 75 # for graticules self._setEdge(**kwargs) self.ax.relim() self.ax.autoscale_view() self._setFrame() self.fig.tight_layout(pad=0.75) def _setFigureAx(self, ax=None, interactive=True): if ax is None: self.fig = matplotlib.pyplot.figure() self.ax = self.fig.add_subplot(111, aspect='equal') else: self.ax = ax self.ax.set_aspect('equal') self.fig = self.ax.get_figure() self.ax.set_axis_off() # do not unset the x/y ticks by e.g. xticks([]), we need them for tight_layout self.ax.xaxis.set_ticks_position('none') self.ax.yaxis.set_ticks_position('none') # attach event handlers if interactive: self.fig.show() self._press_evt = self.fig.canvas.mpl_connect('button_press_event', self._pressHandler) self._release_evt = self.fig.canvas.mpl_connect('button_release_event', self._releaseHandler) self._scroll_evt = self.fig.canvas.mpl_connect('scroll_event', self._scrollHandler) def clone(self, ax=None): """Clone map Args: ax: `matplotlib.axes` instance, will be created otherwise Returns: New map using the same projections and configuration """ config = dict(self._config) config['xlim'] = self.ax.get_xlim() config['ylim'] = self.ax.get_ylim() return Map._create(config, ax=ax) def save(self, filename): """Save map configuration to file All aspects necessary to reproduce a map are stored in a pickle file. Args: filename: name for pickle file Returns: None """ try: with open(filename, 'wb') as fp: config = dict(self._config) config['xlim'] = self.ax.get_xlim() config['ylim'] = self.ax.get_ylim() pickle.dump(config, fp) except IOError as e: raise @staticmethod def load(filename, ax=None): """Load map from pickled file Args: filename: name for pickle file ax: `matplotlib.axes` instance, will be created otherwise Returns: `skymapper.Map` """ try: with open(filename, 'rb') as fp: config = pickle.load(fp) fp.close() return Map._create(config, ax=ax) except IOError as e: raise @staticmethod def _create(config, ax=None): init_args = config.pop('__init__') init_args['ax'] = ax map = Map(**init_args) xlim = config.pop('xlim') ylim = config.pop('ylim') map.ax.set_xlim(xlim) map.ax.set_ylim(ylim) map._setFrame() meridian_args = config.pop('labelMeridianAtParallel', {}) parallel_args = config.pop('labelParallelAtMeridian', {}) for method in config.keys(): getattr(map, method)(**config[method]) for args in meridian_args.values(): map.labelMeridianAtParallel(**args) for args in parallel_args.values(): map.labelParallelAtMeridian(**args) map.fig.tight_layout(pad=0.75) return map @property def parallels(self): """Get the location of the drawn parallels""" return [ float(m.group(1)) for c,m in self.artists(r'grid-parallel-([\-\+0-9.]+)', regex=True) ] @property def meridians(self): """Get the location of the drawn meridians""" return [ float(m.group(1)) for c,m in self.artists(r'grid-meridian-([\-\+0-9.]+)', regex=True) ] def artists(self, gid, regex=False): """Get the `matplotlib` artists used in the map Args: gid: `gid` string of the artist regex: if regex matching is done Returns: list of matching artists if `regex==True`, returns list of (artist, match) """ if regex: matches = [ re.match(gid, c.get_gid()) if c.get_gid() is not None else None for c in self.ax.get_children() ] return [ (c,m) for c,m in zip(self.ax.get_children(), matches) if m is not None ] else: # direct match return [ c for c in self.ax.get_children() if c.get_gid() is not None and c.get_gid().find(gid) != -1 ] def _getParallel(self, p, reverse=False): if not reverse: return self.proj.transform(self._ra_range, p*np.ones(len(self._ra_range))) return self.proj.transform(self._ra_range[::-1], p*np.ones(len(self._ra_range))) def _getMeridian(self, m, reverse=False): if not reverse: return self.proj.transform(m*np.ones(len(self._dec_range)), self._dec_range) return self.proj.transform(m*np.ones(len(self._dec_range)), self._dec_range[::-1]) def _setParallel(self, p, **kwargs): x, y = self._getParallel(p) artist = Line2D(x, y, **kwargs) self.ax.add_line(artist) return artist def _setMeridian(self, m, **kwargs): x, y = self._getMeridian(m) artist = Line2D(x, y, **kwargs) self.ax.add_line(artist) return artist def _setEdge(self, **kwargs): self._dec_range = np.linspace(-90, 90, self._resolution) self._ra_range = np.linspace(-180, 180, self._resolution) + self.proj.ra_0 # styling: frame needs to be on top of everything, must be transparent facecolor = 'None' zorder = 1000 lw = kwargs.pop('lw', 0.7) edgecolor = kwargs.pop('edgecolor', 'k') # if there is facecolor: clone the polygon and put it in as bottom layer facecolor_ = kwargs.pop('facecolor', '#dddddd') # polygon of the map edge: top, left, bottom, right # don't draw poles if that's a single point lines = [self._getMeridian(self.proj.ra_0 + 180, reverse=True), self._getMeridian(self.proj.ra_0 - 180)] if not self.proj.poleIsPoint[-90]: lines.insert(1, self._getParallel(-90, reverse=True)) if not self.proj.poleIsPoint[90]: lines.insert(0, self._getParallel(90)) xy = np.concatenate(lines, axis=1).T self._edge = Polygon(xy, closed=True, edgecolor=edgecolor, facecolor=facecolor, lw=lw, zorder=zorder,gid="edge", **kwargs) self.ax.add_patch(self._edge) if facecolor_ is not None: zorder = -1000 edgecolor = 'None' poly = Polygon(xy, closed=True, edgecolor=edgecolor, facecolor=facecolor_, zorder=zorder, gid="edge-background") self.ax.add_patch(poly) def contains(self, x, y): xy = np.dstack((x,y)).reshape(-1,2) return self._edge.get_path().contains_points(xy).reshape(x.shape) def xlim(self, left=None, right=None, **kw): """Get/set the map limits in x-direction""" if left is None and right is None: return (self._edge.xy[:, 0].min(), self._edge.xy[:, 0].max()) else: self.ax.set_xlim(left=left, right=right, **kw) self._resetFrame() def ylim(self, bottom=None, top=None, **kw): """Get/set the map limits in x-direction""" if bottom is None and top is None: return (self._edge.xy[:, 1].min(), self._edge.xy[:, 1].max()) else: self.ax.set_ylim(bottom=bottom, top=top, **kw) self._resetFrame() def grid(self, sep=30, parallel_fmt=degPMFormatter, meridian_fmt=deg360Formatter, dec_min=-90, dec_max=90, ra_min=-180, ra_max=180, **kwargs): """Set map grid / graticules Args: sep: distance between graticules in deg parallel_fmt: formatter for parallel labels meridian_fmt: formatter for meridian labels dec_min: minimum declination for graticules dec_max: maximum declination for graticules ra_min: minimum declination for graticules (for which reference RA=0) ra_max: maximum declination for graticules (for which reference RA=0) **kwargs: styling of `matplotlib.lines.Line2D` for the graticules """ self._config['grid'] = _parseArgs(locals()) self._dec_range = np.linspace(dec_min, dec_max, self._resolution) self._ra_range = np.linspace(ra_min, ra_max, self._resolution) + self.proj.ra_0 _parallels = np.arange(-90+sep,90,sep) if self.proj.ra_0 % sep == 0: _meridians = np.arange(sep * ((self.proj.ra_0 + 180) // sep), sep * ((self.proj.ra_0 - 180) // sep - 1), -sep) else: _meridians = np.arange(sep * ((self.proj.ra_0 + 180) // sep), sep * ((self.proj.ra_0 - 180) // sep), -sep) # clean up previous grid artists = self.artists('grid-meridian') + self.artists('grid-parallel') for artist in artists: artist.remove() # clean up frame meridian and parallel labels because they're tied to the grid artists = self.artists('meridian-label') + self.artists('parallel-label') for artist in artists: artist.remove() # styling: based on edge ls = kwargs.pop('ls', '-') lw = kwargs.pop('lw', self._edge.get_linewidth() / 2) c = kwargs.pop('c', self._edge.get_edgecolor()) alpha = kwargs.pop('alpha', 0.2) zorder = kwargs.pop('zorder', self._edge.get_zorder() - 1) for p in _parallels: self._setParallel(p, gid='grid-parallel-%r' % p, lw=lw, c=c, alpha=alpha, zorder=zorder, **kwargs) for m in _meridians: self._setMeridian(m, gid='grid-meridian-%r' % m, lw=lw, c=c, alpha=alpha, zorder=zorder, **kwargs) # (re)generate the frame labels for method in ['labelMeridiansAtFrame', 'labelParallelsAtFrame']: if method in self._config.keys(): getattr(self, method)(**self._config[method]) else: getattr(self, method)() # (re)generate edge labels for method in ['labelMeridianAtParallel', 'labelParallelAtMeridian']: if method in self._config.keys(): args_list = self._config.pop(method, []) for args in args_list.values(): getattr(self, method)(**args) else: # label meridians: at the poles if they are not points if method == 'labelMeridianAtParallel': # determine the parallel that has the most space for labels dec = [-90, 0, 90] ra = [self.proj.ra_0,] * 3 jac = np.sum(self.proj.gradient(ra, dec)**2, axis=1) p = dec[np.argmax(jac)] # remove outer meridians to prevent overlap with parallel label if p == 0 and self.proj.ra_0 % sep == 0: _meridians = _meridians[1:-1] getattr(self, method)(p, meridians=_meridians) # label both outer meridians else: degs = [self.proj.ra_0 + 180, self.proj.ra_0 - 180] for deg in degs: getattr(self, method)(deg) def _negateLoc(self, loc): if loc == "bottom": return "top" if loc == "top": return "bottom" if loc == "left": return "right" if loc == "right": return "left" def labelMeridianAtParallel(self, p, loc=None, meridians=None, pad=None, direction='parallel', **kwargs): """Label the meridians intersecting a given parallel The method is called by `grid()` but can be used to overwrite the defaults. Args: p: parallel in deg loc: location of the label with respect to `p`, from `['top', 'bottom']` meridians: list of meridians to label, if None labels all of them pad: padding of annotation, in units of fontsize direction: tangent of the label, from `['parallel', 'meridian']` **kwargs: styling of `matplotlib` annotations for the graticule labels """ arguments = _parseArgs(locals()) if p in self.proj.poleIsPoint.keys() and self.proj.poleIsPoint[p]: return myname = 'labelMeridianAtParallel' if myname not in self._config.keys(): self._config[myname] = dict() # remove exisiting labels at p gid = 'meridian-label-%r' % p if p in self._config[myname].keys(): artists = self.artists(gid) for artist in artists: artist.remove() self._config[myname][p] = arguments if meridians is None: meridians = self.meridians # determine rot_base so that central label is upright rotation = kwargs.pop('rotation', None) if rotation is None or loc is None: m = self.proj.ra_0 dxy = self.proj.gradient(m, p, direction=direction) angle = np.arctan2(*dxy) / DEG2RAD options = [-90, 90] closest = np.argmin(np.abs(options - angle)) rot_base = options[closest] if loc is None: if p >= 0: loc = 'top' else: loc = 'bottom' assert loc in ['top', 'bottom'] horizontalalignment = kwargs.pop('horizontalalignment', 'center') verticalalignment = kwargs.pop('verticalalignment', self._negateLoc(loc)) zorder = kwargs.pop('zorder', 20) size = kwargs.pop('size', matplotlib.rcParams['font.size']) # styling consistent with frame, i.e. with edge color = kwargs.pop('color', self._edge.get_edgecolor()) alpha = kwargs.pop('alpha', self._edge.get_alpha()) zorder = kwargs.pop('zorder', self._edge.get_zorder() + 1) # on top of edge if pad is None: pad = size / 3 for m in meridians: # move label along meridian xp, yp = self.proj.transform(m, p) dxy = self.proj.gradient(m, p, direction="meridian") dxy *= pad / np.sqrt((dxy**2).sum()) if loc == 'bottom': # dxy in positive RA dxy *= -1 if rotation is None: dxy_ = self.proj.gradient(m, p, direction=direction) angle = rot_base - np.arctan2(*dxy_) / DEG2RAD else: angle = rotation self.ax.annotate(self._config['grid']['meridian_fmt'](m), (xp, yp), xytext=dxy, textcoords='offset points', rotation=angle, rotation_mode='anchor', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, gid=gid, **kwargs) def labelParallelAtMeridian(self, m, loc=None, parallels=None, pad=None, direction='parallel', **kwargs): """Label the parallel intersecting a given meridian The method is called by `grid()` but can be used to overwrite the defaults. Args: m: meridian in deg loc: location of the label with respect to `m`, from `['left', 'right']` parallel: list of parallels to label, if None labels all of them pad: padding of annotation, in units of fontsize direction: tangent of the label, from `['parallel', 'meridian']` **kwargs: styling of `matplotlib` annotations for the graticule labels """ arguments = _parseArgs(locals()) myname = 'labelParallelAtMeridian' if myname not in self._config.keys(): self._config[myname] = dict() # remove exisiting labels at m gid = 'parallel-label-%r' % m if m in self._config[myname].keys(): artists = self.artists(gid) for artist in artists: artist.remove() self._config[myname][m] = arguments # determine rot_base so that central label is upright rotation = kwargs.pop('rotation', None) if rotation is None or loc is None: p = 0 dxy = self.proj.gradient(m, p, direction=direction) angle = np.arctan2(*dxy) / DEG2RAD options = [-90, 90] closest = np.argmin(np.abs(options - angle)) rot_base = options[closest] if loc is None: if m < self.proj.ra_0: # meridians on the left: dx goes in positive RA dxy *= -1 if dxy[0] > 0: loc = 'right' else: loc = 'left' assert loc in ['left', 'right'] if parallels is None: parallels = self.parallels horizontalalignment = kwargs.pop('horizontalalignment', self._negateLoc(loc)) verticalalignment = kwargs.pop('verticalalignment', 'center') zorder = kwargs.pop('zorder', 20) size = kwargs.pop('size', matplotlib.rcParams['font.size']) # styling consistent with frame, i.e. with edge color = kwargs.pop('color', self._edge.get_edgecolor()) alpha = kwargs.pop('alpha', self._edge.get_alpha()) zorder = kwargs.pop('zorder', self._edge.get_zorder() + 1) # on top of edge if pad is None: pad = size/2 # more space for horizontal parallels for p in parallels: # move label along parallel xp, yp = self.proj.transform(m, p) dxy = self.proj.gradient(m, p, direction="parallel") dxy *= pad / np.sqrt((dxy**2).sum()) if m < self.proj.ra_0: # meridians on the left: dx goes in positive RA dxy *= -1 if rotation is None: dxy_ = self.proj.gradient(m, p, direction=direction) angle = rot_base - np.arctan2(*dxy_) / DEG2RAD else: angle = rotation self.ax.annotate(self._config['grid']['parallel_fmt'](p), (xp, yp), xytext=dxy, textcoords='offset points', rotation=angle, rotation_mode='anchor', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, gid=gid, **kwargs) def labelMeridiansAtFrame(self, loc='top', meridians=None, pad=None, description=None, **kwargs): """Label the meridians on rectangular frame of the map If the view only shows a fraction of the map, a segment or an entire rectangular frame is shown and the graticule labels are moved to outside that frame. This method is implicitly called, but can be used to overwrite the defaults. Args: loc: location of the label with respect to frame, from `['top', 'bottom']` meridians: list of meridians to label, if None labels all of them pad: padding of annotation, in units of fontsize description: equivalent to `matplotlib` axis label **kwargs: styling of `matplotlib` annotations for the graticule labels """ assert loc in ['bottom', 'top'] arguments = _parseArgs(locals()) myname = 'labelMeridiansAtFrame' # need extra space for tight_layout to consider the frame annnotations # we can't get the actual width, but we can make use the of the default width of the axes tick labels if myname not in self._config.keys() or self._config[myname]['loc'] != loc: self.ax.xaxis.set_ticks_position(loc) self.ax.xaxis.set_label_position(loc) # remove existing frame_artists = self.artists('frame-meridian-label') for artist in frame_artists: artist.remove() self.fig.tight_layout(pad=0.75) self._config[myname] = arguments size = kwargs.pop('size', matplotlib.rcParams['font.size']) # styling consistent with frame, i.e. with edge color = kwargs.pop('color', self._edge.get_edgecolor()) alpha = kwargs.pop('alpha', self._edge.get_alpha()) zorder = kwargs.pop('zorder', self._edge.get_zorder() + 1) # on top of edge horizontalalignment = kwargs.pop('horizontalalignment', 'center') verticalalignment = self._negateLoc(loc) # no option along the frame _ = kwargs.pop('verticalalignment', None) if pad is None: pad = size / 3 if meridians is None: meridians = self.meridians # check if loc has frame poss = {"bottom": 0, "top": 1} pos = poss[loc] xlim, ylim = self.ax.get_xlim(), self.ax.get_ylim() xticks = [] frame_artists = self.artists(r'frame-([a-zA-Z]+)', regex=True) frame_locs = [match.group(1) for c,match in frame_artists] if loc in frame_locs: # find all parallel grid lines m_artists = self.artists(r'grid-meridian-([\-\+0-9.]+)', regex=True) for c,match in m_artists: m = float(match.group(1)) if m in meridians: # intersect with axis xm, ym = c.get_xdata(), c.get_ydata() xm_at_ylim = extrap(ylim, ym, xm)[pos] if xm_at_ylim >= xlim[0] and xm_at_ylim <= xlim[1] and self.contains(xm_at_ylim, ylim[pos]): m_, p_ = self.proj.invert(xm_at_ylim, ylim[pos]) dxy = self.proj.gradient(m_, p_, direction="meridian") dxy /= np.sqrt((dxy**2).sum()) dxy *= pad / dxy[1] # same pad from frame if loc == "bottom": dxy *= -1 angle = 0 # no option along the frame x_im = (xm_at_ylim - xlim[0])/(xlim[1]-xlim[0]) y_im = (ylim[pos] - ylim[0])/(ylim[1]-ylim[0]) # these are set as annotations instead of simple axis ticks # because those cannot be shifted by a constant point amount to # follow the graticule self.ax.annotate(self._config['grid']['meridian_fmt'](m), (x_im, y_im), xycoords='axes fraction', xytext=dxy, textcoords='offset points', annotation_clip=False, gid='frame-meridian-label', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, **kwargs) xticks.append(x_im) if description is not None: # find gap in middle of axis xticks.insert(0, 0) xticks.append(1) xticks = np.array(xticks) gaps = (xticks[1:] + xticks[:-1]) / 2 center_gap = np.argmin(np.abs(gaps - 0.5)) x_im = gaps[center_gap] y_im = (ylim[pos] - ylim[0])/(ylim[1]-ylim[0]) dxy = [0, pad] if loc == "bottom": dxy[1] *= -1 self.ax.annotate(description, (x_im, y_im), xycoords='axes fraction', xytext=dxy, textcoords='offset points', annotation_clip=False, gid='frame-meridian-label-description', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, **kwargs) def labelParallelsAtFrame(self, loc='left', parallels=None, pad=None, description=None, **kwargs): """Label the parallels on rectangular frame of the map If the view only shows a fraction of the map, a segment or an entire rectangular frame is shown and the graticule labels are moved to outside that frame. This method is implicitly called, but can be used to overwrite the defaults. Args: loc: location of the label with respect to frame, from `['left', 'right']` parallels: list of parallels to label, if None labels all of them pad: padding of annotation, in units of fontsize description: equivalent to `matplotlib` axis label **kwargs: styling of `matplotlib` annotations for the graticule labels """ assert loc in ['left', 'right'] arguments = _parseArgs(locals()) myname = 'labelParallelsAtFrame' # need extra space for tight_layout to consider the frame annnotations # we can't get the actual width, but we can make use the of the default width of the axes tick labels if myname not in self._config.keys() or self._config[myname]['loc'] != loc: self.ax.yaxis.set_ticks_position(loc) self.ax.yaxis.set_label_position(loc) # remove existing frame_artists = self.artists('frame-parallel-label') for artist in frame_artists: artist.remove() self.fig.tight_layout(pad=0.75) self._config[myname] = arguments size = kwargs.pop('size', matplotlib.rcParams['font.size']) # styling consistent with frame, i.e. with edge color = kwargs.pop('color', self._edge.get_edgecolor()) alpha = kwargs.pop('alpha', self._edge.get_alpha()) zorder = kwargs.pop('zorder', self._edge.get_zorder() + 1) # on top of edge verticalalignment = kwargs.pop('verticalalignment', 'center') horizontalalignment = self._negateLoc(loc) # no option along the frame _ = kwargs.pop('horizontalalignment', None) if pad is None: pad = size / 3 if parallels is None: parallels = self.parallels # check if loc has frame poss = {"left": 0, "right": 1} pos = poss[loc] xlim, ylim = self.ax.get_xlim(), self.ax.get_ylim() yticks = [] frame_artists = self.artists(r'frame-([a-zA-Z]+)', regex=True) frame_locs = [match.group(1) for c,match in frame_artists] if loc in frame_locs: # find all parallel grid lines m_artists = self.artists(r'grid-parallel-([\-\+0-9.]+)', regex=True) for c,match in m_artists: p = float(match.group(1)) if p in parallels: # intersect with axis xp, yp = c.get_xdata(), c.get_ydata() yp_at_xlim = extrap(xlim, xp, yp)[pos] if yp_at_xlim >= ylim[0] and yp_at_xlim <= ylim[1] and self.contains(xlim[pos], yp_at_xlim): m_, p_ = self.proj.invert(xlim[pos], yp_at_xlim) dxy = self.proj.gradient(m_, p_, direction='parallel') dxy /= np.sqrt((dxy**2).sum()) dxy *= pad / dxy[0] # same pad from frame if loc == "left": dxy *= -1 angle = 0 # no option along the frame x_im = (xlim[pos] - xlim[0])/(xlim[1]-xlim[0]) y_im = (yp_at_xlim - ylim[0])/(ylim[1]-ylim[0]) # these are set as annotations instead of simple axis ticks # because those cannot be shifted by a constant point amount to # follow the graticule self.ax.annotate(self._config['grid']['parallel_fmt'](p), (x_im, y_im), xycoords='axes fraction', xytext=dxy, textcoords='offset points', annotation_clip=False, gid='frame-parallel-label', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, **kwargs) yticks.append(y_im) if description is not None: # find gap in middle of axis yticks.insert(0, 0) yticks.append(1) yticks = np.array(yticks) gaps = (yticks[1:] + yticks[:-1]) / 2 center_gap = np.argmin(np.abs(gaps - 0.5)) y_im = gaps[center_gap] x_im = (xlim[pos] - xlim[0])/(xlim[1]-xlim[0]) dxy = [pad, 0] if loc == "left": dxy[0] *= -1 self.ax.annotate(description, (x_im, y_im), xycoords='axes fraction', xytext=dxy, textcoords='offset points', annotation_clip=False, gid='frame-parallel-label-description', horizontalalignment=horizontalalignment, verticalalignment=verticalalignment, size=size, color=color, alpha=alpha, zorder=zorder, **kwargs) def _setFrame(self): # clean up existing frame frame_artists = self.artists(r'frame-([a-zA-Z]+)', regex=True) for c,m in frame_artists: c.remove() locs = ['left', 'bottom', 'right', 'top'] xlim, ylim = self.ax.get_xlim(), self.ax.get_ylim() # use styling of edge for consistent map borders ls = '-' lw = self._edge.get_linewidth() c = self._edge.get_edgecolor() alpha = self._edge.get_alpha() zorder = self._edge.get_zorder() - 1 # limits imprecise, hide underneath edge precision = 1000 const = np.ones(precision) for loc in locs: # define line along axis if loc == "left": line = xlim[0]*const, np.linspace(ylim[0], ylim[1], precision) if loc == "right": line = xlim[1]*const, np.linspace(ylim[0], ylim[1], precision) if loc == "bottom": line = np.linspace(xlim[0], xlim[1], precision), ylim[0]*const if loc == "top": line = np.linspace(xlim[0], xlim[1], precision), ylim[1]*const # show axis lines only where line is inside of map edge inside = self.contains(*line) if (~inside).all(): continue if inside.all(): startpos, stoppos = 0, -1 xmin = (line[0][startpos] - xlim[0])/(xlim[1]-xlim[0]) ymin = (line[1][startpos] - ylim[0])/(ylim[1]-ylim[0]) xmax = (line[0][stoppos] - xlim[0])/(xlim[1]-xlim[0]) ymax = (line[1][stoppos] - ylim[0])/(ylim[1]-ylim[0]) self.ax.plot([xmin,xmax], [ymin, ymax], c=c, ls=ls, lw=lw, alpha=alpha, zorder=zorder, clip_on=False, transform=self.ax.transAxes, gid='frame-%s' % loc) continue # for piecewise inside: determine limits where it's inside # by checking for jumps in inside inside = inside.astype("int") diff = inside[1:] - inside[:-1] jump = np.flatnonzero(diff) start = 0 if inside[0]: jump = np.concatenate(((0,),jump)) while True: startpos = jump[start] if start+1 < len(jump): stoppos = jump[start + 1] else: stoppos = -1 xmin = (line[0][startpos] - xlim[0])/(xlim[1]-xlim[0]) ymin = (line[1][startpos] - ylim[0])/(ylim[1]-ylim[0]) xmax = (line[0][stoppos] - xlim[0])/(xlim[1]-xlim[0]) ymax = (line[1][stoppos] - ylim[0])/(ylim[1]-ylim[0]) artist = Line2D([xmin,xmax], [ymin, ymax], c=c, ls=ls, lw=lw, alpha=alpha, zorder=zorder, clip_on=False, transform=self.ax.transAxes, gid='frame-%s' % loc) self.ax.add_line(artist) if start + 2 < len(jump): start += 2 else: break def _clearFrame(self): frame_artists = self.artists('frame-') for artist in frame_artists: artist.remove() def _resetFrame(self): self._setFrame() for method in ['labelMeridiansAtFrame', 'labelParallelsAtFrame']: if method in self._config.keys(): getattr(self, method)(**self._config[method]) def _pressHandler(self, evt): if evt.button != 1: return if evt.dblclick: return # remove frame and labels self._clearFrame() self.fig.canvas.draw() def _releaseHandler(self, evt): if evt.button != 1: return if evt.dblclick: return self._resetFrame() self.fig.canvas.draw() def _scrollHandler(self, evt): # mouse scroll for zoom if evt.inaxes != self.ax: return if evt.step == 0: return # remove frame and labels self._clearFrame() # scroll to fixed pointer position: google maps style factor = 0.25 c = 1 - evt.step*factor # scaling factor xlim, ylim = self.ax.get_xlim(), self.ax.get_ylim() xdiff, ydiff = xlim[1] - xlim[0], ylim[1] - ylim[0] x, y = evt.xdata, evt.ydata fx, fy = (x - xlim[0])/xdiff, (y - ylim[0])/ydiff # axis units xlim_, ylim_ = x - fx*c*xdiff, y - fy*c*ydiff xlim__, ylim__ = xlim_ + c*xdiff, ylim_ + c*ydiff self.ax.set_xlim(xlim_, xlim__) self.ax.set_ylim(ylim_, ylim__) self._resetFrame() self.fig.canvas.draw() #### common plot type for maps: follow mpl convention #### def plot(self, ra, dec, *args, **kwargs): """Matplotlib `plot` with `ra/dec` points transformed according to map projection""" x, y = self.proj.transform(ra, dec) return self.ax.plot(x, y, *args, **kwargs) def scatter(self, ra, dec, **kwargs): """Matplotlib `scatter` with `ra/dec` points transformed according to map projection""" x, y = self.proj.transform(ra, dec) return self.ax.scatter(x, y, **kwargs) def hexbin(self, ra, dec, C=None, **kwargs): """Matplotlib `hexbin` with `ra/dec` points transformed according to map projection""" x, y = self.proj.transform(ra, dec) # determine proper gridsize: by default x is only needed, y is chosen accordingly gridsize = kwargs.pop("gridsize", None) mincnt = kwargs.pop("mincnt", 1) clip_path = kwargs.pop('clip_path', self._edge) if gridsize is None: xlim, ylim = (x.min(), x.max()), (y.min(), y.max()) per_sample_volume = (xlim[1]-xlim[0])**2 / x.size * 10 gridsize = int(np.ceil((xlim[1]-xlim[0]) / np.sqrt(per_sample_volume))) # styling: use same default colormap as density for histogram if C is None: cmap = kwargs.pop("cmap", "YlOrRd") else: cmap = kwargs.pop("cmap", None) zorder = kwargs.pop("zorder", 0) # same as for imshow: underneath everything artist = self.ax.hexbin(x, y, C=C, gridsize=gridsize, mincnt=mincnt, cmap=cmap, zorder=zorder, **kwargs) artist.set_clip_path(clip_path) return artist def text(self, ra, dec, s, rotation=None, direction="parallel", **kwargs): """Matplotlib `text` with coordinates given by `ra/dec`. Args: ra: rectascension of text dec: declination of text s: string rotation: if text should be rotated to tangent direction direction: tangent direction, from `['parallel', 'meridian']` **kwargs: styling arguments for `matplotlib.text` """ x, y = self.proj.transform(ra, dec) if rotation is None: dxy_ = self.proj.gradient(ra, dec, direction=direction) angle = 90-np.arctan2(*dxy_) / DEG2RAD else: angle = rotation return self.ax.text(x, y, s, rotation=angle, rotation_mode="anchor", clip_on=True, **kwargs) def colorbar(self, cb_collection, cb_label="", orientation="vertical", size="2%", pad="1%"): """Add colorbar to side of map. The location of the colorbar will be chosen automatically to not interfere with the map frame labels. Args: cb_collection: a `matplotlib` mappable collection cb_label: string for colorbar label orientation: from ["vertical", "horizontal"] size: fraction of ax size to use for colorbar pad: fraction of ax size to use as pad to map frame """ assert orientation in ["vertical", "horizontal"] # pick the side that does not have the tick labels if orientation == "vertical": frame_loc = self._config['labelParallelsAtFrame']['loc'] else: frame_loc = self._config['labelMeridiansAtFrame']['loc'] loc = self._negateLoc(frame_loc) from mpl_toolkits.axes_grid1 import make_axes_locatable divider = make_axes_locatable(self.ax) cax = divider.append_axes(loc, size=size, pad=pad) cb = self.fig.colorbar(cb_collection, cax=cax, orientation=orientation, ticklocation=loc) cb.solids.set_edgecolor("face") cb.set_label(cb_label) self.fig.tight_layout(pad=0.75) return cb def focus(self, ra, dec, pad=0.025): """Focus onto region of map covered by `ra/dec` Adjusts x/y limits to encompass given `ra/dec`. Args: ra: list of rectascensions dec: list of declinations pad: distance to edge of the frame, in units of axis size """ # to replace the autoscale function that cannot zoom in x, y = self.proj.transform(ra, dec) xlim = [x.min(), x.max()] ylim = [y.min(), y.max()] xrange = xlim[1]-xlim[0] yrange = ylim[1]-ylim[0] xlim[0] -= pad * xrange xlim[1] += pad * xrange ylim[0] -= pad * yrange ylim[1] += pad * yrange self.ax.set_xlim(xlim) self.ax.set_ylim(ylim) self._resetFrame() self.fig.canvas.draw() def defocus(self, pad=0.025): """Show entire map. Args: pad: distance to edge of the map, in units of axis size """ # to replace the autoscale function that cannot zoom in xlim, ylim = list(self.xlim()), list(self.ylim()) xrange = xlim[1]-xlim[0] yrange = ylim[1]-ylim[0] xlim[0] -= pad * xrange xlim[1] += pad * xrange ylim[0] -= pad * yrange ylim[1] += pad * yrange self.ax.set_xlim(xlim) self.ax.set_ylim(ylim) self._clearFrame() self.fig.canvas.draw() def show(self, *args, **kwargs): """Show `matplotlib` figure""" self.fig.show(*args, **kwargs) def savefig(self, *args, **kwargs): """Save `matplotlib` figure""" self.fig.savefig(*args, **kwargs) #### special plot types for maps #### def footprint(self, surveyname, **kwargs): """Plot survey footprint polygon onto map Uses `get_footprint()` method of a `skymapper.Survey` derived class instance The name of the survey is indentical to the class name. All available surveys are listed in `skymapper.survey_register`. Args: surveyname: name of the survey, must be in keys of `skymapper.survey_register` **kwargs: styling of `matplotlib.collections.PolyCollection` """ # search for survey in register ra, dec = survey_register[surveyname].getFootprint() x,y = self.proj.transform(ra, dec) poly = Polygon(np.dstack((x,y))[0], closed=True, **kwargs) self.ax.add_patch(poly) return poly def vertex(self, vertices, color=None, vmin=None, vmax=None, **kwargs): """Plot polygons (e.g. Healpix vertices) Args: vertices: cell boundaries in RA/Dec, from getCountAtLocations() color: string or matplib color, or numeric array to set polygon colors vmin: if color is numeric array, use vmin to set color of minimum vmax: if color is numeric array, use vmin to set color of minimum **kwargs: matplotlib.collections.PolyCollection keywords Returns: matplotlib.collections.PolyCollection """ vertices_ = np.empty_like(vertices) vertices_[:,:,0], vertices_[:,:,1] = self.proj.transform(vertices[:,:,0], vertices[:,:,1]) from matplotlib.collections import PolyCollection zorder = kwargs.pop("zorder", 0) # same as for imshow: underneath everything clip_path = kwargs.pop('clip_path', self._edge) coll = PolyCollection(vertices_, array=color, zorder=zorder, clip_path=clip_path, **kwargs) coll.set_clim(vmin=vmin, vmax=vmax) coll.set_edgecolor("face") self.ax.add_collection(coll) return coll def healpix(self, m, nest=False, color_percentiles=[10,90], **kwargs): """Plot HealPix map Args: m: Healpix map array nest: HealPix nest color_percentiles: lower and higher cutoff percentile for map coloring """ # determine ra, dec of map; restrict to non-empty cells pixels = np.flatnonzero(m) nside = healpix.hp.npix2nside(m.size) vertices = healpix.getHealpixVertices(pixels, nside, nest=nest) color = m[pixels] # styling cmap = kwargs.pop("cmap", "YlOrRd") vmin = kwargs.pop("vmin", None) vmax = kwargs.pop("vmax", None) zorder = kwargs.pop("zorder", 0) # same as for imshow: underneath everything if vmin is None or vmax is None: vlim = np.percentile(color, color_percentiles) if vmin is None: vmin = vlim[0] if vmax is None: vmax = vlim[1] # make a map of the vertices return self.vertex(vertices, color=color, vmin=vmin, vmax=vmax, cmap=cmap, zorder=zorder, **kwargs) def density(self, ra, dec, nside=1024, color_percentiles=[10,90], **kwargs): """Plot sample density using healpix binning Args: ra: list of rectascensions dec: list of declinations nside: HealPix nside color_percentiles: lower and higher cutoff percentile for map coloring """ # get count in healpix cells, restrict to non-empty cells bc, _, _, vertices = healpix.getCountAtLocations(ra, dec, nside=nside, return_vertices=True) color = bc # styling cmap = kwargs.pop("cmap", "YlOrRd") vmin = kwargs.pop("vmin", None) vmax = kwargs.pop("vmax", None) zorder = kwargs.pop("zorder", 0) # same as for imshow: underneath everything if vmin is None or vmax is None: vlim = np.percentile(color, color_percentiles) if vmin is None: vmin = vlim[0] if vmax is None: vmax = vlim[1] # make a map of the vertices return self.vertex(vertices, color=color, vmin=vmin, vmax=vmax, cmap=cmap, zorder=zorder, **kwargs) def interpolate(self, ra, dec, value, method='cubic', fill_value=np.nan, **kwargs): """Interpolate ra,dec samples over covered region in the map Requires scipy, uses `scipy.interpolate.griddata` with `method='cubic'`. Args: ra: list of rectascensions dec: list of declinations value: list of sample values **kwargs: arguments for matplotlib.imshow """ x, y = self.proj.transform(ra, dec) # evaluate interpolator over the range covered by data xlim, ylim = (x.min(), x.max()), (y.min(), y.max()) per_sample_volume = min(xlim[1]-xlim[0], ylim[1]-ylim[0])**2 / x.size dx = np.sqrt(per_sample_volume) xline = np.arange(xlim[0], xlim[1], dx) yline = np.arange(ylim[0], ylim[1], dx) xp, yp = np.meshgrid(xline, yline) + dx/2 # evaluate center pixel vp = scipy.interpolate.griddata(np.dstack((x,y))[0], value, (xp,yp), method=method, fill_value=fill_value) # remember axes limits ... xlim_, ylim_ = self.ax.get_xlim(), self.ax.get_ylim() _ = kwargs.pop('extend', None) zorder = kwargs.pop("zorder", 0) # default for imshow: underneath everything clip_path = kwargs.pop('clip_path', self._edge) artist = self.ax.imshow(vp, extent=(xlim[0], xlim[1], ylim[0], ylim[1]), zorder=zorder, **kwargs) artist.set_clip_path(clip_path) # ... because imshow focusses on extent self.ax.set_xlim(xlim_) self.ax.set_ylim(ylim_) return artist def extrapolate(self, ra, dec, value, resolution=100, clean_edge=True, **kwargs): """Extrapolate ra,dec samples on the entire sphere and project on the map Requires scipy, uses default `scipy.interpolate.Rbf`. Args: ra: list of rectascensions dec: list of declinations value: list of sample values resolution: number of evaluated cells per linear map dimension clean_edge: use another interpolation to generate clean edge of the map **kwargs: arguments for matplotlib.imshow """ # interpolate samples in RA/DEC rbfi = scipy.interpolate.Rbf(ra, dec, value, norm=skyDistance) # make grid in x/y over the limits of the map or the clip_path clip_path = kwargs.pop('clip_path', self._edge) if clip_path is None: xlim, ylim = self.xlim(), self.ylim() else: xlim = clip_path.xy[:, 0].min(), clip_path.xy[:, 0].max() ylim = clip_path.xy[:, 1].min(), clip_path.xy[:, 1].max() if resolution % 1 == 0: resolution += 1 dx = (xlim[1]-xlim[0])/resolution xline = np.arange(xlim[0], xlim[1], dx) yline = np.arange(ylim[0], ylim[1], dx) xp, yp = np.meshgrid(xline, yline) + dx/2 # evaluate center pixel inside = self.contains(xp,yp) vp = np.ma.array(np.empty(xp.shape), mask=~inside) rap, decp = self.proj.invert(xp[inside], yp[inside]) vp[inside] = rbfi(rap, decp) # construct another rbf in pixel space to populate the values # outside of the map region, ordinary Euclidean distance now if clean_edge: rbfi = scipy.interpolate.Rbf(xp[inside], yp[inside], vp[inside]) vp[~inside] = rbfi(xp[~inside], yp[~inside]) zorder = kwargs.pop("zorder", 0) # same as for imshow: underneath everything xlim_, ylim_ = self.ax.get_xlim(), self.ax.get_ylim() artist = self.ax.imshow(vp, extent=(xlim[0], xlim[1], ylim[0], ylim[1]), zorder=zorder, **kwargs) artist.set_clip_path(clip_path) # ... because imshow focusses on extent self.ax.set_xlim(xlim_) self.ax.set_ylim(ylim_) return artist