Python dask.array.ones() Examples

def test_multiple_sensors(self):
        """Test the background compositing from multiple sensor data."""
        from satpy.composites import BackgroundCompositor
        import numpy as np
        comp = BackgroundCompositor("name")

        # L mode images
        attrs = {'mode': 'L', 'area': 'foo'}
        foreground = xr.DataArray(np.array([[[1., 0.5],
                                             [0., np.nan]]]),
                                  dims=('bands', 'y', 'x'),
                                  coords={'bands': [c for c in attrs['mode']]},
                                  attrs=attrs.copy())
        foreground.attrs['sensor'] = 'abi'
        background = xr.DataArray(np.ones((1, 2, 2)), dims=('bands', 'y', 'x'),
                                  coords={'bands': [c for c in attrs['mode']]},
                                  attrs=attrs.copy())
        background.attrs['sensor'] = 'glm'
        res = comp([foreground, background])
        self.assertEqual(res.attrs['area'], 'foo')
        self.assertTrue(np.all(res == np.array([[1., 0.5], [0., 1.]])))
        self.assertEqual(res.attrs['mode'], 'L')
        self.assertEqual(res.attrs['sensor'], {'abi', 'glm'}) 

def test_simple(self):
        array0 = da.ones(shape=(10, 10, 40, 40), chunks=(5, 5, 5, 5))
        s0 = LazyDiffraction2D(array0)
        s0_r = s0.radial_average()
        assert (s0_r.data[:, :, :-1] == 1).all()

        data_shape = 2, 2, 11, 11
        array1 = np.zeros(data_shape)
        array1[:, :, 5, 5] = 1
        dask_array = da.from_array(array1, chunks=(1, 1, 1, 1))
        s1 = LazyDiffraction2D(dask_array)
        s1.axes_manager.signal_axes[0].offset = -5
        s1.axes_manager.signal_axes[1].offset = -5
        s1_r = s1.radial_average()
        assert np.all(s1_r.data[:, :, 0] == 1)
        assert np.all(s1_r.data[:, :, 1:] == 0) 

def test_add_missing_coordinates(self):
        """Test coordinate updating."""
        import dask.array as da
        from xarray import DataArray
        from pyresample.bilinear.xarr import XArrayResamplerBilinear

        resampler = XArrayResamplerBilinear(self.source_def, self.target_def,
                                            self.radius)
        bands = ['R', 'G', 'B']
        data = DataArray(da.ones((3, 10, 10)), dims=('bands', 'y', 'x'),
                         coords={'bands': bands,
                                 'y': np.arange(10), 'x': np.arange(10)})
        resampler._add_missing_coordinates(data)
        # X and Y coordinates should not change
        self.assertIsNone(resampler.out_coords_x)
        self.assertIsNone(resampler.out_coords_y)
        self.assertIsNone(resampler.out_coords['x'])
        self.assertIsNone(resampler.out_coords['y'])
        self.assertTrue('bands' in resampler.out_coords)
        self.assertTrue(np.all(resampler.out_coords['bands'] == bands)) 

def _mean4(data, offset=(0, 0), block_id=None):
    rows, cols = data.shape
    # we assume that the chunks except the first ones are aligned
    if block_id[0] == 0:
        row_offset = offset[0] % 2
    else:
        row_offset = 0
    if block_id[1] == 0:
        col_offset = offset[1] % 2
    else:
        col_offset = 0
    row_after = (row_offset + rows) % 2
    col_after = (col_offset + cols) % 2
    pad = ((row_offset, row_after), (col_offset, col_after))

    rows2 = rows + row_offset + row_after
    cols2 = cols + col_offset + col_after

    av_data = np.pad(data, pad, 'edge')
    new_shape = (int(rows2 / 2.), 2, int(cols2 / 2.), 2)
    data_mean = np.nanmean(av_data.reshape(new_shape), axis=(1, 3))
    data_mean = np.repeat(np.repeat(data_mean, 2, axis=0), 2, axis=1)
    data_mean = data_mean[row_offset:row_offset + rows, col_offset:col_offset + cols]
    return data_mean 

def test_compute(self):
        """Test count bucket resampler computation."""
        import dask.array as da
        # 1D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((5,))
        self.bucket.resampler.get_count.return_value = data
        res = self.bucket.compute(data)
        self.bucket.resampler.get_count.assert_called_once_with()
        self.assertEqual(res.shape, (1, 5))
        # 2D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((5, 5))
        self.bucket.resampler.get_count.return_value = data
        res = self.bucket.compute(data)
        self.bucket.resampler.get_count.assert_called_once_with()
        self.assertEqual(res.shape, (1, 5, 5))
        # 3D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((3, 5, 5))
        self.bucket.resampler.get_count.return_value = data[0, :, :]
        res = self.bucket.compute(data)
        self.assertEqual(res.shape, (3, 5, 5)) 

