AcidificationModel/api/app/model.py

import datetime
import os
import pickle
import pandas as pd
from flask_sqlalchemy.session import Session
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split, cross_val_score
from sqlalchemy import text

from .database_models import Models, Datasets



# 加载模型
def load_model(model_name):
    file_path = f'model_optimize/pkl/{model_name}.pkl'
    with open(file_path, 'rb') as f:
        return pickle.load(f)


# 模型预测
def predict(input_data: pd.DataFrame, model_name):
    # 初始化模型
    model = load_model(model_name)  # 根据指定的模型名加载模型
    predictions = model.predict(input_data)
    return predictions.tolist()



def train_and_save_model(session, model_type, model_name, model_description, data_type, dataset_id=None):
    if not dataset_id:
        # 直接创建新的数据集并复制数据
        dataset_id = save_current_dataset(session, data_type)

    # 从新复制的数据集表中加载数据
    dataset_table_name = f"dataset_{dataset_id}"
    dataset = pd.read_sql_table(dataset_table_name, session.bind)

    if dataset.empty:
        raise ValueError(f"Dataset {dataset_id} is empty or not found.")

    # 数据准备
    X = dataset.iloc[:, :-1]
    y = dataset.iloc[:, -1]

    # 训练模型
    model = train_model_by_type(X, y, model_type)

    # 保存模型到数据库
    save_model(session, model, model_name, model_type, model_description, dataset_id, data_type)


# # 保存模型参数
    # save_model_parameters(model, saved_model.ModelID)

    # # 计算评估指标（如MSE）
    # y_pred = model.predict(X)
    # mse = mean_squared_error(y, y_pred)
    #
    # return saved_model, mse

def save_current_dataset(session, data_type):
    """
    创建一个新的数据集条目，并复制对应的数据类型表的数据。

    Args:
    session (Session): SQLAlchemy session对象。
    data_type (str): 数据集的类型，如 'reduce' 或 'reflux'。

    Returns:
    int: 新保存的数据集的ID。
    """
    # 创建一个新的数据集条目
    new_dataset = Datasets(
        Dataset_name=f"{data_type}_dataset_{datetime.datetime.now():%Y%m%d_%H%M%S}",  # 使用当前时间戳生成独特的名称
        Dataset_description=f"Automatically generated dataset for type {data_type}",
        Row_count=0,  # 初始行数为0，将在复制数据后更新
        Status='pending',  # 初始状态为pending
        Dataset_type=data_type
    )

    # 添加到数据库并提交以获取ID
    session.add(new_dataset)
    session.flush()  # flush用于立即执行SQL并获取ID，但不提交事务

    # 获取新数据集的ID
    dataset_id = new_dataset.Dataset_ID

    # 复制数据到新表
    source_table = data_type_table_mapping(data_type)  # 假设有函数映射数据类型到表名
    new_table_name = f"dataset_{dataset_id}"
    copy_table_sql = f"CREATE TABLE {new_table_name} AS SELECT * FROM {source_table};"
    session.execute(text(copy_table_sql))

    # 更新新数据集的状态和行数
    update_sql = f"UPDATE datasets SET status='processed', row_count=(SELECT count(*) FROM {new_table_name}) WHERE dataset_id={dataset_id};"
    session.execute(text(update_sql))

    session.commit()

    return dataset_id

def data_type_table_mapping(data_type):
    """映射数据类型到对应的数据库表名"""
    if data_type == 'reduce':
        return 'current_reduce'
    elif data_type == 'reflux':
        return 'current_reflux'
    else:
        raise ValueError("Invalid data type provided.")


def train_model_by_type(X, y, model_type):
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    if model_type == 'RandomForest':
        # 随机森林的参数优化
        return train_random_forest(X_train, y_train)
    elif model_type == 'XGBR':
        # XGBoost的参数优化
        return train_xgboost(X_train, y_train)
    elif model_type == 'GBSTR':
        # 梯度提升树的参数优化
        return train_gradient_boosting(X_train, y_train)
    else:
        raise ValueError(f"Unsupported model type: {model_type}")


