def test_upsample_multi_channel(tmp_path: Path) -> None:
    num_channels = 3
    size = (32, 32, 10)
    source_data = (
        128 * np.random.randn(num_channels, size[0], size[1], size[2])
    ).astype("uint8")

    ds = Dataset(tmp_path / "multi-channel-test", (1, 1, 1))
    l = ds.add_layer(
        "color",
        COLOR_CATEGORY,
        dtype_per_channel="uint8",
        num_channels=num_channels,
    )
    mag2 = l.add_mag("2", chunks_per_shard=32)

    mag2.write(source_data)
    assert np.any(source_data != 0)

    l._initialize_mag_from_other_mag("1", mag2, False)

    upsample_cube_job(
        (mag2.get_view(), l.get_mag("1").get_view(), 0),
        [0.5, 0.5, 0.5],
        BUFFER_SHAPE,
    )

    channels = []
    for channel_index in range(num_channels):
        channels.append(upsample_cube(source_data[channel_index], [2, 2, 2]))
    joined_buffer = np.stack(channels)

    target_buffer = l.get_mag("1").read()
    assert np.any(target_buffer != 0)
    assert np.all(target_buffer == joined_buffer)
def test_default_anisotropic_voxel_size(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "default_anisotropic_voxel_size",
                 voxel_size=(85, 85, 346))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1)
    mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))

    layer.downsample(Mag(1), None, "median", True)
    assert sorted(layer.mags.keys()) == [Mag("1"), Mag("2-2-1"), Mag("4-4-1")]
def test_downsample_with_invalid_mag_list(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1)
    mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))

    with pytest.raises(AssertionError):
        layer.downsample_mag_list(
            from_mag=Mag(1),
            target_mags=[Mag(1),
                         Mag([1, 1, 2]),
                         Mag([2, 2, 1]),
                         Mag(2)],
        )
def test_downsample_mag_list(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1)
    mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))

    target_mags = [Mag([4, 4, 8]),
                   Mag(2), Mag([32, 32, 8]),
                   Mag(32)]  # unsorted list

    layer.downsample_mag_list(from_mag=Mag(1), target_mags=target_mags)

    for m in target_mags:
        assert m in layer.mags
def test_default_parameter(tmp_path: Path) -> None:
    target_path = tmp_path / "downsaple_default"

    ds = Dataset(target_path, voxel_size=(1, 1, 1))
    layer = ds.add_layer("color",
                         COLOR_CATEGORY,
                         dtype_per_channel="uint8",
                         num_channels=3)
    mag = layer.add_mag("2")
    mag.write(data=(np.random.rand(3, 10, 20, 30) * 255).astype(np.uint8))
    layer.downsample()

    # The max_mag is Mag(4) in this case (see test_default_max_mag)
    assert sorted(layer.mags.keys()) == [Mag("2"), Mag("4")]
def test_upsampling(tmp_path: Path) -> None:
    ds = Dataset(tmp_path, voxel_size=(1, 1, 1))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag([4, 4, 2])
    mag.write(
        absolute_offset=(10 * 4, 20 * 4, 40 * 2),
        data=(np.random.rand(46, 45, 27) * 255).astype(np.uint8),
    )
    layer.upsample(
        from_mag=Mag([4, 4, 2]),
        finest_mag=Mag(1),
        compress=False,
        sampling_mode=SamplingModes.ANISOTROPIC,
        buffer_edge_len=64,
        args=None,
    )
def test_downsample_2d(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "downsample_compressed", voxel_size=(1, 1, 2))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1, chunk_size=8, chunks_per_shard=8)
    # write 2D data with all values set to "123"
    mag.write(data=(np.ones((100, 100, 1)) * 123).astype(np.uint8))
    with pytest.warns(Warning):
        # This call produces a warning because only the mode "CONSTANT_Z" makes sense for 2D data.
        layer.downsample(
            from_mag=Mag(1),
            coarsest_mag=Mag(2),
            sampling_mode=SamplingModes.
            ISOTROPIC,  # this mode is intentionally not "CONSTANT_Z" for this test
        )
    assert Mag("2-2-1") in layer.mags
    assert np.all(layer.get_mag(
        Mag("2-2-1")).read() == 123)  # The data is not darkened
