Apply patch generator (#621)

* Apply_patch_generator to test_parsing

* Add generate_eopatch to test eodata_io

* Update test_eodata with generate_eopatch

* Update test_core_tasks with generate_eopatch

* Update test_train_split with generate_eopatch

* Update test_sampling with generate_eopatch

* Update mini_eopatch feature names

* Split patch_fixture on two

* Rename fixtuer patch_bands to eopatch_to_explode and update CLP_S2C in patch_fixture as suggested

* Update CLP_S2C

* Update test_get_spatial_dimension, update fixtuer in test_sampling

* Update core/eolearn/tests/test_core_tasks.py

Co-authored-by: Žiga Lukšič <[email protected]>

* Update core/eolearn/tests/test_eodata.py

Co-authored-by: Žiga Lukšič <[email protected]>

---------

Co-authored-by: Žiga Lukšič <[email protected]>

Author: jgersak

core/eolearn/tests/test_core_tasks.py

@@ -44,34 +44,32 @@
 from eolearn.core.core_tasks import ExplodeBandsTask
 from eolearn.core.types import FeatureRenameSpec, FeatureSpec, FeaturesSpecification
 from eolearn.core.utils.parsing import parse_features
-from eolearn.core.utils.testing import assert_feature_data_equal
+from eolearn.core.utils.testing import PatchGeneratorConfig, assert_feature_data_equal, generate_eopatch
 
 DUMMY_BBOX = BBox((0, 0, 1, 1), CRS(3857))
 
 
 @pytest.fixture(name="patch")
 def patch_fixture() -> EOPatch:
-    patch = EOPatch(bbox=BBox((324.54, 546.45, 955.4, 63.43), CRS(3857)))
-    patch.data["bands"] = np.arange(5 * 3 * 4 * 8).reshape(5, 3, 4, 8)
-    patch.data["CLP"] = np.full((5, 3, 4, 1), 0.7)
-    patch.data["CLP_S2C"] = np.zeros((5, 3, 4, 1), dtype=np.int64)
-    patch.mask["CLM"] = np.full((5, 3, 4, 1), True)
-    patch.mask_timeless["mask"] = np.arange(3 * 4 * 2).reshape(3, 4, 2)
-    patch.mask_timeless["LULC"] = np.zeros((3, 4, 1), dtype=np.uint16)
-    patch.mask_timeless["RANDOM_UINT8"] = np.random.randint(0, 100, size=(3, 4, 1), dtype=np.int8)
-    patch.scalar["values"] = np.arange(10 * 5).reshape(5, 10)
-    patch.scalar["CLOUD_COVERAGE"] = np.ones((5, 10))
-    patch.timestamps = [
-        datetime(2017, 1, 14, 10, 13, 46),
-        datetime(2017, 2, 10, 10, 1, 32),
-        datetime(2017, 2, 20, 10, 6, 35),
-        datetime(2017, 3, 2, 10, 0, 20),
-        datetime(2017, 3, 12, 10, 7, 6),
-    ]
-    patch.meta_info["something"] = np.random.rand(10, 1)
+    patch = generate_eopatch(
+        {
+            FeatureType.DATA: ["bands", "CLP"],
+            FeatureType.MASK: ["CLM"],
+            FeatureType.MASK_TIMELESS: ["mask", "LULC", "RANDOM_UINT8"],
+            FeatureType.SCALAR: ["values", "CLOUD_COVERAGE"],
+        }
+    )
+    patch.data["CLP_S2C"] = np.zeros_like(patch.data["CLP"])
+
+    patch.meta_info["something"] = "beep boop"
     return patch
 
 
+@pytest.fixture(name="eopatch_to_explode")
+def eopatch_to_explode_fixture() -> EOPatch:
+    return generate_eopatch((FeatureType.DATA, "bands"), config=PatchGeneratorConfig(depth_range=(8, 9)))
+
+
 @pytest.mark.parametrize("task", [DeepCopyTask, CopyTask])
 def test_copy(task: Type[CopyTask], patch: EOPatch) -> None:
     patch_copy = task().execute(patch)
@@ -413,11 +411,11 @@ def kwargs_map(data, *, some=3, **kwargs) -> tuple:
     ],
 )
 def test_explode_bands(
-    patch: EOPatch,
+    eopatch_to_explode: EOPatch,
     feature: Tuple[FeatureType, str],
     task_input: Dict[Tuple[FeatureType, str], Union[int, Iterable[int]]],
 ) -> None:
