Step3_Process.py

import time
import numpy as np
import xarray as xr
import pandas as pd
import glob, os
import os, math
from config import get_tabpfn_args


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) / 3600.
        print('%s cost time: %.3f hours' % (func.__name__, time_spend))
        return result

    return func_wrapper


class Process:
    def __init__(self, cfg, prefix):
        self.cfg = cfg
        self.prefix = prefix

    # def filter_nan(self, s=np.array([]), o=np.array([])):
    #     """
    #     this functions removed the data from simulated and observed data
    #     whereever the observed data contains nan
    #
    #     this is used by all other functions, otherwise they will produce nan as
    #     output
    #     """
    #     data = np.array([s.flatten(), o.flatten()])
    #     data = np.transpose(data)
    #     data = data[~np.isnan(data).any(1)]
    #     return data[:, 0], data[:, 1]
    #
    # def r2_score(self, o, s):
    #     """
    #     correlation coefficient R2
    #     input:
    #         s: simulated
    #         o: observed
    #     output:
    #         correlation: correlation coefficient
    #     """
    #     s, o = self.filter_nan(s, o)
    #     if o.size == 0:
    #         corr = np.nan
    #         return corr
    #     else:
    #         corr = np.corrcoef(o, s)[0, 1]
    #         return corr ** 2
    #
    # def rmse_score(self, y_true, y_pred):
    #     """
    #     计算RMSE分数
    #     """
    #     y_pred, y_true = filter_nan(y_pred, y_true)
    #     if y_true.size == 0:
    #         return np.nan
    #     else:
    #         mse = np.mean((y_pred - y_true) ** 2)
    #         RMSE = np.sqrt(mse)
    #         return RMSE
    #
    # def kge_score(self, y_true, y_pred):
    #     """
    #     Kling-Gupta Efficiency
    #     input:
    #         y_pred: simulated
    #         y_true: observed
    #     output:
    #         kge: Kling-Gupta Efficiency
    #         r: correlation
    #         std_ratio: ratio of the standard deviation
    #         bias_ratio: ratio of the mean
    #     """
    #     y_pred, y_true = self.filter_nan(y_pred, y_true)
    #     if y_true.size == 0:
    #         return np.nan
    #     else:
    #         r = np.corrcoef(y_true, y_pred)[0, 1]
    #         std_ratio = np.std(y_pred) / np.std(y_true)
    #         bias_ratio = np.mean(y_pred) / np.mean(y_true)
    #         KGE = 1 - np.sqrt((r - 1) ** 2 + (std_ratio - 1) ** 2 + (bias_ratio - 1) ** 2)
    #         return KGE
    #
    # def bias(self, o, s):
    #     """
    #     Bias
    #     input:
    #         s: simulated
    #         o: observed
    #     output:
    #         bias: bias
    #     """
    #     # np.mean(s-o)
    #     var = np.mean(s - o)  # (s - o).mean(dim='time')
    #     return var
    #
    # def percent_bias(self, o, s):
    #     """
    #     Calculate Percent Bias.
    #
    #     Args:
    #         s (xr.DataArray): Simulated data
    #         o (xr.DataArray): Observed data
    #
    #     Returns:
    #         xr.DataArray: Percent bias
    #     """
    #     s, o = self.filter_nan(s, o)
    #     if np.sum(o) < 0.01:
    #         return np.nan
    #     else:
    #         return 100.0 * np.sum(s - o) / np.sum(o)
    #
    #     # return 100.0 * (s - o).sum(dim='time') / o.sum(dim='time')

    def _validate_inputs(self, s, o):
        """
        Validate and align input DataArrays.

        Args:
            s (xr.DataArray): Simulated data
            o (xr.DataArray): Observed data

        Returns:
            tuple: Aligned and validated DataArrays
        """
        # Ensure inputs are xarray DataArrays
        if not isinstance(s, xr.DataArray) or not isinstance(o, xr.DataArray):
            logging.error('Inputs must be xarray DataArrays')
            raise TypeError("Inputs must be xarray DataArrays")

        # Align time dimensions
        s, o = xr.align(s, o, join='inner')

        # Remove NaN values
        mask = np.isfinite(s) & np.isfinite(o)
        return s.where(mask), o.where(mask)

    def percent_bias(self, s, o):
        """
        Calculate Percent Bias.

        Args:
            s (xr.DataArray): Simulated data
            o (xr.DataArray): Observed data

        Returns:
            xr.DataArray: Percent bias
        """
        s, o = self._validate_inputs(s, o)
        return 100.0 * (s - o).sum(dim='time') / o.sum(dim='time')

    def correlation(self, s, o):
        """
        correlation coefficient
        input:
            s: simulated
            o: observed
        output:
            correlation: correlation coefficient
        """
        corr = xr.corr(s, o, dim=['time'])

        return corr
    def correlation_R2(self, s, o):
        """
        correlation coefficient R2
        input:
            s: simulated
            o: observed
        output:
            correlation: correlation coefficient
        """

        return xr.corr(s, o, dim=['time']) ** 2

    def KGE(self, s, o):
        """
        Kling-Gupta Efficiency
        input:
            s: simulated
            o: observed
        output:
            kge: Kling-Gupta Efficiency
            cc: correlation
            alpha: ratio of the standard deviation
            beta: ratio of the mean
        """
        cc = self.correlation(s, o)
        alpha = s.std(dim='time') / o.std(dim='time')
        # alpha = np.std(s)/np.std(o)
        beta = s.mean(dim='time') / o.mean(dim='time')
        # beta = np.sum(s)/np.sum(o)
        kge = 1 - ((cc - 1) ** 2 + (alpha - 1) ** 2 + (beta - 1) ** 2) ** 0.5
        return kge  # , cc, alpha, beta

    def get_index_info(self, year, lat, lon):
        time_1D = pd.date_range(f"{year}-01-01", f"{year}-12-31", freq="1D")
        time = xr.DataArray(time_1D, dims='time')
        time.attrs['standard_name'] = 'time'
        time.attrs['long_name'] = 'time'
        time.attrs['axis'] = 'T'
        n = len(time_1D)

        latitude = xr.DataArray(lat, dims='latitude')
        latitude.attrs['_FillValue'] = float('nan')
        latitude.attrs['standard_name'] = 'latitude'
        latitude.attrs['long_name'] = 'latitude'
        latitude.attrs['units'] = 'degrees_north'
        latitude.attrs['axis'] = 'Y'

        longitude = xr.DataArray(lon, dims='longitude')
        longitude.attrs['_FillValue'] = float('nan')
        longitude.attrs['standard_name'] = 'longitude'
        longitude.attrs['long_name'] = 'longitude'
        longitude.attrs['units'] = 'degrees_east'
        longitude.attrs['axis'] = 'X'
        return n, time, latitude, longitude

    @timer
    def main(self, year):
        print(' [DataML] loading lat/lon')
        path = self.cfg["grid_root"]
        outpath = self.cfg["grid_output_root"] + f'/{self.prefix}_TabPFN/'

        lat_file_name = 'lat_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lon_file_name = 'lon_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lat, lon = np.load(path + lat_file_name), np.load(path + lon_file_name)
        n, time, latitude, longitude = self.get_index_info(year, lat, lon)
        print(latitude.shape, longitude.shape)

        print('Preparing Data')
        time_monthly_groups = time.groupby("time.month")

        i_length = 0
        for month_num, month_group in time_monthly_groups:
            print(month_num)
            month_times = month_group.values[:24]

            pred = xr.Dataset({'Runoff': (('time', 'latitude', 'longitude'),
                                          np.full((len(month_times), len(latitude), len(longitude)), np.nan)),
                               'GRFR': (('time', 'latitude', 'longitude'),
                                        np.full((len(month_times), len(latitude), len(longitude)), np.nan))
                               },
                              coords={"time": month_times, "latitude": latitude, "longitude": longitude})

            for i, i_mub in enumerate(range(i_length, i_length + len(month_times))):  #
                if os.path.exists(outpath + f"Runoff_predict_{year}_{i_mub}.npy"):
                    data = np.load(outpath + f"Runoff_predict_{year}_{i_mub}.npy")
                    print(i, data.shape)
                    pred["Runoff"][i, :, :] = data[0]
                else:
                    exit()

            runoff_input = np.load(path + 'Runoff/' + 'Runoff_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
            print(runoff_input.shape)
            pred["GRFR"][:] = runoff_input[i_length: i_length + len(month_times), :751, :1350]
            del runoff_input
            pred['Runoff'].attrs['_FillValue'] = float('nan')
            pred['Runoff'].attrs['long_name'] = "Runoff"
            pred['Runoff'].attrs['units'] = "mm/day"
            pred['Runoff'].attrs['missing_value'] = float('nan')
            pred.to_netcdf(outpath + "Runoff_predict_{year}_{month}_{tr}_{sr}.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month_num,
                                                                                         sr=self.cfg["spatial_resolution"]), engine='netcdf4')
            i_length = i_length + len(month_times)

    def after_process(self, year):
        print(' [DataML] loading lat/lon')
        path = self.cfg["grid_root"]
        outpath = self.cfg["grid_output_root"] + f'/{self.prefix}_TabPFN/'

        lat_file_name = 'lat_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lon_file_name = 'lon_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lat, lon = np.load(path + lat_file_name), np.load(path + lon_file_name)
        n, time, latitude, longitude = self.get_index_info(year, lat, lon)
        print(latitude.shape, longitude.shape)

        for month in range(1, 13):
            r2 = np.full((len(latitude), len(longitude)), np.nan)
            percent_bias = np.full((len(latitude), len(longitude)), np.nan)
            kge = np.full((len(latitude), len(longitude)), np.nan)

            filename = outpath + 'merge/' + "Runoff_predict_{year}_{month}_{tr}_{sr}.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month,
                                                                                                sr=self.cfg["spatial_resolution"])
            if os.path.exists(filename):
                metrics_name = "Metrics_{year}_{month}_{tr}_{sr}.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month,
                                                                            sr=self.cfg["spatial_resolution"])
                ds = xr.open_dataset(filename)
                r2 = self.correlation_R2(ds['Runoff'], ds['GRFR'])
                percent_bias = self.percent_bias(ds['Runoff'], ds['GRFR'])
                kge = self.KGE(ds['Runoff'], ds['GRFR'])
                # Runoff = ds['Runoff'].values
                # GRFR = ds['GRFR'].values
                # for i in range(len(latitude)):
                #     for j in range(len(longitude)):
                #         d, c = GRFR[:, i, j], Runoff[:, i, j]
                #         if not np.isnan(d).all():
                #             r2[i, j] = self.r2_score(d, c)
                #             percent_bias[i, j] = self.percent_bias(d, c)
                #             kge[i, j] = self.kge_score(d, c)
                metric = xr.Dataset(data_vars={'r2': (['latitude', 'longitude'],
                                                      r2.values),
                                               'percent_bias': (['latitude', 'longitude'],
                                                                percent_bias.values),
                                               'kge': (['latitude', 'longitude'],
                                                       kge.values)},
                                    coords={"latitude": latitude, "longitude": longitude})
                metric.to_netcdf(outpath + 'merge/' + f"{metrics_name}.nc", engine='netcdf4')
                del kge, r2, percent_bias

    def after_mean(self, year):
        outpath = self.cfg["grid_output_root"] + f'/{self.prefix}_TabPFN/'

        for month in range(1, 13):
            filename = outpath + 'merge/' + "Runoff_predict_{year}_{month}_{tr}_{sr}.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month,
                                                                                                sr=self.cfg["spatial_resolution"])
            if os.path.exists(filename):
                ds = xr.open_dataset(filename)
                ds = ds.resample(time='1M').mean()
                ds.to_netcdf(
                    outpath + 'merge/' + "Runoff_predict_{year}_{month}_{tr}_{sr}_mean.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month,
                                                                                                  sr=self.cfg["spatial_resolution"]), engine='netcdf4')

    def forcing_check(self, year):
        path = self.cfg["grid_root"]

        outpath = self.cfg["grid_output_root"] + f'/{self.prefix}_TabPFN/'
        os.makedirs(outpath, exist_ok=True)
        print('Grid output path:', outpath)
        print('=' * 50, '\n\n')

        print('[DataML] loading input data')
        print(' [DataML] loading lat/lon')
        lat_file_name = 'lat_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lon_file_name = 'lon_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        lat, lon = np.load(path + lat_file_name), np.load(path + lon_file_name)
        n, time_info, latitude, longitude = self.get_index_info(year, lat, lon)
        print('Preparing Data')

        for i in range(2):  #
            if not os.path.exists(outpath + f"forcing_check_{year}_{i}.nc"):
                print(f'Time to: {time_info[i].values}')
                forcing_data = xr.Dataset({'forcing': (('latitude', 'longitude', 'n'),
                                                       np.full((len(latitude), len(longitude), 73), np.nan))
                                           },
                                          coords={"latitude": latitude, "longitude": longitude, 'n': range(73)})

                print(' [DataML] loading LAI')
                lai_input = np.load(path + 'LAI/' + 'MODIS_LAI_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
                lai = lai_input[i].copy()
                lai_input._mmap.close()

                print(' [DataML] loading static')
                static_input = np.load(path + 'ancillary/' + f'static_1D_0p1_{year}.npy', mmap_mode='r')
                static = static_input[0].copy()
                static_input._mmap.close()

                print(' [DataML] loading forcing')
                forcing_input = np.load(path + 'forcing/' + 'ERA5-Land_forcing_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
                forcing = forcing_input[i].copy()
                forcing_input._mmap.close()

                print(' [DataML] loading precipitation')
                precipitation_input = np.load(path + 'forcing/' + 'Precipitation_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
                precipitation = precipitation_input[i].copy()
                precipitation_input._mmap.close()
                print(precipitation.shape)

                print(' [DataML] loading temperature')
                temperature_input = np.load(path + 'forcing/' + 'Temperature_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
                temperature = temperature_input[i].copy()
                temperature_input._mmap.close()
                print(temperature.shape)

                forcing_data["forcing"][:] = np.concatenate([forcing, precipitation, temperature, lai, static], axis=-1)
                forcing_data.to_netcdf(outpath + "forcing_check_{year}_{i}.nc".format(year=year, i=i), engine='netcdf4')

    def main_process(self):
        for year in range(self.cfg["test_begin_year"], self.cfg["test_end_year"] + 1):
            # self.main(year)
            self.after_process(year)
            self.after_mean(year)
            # self.forcing_check(year)

    def mixed_process(self):
        # for year in range(self.cfg["test_begin_year"], self.cfg["test_end_year"] + 1):
        #     print(' [DataML] loading lat/lon')
        #     path = self.cfg["grid_root"]
        #     outpath = self.cfg["grid_output_root"] + f'/{self.prefix}_TabPFN/'
        #
        #     lat_file_name = 'lat_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        #     lon_file_name = 'lon_{tr}_{sr}.npy'.format(tr=self.cfg["temporal_resolution"], sr=self.cfg["spatial_resolution"])
        #     lat, lon = np.load(path + lat_file_name), np.load(path + lon_file_name)
        #     n, time, latitude, longitude = self.get_index_info(year, lat, lon)
        #     print(latitude.shape, longitude.shape)
        #
        #     print('Preparing Data')
        #     time_monthly_groups = time.groupby("time.month")
        #
        #     i_length = 0
        #     for month_num, month_group in time_monthly_groups:
        #         print(month_num)
        #         month_times = month_group.values[:27]
        #
        #         pred = xr.Dataset({'Runoff': (('time', 'latitude', 'longitude'),
        #                                       np.full((len(month_times), len(latitude), len(longitude)), np.nan)),
        #                            'GRFR': (('time', 'latitude', 'longitude'),
        #                                     np.full((len(month_times), len(latitude), len(longitude)), np.nan))
        #                            },
        #                           coords={"time": month_times, "latitude": latitude, "longitude": longitude})
        #
        #         for i, i_mub in enumerate(range(i_length, i_length + len(month_times))):  #
        #             if os.path.exists(outpath + f"Runoff_predict_{year}_{i_mub}.npy"):
        #                 data = np.load(outpath + f"Runoff_predict_{year}_{i_mub}.npy")
        #                 print(i, data.shape)
        #                 pred["Runoff"][i, :, :] = data[0]
        #             elif os.path.exists(outpath + f"Runoff_predict_{year}_{i_mub}_clf.npy") & os.path.exists(
        #                     outpath + f"Runoff_predict_{year}_{i_mub}_reg.npy"):
        #                 clfdata = np.load(outpath + f"Runoff_predict_{year}_{i_mub}_clf.npy")
        #                 regdata = np.load(outpath + f"Runoff_predict_{year}_{i_mub}_reg.npy")
        #                 print('mix', i, regdata.shape)
        #                 pred["Runoff"][i, :, :] = clfdata[0] * regdata[0]
        #             else:
        #                 exit()
        #
        #         runoff_input = np.load(path + 'Runoff/' + 'Runoff_1D_0p1_{year}.npy'.format(year=year), mmap_mode='r')
        #         print(runoff_input.shape)
        #         pred["GRFR"][:] = runoff_input[i_length: i_length + len(month_times), :751, :1350]
        #         del runoff_input
        #         pred['Runoff'].attrs['_FillValue'] = float('nan')
        #         pred['Runoff'].attrs['long_name'] = "Runoff"
        #         pred['Runoff'].attrs['units'] = "mm/day"
        #         pred['Runoff'].attrs['missing_value'] = float('nan')
        #         pred.to_netcdf(
        #             outpath + 'merge/' + "Runoff_predict_{year}_{month}_{tr}_{sr}.nc".format(year=year, tr=self.cfg["temporal_resolution"], month=month_num,
        #                                                                                      sr=self.cfg["spatial_resolution"]), engine='netcdf4')
        #         i_length = i_length + len(month_times)

        for year in range(self.cfg["test_begin_year"], self.cfg["test_end_year"] + 1):
            self.after_process(year)
            self.after_mean(year)