Skip to content

Commit

Permalink
fixup viz
Browse files Browse the repository at this point in the history
  • Loading branch information
@cbegeman
cbegeman committed Sep 26, 2023
1 parent 969914c commit 16eee6b
Showing 1 changed file with 30 additions and 33 deletions.
63 changes: 30 additions & 33 deletions compass/ocean/tests/drying_slope/viz/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import os
import xarray
import numpy
import subprocess

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import subprocess
import xarray

from compass.step import Step

Expand Down Expand Up @@ -62,9 +63,7 @@ def run(self):
Run this step of the test case
"""
section = self.config['paths']
datapath = section.get('ocean_database_root')
section = self.config['vertical_grid']
vert_levels = section.get('vert_levels')
section = self.config['drying_slope_viz']
generate_movie = section.getboolean('generate_movie')

Expand All @@ -84,7 +83,7 @@ def run(self):
outFolder=outFolder, extension=movie_format)

def _forcing(self, t):
ssh = 10. * numpy.sin(t * numpy.pi / 12.) - 10.
ssh = 10. * np.sin(t * np.pi / 12.) - 10.
return ssh

def _plot_ssh_time_series(self, outFolder='.'):
Expand All @@ -95,11 +94,10 @@ def _plot_ssh_time_series(self, outFolder='.'):
(2013) test case.
"""
colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'}
xSsh = numpy.linspace(0, 12.0, 100)
ySsh = 10.0*numpy.sin(xSsh*numpy.pi/12.0) - 10.0
xSsh = np.linspace(0, 12.0, 100)
ySsh = 10.0 * np.sin(xSsh * np.pi / 12.0) - 10.0

figsize = [6.4, 4.8]
markersize = 20

damping_coeffs = self.damping_coeffs
if damping_coeffs is None:
Expand All @@ -112,11 +110,10 @@ def _plot_ssh_time_series(self, outFolder='.'):
fig, _ = plt.subplots(nrows=naxes, ncols=1, figsize=figsize, dpi=100)

for i in range(naxes):
ax = plt.subplot(naxes, 1, i+1)
ax = plt.subplot(naxes, 1, i + 1)
ds = xarray.open_dataset(ncFilename[i])
ssh = ds.ssh
ympas = ds.ssh.where(ds.tidalInputMask).mean('nCells').values
xmpas = numpy.linspace(0, 1.0, len(ds.xtime))*12.0
xmpas = np.linspace(0, 1.0, len(ds.xtime)) * 12.0
ax.plot(xmpas, ympas, marker='o', label='MPAS-O forward',
color=colors['MPAS-O'])
ax.plot(xSsh, ySsh, lw=3, label='analytical',
Expand All @@ -140,7 +137,7 @@ def _plot_ssh_validation(self, outFolder='.'):
colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'}

locs = [7.2, 2.2, 0.2, 1.2, 4.2, 9.3]
locs = 0.92 - numpy.divide(locs, 11.)
locs = 0.92 - np.divide(locs, 11.)

damping_coeffs = self.damping_coeffs
times = self.times
Expand All @@ -156,16 +153,16 @@ def _plot_ssh_validation(self, outFolder='.'):
ncFilename = [f'output_{damping_coeff}.nc'
for damping_coeff in damping_coeffs]

xBed = numpy.linspace(0, 25, 100)
yBed = 10.0/25.0*xBed
xBed = np.linspace(0, 25, 100)
yBed = 10.0 / 25.0 * xBed

fig, _ = plt.subplots(nrows=naxes, ncols=1, sharex=True)

for i in range(naxes):
ax = plt.subplot(naxes, 1, i+1)
ax = plt.subplot(naxes, 1, i + 1)
ds = xarray.open_dataset(ncFilename[i])
ds = ds.drop_vars(numpy.setdiff1d([j for j in ds.variables],
['yCell', 'ssh']))
ds = ds.drop_vars(np.setdiff1d([j for j in ds.variables],
['yCell', 'ssh']))

ax.plot(xBed, yBed, '-k', lw=3)
ax.set_xlim(0, 25)
Expand All @@ -180,13 +177,13 @@ def _plot_ssh_validation(self, outFolder='.'):
# Plot MPAS-O data
# factor of 1e- needed to account for annoying round-off issue
# to get right time slices
plottime = int((float(atime)/0.2 + 1e-16)*24.0)
plottime = int((float(atime) / 0.2 + 1e-16) * 24.0)
ymean = ds.isel(Time=plottime).groupby('yCell').mean(
dim=xarray.ALL_DIMS)
x = ymean.yCell.values/1000.0
dim=xarray.ALL_DIMS)
x = ymean.yCell.values / 1000.0
y = ymean.ssh.values

mpas = ax.plot(x, -y, label='MPAS-O', color=colors['MPAS-O'])
ax.plot(x, -y, label='MPAS-O', color=colors['MPAS-O'])
ax.text(1, ay, atime + ' days', size=8,
transform=ax.transAxes)
if damping_coeffs is not None:
Expand Down Expand Up @@ -225,7 +222,7 @@ def _plot_ssh_validation_for_movie(self, outFolder='.'):
colors = {'MPAS-O': 'k', 'analytical': 'b', 'ROMS': 'g'}

locs = [7.2, 2.2, 0.2, 1.2, 4.2, 9.3]
locs = 0.92 - numpy.divide(locs, 11.)
locs = 0.92 - np.divide(locs, 11.)

damping_coeffs = self.damping_coeffs
if damping_coeffs is None:
Expand All @@ -241,32 +238,32 @@ def _plot_ssh_validation_for_movie(self, outFolder='.'):
times = self.times
datatypes = self.datatypes

xBed = numpy.linspace(0, 25, 100)
yBed = 10.0/25.0*xBed
xBed = np.linspace(0, 25, 100)
yBed = 10.0 / 25.0 * xBed

ii = 0
# Plot profiles over the 12h simulation duration
for itime in numpy.linspace(0, 0.5, 5*12+1):
for itime in np.linspace(0, 0.5, 5 * 12 + 1):

plottime = int((float(itime)/0.2 + 1e-16)*24.0)
plottime = int((float(itime) / 0.2 + 1e-16) * 24.0)

fig, _ = plt.subplots(nrows=naxes, ncols=1, sharex=True)

for i in range(naxes):
ax = plt.subplot(naxes, 1, i+1)
ax = plt.subplot(naxes, 1, i + 1)
ds = xarray.open_dataset(ncFilename[i])
ds = ds.drop_vars(numpy.setdiff1d([j for j in ds.variables],
['yCell', 'ssh']))
ds = ds.drop_vars(np.setdiff1d([j for j in ds.variables],
['yCell', 'ssh']))

# Plot MPAS-O snapshots
# factor of 1e- needed to account for annoying round-off issue
# to get right time slices
ymean = ds.isel(Time=plottime).groupby('yCell').mean(
dim=xarray.ALL_DIMS)
x = ymean.yCell.values/1000.0
dim=xarray.ALL_DIMS)
x = ymean.yCell.values / 1000.0
y = ymean.ssh.values
ax.plot(xBed, yBed, '-k', lw=3)
mpas = ax.plot(x, -y, label='MPAS-O', color=colors['MPAS-O'])
ax.plot(x, -y, label='MPAS-O', color=colors['MPAS-O'])

ax.set_ylim(-1, 11)
ax.set_xlim(0, 25)
Expand Down

0 comments on commit 16eee6b

Please sign in to comment.