AcidMap/Cd_Prediction_Integrated_System/config.py

"""
配置文件 - Cd预测集成系统
Configuration file for Cd Prediction Integrated System
"""

import os

# 项目根目录
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

# 模型相关路径
MODELS_DIR = os.path.join(PROJECT_ROOT, "models")

# 作物Cd模型配置
CROP_CD_MODEL = {
    "model_dir": os.path.join(MODELS_DIR, "crop_cd_model"),
    "model_file": "cropCdNN.pth",
    "mean_file": "cropCd_mean.npy", 
    "scale_file": "cropCd_scale.npy",
    "constraint_module": "constrained_nn6",
    "input_data": "areatest.csv",
    "output_file": "combined_pH.csv",
    "special_feature": "solution"
}

# 有效态Cd模型配置
EFFECTIVE_CD_MODEL = {
    "model_dir": os.path.join(MODELS_DIR, "effective_cd_model"),
    "model_file": "EffCdNN6C.pth",
    "mean_file": "EffCd_mean.npy",
    "scale_file": "EffCd_scale.npy", 
    "constraint_module": "constrained_nn6C",
    "input_data": "areatest.csv",
    "output_file": "pHcombined.csv",
    "special_feature": "Cdsolution"
}

# 数据路径配置
DATA_PATHS = {
    "coordinates_dir": os.path.join(PROJECT_ROOT, "data", "coordinates"),
    "predictions_dir": os.path.join(PROJECT_ROOT, "data", "predictions"),
    "final_dir": os.path.join(PROJECT_ROOT, "data", "final"),
    "coordinate_file": "坐标.csv"
}

# 分析模块配置
ANALYSIS_CONFIG = {
    "template_tif": os.path.join(PROJECT_ROOT, "output", "raster", "meanTemp.tif"),
    "boundary_shp": os.path.join(PROJECT_ROOT, "output", "raster", "lechang.shp"),
    "output_raster": os.path.join(PROJECT_ROOT, "output", "raster", "output.tif"),
    "temp_shapefile": os.path.join(PROJECT_ROOT, "output", "raster", "points666.shp")
}

# 输出路径配置
OUTPUT_PATHS = {
    "raster_dir": os.path.join(PROJECT_ROOT, "output", "raster"),
    "figures_dir": os.path.join(PROJECT_ROOT, "output", "figures"),
    "reports_dir": os.path.join(PROJECT_ROOT, "output", "reports")
}

# 可视化配置
VISUALIZATION_CONFIG = {
    "color_maps": {
        "colormap1": ['#FFFECE', '#FFF085', '#FEBA17','#BE3D2A','#74512D', '#4E1F00'],
        "colormap2": ['#F6F8D5', '#98D2C0', '#4F959D', '#205781', '#143D60','#2A3335'],
        "colormap3": ['#FFEFC8', '#F8ED8C', '#D3E671', '#89AC46', '#5F8B4C', '#355F2E'],
        "colormap4": ['#F0F1C5', '#BBD8A3', '#6F826A', '#BF9264', '#735557', '#604652'],
        "colormap5": ['#FCFAEE', '#FBF3B9', '#FFDCCC', '#FDB7EA', '#B7B1F2', '#8D77AB'],
        "colormap6": ['#15B392', '#73EC8B', '#FFEB55', '#EE66A6', '#D91656', '#640D5F']
    },
    "default_colormap": "colormap6",
    "figure_size": 12,
    "dpi": 600
}

# 执行流程配置
_DEFAULT_WORKFLOW_CONFIG = {'run_crop_model': True, 'run_effective_model': False, 'combine_predictions': True, 'generate_raster': True, 'create_visualization': True, 'create_histogram': True}

def get_workflow_config():
    """
    获取工作流配置
    优先从环境变量获取，如果没有则使用默认配置
    
    @returns {dict} 工作流配置字典
    """
    import json
    env_config = os.environ.get('CD_WORKFLOW_CONFIG')
    if env_config:
        try:
            return json.loads(env_config)
        except json.JSONDecodeError:
            pass
    return _DEFAULT_WORKFLOW_CONFIG.copy()

# 为了向后兼容，保留WORKFLOW_CONFIG变量（使用默认配置）
WORKFLOW_CONFIG = _DEFAULT_WORKFLOW_CONFIG.copy()

def ensure_directories():
    """确保所有必要的目录存在"""
    directories = [
        MODELS_DIR,
        DATA_PATHS["coordinates_dir"],
        DATA_PATHS["predictions_dir"], 
        DATA_PATHS["final_dir"],
        OUTPUT_PATHS["raster_dir"],
        OUTPUT_PATHS["figures_dir"],
        OUTPUT_PATHS["reports_dir"]
    ]
    
    for directory in directories:
        os.makedirs(directory, exist_ok=True)