test_interpolatecdo.py

# coding=utf-8
from unittest import TestCase

from earthdiagnostics.ocean.interpolatecdo import InterpolateCDO
from mock import Mock, patch

from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.frequency import Frequencies
from earthdiagnostics.diagnostic import (
    DiagnosticVariableListOption,
    DiagnosticOptionError,
)


class TestInterpolate(TestCase):
    def setUp(self):
        self.data_manager = Mock()

        self.diags = Mock()
        self.diags.model_version = "model_version"
        self.diags.config.experiment.get_chunk_list.return_value = (
            ("20010101", 0, 0),
            ("20010101", 0, 1),
        )
        self.diags.config.experiment.model_version = "model_version"
        self.diags.config.experiment.atmos_grid = "atmos_grid"
        self.diags.config.frequency = Frequencies.daily

    def fake_parse(self, value):
        if not value:
            raise DiagnosticOptionError
        return value.split("-")

    @patch(
        "earthdiagnostics.ocean.interpolatecdo.InterpolateCDO.compute_weights"
    )
    @patch(
        "earthdiagnostics.ocean.interpolatecdo.InterpolateCDO."
        "get_sample_grid_file"
    )
    @patch.object(DiagnosticVariableListOption, "parse", fake_parse)
    @patch("os.remove")
    @patch("earthdiagnostics.utils.TempFile.get")
    def test_generate_jobs(
        self, mock_weights, mock_grid_file, mock_remove, mock_get
    ):
        mock_weights.return_value = None
        mock_get.return_value = "path_to_weights"

        jobs = InterpolateCDO.generate_jobs(
            self.diags, ["interpcdo", "ocean", "var"]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "atmos_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "atmos_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags, ["interpcdo", "ocean", "var", "target_grid"]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags, ["interpcdo", "ocean", "var", "target_grid", "bicubic"]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                True,
                "",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags,
            ["interpcdo", "ocean", "var", "target_grid", "bicubic", "false"],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags,
            [
                "interpcdo",
                "ocean",
                "var",
                "target_grid",
                "bicubic",
                "false",
                "orig",
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags,
            [
                "interpcdo",
                "ocean",
                "var",
                "target_grid",
                "bicubic",
                "false",
                "orig",
                "false",
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.daily,
            ),
        )

        jobs = InterpolateCDO.generate_jobs(
            self.diags,
            [
                "interpcdo",
                "ocean",
                "var",
                "target_grid",
                "bicubic",
                "false",
                "orig",
                "false",
                "mon"
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            InterpolateCDO(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                "target_grid",
                "model_version",
                False,
                "orig",
                None,
                Frequencies.monthly,
            ),
        )

        with self.assertRaises(DiagnosticOptionError):
            InterpolateCDO.generate_jobs(self.diags, ["interp"])

        with self.assertRaises(DiagnosticOptionError):
            InterpolateCDO.generate_jobs(
                self.diags,
                [
                    "interpcdo",
                    "ocean",
                    "var",
                    "bicubic",
                    "false",
                    "orig",
                    "false",
                    "daily,"
                    "extra",
                ],
            )

    def test_str(self):
        diag = InterpolateCDO(
            self.data_manager,
            "20010101",
            0,
            0,
            ModelingRealms.ocean,
            "var",
            "atmos_grid",
            "model_version",
            False,
            "orig",
            None,
            Frequencies.monthly,
        )
        self.assertEqual(
            str(diag),
            "Interpolate with CDO Startdate: 20010101 Member: 0 Chunk: 0 "
            "Variable: ocean:var Frequency: mon "
            "Target grid: atmos_grid Original grid: orig Mask ocean: False "
            "Model: model_version",
        )