Python ocgis.RequestDataset类代码示例

    def __iter__(self):
        non_iterables = [AbstractRequestObject, dict, Field]
        if env.USE_ESMF:
            from ocgis.regrid.base import ESMF
            non_iterables.append(ESMF.Field)

        if isinstance(self._value, tuple(non_iterables)):
            to_itr = [self._value]
        else:
            to_itr = self._value
        for uid, element in enumerate(to_itr, start=1):
            if isinstance(element, dict):
                element = RequestDataset(**element)

            if env.USE_ESMF and isinstance(element, ESMF.Field):
                from ocgis.regrid.base import get_ocgis_field_from_esmf_field
                element = get_ocgis_field_from_esmf_field(element)

            try:
                element = element.copy()
            except AttributeError:
                element = copy(element)

            if element.uid is None:
                element.uid = uid
                # TODO: Remove me once the driver does not accept request datasets at initialization.
                # Try to change the driver UID.
                try:
                    element.driver.rd.uid = uid
                except AttributeError:
                    # The field driver does not keep a copy of the request dataset.
                    if hasattr(element.driver, 'rd'):
                        raise

            yield element

 def test_load_bounds_datetime_after_slicing(self):
     ref_test = self.test_data['cancm4_tas']
     uri = self.test_data.get_uri('cancm4_tas')
     rd = RequestDataset(variable=ref_test['variable'],uri=uri)
     field = rd.get()
     slced = field[:,10:130,:,4:7,100:37]
     self.assertEqual(slced.temporal.bounds_datetime.shape,(120,2))

    def test_load_geometry_subset(self):
        ref_test = self.test_data['cancm4_tas']
        uri = self.test_data.get_uri('cancm4_tas')

        states = self.get_2d_state_boundaries_sdim()
        ca = states[:,states.properties['STATE_NAME'] == 'California']
        self.assertTrue(ca.properties['STATE_NAME'] == 'California')
        ca.crs.unwrap(ca)
        ca = ca.geom.polygon.value[0,0]

        for u in [True,False]:
            try:
                rd = RequestDataset(variable=ref_test['variable'],uri=uri,alias='foo')
                field = rd.get()
                ca_sub = field.get_intersects(ca,use_spatial_index=u)
                self.assertEqual(ca_sub.shape,(1, 3650, 1, 5, 4))
                self.assertTrue(ca_sub.variables['foo'].value.mask.any())
                self.assertFalse(field.spatial.uid.mask.any())
                self.assertFalse(field.spatial.get_mask().any())

                ca_sub = field.get_intersects(ca.envelope,use_spatial_index=u)
                self.assertEqual(ca_sub.shape,(1, 3650, 1, 5, 4))
                self.assertFalse(ca_sub.variables['foo'].value.mask.any())

                rd = RequestDataset(variable=ref_test['variable'],uri=uri,alias='foo',time_region={'year':[2007]})
                field = rd.get()
                ca_sub = field.get_intersects(ca,use_spatial_index=u)
                self.assertEqual(ca_sub.shape,(1, 365, 1, 5, 4))
                self.assertEqual(set([2007]),set([d.year for d in ca_sub.temporal.value_datetime]))
            except ImportError:
                with self.assertRaises(ImportError):
                    import_module('rtree')

    def test_insert_weighted(self):
        gs = self.fixture_grid_chunker()

        dst_master_path = self.get_temporary_file_path('out.nc')
        gs.dst_grid.parent.write(dst_master_path)

        dst_master = RequestDataset(dst_master_path).get()
        desired_sums = {}
        for data_variable in dst_master.data_variables:
            dv_sum = data_variable.get_value().sum()
            desired_sums[data_variable.name] = dv_sum
            self.assertNotEqual(dv_sum, 0)
            data_variable.get_value()[:] = 0
        dst_master.write(dst_master_path, write_mode=MPIWriteMode.FILL)
        dst_master = RequestDataset(dst_master_path).get()
        for data_variable in dst_master.data_variables:
            self.assertEqual(data_variable.get_value().sum(), 0)

        gs.write_chunks()

        index_path = gs.create_full_path_from_template('index_file')
        gs.insert_weighted(index_path, self.current_dir_output, dst_master_path)

        actual_sums = {}
        dst_master_inserted = RequestDataset(dst_master_path).get()
        for data_variable in dst_master_inserted.data_variables:
            dv_value = data_variable.get_value()
            dv_sum = dv_value.sum()
            actual_sums[data_variable.name] = dv_sum
        for k, v in list(actual_sums.items()):
            self.assertAlmostEqual(v, desired_sums[k])

    def test_system_create_field_dimensioned_variables(self):
        """Test data is appropriately tagged to identify dimensioned variables."""

        path = self.get_temporary_file_path('foo.nc')
        time = TemporalVariable(value=[1, 2, 3], dimensions='time')
        x = Variable(name='x', value=[10, 20], dimensions='x')
        y = Variable(name='y', value=[30, 40, 50, 60], dimensions='y')
        data1 = Variable(name='data1', value=np.random.rand(3, 4, 2), dimensions=['time', 'y', 'x'])
        data2 = Variable(name='data2', value=np.random.rand(3, 4, 2), dimensions=['time', 'y', 'x'])
        data3 = Variable(name='data3', value=[11, 12, 13], dimensions=['time'])
        field = Field(time=time, grid=Grid(x, y), variables=[data1, data2, data3])
        field.write(path)

        # Test dimensioned variables are read from a file with appropriate metadata.
        rd = RequestDataset(path)
        self.assertEqual(rd.variable, ('data1', 'data2'))
        read_field = rd.get()
        actual = get_variable_names(read_field.data_variables)
        self.assertEqual(actual, ('data1', 'data2'))

        # Test dimensioned variables are overloaded.
        rd = RequestDataset(path, variable='data2')
        read_field = rd.get()
        actual = get_variable_names(read_field.data_variables)
        self.assertEqual(actual, ('data2',))

    def test_system_rotated_pole_spherical_subsetting(self):
        """Test rotated pole coordinates are left alone during a subset (no mask applied)."""

        def _run_mask_test_(target):
            for vn in ['rlat', 'rlon']:
                self.assertIsNone(target[vn].get_mask())

        rd = RequestDataset(metadata=self.fixture_rotated_spherical_metadata)
        field = rd.create_field()
        _run_mask_test_(field)
        for ctr, vn in enumerate(['lat', 'lon', 'rlat', 'rlon']):
            var = field[vn]
            var.get_value()[:] = np.arange(var.size).reshape(var.shape) + (ctr * 10)
        path = self.get_temporary_file_path('foo.nc')
        field.write(path)

        new_field = RequestDataset(path).create_field()
        _run_mask_test_(new_field)
        subset_geom = box(*new_field.grid.extent)
        subset_field = new_field.grid.get_intersects(subset_geom, optimized_bbox_subset=True).parent
        _run_mask_test_(subset_field)

        path2 = self.get_temporary_file_path('foo2.nc')
        subset_field.write(path2)
        in_subset_field = RequestDataset(path2).create_field()
        _run_mask_test_(in_subset_field)

    def test_write_variable_fill_value_is_maintained(self):

        if vm.size != 4:
            raise SkipTest('vm.size != 4')

        dist = OcgDist()
        dim = dist.create_dimension('dim', 8, dist=True)
        dist.update_dimension_bounds()

        var = Variable(name='var', dimensions=dim, fill_value=2.)
        var.v()[0] = 1
        var.v()[1] = 2
        var.get_mask(create=True, check_value=True)

        if vm.rank == 0:
            path = self.get_temporary_file_path('foo.nc')
        else:
            path = None
        path = vm.bcast(path)

        var.parent.write(path)

        # if vm.rank == 0:
        #     self.ncdump(path, header_only=False)

        with vm.scoped('read test', [0]):
            if not vm.is_null:
                invar = RequestDataset(path).create_field()['var']
                self.assertEqual(invar.get_mask().sum(), 4)
                self.assertEqual(invar.fill_value, 2.)

    def test_write_variable_collection_netcdf4_mpi(self):
        # TODO: TEST: Test writing a grouped netCDF file in parallel.

        self.add_barrier = False
        if not env.USE_NETCDF4_MPI:
            raise SkipTest('not env.USE_NETCDF4_MPI')
        path = self.create_rank_valued_netcdf()

        # if vm.rank == 0:
        #     self.ncdump(path, header_only=False)

        rd = RequestDataset(path, driver='netcdf')
        rd.metadata['dimensions']['dist_dim']['dist'] = True

        field = rd.get()

        # self.barrier_print(field['data'].get_value())

        if vm.rank == 0:
            actual_path = self.get_temporary_file_path('actual_mpi.nc')
        else:
            actual_path = None
        actual_path = vm.bcast(actual_path)

        # self.barrier_print('before field.write')
        field.write(actual_path)
        # self.barrier_print('after field.write')

        if vm.rank == 0:
            # self.ncdump(actual_path, header_only=False)
            self.assertNcEqual(actual_path, path)

    def test_load(self):
        ref_test = self.test_data['cancm4_tas']
        uri = self.test_data.get_uri('cancm4_tas')
        rd = RequestDataset(variable=ref_test['variable'],uri=uri)
        field = rd.get()
        ds = nc.Dataset(uri,'r')

        self.assertEqual(field.level,None)
        self.assertEqual(field.spatial.crs,WGS84())

        tv = field.temporal.value
        test_tv = ds.variables['time'][:]
        self.assertNumpyAll(tv,test_tv)
        self.assertNumpyAll(field.temporal.bounds,ds.variables['time_bnds'][:])

        tdt = field.temporal.value_datetime
        self.assertEqual(tdt[4],dt(2001,1,5,12))
        self.assertNumpyAll(field.temporal.bounds_datetime[1001],np.array([dt(2003,9,29),dt(2003,9,30)]))

        rv = field.temporal.value_datetime[100]
        rb = field.temporal.bounds_datetime[100]
        self.assertTrue(all([rv > rb[0],rv < rb[1]]))

        self.assertEqual(field.temporal.extent_datetime,(datetime.datetime(2001,1,1),datetime.datetime(2011,1,1)))

        ds.close()

    def test_system_changing_field_name(self):
        path1 = self.get_temporary_file_path('foo1.nc')
        path2 = self.get_temporary_file_path('foo2.nc')

        vc1 = VariableCollection(name='vc1')
        var1 = Variable('var1', value=[1, 2, 3], dimensions='three', parent=vc1)

        vc2 = VariableCollection(name='vc2')
        vc1.add_child(vc2)
        var2 = Variable('var2', value=[4, 5, 6, 7], dimensions='four', parent=vc2)

        vc1.write(path1)

        rd = RequestDataset(path1)
        # rd.inspect()
        nvc = rd.create_raw_field()
        nvc2 = nvc.children['vc2']
        self.assertIsNone(nvc2['var2']._value)
        self.assertEqual(nvc2.name, 'vc2')
        nvc2.set_name('extraordinary')
        self.assertIsNotNone(nvc2['var2'].get_value())
        self.assertEqual(nvc2['var2'].get_value().tolist(), [4, 5, 6, 7])

        nvc.write(path2)
        rd2 = RequestDataset(path2)
        # rd2.inspect()
        n2vc = rd2.create_raw_field()
        self.assertEqual(n2vc.children[nvc2.name].name, nvc2.name)

 def test_load_time_range(self):
     ref_test = self.test_data['cancm4_tas']
     uri = self.test_data.get_uri('cancm4_tas')
     rd = RequestDataset(variable=ref_test['variable'],uri=uri,time_range=[dt(2005,2,15),dt(2007,4,18)])
     field = rd.get()
     self.assertEqual(field.temporal.value_datetime[0],dt(2005, 2, 15, 12, 0))
     self.assertEqual(field.temporal.value_datetime[-1],dt(2007, 4, 18, 12, 0))
     self.assertEqual(field.shape,(1,793,1,64,128))

 def test_init(self):
     path = self.get_temporary_file_path('foo.nc')
     with self.nc_scope(path, 'w') as ds:
         ds.createDimension('a', 2)
     rd = RequestDataset(uri=path, driver='netcdf')
     self.assertIsInstance(rd.driver, DriverNetcdf)
     vc = rd.create_raw_field()
     self.assertEqual(len(vc), 0)

 def test_init_metadata_only(self):
     metadata = {'variables': {'foo': {}}}
     rd = RequestDataset(metadata=metadata)
     self.assertEqual(rd.driver.key, DriverKey.NETCDF_CF)
     self.assertIsNone(rd.uri)
     self.assertEqual(rd.metadata, metadata)
     field = rd.create_field()
     self.assertIn('foo', field.keys())

    def test_redistribute_by_src_idx(self):
        if vm.size != 4:
            raise SkipTest('vm.size != 4')

        dist = OcgDist()
        dim1 = dist.create_dimension('dim1', 5 * vm.size, dist=True)
        dim2 = dist.create_dimension('dim2', 2, dist=False)
        dist.update_dimension_bounds()

        rank_value = np.arange(5) + (10 * (vm.rank + 1))
        var1 = Variable(name='dvar1', value=rank_value, dimensions=dim1)
        var2 = Variable(name='dvar2', dimensions=[dim1, dim2])
        var1.parent.add_variable(var2)
        path = self.get_temporary_file_path('out.nc')
        var1.parent.write(path)

        desired_idx = np.array([1, 7, 9, 10, 14])
        vdesired_value = variable_gather(var1)
        if vm.rank == 0:
            desired_value = vdesired_value.get_value()[desired_idx]

        desired_idx_ranks = {0: slice(1, 2),
                             1: [2, 4],
                             2: [0, 4]}

        rd = RequestDataset(path)
        rd.metadata['dimensions'][dim1.name]['dist'] = True
        field = rd.create_field()

        indvar = field[var1.name]
        field[var2.name].load()

        try:
            rank_slice = desired_idx_ranks[vm.rank]
        except KeyError:
            sub = Variable(is_empty=True)
        else:
            sub = indvar[rank_slice]

        self.barrier_print(sub.is_empty)

        redistribute_by_src_idx(indvar, dim1.name, sub.dimensions_dict.get(dim1.name))

        with vm.scoped_by_emptyable('gather for test', indvar):
            if vm.is_null:
                self.assertIn(vm.rank_global, [2, 3])
            else:
                self.assertIn(vm.rank_global, [0, 1])
                for v in [indvar, indvar.parent[var2.name]]:
                    self.assertIsNone(v._value)
                    self.assertIsNone(v._mask)
                    self.assertIsNone(v._is_empty)
                    self.assertFalse(v._has_initialized_value)
                self.rank_print(indvar)
                actual_value = variable_gather(indvar)
                if vm.rank == 0:
                    actual_value = actual_value.get_value()
                    self.assertNumpyAll(actual_value, desired_value)

    def test_esmf(self):
        rd1 = RequestDataset(**self.get_dataset())
        rd2 = deepcopy(rd1)
        ops = OcgOperations(dataset=rd1, regrid_destination=rd2, output_format='nc')
        ret = ops.execute()

        actual_value = RequestDataset(ret).get().data_variables[0].get_value()
        desired_value = rd1.get().data_variables[0].get_value()
        self.assertNumpyAllClose(actual_value, desired_value)