forecast/forecast-algorithm/短期中期部署/部署/OuLiPowerV1.0/test/prediction_test.py

import numpy as np
import pandas as pd
import torch
import os
from datetime import datetime, timedelta
import sys
import argparse
import posixpath
import getpass
import paramiko

weather_utils_path = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..")
)
if weather_utils_path not in sys.path:
    sys.path.insert(0, weather_utils_path)

from scripts.ModelLoading import ModelLoading,SolorFeature


_LOCAL_COLS = ("time_shanghai", "time_local", "time_bj", "时间戳")

def parse_init_time(init_time) -> pd.Timestamp:
    """起报时间 -> tz-aware UTC Timestamp。支持 'YYMMDDHH' / 'YYYYMMDDHH' / ISO / datetime / Timestamp"""
    if isinstance(init_time, pd.Timestamp):
        ts = init_time
    elif hasattr(init_time, "year") and not isinstance(init_time, str):
        ts = pd.Timestamp(init_time)
    else:
        s = str(init_time).strip()
        digits = "".join(ch for ch in s if ch.isdigit())
        if len(s) == 8 and s.isdigit():                     # YYMMDDHH
            ts = pd.to_datetime("20" + s, format="%Y%m%d%H")
        elif len(digits) == 10 and s.isdigit():             # YYYYMMDDHH
            ts = pd.to_datetime(s, format="%Y%m%d%H")
        elif "T" in s or " " in s or "-" in s:
            ts = pd.to_datetime(s, utc=True).tz_localize(None)
        else:
            raise ValueError(f"无法识别的 init_time: {init_time!r}")
    ts = pd.Timestamp(ts)
    return ts.tz_localize("UTC") if ts.tz is None else ts.tz_convert("UTC")

def build_forecast_df(shortTerm_result, init_time, *, freq="15min", tz="Asia/Shanghai",
                      pred_col="power", init_fmt="%y%m%d%H",
                      time_fmt="%Y-%m-%dT%H:%MZ", local_fmt="%Y-%m-%d %H:%M:%S"):
    """
    shortTerm_result : 模型输出, 自动 reshape(-1); 长度 = 预报步数(短期 288 / 中期 960)
    init_time        : 起报时间(UTC)
    返回: DataFrame, 列顺序 init_time(YYMMDDHH) | time(UTC) | time_shanghai | <pred_col>
          时间从起报时刻起, 共 N 个点, 步长 15min
    """
    pred = np.asarray(shortTerm_result).reshape(-1)
    n = pred.size
    if n == 0:
        raise ValueError("shortTerm_result 为空")

    t0 = parse_init_time(init_time)
    utc = pd.date_range(t0, periods=n, freq=freq)
    return pd.DataFrame({
        "init_time": [t0.strftime(init_fmt)] * n,
        "time": utc.strftime(time_fmt),
        "time_shanghai": utc.tz_convert(tz).strftime(local_fmt),
        pred_col: np.asarray(pred, dtype=float),
    })

