Skip to content

Commit

Permalink
black and fix failure for singleton dimensions
Browse files Browse the repository at this point in the history
  • Loading branch information
d-v-b committed Jan 24, 2022
1 parent 229d8bd commit d945b29
Show file tree
Hide file tree
Showing 3 changed files with 107 additions and 69 deletions.
10 changes: 4 additions & 6 deletions src/xarray_multiscale/multiscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,16 +333,14 @@ def get_downscale_depth(

_scale_factors = np.array(scale_factors).astype("int")
_shape = np.array(shape).astype("int")
if np.all(_scale_factors == 1):
result = 0
elif np.any(_scale_factors > _shape):
valid = (_scale_factors > 1)
if not valid.any():
result = 0
else:
if pad:
depths = np.ceil(logn(shape, scale_factors)).astype("int")
depths = np.ceil(logn(_shape[valid], _scale_factors[valid])).astype("int")
else:
lg = logn(shape, scale_factors)
depths = np.floor(logn(shape, scale_factors)).astype("int")
depths = np.floor(logn(_shape[valid], _scale_factors[valid])).astype("int")
result = min(depths)
return result

Expand Down
147 changes: 91 additions & 56 deletions tests/test_multiscale.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,13 @@
multiscale,
get_downscale_depth,
normalize_chunks,
ensure_minimum_chunks
ensure_minimum_chunks,
)
from xarray_multiscale.reducers import (
reshape_with_windows,
windowed_mean,
windowed_mode,
)
from xarray_multiscale.reducers import reshape_with_windows, windowed_mean, windowed_mode
import dask.array as da
import numpy as np
from xarray import DataArray
Expand All @@ -20,6 +24,8 @@

def test_downscale_depth():
assert get_downscale_depth((1,), (1,)) == 0
assert get_downscale_depth((2,), (3,)) == 0
assert get_downscale_depth((2, 1), (2, 1)) == 1
assert get_downscale_depth((2, 2, 2), (2, 2, 2)) == 1
assert get_downscale_depth((1, 2, 2), (2, 2, 2)) == 0
assert get_downscale_depth((4, 4, 4), (2, 2, 2)) == 2
Expand All @@ -31,72 +37,88 @@ def test_downscale_depth():
assert get_downscale_depth((1500, 5495, 5200), (2, 2, 2)) == 10


@pytest.mark.parametrize(("size", "scale"), ((10, 2), (11, 2), ((10,11), (2,3))))
@pytest.mark.parametrize(("size", "scale"), ((10, 2), (11, 2), ((10, 11), (2, 3))))
def test_adjust_shape(size, scale):
arr = DataArray(np.zeros(size))
padded = adjust_shape(arr, scale, mode="constant")
scale_array = np.array(scale)
old_shape_array = np.array(arr.shape)
new_shape_array = np.array(padded.shape)

if np.all((old_shape_array % scale_array) == 0):
assert np.array_equal(new_shape_array, old_shape_array)
else:
assert np.array_equal(new_shape_array, old_shape_array + ((scale_array - (old_shape_array % scale_array))))
assert np.array_equal(
new_shape_array,
old_shape_array + ((scale_array - (old_shape_array % scale_array))),
)

cropped = adjust_shape(arr, scale, mode="crop")
new_shape_array = np.array(cropped.shape)
if np.all((old_shape_array % scale_array) == 0):
assert np.array_equal(new_shape_array, old_shape_array)
else:
assert np.array_equal(new_shape_array, old_shape_array - (old_shape_array % scale_array))
assert np.array_equal(
new_shape_array, old_shape_array - (old_shape_array % scale_array)
)


def test_downscale_2d():
chunks = (2, 2)
scale = (2, 1)

