Exemplo n.º 1
0
 def test_set_new_value_mask_for_field(self):
     field = self.get_field()
     self.assertFalse(field.variables.first().value.mask.any())
     mask = np.array([True, False, True, False]).reshape(2, 2)
     set_new_value_mask_for_field(field, mask)
     self.assertTrue(field.variables.first().value.mask.any())
     self.assertEqual(field.variables.first().value.mask.sum(), 4)
Exemplo n.º 2
0
    def _get_spatial_operation_(self, attr, polygon, use_spatial_index=True, select_nearest=False):
        ref = getattr(self.spatial, attr)
        ret = copy(self)
        ret.spatial, slc = ref(polygon, return_indices=True, use_spatial_index=use_spatial_index,
                               select_nearest=select_nearest)
        slc = [slice(None), slice(None), slice(None)] + list(slc)
        ret.variables = self.variables.get_sliced_variables(slc)

        # Update the value mask with the geometry mask.
        set_new_value_mask_for_field(ret, ret.spatial.get_mask())

        return ret
Exemplo n.º 3
0
    def test_get_ocgis_field_from_esmpy_field(self):
        np.random.seed(1)
        temporal = TemporalDimension(value=[3000.0, 4000.0, 5000.0])
        level = VectorDimension(value=[10, 20, 30, 40])
        realization = VectorDimension(value=[100, 200])

        kwds = dict(
            crs=[None, CoordinateReferenceSystem(epsg=4326), Spherical()],
            with_mask=[False, True],
            with_corners=[False, True],
            dimensions=[False, True],
            drealization=[False, True],
            dtemporal=[False, True],
            dlevel=[False, True],
        )

        for k in self.iter_product_keywords(kwds):
            row = VectorDimension(value=[1.0, 2.0])
            col = VectorDimension(value=[3.0, 4.0])
            if k.with_corners:
                row.set_extrapolated_bounds()
                col.set_extrapolated_bounds()

            value_tmin = np.random.rand(2, 3, 4, 2, 2)
            tmin = Variable(value=value_tmin, name="tmin")
            variables = VariableCollection([tmin])
            grid = SpatialGridDimension(row=row, col=col)
            sdim = SpatialDimension(grid=grid, crs=k.crs)
            field = Field(variables=variables, spatial=sdim, temporal=temporal, level=level, realization=realization)
            if k.with_mask:
                mask = np.zeros(value_tmin.shape[-2:], dtype=bool)
                mask[0, 1] = True
                set_new_value_mask_for_field(field, mask)
                sdim.set_mask(mask)
                self.assertTrue(tmin.value.mask.any())
                self.assertTrue(sdim.get_mask().any())
            else:
                self.assertFalse(tmin.value.mask.any())
                self.assertFalse(sdim.get_mask().any())
            coll = SpatialCollection()
            coll[1] = {field.name: field}
            conv = ESMPyConverter([coll])
            efield = conv.write()

            if k.dimensions:
                dimensions = {}
                if k.drealization:
                    dimensions["realization"] = realization
                if k.dtemporal:
                    dimensions["temporal"] = temporal
                if k.dlevel:
                    dimensions["level"] = level
            else:
                dimensions = None

            ofield = get_ocgis_field_from_esmf_field(efield, crs=k.crs, dimensions=dimensions)

            self.assertIsInstance(ofield, Field)
            self.assertEqual(ofield.shape, efield.data.shape)

            # Test a default CRS is applied for the spherical case.
            if k.crs is None:
                self.assertEqual(ofield.spatial.crs, Spherical())

            if k.drealization and k.dimensions:
                target = realization.value
            else:
                target = np.array([1, 2])
            self.assertNumpyAll(ofield.realization.value, target)

            if k.dtemporal and k.dimensions:
                target = temporal.value
            else:
                target = np.array([1, 1, 1])
                with self.assertRaises(CannotFormatTimeError):
                    ofield.temporal.value_datetime
                self.assertFalse(ofield.temporal.format_time)
            self.assertNumpyAll(ofield.temporal.value, target)

            if k.dlevel and k.dimensions:
                target = level.value
            else:
                target = np.array([1, 2, 3, 4])
            self.assertNumpyAll(ofield.level.value, target)

            self.assertNumpyAll(field.spatial.grid.value, ofield.spatial.grid.value)
            if k.with_corners:
                self.assertIsNotNone(ofield.spatial.grid.corners)
                self.assertNumpyAll(field.spatial.grid.corners, ofield.spatial.grid.corners)

            try:
                self.assertEqual(ofield.spatial.crs, sdim.crs)
            except AssertionError:
                # A "None" "crs" argument results in a default coordinate system applied to the output OCGIS field.
                self.assertIsNone(k.crs)

            ofield_tmin_value = ofield.variables[efield.name].value
            for arr1, arr2 in itertools.combinations([tmin.value.data, efield.data, ofield_tmin_value.data], r=2):
                self.assertNumpyAll(arr1, arr2, check_arr_type=False)

            rows = list(ofield.get_iter())
            try:
                self.assertEqual(len(rows), len(value_tmin.flatten()))
            except AssertionError:
                self.assertTrue(k.with_mask)
                self.assertEqual(len(rows), len(tmin.value.compressed()))

            self.assertTrue(np.may_share_memory(ofield_tmin_value, efield.data))
            self.assertFalse(np.may_share_memory(ofield_tmin_value, tmin.value))