Exemple #1
0
def main(inputfile, outputfile, **kwargs):

    # output figure type and name
    fig_name = os.path.splitext(outputfile)[0]
    fig_type = os.path.splitext(outputfile)[1][1:]

    # create the plot workstation
    wks = ngl.open_wks(fig_type, fig_name)

    # Increase maximum buffer memory
    increase_workspace_memory(value=10000000000)

    # Read image data
    ds = xr.open_dataset(inputfile)
    band1 = ds['Band1'].values
    band2 = ds['Band2'].values
    band3 = ds['Band3'].values

    # Plot image
    # mpCenterLonF=150 for Pacific; 340 for Atlantic; 200 Pacific v2
    # mpCenterLatF=30 for Pacific; 45 for Atlantic; 40 Pacific v2
    pl = plot_image_gdal(wks, band1, band2, band3, mpCenterLonF=340, mpCenterLatF=45)

    # Finalize plot
    ngl.frame(wks)
    ngl.end()
Exemple #2
0
def main(test_files,
         cntl_files,
         vname,
         fig_name,
         test_name='Model',
         cntl_name='Obs',
         **kwargs):

    # Load datasets
    datasets = [open_dataset(f) for f in (test_files, cntl_files)]
    data_arrays = [get_data(ds, vname) for ds in datasets]
    coords = [get_coords(ds) for ds in datasets]
    weights = [get_area_weights(ds) for ds in datasets]

    # Compute area averages
    means = [
        area_average(da, wgt, dims=[d for d in da.dims if d != 'time'])
        for (da, wgt) in zip(data_arrays, weights)
    ]

    # Make line plots of area averages
    times = [da['time'] for da in data_arrays]
    labels = (test_name, cntl_name)
    plot_format = fig_name.split('.')[-1]
    wks = ngl.open_wks(plot_format, fig_name)
def main(datafile, varname, gridfile=None):

    # Read data
    ds_data = xarray.open_dataset(datafile)
    data = ds_data[varname]
    if gridfile is not None:
        ds_grid = xarray.open_dataset(gridfile).rename({'grid_size': 'ncol'})
        lon = ds_grid['grid_corner_lon']
        lat = ds_grid['grid_corner_lat']
    else:
        lon = ds_data['lon']
        lat = ds_data['lat']

    # Make sure we don't have time or level dimensions
    if 'time' in data.dims:
        data = data.isel(time=0).squeeze()
    if 'lev' in data.dims:
        data = data.isel(lev=-1).squeeze()

    if 'lon' in ds_grid and 'lat' in ds_grid:
        x = ds_grid['lon']
        y = ds_grid['lat']
    elif 'grid_corner_lon' in ds_grid and 'grid_corner_lat' in ds_grid:
        x = ds_grid['grid_corner_lon'].rename({'grid_size': 'ncol'})
        y = ds_grid['grid_corner_lat'].rename({'grid_size': 'ncol'})

    # Setup the canvas
    plot_format='png'
    plot_name='./' + varname
    wks = ngl.open_wks(plot_format, plot_name)

    # Make plot
    plot = plot_unstructured(
        wks, x.values, y.values, data.values,
        mpGeophysicalLineColor='white',
        lbOrientation='horizontal', 
        lbTitleString='%s (%s)'%(data.long_name, data.units),
        cnFillMode='RasterFill',
        cnLineLabelsOn=False, cnLinesOn=False,
    )

    # Close things
    ngl.end()
Exemple #4
0
def plot_frame(lon, lat, data, frame_name,
               lat_min=None, lat_max=None, lon_min=None, lon_max=None, 
               **kwargs):

    # Open figure
    wks = ngl.open_wks('png', os.path.splitext(frame_name)[0])

    # Plot data
    pl = plot_map(
        wks, lon, lat, data,
        #tiMainString=f'{str(data.time.values)}',
        lbOrientation='horizontal',
        #lbTitleString='%s (%s)'%(data.long_name, data.units),
        cnFillMode='RasterFill',
        **kwargs
    )

    # Save figure
    ngl.draw(pl)
    ngl.destroy(wks)

    return frame_name
