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)