data = DataArray(da.from_array(np.array(
[[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], dtype="uint8"
), chunks=chunks))
data = DataArray(
da.from_array(
np.array(
[[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], dtype="uint8"
),
chunks=chunks,
)
)
answer = DataArray(np.array([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]]))
downscaled = downscale(data, windowed_mean, scale, pad_mode='crop').compute()
downscaled = downscale(data, windowed_mean, scale, pad_mode="crop").compute()
assert np.array_equal(downscaled, answer)


def test_downscale_coords():
data = DataArray(np.zeros((10, 10)), dims=('x','y'), coords={'x': np.arange(10)})
scale_factors = (2,1)
data = DataArray(np.zeros((10, 10)), dims=("x", "y"), coords={"x": np.arange(10)})
scale_factors = (2, 1)
downscaled = downscale_coords(data, scale_factors)
answer = {'x': data['x'].coarsen({'x' : scale_factors[0]}).mean()}
answer = {"x": data["x"].coarsen({"x": scale_factors[0]}).mean()}

assert downscaled.keys() == answer.keys()
for k in downscaled:
assert_equal(answer[k], downscaled[k])

data = DataArray(np.zeros((10, 10)),
dims=('x','y'),
coords={'x': np.arange(10),
'y': 5 + np.arange(10)})
scale_factors = (2,1)
data = DataArray(
np.zeros((10, 10)),
dims=("x", "y"),
coords={"x": np.arange(10), "y": 5 + np.arange(10)},
)
scale_factors = (2, 1)
downscaled = downscale_coords(data, scale_factors)
answer = {'x': data['x'].coarsen({'x' : scale_factors[0]}).mean(),
'y' : data['y'].coarsen({'y' : scale_factors[1]}).mean()}

answer = {
"x": data["x"].coarsen({"x": scale_factors[0]}).mean(),
"y": data["y"].coarsen({"y": scale_factors[1]}).mean(),
}

assert downscaled.keys() == answer.keys()
for k in downscaled:
assert_equal(answer[k], downscaled[k])

data = DataArray(np.zeros((10, 10)),
dims=('x','y'),
coords={'x': np.arange(10),
'y': 5 + np.arange(10),
'foo' : 5})
scale_factors = (2,2)
data = DataArray(
np.zeros((10, 10)),
dims=("x", "y"),
coords={"x": np.arange(10), "y": 5 + np.arange(10), "foo": 5},
)
scale_factors = (2, 2)
downscaled = downscale_coords(data, scale_factors)
answer = {'x': data['x'].coarsen({'x' : scale_factors[0]}).mean(),
'y' : data['y'].coarsen({'y' : scale_factors[1]}).mean(),
'foo': data['foo']}

answer = {
"x": data["x"].coarsen({"x": scale_factors[0]}).mean(),
"y": data["y"].coarsen({"y": scale_factors[1]}).mean(),
"foo": data["foo"],
}

assert downscaled.keys() == answer.keys()
for k in downscaled:
assert_equal(answer[k], downscaled[k])
Expand All @@ -106,7 +128,7 @@ def test_invalid_multiscale():
with pytest.raises(ValueError):
downscale_dask(np.arange(10), windowed_mean, (3,))
with pytest.raises(ValueError):
downscale_dask(np.arange(16).reshape(4,4), windowed_mean, (3,3))
downscale_dask(np.arange(16).reshape(4, 4), windowed_mean, (3, 3))


def test_multiscale():
Expand Down Expand Up @@ -163,14 +185,20 @@ def test_chunking():
assert m.data.chunksize == chunks or m.data.chunksize == m.data.shape

chunks = (3,) * ndim
multi = multiscale(base_array, reducer, 2, chunks=chunks, chunk_mode='minimum')
multi = multiscale(base_array, reducer, 2, chunks=chunks, chunk_mode="minimum")
for m in multi:
assert np.greater_equal(m.data.chunksize, chunks).all() or m.data.chunksize == m.data.shape
assert (
np.greater_equal(m.data.chunksize, chunks).all()
or m.data.chunksize == m.data.shape
)

chunks = 3
multi = multiscale(base_array, reducer, 2, chunks=chunks, chunk_mode='minimum')
multi = multiscale(base_array, reducer, 2, chunks=chunks, chunk_mode="minimum")
for m in multi:
assert np.greater_equal(m.data.chunksize, (chunks,) * ndim).all() or m.data.chunksize == m.data.shape
assert (
np.greater_equal(m.data.chunksize, (chunks,) * ndim).all()
or m.data.chunksize == m.data.shape
)


def test_depth():
Expand All @@ -182,16 +210,16 @@ def test_depth():
assert len(full) == 5

partial = multiscale(base_array, reducer, 2, depth=-2)
assert len(partial) == len(full) - 1
[assert_equal(a,b) for a,b in zip(full, partial)]
assert len(partial) == len(full) - 1
[assert_equal(a, b) for a, b in zip(full, partial)]

partial = multiscale(base_array, reducer, 2, depth=2)
assert len(partial) == 3
[assert_equal(a,b) for a,b in zip(full, partial)]
assert len(partial) == 3
[assert_equal(a, b) for a, b in zip(full, partial)]

partial = multiscale(base_array, reducer, 2, depth=0)
assert len(partial) == 1
[assert_equal(a,b) for a,b in zip(full, partial)]
assert len(partial) == 1
[assert_equal(a, b) for a, b in zip(full, partial)]


def test_coords():
Expand All @@ -215,23 +243,23 @@ def test_coords():


def test_normalize_chunks():
data = DataArray(da.zeros((4,6), chunks=(1,1)))
assert normalize_chunks(data, {'dim_0' : 2, 'dim_1' : 1}) == (2,1)
data = DataArray(da.zeros((4, 6), chunks=(1, 1)))
assert normalize_chunks(data, {"dim_0": 2, "dim_1": 1}) == (2, 1)


def test_ensure_minimum_chunks():
data = da.zeros((4,6), chunks=(1,1))
assert ensure_minimum_chunks(data, (2,2)) == (2,2)
data = da.zeros((4, 6), chunks=(1, 1))
assert ensure_minimum_chunks(data, (2, 2)) == (2, 2)

data = da.zeros((4,6), chunks=(4,1))
assert ensure_minimum_chunks(data, (2,2)) == (4,2)
data = da.zeros((4, 6), chunks=(4, 1))
assert ensure_minimum_chunks(data, (2, 2)) == (4, 2)


def test_broadcast_to_rank():
assert broadcast_to_rank(2, 1) == (2,)
assert broadcast_to_rank(2, 2) == (2,2)
assert broadcast_to_rank((2,3), 2) == (2,3)
assert broadcast_to_rank({0 : 2}, 3) == (2,1,1)
assert broadcast_to_rank(2, 2) == (2, 2)
assert broadcast_to_rank((2, 3), 2) == (2, 3)
assert broadcast_to_rank({0: 2}, 3) == (2, 1, 1)


def test_align_chunks():
Expand All @@ -243,11 +271,18 @@ def test_align_chunks():
rechunked = align_chunks(data, scale_factors=(2,))
assert rechunked.chunks == ((2,) * 5,)

data = da.arange(10, chunks=(1,1,3,5))
data = da.arange(10, chunks=(1, 1, 3, 5))
rechunked = align_chunks(data, scale_factors=(2,))
assert rechunked.chunks == ((2, 2, 2, 4,),)
assert rechunked.chunks == (
(
2,
2,
2,
4,
),
)


def test_reshape_with_windows():
data = np.arange(36).reshape(6,6)
assert reshape_with_windows(data, (2,2)).shape == (3,2,3,2)
data = np.arange(36).reshape(6, 6)
assert reshape_with_windows(data, (2, 2)).shape == (3, 2, 3, 2)
19 changes: 12 additions & 7 deletions tests/test_reducers.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,18 @@ def test_windowed_mode():
results = windowed_mode(data, (4,))
assert np.array_equal(results, answer)

data = np.arange(16).reshape(4,4) % 3
answer = np.array([[1,0],[0,2]])
results = windowed_mode(data, (2,2))
data = np.arange(16).reshape(4, 4) % 3
answer = np.array([[1, 0], [0, 2]])
results = windowed_mode(data, (2, 2))
assert np.array_equal(results, answer)


def test_windowed_mean():
data = np.arange(16).reshape(4,4) % 2
answer = np.array([[0.5, 0.5],[0.5, 0.5]])
results = windowed_mean(data, (2,2))
assert np.array_equal(results, answer)
data = np.arange(16).reshape(4, 4) % 2
answer = np.array([[0.5, 0.5], [0.5, 0.5]])
results = windowed_mean(data, (2, 2))
assert np.array_equal(results, answer)

data = np.arange(16).reshape(4, 4, 1) % 2
answer = np.array([[0.5, 0.5], [0.5, 0.5]]).reshape((2,2,1))
results = windowed_mean(data, (2, 2, 1))

0 comments on commit d945b29

Please sign in to comment.