def remove_night_power(MediumTerm, lat, lon, *, power_col="power", time_col="time",
                       elev_min=0.0, tz="Asia/Shanghai", clip_negative=True, return_info=False):
    """
    MediumTerm    : DataFrame, 含时间列(默认 'time', UTC)与功率列(默认 'power')
    lat, lon      : 场站纬度、经度(度), 正东经/正北纬
    power_col     : 功率列名, 默认 'power'
    time_col      : 时间列名, 默认 'time'(UTC); 传 'time_shanghai' 按北京时间解释;
                    指定列不存在而存在 'time_shanghai' 时自动回退
    elev_min      : 夜间阈值(度), 默认 0(含大气折射修正, 即视地平线)
    clip_negative : 是否同时把负功率置 0
    return_info   : True 时返回 (df, info)
    返回: 处理后的 DataFrame(副本), 列与原表一致, 仅 power 列被改写
    """
    if not isinstance(MediumTerm, pd.DataFrame):
        raise TypeError("MediumTerm 需要是 DataFrame(含时间列); 若只有数值请先自行补上时间列")
    df = MediumTerm.copy()

    # 1. 时间列
    if time_col not in df.columns:
        alt = next((c for c in _LOCAL_COLS if c in df.columns), None)
        if alt is None:
            raise ValueError(f"时间列 {time_col!r} 不存在, 现有列: {list(df.columns)}")
        time_col = alt

    t = pd.to_datetime(df[time_col], errors="coerce")
    if t.isna().any():
        raise ValueError(f"时间列 {time_col!r} 存在无法解析的值({int(t.isna().sum())} 个)")

    is_local = (time_col in _LOCAL_COLS) or time_col.lower().endswith("_shanghai")
    if getattr(t.dt, "tz", None) is not None:
        t_utc = t.dt.tz_convert("UTC")                      # 已带时区 -> 转 UTC
    elif is_local:
        t_utc = t.dt.tz_localize(tz).dt.tz_convert("UTC")   # time_shanghai(北京时) -> UTC
    else:
        t_utc = t.dt.tz_localize("UTC")                     # 'time' 列按 UTC 约定

    # 2. 功率列
    if power_col not in df.columns:
        raise ValueError(f"功率列 {power_col!r} 不在列中: {list(df.columns)}")

    # 3. 太阳高度角 -> 夜间掩码
    d = pd.DatetimeIndex(t_utc.dt.tz_localize(None))
    jd = d.values.astype("datetime64[s]").astype("int64") / 86400.0 + 2440587.5
    T = (jd - 2451545.0) / 36525.0
    L0 = 280.46646 + T * (36000.76983 + T * 0.0003032)
    M = 357.52911 + T * (35999.05029 - 0.0001537 * T)
    e = 0.016708634 - T * (0.000042037 + 0.0000001267 * T)
    Mr = np.radians(M)
    C = (np.sin(Mr) * (1.914602 - T * (0.004817 + 0.000014 * T))
         + np.sin(2 * Mr) * (0.019993 - 0.000101 * T) + np.sin(3 * Mr) * 0.000289)
    omega = 125.04 - 1934.136 * T
    lam = L0 + C - 0.00569 - 0.00478 * np.sin(np.radians(omega))
    eps = (23 + (26 + (21.448 - T * (46.815 + T * (0.00059 - T * 0.001813))) / 60) / 60
           + 0.00256 * np.cos(np.radians(omega)))
    decl = np.degrees(np.arcsin(np.sin(np.radians(eps)) * np.sin(np.radians(lam))))
    y = np.tan(np.radians(eps / 2)) ** 2
    eot = 4 * np.degrees(y * np.sin(2 * np.radians(L0)) - 2 * e * np.sin(Mr)
                         + 4 * e * y * np.sin(Mr) * np.cos(2 * np.radians(L0))
                         - 0.5 * y * y * np.sin(4 * np.radians(L0))
                         - 1.25 * e * e * np.sin(2 * Mr))
    mins = d.hour.values * 60 + d.minute.values + d.second.values / 60.0
    ha = np.radians((mins + eot + 4.0 * lon) / 4.0 - 180.0)          # 时角(经度正东)
    latr, dr = np.radians(lat), np.radians(decl)
    cosz = np.sin(latr) * np.sin(dr) + np.cos(latr) * np.cos(dr) * np.cos(ha)
    el = 90.0 - np.degrees(np.arccos(np.clip(cosz, -1, 1)))
    el = el + np.where(el > -0.575,                                   # 大气折射修正
                       1.02 / (60 * np.tan(np.radians(el + 10.3 / (el + 5.11)))), 0.0)
    night = el <= elev_min

    # 4. 夜间非 0 值置 0
    p = pd.to_numeric(df[power_col], errors="coerce")
    bad = night & p.notna()
    if clip_negative:
        bad = bad | (p < 0).fillna(False).to_numpy()
    changed = bad & p.notna() & (p.to_numpy() != 0)
    info = {"time_col": time_col, "power_col": power_col, "n_rows": len(df),
            "n_night": int(night.sum()), "n_night_nonzero": int((night & (p > 0)).sum()),
            "n_zeroed": int(changed.sum()), "sum_removed": float(p[changed].sum())}
    df[power_col] = p.mask(bad, 0.0)
    return (df, info) if return_info else df

def mkdir_p(sftp, remote_dir):
    """递归创建远程目录，已存在则跳过"""
    dirs = []
    current = remote_dir

    while current not in ("", ".", "/"):
        dirs.append(current)
        current = posixpath.dirname(current)

    for directory in reversed(dirs):
        try:
            sftp.stat(directory)
        except OSError:
            sftp.mkdir(directory)

