def test_replace_coords_to_nearest_arakawa_a_grid(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
prep_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)
reindexed_ds = cosmo._interp_vgrid(prep_ds)
replaced = cosmo._replace_coords(reindexed_ds)
self.assertNotIn('srlon', replaced['U'].dims)
def test_expand_no_grid_vars(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
ret_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)
self.assertTupleEqual(tuple(ret_ds['T_S'].dims),
('time', 'no_vgrid', 'rlat', 'rlon'))
np.testing.assert_equal(ret_ds.no_vgrid.values, np.array(0))
def test_replace_coords_replaces_vertical_coords_with_vgrid(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
prep_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)
reindexed_ds = cosmo._interp_vgrid(prep_ds)
replaced = cosmo._replace_coords(reindexed_ds)
for var in self.assim_vars:
self.assertIn('vgrid', replaced[var].dims)
def test_precosmo_calls_replace_coords(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
prep_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)
reindexed_ds = cosmo._interp_vgrid(prep_ds)
replaced_ds = cosmo._replace_coords(reindexed_ds)
with patch('pytassim.model.terrsysmp.cosmo._replace_coords',
return_value=replaced_ds) as vgrid_patch:
_ = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
vgrid_patch.assert_called_once()
self.assertEqual(len(vgrid_patch.call_args[0]), 1)
xr.testing.assert_identical(vgrid_patch.call_args[0][0], reindexed_ds)
def test_interp_remaps_to_right_vcoords(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
prep_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)
reindexed_ds = cosmo._interp_vgrid(prep_ds)
expected_values = {
'no_vgrid': ('T_S', np.array(0)),
'height_2m': ('T_2M', np.array(0)),
'height_10m': ('U_10M', np.array(20)),
'height_toa': ('ASOB_T', reindexed_ds.vgrid.values[0]),
'soil1': ('W_SO', reindexed_ds.vgrid.values[-8:]),
'level1': ('W', reindexed_ds.vgrid.values[:51]),
'level': ('T', reindexed_ds.vgrid.values[:50])
}
for coord, val in expected_values.items():
dropped_arr = reindexed_ds[val[0]].dropna(coord, how='all')
np.testing.assert_equal(dropped_arr[coord].values, val[1])
def test_expand_no_grid_vars_works_with_no_vars(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
prep_ds = cosmo._prepare_vgrid(ds, vcoord)
del prep_ds['T_S']
ret_ds = common.add_no_vgrid(prep_ds, cosmo._cosmo_vcoords)