Exemplo n.º 1
0
def yzplot(field, y=None, z=None, ylabel=None, yunits=None, zlabel=None, zunits=None,
  splitscale=None, title='', suptitle='', clim=None, colormap=None, extend=None, centerlabels=False,
  nbins=None, landcolor=[.5,.5,.5], show_stats=2, aspect=[16,9], resolution=576, axis=None,
  ignore=None, debug=False, cbar=True, cbar_hor=False, cbar_label=''):
      # Create coordinates if not provided
  ylabel, yunits, zlabel, zunits = createYZlabels(y, z, ylabel, yunits, zlabel, zunits)
  if debug: print('y,z label/units=',ylabel,yunits,zlabel,zunits)
  if len(y)==z.shape[-1]: y = expand(y)
  elif len(y)==z.shape[-1]+1: y = y
  else: raise Exception('Length of y coordinate should be equal or 1 longer than horizontal length of z')
  if ignore is not None: maskedField = np.ma.masked_array(field, mask=[field==ignore])
  else: maskedField = field.copy()
  yCoord, zCoord, field2 = section2quadmesh(y, z, maskedField)

  # Diagnose statistics
  sMin, sMax, sMean, sStd, sRMS = myStats(maskedField, yzWeight(y, z), debug=debug)
  yLims = np.amin(yCoord), np.amax(yCoord)
  zLims = boundaryStats(zCoord)

  # Choose colormap
  if nbins is None and (clim is None or len(clim)==2): nbins=35
  if colormap is None: colormap = chooseColorMap(sMin, sMax)
  cmap, norm, extend = chooseColorLevels(sMin, sMax, colormap, clim=clim, nbins=nbins, extend=extend)

  if axis is None:
    setFigureSize(aspect, resolution, debug=debug)
    axis = plt.gca()

  cs = axis.pcolormesh(yCoord, zCoord, field2, cmap=cmap, norm=norm)
  if cbar:
    if cbar_hor:  
      cb = plt.colorbar(cs,ax=axis, fraction=.08, pad=0.02, extend=extend, orientation='horizontal')
    else:
      cb = plt.colorbar(cs,ax=axis, fraction=.08, pad=0.02, extend=extend)
    if len(cbar_label)>0:
      cb.set_label(cbar_label)
    
  if centerlabels and len(clim)>2: cb.set_ticks(  0.5*(clim[:-1]+clim[1:]) )
  axis.set_facecolor(landcolor)
  if splitscale is not None:
    for zzz in splitscale[1:-1]: axis.axhline(zzz,color='k',linestyle='--')
    axis.set_yscale('splitscale', zval=splitscale)
  axis.set_xlim( yLims ); axis.set_ylim( zLims )
  if show_stats > 0:
    axis.annotate('max=%.5g\nmin=%.5g'%(sMax,sMin), xy=(0.0,1.01), xycoords='axes fraction', verticalalignment='bottom', fontsize=10)
    if show_stats > 1:
      if sMean is not None:
        axis.annotate('mean=%.5g\nrms=%.5g'%(sMean,sRMS), xy=(1.0,1.01), xycoords='axes fraction', verticalalignment='bottom', horizontalalignment='right', fontsize=10)
        axis.annotate(' sd=%.5g\n'%(sStd), xy=(1.0,1.01), xycoords='axes fraction', verticalalignment='bottom', horizontalalignment='left', fontsize=10)

  if len(ylabel+yunits)>0: axis.set_xlabel(label(ylabel, yunits))
  if len(zlabel+zunits)>0: axis.set_ylabel(label(zlabel, zunits))
  if len(title)>0: axis.set_title(title)
  if len(suptitle)>0: plt.suptitle(suptitle)
  return cs
