def test_values_and_endianness(all_mds_datadirs): """Make sure we read all the grid variables.""" dirname, expected = all_mds_datadirs if expected['geometry'] == 'llc' and (dask.__version__ < '0.11.2'): pytest.xfail("LLC value tests require fixed dask") # default endianness ds = xmitgcm.open_mdsdataset(dirname, iters=None, read_grid=True, swap_dims=False, geometry=expected['geometry']) # now reverse endianness ds_le = xmitgcm.open_mdsdataset(dirname, iters=None, read_grid=True, endian='<', swap_dims=False, geometry=expected['geometry']) for vname, (idx, val) in expected['expected_values'].items(): np.testing.assert_allclose(ds[vname].values[idx], val) # dask arrays that have been concatenated revert to native endianness # https://github.com/dask/dask/issues/1647 if ds[vname].dtype.byteorder == '>': val_le = ds[vname].values.newbyteorder('<')[idx] np.testing.assert_allclose(ds_le[vname].values[idx], val_le)
def test_multiple_iters(multidim_mds_datadirs): """Test ability to load multiple iters into a single dataset.""" dirname, expected = multidim_mds_datadirs # first try specifying the iters ds = xmitgcm.open_mdsdataset( dirname, read_grid=False, geometry=expected['geometry'], iters=expected['all_iters'], prefix=expected['prefixes']) assert list(ds.iter.values) == expected['all_iters'] # now infer the iters, should be the same ds2 = xmitgcm.open_mdsdataset( dirname, read_grid=False, iters='all', geometry=expected['geometry'], prefix=expected['prefixes']) assert ds.equals(ds2) # In the test datasets, there is no PHL.0000000000.data file. # By default we infer the prefixes from the first iteration number, so this # leads to an error. # (Need to specify iters because there is some diagnostics output with # weird iterations numbers present in some experiments.) with pytest.raises(IOError): ds = xmitgcm.open_mdsdataset( dirname, read_grid=False, iters=expected['all_iters'], geometry=expected['geometry']) # now hide all the PH and PHL files: should be able to infer prefixes fine missing_files = [os.path.basename(f) for f in glob(os.path.join(dirname, 'PH*.0*data'))] with hide_file(dirname, *missing_files): ds = xmitgcm.open_mdsdataset( dirname, read_grid=False, iters=expected['all_iters'], geometry=expected['geometry'])
def get_compressed_level_index(grid_dir, index_fname='llc4320_compressed_level_index.nc', geometry='llc'): ''' Some doc ''' # ds = xm.open_mdsdataset( '', grid_dir=grid_dir, iters=None, geometry=geometry, read_grid=True, default_dtype=np.dtype('>f4'), ignore_unknown_vars=True) #, extra_metadata=llc4320) # get shape #nz, nface, ny, nx = ds.hFacC.shape #shape = (1, nface, ny, nx) try: ds_index = xr.open_dataset(grid_dir + index_fname) except OSError: # compute and save mask indices print('Create llc4320_compressed_level_index.nc in grid_dir') ds_index = ((ds.reset_coords()[['hFacC', 'hFacW', 'hFacS']] > 0).sum( axis=(1, 2, 3))) ds_index.coords['k'] = ds.k ds_index.load().to_netcdf(grid_dir + index_fname) print('done') return ds_index, ds
def test_open_dataset_2D_diags(all_mds_datadirs): # convert 3D fields with only 2D diagnostic output # https://github.com/xgcm/xmitgcm/issues/140 dirname, expected = all_mds_datadirs shape = expected['shape'] nz = shape[0] ny, nx = shape[-2:] shape_2d = shape[1:] dims_2d = ('j', 'i') if expected['geometry'] == 'llc': dims_2d = ('face', ) + dims_2d ny = nx * shape[-3] dims_3d = dims_2d if nz == 1 else ('k', ) + dims_2d dims_2d = ('time', ) + dims_2d dims_3d = ('time', ) + dims_3d it = expected['test_iternum'] kwargs = dict(iters=it, geometry=expected['geometry'], read_grid=False, swap_dims=False) to_hide = ['T.%010d.meta' % it, 'T.%010d.data' % it] with hide_file(dirname, *to_hide): ldir = py.path.local(dirname) old_prefix = 'Eta.%010d' % it new_prefix = 'T.%010d' % it for suffix in ['.data', '.meta']: lp = ldir.join(old_prefix + suffix) lp.copy(ldir.join(new_prefix + suffix)) ds = xmitgcm.open_mdsdataset(dirname, prefix=['T'], **kwargs)
def test_llc_dims(llc_mds_datadirs, method, with_refdate): """Check that the LLC file dimensions are correct.""" dirname, expected = llc_mds_datadirs if with_refdate: kwargs = { 'delta_t': expected['delta_t'], 'ref_date': expected['ref_date'] } else: kwargs = {} ds = xmitgcm.open_mdsdataset(dirname, iters=expected['test_iternum'], geometry=expected['geometry'], llc_method=method, **kwargs) nz, nface, ny, nx = expected['shape'] nt = 1 assert ds.dims['face'] == 13 assert ds.rA.dims == ('face', 'j', 'i') assert ds.rA.values.shape == (nface, ny, nx) assert ds.U.dims == ('time', 'k', 'face', 'j', 'i_g') assert ds.U.values.shape == (nt, nz, nface, ny, nx) assert ds.V.dims == ('time', 'k', 'face', 'j_g', 'i') assert ds.V.values.shape == (nt, nz, nface, ny, nx) print(ds.U.chunks) if method == "smallchunks": assert ds.U.chunks == (nt * (1, ), nz * (1, ), nface * (1, ), (ny, ), (nx, ))
def test_llc_extra_metadata(llc_mds_datadirs, method): """Check that the LLC reads properly when using extra_metadata.""" dirname, expected = llc_mds_datadirs nz, nface, ny, nx = expected['shape'] nt = 1 llc = xmitgcm.utils.get_extra_metadata(domain='llc', nx=nx) ds = xmitgcm.open_mdsdataset(dirname, iters=expected['test_iternum'], geometry=expected['geometry'], llc_method=method, extra_metadata=llc) assert ds.dims['face'] == 13 assert ds.rA.dims == ('face', 'j', 'i') assert ds.rA.values.shape == (nface, ny, nx) assert ds.U.dims == ('time', 'k', 'face', 'j', 'i_g') assert ds.U.values.shape == (nt, nz, nface, ny, nx) assert ds.V.dims == ('time', 'k', 'face', 'j_g', 'i') assert ds.V.values.shape == (nt, nz, nface, ny, nx) if method == "smallchunks": assert ds.U.chunks == (nt * (1, ), nz * (1, ), nface * (1, ), (ny, ), (nx, ))
def test_open_mdsdataset_minimal(all_mds_datadirs): """Create a minimal xarray object with only dimensions in it.""" dirname, expected = all_mds_datadirs ds = xmitgcm.open_mdsdataset(dirname, iters=None, read_grid=False, swap_dims=False, geometry=expected['geometry']) # the expected dimensions of the dataset eshape = expected['shape'] if len(eshape) == 3: nz, ny, nx = eshape nface = None elif len(eshape) == 4: nz, nface, ny, nx = eshape else: raise ValueError("Invalid expected shape") coords = { 'i': np.arange(nx), 'i_g': np.arange(nx), # 'i_z': np.arange(nx), 'j': np.arange(ny), 'j_g': np.arange(ny), # 'j_z': np.arange(ny), 'k': np.arange(nz), 'k_u': np.arange(nz), 'k_l': np.arange(nz), 'k_p1': np.arange(nz + 1) } if nface is not None: coords['face'] = np.arange(nface) if 'layers' in expected: for layer_name, n_layers in expected['layers'].items(): for suffix, offset in zip(['bounds', 'center', 'interface'], [0, -1, -2]): dimname = 'l' + layer_name[0] + '_' + suffix[0] index = np.arange(n_layers + offset) coords[dimname] = index ds_expected = xr.Dataset(coords=coords) assert ds_expected.equals(ds) # check that the datatypes are correct for key in coords: assert ds[key].dtype == np.int64 # check for comodo metadata needed by xgcm assert ds['i'].attrs['axis'] == 'X' assert ds['i_g'].attrs['axis'] == 'X' assert ds['i_g'].attrs['c_grid_axis_shift'] == -0.5 assert ds['j'].attrs['axis'] == 'Y' assert ds['j_g'].attrs['axis'] == 'Y' assert ds['j_g'].attrs['c_grid_axis_shift'] == -0.5 assert ds['k'].attrs['axis'] == 'Z' assert ds['k_l'].attrs['axis'] == 'Z' assert ds['k_l'].attrs['c_grid_axis_shift'] == -0.5
def test_llc_facets_2d_to_compact(llc_mds_datadirs): from xmitgcm.utils import llc_facets_2d_to_compact, get_extra_metadata from xmitgcm.utils import rebuild_llc_facets, read_raw_data from xmitgcm.utils import write_to_binary from xmitgcm import open_mdsdataset dirname, expected = llc_mds_datadirs # open dataset ds = open_mdsdataset(dirname, iters=expected['test_iternum'], geometry=expected['geometry']) nt, nfaces, ny, nx = expected['shape'] md = get_extra_metadata(domain=expected['geometry'], nx=nx) # split in facets facets = rebuild_llc_facets(ds['XC'], md) flatdata = llc_facets_2d_to_compact(facets, md) # compare with raw data raw = read_raw_data(dirname + '/XC.data', np.dtype('>f'), (nfaces, ny, nx)) flatraw = raw.flatten() assert len(flatdata) == len(flatraw) assert flatdata.min() == flatraw.min() assert flatdata.max() == flatraw.max() # write new file write_to_binary(flatdata, 'tmp.bin', dtype=np.dtype('f')) md5new = file_md5_checksum('tmp.bin') md5old = file_md5_checksum(dirname + '/XC.data') assert md5new == md5old os.remove('tmp.bin')
def plot_yield(name_file, data_dir='./', save=False, e=2, f=1.0, t=0.0, tstp=-1): data = open_mdsdataset(data_dir, prefix=[name_file], chunks=None) sig1 = data['SIsig1'][tstp] sig2 = data['SIsig2'][tstp] sig1f = sig1.values.flatten() sig2f = sig2.values.flatten() print len(sig1f) plt.figure(1) ax = plt.gca() elli = Ellipse(xy=((-f + t) / 2, (-f + t) / 2), width=(f + t) / e * np.sqrt(2), height=(f + t) * np.sqrt(2), angle=-45, edgecolor='b', fc='None', lw=0.5) ax.add_patch(elli) plt.plot(sig1f, sig2f, 'r.', markersize=1) xlim = ax.get_xlim() ylim = ax.get_ylim() ax.set(xlim=xlim, ylim=ylim) cordx = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])] cordy = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])] plt.plot(cordx, cordy, '-k', lw=1) coradx = [-min(-xlim[0], ylim[1]), min(xlim[1], -ylim[0])] corady = [min(-xlim[0], ylim[1]), -min(xlim[1], -ylim[0])] plt.plot(coradx, corady, '-k', lw=1) plt.axvline(linewidth=1, color='k') plt.axhline(linewidth=1, color='k') plt.title(r" Yield curve ") plt.axes().set_aspect('equal') plt.grid() if save == True: plt.savefig("yield_curve.png", bbox_inches='tight', dpi=200) plt.show() # plt.close('all') return None
def test_levels_diagnostics(mds_datadirs_with_inputfiles): dirname, expected = mds_datadirs_with_inputfiles for diagname, (levels, (idx, value)) in expected['diag_levels'].items(): ds = xmitgcm.open_mdsdataset(dirname, prefix=[diagname], levels=levels, geometry=expected['geometry']) assert ds['Zl'].values[idx] == value with pytest.warns(UserWarning, match='nz will be ignored'): xmitgcm.open_mdsdataset(dirname, prefix=[diagname], levels=levels, geometry=expected['geometry'], nz=12)
def plot_comp(name_file, dirs, varnames, loga=0., beg=0, end=-1, step=1, dpi=100): data_dir1 = dirs[0] data1 = open_mdsdataset(data_dir1, prefix=[name_file]) data_dir2 = dirs[1] data2 = open_mdsdataset(data_dir2, prefix=[name_file]) if loga == 0.: loga = np.full((len(varnames)), False, dtype=bool) if beg == 0 and end == -1 and step == 1: nbr = data1.dims['time'] times = xrange(nbr) else: var_cho = np.array(var[beg:end + 1:step, :]) for ivar in range(len(varnames)): name_var = varnames[ivar] var_dir = str('./plots_comp_' + name_var + '_' + str(data1.coords['iter'][beg].values) + '_' + str(data1.coords['iter'][end].values)) var_dirp = str(var_dir + '/') if os.path.exists(var_dir): shutil.rmtree(var_dir) os.mkdir(var_dir) var_max = 0.9 var_min = 1.1 for tstep in times: plot_comp_forloop(data1, data2, name_var, tstep, loga[ivar], var_dirp, var_min, var_max, dpi) return None
def test_drc_length(all_mds_datadirs): """Test that open_mdsdataset is adding an extra level to drC if it has length nr""" dirname, expected = all_mds_datadirs #Only older versions of the gcm have len(drC) = nr, so force len(drC) = nr for the test copyfile(os.path.join(dirname, 'DRF.data'), os.path.join(dirname, 'DRC.data')) copyfile(os.path.join(dirname, 'DRF.meta'), os.path.join(dirname, 'DRC.meta')) ds = xmitgcm.open_mdsdataset( dirname, iters=None, read_grid=True, geometry=expected['geometry']) assert len(ds.drC)==(len(ds.drF)+1)
def test_read_grid(all_mds_datadirs): """Make sure we read all the grid variables.""" dirname, expected = all_mds_datadirs ds = xmitgcm.open_mdsdataset(dirname, iters=None, read_grid=True, geometry=expected['geometry']) for vname in _EXPECTED_GRID_VARS: assert vname in ds
def mitgcm_Movie(ddir, prefix=["tracer_snapshots"], maskname="hFacC", varname="TRAC01", clim=[-1, 35]): ds = open_mdsdataset(ddir, prefix=prefix, swap_dims=False) ds = ds[varname].where(ds[maskname] == 1) run = os.path.basename(os.path.normpath(ddir)) odir = ddir + "/movie" Movie(ds, odir, clim=clim, moviename=run)
def test_open_with_ref_date_grid(mds_datadirs_with_refdate, method): """With ref_date, without grid.""" dirname, expected = mds_datadirs_with_refdate ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=expected['ref_date'], read_grid=method, delta_t=expected['delta_t'], geometry=expected['geometry'])
def test_read_grid(all_mds_datadirs): """Make sure we read all the grid variables.""" dirname, expected = all_mds_datadirs ds = xmitgcm.open_mdsdataset( dirname, iters=None, read_grid=True, geometry=expected['geometry']) for vname in _EXPECTED_GRID_VARS: assert vname in ds # actually load the data, to check for dask-related errors ds.load()
def test_prefixes(all_mds_datadirs): """Make sure we read all the grid variables.""" dirname, expected = all_mds_datadirs prefixes = ['U', 'V', 'W', 'T', 'S', 'PH'] # , 'PHL', 'Eta'] iters = [expected['test_iternum']] ds = xmitgcm.open_mdsdataset( dirname, iters=iters, prefix=prefixes, read_grid=False, geometry=expected['geometry']) for p in prefixes: assert p in ds
def plot_yield_vid(name_file, data_dir='./', e=2, f=1.0, t=0.0): FFMpegWriter = manimation.writers['ffmpeg'] metadata = dict(title='yield', artist='dringeis') writer = FFMpegWriter(fps=30, metadata=metadata, bitrate=-1) data = open_mdsdataset(data_dir, prefix=[name_file], chunks=None) sig1 = data['SIsig1'] sig2 = data['SIsig2'] fig = plt.figure() ax = fig.gca() elli = Ellipse(xy=((-f + t) / 2, (-f + t) / 2), width=(f + t) / e * np.sqrt(2), height=(f + t) * np.sqrt(2), angle=-45, edgecolor='b', fc='None', lw=0.5) ax.add_patch(elli) im, = plt.plot([], [], 'r.', markersize=1) xlim = ax.get_xlim() ylim = ax.get_ylim() ax.set(xlim=xlim, ylim=ylim) cordx = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])] cordy = [min(xlim[0], ylim[0]), min(xlim[1], ylim[1])] plt.plot(cordx, cordy, '-k', lw=0.5) coradx = [-min(-xlim[0], ylim[1]), min(xlim[1], -ylim[0])] corady = [min(-xlim[0], ylim[1]), -min(xlim[1], -ylim[0])] plt.plot(coradx, corady, '-k', lw=0.5) plt.axvline(linewidth=0.5, color='k') plt.axhline(linewidth=0.5, color='k') plt.title(r" Yield curve ") plt.axes().set_aspect('equal') plt.grid() with writer.saving(fig, "yield.mp4", 200): for t in xrange(data.dims['time']): sig1f = sig1[t, :, :].values.flatten() sig2f = sig2[t, :, :].values.flatten() im.set_data(sig1f, sig2f) writer.grab_frame() return None
def test_swap_dims(all_mds_datadirs): """See if we can swap dimensions.""" dirname, expected = all_mds_datadirs kwargs = dict(iters=None, read_grid=True, geometry=expected['geometry']) expected_dims = ['XC', 'XG', 'YC', 'YG', 'Z', 'Zl', 'Zp1', 'Zu'] # make sure we never swap if not reading grid assert 'i' in xmitgcm.open_mdsdataset(dirname, iters=None, read_grid=False, geometry=expected['geometry']) if expected['geometry'] in ('llc', 'curvilinear'): # make sure swapping is not the default ds = xmitgcm.open_mdsdataset(dirname, **kwargs) assert 'i' in ds # and is impossible with pytest.raises(ValueError) as excinfo: ds = xmitgcm.open_mdsdataset(dirname, swap_dims=True, **kwargs) else: # make sure swapping *IS* the default assert 'i' not in xmitgcm.open_mdsdataset(dirname, **kwargs) ds = xmitgcm.open_mdsdataset(dirname, geometry=expected['geometry'], iters=None, read_grid=True, swap_dims=True, grid_vars_to_coords=True) # check for comodo metadata needed by xgcm assert ds['XC'].attrs['axis'] == 'X' assert ds['XG'].attrs['axis'] == 'X' assert ds['XG'].attrs['c_grid_axis_shift'] == -0.5 assert ds['YC'].attrs['axis'] == 'Y' assert ds['YG'].attrs['axis'] == 'Y' assert ds['YG'].attrs['c_grid_axis_shift'] == -0.5 assert ds['Z'].attrs['axis'] == 'Z' assert ds['Zl'].attrs['axis'] == 'Z' assert ds['Zl'].attrs['c_grid_axis_shift'] == -0.5 # add extra layers dimensions if needed if 'layers' in expected: for layer_name in expected['layers']: extra_dims = [ 'layer_' + layer_name + suffix for suffix in ['_bounds', '_center', '_interface'] ] expected_dims += extra_dims assert list(ds.dims.keys()) == expected_dims # make sure swapping works with multiple iters ds = xmitgcm.open_mdsdataset(dirname, geometry=expected['geometry'], prefix=['S']) #print(ds) ds.load() assert 'XC' in ds['S'].dims assert 'YC' in ds['S'].dims
def test_tload_with_ref_date_swap(mds_datadirs_with_refdate, method): """With ref_date, without grid.""" dirname, expected = mds_datadirs_with_refdate ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=expected['ref_date'], read_grid=True, delta_t=expected['delta_t'], swap_dims=method, geometry=expected['geometry']) ds.time.load()
def load_masks_from_mds(grid_dir): import xmitgcm ds = xmitgcm.open_mdsdataset(grid_dir, iters=None, geometry='llc') points = ['C', 'W', 'S'] masks = [ ds['hfac' + point].reset_coords(drop=True).rename('mask' + point) for point in points ] ds_mask = xr.merge(masks) for c in ds_mask.coords: ds_mask[c] = ds_mask[c].astype('i2') return ds_mask.chunk({'k': 1, 'face': -1})
def test_mask_values(all_mds_datadirs): """Test that open_mdsdataset generates binary masks with correct values""" dirname, expected = all_mds_datadirs ds = xmitgcm.open_mdsdataset( dirname, iters=None, read_grid=True, geometry=expected['geometry']) hFac_list = ['hFacC', 'hFacW', 'hFacS'] mask_list = ['maskC', 'maskW', 'maskS'] for hFac, mask in zip(hFac_list, mask_list): xr.testing.assert_equal(ds[hFac] * ds[mask], ds[hFac])
def test_serialize_nonstandard_calendar(multidim_mds_datadirs, tmp_path): dirname, expected = multidim_mds_datadirs ref_date = '2680-01-01 00:00:00' calendar = '360_day' ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=ref_date, calendar=calendar, read_grid=False, delta_t=expected['delta_t'], geometry=expected['geometry']) ds.to_zarr(tmp_path / 'test.zarr')
def test_date_parsing(mds_datadirs_with_refdate): """Verify that time information is decoded properly.""" dirname, expected = mds_datadirs_with_refdate ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=expected['ref_date'], read_grid=False, delta_t=expected['delta_t'], geometry=expected['geometry']) for i, date in expected['expected_time']: assert ds.time[i].values == date
def test_extra_variables(all_mds_datadirs): """Test that open_mdsdataset reads extra_variables correctly""" dirname, expected = all_mds_datadirs extra_variable_data = dict( # use U,V to test with attrs (including mate) testdataU=dict(dims=['k', 'j', 'i_g'], attrs=dict(standard_name='sea_water_x_velocity', mate='testdataV', long_name='Zonal Component of Velocity', units='m s-1')), testdataV=dict(dims=['k', 'j_g', 'i'], attrs=dict(standard_name='sea_water_y_velocity', mate='testdataU', long_name='Meridional Component of Velocity', units='m s-1')), # use T to test without attrs testdataT=dict(dims=['k', 'j', 'i'], attrs=dict())) for var in ["U", "V", "T"]: input_dir = os.path.join( dirname, '{}.{:010d}'.format(var, expected['test_iternum'])) test_dir = os.path.join( dirname, 'testdata{}.{:010d}'.format(var, expected['test_iternum'])) if input_dir + ".meta" not in os.listdir( ) or input_dir + ".data" not in os.listdir(): return copyfile(input_dir + ".meta", test_dir + ".meta") copyfile(input_dir + ".data", test_dir + ".data") assert test_dir + ".meta" not in os.listdir( ), f"{var} did not copy meta!" assert test_dir + ".data" not in os.listdir( ), f"{var} did not copy data!" ds = xmitgcm.open_mdsdataset(dirname, read_grid=False, iters=expected['test_iternum'], geometry=expected['geometry'], prefix=list(extra_variable_data.keys()), extra_variables=extra_variable_data) for var in extra_variable_data.keys(): assert var in ds if 'mate' in ds[var].attrs: mate = ds[var].attrs['mate'] assert ds[mate].attrs['mate'] == var
def test_separate_grid_dir(all_mds_datadirs): """Make sure we can have the grid files in a separate directory.""" dirname, expected = all_mds_datadirs prefixes = ['U', 'V', 'W', 'T', 'S', 'PH'] # , 'PHL', 'Eta'] iters = [expected['test_iternum']] with hide_file(dirname, *['XC.meta', 'XC.data', 'RC.meta', 'RC.data']) as grid_dir: ds = xmitgcm.open_mdsdataset( dirname, grid_dir=grid_dir, iters=iters, prefix=prefixes, read_grid=False, geometry=expected['geometry']) for p in prefixes: assert p in ds
def test_ref_date(mds_datadirs_with_refdate, swap_dims, read_grid, load): """With ref_date, without grid.""" dirname, expected = mds_datadirs_with_refdate if expected['geometry']=='llc' and swap_dims: pytest.skip("can't swap_dims with geometry=='llc'") ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=expected['ref_date'], delta_t=expected['delta_t'], geometry=expected['geometry'], read_grid=read_grid, swap_dims=swap_dims) if load: ds.time.load()
def calc_mean_EK(start=100, end=1000, step=100): times = np.arange(start, end + 1, step, dtype=np.int) var_list = ['U', 'V', 'W'] ds2 = open_mdsdataset('.', prefix=['CN', 'SN', 'XC', 'YC'], iters=100, geometry='curvilinear', ignore_unknown_vars=True) XC = np.array(ds2.XC) YC = np.array(ds2.YC) #EK = np.zeros(XC.shape) for i in range(len(times)): time = times[i] print("TIME: ", time / times[-1]) ds = open_mdsdataset('.', prefix=var_list, iters=time, geometry='curvilinear', ignore_unknown_vars=True, default_dtype=np.float) U = np.array(ds.U)[0, :, :] V = np.array(ds.V)[0, :, :] W = np.array(ds.W)[0, :, :] U2 = np.sqrt(U * U + V * V + W * W) if i == 0: EK = U2.copy() else: EK += U2 return EK
def test_date_parsing(mds_datadirs_with_refdate): """Verify that time information is decoded properly.""" dirname, expected = mds_datadirs_with_refdate ds = xmitgcm.open_mdsdataset(dirname, iters='all', prefix=['S'], ref_date=expected['ref_date'], read_grid=False, delta_t=expected['delta_t'], geometry=expected['geometry']) for i, date in expected['expected_time']: assert ds.time[i].values == date # since time was decoded, this encoding should be removed from attributes assert 'units' not in ds.time.attrs assert 'calendar' not in ds.time.attrs
def test_diagnostics(mds_datadirs_with_diagnostics): """Try reading dataset with diagnostics output.""" dirname, expected = mds_datadirs_with_diagnostics diag_prefix, expected_diags = expected['diagnostics'] ds = xmitgcm.open_mdsdataset(dirname, read_grid=False, iters=expected['test_iternum'], prefix=[diag_prefix], geometry=expected['geometry']) for diagname in expected_diags: assert diagname in ds # check vector mates if 'mate' in ds[diagname].attrs: mate = ds[diagname].attrs['mate'] assert ds[mate].attrs['mate'] == diagname