test_bug_horizontal.py

import netCDF4 as nc
import numpy as np
from scipy.signal import resample
import seapy
import inspect
import os
import matplotlib.pyplot as plt

this_folder = os.path.dirname(inspect.getfile(inspect.currentframe()))

output_folder = this_folder+'/outputs/'

gridfiles = {}
gridfiles[
    'WC13'] = this_folder + '/wc13_grd.nc'

def genfile_fake_sfc_obs_wc13(i0=20,
                              i1=23,
                              dx=0.0,
                              outfile='sobs_fake_sfc_wc13_basecase.nc'):
    grid = gridfiles['WC13']
    dt = 3600.0

    wcgrd = nc.Dataset(gridfiles['WC13'])
    olon = wcgrd['lon_rho'][i0, i0]
    olat = wcgrd['lat_rho'][i0, i0]

    olon2 = wcgrd['lon_rho'][i1, i1]
    olat2 = wcgrd['lat_rho'][i1, i1]

    dlon = [olon - dx, olon - dx, olon + dx, olon + dx]
    dlat = [olat - dx, olat - dx, olat + dx, olat + dx]
    dtemp = np.array([22.2, 22.2, 21.8, 21.8])

    qlon = [olon2 - dx, olon2 + dx, olon2, olon2]
    qlat = [olat2, olat2, olat2 + dx, olat2 - dx]
    qtemp = dtemp
    lon = np.ma.masked_invalid(np.hstack([dlon, qlon]))
    lat = np.ma.masked_invalid(np.hstack([dlat, qlat]))
    temp = np.ma.masked_invalid(np.hstack([dtemp, qtemp]))

    data = [
        seapy.roms.obs.raw_data("TEMP", "SST_OSTIA", temp.compressed(),
                                temp * 0 + 0.001, 0.01)
    ]
    obs = seapy.roms.obs.gridder(grid, 1.0, lon, lat, None, data, dt,
                                 "OSTIA SST Obs")
    obs.to_netcdf(output_folder + outfile)
    return i0, i1, lon, lat, temp, outfile


def genfile_line_wc13(outfile='sobs_fake_sfc_wc13_linecase.nc',
                      nx=1000,
                      ny=1000):
    
    grid = gridfiles['WC13']
    dt = 3600.0
    i0 = 10
    i1 = 31
    gobj = nc.Dataset(grid)
    lon = gobj['lon_rho'][20, 10:30]
    lat = gobj['lat_rho'][20, 10:30]

    llon = np.ma.masked_invalid(resample([lon[0], lon[-1]], nx))
    llat = np.ma.masked_invalid(resample([lat[0], lat[-1]], ny))
    temp = llon * 0 + 22
    data = [
        seapy.roms.obs.raw_data("TEMP", "SST_OSTIA", temp.compressed(),
                                temp * 0 + 0.001, 0.01)
    ]
    obs = seapy.roms.obs.gridder(grid, 1.0, llon, llat, None, data, dt,
                                 "OSTIA SST Obs")
    obs.to_netcdf(output_folder + outfile)
    return i0, i1, llon, llat, temp, outfile


def genfile_grid_wc13(outfile='sobs_fake_sfc_wc13_gridcase.nc',
                      nx=1000,
                      ny=1000,
                      dy=0.14):
    from scipy.signal import resample
    grid = gridfiles['WC13']
    dt = 3600.0
    i0 = 10
    i1 = 31
    gobj = nc.Dataset(grid)
    lon = gobj['lon_rho'][20, 10:30]
    lat = gobj['lat_rho'][20, 10:30]

    llon = np.ma.masked_invalid(resample([lon[0], lon[-1]], nx))
    llat = np.ma.masked_invalid(resample([lat[0], lat[-1]], ny))
    llon = np.hstack([llon, llon, llon])
    llat = np.hstack([llat - dy, llat, llat + dy])
    #    llon = np.hstack([lon + 0.1, lon - 0.1])
    #    llat = np.hstack([lat + 0.1, lat - 0.1])
    temp = llon * 0 + 22
    data = [
        seapy.roms.obs.raw_data("TEMP", "SST_OSTIA", temp.compressed(),
                                temp * 0 + 0.001, 0.01)
    ]
    obs = seapy.roms.obs.gridder(grid, 1.0, llon, llat, None, data, dt,
                                 "OSTIA SST Obs")
    obs.to_netcdf(output_folder + outfile)
    return i0, i1, llon, llat, temp, outfile


