tests/test_dataset.py

import tempfile
import unittest
from pathlib import Path

import h5py
import numpy as np
import torch

from src.data.dataset import (
    SPACE_BANDS,
    SPACE_TIME_BANDS,
    STATIC_BANDS,
    TIME_BANDS,
    Dataset,
    Normalizer,
    to_cartesian,
)

BROKEN_FILE = "min_lat=24.7979_min_lon=-105.1508_max_lat=24.8069_max_lon=-105.141_dates=2022-01-01_2023-12-31.tif"
TEST_FILENAMES = [
    "min_lat=5.4427_min_lon=101.4016_max_lat=5.4518_max_lon=101.4107_dates=2022-01-01_2023-12-31.tif",
    "min_lat=-27.6721_min_lon=25.6796_max_lat=-27.663_max_lon=25.6897_dates=2022-01-01_2023-12-31.tif",
]
TIFS_FOLDER = Path(__file__).parents[1] / "data/tifs"
TEST_FILES = [TIFS_FOLDER / x for x in TEST_FILENAMES]


class TestDataset(unittest.TestCase):
    def test_tif_to_array(self):
        ds = Dataset(TIFS_FOLDER, download=False)
        for test_file in TEST_FILES:
            s_t_x, sp_x, t_x, st_x, months = ds._tif_to_array(test_file)
            self.assertFalse(np.isnan(s_t_x).any())
            self.assertFalse(np.isnan(sp_x).any())
            self.assertFalse(np.isnan(t_x).any())
            self.assertFalse(np.isnan(st_x).any())
            self.assertFalse(np.isinf(s_t_x).any())
            self.assertFalse(np.isinf(sp_x).any())
            self.assertFalse(np.isinf(t_x).any())
            self.assertFalse(np.isinf(st_x).any())
            self.assertEqual(sp_x.shape[0], s_t_x.shape[0])
            self.assertEqual(sp_x.shape[1], s_t_x.shape[1])
            self.assertEqual(t_x.shape[0], s_t_x.shape[2])
            self.assertEqual(len(SPACE_TIME_BANDS), s_t_x.shape[-1])
            self.assertEqual(len(SPACE_BANDS), sp_x.shape[-1])
            self.assertEqual(len(TIME_BANDS), t_x.shape[-1])
            self.assertEqual(len(STATIC_BANDS), st_x.shape[-1])
            self.assertEqual(months[0], 0)

    def test_files_are_replaced(self):
        ds = Dataset(TIFS_FOLDER, download=False)
        assert TIFS_FOLDER / BROKEN_FILE in ds.tifs

        for b in ds:
            assert len(b) == 5
        assert TIFS_FOLDER / BROKEN_FILE not in ds.tifs

    def test_normalization(self):
        ds = Dataset(TIFS_FOLDER, download=False)
        o = ds.load_normalization_values(path=Path("config/normalization.json"))
        for t in [len(SPACE_TIME_BANDS), len(SPACE_BANDS), len(STATIC_BANDS), len(TIME_BANDS)]:
            subdict = o[t]
            self.assertTrue("mean" in subdict)
            self.assertTrue("std" in subdict)
            self.assertTrue(len(subdict["mean"]) == len(subdict["std"]))
        normalizer = Normalizer(normalizing_dicts=o)
        ds.normalizer = normalizer
        for b in ds:
            for t in b:
                self.assertFalse(np.isnan(t).any())

    def test_subset_image_with_minimum_size(self):
        input = np.ones((3, 3, 1))
        months = static = np.ones(1)
        output = Dataset.subset_image(input, input, months, static, months, 3, 1)
        self.assertTrue(np.equal(input, output[0]).all())
        self.assertTrue(np.equal(input, output[1]).all())
        self.assertTrue(np.equal(months, output[2]).all())

    def test_subset_with_too_small_image(self):
        input = np.ones((2, 2, 1))
        months = static = np.ones(1)
        self.assertRaises(
            AssertionError, Dataset.subset_image, input, input, months, static, months, 3, 1
        )

    def test_subset_with_larger_images(self):
        input = np.ones((5, 5, 1))
        months = static = np.ones(1)
        output = Dataset.subset_image(input, input, months, static, months, 3, 1)
        self.assertTrue(np.equal(np.ones((3, 3, 1)), output[0]).all())
        self.assertTrue(np.equal(np.ones((3, 3, 1)), output[1]).all())
        self.assertTrue(np.equal(months, output[2]).all())

    def test_latlon_checks_float(self):
        # just checking it runs
        _ = to_cartesian(
            30.0,
            40.0,
        )
        with self.assertRaises(AssertionError):
            to_cartesian(1000.0, 1000.0)

    def test_latlon_checks_np(self):
        # just checking it runs
        _ = to_cartesian(np.array([30.0]), np.array([40.0]))
        with self.assertRaises(AssertionError):
            to_cartesian(np.array([1000.0]), np.array([1000.0]))

    def test_latlon_checks_tensor(self):
        # just checking it runs
        _ = to_cartesian(torch.tensor([30.0]), torch.tensor([40.0]))
        with self.assertRaises(AssertionError):
            to_cartesian(torch.tensor([1000.0]), torch.tensor([1000.0]))

    def test_process_h5pys(self):
        with tempfile.TemporaryDirectory() as tempdir_str:
            tempdir = Path(tempdir_str)
            dataset = Dataset(
                TIFS_FOLDER,
                download=False,
                h5py_folder=tempdir,
                h5pys_only=False,
            )
            dataset.process_h5pys()

            h5py_files = list(tempdir.glob("*.h5"))
            self.assertEqual(len(h5py_files), 2)
            for h5_file in h5py_files:
                with h5py.File(h5_file, "r") as f:
                    # mostly checking it can be read
                    self.assertEqual(f["t_x"].shape[0], 24)