Skip to content

Commit 7a8cd96

Browse files
committed
viz step always plots from simulation step
1 parent ef39cde commit 7a8cd96

File tree

1 file changed

+239
-70
lines changed

1 file changed

+239
-70
lines changed

compass/ocean/tests/isomip_plus/viz/__init__.py

Lines changed: 239 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ def run(self):
5555
plot_streamfunctions = section.getboolean('plot_streamfunctions')
5656
plot_haney = section.getboolean('plot_haney')
5757
frames_per_second = section.getint('frames_per_second')
58+
# output_slices = section.getint('output_slices')
5859
movie_format = section.get('movie_format')
5960
section_y = section.getfloat('section_y')
6061

@@ -65,47 +66,167 @@ def run(self):
6566
show_progress = self.log_filename is None
6667

6768
expt = self.experiment
68-
if self.tidal_forcing:
69-
sim_dir = '../performance'
70-
else:
71-
sim_dir = '../simulation'
72-
if not os.path.exists(f'{sim_dir}/timeSeriesStatsMonthly.0001-01-01.nc'): # noqa: E501
73-
sim_dir = '../performance'
69+
sim_dir = '../simulation'
7470
streamfunction_dir = '../streamfunction'
7571
out_dir = '.'
7672

7773
dsMesh = xarray.open_dataset(f'{sim_dir}/init.nc')
74+
maxBottomDepth = dsMesh.bottomDepth.max().values
7875
dsOut = xarray.open_dataset(f'{sim_dir}/output.nc')
79-
76+
dsIce = xarray.open_dataset(f'{sim_dir}/land_ice_fluxes.nc')
77+
dsPerf = xarray.open_dataset('../performance/output.nc')
78+
# ds = xarray.open_mfdataset(
79+
# f'{sim_dir}/land_ice_fluxes.nc',
80+
# concat_dim='Time', combine='nested')
81+
if dsOut.sizes['Time'] < 50:
82+
print('Cropping dsOut')
83+
dsOut = dsOut.isel(Time=slice(-31, -1))
84+
nTime = dsOut.sizes['Time']
85+
if nTime != dsIce.sizes['Time']:
86+
print('Cropping dsIce to match dsOut')
87+
dsIce = dsIce.isel(Time=slice(-nTime - 1, -1))
88+
dsOut['landIceDraft'] = dsIce.landIceDraft
8089
plotter = MoviePlotter(inFolder=sim_dir,
8190
streamfunctionFolder=streamfunction_dir,
8291
outFolder=f'{out_dir}/plots',
8392
expt=expt, sectionY=section_y,
8493
dsMesh=dsMesh, ds=dsOut,
8594
showProgress=show_progress)
95+
plotter.plot_3d_field_top_bot_section(dsOut.velocityX, 'u', 'u', True,
96+
cmap='cmo.balance',
97+
vmin=-5e-1, vmax=5e-1)
98+
plotter.plot_3d_field_top_bot_section(dsOut.velocityY, 'v', 'v', True,
99+
cmap='cmo.balance',
100+
vmin=-5e-1, vmax=5e-1)
101+
102+
if dsOut.sizes['Time'] > 24:
103+
days = 30
104+
samples_per_day = 1
105+
dsAve = dsOut.isel(Time=slice(-int(days * samples_per_day), -1))
106+
velocityMagSq = (dsAve.velocityX.isel(nVertLevels=0)**2. +
107+
dsAve.velocityX.isel(nVertLevels=0)**2.)
108+
rmsVelocity = xarray.DataArray(
109+
data=np.sqrt(velocityMagSq.mean('Time').values),
110+
dims=['nCells']).expand_dims(dim='Time', axis=0)
111+
plotter.plot_horiz_series(rmsVelocity,
112+
'rmsVelocity', 'rmsVelocity',
113+
False, vmin=1.e-2, vmax=1.e0,
114+
cmap='cmo.speed', cmap_scale='log')
115+
116+
cavityMask = dsMesh.landIceFloatingMask == 1
117+
if 'xIsomipCell' in dsMesh.keys():
118+
xCell = dsMesh.xIsomipCell / 1.e3
119+
else:
120+
xCell = dsMesh.xCell / 1.e3
121+
# salinity = dsIce.landIceInterfaceSalinity
122+
salinity = dsOut.salinity.isel(nVertLevels=0)
123+
pressure = dsIce.landIcePressure
124+
# temperature = dsIce.landIceInterfaceTemperature
125+
temperature = dsOut.temperature.isel(nVertLevels=0)
126+
coeff_0 = 6.22e-2
127+
coeff_S = -5.63e-2
128+
coeff_p = -7.43e-8
129+
coeff_pS = -1.74e-10
130+
coeff_mushy = 0.
131+
ocn_freezing_temperature = \
132+
coeff_0 + coeff_S * salinity + \
133+
coeff_p * pressure + \
134+
coeff_pS * pressure * salinity + \
135+
coeff_mushy * salinity / (1.0 - salinity / 1e3)
136+
dsIce['thermalForcing'] = temperature - ocn_freezing_temperature
137+
138+
dx = 2.
139+
xbins = np.arange(xCell.isel(nCells=cavityMask).min(),
140+
xCell.isel(nCells=cavityMask).max(), dx)
141+
# time_slice = [-1]
142+
# time_slice = slice(-output_slices - 1,-1)
143+
if dsIce.sizes['Time'] > 14:
144+
dsIce = dsIce.isel(Time=slice(-13, -12))
145+
for time_slice in range(dsIce.sizes['Time']):
146+
title = dsIce.xtime.isel(Time=time_slice).values
147+
meltTransect = np.zeros_like(xbins)
148+
landIceDraftTransect = np.zeros_like(xbins)
149+
frictionVelocityTransect = np.zeros_like(xbins)
150+
temperatureTransect = np.zeros_like(xbins)
151+
thermalForcingTransect = np.zeros_like(xbins)
152+
rmsVelocityTransect = np.zeros_like(xbins)
153+
for i, xmin in enumerate(xbins):
154+
binMask = np.logical_and(cavityMask,
155+
np.logical_and(xCell >= xmin,
156+
xCell < xmin + dx))
157+
dsTransect = dsIce.isel(nCells=binMask, Time=time_slice)
158+
meltTransect[i] = dsTransect.landIceFreshwaterFlux.mean()
159+
landIceDraftTransect[i] = dsTransect.landIceDraft.mean()
160+
frictionVelocityTransect[i] = \
161+
dsTransect.landIceFrictionVelocity.mean()
162+
temperatureTransect[i] = \
163+
dsTransect.landIceInterfaceTemperature.mean()
164+
thermalForcingTransect[i] = dsTransect.thermalForcing.mean()
165+
if dsOut.sizes['Time'] > 24:
166+
rmsVelocityTransect[i] = np.mean(rmsVelocity[binMask])
167+
168+
fig, ax = plt.subplots(4, 1, sharex=True, figsize=(4, 8))
169+
color = 'darkgreen'
170+
x = xbins - xbins[0]
171+
ax[0].set_title(title)
172+
ax[0].plot(x, landIceDraftTransect, color=color)
173+
ax[0].set_ylabel('Land ice draft (m)')
174+
secPerYear = 365 * 24 * 60 * 60
175+
density = 1026.
176+
ax[1].plot(x, meltTransect * secPerYear / density, color=color)
177+
ax[1].set_ylabel('Melt rate (m/yr)')
178+
ax[2].plot(x, frictionVelocityTransect, color=color)
179+
ax[2].set_ylabel('Friction velocity (m/s)')
180+
ax[3].plot(x, thermalForcingTransect, color=color)
181+
ax[3].set_ylabel(r'ThermalForcing ($^{\circ}$C)')
182+
# if dsOut.sizes['Time'] > 24:
183+
# ax2 = ax[2].twinx()
184+
# color = 'DarkGreen'
185+
# ax2.set_ylabel('RMS velocity (m/s)', color=color)
186+
# ax2.plot(x, rmsVelocityTransect, color=color)
187+
# ax2.tick_params(axis ='y', labelcolor=color)
188+
ax[3].set_xlabel('Distance along ice flow (km)')
189+
plt.tight_layout()
190+
plt.savefig(f'meanTransect_tend{time_slice:03g}.png', dpi=600,
191+
bbox_inches='tight', transparent=True)
192+
plt.close(fig)
86193