def plot_fake_sfc_obs_wc13(i0, i1, lon, lat, temp, outfile, edges='psi'):
    grid = nc.Dataset(gridfiles['WC13'])

    sobs = nc.Dataset(output_folder + outfile)

    fig, ax = plt.subplots()
    plt.pcolor(
        grid['lon_'+edges][:],
        grid['lat_'+edges][:],
        grid['lon_'+edges][:] * 0,
        facecolor='none',
        edgecolor='k',
        linewidth=1.,
        zorder=9999)
    plt.plot(grid['lon_rho'][:], grid['lat_rho'][:], 'k*')
    plt.plot(lon, lat, 'b.', markersize=7)
    plt.plot(sobs['obs_lon'][:], sobs['obs_lat'][:], 'r.', markersize=3)
    n_obs = len(lon)
    n_sobs = len(sobs['obs_lon'][:])
    tstr = 'Total rawobs %d, total sobs %d' % (n_obs, n_sobs)
    plt.title(tstr)
    di = (i1 - i0) * 2
    plt.axis([
        grid['lon_rho'][i0 - di, i0 - di], grid['lon_rho'][i0 + di, i0 + di],
        grid['lat_rho'][i0 - di, i0 - di], grid['lat_rho'][i0 + di, i0 + di]
    ])

def test_wc13_basecase():
    plot_fake_sfc_obs_wc13(*genfile_fake_sfc_obs_wc13())
    plt.savefig('sobs_fake_sfc_wc13_basecase.png')


def test_wc13_standardcase():
    plot_fake_sfc_obs_wc13(*genfile_fake_sfc_obs_wc13(
        dx=0.1, outfile='sobs_fake_sfc_wc13_standardcase.nc'))
    plt.savefig('sobs_fake_sfc_wc13_standardcase.png')

def test_wc13_cornercase():
    plot_fake_sfc_obs_wc13(*genfile_fake_sfc_obs_wc13(
        dx=0.32, outfile='sobs_fake_sfc_wc13_cornercase.nc'))
    plt.savefig('sobs_fake_sfc_wc13_cornercase.png')


def test_wc13_beyondcase():
    plot_fake_sfc_obs_wc13(*genfile_fake_sfc_obs_wc13(
        dx=-0.16, outfile='sobs_fake_sfc_wc13_beyondcase.nc'),edges='rho')
    plt.savefig('sobs_fake_sfc_wc13_beyondcase_asrho.png')


def test_wc13_linecase():
    plot_fake_sfc_obs_wc13(
        *genfile_line_wc13(outfile='sobs_fake_sfc_wc13_linecase.nc'))
    ax = plt.gca()
    ax.axis([-131, -124, 36.0, 37.5])
    plt.savefig('sobs_fake_sfc_wc13_linecase.png')


def test_wc13_gridcase():
    plot_fake_sfc_obs_wc13(
        *genfile_grid_wc13(outfile='sobs_fake_sfc_wc13_gridcase.nc'))
    ax = plt.gca()
    ax.axis([-131, -124, 36.0, 37.5])
    plt.savefig('sobs_fake_sfc_wc13_gridcase.png')

def test_wc13_gridcase_asrho():
    plot_fake_sfc_obs_wc13(
        *genfile_grid_wc13(outfile='sobs_fake_sfc_wc13_gridcase.nc'), edges='rho')
    ax = plt.gca()
    ax.axis([-131, -124, 36.0, 37.5])
    plt.savefig('sobs_fake_sfc_wc13_gridcase_rhoascorners.png')