xarray-spatial/xrspatial/datasets/__init__.py at a764e8897501a0440a214324e217874ec2dcbcbc · xarray-contrib/xarray-spatial · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import os

try:
    import dask.array as da
except ImportError:
    da = None

import datashader as ds
import noise
import numpy as np
import pandas as pd
import xarray as xr

__all__ = [
    'available',
    'get_data',
    'make_terrain',
]

_module_path = os.path.dirname(os.path.abspath(__file__))
_available_datasets = [p for p in next(os.walk(_module_path))[1]
                       if not p.startswith("__")]
available_datasets = _available_datasets


def get_data(dataset):
    """
    Open example multispectral band data.
    Parameters
    ----------
    dataset : str
        The name of the dataset. See ``xrspatial.datasets.available`` for
        all options.
    Examples
    --------
    >>>     xrspatial.datasets.get_data("sentinel-2")
    """
    data = {}
    if dataset in _available_datasets:
        folder_path = os.path.abspath(os.path.join(_module_path, dataset))
        band_files = [p for p in next(os.walk(folder_path))[2]]
        for band_file in band_files:
            array = xr.open_dataarray(os.path.join(folder_path, band_file))
            data[array.Name] = array
    else:
        msg = f'The dataset {dataset} is not available. '
        msg += f'Available folders are {available_datasets}.'
        raise ValueError(msg)
    return data


def make_terrain(
    shape=(1024, 1024),
    scale=100.0,
    octaves=6,
    persistence=0.5,
    lacunarity=2.0,
    chunks=(512, 512)
):
    """
    Generate a pseudo-random terrain data dask array.

    Parameters
    ----------
    shape : int or tuple of int, default=(1024, 1024)
        Output array shape.
    scale : float, default=100.0
        Noise factor scale.
    octaves : int, default=6
        Number of waves when generating the noise.
    persistence : float, default=0.5
        Amplitude of each successive octave relative.
    lacunarity : float, default=2.0
        Frequency of each successive octave relative.
    chunks : int or tuple of int, default=(512, 512)
        Number of samples on each block.

    Returns
    -------
    terrain : xarray.DataArray
        2D array of generated terrain values.
    """

    if da is None:
        raise Exception("make terrain requires dask.Array (pip install dask)")

    def _func(arr, block_id=None):
        block_ystart = block_id[0] * arr.shape[0]
        block_xstart = block_id[1] * arr.shape[1]
        out = np.zeros(arr.shape)
        for i in range(out.shape[0]):
            for j in range(out.shape[1]):
                out[i][j] = noise.pnoise2(
                    (block_ystart + i)/scale,
                    (block_xstart + j)/scale,
                    octaves=octaves,
                    persistence=persistence,
                    lacunarity=lacunarity,
                    repeatx=1024,
                    repeaty=1024,
                    base=42,
                )
        return out

    data = (
        da.zeros(shape=shape, chunks=chunks, dtype=np.float32)
        .map_blocks(_func, dtype=np.float32)
    )

    cvs = ds.Canvas(
        x_range=(0, 500),
        y_range=(0, 500),
        plot_width=shape[1],
        plot_height=shape[0],
    )

    hack_agg = cvs.points(pd.DataFrame({'x': [], 'y': []}), 'x', 'y')

    agg = xr.DataArray(
        data,
        name='terrain',
        coords=hack_agg.coords,
        dims=hack_agg.dims,
        attrs={'res': 1},
    )

    return agg