87-
if 'time_varying' in expt:
88-
plotter.plot_horiz_series(dsOut.ssh, 'ssh', 'ssh', True,
89-
cmap='cmo.curl')
90-
delice = dsOut.landIcePressure - dsOut.landIcePressure[0, :]
91-
plotter.plot_horiz_series(delice, 'delLandIcePressure',
92-
'delLandIcePressure', True,
93-
cmap='cmo.curl')
94-
plotter.plot_horiz_series(dsOut.velocityX[:, :, 0], 'u', 'u', True,
95-
cmap='cmo.balance', vmin=-5e-1, vmax=5e-1)
96-
plotter.plot_horiz_series(dsOut.velocityY[:, :, 0], 'v', 'v', True,
97-
cmap='cmo.balance', vmin=-5e-1, vmax=5e-1)
194+
plotter.plot_horiz_series(dsOut.ssh, 'ssh', 'ssh', True,
195+
vmin=-maxBottomDepth, vmax=0)
196+
if 'layerThickness' in dsOut.keys():
197+
plotter.plot_layer_interfaces()
198+
delssh = dsOut.ssh - dsMesh.ssh.isel(Time=0)
199+
plotter.plot_horiz_series(delssh, 'ssh - ssh_init', 'delssh_init',
200+
True, cmap='cmo.curl', vmin=-1, vmax=1)
98201
plotter.plot_horiz_series(dsOut.ssh + dsMesh.bottomDepth, 'H', 'H',
99202
True, vmin=min_column_thickness + 1e-10,
100-
vmax=700, cmap_set_under='r',
203+
vmax=500, cmap_set_under='k',
101204
cmap_scale='log')
205+
plotter.plot_horiz_series(dsOut.landIcePressure,
206+
'landIcePressure', 'landIcePressure',
207+
True, vmin=1e5, vmax=1e7, cmap_scale='log')
208+
delLandIcePressureApplied = dsOut.landIcePressureApplied - \
209+
dsPerf.landIcePressureApplied.isel(Time=0)
210+
plotter.plot_horiz_series(delLandIcePressureApplied,
211+
'delLandIcePressureApplied',
212+
'delLandIcePressureApplied',
213+
True, vmin=-1e6, vmax=1e6,
214+
cmap='cmo.balance')
215+
plotter.plot_horiz_series(dsOut.landIcePressureApplied,
216+
'landIcePressureApplied',
217+
'landIcePressureApplied',
218+
True, vmin=1e5, vmax=1e7, cmap_scale='log')
102219
plotter.plot_horiz_series(dsOut.surfacePressure,
103220
'surfacePressure', 'surfacePressure',
104221
True, vmin=1e5, vmax=1e7, cmap_scale='log')
105222
delSurfacePressure = dsOut.landIcePressure - dsOut.surfacePressure
106223
plotter.plot_horiz_series(delSurfacePressure,
107224
'delSurfacePressure', 'delSurfacePressure',
108225
True, vmin=1e5, vmax=1e7, cmap_scale='log')
226+
delice = dsOut.landIcePressure - dsMesh.landIcePressure.isel(Time=0)
227+
plotter.plot_horiz_series(delice, 'delLandIcePressure',
228+
'delLandIcePressure', True,
229+
vmin=-1e6, vmax=1e6, cmap='cmo.balance')
109230
# Cannot currently plot because it is on edges
110231
for t in range(dsOut.sizes['Time']):
111232
fig = plt.figure()
@@ -121,73 +242,121 @@ def run(self):
121242
plt.savefig(f'wettingVelocityFactor_{t:02g}.png')
122243
plt.close(fig)
123244

