Skip to content

Commit 2a116ba

Browse files
cbegemancbegeman
committed
fixup viz
1 parent d65be6a commit 2a116ba

File tree

1 file changed

+110
-106
lines changed

1 file changed

+110
-106
lines changed

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

Lines changed: 110 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,6 @@ def run(self):
7171
out_dir = '.'
7272

7373
dsMesh = xarray.open_dataset(f'{sim_dir}/init.nc')
74-
maxBottomDepth = dsMesh.bottomDepth.max().values
7574
dsOut = xarray.open_dataset(f'{sim_dir}/output.nc')
7675
dsIce = xarray.open_dataset(f'{sim_dir}/land_ice_fluxes.nc')
7776
if dsOut.sizes['Time'] < 50:
@@ -95,25 +94,6 @@ def run(self):
9594
cmap='cmo.balance',
9695
vmin=-5e-1, vmax=5e-1)
9796

98-
if dsOut.sizes['Time'] > 24:
99-
days = 30
100-
samples_per_day = 1
101-
dsAve = dsOut.isel(Time=slice(-int(days * samples_per_day), -1))
102-
velocityMagSq = (dsAve.velocityX.isel(nVertLevels=0)**2. +
103-
dsAve.velocityX.isel(nVertLevels=0)**2.)
104-
rmsVelocity = xarray.DataArray(
105-
data=np.sqrt(velocityMagSq.mean('Time').values),
106-
dims=['nCells']).expand_dims(dim='Time', axis=0)
107-
plotter.plot_horiz_series(rmsVelocity,
108-
'rmsVelocity', 'rmsVelocity',
109-
False, vmin=1.e-2, vmax=1.e0,
110-
cmap='cmo.speed', cmap_scale='log')
111-
112-
cavityMask = dsMesh.landIceFloatingMask == 1
113-
if 'xIsomipCell' in dsMesh.keys():
114-
xCell = dsMesh.xIsomipCell / 1.e3
115-
else:
116-
xCell = dsMesh.xCell / 1.e3
11797
# salinity = dsIce.landIceInterfaceSalinity
11898
salinity = dsOut.salinity.isel(nVertLevels=0)
11999
pressure = dsIce.landIcePressure
@@ -131,61 +111,35 @@ def run(self):
131111
coeff_mushy * salinity / (1.0 - salinity / 1e3)
132112
dsIce['thermalForcing'] = temperature - ocn_freezing_temperature
133113

134-
dx = 2.
135-
xbins = np.arange(xCell.isel(nCells=cavityMask).min(),
136-
xCell.isel(nCells=cavityMask).max(), dx)
137114
if dsIce.sizes['Time'] > 14:
138115
dsIce = dsIce.isel(Time=slice(-13, -12))
139-
for time_slice in range(dsIce.sizes['Time']):
140-
title = dsIce.xtime.isel(Time=time_slice).values
141-
meltTransect = np.zeros_like(xbins)
142-
landIceDraftTransect = np.zeros_like(xbins)
143-
frictionVelocityTransect = np.zeros_like(xbins)
144-
temperatureTransect = np.zeros_like(xbins)
145-
thermalForcingTransect = np.zeros_like(xbins)
146-
rmsVelocityTransect = np.zeros_like(xbins)
147-
for i, xmin in enumerate(xbins):
148-
binMask = np.logical_and(cavityMask,
149-
np.logical_and(xCell >= xmin,
150-
xCell < xmin + dx))
151-
dsTransect = dsIce.isel(nCells=binMask, Time=time_slice)
152-
meltTransect[i] = dsTransect.landIceFreshwaterFlux.mean()
153-
landIceDraftTransect[i] = dsTransect.landIceDraft.mean()
154-
frictionVelocityTransect[i] = \
155-
dsTransect.landIceFrictionVelocity.mean()
156-
temperatureTransect[i] = \
157-
dsTransect.landIceInterfaceTemperature.mean()
158-
thermalForcingTransect[i] = dsTransect.thermalForcing.mean()
159-
if dsOut.sizes['Time'] > 24:
160-
rmsVelocityTransect[i] = np.mean(rmsVelocity[binMask])
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
161121