def test_downsample_multi_channel(tmp_path: Path) -> None:
    num_channels = 3
    size = (32, 32, 10)
    source_data = (128 * np.random.randn(num_channels, size[0], size[1],
                                         size[2])).astype("uint8")

    ds = Dataset(tmp_path / "multi-channel-test", (1, 1, 1))
    l = ds.add_layer(
        "color",
        COLOR_CATEGORY,
        dtype_per_channel="uint8",
        num_channels=num_channels,
    )
    mag1 = l.add_mag("1", chunks_per_shard=32)

    print("writing source_data shape", source_data.shape)
    mag1.write(source_data)
    assert np.any(source_data != 0)

    mag2 = l._initialize_mag_from_other_mag("2", mag1, False)

    downsample_cube_job(
        (l.get_mag("1").get_view(), l.get_mag("2").get_view(), 0),
        Vec3Int(2, 2, 2),
        InterpolationModes.MAX,
        BUFFER_SHAPE,
    )

    channels = []
    for channel_index in range(num_channels):
        channels.append(
            downsample_cube(source_data[channel_index], [2, 2, 2],
                            InterpolationModes.MAX))
    joined_buffer = np.stack(channels)

    target_buffer = mag2.read()
    assert np.any(target_buffer != 0)
    assert np.all(target_buffer == joined_buffer)
def upsample_test_helper(tmp_path: Path, use_compress: bool) -> None:
    ds = Dataset(tmp_path, voxel_size=(10.5, 10.5, 5))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag2 = layer.add_mag([2, 2, 2])

    offset = Vec3Int(WKW_CUBE_SIZE, 2 * WKW_CUBE_SIZE, 0)

    mag2.write(
        absolute_offset=offset,
        data=(np.random.rand(*BUFFER_SHAPE) * 255).astype(np.uint8),
    )
    mag1 = layer._initialize_mag_from_other_mag("1-1-2", mag2, use_compress)

    source_buffer = mag2.read(
        absolute_offset=offset,
        size=BUFFER_SHAPE,
    )[0]
    assert np.any(source_buffer != 0)

    upsample_cube_job(
        (
            mag2.get_view(absolute_offset=offset, size=BUFFER_SHAPE),
            mag1.get_view(
                absolute_offset=offset,
                size=BUFFER_SHAPE,
            ),
            0,
        ),
        [0.5, 0.5, 1.0],
        BUFFER_SHAPE,
    )

    assert np.any(source_buffer != 0)

    target_buffer = mag1.read(absolute_offset=offset, size=BUFFER_SHAPE)[0]
    assert np.any(target_buffer != 0)

    assert np.all(target_buffer == upsample_cube(source_buffer, [2, 2, 1]))
def test_downsample_compressed(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "downsample_compressed", voxel_size=(1, 1, 2))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1, chunk_size=8, chunks_per_shard=8)
    mag.write(data=(np.random.rand(80, 240, 15) * 255).astype(np.uint8))

    assert not mag._is_compressed()
    mag.compress()
    assert mag._is_compressed()

    layer.downsample(
        from_mag=Mag(1),
        coarsest_mag=Mag(
            4
        ),  # Setting max_mag to "4" covers an edge case because the z-dimension (15) has to be rounded
    )

    # Note: this test does not check if the data is correct. This is already covered by other test cases.

    assert len(layer.mags) == 3
    assert Mag("1") in layer.mags.keys()
    assert Mag("2-2-1") in layer.mags.keys()
    assert Mag("4-4-2") in layer.mags.keys()
def test_downsample_mag_list_with_only_setup_mags(tmp_path: Path) -> None:
    ds = Dataset(tmp_path / "downsample_mag_list", voxel_size=(1, 1, 2))
    layer = ds.add_layer("color", COLOR_CATEGORY)
    mag = layer.add_mag(1)
    mag.write(data=(np.random.rand(10, 20, 30) * 255).astype(np.uint8))

    target_mags = [Mag([4, 4, 8]),
                   Mag(2), Mag([32, 32, 8]),
                   Mag(32)]  # unsorted list

    layer.downsample_mag_list(from_mag=Mag(1),
                              target_mags=target_mags,
                              only_setup_mags=True)

    for m in target_mags:
        assert np.all(
            layer.get_mag(m).read() == 0), "The mags should be empty."

    layer.downsample_mag_list(from_mag=Mag(1),
                              target_mags=target_mags,
                              allow_overwrite=True)

    for m in target_mags:
        assert m in layer.mags