call_model.py

# 对应 ccx_process_caller_call_models_once.txt
from plasma_geometry import plasma_geom
from plasma_physics import PlasmaPhysics
from build import BuildModule
from TFcoil import ResisTFCoil
import power_component
import numpy as np


# 内外迭代耦合量：相邻两次 call_model_once 后应基本不变才算收敛
_COUPLED_SNAPSHOT_FIELDS = (
    ("comp_data", "f_density_proton_electron"),
    ("power_data", "alpha_produce_rate_total"),
    ("power_data", "p_produce_rate"),
)


def _coupled_state(model):
    """抽取耦合相关标量，用于判断模型是否基本不变。"""
    vals = []
    for container, field in _COUPLED_SNAPSHOT_FIELDS:
        v = getattr(getattr(model, container), field, None)
        vals.append(np.nan if v is None else float(v))
    return np.asarray(vals, dtype=float)


def _model_basically_unchanged(prev, curr, rtol=1e-6):
    """相邻两次物理更新后，耦合量是否在相对容差内一致。"""
    if prev is None:
        return False
    return bool(np.allclose(prev, curr, rtol=rtol, equal_nan=True))


def call_model_once(model):
    # calculate geometry parameters
    # input: i_plasma_geometry, LCFS
    plasma_geom(model.geom_data, model.input_data.i_plasma_geometry)

    # build
    # build_module = BuildModule()
    # build_module.build(model)

    # calculate physics parameters
    physics_module = PlasmaPhysics()
    physics_module.plasma_physics(model)
    # plasma_power_and_energy(model)
    # plasma_limit(model)
    # divertor_target(model)
    # power_component(model)

    # tf coil
    # tf_coil_module = ResisTFCoil()
    # tf_coil_module.run(model)

def call_model(model, max_iter=10, rtol=1e-6, iter_idx=None):
    """多次调用直至耦合量基本不变（幂等），最多 max_iter 次。

    Parameters
    ----------
    iter_idx :
        外层迭代编号（如约束求解第几次物理评估）。为 None 时只打印内层调用次数。
    """
    prev = None
    for j in range(1, max_iter + 1):
        if iter_idx is None:
            print(f"第 {j} 次调用模型")
        else:
            print(f"第 {iter_idx} 次迭代的第 {j} 次调用模型")
        call_model_once(model)
        curr = _coupled_state(model)
        if _model_basically_unchanged(prev, curr, rtol=rtol):
            break
        prev = curr


## 注意：
# 换一个思路，通过轴电子密度和温度计算其他值，如果其他值有输入，将输入作为约束条件