Exemplo n.º 2
0
def main():
    # Get options
    args = options()

    nw = args.number_of_workers
    if not os.path.isdir('PNG/MOC'):
        print('Creating a directory to place figures (PNG/MOC)... \n')
        os.system('mkdir -p PNG/MOC')
    if not os.path.isdir('ncfiles'):
        print('Creating a directory to place figures (ncfiles)... \n')
        os.system('mkdir ncfiles')

    # Read in the yaml file
    diag_config_yml = yaml.load(open(args.diag_config_yml_path, 'r'),
                                Loader=yaml.Loader)

    # Create the case instance
    dcase = DiagsCase(diag_config_yml['Case'])
    args.case_name = dcase.casename
    args.savefigs = True
    args.outdir = 'PNG/MOC/'
    RUNDIR = dcase.get_value('RUNDIR')
    print('Run directory is:', RUNDIR)
    print('Casename is:', dcase.casename)
    print('Number of workers to be used:', nw)

    # set avg dates
    avg = diag_config_yml['Avg']
    if not args.start_date: args.start_date = avg['start_date']
    if not args.end_date: args.end_date = avg['end_date']

    # read grid info
    grd = MOM6grid(RUNDIR + '/' + dcase.casename + '.mom6.static.nc')
    depth = grd.depth_ocean
    # remote Nan's, otherwise genBasinMasks won't work
    depth[np.isnan(depth)] = 0.0
    basin_code = m6toolbox.genBasinMasks(grd.geolon, grd.geolat, depth)

    parallel, cluster, client = m6toolbox.request_workers(nw)

    print('Reading {} dataset...'.format(args.file_name))
    startTime = datetime.now()

    # load data
    def preprocess(ds):
        variables = ['vmo', 'vhml', 'vhGM']
        for v in variables:
            if v not in ds.variables:
                ds[v] = xr.zeros_like(ds.vo)
        return ds[variables]

    if parallel:
        ds = xr.open_mfdataset(
            RUNDIR + '/' + dcase.casename + args.file_name,
            parallel=True,
            combine="nested",  # concatenate in order of files
            concat_dim="time",  # concatenate along time
            preprocess=preprocess,
        ).chunk({"time": 12})

    else:
        ds = xr.open_mfdataset(RUNDIR+'/'+dcase.casename+args.file_name, data_vars='minimal', \
                               coords='minimal', compat='override', preprocess=preprocess)

    print('Time elasped: ', datetime.now() - startTime)

    # compute yearly means first since this will be used in the time series
    print('Computing yearly means...')
    startTime = datetime.now()
    ds_yr = ds.resample(time="1Y", closed='left').mean('time')
    print('Time elasped: ', datetime.now() - startTime)

    print('Selecting data between {} and {}...'.format(args.start_date,
                                                       args.end_date))
    startTime = datetime.now()
    ds_sel = ds_yr.sel(time=slice(args.start_date, args.end_date))
    print('Time elasped: ', datetime.now() - startTime)

    print('Computing time mean...')
    startTime = datetime.now()
    ds_mean = ds_sel.mean('time').compute()
    print('Time elasped: ', datetime.now() - startTime)

    # create a ndarray subclass
    class C(np.ndarray):
        pass

    varName = 'vmo'
    conversion_factor = 1.e-9
    tmp = np.ma.masked_invalid(ds_mean[varName].values)
    tmp = tmp[:].filled(0.)
    VHmod = tmp.view(C)
    VHmod.units = ds[varName].units
    Zmod = m6toolbox.get_z(ds, depth, varName)  # same here

    if args.case_name != '': case_name = args.case_name
    else: case_name = ''

    # Global MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    axis = plt.gca()
    cmap = plt.get_cmap('dunnePM')
    zg = Zmod.min(axis=-1)
    psiPlot = MOCpsi(VHmod) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yyg = grd.geolat_c[:, :].max(axis=-1) + 0 * zg
    ci = m6plot.pmCI(0., 40., 5.)
    plotPsi(
        yyg, zg, psiPlot, ci, 'Global MOC [Sv],' + 'averaged between ' +
        args.start_date + ' and ' + args.end_date)
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    findExtrema(yyg, zg, psiPlot, max_lat=-30.)
    findExtrema(yyg, zg, psiPlot, min_lat=25., min_depth=250.)
    findExtrema(yyg, zg, psiPlot, min_depth=2000., mult=-1.)
    objOut = args.outdir + str(case_name) + '_MOC_global.png'
    plt.savefig(objOut)

    if 'zl' in ds:
        zl = ds.zl.values
    elif 'z_l' in ds:
        zl = ds.z_l.values
    else:
        raise ValueError("Dataset does not have vertical coordinate zl or z_l")

    # create dataset to store results
    moc = xr.Dataset(data_vars={
        'moc': (('zl', 'yq'), psiPlot),
        'amoc': (('zl', 'yq'), np.zeros((psiPlot.shape))),
        'moc_FFM': (('zl', 'yq'), np.zeros((psiPlot.shape))),
        'moc_GM': (('zl', 'yq'), np.zeros((psiPlot.shape))),
        'amoc_45': (('time'), np.zeros((ds_yr.time.shape))),
        'moc_GM_ACC': (('time'), np.zeros((ds_yr.time.shape))),
        'amoc_26': (('time'), np.zeros((ds_yr.time.shape)))
    },
                     coords={
                         'zl': zl,
                         'yq': ds.yq,
                         'time': ds_yr.time
                     })
    attrs = {
        'description': 'MOC time-mean sections and time-series',
        'units': 'Sv',
        'start_date': avg['start_date'],
        'end_date': avg['end_date'],
        'casename': dcase.casename
    }
    m6toolbox.add_global_attrs(moc, attrs)

    # Atlantic MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    cmap = plt.get_cmap('dunnePM')
    m = 0 * basin_code
    m[(basin_code == 2) | (basin_code == 4) | (basin_code == 6) |
      (basin_code == 7) | (basin_code == 8)] = 1
    ci = m6plot.pmCI(0., 22., 2.)
    z = (m * Zmod).min(axis=-1)
    psiPlot = MOCpsi(VHmod,
                     vmsk=m * np.roll(m, -1, axis=-2)) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    plotPsi(
        yy, z, psiPlot, ci, 'Atlantic MOC [Sv],' + 'averaged between ' +
        args.start_date + ' and ' + args.end_date)
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    findExtrema(yy, z, psiPlot, min_lat=26.5, max_lat=27.,
                min_depth=250.)  # RAPID
    findExtrema(yy, z, psiPlot, max_lat=-33.)
    findExtrema(yy, z, psiPlot)
    findExtrema(yy, z, psiPlot, min_lat=5.)
    objOut = args.outdir + str(case_name) + '_MOC_Atlantic.png'
    plt.savefig(objOut, format='png')
    moc['amoc'].data = psiPlot

    print('Plotting AMOC profile at 26N...')
    rapid_vertical = xr.open_dataset(
        '/glade/work/gmarques/cesm/datasets/RAPID/moc_vertical.nc')
    fig, ax = plt.subplots(nrows=1, ncols=1)
    ax.plot(rapid_vertical.stream_function_mar.mean('time'),
            rapid_vertical.depth,
            'k',
            label='RAPID')
    ax.plot(moc['amoc'].sel(yq=26, method='nearest'), zl, label=case_name)
    ax.legend()
    plt.gca().invert_yaxis()
    plt.grid()
    ax.set_xlabel('AMOC @ 26N [Sv]')
    ax.set_ylabel('Depth [m]')
    objOut = args.outdir + str(case_name) + '_MOC_profile_26N.png'
    plt.savefig(objOut, format='png')

    print('Computing time series...')
    startTime = datetime.now()
    # time-series
    dtime = ds_yr.time
    amoc_26 = np.zeros(len(dtime))
    amoc_45 = np.zeros(len(dtime))
    moc_GM_ACC = np.zeros(len(dtime))
    if args.debug: startTime = datetime.now()
    # loop in time
    for t in range(len(dtime)):
        tmp = np.ma.masked_invalid(ds_yr[varName][t, :].values)
        tmp = tmp[:].filled(0.)
        # m is still Atlantic ocean
        psi = MOCpsi(tmp, vmsk=m * np.roll(m, -1, axis=-2)) * conversion_factor
        psi = 0.5 * (psi[0:-1, :] + psi[1::, :])
        amoc_26[t] = findExtrema(yy,
                                 z,
                                 psi,
                                 min_lat=26.,
                                 max_lat=27.,
                                 plot=False,
                                 min_depth=250.)
        amoc_45[t] = findExtrema(yy,
                                 z,
                                 psi,
                                 min_lat=44.,
                                 max_lat=46.,
                                 plot=False,
                                 min_depth=250.)
        tmp_GM = np.ma.masked_invalid(ds_yr['vhGM'][t, :].values)
        tmp_GM = tmp_GM[:].filled(0.)
        psiGM = MOCpsi(tmp_GM) * conversion_factor
        psiGM = 0.5 * (psiGM[0:-1, :] + psiGM[1::, :])
        moc_GM_ACC[t] = findExtrema(yyg,
                                    zg,
                                    psiGM,
                                    min_lat=-65.,
                                    max_lat=-30,
                                    mult=-1.,
                                    plot=False)
    print('Time elasped: ', datetime.now() - startTime)

    # add dataarays to the moc dataset
    moc['amoc_26'].data = amoc_26
    moc['amoc_45'].data = amoc_45
    moc['moc_GM_ACC'].data = moc_GM_ACC

    if parallel:
        print('Releasing workers ...')
        client.close()
        cluster.close()

    print('Plotting...')
    # load AMOC time series data (5th) cycle used in Danabasoglu et al., doi:10.1016/j.ocemod.2015.11.007
    path = '/glade/p/cesm/omwg/amoc/COREII_AMOC_papers/papers/COREII.variability/data.original/'
    amoc_core_26 = xr.open_dataset(path + 'AMOCts.cyc5.26p5.nc')
    # load AMOC from POP JRA-55
    amoc_pop_26 = xr.open_dataset(
        '/glade/u/home/bryan/MOM6-modeloutputanalysis/'
        'AMOC_series_26n.g210.GIAF_JRA.v13.gx1v7.01.nc')
    # load RAPID time series
    rapid = xr.open_dataset(
        '/glade/work/gmarques/cesm/datasets/RAPID/moc_transports.nc').resample(
            time="1Y", closed='left', keep_attrs=True).mean('time',
                                                            keep_attrs=True)
    # plot
    fig = plt.figure(figsize=(12, 6))
    plt.plot(np.arange(len(moc.time)) + 1958.5,
             moc['amoc_26'].values,
             color='k',
             label=case_name,
             lw=2)
    # core data
    core_mean = amoc_core_26['MOC'].mean(axis=0).data
    core_std = amoc_core_26['MOC'].std(axis=0).data
    plt.plot(amoc_core_26.time,
             core_mean,
             'k',
             label='CORE II (group mean)',
             color='#1B2ACC',
             lw=1)
    plt.fill_between(amoc_core_26.time,
                     core_mean - core_std,
                     core_mean + core_std,
                     alpha=0.25,
                     edgecolor='#1B2ACC',
                     facecolor='#089FFF')
    # pop data
    plt.plot(np.arange(len(amoc_pop_26.time)) + 1958.5,
             amoc_pop_26.AMOC_26n.values,
             color='r',
             label='POP',
             lw=1)
    # rapid
    plt.plot(np.arange(len(rapid.time)) + 2004.5,
             rapid.moc_mar_hc10.values,
             color='green',
             label='RAPID',
             lw=1)

    plt.title('AMOC @ 26 $^o$ N', fontsize=16)
    plt.ylim(5, 20)
    plt.xlim(1948, 1958.5 + len(moc.time))
    plt.xlabel('Time [years]', fontsize=16)
    plt.ylabel('Sv', fontsize=16)
    plt.legend(fontsize=13, ncol=2)
    objOut = args.outdir + str(case_name) + '_MOC_26N_time_series.png'
    plt.savefig(objOut, format='png')

    amoc_core_45 = xr.open_dataset(path + 'AMOCts.cyc5.45.nc')
    amoc_pop_45 = xr.open_dataset(
        '/glade/u/home/bryan/MOM6-modeloutputanalysis/'
        'AMOC_series_45n.g210.GIAF_JRA.v13.gx1v7.01.nc')
    # plot
    fig = plt.figure(figsize=(12, 6))
    plt.plot(np.arange(len(moc.time)) + 1958.5,
             moc['amoc_45'],
             color='k',
             label=case_name,
             lw=2)
    # core data
    core_mean = amoc_core_45['MOC'].mean(axis=0).data
    core_std = amoc_core_45['MOC'].std(axis=0).data
    plt.plot(amoc_core_45.time,
             core_mean,
             'k',
             label='CORE II (group mean)',
             color='#1B2ACC',
             lw=2)
    plt.fill_between(amoc_core_45.time,
                     core_mean - core_std,
                     core_mean + core_std,
                     alpha=0.25,
                     edgecolor='#1B2ACC',
                     facecolor='#089FFF')
    # pop data
    plt.plot(np.arange(len(amoc_pop_45.time)) + 1958.5,
             amoc_pop_45.AMOC_45n.values,
             color='r',
             label='POP',
             lw=1)

    plt.title('AMOC @ 45 $^o$ N', fontsize=16)
    plt.ylim(5, 20)
    plt.xlim(1948, 1958 + len(moc.time))
    plt.xlabel('Time [years]', fontsize=16)
    plt.ylabel('Sv', fontsize=16)
    plt.legend(fontsize=14)
    objOut = args.outdir + str(case_name) + '_MOC_45N_time_series.png'
    plt.savefig(objOut, format='png')

    # Submesoscale-induced Global MOC
    class C(np.ndarray):
        pass

    varName = 'vhml'
    conversion_factor = 1.e-9
    tmp = np.ma.masked_invalid(ds_mean[varName].values)
    tmp = tmp[:].filled(0.)
    VHml = tmp.view(C)
    VHml.units = ds[varName].units
    Zmod = m6toolbox.get_z(ds, depth, varName)  # same here
    m6plot.setFigureSize([16, 9], 576, debug=False)
    axis = plt.gca()
    cmap = plt.get_cmap('dunnePM')
    z = Zmod.min(axis=-1)
    psiPlot = MOCpsi(VHml) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    ci = m6plot.pmCI(0., 20., 2.)
    plotPsi(yy,
            z,
            psiPlot,
            ci,
            'Global FFH MOC [Sv],' + 'averaged between ' + args.start_date +
            ' and ' + args.end_date,
            zval=[0., -400., -1000.])
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    objOut = args.outdir + str(case_name) + '_FFH_MOC_global.png'
    plt.savefig(objOut)
    moc['moc_FFM'].data = psiPlot

    # GM-induced Global MOC
    class C(np.ndarray):
        pass

    varName = 'vhGM'
    conversion_factor = 1.e-9
    tmp = np.ma.masked_invalid(ds_mean[varName].values)
    tmp = tmp[:].filled(0.)
    VHGM = tmp.view(C)
    VHGM.units = ds[varName].units
    Zmod = m6toolbox.get_z(ds, depth, varName)  # same here
    m6plot.setFigureSize([16, 9], 576, debug=False)
    axis = plt.gca()
    cmap = plt.get_cmap('dunnePM')
    z = Zmod.min(axis=-1)
    psiPlot = MOCpsi(VHGM) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    ci = m6plot.pmCI(0., 20., 2.)
    plotPsi(
        yy, z, psiPlot, ci, 'Global GM MOC [Sv],' + 'averaged between ' +
        args.start_date + ' and ' + args.end_date)
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    findExtrema(yy, z, psiPlot, min_lat=-65., max_lat=-30, mult=-1.)
    objOut = args.outdir + str(case_name) + '_GM_MOC_global.png'
    plt.savefig(objOut)
    moc['moc_GM'].data = psiPlot

    print('Saving netCDF files...')
    moc.to_netcdf('ncfiles/' + str(case_name) + '_MOC.nc')
    return