def test_compute(self):
        """Test fraction bucket resampler computation."""
        import dask.array as da
        import numpy as np

        self.bucket.resampler = mock.MagicMock()
        data = da.ones((3, 3))

        # No kwargs given
        _ = self.bucket.compute(data)
        self.bucket.resampler.get_fractions.assert_called_with(
            data,
            categories=None,
            fill_value=np.nan)
        # Custom kwargs
        _ = self.bucket.compute(data, categories=[1, 2], fill_value=0)
        self.bucket.resampler.get_fractions.assert_called_with(
            data,
            categories=[1, 2],
            fill_value=0)

        # Too many dimensions
        data = da.ones((3, 5, 5))
        with self.assertRaises(ValueError):
            _ = self.bucket.compute(data) 

def test_resample(self, pyresample_bucket):
        """Test fraction bucket resamplers resample method."""
        import xarray as xr
        import dask.array as da
        import numpy as np

        self.bucket.resampler = mock.MagicMock()
        self.bucket.precompute = mock.MagicMock()
        self.bucket.compute = mock.MagicMock()

        # Fractions return a dict
        data = xr.DataArray(da.ones((1, 5, 5)), dims=('bands', 'y', 'x'))
        arr = da.ones((5, 5))
        self.bucket.compute.return_value = {0: arr, 1: arr, 2: arr}
        res = self.bucket.resample(data)
        self.assertTrue('categories' in res.coords)
        self.assertTrue('categories' in res.dims)
        self.assertTrue(np.all(res.coords['categories'] == np.array([0, 1, 2]))) 

def test_delitem(comm):

    source = UniformCatalog(nbar=2e-4, BoxSize=512., seed=42, comm=comm)

    # add a test column
    test = numpy.ones(source.size)
    source['test'] = test

    # cannot delete hard coded column
    with pytest.raises(ValueError):
        del source['Position']

    # cannot delete missing column
    with pytest.raises(ValueError):
        del source['BAD_COLUMN']

    assert 'test' in source
    del source['test']
    assert 'test' not in source 

def test_transform(comm):
    cosmo = cosmology.Planck15

    data = numpy.ones(100, dtype=[
            ('Position', ('f4', 3)),
            ('Velocity', ('f4', 3))]
            )

    source = ArrayCatalog(data, BoxSize=100, Nmesh=32, comm=comm)

    source['Velocity'] = source['Position'] + source['Velocity']

    source['Position'] = source['Position'] + source['Velocity']

    # Position triggers  Velocity which triggers Position and Velocity
    # which resolves to the true data.
    # so total is 3.
    assert_allclose(source['Position'], 3)

    mesh = source.to_mesh() 

def setUp(self):
        """Create test data."""
        from satpy.composites import GenericCompositor
        self.comp = GenericCompositor(name='test')
        self.comp2 = GenericCompositor(name='test2', common_channel_mask=False)

        all_valid = np.ones((1, 2, 2))
        self.all_valid = xr.DataArray(all_valid, dims=['bands', 'y', 'x'])
        first_invalid = np.reshape(np.array([np.nan, 1., 1., 1.]), (1, 2, 2))
        self.first_invalid = xr.DataArray(first_invalid,
                                          dims=['bands', 'y', 'x'])
        second_invalid = np.reshape(np.array([1., np.nan, 1., 1.]), (1, 2, 2))
        self.second_invalid = xr.DataArray(second_invalid,
                                           dims=['bands', 'y', 'x'])
        wrong_shape = np.reshape(np.array([1., 1., 1.]), (1, 3, 1))
        self.wrong_shape = xr.DataArray(wrong_shape, dims=['bands', 'y', 'x']) 

