def test_preprocess_raises_logger_warning_for_filtered_vars(self):
with self.assertLogs(level=logging.WARNING) as log:
_ = cosmo.preprocess_cosmo(self.dataset,
self.assim_vars + ['TEMP'])
with self.assertRaises(AssertionError):
with self.assertLogs(level=logging.WARNING):
_ = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
def test_precosmo_calls_prepare_vgrid(self):
vcoord = self.dataset['vcoord']
ds = self.dataset[self.assim_vars]
ret_ds = cosmo._prepare_vgrid(ds, vcoord)
with patch('pytassim.model.terrsysmp.cosmo._prepare_vgrid',
return_value=ret_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]), 2)
xr.testing.assert_identical(vgrid_patch.call_args[0][0], ds)
xr.testing.assert_identical(vgrid_patch.call_args[0][1], vcoord)
def test_post_cosmo_logs_warning_if_var_not_found(self):
prepared_arr = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
tmp_arr = prepared_arr.sel(var_name=[
'T',
])
tmp_arr['var_name'] = [
'TEMP',
]
concatenated_arr = xr.concat([prepared_arr, tmp_arr], dim='var_name')
with self.assertLogs(level=logging.WARNING) as log:
_ = cosmo.postprocess_cosmo(concatenated_arr, self.dataset)
self.assertEqual(log.records[0].msg, 'Var: TEMP is not found')
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_preprocess_selects_only_available_vars(self):
with self.assertLogs(level=logging.WARNING) as log:
ens_arr = cosmo.preprocess_cosmo(self.dataset,
self.assim_vars + ['TEMP'])
self.assertListEqual(list(ens_arr.var_name), self.assim_vars)
def test_dataset_can_be_reconstructed(self):
pre_arr = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
post_ds = cosmo.postprocess_cosmo(pre_arr, self.dataset)
xr.testing.assert_equal(post_ds, self.dataset)
def test_precosmo_tranposes_for_valid_state(self):
preprocessed = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
self.assertTrue(preprocessed.state.valid)
def test_precosmo_expands_ensemble(self):
self.assertNotIn('ensemble', self.dataset)
preprocessed = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
self.assertIn('ensemble', preprocessed.dims)
np.testing.assert_equal(preprocessed.ensemble.values, np.array(0))
def test_precosmo_stacks_grid_coordinates(self):
preprocessed = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
self.assertIn('grid', preprocessed.dims)
self.assertListEqual(list(preprocessed.indexes['grid'].names),
['rlat', 'rlon', 'vgrid'])
def test_precosmo_makes_ds_to_array(self):
preprocessed = cosmo.preprocess_cosmo(self.dataset, self.assim_vars)
self.assertIsInstance(preprocessed, xr.DataArray)
self.assertIn('var_name', preprocessed.dims)
self.assertListEqual(list(preprocessed.var_name), self.assim_vars)