def upload_files(
    host,
    port,
    username,
    local_files,
    remote_files,
    password=None,
):
    """
    通过 SFTP 批量上传文件。

    参数：
        host:         SFTP 服务器地址
        port:         端口
        username:     用户名
        local_files:  本地文件路径数组
        remote_files: 远程文件路径数组
        password:     密码，可选；不传则运行时输入
    """

    if len(local_files) != len(remote_files):
        raise ValueError("local_files 和 remote_files 数量必须一致")

    if password is None:
        password = getpass.getpass("请输入 SFTP 密码: ")

    transport = paramiko.Transport((host, port))

    try:
        transport.connect(
            username=username,
            password=password
        )

        sftp = paramiko.SFTPClient.from_transport(transport)

        try:
            for local_file, remote_file in zip(local_files, remote_files):

                remote_dir = posixpath.dirname(remote_file)

                # 自动创建远程目录
                mkdir_p(sftp, remote_dir)

                print(f"正在上传:")
                print(f"  本地: {local_file}")
                print(f"  远程: {remote_file}")

                sftp.put(
                    local_file,
                    remote_file,
                    confirm=False
                )

                print("上传成功\n")

        finally:
            sftp.close()

    finally:
        transport.close()

if __name__ == '__main__':

    # 配置光伏场站信息
    parser = argparse.ArgumentParser(description='OuLiPoewev1.0 Prediction')
    parser.add_argument('--projs', type=str, default='JMJD01', help='场站名')
    parser.add_argument('--model_name', type=str, default='GRU', help='模型名称')
    parser.add_argument('--release', type=str, default='v24', help='特征集')
    parser.add_argument('--lat', type=float, default=22.292, help='场站纬度')
    parser.add_argument('--lon', type=float, default=120.290, help='场站经度')
    parser.add_argument('--short_term_seq', type=int, default=288, help='输入序列长度-短期')
    parser.add_argument('--medium_term_seq', type=int, default=960, help='输入序列长度-中期')
    parser.add_argument('--capacity', type=int, default=4050, help='装机容量(kw)')
    parser.add_argument('--saved_model_root', type=str,
                            default='/home/jinhuaiming/OuLiPowerV1.0/saved_model', 
                            help='模型权重保存根路径')
    parser.add_argument('--saved_result_root', type=str,
                            default='/home/jinhuaiming/OuLiPowerV1.0/archive', 
                            help='预报结果存档根路径')
    parser.add_argument('--init_time', type=str, default='2026060200', help='预报时间(UTC)')
    args = parser.parse_args()

    #=====读取气象数据--气象读取部分=========================================================
    # 气象变量特征
        # 短期：v24-【'ssrd','tcc','lcc','mcc','hcc','2t','2d','10u','10v'】 
        # 中期：v24-【'ssrd','tcc','lcc','mcc','hcc','2t','2d','10u','10v'】
        # 单位：       w/m2, 0-100,0-100,0-100,0-100,°C,  °C,   m/s,  m/s
    # 模型【输入输出】维度
        # 短期：输入 =（1，288，9），输出 =（1，288）
        # 中期：输出 =（1，960，9），输出 =（1，960）
        # 注意⚠️：特征顺序不要改变，保持上述描述顺序。

    # 根据 某个站点(args.projs)，预报时间(args.init_time)取的气象数据
    init_time = args.init_time  #YYYYMMDDHH
    ShortTerm_input = torch.randn(1, 288, 9)
    MediumTerm_input = torch.randn(1, 960, 9)

    #=====加载模型=======================================================================
    # 加载短期功率预测模型
    ShortTerm_Modelloader = ModelLoading(
                    trial = 'None',
                    projs = args.projs,
                    model_name = args.model_name,
                    release = args.release,
                    input_seq = args.short_term_seq,
                    saved_model_path = args.saved_model_root,
                    predict_mode = True,
                    capacity = args.capacity
                    )
    # 加载中期功率预测模型
    MediumTerm_Modelloader = ModelLoading(
                    trial = 'None',
                    projs = args.projs,
                    model_name = args.model_name,
                    release = args.release,
                    input_seq = args.medium_term_seq,
                    saved_model_path = args.saved_model_root,
                    predict_mode = True,
                    capacity = args.capacity
                    )

    #=====预测==========================================================================
    # 短期预测 -输出（288,）
    shortTerm_result = ShortTerm_Modelloader.predict(ShortTerm_input).detach().numpy().reshape(-1)
    ShortTerm_df = build_forecast_df(shortTerm_result, args.init_time)
    # 中期预测 - 输出（960,）
    MediumTerm_result = MediumTerm_Modelloader.predict(MediumTerm_input).detach().numpy().reshape(-1)
    MediumTerm_df = build_forecast_df(MediumTerm_result, args.init_time)

    #=====后处理-消除夜间非0值和可能存在的赋值===============================================

    ShortTerm_df = remove_night_power(ShortTerm_df, lat=args.lat, lon=args.lon)
    MediumTerm_df = remove_night_power(MediumTerm_df, lat=args.lat, lon=args.lon)

    print('测试成功🏅')