compete/app/predict/tf_model_pre.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :tf_lstm_pre.py
# @Time      :2025/2/13 10:52
# @Author    :David
# @Company: shenyang JY
import os.path
import pandas as pd
import numpy as np
import logging, argparse, traceback
from app.common.data_handler import DataHandler, write_number_to_file
from threading import Lock
import time, json

model_lock = Lock()
from itertools import chain
from typing import Dict, Any
from app.common.config import logger, parser
from app.common.tf_lstm import TSHandler
from app.common.dbmg import MongoUtils

np.random.seed(42)  # NumPy随机种子

mgUtils = MongoUtils(logger)
class ModelPre(object):
    """模型训练器封装类"""

    def __init__(self,
                 pre_data: pd.DataFrame,
                 capacity: float,
                 config: Dict[str, Any] = None,
                 ):
        self.config = config
        self.logger = logger
        self.pre_data = pre_data
        self.capacity = capacity
        self.gpu_id = config.get('gpu_assignment')
        self._setup_resources()

        # 初始化组件
        self.input_file = config.get("input_file")
        self.opt = argparse.Namespace(**config)
        self.dh = DataHandler(logger, self.opt)
        self.ts = TSHandler(logger, self.opt)
        self.mgUtils = MongoUtils(logger)

    def _setup_resources(self):
        """GPU资源分配"""
        if self.gpu_id is not None:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id)
            self.logger.info(f"GPU {self.gpu_id} allocated")

    def predict(self, pre_area=False):
        # 获取程序开始时间
        start_time = time.time()
        success = 0
        print("Program starts execution!")
        pre_id = self.config['area_id'] if pre_area else self.config['station_id']
        pre_type = 'a' if pre_area else 's'
        output_file = os.path.join(self.opt.dqyc_base_path, self.input_file)
        output_file = output_file.replace('IN', 'OUT') if pre_area else os.path.join(str(output_file), pre_id)
        file = 'DQYC_OUT_PREDICT_POWER.txt'
        status_file = 'STATUS.TXT'
        try:
            self.config['model_table'] = self.config['model_table'] + f'_{pre_type}_'+str(pre_id)
            self.config['scaler_table'] = self.config['scaler_table'] + f'_{pre_type}_'+ str(pre_id)
            feature_scaler, target_scaler = mgUtils.get_scaler_model_from_mongo(self.config)
            self.ts.opt.cap = round(target_scaler.transform(np.array([[self.capacity]]))[0, 0], 2)
            self.ts.get_model(self.config)
            print("!!!!", self.ts.model_params)
            self.dh.opt.features = json.loads(self.ts.model_params).get('Model').get('features', ','.join(self.ts.opt.features)).split(',')
            scaled_pre_x, pre_data = self.dh.pre_data_handler(self.pre_data, feature_scaler)

            success = 1
            # 更新算法状态：1. 启动成功
            write_number_to_file(os.path.join(str(output_file), status_file), 1, 1, 'rewrite')
            logger.info("算法启动成功")
            res = list(chain.from_iterable(target_scaler.inverse_transform([self.ts.predict(scaled_pre_x).flatten()])))
            pre_data['Power'] = res[:len(pre_data)]
            pre_data['PlantID'] = pre_id
            pre_data = pre_data[['PlantID', self.config['col_time'], 'Power']]

            pre_data.loc[:, 'Power'] = pre_data['Power'].round(2)
            pre_data.loc[pre_data['Power'] > self.capacity, 'Power'] = self.capacity
            pre_data.loc[pre_data['Power'] < 0, 'Power'] = 0
            pre_data.to_csv(os.path.join(str(output_file), file), sep=' ', index=False)
            # 更新算法状态：正常结束
            write_number_to_file(os.path.join(str(output_file), status_file), 2, 2)
            logger.info("算法正常结束")
        except Exception as e:
            # 如果算法状态没启动，不更新
            if success:
                write_number_to_file(os.path.join(str(output_file), status_file), 2, 3)
            my_exception = traceback.format_exc()
            my_exception.replace("\n", "\t")
            logger.info("算法状态异常：{}".format(my_exception))
        end_time = time.time()
        logger.info("lstm预测任务：用了 %s 秒 " % (end_time - start_time))


if __name__ == "__main__":
    print("Program starts execution!")
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger = logging.getLogger("model_training_bp log")

    # serve(app, host="0.0.0.0", port=1010x, threads=4)
    print("server start!")

    # ------------------------测试代码------------------------
    args_dict = {"mongodb_database": 'david_test', 'scaler_table': 'j00083_scaler', 'model_name': 'bp1.0.test',
                 'model_table': 'j00083_model', 'mongodb_read_table': 'j00083_test', 'col_time': 'date_time', 'mongodb_write_table': 'j00083_rs',
                 'features': 'speed10,direction10,speed30,direction30,speed50,direction50,speed70,direction70,speed90,direction90,speed110,direction110,speed150,direction150,speed170,direction170'}
    args_dict['features'] = args_dict['features'].split(',')
    arguments.update(args_dict)
    dh = DataHandler(logger, arguments)
    ts = TSHandler(logger)
    opt = argparse.Namespace(**arguments)

    opt.Model['input_size'] = len(opt.features)
    pre_data = get_data_from_mongo(args_dict)
    feature_scaler, target_scaler = get_scaler_model_from_mongo(arguments)
    pre_x = dh.pre_data_handler(pre_data, feature_scaler, opt)
    ts.get_model(arguments)
    result = ts.predict(pre_x)
    result1 = list(chain.from_iterable(target_scaler.inverse_transform([result.flatten()])))
    pre_data['power_forecast'] = result1[:len(pre_data)]
    pre_data['farm_id'] = 'J00083'
    pre_data['cdq'] = 1
    pre_data['dq'] = 1
    pre_data['zq'] = 1
    pre_data.rename(columns={arguments['col_time']: 'date_time'}, inplace=True)
    pre_data = pre_data[['date_time', 'power_forecast', 'farm_id', 'cdq', 'dq', 'zq']]

    pre_data['power_forecast'] = pre_data['power_forecast'].round(2)
    pre_data.loc[pre_data['power_forecast'] > opt.cap, 'power_forecast'] = opt.cap
    pre_data.loc[pre_data['power_forecast'] < 0, 'power_forecast'] = 0

    insert_data_into_mongo(pre_data, arguments)