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 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 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
