| 2157 | |
| 2158 | |
| 2159 | def test_store_regions(): |
| 2160 | d = da.ones((4, 4, 4), dtype=int, chunks=(2, 2, 2)) |
| 2161 | a, b = d + 1, d + 2 |
| 2162 | a = a[:, 1:, :].astype(float) |
| 2163 | |
| 2164 | region = (slice(None, None, 2), slice(None), [1, 2, 4, 5]) |
| 2165 | |
| 2166 | # Single region: |
| 2167 | at = np.zeros(shape=(8, 3, 6)) |
| 2168 | bt = np.zeros(shape=(8, 4, 6)) |
| 2169 | v = store([a, b], [at, bt], regions=region, compute=False) |
| 2170 | assert all([isinstance(a, Array) for a in v]) |
| 2171 | assert (at == 0).all() and (bt[region] == 0).all() |
| 2172 | results = dask.compute(*v) |
| 2173 | assert all([ev.size == 0 for ev in results]) |
| 2174 | assert (at[region] == 2).all() and (bt[region] == 3).all() |
| 2175 | assert not (bt == 3).all() and not (bt == 0).all() |
| 2176 | assert not (at == 2).all() and not (at == 0).all() |
| 2177 | |
| 2178 | # Multiple regions: |
| 2179 | at = np.zeros(shape=(8, 3, 6)) |
| 2180 | bt = np.zeros(shape=(8, 4, 6)) |
| 2181 | v = store([a, b], [at, bt], regions=[region, region], compute=False) |
| 2182 | assert (at == 0).all() and (bt[region] == 0).all() |
| 2183 | results = dask.compute(*v) |
| 2184 | assert all([ev.size == 0 for ev in results]) |
| 2185 | assert (at[region] == 2).all() and (bt[region] == 3).all() |
| 2186 | assert not (bt == 3).all() and not (bt == 0).all() |
| 2187 | assert not (at == 2).all() and not (at == 0).all() |
| 2188 | |
| 2189 | # Single region (keep result): |
| 2190 | for st_compute in [False, True]: |
| 2191 | at = np.zeros(shape=(8, 3, 6)) |
| 2192 | bt = np.zeros(shape=(8, 4, 6)) |
| 2193 | v = store( |
| 2194 | [a, b], [at, bt], regions=region, compute=st_compute, return_stored=True |
| 2195 | ) |
| 2196 | assert isinstance(v, tuple) |
| 2197 | assert all([isinstance(e, da.Array) for e in v]) |
| 2198 | if st_compute: |
| 2199 | for arr in v: |
| 2200 | assert_has_persisted_data(arr) |
| 2201 | else: |
| 2202 | assert (at == 0).all() and (bt[region] == 0).all() |
| 2203 | |
| 2204 | ar, br = v |
| 2205 | assert ar.dtype == a.dtype |
| 2206 | assert br.dtype == b.dtype |
| 2207 | assert ar.shape == a.shape |
| 2208 | assert br.shape == b.shape |
| 2209 | assert ar.chunks == a.chunks |
| 2210 | assert br.chunks == b.chunks |
| 2211 | |
| 2212 | ar, br = da.compute(ar, br) |
| 2213 | assert (at[region] == 2).all() and (bt[region] == 3).all() |
| 2214 | assert not (bt == 3).all() and not (bt == 0).all() |
| 2215 | assert not (at == 2).all() and not (at == 0).all() |
| 2216 | assert (br == 3).all() |