def add_bands(data, bands):
    """Add bands so that they match *bands*."""
    # Add R, G and B bands, remove L band
    if 'L' in data['bands'].data and 'R' in bands.data:
        lum = data.sel(bands='L')
        # Keep 'A' if it was present
        if 'A' in data['bands']:
            alpha = data.sel(bands='A')
            new_data = (lum, lum, lum, alpha)
            new_bands = ['R', 'G', 'B', 'A']
            mode = 'RGBA'
        else:
            new_data = (lum, lum, lum)
            new_bands = ['R', 'G', 'B']
            mode = 'RGB'
        data = xr.concat(new_data, dim='bands', coords={'bands': new_bands})
        data['bands'] = new_bands
        data.attrs['mode'] = mode
    # Add alpha band
    if 'A' not in data['bands'].data and 'A' in bands.data:
        new_data = [data.sel(bands=band) for band in data['bands'].data]
        # Create alpha band based on a copy of the first "real" band
        alpha = new_data[0].copy()
        alpha.data = da.ones((data.sizes['y'],
                              data.sizes['x']),
                             chunks=new_data[0].chunks)
        # Rename band to indicate it's alpha
        alpha['bands'] = 'A'
        new_data.append(alpha)
        new_data = xr.concat(new_data, dim='bands')
        new_data.attrs['mode'] = data.attrs['mode'] + 'A'
        data = new_data

    return data 

def setUp(self):
        """Create test data."""
        from pyresample.geometry import AreaDefinition
        area = AreaDefinition('test', 'test', 'test',
                              {'proj': 'merc'}, 2, 2,
                              (-2000, -2000, 2000, 2000))
        bigger_area = AreaDefinition('test', 'test', 'test',
                                     {'proj': 'merc'}, 4, 4,
                                     (-2000, -2000, 2000, 2000))
        attrs = {'area': area,
                 'start_time': datetime(2018, 1, 1, 18),
                 'modifiers': tuple(),
                 'name': 'test_vis'}
        ds1 = xr.DataArray(da.ones((2, 2), chunks=2, dtype=np.float64),
                           attrs=attrs, dims=('y', 'x'),
                           coords={'y': [0, 1], 'x': [0, 1]})
        self.ds1 = ds1
        ds2 = xr.DataArray(da.ones((4, 4), chunks=2, dtype=np.float64),
                           attrs=attrs, dims=('y', 'x'),
                           coords={'y': [0, 0.5, 1, 1.5], 'x': [0, 0.5, 1, 1.5]})
        ds2.attrs['area'] = bigger_area
        self.ds2 = ds2
        self.sza = xr.DataArray(
            np.rad2deg(np.arccos(da.from_array([[0.0149581333, 0.0146694376], [0.0150812684, 0.0147925727]],
                                               chunks=2))),
            attrs={'area': area},
            dims=('y', 'x'),
            coords={'y': [0, 1], 'x': [0, 1]},
        ) 

def test_resample(self, pyresample_bucket):
        """Test bucket resamplers resample method."""
        import xarray as xr
        import dask.array as da
        self.bucket.resampler = mock.MagicMock()
        self.bucket.precompute = mock.MagicMock()
        self.bucket.compute = mock.MagicMock()

        # 1D input data
        data = xr.DataArray(da.ones((5,)), dims=('foo'), attrs={'bar': 'baz'})
        self.bucket.compute.return_value = da.ones((5, 5))
        res = self.bucket.resample(data)
        self.bucket.precompute.assert_called_once()
        self.bucket.compute.assert_called_once()
        self.assertEqual(res.shape, (5, 5))
        self.assertEqual(res.dims, ('y', 'x'))
        self.assertTrue('bar' in res.attrs)
        self.assertEqual(res.attrs['bar'], 'baz')

        # 2D input data
        data = xr.DataArray(da.ones((5, 5)), dims=('foo', 'bar'))
        self.bucket.compute.return_value = da.ones((5, 5))
        res = self.bucket.resample(data)
        self.assertEqual(res.shape, (5, 5))
        self.assertEqual(res.dims, ('y', 'x'))

        # 3D input data with 'bands' dim
        data = xr.DataArray(da.ones((1, 5, 5)), dims=('bands', 'foo', 'bar'),
                            coords={'bands': ['L']})
        self.bucket.compute.return_value = da.ones((1, 5, 5))
        res = self.bucket.resample(data)
        self.assertEqual(res.shape, (1, 5, 5))
        self.assertEqual(res.dims, ('bands', 'y', 'x'))
        self.assertEqual(res.coords['bands'], ['L'])

        # 3D input data with misc dim names
        data = xr.DataArray(da.ones((3, 5, 5)), dims=('foo', 'bar', 'baz'))
        self.bucket.compute.return_value = da.ones((3, 5, 5))
        res = self.bucket.resample(data)
        self.assertEqual(res.shape, (3, 5, 5))
        self.assertEqual(res.dims, ('foo', 'bar', 'baz')) 