Exemple #5
0
def main(varname,
         plotname,
         *datafiles,
         gridfile=None,
         time_index=None,
         vmin=None,
         vmax=None,
         **kwargs):

    # Read data
    ds_data = xarray.open_mfdataset(
        sorted(datafiles),
        chunks={'time': 1},
        drop_variables=('P3_input_dim', 'P3_output_dim'),
    )
    data = get_data(ds_data, varname)
    if gridfile is not None:
        ds_grid = xarray.open_dataset(gridfile).rename({'grid_size': 'ncol'})
        if 'lon' in ds_grid and 'lat' in ds_grid:
            x = ds_grid['lon']
            y = ds_grid['lat']
        elif 'grid_corner_lon' in ds_grid and 'grid_corner_lat' in ds_grid:
            x = ds_grid['grid_corner_lon']
            y = ds_grid['grid_corner_lat']
        else:
            raise RuntimeError('No valid coordinates in grid file.')
    else:
        x = ds_data['lon']
        y = ds_data['lat']

    # Make sure we don't have time or level dimensions
    if 'time' in data.dims:
        if time_index is None:
            data = data.mean(dim='time', keep_attrs=True).squeeze()
        else:
            data = data.isel(time=int(time_index))
    if 'lev' in data.dims:
        data = data.isel(lev=-1).squeeze()
    if 'time' in x.dims: x = x.isel(time=0)
    if 'time' in y.dims: y = y.isel(time=0)

    # Setup the canvas
    wks = ngl.open_wks(
        os.path.splitext(plotname)[1][1:],
        os.path.splitext(plotname)[0])

    # Get contour levels; the explicit type casting deals with problems calling
    # this standalone code using subprocess.run() with string arguments, where
    # all kwargs are going to be interpreted as strings
    if vmin is None: vmin = data.min().values
    if vmax is None: vmax = data.max().values
    if float(vmin) < 0 and float(vmax) > 0:
        *__, clevels = nice_cntr_levels(float(vmin),
                                        float(vmax),
                                        returnLevels=True,
                                        max_steps=13,
                                        aboutZero=True)
    else:
        *__, clevels = nice_cntr_levels(float(vmin),
                                        float(vmax),
                                        returnLevels=True,
                                        max_steps=13)
    kwargs['cnLevels'] = clevels  #get_contour_levels(data)
    kwargs['cnLevelSelectionMode'] = 'ExplicitLevels'

    # Make plot
    if 'lbTitleString' not in kwargs.keys():
        kwargs['lbTitleString'] = f'{data.long_name} ({data.units})'
    plot = plot_map(wks,
                    x.values,
                    y.values,
                    data.values,
                    mpGeophysicalLineColor='white',
                    lbOrientation='horizontal',
                    cnFillMode='RasterFill',
                    cnLineLabelsOn=False,
                    cnLinesOn=False,
                    **kwargs)

    ngl.destroy(wks)
