Python xray.Dataset类代码示例

 def test_open_and_do_math(self):
     original = Dataset({'foo': ('x', np.random.randn(10))})
     with create_tmp_file() as tmp:
         original.to_netcdf(tmp)
         with open_mfdataset(tmp) as ds:
             actual = 1.0 * ds
             self.assertDatasetAllClose(original, actual)

    def test_coordinates_encoding(self):
        def equals_latlon(obj):
            return obj == 'lat lon' or obj == 'lon lat'

        original = Dataset({'temp': ('x', [0, 1]), 'precip': ('x', [0, -1])},
                           {'lat': ('x', [2, 3]), 'lon': ('x', [4, 5])})
        with self.roundtrip(original) as actual:
            self.assertDatasetIdentical(actual, original)
        with create_tmp_file() as tmp_file:
            original.to_netcdf(tmp_file)
            with open_dataset(tmp_file, decode_coords=False) as ds:
                self.assertTrue(equals_latlon(ds['temp'].attrs['coordinates']))
                self.assertTrue(equals_latlon(ds['precip'].attrs['coordinates']))
                self.assertNotIn('coordinates', ds.attrs)
                self.assertNotIn('coordinates', ds['lat'].attrs)
                self.assertNotIn('coordinates', ds['lon'].attrs)

        modified = original.drop(['temp', 'precip'])
        with self.roundtrip(modified) as actual:
            self.assertDatasetIdentical(actual, modified)
        with create_tmp_file() as tmp_file:
            modified.to_netcdf(tmp_file)
            with open_dataset(tmp_file, decode_coords=False) as ds:
                self.assertTrue(equals_latlon(ds.attrs['coordinates']))
                self.assertNotIn('coordinates', ds['lat'].attrs)
                self.assertNotIn('coordinates', ds['lon'].attrs)

    def test_roundtrip_strings_with_fill_value(self):
        values = np.array(['ab', 'cdef', np.nan], dtype=object)
        encoding = {'_FillValue': np.string_('X'), 'dtype': np.dtype('S1')}
        original = Dataset({'x': ('t', values, {}, encoding)})
        expected = original.copy(deep=True)
        expected['x'][:2] = values[:2].astype('S')
        with self.roundtrip(original) as actual:
            self.assertDatasetIdentical(expected, actual)

        original = Dataset({'x': ('t', values, {}, {'_FillValue': '\x00'})})
        if not isinstance(self, Only32BitTypes):
            # these stores can save unicode strings
            expected = original.copy(deep=True)
        if isinstance(self, BaseNetCDF4Test):
            # netCDF4 can't keep track of an empty _FillValue for VLEN
            # variables
            expected['x'][-1] = ''
        elif (isinstance(self, (NetCDF3ViaNetCDF4DataTest,
                                NetCDF4ClassicViaNetCDF4DataTest)) or
              (has_netCDF4 and type(self) is GenericNetCDFDataTest)):
            # netCDF4 can't keep track of an empty _FillValue for nc3, either:
            # https://github.com/Unidata/netcdf4-python/issues/273
            expected['x'][-1] = np.string_('')
        with self.roundtrip(original) as actual:
            self.assertDatasetIdentical(expected, actual)

 def test_roundtrip_object_dtype(self):
     floats = np.array([0.0, 0.0, 1.0, 2.0, 3.0], dtype=object)
     floats_nans = np.array([np.nan, np.nan, 1.0, 2.0, 3.0], dtype=object)
     letters = np.array(['ab', 'cdef', 'g'], dtype=object)
     letters_nans = np.array(['ab', 'cdef', np.nan], dtype=object)
     all_nans = np.array([np.nan, np.nan], dtype=object)
     original = Dataset({'floats': ('a', floats),
                         'floats_nans': ('a', floats_nans),
                         'letters': ('b', letters),
                         'letters_nans': ('b', letters_nans),
                         'all_nans': ('c', all_nans),
                         'nan': ([], np.nan)})
     if PY3 and type(self) is ScipyDataTest:
         # see the note under test_zero_dimensional_variable
         del original['nan']
     expected = original.copy(deep=True)
     if type(self) in [NetCDF3ViaNetCDF4DataTest, ScipyDataTest]:
         # for netCDF3 tests, expect the results to come back as characters
         expected['letters_nans'] = expected['letters_nans'].astype('S')
         expected['letters'] = expected['letters'].astype('S')
     with self.roundtrip(original) as actual:
         try:
             self.assertDatasetIdentical(expected, actual)
         except AssertionError:
             # Most stores use '' for nans in strings, but some don't
             # first try the ideal case (where the store returns exactly)
             # the original Dataset), then try a more realistic case.
             # ScipyDataTest, NetCDF3ViaNetCDF4DataTest and NetCDF4DataTest
             # all end up using this case.
             expected['letters_nans'][-1] = ''
             self.assertDatasetIdentical(expected, actual)

    def test_groupby(self):
        data = Dataset({'x': ('x', list('abc')),
                        'c': ('x', [0, 1, 0]),
                        'z': (['x', 'y'], np.random.randn(3, 5))})
        groupby = data.groupby('x')
        self.assertEqual(len(groupby), 3)
        expected_groups = {'a': 0, 'b': 1, 'c': 2}
        self.assertEqual(groupby.groups, expected_groups)
        expected_items = [('a', data.indexed(x=0)),
                          ('b', data.indexed(x=1)),
                          ('c', data.indexed(x=2))]
        self.assertEqual(list(groupby), expected_items)

        identity = lambda x: x
        for k in ['x', 'c', 'y']:
            actual = data.groupby(k, squeeze=False).apply(identity)
            self.assertEqual(data, actual)

        data = create_test_data()
        for n, (t, sub) in enumerate(list(data.groupby('dim1'))[:3]):
            self.assertEqual(data['dim1'][n], t)
            self.assertVariableEqual(data['var1'][n], sub['var1'])
            self.assertVariableEqual(data['var2'][n], sub['var2'])
            self.assertVariableEqual(data['var3'][:, n], sub['var3'])

        # TODO: test the other edge cases
        with self.assertRaisesRegexp(ValueError, 'must be 1 dimensional'):
            data.groupby('var1')
        with self.assertRaisesRegexp(ValueError, 'length does not match'):
            data.groupby(data['dim1'][:3])

 def test_preprocess_mfdataset(self):
     original = Dataset({'foo': ('x', np.random.randn(10))})
     with create_tmp_file() as tmp:
         original.to_netcdf(tmp)
         preprocess = lambda ds: ds.assign_coords(z=0)
         expected = preprocess(original)
         with open_mfdataset(tmp, preprocess=preprocess) as actual:
             self.assertDatasetIdentical(expected, actual)

    def test_pipe_tuple_error(self):
        df = Dataset({"A": ("x", [1, 2, 3])})
        f = lambda x, y: y
        with self.assertRaises(ValueError):
            df.pipe((f, "y"), x=1, y=0)

        with self.assertRaises(ValueError):
            df.A.pipe((f, "y"), x=1, y=0)

    def test_pipe_tuple(self):
        df = Dataset({"A": ("x", [1, 2, 3])})
        f = lambda x, y: y
        result = df.pipe((f, "y"), 0)
        self.assertDatasetIdentical(result, df)

        result = df.A.pipe((f, "y"), 0)
        self.assertDataArrayIdentical(result, df.A)

 def test_save_mfdataset_roundtrip(self):
     original = Dataset({'foo': ('x', np.random.randn(10))})
     datasets = [original.isel(x=slice(5)),
                 original.isel(x=slice(5, 10))]
     with create_tmp_file() as tmp1:
         with create_tmp_file() as tmp2:
             save_mfdataset(datasets, [tmp1, tmp2])
             with open_mfdataset([tmp1, tmp2]) as actual:
                 self.assertDatasetIdentical(actual, original)

    def test_reduce_argmin(self):
        # regression test for #205
        ds = Dataset({'a': ('x', [0, 1])})
        expected = Dataset({'a': ([], 0)})
        actual = ds.argmin()
        self.assertDatasetIdentical(expected, actual)

        actual = ds.argmin('x')
        self.assertDatasetIdentical(expected, actual)

    def test_pipe(self):
        df = Dataset({"A": ("x", [1, 2, 3])})
        f = lambda x, y: x ** y
        result = df.pipe(f, 2)
        expected = Dataset({"A": ("x", [1, 4, 9])})
        self.assertDatasetIdentical(result, expected)

        result = df.A.pipe(f, 2)
        self.assertDataArrayIdentical(result, expected.A)

    def test_weakrefs(self):
        example = Dataset({'foo': ('x', np.arange(5.0))})
        expected = example.rename({'foo': 'bar', 'x': 'y'})

        with create_tmp_file() as tmp_file:
            example.to_netcdf(tmp_file, engine='scipy')
            on_disk = open_dataset(tmp_file, engine='pynio')
            actual = on_disk.rename({'foo': 'bar', 'x': 'y'})
            del on_disk  # trigger garbage collection
            self.assertDatasetIdentical(actual, expected)

    def test_groupby_returns_new_type(self):
        data = Dataset({'z': (['x', 'y'], np.random.randn(3, 5))})

        actual = data.groupby('x').apply(lambda ds: ds['z'])
        expected = data['z']
        self.assertDataArrayIdentical(expected, actual)

        actual = data['z'].groupby('x').apply(lambda x: x.to_dataset())
        expected = data
        self.assertDatasetIdentical(expected, actual)

 def test_lock(self):
     original = Dataset({'foo': ('x', np.random.randn(10))})
     with create_tmp_file() as tmp:
         original.to_netcdf(tmp)
         with open_dataset(tmp, chunks=10) as ds:
             task = ds.foo.data.dask[ds.foo.data.name, 0]
             self.assertIsInstance(task[-1], type(Lock()))
         with open_mfdataset(tmp) as ds:
             task = ds.foo.data.dask[ds.foo.data.name, 0]
             self.assertIsInstance(task[-1], type(Lock()))

    def test_variable_order(self):
        # doesn't work with scipy or h5py :(
        ds = Dataset()
        ds['a'] = 1
        ds['z'] = 2
        ds['b'] = 3
        ds.coords['c'] = 4

        with self.roundtrip(ds) as actual:
            self.assertEqual(list(ds), list(actual))