124-
if 'tidal' in expt:
125-
delssh = dsOut.ssh - dsOut.ssh[0, :]
126-
plotter.plot_horiz_series(delssh, 'delssh', 'delssh', True,
127-
cmap='cmo.curl', vmin=-1, vmax=1)
128-
129245
wct = dsOut.ssh + dsMesh.bottomDepth
130-
idx_thin = np.logical_and(wct[0, :] > 1e-1,
131-
wct[0, :] < 1)
132-
wct_thin = wct[:, idx_thin]
133-
wct_mean = wct_thin.mean(dim='nCells').values
134-
time = dsOut.daysSinceStartOfSim.values
135-
plt.figure()
136-
plt.plot(time, wct_mean, '.')
137-
plt.xlabel('Time (days)')
138-
plt.ylabel('Mean thickness of thin film (m)')
139-
plt.savefig('wct_thin_t.png')
140-
plt.close()
246+
# idx_thin = np.logical_and(wct[0, :] > 1e-1,
247+
# wct[0, :] < 1)
248+
# wct_thin = wct[:, idx_thin]
249+
# wct_mean = wct_thin.mean(dim='nCells').values
250+
# time = dsOut.daysSinceStartOfSim.values
251+
# plt.figure()
252+
# plt.plot(time, wct_mean, '.')
253+
# plt.xlabel('Time (days)')
254+
# plt.ylabel('Mean thickness of thin film (m)')
255+
# plt.savefig('wct_thin_t.png')
256+
# plt.close()
141257

