Source code for hydroflows.methods.wflow.wflow_update_factors

"""Downscale a climate change factor dataset to the resolution of a wflow model and add to settings file."""

from os.path import relpath
from pathlib import Path

import hydromt  # noqa: F401
import xarray as xr
from hydromt_wflow import WflowModel

from hydroflows._typing import FileDirPath, OutputDirPath
from hydroflows.methods.utils.io import to_netcdf
from hydroflows.methods.wflow.wflow_utils import copy_wflow_model
from hydroflows.workflow.method import Method
from hydroflows.workflow.method_parameters import Parameters

__all__ = ["WflowUpdateChangeFactors", "Input", "Output", "Params"]




class Input(Parameters):
    """Input parameters.

    This class represents the input data
    required for the :py:class:`WflowUpdateChangeFactors` method.
    """

    change_factor_dataset: Path
    """
    The path to the to be change factor dataset.
    """

    wflow_toml: FileDirPath
    """
    The path to the wflow settings toml that will be updated.
    """





class Output(Parameters):
    """output parameters.

    this class represents the output data
    generated by the :py:class:`WflowUpdateChangeFactors` method.
    """

    wflow_change_factors: Path
    """
    Path to the change factor dataset at wflow model resolution.
    """

    wflow_out_toml: FileDirPath
    """
    The path to the updated wflow settings toml.
    """





class Params(Parameters):
    """Parameters for the :py:class:`WflowUpdateChangeFactors`.

    Instances of this class are used in the :py:class:`WflowUpdateChangeFactors`
    method to define the required settings.
    """

    output_dir: OutputDirPath
    """Output location relative to the workflow root. The updated model will be stored in <output_dir>."""

    copy_model: bool = False
    """Create full copy of model or create rel paths in model config."""

    resample_method: str = "nearest"
    """Method of resampling the low(er) res dataset."""





class WflowUpdateChangeFactors(Method):
    """Downscale a climate change factor dataset to the resolution of a wflow model and add to settings file.

    Parameters
    ----------
    change_factor_dataset : Path
        Path to the to be downscaled dataset.
    wflow_toml : Path
        Path to the wflow settings toml that needs to be adjusted.
    **params
        Additional parameters to pass to the WflowDownscale instance.
        See :py:class:`wflow_update_factors Params <hydroflows.methods.wflow.wflow_update_factors.Params>`.

    See Also
    --------
    :py:class:`wflow_update_factors Input <~hydroflows.methods.wflow.wflow_update_factors.Input>`
    :py:class:`wflow_update_factors Output <~hydroflows.methods.wflow.wflow_update_factors.Output>`
    :py:class:`wflow_update_factors Params <~hydroflows.methods.wflow.wflow_update_factors.Params>`
    """

    name: str = "wflow_update_factors"

    _test_kwargs = {
        "change_factor_dataset": Path("dataset.nc"),
        "wflow_toml": Path("wflow_sbm.toml"),
        "output_dir": Path("data"),
    }

    def __init__(
        self,
        change_factor_dataset: Path,
        wflow_toml: Path,
        output_dir: Path,
        **params,
    ):
        self.params: Params = Params(output_dir=output_dir, **params)
        self.input: Input = Input(
            change_factor_dataset=change_factor_dataset,
            wflow_toml=wflow_toml,
        )
        fname = self.input.change_factor_dataset.stem
        fsuffix = self.input.change_factor_dataset.suffix

        if not self.params.copy_model and not self.params.output_dir.is_relative_to(
            self.input.wflow_toml.parent
        ):
            raise ValueError(
                "Output directory must be relative to input directory when not copying model."
            )

        self.output: Output = Output(
            wflow_change_factors=self.params.output_dir
            / f"{fname}_downscaled{fsuffix}",
            wflow_out_toml=self.params.output_dir / "wflow_sbm.toml",
        )

    def _run(self):
        """Run the downscale dataset method."""
        if self.params.copy_model:
            copy_wflow_model(
                src=self.input.wflow_toml.parent,
                dest=self.output.wflow_out_toml.parent,
                copy_forcing=True,
            )

        # Open input files
        ds = xr.open_dataset(self.input.change_factor_dataset, lock=False)
        w = WflowModel(root=self.input.wflow_toml.parent, mode="r+")

        # squeeze
        ds = ds.squeeze()

        # Downscale the data
        ds_downscaled = ds.raster.reproject_like(
            w.grid,
            method=self.params.resample_method,
        )

        # rename from month to time for wflow
        if "month" in ds_downscaled.coords:
            ds_downscaled = ds_downscaled.rename({"month": "time"})

        # Update the config
        redirect = ["input.path_forcing", "input.path_static"]
        # Redirect paths to forcing and staticmaps
        if not self.params.copy_model:
            for item in redirect:
                value = w.get_config(item)
                full_path = Path(value)
                if not full_path.is_absolute():
                    full_path = Path(self.input.wflow_toml.parent, full_path)
                else:
                    continue
                new_path = Path(relpath(full_path, self.output.wflow_out_toml.parent))
                w.set_config(item, new_path.as_posix())
        # Put the downscaled dataset in the toml
        w.set_config("input.path_forcing_scale", self.output.wflow_change_factors.name)

        # Write the config
        w.write_config(
            config_name=self.output.wflow_out_toml.name,
            config_root=self.params.output_dir,
        )

        # write netcdf
        to_netcdf(
            ds_downscaled,
            file_name=self.output.wflow_change_factors.name,
            output_dir=self.params.output_dir,
        )