smashbox 1.0__py2.py3-none-any.whl

smashbox/read_inputdata/smashmodel.py

@@ -0,0 +1,395 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Sep 27 12:58:19 2024
+
+@author: maxime
+"""
+
+import smash
+import yaml
+
+import os
+import pandas as pd
+import datetime
+import shutil
+import glob
+import random
+from . import read_data
+
+update_default_setup = {
+    "prcp_date_pattern": "",
+    "pet_date_pattern": "",
+    "snow_date_pattern": "",
+    "temp_date_pattern": "",
+    "prcp_directories": {},
+    "pet_directories": {},
+    "snow_directories": {},
+    "temp_directories": {},
+    "continuous_pet": {},
+    "timezone": "UTC",
+}
+
+
+def standardize_update_default_setup_options(model, update_default_setup):
+
+    for var in [
+        "prcp_date_pattern",
+        "snow_date_pattern",
+        "temp_date_pattern",
+        "pet_date_pattern",
+    ]:
+        if update_default_setup[var] == "":
+            if model.setup.dt == 86_400:
+                update_default_setup[var] = "%Y%m%d"
+            else:
+                update_default_setup[var] = "%Y%m%d%H%M"
+
+    if model.setup.daily_interannual_pet:
+        update_default_setup["pet_date_pattern"] = "%Y%m%d"
+
+
+def standardize_model_date_pattern(key: str, value: str) -> str:
+    if len(value) > 0:
+        if not isinstance(value, str):
+            raise TypeError(f"{key} model setup must be a str")
+
+        if not value.startswith("%"):
+            raise ValueError(f"{value} must start by character %")
+
+        sample = value.split("%")
+        sample.pop(0)
+
+        for s in sample:
+            if not isinstance(s, str) and len(s) > 1:
+                raise ValueError(
+                    f"%{s} in {value} must start by character % following by a unique letter: Ex: %Y%m%d%H%M"
+                )
+
+    return value
+
+
+def standardize_model_data_directories(key: str, value: dict) -> str:
+
+    if not isinstance(value, dict):
+        raise TypeError(f"{key} model setup must be a dict")
+
+    key_to_pop = []
+    for subkey, subvalue in value.items():
+
+        if isinstance(subkey, str):
+            try:
+                int(subkey)
+            except:
+                raise ValueError(
+                    f"{subkey}, must represent an integer as the dict key (order priority of the data)."
+                )
+        elif not isinstance(subkey, int):
+            raise ValueError(
+                f"{subkey}, must be an integer as the dict key (order priority of the data)."
+            )
+
+        if not isinstance(subvalue, str):
+            raise ValueError(f"{subvalue}, must be a str, a pathlike.")
+
+        if len(subvalue) == 0:
+            key_to_pop.append(subkey)
+
+        if len(subvalue) > 0 and not os.path.exists(subvalue):
+            raise ValueError(f"'{subvalue}', is not an existing path.")
+
+    # remove empty keys
+    for subkey in key_to_pop:
+        value.pop(subkey)
+
+    return value
+
+
+def standardize_model_continuous_pet(key: str, value: dict) -> str:
+
+    if not isinstance(value, dict):
+        raise TypeError(f"{key} model setup must be a dict")
+
+    true_count = 0
+    for subkey, subvalue in value.items():
+
+        if isinstance(subkey, str):
+            try:
+                int(subkey)
+            except:
+                raise ValueError(
+                    f"{subkey}, must represent an integer as the dict key (order priority of the data)."
+                )
+        elif not isinstance(subkey, int):
+            raise ValueError(
+                f"{subkey}, must be an integer as the dict key (order priority of the pet data)."
+            )
+
+        if not isinstance(subvalue, bool):
+            raise ValueError(f"{subvalue}, must be a bool.")
+
+        if subvalue == True:
+            true_count = true_count + 1
+
+        if true_count > 1:
+            raise ValueError("Only one True data is allowed in continuous pet")
+
+    return value
+
+
+def update_model_setup(model, update_default_setup):
+
+    for key, value in update_default_setup.items():
+        setattr(model.setup, key, value)
+
+
+# Build a smash model from a setup and a mesh with external new options
+# path_setup: path to the setup file
+# path_mesh: path to the mesh file
+# setup_options: dictionnary to pass setup options that will overwrite options in path_setup
+# new options can be defined in path_setup:
+# prcp_directories: dict
+# pet_directories: dict
+# snow_directories: dict
+# temp_directories: dict
+# prcp_date_pattern : str
+# pet_date_pattern : str
+# snow_date_pattern : str
+# temp_date_pattern: str
+# continuous_pet : dict of bool
+# timezone: str
+
+
+def SmashModel(input_setup, input_mesh, setup_options={}):
+    """
+
+    Description
+    -----------
+
+    Build a model Smash like smash.Model() but extend capabilities of the atmos data reading
+
+    Paramerters
+    -----------
+
+    input_setup: str() | dict()
+        the path to the setup.yaml file or the setup dictionary
+
+    input_mesh: str() | dict()
+        the path to the mesh.hdf5 file or the mesh dictionary
+
+    setup_options: dict()
+        dictionnary of setup options that will erase setup.yaml
+
+    Return
+    ------
+
+    smash.Model()
+        Un objet model smash.
+
+    """
+
+    if isinstance(input_setup, str):
+        setup = smash.io.read_setup(input_setup)
+    elif isinstance(input_setup, dict):
+        setup = input_setup.copy()
+    else:
+        raise ValueError("Wrong setup data type: input_setup must be a path or a dict")
+
+    sending_options = setup_options.copy()
+
+    for key, value in setup_options.items():
+        setup.update({key: value})
+
+    # filter setup options not allowed in smash to prevent warnings
+    for key in update_default_setup:
+        if key in setup:
+            setup.pop(key)
+
+    if isinstance(input_mesh, str):
+        mesh = smash.io.read_mesh(input_mesh)
+    elif isinstance(input_mesh, dict):
+        mesh = input_mesh.copy()
+    else:
+        raise ValueError("Wrong mesh data type: setup must be a path or a dict")
+
+    # force to prevent reading data
+    if "read_prcp" in setup:
+        setup["read_prcp"] = False
+    if "read_pet" in setup:
+        setup["read_pet"] = False
+    if "read_temp" in setup:
+        setup["read_temp"] = False
+    if "read_snow" in setup:
+        setup["read_snow"] = False
+    if "compute_mean_atmos" in setup:
+        setup["compute_mean_atmos"] = False
+
+    print(
+        f"</> Building model from {setup['start_time']} to {setup['end_time']} with time-step {setup['dt']}"
+    )
+
+    # build the model
+    model = smash.Model(setup, mesh)
+
+    # Updtade model.setup with new options
+    standardize_update_default_setup_options(model, update_default_setup)
+    update_model_setup(model, update_default_setup)
+
+    # read all options in setup, not only those for smash:
+    if isinstance(input_setup, str):
+        with open(input_setup, "r") as file:
+            setup_overwrite = yaml.safe_load(file)
+    elif isinstance(input_setup, dict):
+        setup_overwrite = input_setup.copy()
+    else:
+        raise ValueError("Wrong setup data type: input_setup must be a path or a dict")
+
+    # overwrite setup_overwrite with setup_options
+    for key, value in setup_options.items():
+        setup_overwrite.update({key: value})
+
+    # Overight options with the config file
+    for key, value in setup_overwrite.items():
+        if hasattr(model.setup, key):
+            setattr(model.setup, key, value)
+
+    # print(setup_overwrite)
+    # print(model.setup)
+    # standardize new user options
+    standardize_model_date_pattern("prcp_date_pattern", model.setup.prcp_date_pattern)
+    standardize_model_date_pattern("pet_date_pattern", model.setup.pet_date_pattern)
+    standardize_model_date_pattern("snow_date_pattern", model.setup.snow_date_pattern)
+    standardize_model_date_pattern("temp_date_pattern", model.setup.temp_date_pattern)
+    standardize_model_data_directories("prcp_directories", model.setup.prcp_directories)
+    standardize_model_data_directories("pet_directories", model.setup.pet_directories)
+    standardize_model_data_directories("temp_directories", model.setup.temp_directories)
+    standardize_model_data_directories("snow_directories", model.setup.snow_directories)
+    standardize_model_continuous_pet("continuous_pet", model.setup.continuous_pet)
+
+    # here we must manage etpc... => must be moved into read_data
+    # Get ready with Continuous ETP
+    path_to_smashbox = os.path.join(
+        os.path.expandvars("$HOME"), ".smashbox", "simlink_pet"
+    )
+
+    if not os.path.exists(path_to_smashbox):
+        os.makedirs(path_to_smashbox)
+
+    path_to_symlink_etp = None
+
+    if "continuous_pet" in setup_overwrite:
+
+        if "pet_directories" in setup_overwrite:
+            dir_etp = setup_overwrite["pet_directories"]
+        elif "pet_directory" in setup_overwrite:
+            dir_etp = {1: setup_overwrite["pet_directory"]}
+        else:
+            dir_etp = {}
+
+        continuous_pet = setup_overwrite["continuous_pet"]
+        list_dir = os.listdir(path_to_smashbox)
+        path_to_symlink_etp = os.path.join(
+            path_to_smashbox,
+            f"symlinked_etpc_{len(list_dir)+1}_{int(random.random()*1000000)}",
+        )
+        print("</> Path to symlinked pet is set to ", path_to_symlink_etp)
+    else:
+        continuous_pet = {}
+
+    for key, value in continuous_pet.items():
+        if value:
+            prepare_etp_c(
+                path_to_symlink_etp=path_to_symlink_etp,
+                start_time=model.setup.start_time,
+                end_time=model.setup.end_time,
+                pet_date_pattern=model.setup.pet_date_pattern,
+                fmt=".tif",
+                dir_etp=dir_etp[key],
+            )
+            dir_etp[key] = path_to_symlink_etp
+            sending_options.update({"pet_directories": dir_etp})
+            break  # only one source of continuous PET
+
+    # print(model.setup)
+    # read atmos data with updated func
+    print("</> Reading atmospheric data ...")
+    read_data.read_atmos_data(model, setup_options=sending_options)
+
+    # clean path to symplink pet files
+    if path_to_symlink_etp is not None:
+        if os.path.exists(path_to_symlink_etp):
+            print("</> Removing path to symplink pet ", path_to_symlink_etp)
+            shutil.rmtree(path_to_symlink_etp)
+
+    return model
+
+
+def prepare_etp_c(
+    path_to_symlink_etp="symlinked_etp_c",
+    start_time=None,
+    end_time=None,
+    pet_date_pattern="%Y%m%d",
+    fmt=".tif",
+    dir_etp="",
+):
+
+    if (start_time is None) or (end_time is None):
+        print("start_time or end_time are empty... ")
+        return None
+
+    if dir_etp == "":
+        print("dir_etp empty...")
+        return None
+
+    date_range = pd.date_range(
+        start=datetime.datetime.fromisoformat(start_time),
+        end=datetime.datetime.fromisoformat(end_time),
+        freq=datetime.timedelta(days=1),
+    )
+
+    print("</> Prepare continuous PET ...")
+
+    if os.path.exists(path_to_symlink_etp):
+        shutil.rmtree(path_to_symlink_etp)
+
+    os.makedirs(path_to_symlink_etp)
+
+    warn_list = []
+    for date in date_range:
+
+        date_etp_c_to_read = (date - datetime.timedelta(days=1)).strftime(
+            pet_date_pattern
+        )
+        date_etp_c_j0 = (date).strftime("%m%d")
+        date_etp_c_j1 = (date + datetime.timedelta(days=1)).strftime("%m%d")
+
+        file_etp_c = glob.glob(f"{dir_etp}/**/*{date_etp_c_to_read}*.tif", recursive=True)
+
+        if len(file_etp_c) > 0:
+            file_etp_c = file_etp_c[0]
+
+            if os.path.exists(file_etp_c):
+                mylink = os.path.join(path_to_symlink_etp, f"ETPjc_{date_etp_c_j0}{fmt}")
+
+                if os.path.exists(mylink):
+                    os.remove(mylink)
+
+                os.symlink(file_etp_c, mylink)
+
+                mylink = os.path.join(path_to_symlink_etp, f"ETPjc_{date_etp_c_j1}{fmt}")
+
+                if os.path.exists(mylink):
+                    os.remove(mylink)
+
+                os.symlink(file_etp_c, mylink)
+
+            warn_list.append(
+                "Linked "
+                + f"ETPjc_{date_etp_c_j0}{fmt} and ETPjc_{date_etp_c_j1}{fmt} "
+                + f"to file {os.path.basename(file_etp_c)}"
+            )
+
+    if len(warn_list) > 1:
+        print(f"</> {warn_list[0]}, ..., {warn_list[-1]}")
+    elif len(warn_list) > 0:
+        print(f"</> {warn_list[0]}")

smashbox/stats/__init__.py

@@ -0,0 +1 @@
+