Exemple #6
0
def main(test_files,
         cntl_files,
         varname,
         plotname,
         test_name="Model",
         cntl_name="Obs",
         time_offsets=[None, None],
         map_file=None,
         test_grid=None,
         cntl_grid=None,
         t1=None,
         t2=None,
         verbose=False,
         **kwargs):

    # Read data
    if verbose:
        print('Open datasets...')
        sys.stdout.flush()
    datasets = [
        open_dataset(files, time_offset=time_offset, chunks={'time': 1})
        for files, time_offset in zip((test_files, cntl_files), time_offsets)
    ]
    grids = [
        xarray.open_mfdataset(f) if f is not None else None
        for f in (test_grid, cntl_grid)
    ]

    # Subset data
    if verbose:
        print('Subset consistent time periods...')
        sys.stdout.flush()
    if t1 is None: t1 = max([ds.time[0].values for ds in datasets])
    if t2 is None: t2 = min([ds.time[-1].values for ds in datasets])
    if verbose:
        print('Comparing period {} to {}'.format(str(t1), str(t2)))
        sys.stdout.flush()
    datasets_subset = [ds.sel(time=slice(str(t1), str(t2))) for ds in datasets]

    # Compute time average
    if verbose:
        print('Compute time averages...')
        sys.stdout.flush()
    datasets_ta = [
        ds.mean(dim='time', keep_attrs=True) for ds in datasets_subset
    ]

    # Select data
    if verbose:
        print('Select data...')
        sys.stdout.flush()
    data_arrays = [mask_all_zero(get_data(ds, varname)) for ds in datasets_ta]
    #data_arrays = [get_data(ds, varname) for ds in datasets]
    coords = [
        get_coords(ds, ds_grid) for ds, ds_grid in zip(datasets_ta, grids)
    ]
    #coords = [get_coords(ds) for ds in datasets]
    weights = [get_area_weights(ds) for ds in datasets_ta]

    # Try explicitly evaluating dask graph to avoid dask getting confused later
    # and trying to load entire dataset into memory. By this point we should
    # have reduced the data arrays to a manageable size, so loading this into
    # memory should not be a problem. TODO: this feels like a hack I should not
    # have to do, is something else going on here?
    # NOTE: doing this here seems to speed things up quite a bit too.
    #print('Execute dask computations (data_arrays)...'); sys.stdout.flush()
    #data_arrays = [d.compute() for d in data_arrays]
    #print('Execute dask computations (coords)...'); sys.stdout.flush()
    #coords = [[c.compute() for c in coord] for coord in coords]
    #print('Execute dask computations (weights)...'); sys.stdout.flush()
    #weights = [w.compute() for w in weights]
    for i in range(len(data_arrays)):
        data_arrays[i].load()
        coords[i][0].load()
        coords[i][1].load()
        weights[i].load()

    # Compute differences
    if verbose:
        print('Compute differences...')
        sys.stdout.flush()
    data_arrays.append(
        compute_differences(data_arrays[0], data_arrays[1], *coords[0],
                            *coords[1], map_file))
    coords.append(coords[1])
    weights.append(weights[1])

    # Compute global statistics for plot labels
    if verbose:
        print('Compute statistics...')
        sys.stdout.flush()
    means = [area_average(d, w) for (d, w) in zip(data_arrays, weights)]

    # Get common contour levels
    if verbose:
        print('Find good clevels for plotting...')
        sys.stdout.flush()
    clevels = get_contour_levels([d for d in data_arrays[:-1]])
    dlevels = get_contour_levels([
        data_arrays[-1],
    ], aboutZero=True)
    levels_list = [clevels, clevels, dlevels]

    # Make plots
    if verbose:
        print('Make plots...')
        sys.stdout.flush()
    plot_format = plotname.split('.')[-1]
    wks = ngl.open_wks(plot_format, plotname)
    diff_name = f'{test_name} minus {cntl_name}'
    labels = [
        f'{label} ({m.values:.1f})'
        for (label, m) in zip((test_name, cntl_name, diff_name), means)
    ]
    plots = [
        plot_panel(wks, d, *c, tiMainString=label, cnLevels=levels, **kwargs)
        for (d, c, label,
             levels) in zip(data_arrays, coords, labels, levels_list)
    ]

    # Panel plots
    if verbose:
        print('Closing plot workspace and writing figures...')
        sys.stdout.flush()
    ngl.panel(wks, plots, [1, len(plots)])
    ngl.destroy(wks)

    # Finally, trim whitespace from our figures
    if verbose:
        print('Trimming whitespace from figure...')
        sys.stdout.flush()
    subprocess.call(f'convert -trim {plotname} {plotname}'.split(' '))