Exemplo n.º 3
0
def main(stream=False):
    # Get options
    args = options()
    # mom6 grid
    grd = MOM6grid(args.infile + args.static)
    depth = grd.depth_ocean
    # remote Nan's, otherwise genBasinMasks won't work
    depth[numpy.isnan(depth)] = 0.0
    basin_code = m6toolbox.genBasinMasks(grd.geolon, grd.geolat, depth)

    # load data
    ds = xr.open_mfdataset(args.infile + args.monthly, decode_times=False)
    # convert time in years
    ds['time'] = ds.time / 365.
    ti = args.year_start
    tf = args.year_end
    # check if data includes years between ti and tf
    m6toolbox.check_time_interval(ti, tf, ds)

    # create a ndarray subclass
    class C(numpy.ndarray):
        pass

    if 'vmo' in ds.variables:
        varName = 'vmo'
        conversion_factor = 1.e-9
    elif 'vh' in ds.variables:
        varName = 'vh'
        conversion_factor = 1.e-6
        if 'zw' in ds.variables:
            conversion_factor = 1.e-9  # Backwards compatible for when we had wrong units for 'vh'
    else:
        raise Exception('Could not find "vh" or "vmo" in file "%s"' %
                        (args.infile + args.static))

    tmp = numpy.ma.masked_invalid(
        ds[varName].sel(time=slice(ti, tf)).mean('time').data)
    tmp = tmp[:].filled(0.)
    VHmod = tmp.view(C)
    VHmod.units = ds[varName].units

    Zmod = m6toolbox.get_z(ds, depth, varName)

    if args.case_name != '': case_name = args.case_name + ' ' + args.label
    else: case_name = ds.title + ' ' + args.label

    imgbufs = []

    # Global MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    axis = plt.gca()
    cmap = plt.get_cmap('dunnePM')
    z = Zmod.min(axis=-1)
    psiPlot = MOCpsi(VHmod) * conversion_factor
    #yy = y[1:,:].max(axis=-1)+0*z
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    print(z.shape, yy.shape, psiPlot.shape)
    ci = m6plot.pmCI(0., 40., 5.)
    plotPsi(yy, z, psiPlot[1::, :], ci, 'Global MOC [Sv]')
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    findExtrema(yy, z, psiPlot, max_lat=-30.)
    findExtrema(yy, z, psiPlot, min_lat=25.)
    findExtrema(yy, z, psiPlot, min_depth=2000., mult=-1.)
    if stream is True: objOut = io.BytesIO()
    else: objOut = args.outdir + '/MOC_global.png'
    plt.savefig(objOut)
    if stream is True: imgbufs.append(objOut)

    # Atlantic MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    cmap = plt.get_cmap('dunnePM')
    m = 0 * basin_code
    m[(basin_code == 2) | (basin_code == 4) | (basin_code == 6) |
      (basin_code == 7) | (basin_code == 8)] = 1
    ci = m6plot.pmCI(0., 22., 2.)
    z = (m * Zmod).min(axis=-1)
    psiPlot = MOCpsi(VHmod,
                     vmsk=m * numpy.roll(m, -1, axis=-2)) * conversion_factor
    #yy = y[1:,:].max(axis=-1)+0*z
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    plotPsi(yy, z, psiPlot[1::, :], ci, 'Atlantic MOC [Sv]')
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    findExtrema(yy, z, psiPlot, min_lat=26.5, max_lat=27.)  # RAPID
    findExtrema(yy, z, psiPlot, max_lat=-33.)
    findExtrema(yy, z, psiPlot)
    findExtrema(yy, z, psiPlot, min_lat=5.)
    if stream is True: objOut = io.BytesIO()
    else: objOut = args.outdir + '/MOC_Atlantic.png'
    plt.savefig(objOut, format='png')
    if stream is True: imgbufs.append(objOut)

    if stream is True:
        return imgbufs
