Step2_Process_grid_input.py

from locale import normalize
import os
import glob
import json
import datetime
import numpy as np
import xarray as xr
import pandas as pd


def timer(func):
    def func_wrapper(*args, **kwargs):
        from time import time
        time_start = time()
        result = func(*args, **kwargs)
        time_end = time()
        time_spend = (time_end - time_start) / 60.
        print('%s cost time: %.3f min' % (func.__name__, time_spend))
        return result

    return func_wrapper


class Dataset():
    __name__ = ['fit', 'clip_by_date']

    def __init__(self, cfg: dict):
        """_summary_

        Args:
            forcing_root (str): root path of forcing data;
            lai_root (str): root path of lai data;
            ancillary_root (str): root path of ancillary data;
            et_product (str): product name of target et data;
            begin_year (int): begin year for training;
            end_year (int): end year for training;
            t_resolution (str): temporal resolution of target et data;
            s_resolution (str): spatial resolution of target et data;
        """
        self.t_resolution = cfg["temporal_resolution"]
        self.s_resolution = cfg["spatial_resolution"]
        self.data_path = cfg["grid_root"]
        self.cfg = cfg

        self.precipitation_product = cfg["precipitation_product"]
        self.KGE_threshold = cfg["KGE_threshold"]

        self.forcing_list = cfg["variables"]
        self.var_list = ['Q', 'sp', 'ssrd', 'strd', 't2m', 'u10', 'v10', 'swvl1', 'swvl2']

    @timer
    def fit(self, year):
        """main process for making training/test data"""
        # get path of target et data
        print('Now we predict on:{begin_year}'.format(begin_year=year))
        lai_path = glob.glob(self.data_path + 'ERA5Land_1D_0p1.nc')[0]
        PATH = self.data_path

        lat_file_name = 'lat_{t}_{s}.npy'.format(
            t=self.t_resolution, s=self.s_resolution)
        lon_file_name = 'lon_{t}_{s}.npy'.format(
            t=self.t_resolution, s=self.s_resolution)
        if os.path.exists(PATH + lat_file_name):
            print('  [DataML] lat/lon grids exists')
            self.lat_idx = np.load(PATH + 'lat_idx.npy', allow_pickle=True)
            self.lon_idx = np.load(PATH + 'lon_idx.npy', allow_pickle=True)

        else:
            print('  [DataML] making lat/lon grids')
            with xr.open_dataset(lai_path) as f:
                lat, lon = np.array(f.latitude), np.array(f.longitude)
            # 创建一个布尔掩码，标记哪些位置满足条件
            mask_lat = (lat >= self.cfg['min_lat']) & (lat <= self.cfg['max_lat'])
            mask_lon = (lon >= self.cfg['min_lon']) & (lon <= self.cfg['max_lon'])

            # 使用掩码直接获取满足条件的值
            lat = lat[mask_lat]
            lon = lon[mask_lon]
            # 使用 np.where 获取满足条件的元素的索引
            self.lat_idx = np.where(mask_lat)[0]  # np.where 返回一个元组，我们取第一个元素
            self.lon_idx = np.where(mask_lon)[0]  # np.where 返回一个元组，我们取第一个元素
            np.save(PATH + lat_file_name, lat)
            np.save(PATH + lon_file_name, lon)
            np.save(PATH + 'lat_idx.npy', self.lat_idx)
            np.save(PATH + 'lon_idx.npy', self.lon_idx)

        file_name = 'Runoff_{tr}_{sr}_{begin_year}.npy'.format(
            tr=self.t_resolution,
            sr=self.s_resolution,
            begin_year=year)
        if os.path.exists(PATH + 'Runoff/' + file_name):
            print('  [DataML] Runoff exists')
        else:
            print('  [DataML] making Runoff')
            Runoff = self._load_runoff(self.data_path + 'Runoff/',
                                       year,
                                       self.t_resolution,
                                       self.s_resolution)
            np.save(PATH + 'Runoff/' + file_name, Runoff)
            print('Runoff shape is {shape}'.format(shape=Runoff.shape))
            del Runoff

        file_name = 'static_{tr}_{sr}_{year}.npy'.format(year=year, tr=self.t_resolution, sr=self.s_resolution)
        if os.path.exists(self.data_path + 'ancillary/' + file_name):
            print('  [DataML] ancillary exists')
        else:
            print('  [DataML] making ancillary')
            static = self._load_ancillary(
                self.data_path + "ancillary", self.s_resolution, year)
            np.save(PATH + 'ancillary/' + file_name, static)
            print('static shape is {shape}'.format(shape=static.shape))
            del static

        file_name = 'MODIS_LAI_{tr}_{sr}_{begin_year}.npy'.format(
            tr=self.t_resolution,
            sr=self.s_resolution,
            begin_year=year)
        if os.path.exists(PATH + 'LAI/' + file_name):  # 'LAI/' +
            print('  [DataML] LAI exists')
        else:
            print('  [DataML] making LAI')
            lai = self._load_lai(self.data_path + "LAI/",
                                 year,
                                 t_resolution=self.t_resolution,
                                 s_resolution=self.s_resolution)
            np.save(PATH + 'LAI/' + file_name, lai)
            print('LAI shape is {shape}'.format(shape=lai.shape))
            del lai
        # 加载LAI

        file_name = 'ERA5-Land_forcing_{tr}_{sr}_{begin_year}.npy'.format(
            tr=self.t_resolution,
            sr=self.s_resolution,
            begin_year=year)
        if os.path.exists(PATH + 'forcing/' + file_name):
            print('  [DataML] forcing exists')
        else:
            print('  [DataML] making forcing')
            forcing = self._load_forcing(self.data_path + f'forcing/',
                                         self.forcing_list,
                                         year,
                                         self.t_resolution,
                                         self.s_resolution)
            np.save(PATH + 'forcing/' + file_name, forcing)
            print('forcing shape is {shape}'.format(shape=forcing.shape))
            del forcing

        file_name = 'Precipitation_{tr}_{sr}_{begin_year}.npy'.format(
            tr=self.t_resolution,
            sr=self.s_resolution,
            begin_year=year)
        if os.path.exists(PATH + 'forcing/' + file_name):
            print('  [DataML] precipitation exists')
        else:
            print('  [DataML] making precipitation')
            precipitation = self._load_pr_variable(self.data_path + f'forcing/',
                                                   year,
                                                   self.t_resolution,
                                                   self.s_resolution)
            np.save(PATH + 'forcing/' + file_name, precipitation)
            print('precipitation shape is {shape}'.format(shape=precipitation.shape))
            del precipitation

        file_name = 'Temperature_{tr}_{sr}_{begin_year}.npy'.format(
            tr=self.t_resolution,
            sr=self.s_resolution,
            begin_year=year)
        if os.path.exists(PATH + 'forcing/' + file_name):
            print('  [DataML] temperature exists')
        else:
            print('  [DataML] making temperature')
            temperature = self._load_t2m_variable(self.data_path + f'forcing/',
                                                  year,
                                                  self.t_resolution,
                                                  self.s_resolution)
            np.save(PATH + 'forcing/' + file_name, temperature)
            print('temperature shape is {shape}'.format(shape=temperature.shape))
            del temperature

        # print('Now we predict on:{begin_year}'.format(begin_year=year))
        # print('forcing shape is {shape}'.format(shape=forcing.shape))
        # print('LAI shape is {shape}'.format(shape=lai.shape))
        # print('static shape is {shape}'.format(shape=static.shape))
        # print('ET shape is {shape}'.format(shape=ET.shape))
        print('[DataML] preprocessing')

    @timer
    def _load_forcing(self, forcing_root, forcing_list, year, t_resolution, s_resolution):
        fold = '{tr}_{sr}'.format(tr=t_resolution, sr=s_resolution)
        forcing = []
        tmp = []
        file = forcing_root + fold + f"/Precipitation_{self.precipitation_product}_{year}.nc"
        print(file)
        with xr.open_dataset(file) as f:
            tmp.append(f['precipitation'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        for i in range(len(forcing_list)):
            file = forcing_root + fold + "/{var}_{year}.nc".format(year=year, var=self.forcing_list[i])
            with xr.open_dataset(file, drop_variables='time_bnds') as f:
                print(self.var_list[i], f[self.var_list[i]][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1].shape)
                tmp.append(f[self.var_list[i]][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])
        tmp = np.stack(tmp, axis=-1)
        forcing.append(tmp)
        forcing = np.concatenate(forcing, axis=0)
        return forcing

    @timer
    def _load_pr_variable(self,
                          forcing_root,
                          year,
                          t_resolution,
                          s_resolution
                          ):
        fold = '{tr}_{sr}'.format(tr=t_resolution, sr=s_resolution)
        path_in = os.path.join(forcing_root, fold, f"precipitation_{self.precipitation_product}")
        data = []

        file = f"{path_in}/Precipitation_{self.precipitation_product}_intensity_{year}.nc"
        with xr.open_dataset(file, drop_variables='time_bnds') as f:
            data.append(f['precipitation'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        file = f"{path_in}/Precipitation_{self.precipitation_product}_frequency_{year}.nc"
        with xr.open_dataset(file, drop_variables='time_bnds') as f:
            data.append(f['precipitation'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        windows = [3, 7, 10, 15, 30, 90]
        calculates = ['sum', 'mean']
        for calculate in calculates:
            for window in windows:
                file = f"{path_in}/Precipitation_{self.precipitation_product}_{window}_{calculate}_{year}.nc"
                with xr.open_dataset(file, drop_variables='time_bnds') as f:
                    data.append(f['precipitation'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])
        data = np.stack(data, axis=-1)
        # data = np.concatenate(data, axis=0)
        return data

    @timer
    def _load_t2m_variable(self,
                           forcing_root,
                           year,
                           t_resolution,
                           s_resolution
                           ):
        data = []
        fold = '{tr}_{sr}'.format(tr=t_resolution, sr=s_resolution)
        path_in = os.path.join(forcing_root, fold, f"air_temperature_2m")

        # dt---------------------------------------
        file = f"{path_in}/air_temperature_2m_dt_{year}.nc"
        with xr.open_dataset(file, drop_variables='time_bnds') as f:
            data.append(f['t2m'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        windows = [3, 7, 10, 15, 30, 90]
        calculates = ['sum', 'mean']
        for calculate in calculates:
            for window in windows:
                file = f"{path_in}/air_temperature_2m_{window}_{calculate}_{year}.nc"
                with xr.open_dataset(file, drop_variables='time_bnds') as f:
                    data.append(f['t2m'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        windows = [1, 3, 7]
        calculates = ['over', 'beyond']
        for window in windows:
            for calculate in calculates:
                file = f"{path_in}/air_temperature_2m_{calculate}_0_percentage_{window}_{year}.nc"
                with xr.open_dataset(file) as f:
                    data.append(f['t2m'][:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])

        data = np.stack(data, axis=-1)
        print(data.shape)
        # data = np.concatenate(data, axis=0)
        return data

    @timer
    def _load_lai(self, lai_root, year, t_resolution, s_resolution):
        lai_all = []
        with xr.open_dataset(lai_root + 'LAI.nc'.format(tr=t_resolution, sr=s_resolution)) as f:
            lai = np.array(f.lai[:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1])
            if (year % 4 == 0) & (year % 100 != 0) | (year % 400 == 0):
                lai_all.append(lai)
            else:
                idx = np.delete(np.arange(366), 59, axis=0)  # remove 2.29
                lai_all.append(lai[idx])
        lai = np.concatenate(lai_all, axis=0)
        lai = lai[:, :, :, np.newaxis]

        topostd_all = []
        with xr.open_dataset(lai_root + 'topostd.nc'.format(tr=t_resolution, sr=s_resolution)) as f:
            topostd = np.array(f['topostd'].squeeze('sensor'))
            if (year % 4 == 0) & (year % 100 != 0) | (year % 400 == 0):
                new_tmp = np.full((366, topostd.shape[1], topostd.shape[2]), np.nan)
                new_tmp[:59] = topostd[:59]
                new_tmp[60:] = topostd[59:]
                new_tmp[59] = np.mean(topostd[58:60], axis=0)
                topostd_all.append(new_tmp)
            else:
                topostd_all.append(topostd)
        topostd = np.concatenate(topostd_all, axis=0)
        topostd = topostd[:, :, :, np.newaxis]
        data = np.concatenate([lai, topostd], axis=-1)
        return data

    @timer
    def _load_ancillary(self, ancillary_root, s_resolution, year):
        ancillary_root = ancillary_root + '/'
        with xr.open_dataset(ancillary_root + 'kosugi_{s_resolution}.nc'.format(s_resolution=s_resolution)) as f:
            kosugi = np.array(f.kosugi)[:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1]
        with xr.open_dataset(ancillary_root + 'soilgrid_{s_resolution}.nc'.format(s_resolution=s_resolution)) as f:
            soilgrid = np.array(f.soilgrids)[:, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1]

        # tmp = np.stack([soilgrid[0], soilgrid[1],soilgrid[7], soilgrid[8], soilgrid[14], soilgrid[15]], axis=0)
        print(kosugi.shape, soilgrid.shape)
        static = np.concatenate([soilgrid, kosugi], axis=0)
        static = np.transpose(static, (1, 2, 0))[np.newaxis]  # for concat
        time_index = pd.date_range(f"{year}-01-01", f"{year}-12-31", freq="D")
        if year >= 2001 and year <= 2021:
            with xr.open_dataset(ancillary_root + 'LULCC_{year}.nc'.format(year=year)) as f:
                LULCC = np.array(f.LULCC)
        elif year < 2001:
            with xr.open_dataset(ancillary_root + 'LULCC_2001.nc'.format(year=year)) as f:
                LULCC = np.array(f.LULCC)
        else:
            with xr.open_dataset(ancillary_root + 'LULCC_2021.nc'.format(year=year)) as f:
                LULCC = np.array(f.LULCC)
        LULCC = LULCC[np.newaxis, self.lat_idx[0]:self.lat_idx[-1] + 1, self.lon_idx[0]:self.lon_idx[-1] + 1, np.newaxis]
        static = np.concatenate([LULCC, static], axis=-1)
        return static

    @timer
    def _load_runoff(self, data_root, year, t_resolution, s_resolution):
        file = data_root + "/RUNOFF_{year}.nc".format(year=year)
        f = xr.open_dataset(file)
        return f['ro'][:, self.lat_idx[0]: self.lat_idx[-1] + 1, self.lon_idx[0]: self.lon_idx[-1] + 1]