162-
fig, ax = plt.subplots(4, 1, sharex=True, figsize=(4, 8))
163-
color = 'darkgreen'
164-
x = xbins - xbins[0]
165-
ax[0].set_title(title)
166-
ax[0].plot(x, landIceDraftTransect, color=color)
167-
ax[0].set_ylabel('Land ice draft (m)')
168-
secPerYear = 365 * 24 * 60 * 60
169-
density = 1026.
170-
ax[1].plot(x, meltTransect * secPerYear / density, color=color)
171-
ax[1].set_ylabel('Melt rate (m/yr)')
172-
ax[2].plot(x, frictionVelocityTransect, color=color)
173-
ax[2].set_ylabel('Friction velocity (m/s)')
174-
ax[3].plot(x, thermalForcingTransect, color=color)
175-
ax[3].set_ylabel(r'ThermalForcing ($^{\circ}$C)')
176-
ax[3].set_xlabel('Distance along ice flow (km)')
177-
plt.tight_layout()
178-
plt.savefig(f'meanTransect_tend{time_slice:03g}.png', dpi=600,
179-
bbox_inches='tight', transparent=True)
180-
plt.close(fig)
122+
if dsOut.sizes['Time'] > 24:
123+
days = 30
124+
samples_per_day = 1
125+
dsAve = dsOut.isel(Time=slice(-int(days * samples_per_day), -1))
126+
velocityMagSq = (dsAve.velocityX.isel(nVertLevels=0)**2. +
127+
dsAve.velocityX.isel(nVertLevels=0)**2.)
128+
rmsVelocity = xarray.DataArray(
129+
data=np.sqrt(velocityMagSq.mean('Time').values),
130+
dims=['nCells']).expand_dims(dim='Time', axis=0)
131+
plotter.plot_horiz_series(rmsVelocity,
132+
'rmsVelocity', 'rmsVelocity',
133+
False, vmin=1.e-2, vmax=1.e0,
134+
cmap='cmo.speed', cmap_scale='log')
135+
self._plot_transect(dsIce, rmsVelocity, xCell, cavityMask,
136+
outFolder=f'{out_dir}/plots')
181137

138+
maxBottomDepth = dsMesh.bottomDepth.max().values
182139
plotter.plot_horiz_series(dsOut.ssh, 'ssh', 'ssh', True,
183140
vmin=-maxBottomDepth, vmax=0)
184141
if 'layerThickness' in dsOut.keys():
185142
plotter.plot_layer_interfaces()
186-
delssh = dsOut.ssh - dsMesh.ssh.isel(Time=0)
187-
plotter.plot_horiz_series(delssh, 'ssh - ssh_init', 'delssh_init',
188-
True, cmap='cmo.curl', vmin=-1, vmax=1)
189143
plotter.plot_horiz_series(dsOut.ssh + dsMesh.bottomDepth, 'H', 'H',
190144
True, vmin=min_column_thickness + 1e-10,
191145
vmax=500, cmap_set_under='k',
@@ -204,6 +158,9 @@ def run(self):
204158
plotter.plot_horiz_series(delice, 'delLandIcePressure',
205159
'delLandIcePressure', True,
206160
vmin=-1e6, vmax=1e6, cmap='cmo.balance')
161+
delssh = dsOut.ssh - dsMesh.ssh.isel(Time=0)
162+
plotter.plot_horiz_series(delssh, 'ssh - ssh_init', 'delssh_init',
163+
True, cmap='cmo.curl', vmin=-1, vmax=1)
207164
# Cannot currently plot because it is on edges
208165
for t in range(dsOut.sizes['Time']):
209166
fig = plt.figure()
@@ -224,6 +181,26 @@ def run(self):
224181
True, vmin=min_column_thickness + 1e-10,
225182
vmax=700, cmap_set_under='k',
226183
cmap_scale='log')
184+
if os.path.exists('../simulation/land_ice_fluxes.nc'):
185+
ds = xarray.open_mfdataset(
186+
'../simulation/land_ice_fluxes.nc',
187+
concat_dim='Time', combine='nested')
188+
if ds.sizes['Time'] > 50 or \
189+
ds.sizes['Time'] != dsOut.sizes['Time']:
190+
# assumes ds.sizes['Time'] > dsOut.sizes['Time']
191+
ds = ds.isel(Time=slice(-31, -1))
192+
dsOut = dsOut.isel(Time=slice(-31, -1))
193+
if ds.sizes['Time'] != dsOut.sizes['Time']:
194+
print('error in time selection')
195+
return
196+
ds['layerThickness'] = dsOut.layerThickness
197+
mPlotter = MoviePlotter(inFolder=sim_dir,
198+
streamfunctionFolder=streamfunction_dir,
199+
outFolder='{}/plots'.format(out_dir),
200+
expt=expt, sectionY=section_y,
201+
dsMesh=dsMesh, ds=ds,
202+
showProgress=show_progress)
203+
self._plot_land_ice_variables(ds, mPlotter, maxBottomDepth)
227204