Exemplo n.º 4
0
def main():
    # Get options
    args = options()
    # mom6 grid
    grd = MOM6grid(args.infile + args.case_name + '.mom6.static.nc')
    depth = grd.depth_ocean
    # remote Nan's, otherwise genBasinMasks won't work
    depth[numpy.isnan(depth)] = 0.0
    basin_code = m6toolbox.genBasinMasks(grd.geolon, grd.geolat, depth)

    # load data
    ds = xr.open_mfdataset(args.infile + args.case_name + '.mom6.hm_*.nc',
                           combine='by_coords')
    ti = args.start_date
    tf = args.end_date

    # create a ndarray subclass
    class C(numpy.ndarray):
        pass

    if 'vmo' in ds.variables:
        varName = 'vmo'
        conversion_factor = 1.e-9
    elif 'vh' in ds.variables:
        varName = 'vh'
        conversion_factor = 1.e-6
        if 'zw' in ds.variables:
            conversion_factor = 1.e-9  # Backwards compatible for when we had wrong units for 'vh'
    else:
        raise Exception('Could not find "vh" or "vmo" in file "%s"' %
                        (args.infile + args.static))

    # selected dates
    ds_var = ds[varName].sel(time=slice(ti, tf))

    # yearly means
    ds_var_yr = ds_var.resample(time="1Y", closed='left',
                                keep_attrs=True).mean(dim='time',
                                                      keep_attrs=True).load()

    tmp = numpy.ma.masked_invalid(ds_var_yr.mean('time').values)
    tmp = tmp[:].filled(0.)
    VHmod = tmp.view(C)
    VHmod.units = ds[varName].units

    Zmod = m6toolbox.get_z(ds, depth, varName)

    if args.case_name != '': case_name = args.case_name + ' ' + args.label
    else: case_name = ds.title + ' ' + args.label

    imgbufs = []

    # Global MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    axis = plt.gca()
    cmap = plt.get_cmap('dunnePM')
    z = Zmod.min(axis=-1)
    psiPlot = MOCpsi(VHmod) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    ci = m6plot.pmCI(0., 40., 5.)
    plotPsi(yy, z, psiPlot, ci,
            'Global MOC [Sv],' + 'averaged between ' + ti + 'and ' + tf)
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    findExtrema(yy, z, psiPlot, max_lat=-30.)
    findExtrema(yy, z, psiPlot, min_lat=25.)
    findExtrema(yy, z, psiPlot, min_depth=2000., mult=-1.)
    objOut = str(case_name) + '_MOC_global.png'
    plt.savefig(objOut)

    # Atlantic MOC
    m6plot.setFigureSize([16, 9], 576, debug=False)
    cmap = plt.get_cmap('dunnePM')
    m = 0 * basin_code
    m[(basin_code == 2) | (basin_code == 4) | (basin_code == 6) |
      (basin_code == 7) | (basin_code == 8)] = 1
    ci = m6plot.pmCI(0., 22., 2.)
    z = (m * Zmod).min(axis=-1)
    psiPlot = MOCpsi(VHmod,
                     vmsk=m * numpy.roll(m, -1, axis=-2)) * conversion_factor
    psiPlot = 0.5 * (psiPlot[0:-1, :] + psiPlot[1::, :])
    yy = grd.geolat_c[:, :].max(axis=-1) + 0 * z
    plotPsi(yy, z, psiPlot, ci,
            'Atlantic MOC [Sv],' + 'averaged between ' + ti + 'and ' + tf)
    plt.xlabel(r'Latitude [$\degree$N]')
    plt.suptitle(case_name)
    plt.gca().invert_yaxis()
    findExtrema(yy, z, psiPlot, min_lat=26.5, max_lat=27.)  # RAPID
    findExtrema(yy, z, psiPlot, max_lat=-33.)
    findExtrema(yy, z, psiPlot)
    findExtrema(yy, z, psiPlot, min_lat=5.)
    objOut = str(case_name) + '_MOC_Atlantic.png'
    plt.savefig(objOut, format='png')

    # time-series
    dtime = ds_var_yr.time.values
    amoc_26 = numpy.zeros(len(dtime))
    amoc_45 = numpy.zeros(len(dtime))

    # loop in time
    for t in range(len(dtime)):
        tmp = numpy.ma.masked_invalid(ds_var_yr.sel(time=dtime[t]).values)
        tmp = tmp[:].filled(0.)
        psi = MOCpsi(tmp,
                     vmsk=m * numpy.roll(m, -1, axis=-2)) * conversion_factor
        psi = 0.5 * (psi[0:-1, :] + psi[1::, :])
        amoc_26[t] = findExtrema(yy,
                                 z,
                                 psi,
                                 min_lat=26.5,
                                 max_lat=27.,
                                 plot=False)
        amoc_45[t] = findExtrema(yy,
                                 z,
                                 psi,
                                 min_lat=44.,
                                 max_lat=46.,
                                 plot=False)

    # create dataarays
    amoc_26_da = xr.DataArray(amoc_26, dims=['time'], coords={'time': dtime})
    amoc_45_da = xr.DataArray(amoc_45, dims=['time'], coords={'time': dtime})

    # load AMOC time series data (5th) cycle used in Danabasoglu et al., doi:10.1016/j.ocemod.2015.11.007
    path = '/glade/p/cesm/omwg/amoc/COREII_AMOC_papers/papers/COREII.variability/data.original/'
    amoc_core_26 = xr.open_dataset(path + 'AMOCts.cyc5.26p5.nc')
    # plot
    fig = plt.figure(figsize=(12, 6))
    plt.plot(numpy.arange(len(amoc_26_da.time)) + 1948.5,
             amoc_26_da.values,
             color='k',
             label=case_name,
             lw=2)
    # core data
    core_mean = amoc_core_26['MOC'].mean(axis=0).data
    core_std = amoc_core_26['MOC'].std(axis=0).data
    plt.plot(amoc_core_26.time,
             core_mean,
             'k',
             label='CORE II (group mean)',
             color='#1B2ACC',
             lw=2)
    plt.fill_between(amoc_core_26.time,
                     core_mean - core_std,
                     core_mean + core_std,
                     alpha=0.25,
                     edgecolor='#1B2ACC',
                     facecolor='#089FFF')
    plt.title('AMOC @ 26 $^o$ N', fontsize=16)
    plt.xlabel('Time [years]', fontsize=16)
    plt.ylabel('Sv', fontsize=16)
    plt.legend(fontsize=14)
    objOut = str(case_name) + '_MOC_26N_time_series.png'
    plt.savefig(objOut, format='png')

    amoc_core_45 = xr.open_dataset(path + 'AMOCts.cyc5.45.nc')
    # plot
    fig = plt.figure(figsize=(12, 6))
    plt.plot(numpy.arange(len(amoc_45_da.time)) + 1948.5,
             amoc_45_da.values,
             color='k',
             label=case_name,
             lw=2)
    # core data
    core_mean = amoc_core_45['MOC'].mean(axis=0).data
    core_std = amoc_core_45['MOC'].std(axis=0).data
    plt.plot(amoc_core_45.time,
             core_mean,
             'k',
             label='CORE II (group mean)',
             color='#1B2ACC',
             lw=2)
    plt.fill_between(amoc_core_45.time,
                     core_mean - core_std,
                     core_mean + core_std,
                     alpha=0.25,
                     edgecolor='#1B2ACC',
                     facecolor='#089FFF')
    plt.title('AMOC @ 45 $^o$ N', fontsize=16)
    plt.xlabel('Time [years]', fontsize=16)
    plt.ylabel('Sv', fontsize=16)
    plt.legend(fontsize=14)
    objOut = str(case_name) + '_MOC_45N_time_series.png'
    plt.savefig(objOut, format='png')
    return