def test_compute(self):
        """Test bucket resampler computation."""
        import dask.array as da
        # 1D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((5,))
        self.bucket.resampler.get_average.return_value = data
        res = self.bucket.compute(data, fill_value=2)
        self.bucket.resampler.get_average.assert_called_once_with(
            data,
            fill_value=2,
            mask_all_nan=False)
        self.assertEqual(res.shape, (1, 5))
        # 2D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((5, 5))
        self.bucket.resampler.get_average.return_value = data
        res = self.bucket.compute(data, fill_value=2)
        self.bucket.resampler.get_average.assert_called_once_with(
            data,
            fill_value=2,
            mask_all_nan=False)
        self.assertEqual(res.shape, (1, 5, 5))
        # 3D data
        self.bucket.resampler = mock.MagicMock()
        data = da.ones((3, 5, 5))
        self.bucket.resampler.get_average.return_value = data[0, :, :]
        res = self.bucket.compute(data, fill_value=2)
        self.assertEqual(res.shape, (3, 5, 5)) 

def __call__(self, projectables, *args, **kwargs):
        """Call the compositor."""
        if len(projectables) != 2:
            raise ValueError("Expected 2 datasets, got %d" % (len(projectables),))
        projectables = self.match_data_arrays(projectables)
        data_in = projectables[0]
        mask_in = projectables[1]
        mask_data = mask_in.data

        alpha_attrs = data_in.attrs.copy()
        if 'bands' in data_in.dims:
            data = [data_in.sel(bands=b) for b in data_in['bands'] if b != 'A']
        else:
            data = [data_in]

        # Create alpha band
        alpha = da.ones((data[0].sizes['y'],
                         data[0].sizes['x']),
                        chunks=data[0].chunks)

        for condition in self.conditions:
            method = condition['method']
            value = condition.get('value', None)
            if isinstance(value, str):
                value = _get_flag_value(mask_in, value)
            transparency = condition['transparency']
            mask = self._get_mask(method, value, mask_data)

            if transparency == 100.0:
                data = self._set_data_nans(data, mask, alpha_attrs)
            alpha_val = 1. - transparency / 100.
            alpha = da.where(mask, alpha_val, alpha)

        alpha = xr.DataArray(data=alpha, attrs=alpha_attrs,
                             dims=data[0].dims, coords=data[0].coords)
        data.append(alpha)

        res = super(MaskingCompositor, self).__call__(data, **kwargs)
        return res 

def test_mask(self, mask):
        s = Diffraction2D(np.ones((10, 10, 10, 10)))
        s1 = s.threshold_and_mask(mask=mask)
        slice0 = np.s_[:, :, mask[1] - mask[2] : mask[1] + mask[2] + 1, mask[0]]
        assert (s1.data[slice0] == 1.0).all()
        slice1 = np.s_[:, :, mask[1], mask[0] - mask[2] : mask[0] + mask[2] + 1]
        assert (s1.data[slice1] == 1.0).all()
        s1.data[slice0] = 0
        s1.data[slice1] = 0
        assert (s1.data == 0).all() 

def test_get_dataset_angles(self, get_angle, *mocks):
        from satpy.dataset import DatasetID
        from satpy.readers.avhrr_l1b_gaclac import ANGLES

        ones = np.ones((3, 3))
        get_angle.return_value = ones
        reader = self._get_reader_mocked()
        fh = self._get_fh_mocked(
            reader=reader,
            start_line=None,
            end_line=None,
            strip_invalid_coords=False,
            interpolate_coords=True
        )

        # With interpolation of coordinates
        for angle in ANGLES:
            key = DatasetID(angle)
            info = {'name': angle, 'standard_name': 'my_standard_name'}
            res = fh.get_dataset(key=key, info=info)
            exp = xr.DataArray(ones,
                               name=res.name,
                               dims=('y', 'x'),
                               coords={'acq_time': ('y', [0, 1, 2])})
            xr.testing.assert_equal(res, exp)

        # Without interpolation of coordinates
        fh.interpolate_coords = False
        for angle in ANGLES:
            key = DatasetID(angle)
            info = {'name': angle, 'standard_name': 'my_standard_name'}
            res = fh.get_dataset(key=key, info=info)
            self.assertTupleEqual(res.dims, ('y', 'x_every_eighth')) 

def test_resample_area_to_area_2d(self):
        """Resample area to area, 2d."""
        data = xr.DataArray(da.ones(self.src_area.shape, dtype=np.float64),
                            dims=['y', 'x'])
        res = self.resampler.compute(
            data, method='bil').compute(scheduler='single-threaded')
        assert res.shape == self.dst_area.shape
        assert np.allclose(res, 1) 

def setUp(self):
        """Create test data."""
        from satpy.composites import SingleBandCompositor
        self.comp = SingleBandCompositor(name='test')

        all_valid = np.ones((2, 2))
        self.all_valid = xr.DataArray(all_valid, dims=['y', 'x']) 