-    patch = ExplodeBandsTask(feature, task_input)(patch)
+    patch = ExplodeBandsTask(feature, task_input)(eopatch_to_explode)
     assert all(new_feature in patch for new_feature in task_input)
 
     for new_feature, bands in task_input.items():
@@ -426,19 +424,23 @@ def test_explode_bands(
         assert_array_equal(patch[new_feature], patch[feature][..., bands])
 
 
-def test_extract_bands(patch: EOPatch) -> None:
+def test_extract_bands(eopatch_to_explode: EOPatch) -> None:
     bands = [2, 4, 6]
-    patch = ExtractBandsTask((FeatureType.DATA, "bands"), (FeatureType.DATA, "EXTRACTED_BANDS"), bands)(patch)
+    patch = ExtractBandsTask((FeatureType.DATA, "bands"), (FeatureType.DATA, "EXTRACTED_BANDS"), bands)(
+        eopatch_to_explode
+    )
     assert_array_equal(patch.data["EXTRACTED_BANDS"], patch.data["bands"][..., bands])
 
     patch.data["EXTRACTED_BANDS"][0, 0, 0, 0] += 1
     assert patch.data["EXTRACTED_BANDS"][0, 0, 0, 0] != patch.data["bands"][0, 0, 0, bands[0]]
 
 
-def test_extract_bands_fails(patch: EOPatch) -> None:
+def test_extract_bands_fails(eopatch_to_explode: EOPatch) -> None:
     with pytest.raises(ValueError):
         # fails because band 16 does not exist
-        ExtractBandsTask((FeatureType.DATA, "bands"), (FeatureType.DATA, "EXTRACTED_BANDS"), [2, 4, 16])(patch)
+        ExtractBandsTask((FeatureType.DATA, "bands"), (FeatureType.DATA, "EXTRACTED_BANDS"), [2, 4, 16])(
+            eopatch_to_explode
+        )
 
 
 @pytest.mark.parametrize(

core/eolearn/tests/test_eodata.py

@@ -18,19 +18,20 @@
 from eolearn.core.eodata_io import FeatureIO
 from eolearn.core.exceptions import EODeprecationWarning
 from eolearn.core.types import FeatureSpec, FeaturesSpecification
-from eolearn.core.utils.testing import assert_feature_data_equal
+from eolearn.core.utils.testing import assert_feature_data_equal, generate_eopatch
 
 DUMMY_BBOX = BBox((0, 0, 1, 1), CRS(3857))
 
 
 @pytest.fixture(name="mini_eopatch")
 def mini_eopatch_fixture() -> EOPatch:
-    eop = EOPatch(bbox=DUMMY_BBOX)
-    eop.data["bands"] = np.arange(2 * 3 * 3 * 2).reshape(2, 3, 3, 2)
-    eop.data["zeros"] = np.zeros((2, 3, 3, 2), dtype=float)
-    eop.mask["ones"] = np.ones((2, 6, 6, 1), dtype=int)
-    eop.mask["twos"] = np.ones((2, 3, 3, 2), dtype=int) * 2
-    eop.mask_timeless["threes"] = np.ones((3, 3, 1), dtype=np.uint8) * 3
+    eop = generate_eopatch(
+        {
+            FeatureType.DATA: ["A", "B"],
+            FeatureType.MASK: ["C", "D"],
+            FeatureType.MASK_TIMELESS: ["E"],
+        }
+    )
     eop.meta_info["beep"] = "boop"
 
     return eop
@@ -161,9 +162,9 @@ def test_simplified_feature_operations() -> None:
 @pytest.mark.parametrize(
     "feature_to_delete",
     [
-        (FeatureType.DATA, "zeros"),
-        (FeatureType.MASK, "ones"),
-        (FeatureType.MASK_TIMELESS, "threes"),
+        (FeatureType.DATA, "A"),
+        (FeatureType.MASK, "C"),
+        (FeatureType.MASK_TIMELESS, "E"),
         (FeatureType.META_INFO, "beep"),
         (FeatureType.TIMESTAMPS, None),
     ],
@@ -317,18 +318,27 @@ def test_equals() -> None:
     assert eop1 != eop2
 
 
+@pytest.fixture(scope="function", name="eopatch_spatial_dim")
+def eopatch_spatial_dim_fixture() -> EOPatch:
+    patch = EOPatch(bbox=DUMMY_BBOX)
+    patch.data["A"] = np.zeros((1, 2, 3, 4))
+    patch.mask["B"] = np.ones((4, 3, 2, 1), dtype=np.uint8)
+    patch.mask_timeless["C"] = np.zeros((4, 5, 1), dtype=np.uint8)
+    return patch
+
+
 @pytest.mark.parametrize(
     "feature, expected_dim",
     [
-        [(FeatureType.DATA, "zeros"), (3, 3)],
-        [(FeatureType.MASK, "ones"), (6, 6)],
-        [(FeatureType.MASK_TIMELESS, "threes"), (3, 3)],
+        [(FeatureType.DATA, "A"), (2, 3)],
+        [(FeatureType.MASK, "B"), (3, 2)],
+        [(FeatureType.MASK_TIMELESS, "C"), (4, 5)],
     ],
 )
 def test_get_spatial_dimension(
-    feature: Tuple[FeatureType, str], expected_dim: Tuple[int, int], mini_eopatch: EOPatch
+    feature: Tuple[FeatureType, str], expected_dim: Tuple[int, int], eopatch_spatial_dim: EOPatch
 ) -> None:
-    assert mini_eopatch.get_spatial_dimension(*feature) == expected_dim
+    assert eopatch_spatial_dim.get_spatial_dimension(*feature) == expected_dim
 
 
 @pytest.mark.parametrize(
@@ -337,13 +347,14 @@ def test_get_spatial_dimension(
         (
             pytest.lazy_fixture("mini_eopatch"),
             [
-                (FeatureType.DATA, "bands"),
-                (FeatureType.DATA, "zeros"),
-                (FeatureType.MASK, "ones"),
-                (FeatureType.MASK, "twos"),
-                (FeatureType.MASK_TIMELESS, "threes"),
+                (FeatureType.DATA, "A"),
+                (FeatureType.DATA, "B"),
+                (FeatureType.MASK, "C"),
+                (FeatureType.MASK, "D"),
+                (FeatureType.MASK_TIMELESS, "E"),
                 (FeatureType.META_INFO, "beep"),
                 (FeatureType.BBOX, None),
+                (FeatureType.TIMESTAMPS, None),
             ],
         ),
         (EOPatch(bbox=DUMMY_BBOX), [(FeatureType.BBOX, None)]),

core/eolearn/tests/test_eodata_io.py

@@ -35,7 +35,7 @@
 from eolearn.core.exceptions import EODeprecationWarning
 from eolearn.core.types import FeaturesSpecification
 from eolearn.core.utils.parsing import FeatureParser
-from eolearn.core.utils.testing import assert_feature_data_equal
+from eolearn.core.utils.testing import assert_feature_data_equal, generate_eopatch
 
 FS_LOADERS = [TempFS, pytest.lazy_fixture("create_mocked_s3fs")]
 
@@ -44,16 +44,16 @@
 
 @pytest.fixture(name="eopatch")
 def eopatch_fixture():
-    eopatch = EOPatch(bbox=DUMMY_BBOX)
-    mask = np.zeros((3, 3, 2), dtype=np.int16)
-    data = np.zeros((2, 3, 3, 2), dtype=np.int16)
-    eopatch.mask_timeless["mask"] = mask
-    eopatch.data["data"] = data
-    eopatch.timestamps = [datetime.datetime(2017, 1, 1, 10, 4, 7), datetime.datetime(2017, 1, 4, 10, 14, 5)]
+    eopatch = generate_eopatch(
+        {
+            FeatureType.DATA: ["data"],
+            FeatureType.MASK_TIMELESS: ["mask"],
+            FeatureType.SCALAR: ["my scalar with spaces"],
+            FeatureType.SCALAR_TIMELESS: ["my timeless scalar with spaces"],
+        }
+    )
     eopatch.meta_info["something"] = "nothing"
     eopatch.meta_info["something-else"] = "nothing"
-    eopatch.scalar["my scalar with spaces"] = np.array([[1, 2, 3], [1, 2, 3]])
-    eopatch.scalar_timeless["my timeless scalar with spaces"] = np.array([1, 2, 3])
     eopatch.vector["my-df"] = GeoDataFrame(
         {
             "values": [1, 2],
@@ -113,7 +113,7 @@ def test_overwriting_non_empty_folder(eopatch, fs_loader):
         eopatch.save("/", filesystem=temp_fs, overwrite_permission=OverwritePermission.OVERWRITE_FEATURES)
         eopatch.save("/", filesystem=temp_fs, overwrite_permission=OverwritePermission.OVERWRITE_PATCH)
 
-        add_eopatch = EOPatch(bbox=DUMMY_BBOX)
+        add_eopatch = EOPatch(bbox=eopatch.bbox)
         add_eopatch.data_timeless["some data"] = np.empty((3, 3, 2))
         add_eopatch.save("/", filesystem=temp_fs, overwrite_permission=OverwritePermission.ADD_ONLY)
         with pytest.raises(ValueError):

core/eolearn/tests/test_utils/test_parsing.py

@@ -1,26 +1,11 @@
-import datetime as dt
 from dataclasses import dataclass
 from typing import Callable, Iterable, List, Optional, Tuple, Union
 
-import numpy as np
 import pytest
 
-from sentinelhub import CRS, BBox
-
 from eolearn.core import EOPatch, FeatureParser, FeatureType
 from eolearn.core.types import EllipsisType, FeatureRenameSpec, FeatureSpec, FeaturesSpecification
-
-
-@pytest.fixture(name="eopatch", scope="module")
-def eopatch_fixture():
-    return EOPatch(
-        data=dict(data=np.zeros((2, 2, 2, 2)), CLP=np.zeros((2, 2, 2, 2))),  # name duplication intentional
-        bbox=BBox((1, 2, 3, 4), CRS.WGS84),
-        timestamps=[dt.datetime(2020, 5, 1), dt.datetime(2020, 5, 25)],
-        mask=dict(data=np.zeros((2, 2, 2, 2), dtype=int), IS_VALID=np.zeros((2, 2, 2, 2), dtype=int)),
-        mask_timeless=dict(LULC=np.zeros((2, 2, 2), dtype=int)),
-        meta_info={"something": "else"},
-    )
+from eolearn.core.utils.testing import generate_eopatch
 
 
 @dataclass
@@ -181,6 +166,15 @@ def test_allowed_feature_types_iterable(test_input: FeaturesSpecification, allow
         FeatureParser(features=test_input, allowed_feature_types=allowed_types)
 
 
+@pytest.fixture(name="eopatch", scope="module")
+def eopatch_fixture():
+    patch = generate_eopatch(
+        {FeatureType.DATA: ["data", "CLP"], FeatureType.MASK: ["data", "IS_VALID"], FeatureType.MASK_TIMELESS: ["LULC"]}
+    )
+    patch.meta_info = {"something": "else"}
+    return patch
+
+
 @pytest.mark.parametrize(
     "test_input, allowed_types",
     [

ml_tools/eolearn/tests/test_sampling.py

@@ -14,9 +14,8 @@
 from pytest import approx
 from shapely.geometry import Point, Polygon
 
-from sentinelhub import CRS, BBox
-
 from eolearn.core import EOPatch, EOTask, FeatureType
+from eolearn.core.utils.testing import PatchGeneratorConfig, generate_eopatch
 from eolearn.ml_tools import BlockSamplingTask, FractionSamplingTask, GridSamplingTask, sample_by_values
 from eolearn.ml_tools.sampling import expand_to_grids, get_mask_of_samples, random_point_in_triangle
 
@@ -134,11 +133,10 @@ def test_get_mask_of_samples(small_image: np.ndarray, n_samples: Dict[int, int])
 
 @pytest.fixture(name="eopatch")
 def eopatch_fixture(small_image: np.ndarray) -> EOPatch:
-    t, h, w, d = 10, *small_image.shape, 5
-    eopatch = EOPatch(bbox=BBox((0, 0, 1, 1), CRS(3857)))
-    eopatch.data["bands"] = np.arange(t * h * w * d).reshape(t, h, w, d)
-    eopatch.mask_timeless["raster"] = small_image.reshape(small_image.shape + (1,))
-    return eopatch
+    config = PatchGeneratorConfig(raster_shape=small_image.shape, depth_range=(5, 6), num_timestamps=10)
+    patch = generate_eopatch([(FeatureType.DATA, "bands")], config=config)
+    patch.mask_timeless["raster"] = small_image.reshape(small_image.shape + (1,))
+    return patch
 
 
 SAMPLING_MASK = FeatureType.MASK_TIMELESS, "sampling_mask"
@@ -159,14 +157,15 @@ def block_task_fixture(request) -> EOTask:
 def test_object_sampling_task_mask(
     eopatch: EOPatch, small_image: np.ndarray, seed: int, block_task: BlockSamplingTask
 ) -> None:
-    t, h, w, d = 10, *small_image.shape, 5
+    t, h, w, d = eopatch.data["bands"].shape
+    dr = eopatch.mask_timeless["raster"].shape[2]
     amount = block_task.amount
 
     block_task.execute(eopatch, seed=seed)
     expected_amount = amount if isinstance(amount, int) else round(np.prod(small_image.shape) * amount)
 
     assert eopatch.data["SAMPLED_DATA"].shape == (t, expected_amount, 1, d)
-    assert eopatch.mask_timeless["SAMPLED_LABELS"].shape == (expected_amount, 1, 1)
+    assert eopatch.mask_timeless["SAMPLED_LABELS"].shape == (expected_amount, 1, dr)
     assert eopatch.mask_timeless["sampling_mask"].shape == (h, w, 1)
 
     sampled_uniques, sampled_counts = np.unique(eopatch.data["SAMPLED_DATA"], return_counts=True)

ml_tools/eolearn/tests/test_train_split.py

@@ -11,13 +11,13 @@
 import pytest
 from numpy.testing import assert_array_equal
 
-from sentinelhub import CRS, BBox
-
 from eolearn.core import EOPatch, FeatureType
+from eolearn.core.utils.testing import PatchGeneratorConfig, generate_eopatch
 from eolearn.ml_tools.train_test_split import TrainTestSplitTask, TrainTestSplitType
 
 INPUT_FEATURE = (FeatureType.MASK_TIMELESS, "TEST")
 OUTPUT_FEATURE = (FeatureType.MASK_TIMELESS, "TEST_TRAIN_MASK")
+INPUT_FEATURE_CONFIG = PatchGeneratorConfig(raster_shape=(1000, 1000), depth_range=(3, 4))
 
 
 @pytest.mark.parametrize(
@@ -35,26 +35,14 @@ def test_bad_args(bad_arg: Any, bad_kwargs: Any) -> None:
         TrainTestSplitTask(INPUT_FEATURE, OUTPUT_FEATURE, bad_arg, **bad_kwargs)
 
 
-SEED = 1
-
-
 @pytest.fixture(name="eopatch1", scope="function")
 def eopatch1_fixture() -> EOPatch:
-    eopatch = EOPatch(bbox=BBox((0, 0, 1, 1), CRS(3857)))
-
-    rng = np.random.default_rng(SEED)
-    eopatch[INPUT_FEATURE] = rng.integers(0, 10, size=(1000, 1000, 3))
-
-    return eopatch
+    return generate_eopatch(INPUT_FEATURE, config=INPUT_FEATURE_CONFIG)
 
 
 @pytest.fixture(name="eopatch2")
 def eopatch2_fixture() -> EOPatch:
-    eopatch = EOPatch(bbox=BBox((0, 0, 1, 1), CRS(3857)))
-    rng = np.random.default_rng(SEED)
-    eopatch[INPUT_FEATURE] = rng.integers(0, 10, size=(1000, 1000, 3), dtype=int)
-
-    return eopatch
+    return generate_eopatch(INPUT_FEATURE, seed=69, config=INPUT_FEATURE_CONFIG)
 
 
 def test_train_split(eopatch1: EOPatch) -> None: