Make Nowcast Timeseries¶

from salishsea_tools import grid_tools, timeseries_tools
import xarray as xr
import pickle

def reshape_GEM(data, grid, mask):
    """Flattens GEM HRDPS wind to 2-D (time-space) and removes land points
    """

    # Initialize storage
    data_flat = {'raw': {'grid': {}}}
    
    # Build GEM coordinates
    mask_flt, coords, ngrid, ngrid_water = timeseries_tools.reshape_coords_GEM(
        grid, mask,
    )

    # Reshape, flatten, and trim
    data_flat['raw']['grid']['u'] = timeseries_tools.reshape_to_ts(
        data.u_wind.values, mask_flt.astype(bool), ngrid, ngrid_water,
    )
    data_flat['raw']['grid']['v'] = timeseries_tools.reshape_to_ts(
        data.v_wind.values, mask_flt.astype(bool), ngrid, ngrid_water,
    )

    # Assign parameters to output dict
    data_flat['mask'] = mask_flt
    data_flat['coords'] = coords
    data_flat['ngrid'] = ngrid
    data_flat['ngrid_water'] = ngrid_water

    return data_flat

def iterate_NEMO_timeseries(timerange, variables, mask, dims, indices):
    """Wrapper for timeseries_tools.load_NEMO_timeseries that loads and flattens
    Nowcast results to 2-D (time-space) and removes land points
    """
    
    # Intialize empty dict
    results = {}

    # Loop through results slices
    for var in variables:  # --- Loop through variables
        for dim, index in zip(dims, indices):  # --- Loop through slices
            
            # Populate dict fields
            results[var] = {}
            results[var][dim] = {}
            results[var][dim][index] = {}

            # Determine filenames and unstaggering from variable names
            dim_in, unstagger = dim, None
            if var is 'vozocrtx':
                file = 'U'
                if dim is 'depth':
                    dim_in, unstagger = 'depthu', 'x'
            elif var is 'vomecrty':
                file = 'V'
                if dim is 'depth':
                    dim_in, unstagger = 'depthv', 'y'
            else:
                file = 'T'
                if dim is 'depth':
                    dim_in = 'deptht'

            # Determine spacing from variable names and slicing
            if (var is 'vozocrtx' or var is 'vomecrty') and dim is 'depth':
                spacing = 5
            else:
                spacing = 1

            # Construct Filename List
            filenames = timeseries_tools.make_filename_list(
                timerange, file, model='nowcast-green', resolution='h'
            )
            
            # Load Results
            data, coords = timeseries_tools.load_NEMO_timeseries(
                filenames, mask, var, dim_in, index=index,
                spacing=spacing, shape='flat', unstagger_dim=unstagger,
            )
            
            # Export results as dict
            results[var][dim][index]['data'] = data
            results[var][dim][index]['coords'] = coords

            # Export timerange
            results[var][dim][index]['coords']['timerange'] = timerange

    return results

# Timerange
timerange = ['2017 Jan 1 00:00', '2017 Jan 31 23:00']

# Load GEM record
GEM_in = xr.open_dataset('https://salishsea.eos.ubc.ca/erddap/griddap/ubcSSaSurfaceAtmosphereFieldsV1')
GEM = reshape_GEM(GEM_in.sel(time=slice(*timerange)), grid_GEM, mask_GEM)