def test_big_value(self):
        data_shape = (5, 40, 40)
        big_value = 50000000
        array0 = np.ones(shape=data_shape) * big_value
        dask_array = da.from_array(array0, chunks=(2, 10, 10))
        s0 = LazyDiffraction2D(dask_array)
        s0_r = s0.radial_average()
        assert s0_r.axes_manager.navigation_shape == data_shape[:1]
        assert (s0_r.data[:, :-1] == big_value).all() 

def test_nav_1(self):
        data_shape = (5, 40, 40)
        array0 = da.ones(shape=data_shape, chunks=(5, 5, 5))
        s0 = LazyDiffraction2D(array0)
        s0_r = s0.radial_average()
        assert s0_r.axes_manager.navigation_shape == data_shape[:1]
        assert (s0_r.data[:, :-1] == 1).all() 

def test_different_shape(self):
        array = da.ones(shape=(7, 9, 30, 40), chunks=(3, 3, 5, 5))
        s = LazyDiffraction2D(array)
        s_r = s.radial_average()
        assert (s_r.data[:, :, :-2] == 1).all() 

def test_big_value(self):
        data_shape = (5, 40, 40)
        big_value = 50000000
        array0 = np.ones(shape=data_shape) * big_value
        s0 = Diffraction2D(array0)
        s0_r = s0.radial_average()
        assert s0_r.axes_manager.navigation_shape == data_shape[:1]
        assert (s0_r.data[:, :-1] == big_value).all() 

def test_nav_0(self):
        data_shape = (40, 40)
        array0 = np.ones(shape=data_shape)
        s0 = Diffraction2D(array0)
        s0_r = s0.radial_average()
        assert s0_r.axes_manager.navigation_dimension == 0
        assert (s0_r.data[:-1] == 1).all() 

def test_different_shape(self):
        array = np.ones(shape=(7, 9, 30, 40))
        s = Diffraction2D(array)
        s_r = s.radial_average()
        assert (s_r.data[:, :, :-2] == 1).all() 

def test_simple(self):
        array0 = np.ones(shape=(10, 10, 40, 40))
        s0 = Diffraction2D(array0)
        s0_r = s0.radial_average()
        assert (s0_r.data[:, :, :-1] == 1).all()

        data_shape = 2, 2, 11, 11
        array1 = np.zeros(data_shape)
        array1[:, :, 5, 5] = 1
        s1 = Diffraction2D(array1)
        s1.axes_manager.signal_axes[0].offset = -5
        s1.axes_manager.signal_axes[1].offset = -5
        s1_r = s1.radial_average()
        assert np.all(s1_r.data[:, :, 0] == 1)
        assert np.all(s1_r.data[:, :, 1:] == 0) 

def test_lazy_result(self):
        data = da.ones((10, 10, 20, 20), chunks=(10, 10, 10, 10))
        s_lazy = LazyDiffraction2D(data)
        s_lazy_com = s_lazy.center_of_mass(lazy_result=True)
        assert s_lazy_com._lazy 

def test_non_lazy_input(self):
        data = np.ones((100, 90))
        data[41, 21] = 0
        data[9, 81] = 50000
        s = Diffraction2D(data)
        s_dead_pixels = s.find_dead_pixels()
        s_hot_pixels = s.find_hot_pixels()
        s_bad_pixels = s_dead_pixels + s_hot_pixels
        s_corr = s.correct_bad_pixels(s_bad_pixels, lazy_result=False)
        assert s_dead_pixels.data.shape == data.shape
        assert not s_corr._lazy
        assert (s_corr.data == 1.0).all() 

def test_2d(self):
        data = np.ones((100, 90))
        data[41, 21] = 0
        data[9, 81] = 50000
        dask_array = da.from_array(data, chunks=(10, 10))
        s = LazyDiffraction2D(dask_array)
        s_dead_pixels = s.find_dead_pixels(lazy_result=True)
        s_hot_pixels = s.find_hot_pixels(lazy_result=True)
        s_bad_pixels = s_dead_pixels + s_hot_pixels
        s_corr = s.correct_bad_pixels(s_bad_pixels)
        assert s_dead_pixels.data.shape == data.shape
        assert s_dead_pixels._lazy
        s_corr.compute()
        assert (s_corr.data == 1.0).all() 

def test_different_chunks(self):
        data_array_dask = da.ones((6, 16, 100, 50), chunks=(6, 4, 50, 25))
        peak_array_dask = da.empty((6, 16), chunks=(3, 2), dtype=np.object)
        dt._intensity_peaks_image(data_array_dask, peak_array_dask, 5)