228205
ds = xarray.open_mfdataset(
229206
'{}/timeSeriesStatsMonthly*.nc'.format(sim_dir),
@@ -270,45 +247,72 @@ def run(self):
270247
framesPerSecond=frames_per_second,
271248
extension=movie_format)
272249

273-
if os.path.exists('../simulation/land_ice_fluxes.nc'):
274-
ds = xarray.open_mfdataset(
275-
'../simulation/land_ice_fluxes.nc',
276-
concat_dim='Time', combine='nested')
277-
if ds.sizes['Time'] != dsOut.sizes['Time']:
278-
# assumes ds.sizes['Time'] > dsOut.sizes['Time']
279-
ds = ds.isel(Time=slice(-31, -1))
280-
dsOut = dsOut.isel(Time=slice(-31, -1))
281-
if ds.sizes['Time'] != dsOut.sizes['Time']:
282-
print('error in time selection')
283-
return
284-
if ds.sizes['Time'] > 50:
285-
return
286-
ds['layerThickness'] = dsOut.layerThickness
287-
mPlotter = MoviePlotter(inFolder=sim_dir,
288-
streamfunctionFolder=streamfunction_dir,
289-
outFolder='{}/plots'.format(out_dir),
290-
expt=expt, sectionY=section_y,
291-
dsMesh=dsMesh, ds=ds,
292-
showProgress=show_progress)
293-
mPlotter.plot_horiz_series(ds.landIceFrictionVelocity,
294-
'landIceFrictionVelocity',
295-
'frictionVelocity',
296-
True)
297-
mPlotter.plot_horiz_series(ds.landIceFreshwaterFlux,
298-
'landIceFreshwaterFlux', 'melt',
299-
True)
300-
mPlotter.plot_horiz_series(ds.landIceFloatingFraction,
301-
'landIceFloatingFraction',
302-
'landIceFloatingFraction',
303-
True, vmin=1e-16, vmax=1,
304-
cmap_set_under='k')
305-
mPlotter.plot_horiz_series(ds.landIceDraft,
306-
'landIceDraft', 'landIceDraft',
307-
True, vmin=-maxBottomDepth, vmax=0.)
308-
mPlotter.plot_melt_rates()
309-
mPlotter.plot_ice_shelf_boundary_variables()
310-
else:
311-
print(f'didnt find {sim_dir}/land_ice_fluxes.nc')
250+
def _plot_land_ice_variables(ds, mPlotter, maxBottomDepth):
251+
mPlotter.plot_horiz_series(ds.landIceFrictionVelocity,
252+
'landIceFrictionVelocity',
253+
'frictionVelocity',
254+
True)
255+
mPlotter.plot_horiz_series(ds.landIceFreshwaterFlux,
256+
'landIceFreshwaterFlux', 'melt',
257+
True)
258+
mPlotter.plot_horiz_series(ds.landIceFloatingFraction,
259+
'landIceFloatingFraction',
260+
'landIceFloatingFraction',
261+
True, vmin=1e-16, vmax=1,
262+
cmap_set_under='k')
263+
mPlotter.plot_horiz_series(ds.landIceDraft,
264+
'landIceDraft', 'landIceDraft',
265+
True, vmin=-maxBottomDepth, vmax=0.)
266+
mPlotter.plot_melt_rates()
267+
mPlotter.plot_ice_shelf_boundary_variables()
268+
269+
def _plot_transect(dsIce, rmsVelocity, xCell, cavityMask, outFolder='.'):
270+
dx = 2.
271+
xbins = np.arange(xCell.isel(nCells=cavityMask).min(),
272+
xCell.isel(nCells=cavityMask).max(), dx)
273+
for time_slice in range(dsIce.sizes['Time']):
274+
title = dsIce.xtime.isel(Time=time_slice).values
275+
meltTransect = np.zeros_like(xbins)
276+
landIceDraftTransect = np.zeros_like(xbins)
277+
frictionVelocityTransect = np.zeros_like(xbins)
278+
temperatureTransect = np.zeros_like(xbins)
279+
thermalForcingTransect = np.zeros_like(xbins)
280+
rmsVelocityTransect = np.zeros_like(xbins)
281+
for i, xmin in enumerate(xbins):
282+
binMask = np.logical_and(cavityMask,
283+
np.logical_and(xCell >= xmin,
284+
xCell < xmin + dx))
285+
if (np.sum(binMask) < 1):
286+
continue
287+
dsTransect = dsIce.isel(nCells=binMask, Time=time_slice)
288+
meltTransect[i] = dsTransect.landIceFreshwaterFlux.mean()
289+
landIceDraftTransect[i] = dsTransect.landIceDraft.mean()
290+
frictionVelocityTransect[i] = \
291+
dsTransect.landIceFrictionVelocity.mean()
292+
temperatureTransect[i] = \
293+
dsTransect.landIceInterfaceTemperature.mean()
294+
thermalForcingTransect[i] = dsTransect.thermalForcing.mean()
295+
rmsVelocityTransect[i] = np.mean(rmsVelocity[binMask])
296+
297+
fig, ax = plt.subplots(4, 1, sharex=True, figsize=(4, 8))
298+
color = 'darkgreen'
299+
x = xbins - xbins[0]
300+
ax[0].set_title(title)
301+
ax[0].plot(x, landIceDraftTransect, color=color)
302+
ax[0].set_ylabel('Land ice draft (m)')
303+
secPerYear = 365 * 24 * 60 * 60
304+
density = 1026.
305+
ax[1].plot(x, meltTransect * secPerYear / density, color=color)
306+
ax[1].set_ylabel('Melt rate (m/yr)')
307+
ax[2].plot(x, frictionVelocityTransect, color=color)
308+
ax[2].set_ylabel('Friction velocity (m/s)')
309+
ax[3].plot(x, thermalForcingTransect, color=color)
310+
ax[3].set_ylabel(r'ThermalForcing ($^{\circ}$C)')
311+
ax[3].set_xlabel('Distance along ice flow (km)')
312+
plt.tight_layout()
313+
plt.savefig(f'{outFolder}/meanTransect_tend{time_slice:03g}.png',
314+
dpi=600, bbox_inches='tight', transparent=True)
315+
plt.close(fig)
312316

313317

314318
def file_complete(ds, fileName):

0 commit comments

Comments
 (0)