def _group_anomaly(group: xr.Dataset, ref: xr.Dataset, monitor: Monitor = Monitor.NONE, step: float = None): """ Calculate anomaly for the given group. :param group: Result of a groupby('time.month') operation :param ref: Reference dataset :param monitor: Monitor of the parent method :param step: Step to add to monitor progress :return: Group dataset with anomaly calculation applied """ # Retrieve the month of the current group month = group['time.month'][0].values ret = diff(group, ref.isel(time=month - 1)) monitor.progress(work=step) return ret
def test_diff(self): # Test nominal dataset = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) expected = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) actual = arithmetics.diff(dataset, dataset * 2) assert_dataset_equal(expected * -1, actual) # Test variable mismatch ds = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'third': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) ds1 = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) expected = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.zeros([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) actual = arithmetics.diff(ds, ds1) assert_dataset_equal(expected, actual) actual = arithmetics.diff(ds1, ds) assert_dataset_equal(expected, actual) # Test date range mismatch ds = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 12])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 12])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90), 'time': [datetime(2000, x, 1) for x in range(1, 13)]}) ds1 = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 12])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 12])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90), 'time': [datetime(2003, x, 1) for x in range(1, 13)]}) expected = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.zeros([45, 90, 12])), 'second': (['lat', 'lon', 'time'], np.zeros([45, 90, 12])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) actual = arithmetics.diff(ds, ds1) assert_dataset_equal(actual, expected) actual = arithmetics.diff(ds, ds1.drop('time')) expected['time'] = [datetime(2000, x, 1) for x in range(1, 13)] assert_dataset_equal(actual, expected) # Test broadcasting ds = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) ds1 = xr.Dataset({ 'first': (['lat', 'lon'], np.ones([45, 90])), 'second': (['lat', 'lon'], np.ones([45, 90])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) expected = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.zeros([45, 90, 3])), 'second': (['lat', 'lon', 'time'], np.zeros([45, 90, 3])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90)}) actual = arithmetics.diff(ds, ds1) assert_dataset_equal(expected, actual) ds['time'] = [datetime(2000, x, 1) for x in range(1, 4)] expected['time'] = [datetime(2000, x, 1) for x in range(1, 4)] actual = arithmetics.diff(ds, ds1) assert_dataset_equal(expected, actual) ds1 = xr.Dataset({ 'first': (['lat', 'lon', 'time'], np.ones([45, 90, 1])), 'second': (['lat', 'lon', 'time'], np.ones([45, 90, 1])), 'lat': np.linspace(-88, 88, 45), 'lon': np.linspace(-178, 178, 90), 'time': [datetime(2001, 1, 1)]}) actual = arithmetics.diff(ds, ds1) assert_dataset_equal(expected, actual) ds1 = ds1.squeeze('time') ds1['time'] = 1 actual = arithmetics.diff(ds, ds1) assert_dataset_equal(expected, actual)