def train_random_forest(X_train, y_train):
    best_score = 0
    best_n_estimators = None
    best_max_depth = None
    random_state = 43

    # 筛选最佳的树的数量
    for n_estimators in range(1, 20, 1):
        model = RandomForestRegressor(n_estimators=n_estimators, random_state=random_state)
        score = cross_val_score(model, X_train, y_train, cv=5).mean()
        if score > best_score:
            best_score = score
            best_n_estimators = n_estimators

    print(f"Best number of trees: {best_n_estimators}, Score: {best_score}")

    # 在找到的最佳树的数量基础上，筛选最佳的最大深度
    best_score = 0  # 重置最佳得分，为最大深度优化做准备
    for max_depth in range(1, 30, 1):
        model = RandomForestRegressor(n_estimators=best_n_estimators, max_depth=max_depth, random_state=random_state)
        score = cross_val_score(model, X_train, y_train, cv=5).mean()
        if score > best_score:
            best_score = score
            best_max_depth = max_depth

    print(f"Best max depth: {best_max_depth}, Score: {best_score}")

    # 使用最佳的树的数量和最大深度训练最终模型
    best_model = RandomForestRegressor(n_estimators=best_n_estimators, max_depth=best_max_depth,
                                       random_state=random_state)
    best_model.fit(X_train, y_train)

    return best_model


def train_xgboost(X_train, y_train, X_test, y_test):
    # XGBoost训练过程
    # (将类似上面的代码添加到这里)
    pass


def train_gradient_boosting(X_train, y_train, X_test, y_test):
    # 梯度提升树训练过程
    # (将类似上面的代码添加到这里)
    pass

def save_model(session, model, model_name, model_type, model_description, dataset_id, data_type, custom_path='pkl'):
    """
    保存模型到数据库，并将模型文件保存到磁盘。

    :param session: 数据库会话
    :param model: 要保存的模型对象
    :param model_name: 模型的名称
    :param model_type: 模型的类型
    :param model_description: 模型的描述信息
    :param dataset_id: 数据集ID
    :param custom_path: 保存模型的路径
    :return: 返回保存的模型文件路径
    """
    # 根据模型类型设置文件名前缀
    prefix_dict = {
        'RandomForest': 'rf_model_',
        'XGBRegressor': 'xgbr_model_',
        'GBSTRegressor': 'gbstr_model_'
    }
    prefix = prefix_dict.get(model_type, 'default_model_')  # 如果model_type不在字典中，默认前缀

    try:
        # 确保路径存在
        os.makedirs(custom_path, exist_ok=True)

        # 获取当前时间戳（格式：月日时分）
        timestamp = datetime.datetime.now().strftime('%m%d_%H%M')

        # 拼接完整的文件名
        file_name = os.path.join(custom_path, f'{prefix}{timestamp}.pkl')

        # 保存模型到文件
        with open(file_name, 'wb') as f:
            pickle.dump(model, f)
        print(f"模型已保存为: {file_name}")

        # 创建模型数据库记录
        new_model = Models(
            Model_name=model_name,
            Model_type=model_type,
            Description=model_description,
            DatasetID=dataset_id,
            Created_at=datetime.datetime.now(),
            ModelFilePath=file_name,
            Data_type=data_type
        )

        # 添加记录到数据库
        session.add(new_model)
        session.commit()

        # 返回文件路径
        return file_name

    except Exception as e:
        session.rollback()
        print(f"Error saving model: {str(e)}")
        raise e  # 显式抛出异常供调用者处理



if __name__ == '__main__':
    # 反酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "organic_matter": 5.2,
        "chloride": 3.1,
        "cec": 25.6,
        "h_concentration": 0.5,
        "hn": 12.4,
        "al_concentration": 0.8,
        "free_alumina": 1.2,
        "free_iron": 0.9,
        "delta_ph": -0.2
    }])

    model_name = 'RF_filt'

    Acid_reflux_result = predict(input_data, model_name)
    print("Acid_reflux_result:", Acid_reflux_result)  # 预测结果

    # 降酸模型预测
    # 测试 predict 函数
    input_data = pd.DataFrame([{
        "pH": 5.2,
        "OM": 3.1,
        "CL": 25.6,
        "H": 0.5,
        "Al": 12.4
    }])

    model_name = 'rf_model_1214_1008'
    Acid_reduce_result = predict(input_data, model_name)
    print("Acid_reduce_result:", Acid_reduce_result)  # 预测结果