def plotB_subplot(p2D,
th2D,
B,
ax,
planet="earth",
levels=60,
cmap='RdBu_r',
proj='Mollweide'):
planet = planet.lower()
bmax = np.abs(B).max()
digits = int(np.log10(bmax)) + 1
if digits > 1:
bmax = np.round(bmax)
else:
bmax = np.round(bmax, decimals=1)
p2D -= np.pi
th2D -= np.pi / 2
th2D = -th2D
cs = np.linspace(-bmax, bmax, levels)
divnorm = colors.TwoSlopeNorm(vmin=-bmax, vcenter=0, vmax=bmax)
try:
import cartopy.crs as ccrs
except:
print("cartopy library not available, using Hammer projection")
proj = 'hammer'
if proj.lower() == 'hammer':
xx, yy = hammer2cart(th2D, p2D)
cont = ax.contourf(xx,
yy,
B,
cs,
cmap=cmap,
norm=divnorm,
extend='both')
else:
if planet == "earth":
ax.coastlines()
cont = ax.contourf(p2D*180/np.pi,th2D*180/np.pi,B,cs, \
transform=ccrs.PlateCarree(),cmap=cmap,norm=divnorm,extend='both')
cbar = plt.colorbar(cont,
orientation='horizontal',
fraction=0.06,
pad=0.04,
ticks=[-bmax, 0, bmax])
cbar.ax.tick_params(labelsize=15)
ax.set_title(planet.capitalize(), fontsize=20)
ax.axis('equal')
ax.axis('off')
def animate3D(u,v,h,H,dt):
print('Saving 3D figures...')
# specify max and min contours
minh, maxh = np.amin(h),np.amax(h)
norm_h = colors.TwoSlopeNorm(vmin=minh,
vcenter=H,
vmax=maxh)
# colormap
cmap = 'cmo.deep'
# add light for shadings
ls = LightSource(azdeg=315, altdeg=45)
# loop to animate
for t in range(round(tmax/frecplot)+1):
# fig params
fig = plt.figure(figsize=(20, 10))
# 3D plot
ax = fig.add_subplot(1, 2, 1, projection='3d')
ax.set_zlim(minh, maxh)
# view zimuth and angle
ax.view_init(50, 260)
# labels in km
ax.set_xlabel('X (km)', fontweight='bold')
ax.set_ylabel('Y (km)', fontweight='bold')
ax.set_zlabel('Height (m)', fontweight='bold')
# plot surface
z = h[t]
rgb = ls.shade(z, cmap=cmo.deep, vert_exag=0.1, blend_mode='soft')
ax.plot_surface(lonsh/1e3,latsh/1e3, z,
rstride=1, cstride=1, facecolors=rgb,
cmap = cmap,
linewidth=0,
antialiased=False,
shade=False,
norm=norm_h)
ax.set_title('Surface elevation h(x,y,t) after t ='\
+str(round(t*dt*frecplot/60/60/24,2))+'d, '+ \
str(round(t*dt/60/60,2))+'h')
# plot vectors
ax = fig.add_subplot(1, 2, 2)
u_mean = (u[t][:,:-1]+u[t][:,1:])/2
v_mean = (v[t][:-1]+v[t][1:])/2
ax.quiver(lonsh[::5]/1e3,latsh[::5]/1e3, u_mean[::5],v_mean[::5],scale=0.5)
ax.set_title('Velocity field u(x,y,t) after t ='\
+str(round(t*dt*frecplot/60/60/24,2))+'d, '+ \
str(round(t*dt/60/60,2))+'h')
print('animating timestep = '+str(t))
pl.savefig('./3Dfigs/'+str(t)+'.png')
def normRange(delta, mmin=-.5, mmax=.5):
"""
"""
# vmin, vcent, vmax = max(delta.min(), mmin), 0., min(delta.max(), mmax)
# if not (vmin < vcent and vcent < vmax):
# vmin, vmax = delta.min(), delta.max()
# vcent = .5 * (vmin + vmax)
vmin, vcent, vmax = mmin, 0., mmax
norm = mcolors.TwoSlopeNorm(vcenter=vcent, vmin=vmin, vmax=vmax)
return norm
def test_twoslope_colorbar():
# Note that the second tick = 20, and should be in the middle
# of the colorbar (white)
# There should be no tick right at the bottom, nor at the top.
fig, ax = plt.subplots()
norm = mcolors.TwoSlopeNorm(20, 5, 95)
pc = ax.pcolormesh(np.arange(1, 11), np.arange(1, 11),
np.arange(100).reshape(10, 10),
norm=norm, cmap='RdBu_r')
fig.colorbar(pc)
def test_twoslope_colorbar():
# Note that the first tick = 20, and should be in the middle
# of the colorbar (white)
fig, ax = plt.subplots()
norm = mcolors.TwoSlopeNorm(20, 0, 100)
pc = ax.pcolormesh(np.arange(1, 11),
np.arange(1, 11),
np.arange(100).reshape(10, 10),
norm=norm,
cmap='RdBu_r')
fig.colorbar(pc)
def radContour(theta, phi, data, idxPlot, grid, levels, cm, proj):
data = data[..., idxPlot]
p2D, th2D = get_grid2D(theta, phi)
p2D = p2D - np.pi
th2D= np.pi/2 - th2D
lon = p2D * 180./np.pi
lat = th2D * 180./np.pi
if proj == "ortho":
fig = plt.figure(figsize=(10, 10))
plotcrs = ccrs.Orthographic(0, 40)
elif proj == "moll":
fig = plt.figure(figsize=(12, 10))
plotcrs = ccrs.Mollweide()
ax = fig.add_subplot(1, 1, 1, projection=plotcrs)
datMax = (np.abs(data)).max()
divnorm = colors.TwoSlopeNorm(vmin=-datMax, vcenter=0, vmax=datMax)
if grid:
ax.gridlines(linewidth=1, color='gray', alpha=0.5, linestyle=':')
if proj == "ortho":
cont = ax.pcolormesh(lon, lat, data, \
transform=ccrs.PlateCarree(), cmap=cm, \
norm=divnorm)
elif proj == "moll":
cont = ax.contourf(lon, lat, data, levels, \
transform=ccrs.PlateCarree(), cmap=cm, \
norm=divnorm)
for c in cont.collections:
c.set_edgecolor("face")
# if vec:
# ut = self.U.Us*cos(self.grid.th3D) - self.U.Uz*sin(self.grid.th3D)
# ut = ut[::vecStride,::vecStride,idxPlot]
# up = self.U.Up[::vecStride,::vecStride,idxPlot]
# lon = lon[::vecStride,::vecStride]
# lat = lat[::vecStride,::vecStride]
# #ax.quiver(lon,lat,up,ut,transform=plotcrs,width=vecWidth,scale=vecScale,regrid_shape=20)
# ax.quiver(lon,lat,up,ut,transform=plotcrs)#,regrid_shape=regrid_shape)
plt.axis('equal')
plt.axis('off')
plt.tight_layout()
def reset(self):
# Colorbar reset
self.norm = matcolors.TwoSlopeNorm(vmin=self.vmin0,
vcenter=0.0,
vmax=self.vmax0)
self.cmap = plt.get_cmap(self.cmapname)
self.cmap.set_bad(color=self.nancolor)
self.mappable = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)
# Update slices and lines
self.update_xslice(self.xl0)
self.update_yslice(self.yl0)
self.update_zslice(self.zl0)
def plotSurf(p2D, th2D, B, levels=60, cmap='RdBu_r', proj='Mollweide'):
bmax = np.abs(B).max()
digits = int(np.log10(bmax)) + 1
if digits > 1:
bmax = np.round(bmax)
else:
bmax = np.round(bmax, decimals=1)
divnorm = colors.TwoSlopeNorm(vmin=-bmax, vcenter=0, vmax=bmax)
cs = np.linspace(-bmax, bmax, levels)
lon2D = p2D - np.pi
lat2D = np.pi / 2 - th2D
try:
import cartopy.crs as ccrs
except:
print("cartopy library not available, using Hammer projection")
proj = 'hammer'
if proj.lower() == 'hammer':
ax = plt.axes()
xx, yy = hammer2cart(lat2D, lon2D)
cont = ax.contourf(xx,
yy,
B,
cs,
cmap=cmap,
norm=divnorm,
extend='both')
else:
projection = eval('ccrs.' + proj + '()')
ax = plt.axes(projection=projection)
cont = ax.contourf(lon2D*180/np.pi,lat2D*180/np.pi,B,cs, \
transform=ccrs.PlateCarree(),cmap=cmap,norm=divnorm,extend='both')
cbar = plt.colorbar(cont,
orientation='horizontal',
fraction=0.06,
pad=0.04,
ticks=[-bmax, 0, bmax])
ax.axis('equal')
ax.axis('off')
return ax, cbar
def update_cmap(self, vmin, vcenter, vmax):
# Colormap settings
self.vmin, self.vmax = [vmin, vmax]
self.norm = matcolors.TwoSlopeNorm(vmin=vmin,
vcenter=vcenter,
vmax=vmax)
self.cmap = plt.cm.coolwarm
self.mappable = cm.ScalarMappable(norm=self.norm, cmap=self.cmap)
self.slice['x'].norm = self.norm
self.slice['x'].set_data(self.V[self.xli, :, :].T)
self.slice['y'].norm = self.norm
self.slice['y'].set_data(self.V[:, self.yli, :].T)
self.slice['z'].norm = self.norm
self.slice['z'].set_data(self.V[:, :, self.zli].T)
self.fig.canvas.draw_idle()
def plot_stream_diff_prob(stream_prob_mat, event_list, *img_name):
mid_norm = clr.TwoSlopeNorm(vcenter=0.)
fig = plt.figure()
fig = plt.figure(figsize=(20, 10))
ax = fig.add_subplot(111)
color = 'bwr_r'
ax.matshow(stream_prob_mat, cmap=color)
fig.colorbar(ax.matshow(stream_prob_mat, cmap=color, norm=mid_norm))
ax.set_yticks(range(len(event_list) + 1))
ax.set_yticklabels(event_list + [''], fontsize=15)
ax.set_xticks(range(len(event_list) + 1))
plt.xticks(rotation=90, ha='right')
ax.set_xticklabels([''] + event_list, fontsize=15)
if type(img_name) == str:
fig.savefig(img_name + '.pdf', bbox_inches="tight")
def scatter(robustness_array, delta_pos_array, delta_vel_array, filename):
fig, ax = plt.subplots(figsize=(10, 5))
customnorm = mcolors.TwoSlopeNorm(0)
sp = ax.scatter(delta_vel_array,
delta_pos_array,
c=robustness_array,
cmap='RdYlGn',
norm=customnorm)
ax.set(xlabel='$\Delta$v between leader and follower ($m/s$)',
ylabel='Distance ($m$)')
cb = fig.colorbar(sp)
cb.ax.set_xlabel('$\\rho$')
fig.suptitle('Initial conditions vs robustness $\\rho$')
fig.savefig(os.path.join(EXP + relpath, filename), dpi=150)
def scatter(robustness_array, cart_pos_array, pole_ang_array, cart_vel_array, pole_ang_vel_array, filename):
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
print(cart_pos_array, "\n", pole_ang_array, "\n", cart_vel_array, "\n", pole_ang_vel_array)
customnorm = mcolors.TwoSlopeNorm(0)
sp0 = ax[0].scatter(cart_pos_array, cart_vel_array, c=robustness_array, cmap='RdYlGn', norm=customnorm)
ax[0].set(xlabel='cart position', ylabel='cart velocity')
# cb = fig.colorbar(sp0)
# cb.ax[0].set_label('$\\rho$')
sp1 = ax[1].scatter(pole_ang_array, pole_ang_vel_array, c=robustness_array, cmap='RdYlGn', norm=customnorm)
ax[1].set(xlabel='pole angle', ylabel='pole angular frequency')
# cb = fig.colorbar(sp1)
# cb.ax[1].set_label('$\\rho$')
fig.suptitle('Initial conditions vs robustness $\\rho$')
fig.savefig(os.path.join(EXP+relpath, filename), dpi=150)
def show_normalisation(df, truth, pred_raw, pred_norm, pred_match, truth_smoothed=None, y_log=False,
y_lbl='# daily fatalities', x_pad=3, cb_center=0, figsize=(15,8)):
'''
Utility function to visualise the effect of the normalisation procedure on a model predictions, where the
`pred_match` values are to be suitably computed depending on the normalising procedure applied.
'''
# find when truth was available
idx = df.index[~df[truth].isnull()]
# plot time series
fig, ax = plt.subplots(figsize=figsize)
ax.scatter(idx, df.loc[idx, truth], marker='x', color='C0', label='reported')
ax.plot(idx, df.loc[idx, pred_raw], color='k', alpha=0.3, label='predicted (raw)')
ax.plot(idx, df.loc[idx, pred_norm], color='k', alpha=0.8, linestyle='--', label='predicted (normalised)')
if truth_smoothed is not None:
ax.plot(idx, df.loc[idx, truth_smoothed],color='C0', alpha=0.4, linestyle='--',
label='reported (smoothed)')
if y_log:
ax.set_yscale('symlog', linthreshy=10)
ax.yaxis.set_major_formatter(mticker.ScalarFormatter())
ax.legend()
# add background colouring
_abs_max = max(np.abs(df.loc[idx, pred_match].values))
divnorm = mcolors.TwoSlopeNorm(vmin=-_abs_max, vcenter=cb_center, vmax=_abs_max)
pcf = ax.pcolorfast(mdates.date2num([idx[0], idx[-1]]), ax.get_ylim(),
df.loc[idx, pred_match].values[np.newaxis],
cmap='bwr', alpha=0.2, norm=divnorm)
t_delta = pd.Timedelta(x_pad, unit='days')
ax.set_xlim(idx[0]-t_delta, idx[-1]+t_delta)
ax.set_ylabel(y_lbl)
# configure colorbar
cbar = fig.colorbar(pcf)
cbar.ax.tick_params(size=0)
cbar.set_ticks([-_abs_max, cb_center, _abs_max])
cbar.set_ticklabels(['too pessimistic\n(reality better)','matches reality','too optimistic\n(reality worse)'])
# configure time axis
ax.xaxis.set_minor_locator(mdates.DayLocator(range(1,32)))
ax.xaxis.set_major_locator(mdates.WeekdayLocator(mdates.MO))
ax.xaxis.set_major_formatter(mdates.DateFormatter("%d %b"))
return fig
def plot_3panel(x, y, model):
plt.figure(figsize=[19.2, 4.85])
# for plotting a colorbar
vmn = np.ma.min([x.data, y.data])
vmx = np.ma.max([x.data, y.data])
if vmx < 0:
vmx = 0.001
if vmn > 0:
vmn = -0.001
divnorm = mcolors.TwoSlopeNorm(vmin=vmn, vcenter=0, vmax=vmx)
# x
plt.subplot(1, 3, 1)
pmesh = iplt.pcolormesh(x, norm=divnorm, cmap='RdBu')
plot_formatting(plt.gca())
plt.title('Mean pr')
# y
plt.subplot(1, 3, 2)
iplt.pcolormesh(y, norm=divnorm, cmap='RdBu')
plot_formatting(plt.gca())
plt.title('Rx1Day')
plt.subplot(1, 3, 3)
plt.scatter(x.data, y.data, s=1)
guide_lines(plt.gca())
plt.xlabel('Mean pr')
plt.ylabel('Rx1day')
# compute simple linear regression
fit = linregress(x.data.flatten(), y.data.flatten())
plt.title(f"R: {fit.rvalue:.2f}. Slope: {fit.slope:.2f}")
plt.suptitle(f"{model} %")
# add colorbar for map plots
fig = plt.gcf()
fig.subplots_adjust(bottom=0.25)
cbax = fig.add_axes([0.125, 0.1, 0.5, 0.075])
fig.colorbar(pmesh, cax=cbax, orientation="horizontal")
# save plot
plt.savefig(f"{PLOT_PATH}{model}.png")
plt.close()
def plot(self, ax):
# Get global u/v wind data
ug, vg = self.chart.get_data(self.gfs_vars, apply_domain=False)
# Get level coordinates for U/V wind components
ug_lv_coord = gfs_utils.get_level_coord(ug, self.level_units)
vg_lv_coord = gfs_utils.get_level_coord(vg, self.level_units)
# Constrain to specified level(s)
ugc = gfs_utils.get_coord_constraint(ug_lv_coord.name(), self.level)
ug = ug.extract(ugc)
vgc = gfs_utils.get_coord_constraint(vg_lv_coord.name(), self.level)
vg = vg.extract(vgc)
# Apply smoothing
ug_data = wrf_smooth2d(ug.data, 6)
ug.data = ug_data.data
vg_data = wrf_smooth2d(vg.data, 6)
vg.data = vg_data.data
# Set up windspharm wind vector object
w = VectorWind(ug, vg)
# Get divergence
div = w.divergence()
# Constrain to specified domain
domain_constraint = gfs_utils.get_domain_constraint(self.chart.domain)
div = div.extract(domain_constraint)
# Mask absolute values below threshold
div_masked = np.ma.masked_where(
np.abs(div.data) < self.thres, div.data)
# Set norm to match centre of colour scale to zero value
self.options['norm'] = mcolors.TwoSlopeNorm(vmin=np.min(div_masked),
vcenter=0,
vmax=np.max(div_masked))
ax.contourf(self.lon, self.lat, div_masked, **self.options)
def anim_forcing(H,lons,lats,dt):
# Pertubation (2D Gaussian in space):
# a: amplitude of the pertubation
# cx: displacement of the peak on x dims
# cy: displacement of the peak on y dims
# nrx: width of the pertubation on x dims
# nry: width of the pertubation on y dis
# dx: dims spacing
nx = 81 # (number of grid points on x)
ny = 81 # (number of grid points on y)
delta = 100e3 # grid spacing (m)
xmax,xmin = 4e6,-4e6
ymax,ymin = 4e6,-4e6
a,cx,cy,nrx,nry,dx = 1,40,41,11,5,delta
xgauss = np.linspace(xmin, xmax,nx+1)
ygauss = np.linspace(ymax, ymin,ny)
longa,latga = np.meshgrid(xgauss,ygauss)
gauss_space = a * np.exp(- (((longa-cx)**2)/(nrx*dx)**2) -\
(((latga-cy)**2)/(nry*dx)**2))
normf = colors.TwoSlopeNorm(vmin=0,
vcenter=0.05,
vmax=0.1)
alpha = 0.01
for t in range(120):
fig, axs = plt.subplots(figsize=(15, 10))
# # Time component of the pertubation
gauss_time = (H/2)*(alpha**3)*(t**2)*math.exp(-alpha*t)
forcing = (gauss_space * gauss_time)
# plot
p1 = axs.pcolormesh(lons,lats,forcing[:,:-1],cmap = 'Reds')
axs.text(-4000000,4000000,'t = '\
+str(round(t*dt/60/60,2))+'h',fontsize=18)
axs.grid(True, color="grey", linestyle='--',lw=2)
divider = make_axes_locatable(axs)
cax = divider.append_axes('right', size='5%', pad=0.05)
plt.colorbar(p1,cax = cax)
print('animating timestep = '+str(t))
pl.savefig('./forcing/'+str(t)+'.png')
def local_laplacian_zscore(venue, results):
"""Plot the local laplacian zscore.
Plot the grid showing the z-score for the local laplacian of the
vertically-averaged head as color-coded markers. The colors are determined
by the z-score and not the quantitative magnitude. However, the colormap
places 0 in the middle. Thus, blue values (negative) indicate local net
infiltration, red values (positive) indicate local net exfiltration.
"""
xtarget = np.array([row[0][0] for row in results])
ytarget = np.array([row[0][1] for row in results])
bdry = venue.boundary()
score = np.empty(xtarget.shape)
for i, row in enumerate(results):
evp = row[2]
varp = row[3]
laplacian = 2 * (evp[0] + evp[1])
stdev = 2 * np.sqrt(varp[0, 0] + varp[1, 1] + 2 * varp[0, 1])
score[i] = min(max(laplacian / stdev, -3), 3)
plt.figure(figsize=FIGSIZE)
plt.axis("equal")
plt.fill(bdry[:, 0], bdry[:, 1], "0.80")
divnorm = colors.TwoSlopeNorm(vmin=-3, vcenter=0, vmax=3)
plt.scatter(xtarget, ytarget, c=score, norm=divnorm, zorder=10, cmap="bwr")
cbar = plt.colorbar()
cbar.ax.set_title("z-score")
plt.xlabel("Easting [m]")
plt.ylabel("Northing [m]")
plt.title(
venue.fullname() + "\n" + "Vertically-averaged Laplacian Z-score",
{"fontsize": 18})
plt.grid(True)
return (xtarget, ytarget, score)
def eqContour(r, phi, data, levels, cm):
phi2D, r2D = np.meshgrid(phi, r, indexing='ij')
xx = r2D * np.cos(phi2D)
yy = r2D * np.sin(phi2D)
plt.figure(figsize=(10, 10))
datMax = (np.abs(data)).max()
divnorm = colors.TwoSlopeNorm(vmin=-datMax, vcenter=0, vmax=datMax)
cont = plt.contourf(xx, yy, data, levels, cmap=cm, norm=divnorm)
plt.plot(r[0]*np.cos(phi), r[0]*np.sin(phi), 'k', lw=1)
plt.plot(r[-1]*np.cos(phi), r[-1]*np.sin(phi), 'k', lw=1)
for c in cont.collections:
c.set_edgecolor("face")
plt.axis('equal')
plt.axis('off')
plt.tight_layout()
def _draw_rt(self, ax, df, draw_key_dates):
rt = df.Rt
cmap = _colors.LinearSegmentedColormap.from_list(
'test', [(0, 'darkGreen'), (0.2, 'forestgreen'),
(0.4, 'yellowgreen'), (0.5, 'gold'), (0.7, 'orange'),
(0.9, 'orangered'), (1, 'crimson')])
norm = _colors.TwoSlopeNorm(1, 0.6, 1.4)
colors = cmap(norm(rt))
line_color = _np.array(cmap(norm(df.Rt[-7:].mean())))
ax.bar(rt.index, rt.values, color=colors, alpha=0.3)
ax.plot(rt.index, rt.rolling(window=3).mean(), color=line_color)
ax.set_ylim(top=2)
self._setup_axes_for_russian_regions_stat(
ax, r"Коэффициент распространения ($R_t$)", False, False,
draw_key_dates)
ax.set_ylim(bottom=0.6)
ax.axhline(1, color='Red', linestyle='dashed', alpha=0.6)
def CB_plotJ(df: pd.DataFrame):
"""Function that plots Counts/Baseline as a 3D plot with detector locations. Counts are
shown by size of point, and C/B is shown using a colourmap
Args:
Dataframe with Time, Count and Baseline columns
"""
df_format = df
forecast_t_min = df_format["measurement_start_utc"].min()
forecast_t_max = df_format["measurement_end_utc"].max()
df_format["C/B"] = df_format["count"] / df_format["baseline"]
df_format["hour_from_start"] = (
df_format["measurement_end_utc"] - df_format["measurement_end_utc"].min()
)
df_format["hour_from_start"] = df_format["hour_from_start"].astype(
dtype="timedelta64[h]"
)
offset = colors.TwoSlopeNorm(vmin=0.5, vcenter=1, vmax=2)
fig = plt.figure(figsize=(20, 10))
ax = fig.add_subplot(111, projection="3d")
p = ax.scatter(
df_format["lon"],
df_format["lat"],
df_format["hour_from_start"],
c=df_format["C/B"],
s=df_format["count"] * 0.1,
norm=offset,
cmap="coolwarm",
)
ax.set_xlabel("lon")
ax.set_ylabel("lat")
ax.set_zlabel("Hours")
fig.colorbar(p)
fig.suptitle("Forecast from {} to {}".format(forecast_t_min, forecast_t_max))
plt.show()
def get_cmap_norm(vmin: float, vmax: float) -> mcolors:
"""
Get the cmap norm key based on limits.
If the boundary limits are positive and negative, returns a 0 divergent
norm.
Args:
vmin (float): Data minimum
vmax (float): Daata Maximum
Returns:
mcolors: norm.
"""
if vmin * vmax >= 0:
cNorm = mcolors.Normalize(vmin=vmin, vmax=vmax)
else:
color_max = np.max((np.abs(vmax), np.abs(vmin)))
cNorm = mcolors.TwoSlopeNorm(vmin=-color_max,
vmax=color_max,
vcenter=0)
return cNorm
topo = dem['topo']
longitude = dem['longitude']
latitude = dem['latitude']
fig, ax = plt.subplots()
# make a colormap that has land and ocean clearly delineated and of the
# same length (256 + 256)
colors_undersea = plt.cm.terrain(np.linspace(0, 0.17, 256))
colors_land = plt.cm.terrain(np.linspace(0.25, 1, 256))
all_colors = np.vstack((colors_undersea, colors_land))
terrain_map = colors.LinearSegmentedColormap.from_list('terrain_map',
all_colors)
# make the norm: Note the center is offset so that the land has more
# dynamic range:
divnorm = colors.TwoSlopeNorm(vmin=-500., vcenter=0, vmax=4000)
pcm = ax.pcolormesh(longitude,
latitude,
topo,
rasterized=True,
norm=divnorm,
cmap=terrain_map,
shading='auto')
# Simple geographic plot, set aspect ratio beecause distance between lines of
# longitude depends on latitude.
ax.set_aspect(1 / np.cos(np.deg2rad(49)))
ax.set_title('TwoSlopeNorm(x)')
fig.colorbar(pcm, shrink=0.6)
plt.show()
# Q_value = data["Q_value"]
R_value = data["R_value"]
grids = data["grids"]
S_M = data["S_M"]
XV = S_M > 0.0
# *
RX_value = vibly.project_Q2S(R_value, grids, proj_opt=np.mean)
# * limit of reward
r_max = RX_value.max()
r_min = RX_value.min()
# * limit of values
alpha = 0.
# mynorm = colors.TwoSlopeNorm(vmin=0., vcenter=0., vmax=2.5)
mynorm = colors.TwoSlopeNorm(vmin=-0.5, vcenter=alpha, vmax=2.5)
mymap = vplot.get_vmap(2)
extent = [
grids['states'][1][0], grids['states'][1][-1], grids['states'][0][0],
grids['states'][0][-1]
]
fig, axs = plt.subplots(1, figsize=(6, 5))
X_value = vibly.project_Q2S(data["Q_value"], grids, proj_opt=np.max)
# lvls = [0., 0.15]
print(X_value.max())
pc1 = vplot.value_function(axs,
X_value,
grids,
def GridStrain(pos, disp, k, par, plotpar, plotst):
'''
GridStrain computes the infinitesimal strain of a network
of stations with displacements in x (east) and y (north).
Strain in z is assumed to be zero (plane strain)
USE: cent,eps,ome,pstrain,rotc = GridStrain(pos,disp,k,par,plotpar,plotst)
pos = nstations x 2 matrix with x (east) and y (north)
positions of stations in meters
disp = nstations x 2 matrix with x (east) and y (north)
displacements of stations in meters
k = Type of computation: Delaunay (k = 0), nearest
neighbor (k = 1), or distance weighted (k = 2)
par = Parameters for nearest neighbor or distance
weighted computation. If Delaunay (k = 0), enter
a scalar corresponding to the minimum internal
angle of a triangle valid for computation.
If nearest neighbor (k = 1), input a 1 x 3 vector
with grid spacing, number of nearest neighbors,
and maximum distance to neighbors. If distance
weighted (k = 2), input a 1 x 2 vector with grid
spacing and distance weighting factor alpha
plotpar = Parameter to color the cells: Max elongation
(plotpar = 0), minimum elongation
(plotpar = 1), rotation (plotpar = 2),
or dilatation (plotpar = 3)
plotst = A flag to plot the stations (1) or not (0)
cent = ncells x 2 matrix with x and y positions of cells
centroids
eps = 3 x 3 x ncells array with strain tensors of
the cells
ome = 3 x 3 x ncells array with rotation tensors of
the cells
pstrain = 3 x 3 x ncells array with magnitude and
orientation of principal strains of
the cells
rotc = ncells x 3 matrix with rotation components
of cells
NOTE: Input/Output angles are in radians. Output
azimuths are given with respect to North
pos, disp, grid spacing, max. distance to
neighbors, and alpha should be in meters
GridStrain uses functions lscov and InfStrain
Python function translated from the Matlab function
GridStrain in Allmendinger et al. (2012)
'''
pi = np.pi
# If Delaunay
if k == 0:
# Indexes of triangles vertices
# Use function Delaunay
tri = Delaunay(pos)
inds = tri.simplices
# Number of cells
ncells = np.size(inds, 0)
# Number of stations per cell = 3
nstat = 3
# Centers of cells
cent = np.zeros((ncells, 2))
for i in range(0, ncells):
# Triangle vertices
v1x = pos[inds[i, 0], 0]
v2x = pos[inds[i, 1], 0]
v3x = pos[inds[i, 2], 0]
v1y = pos[inds[i, 0], 1]
v2y = pos[inds[i, 1], 1]
v3y = pos[inds[i, 2], 1]
# Center of cell
cent[i, 0] = (v1x + v2x + v3x) / 3.0
cent[i, 1] = (v1y + v2y + v3y) / 3.0
# Triangle internal angles
s1 = np.sqrt((v3x - v2x)**2 + (v3y - v2y)**2)
s2 = np.sqrt((v1x - v3x)**2 + (v1y - v3y)**2)
s3 = np.sqrt((v2x - v1x)**2 + (v2y - v1y)**2)
a1 = np.arccos((v2x-v1x)*(v3x-v1x)/(s3*s2)+\
(v2y-v1y)*(v3y-v1y)/(s3*s2))
a2 = np.arccos((v3x-v2x)*(v1x-v2x)/(s1*s3)+\
(v3y-v2y)*(v1y-v2y)/(s1*s3))
a3 = np.arccos((v2x-v3x)*(v1x-v3x)/(s1*s2)+\
(v2y-v3y)*(v1y-v3y)/(s1*s2))
# If any of the internal angles is less than
# specified minimum, invalidate triangle
if a1 < par or a2 < par or a3 < par:
inds[i, :] = np.zeros(3)
# If nearest neighbor or distance weighted
else:
# Construct grid
xmin = min(pos[:, 0])
xmax = max(pos[:, 0])
ymin = min(pos[:, 1])
ymax = max(pos[:, 1])
cellsx = int(np.ceil((xmax - xmin) / par[0]))
cellsy = int(np.ceil((ymax - ymin) / par[0]))
xgrid = np.arange(xmin, (xmin + (cellsx + 1) * par[0]), par[0])
ygrid = np.arange(ymin, (ymin + (cellsy + 1) * par[0]), par[0])
XX, YY = np.meshgrid(xgrid, ygrid)
# Number of cells
ncells = cellsx * cellsy
# Number of stations per cell (nstat) and
# other parameters
# If nearest neighbor
if k == 1:
nstat = par[1] # max neighbors
sqmd = par[2]**2 # max squared distance
# If distance weighted
elif k == 2:
nstat = np.size(pos, 0) # all stations
dalpha = 2.0 * par[1] * par[1] # 2*alpha*alpha
# Cells' centers
cent = np.zeros((ncells, 2))
count = 0
for i in range(0, cellsy):
for j in range(0, cellsx):
cent[count, 0] = (XX[i, j] + XX[i, j + 1]) / 2.0
cent[count, 1] = (YY[i, j] + YY[i + 1, j]) / 2.0
count += 1
# Initialize stations indexes for cells to -1
inds = np.ones((ncells, nstat), dtype=int) * -1
# Initialize weight matrix for distance weighted
wv = np.zeros((ncells, nstat * 2))
# For all cells set stations indexes
for i in range(0, ncells):
# Initialize sq distances to -1.0
sds = np.ones(nstat) * -1.0
# For all stations
for j in range(0, np.size(pos, 0)):
# Sq distance from cell center to station
dx = cent[i, 0] - pos[j, 0]
dy = cent[i, 1] - pos[j, 1]
sd = dx**2 + dy**2
# If nearest neighbor
if k == 1:
# If within the max sq distance
if sd <= sqmd:
minsd = min(sds)
mini = np.argmin(sds)
# If less than max neighbors
if minsd == -1.0:
sds[mini] = sd
inds[i, mini] = j
# If max neighbors
else:
# If sq distance is less
# than neighbors max sq distance
maxsd = max(sds)
maxi = np.argmax(sds)
if sd < maxsd:
sds[maxi] = sd
inds[i, maxi] = j
# If distance weighted
elif k == 2:
# All stations indexes
inds[i, :] = np.arange(nstat)
# Eq. 8.27: Weight factor
weight = np.exp(-sd / dalpha)
wv[i, j * 2] = weight
wv[i, j * 2 + 1] = weight
# Initialize arrays
y = np.zeros(nstat * 2)
M = np.zeros((nstat * 2, 6))
e = np.zeros((3, 3))
eps = np.zeros((3, 3, ncells))
ome = np.zeros((3, 3, ncells))
pstrain = np.zeros((3, 3, ncells))
rotc = np.zeros((ncells, 3))
# For each cell
for i in range(0, ncells):
# If required minimum number of stations
if min(inds[i, :]) >= 0:
# Eq. 8.24: Displacements column vector y
# and design matrix M. X1 = North, X2 = East
for j in range(0, nstat):
ic = inds[i, j]
y[j * 2] = disp[ic, 1]
y[j * 2 + 1] = disp[ic, 0]
M[j * 2, :] = [1., 0., pos[ic, 1], pos[ic, 0], 0., 0.]
M[j * 2 + 1, :] = [0., 1., 0., 0., pos[ic, 1], pos[ic, 0]]
# Eqs. 8.25-8.26: Find x using function lscov
# If Delaunay or nearest neighbor
if k == 0 or k == 1:
x = lscov(M, y)
# If distance weighted
elif k == 2:
x = lscov(M, y, wv[i, :])
# Displacement gradient tensor
for j in range(0, 2):
e[j, 0] = x[j * 2 + 2]
e[j, 1] = x[j * 2 + 3]
# Compute strain
eps[:,:,i],ome[:,:,i],pstrain[:,:,i],\
rotc[i,:],_ = InfStrain(e)
# Variable to plot
# If maximum principal strain
if plotpar == 0:
vp = pstrain[0, 0, :]
lcb = "emax"
# If minimum principal strain
elif plotpar == 1:
vp = pstrain[2, 0, :]
lcb = "emin"
# If rotation:
# For plane strain, rotation = rotc(3)
elif plotpar == 2:
vp = rotc[:, 2] * 180 / pi
lcb = "Rotation (deg)"
# If dilatation
elif plotpar == 3:
vp = pstrain[0, 0, :] + pstrain[1, 0, :] + pstrain[2, 0, :]
lcb = "dilatation"
# Make a figure
fig, ax = plt.subplots()
fig.set_size_inches(15.0, 7.5)
# Patches and colors for cells
patches = []
colors = []
# Fill cells patches and colors
# If Delaunay
if k == 0:
for i in range(0, ncells):
# If minimum number of stations
if min(inds[i, :]) >= 0:
xpyp = [[pos[inds[i,0],0],pos[inds[i,0],1]],\
[pos[inds[i,1],0],pos[inds[i,1],1]],\
[pos[inds[i,2],0],pos[inds[i,2],1]]]
# length in km
xpyp = np.divide(xpyp, 1e3)
polygon = Polygon(xpyp, True)
patches.append(polygon)
colors.append(vp[i])
# If nearest neighbor or distance weighted
if k == 1 or k == 2:
count = 0
for i in range(0, cellsy):
for j in range(0, cellsx):
# If minimum number of stations
if min(inds[count, :]) >= 0:
xpyp = [[XX[i,j],YY[i,j]],[XX[i,j+1],YY[i,j+1]],\
[XX[i+1,j+1],YY[i+1,j+1]],[XX[i+1,j],YY[i+1,j]]]
# length in km
xpyp = np.divide(xpyp, 1e3)
polygon = Polygon(xpyp, True)
patches.append(polygon)
colors.append(vp[count])
count += 1
# Collect cells patches
pcoll = PatchCollection(patches)
# Cells colors
pcoll.set_array(np.array(colors))
# Color map is blue to red
pcoll.set_cmap('bwr')
# Positive values are red, negative are
# blue and zero is white
vmin = min(vp)
vmax = max(vp)
norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
pcoll.set_norm(norm)
# Draw cells
ax.add_collection(pcoll)
# Plot stations
if plotst == 1:
plt.plot(pos[:, 0] * 1e-3, pos[:, 1] * 1e-3, 'k.', markersize=2)
# Axes
plt.axis('equal')
plt.xlabel('x (km)')
plt.ylabel('y (km)')
# Color bar with nice ticks
intv = (vmax - vmin) * 0.25
ticks = [vmin, vmin + intv, vmin + 2 * intv, vmin + 3 * intv, vmax]
lticks = ['{:.2e}'.format(ticks[0]),\
'{:.2e}'.format(ticks[1]),'{:.2e}'.format(ticks[2]),\
'{:.2e}'.format(ticks[3]),'{:.2e}'.format(ticks[4])]
cbar = fig.colorbar(pcoll, label=lcb, ticks=ticks)
cbar.ax.set_yticklabels(lticks)
# Show plot
plt.show()
return cent, eps, ome, pstrain, rotc
### PLOT FITTED MODEL (SCATTER)
vel_z_predicted = model.predict(x)
fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
plot = ax.scatter(theta,
r,
c=vel_z_predicted,
s=10,
vmin=np.min(vel_z),
vmax=np.max(vel_z))
ax.set_rmax(0.33)
fig.colorbar(plot)
### PLOT ERROR (PROFILE - FITTED)
vel_z_error = (vel_z - vel_z_predicted)
divnorm = mcolors.TwoSlopeNorm(vmin=np.min(vel_z_error),
vcenter=0,
vmax=np.max(vel_z_error))
fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})
plot = ax.scatter(theta, r, c=vel_z_error, cmap="seismic", norm=divnorm, s=2)
ax.set_rmax(0.33)
fig.colorbar(plot)
### PLOT FITTED MODEL (CONTOUR MAP)
r_list = np.linspace(np.min(r), np.max(r), 1000)
theta_list = np.linspace(np.min(theta), np.max(theta), 1000)
np.linspace(np.pi / 2, 3 * np.pi / 2, 1000)
r_mesh, theta_mesh = np.meshgrid(r_list, theta_list)
def Z(r, theta):
return r * 0.00001 + theta * 0.00001
scale = 2
l_up = -14.47 * scale + 31.45
l_down = -l_up
x = linspace(l_down, l_up, 101)
y = linspace(l_down, l_up, 101)
xx, yy = meshgrid(x, y)
t = (xx**2 + yy**2)**(1 / 2)
zz = (2 * pi**(-1 / 4) /
3**(1 / 2)) * (1 - (scale * t)**2) * e**(-(scale * t)**2 / 2)
#zz = (1 - (scale*t)**2) * e**(-(scale*t)**2 / 2)
period = 2 * pi / (scale * 2.5**(1 / 2))
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
offset = mcolors.TwoSlopeNorm(vcenter=0, vmin=-0.4, vmax=1)
mapp = ax.plot_surface(xx, yy, zz, cmap='inferno', norm=offset)
fig.colorbar(mapp, ax=ax, extend='both')
fig.show()
input()
for chemLabel, plotLabel in label_Conver.items():
if chemLabel != 'H1_4102A':
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot()
dict_label = f'{plotLabel}/Hdelta'
flux_image = fits.getdata(linesMaps_fits_address, chemLabel,
ver=1) / hbeta_image
print(chemLabel)
print(np.nanmin(flux_image), theoEmis_dict[dict_label],
np.nanmax(flux_image))
print()
divnorm = colors.TwoSlopeNorm(vmin=np.nanmin(flux_image),
vcenter=theoEmis_dict[dict_label],
vmax=np.nanmax(flux_image))
im = ax.imshow(flux5007_image, cmap=halpha_cmap)
im2 = ax.imshow(flux_image, cmap='RdBu_r', norm=divnorm)
cbar = fig.colorbar(im2, ax=ax)
cbar.set_label('Line ratio (white theoretical value)',
rotation=270,
labelpad=50,
fontsize=15)
ratio_label = r'$\frac{{{}}}{{{}}}$'.format(latex_Conver[chemLabel],
latex_Conver['H1_4861A'])
ax.update({
'title': r'Galaxy J0838: {}'.format(ratio_label),
def test_TwoSlopeNorm_premature_scaling():
norm = mcolors.TwoSlopeNorm(vcenter=2)
with pytest.raises(ValueError):
norm.inverse(np.array([0.1, 0.5, 0.9]))
def test_TwoSlopeNorm_VcenterGTVmax():
with pytest.raises(ValueError):
mcolors.TwoSlopeNorm(vmin=10, vcenter=25, vmax=20)
def test_TwoSlopeNorm_VmaxEqualsVcenter():
with pytest.raises(ValueError):
mcolors.TwoSlopeNorm(vmin=-2, vcenter=2, vmax=2)