142258
plotter.plot_horiz_series(wct, 'H', 'H',
143259
True, vmin=min_column_thickness + 1e-10,
144-
vmax=700, cmap_set_under='r',
260+
vmax=700, cmap_set_under='k',
145261
cmap_scale='log')
146262

147-
if os.path.exists(f'{sim_dir}/timeSeriesStatsMonthly.0001-01-01.nc'):
263+
ds = xarray.open_mfdataset(
264+
'{}/timeSeriesStatsMonthly*.nc'.format(sim_dir),
265+
concat_dim='Time', combine='nested')
266+
267+
if plot_haney:
268+
_compute_and_write_haney_number(dsMesh, ds, out_dir,
269+
showProgress=show_progress)
270+
271+
tsPlotter = TimeSeriesPlotter(inFolder=sim_dir,
272+
outFolder='{}/plots'.format(out_dir),
273+
expt=expt)
274+
tsPlotter.plot_melt_time_series()
275+
tsPlotter.plot_melt_time_series(wctMax=20.)
276+
tsPlotter = TimeSeriesPlotter(
277+
inFolder=sim_dir,
278+
outFolder='{}/timeSeriesBelow300m'.format(out_dir),
279+
expt=expt)
280+
tsPlotter.plot_melt_time_series(sshMax=-300.)
281+
282+
mPlotter = MoviePlotter(inFolder=sim_dir,
283+
streamfunctionFolder=streamfunction_dir,
284+
outFolder='{}/plots'.format(out_dir),
285+
expt=expt, sectionY=section_y,
286+
dsMesh=dsMesh, ds=ds,
287+
showProgress=show_progress)
288+
289+
mPlotter.plot_layer_interfaces()
290+
291+
if plot_streamfunctions:
292+
mPlotter.plot_barotropic_streamfunction()
293+
mPlotter.plot_overturning_streamfunction()
294+
295+
if plot_haney:
296+
mPlotter.plot_haney_number(haneyFolder=out_dir)
297+
298+
mPlotter.plot_melt_rates()
299+
mPlotter.plot_ice_shelf_boundary_variables()
300+
mPlotter.plot_temperature()
301+
mPlotter.plot_salinity()
302+
mPlotter.plot_potential_density()
303+
304+
mPlotter.images_to_movies(outFolder=f'{out_dir}/movies',
305+
framesPerSecond=frames_per_second,
306+
extension=movie_format)
307+
308+
if os.path.exists('../simulation/land_ice_fluxes.nc'):
148309
ds = xarray.open_mfdataset(
149-
'{}/timeSeriesStatsMonthly*.nc'.format(sim_dir),
310+
'../simulation/land_ice_fluxes.nc',
150311
concat_dim='Time', combine='nested')
151-
152-
if plot_haney:
153-
_compute_and_write_haney_number(dsMesh, ds, out_dir,
154-
showProgress=show_progress)
155-
156-
tsPlotter = TimeSeriesPlotter(inFolder=sim_dir,
157-
outFolder='{}/plots'.format(out_dir),
158-
expt=expt)
159-
tsPlotter.plot_melt_time_series()
160-
tsPlotter = TimeSeriesPlotter(
161-
inFolder=sim_dir,
162-
outFolder='{}/timeSeriesBelow300m'.format(out_dir),
163-
expt=expt)
164-
tsPlotter.plot_melt_time_series(sshMax=-300.)
165-
312+
if ds.sizes['Time'] != dsOut.sizes['Time']:
313+
# assumes ds.sizes['Time'] > dsOut.sizes['Time']
314+
ds = ds.isel(Time=slice(-31, -1))
315+
dsOut = dsOut.isel(Time=slice(-31, -1))
316+
if ds.sizes['Time'] != dsOut.sizes['Time']:
317+
print('error in time selection')
318+
return
319+
if ds.sizes['Time'] > 50:
320+
return
321+
ds['layerThickness'] = dsOut.layerThickness
166322
mPlotter = MoviePlotter(inFolder=sim_dir,
167323
streamfunctionFolder=streamfunction_dir,
168324
outFolder='{}/plots'.format(out_dir),
169325
expt=expt, sectionY=section_y,
170326
dsMesh=dsMesh, ds=ds,
171327
showProgress=show_progress)
172-
173-
mPlotter.plot_layer_interfaces()
174-
175-
if plot_streamfunctions:
176-
mPlotter.plot_barotropic_streamfunction()
177-
mPlotter.plot_overturning_streamfunction()
178-
179-
if plot_haney:
180-
mPlotter.plot_haney_number(haneyFolder=out_dir)
181-
328+
mPlotter.plot_horiz_series(ds.landIceFrictionVelocity,
329+
'landIceFrictionVelocity',
330+
'frictionVelocity',
331+
True)
332+
mPlotter.plot_horiz_series(ds.landIceFreshwaterFlux,
333+
'landIceFreshwaterFlux', 'melt',
334+
True)
335+
delLandIceDraft = ds.landIceDraft - dsMesh.landIceDraft
336+
vmax = np.max(np.abs(delLandIceDraft.values))
337+
mPlotter.plot_horiz_series(delLandIceDraft,
338+
'landIceDraft - landIceDraftInit',
339+
'delLandIceDraft',
340+
True, vmin=-vmax, vmax=vmax,
341+
cmap='cmo.balance')
342+
mPlotter.plot_horiz_series(ds.landIceFloatingFraction,
343+
'landIceFloatingFraction',
344+
'landIceFloatingFraction',
345+
True, vmin=1e-16, vmax=1,
346+
cmap_set_under='k')
347+
mPlotter.plot_horiz_series(ds.landIceDraft,
348+
'landIceDraft', 'landIceDraft',
349+
True, vmin=-maxBottomDepth, vmax=0.)
350+
delssh = dsOut.ssh - ds.landIceDraft
351+
vmax = np.max(np.abs(delssh.values))
352+
mPlotter.plot_horiz_series(delssh,
353+
'ssh - landIceDraft', 'delssh',
354+
True, vmin=-vmax, vmax=vmax,
355+
cmap='cmo.balance', cmap_scale='symlog')
182356
mPlotter.plot_melt_rates()
183357
mPlotter.plot_ice_shelf_boundary_variables()
184-
mPlotter.plot_temperature()
185-
mPlotter.plot_salinity()
186-
mPlotter.plot_potential_density()
187-
188-
mPlotter.images_to_movies(outFolder='{}/movies'.format(out_dir),
189-
framesPerSecond=frames_per_second,
190-
extension=movie_format)
358+
else:
359+
print(f'didnt find {sim_dir}/land_ice_fluxes.nc')
191360

192361

193362
def file_complete(ds, fileName):

0 commit comments

Comments
 (0)