Upload 6 files

dataset/all_test_pred2.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf6c0136608c8ef425d40ca4031f0af7b160238d9c4f2b5676ebbaf7be21391e
+size 74641224

dataset/combined_test_special.txt

@@ -0,0 +1,240 @@
+min5_m_v1_0_d270_sc2_s15_04123.npz
+min5_m_v1_0_d270_sc2_s21_04130.npz
+min5_m_v1_0_d270_sc2_s52_04164.npz
+min5_m_v1_0_d270_sc2_s61_04174.npz
+min5_m_v1_0_d270_sc2_s72_04186.npz
+min5_m_v1_0_d270_sc2_s83_04198.npz
+min5_m_v1_0_d270_sc2_s87_04202.npz
+min5_m_v1_0_d270_sc2_s93_04209.npz
+min5_m_v1_0_d270_sc4_s15_04221.npz
+min5_m_v1_0_d270_sc4_s21_04228.npz
+min5_m_v1_0_d270_sc4_s52_04262.npz
+min5_m_v1_0_d270_sc4_s61_04272.npz
+min5_m_v1_0_d270_sc4_s72_04284.npz
+min5_m_v1_0_d270_sc4_s83_04296.npz
+min5_m_v1_0_d270_sc4_s87_04300.npz
+min5_m_v1_0_d270_sc4_s93_04307.npz
+min5_m_v1_0_d270_sc6_s15_04319.npz
+min5_m_v1_0_d270_sc6_s21_04326.npz
+min5_m_v1_0_d270_sc6_s52_04360.npz
+min5_m_v1_0_d270_sc6_s61_04370.npz
+min5_m_v1_0_d270_sc6_s72_04382.npz
+min5_m_v1_0_d270_sc6_s83_04394.npz
+min5_m_v1_0_d270_sc6_s87_04398.npz
+min5_m_v1_0_d270_sc6_s93_04405.npz
+min5_m_v1_0_d90_sc2_s15_04417.npz
+min5_m_v1_0_d90_sc2_s21_04424.npz
+min5_m_v1_0_d90_sc2_s52_04458.npz
+min5_m_v1_0_d90_sc2_s61_04468.npz
+min5_m_v1_0_d90_sc2_s72_04480.npz
+min5_m_v1_0_d90_sc2_s83_04492.npz
+min5_m_v1_0_d90_sc2_s87_04496.npz
+min5_m_v1_0_d90_sc2_s93_04503.npz
+min5_m_v1_0_d90_sc4_s15_04515.npz
+min5_m_v1_0_d90_sc4_s21_04522.npz
+min5_m_v1_0_d90_sc4_s52_04556.npz
+min5_m_v1_0_d90_sc4_s61_04566.npz
+min5_m_v1_0_d90_sc4_s72_04578.npz
+min5_m_v1_0_d90_sc4_s83_04590.npz
+min5_m_v1_0_d90_sc4_s87_04594.npz
+min5_m_v1_0_d90_sc4_s93_04601.npz
+min5_m_v1_0_d90_sc6_s15_04613.npz
+min5_m_v1_0_d90_sc6_s21_04620.npz
+min5_m_v1_0_d90_sc6_s52_04654.npz
+min5_m_v1_0_d90_sc6_s61_04664.npz
+min5_m_v1_0_d90_sc6_s72_04676.npz
+min5_m_v1_0_d90_sc6_s83_04688.npz
+min5_m_v1_0_d90_sc6_s87_04692.npz
+min5_m_v1_0_d90_sc6_s93_04699.npz
+min5_m_v2_0_d270_sc2_s15_08827.npz
+min5_m_v2_0_d270_sc2_s21_08834.npz
+min5_m_v2_0_d270_sc2_s52_08868.npz
+min5_m_v2_0_d270_sc2_s61_08878.npz
+min5_m_v2_0_d270_sc2_s72_08890.npz
+min5_m_v2_0_d270_sc2_s83_08902.npz
+min5_m_v2_0_d270_sc2_s87_08906.npz
+min5_m_v2_0_d270_sc2_s93_08913.npz
+min5_m_v2_0_d270_sc4_s15_08925.npz
+min5_m_v2_0_d270_sc4_s21_08932.npz
+min5_m_v2_0_d270_sc4_s52_08966.npz
+min5_m_v2_0_d270_sc4_s61_08976.npz
+min5_m_v2_0_d270_sc4_s72_08988.npz
+min5_m_v2_0_d270_sc4_s83_09000.npz
+min5_m_v2_0_d270_sc4_s87_09004.npz
+min5_m_v2_0_d270_sc4_s93_09011.npz
+min5_m_v2_0_d270_sc6_s15_09023.npz
+min5_m_v2_0_d270_sc6_s21_09030.npz
+min5_m_v2_0_d270_sc6_s52_09064.npz
+min5_m_v2_0_d270_sc6_s61_09074.npz
+min5_m_v2_0_d270_sc6_s72_09086.npz
+min5_m_v2_0_d270_sc6_s83_09098.npz
+min5_m_v2_0_d270_sc6_s87_09102.npz
+min5_m_v2_0_d270_sc6_s93_09109.npz
+min5_m_v2_0_d90_sc2_s15_09121.npz
+min5_m_v2_0_d90_sc2_s21_09128.npz
+min5_m_v2_0_d90_sc2_s52_09162.npz
+min5_m_v2_0_d90_sc2_s61_09172.npz
+min5_m_v2_0_d90_sc2_s72_09184.npz
+min5_m_v2_0_d90_sc2_s83_09196.npz
+min5_m_v2_0_d90_sc2_s87_09200.npz
+min5_m_v2_0_d90_sc2_s93_09207.npz
+min5_m_v2_0_d90_sc4_s15_09219.npz
+min5_m_v2_0_d90_sc4_s21_09226.npz
+min5_m_v2_0_d90_sc4_s52_09260.npz
+min5_m_v2_0_d90_sc4_s61_09270.npz
+min5_m_v2_0_d90_sc4_s72_09282.npz
+min5_m_v2_0_d90_sc4_s83_09294.npz
+min5_m_v2_0_d90_sc4_s87_09298.npz
+min5_m_v2_0_d90_sc4_s93_09305.npz
+min5_m_v2_0_d90_sc6_s15_09317.npz
+min5_m_v2_0_d90_sc6_s21_09324.npz
+min5_m_v2_0_d90_sc6_s52_09358.npz
+min5_m_v2_0_d90_sc6_s61_09368.npz
+min5_m_v2_0_d90_sc6_s72_09380.npz
+min5_m_v2_0_d90_sc6_s83_09392.npz
+min5_m_v2_0_d90_sc6_s87_09396.npz
+min5_m_v2_0_d90_sc6_s93_09403.npz
+min5_m_v3_0_d270_sc2_s15_13531.npz
+min5_m_v3_0_d270_sc2_s21_13538.npz
+min5_m_v3_0_d270_sc2_s52_13572.npz
+min5_m_v3_0_d270_sc2_s61_13582.npz
+min5_m_v3_0_d270_sc2_s72_13594.npz
+min5_m_v3_0_d270_sc2_s83_13606.npz
+min5_m_v3_0_d270_sc2_s87_13610.npz
+min5_m_v3_0_d270_sc2_s93_13617.npz
+min5_m_v3_0_d270_sc4_s15_13629.npz
+min5_m_v3_0_d270_sc4_s21_13636.npz
+min5_m_v3_0_d270_sc4_s52_13670.npz
+min5_m_v3_0_d270_sc4_s61_13680.npz
+min5_m_v3_0_d270_sc4_s72_13692.npz
+min5_m_v3_0_d270_sc4_s83_13704.npz
+min5_m_v3_0_d270_sc4_s87_13708.npz
+min5_m_v3_0_d270_sc4_s93_13715.npz
+min5_m_v3_0_d270_sc6_s15_13727.npz
+min5_m_v3_0_d270_sc6_s21_13734.npz
+min5_m_v3_0_d270_sc6_s52_13768.npz
+min5_m_v3_0_d270_sc6_s61_13778.npz
+min5_m_v3_0_d270_sc6_s72_13790.npz
+min5_m_v3_0_d270_sc6_s83_13802.npz
+min5_m_v3_0_d270_sc6_s87_13806.npz
+min5_m_v3_0_d270_sc6_s93_13813.npz
+min5_m_v3_0_d90_sc2_s15_13825.npz
+min5_m_v3_0_d90_sc2_s21_13832.npz
+min5_m_v3_0_d90_sc2_s52_13866.npz
+min5_m_v3_0_d90_sc2_s61_13876.npz
+min5_m_v3_0_d90_sc2_s72_13888.npz
+min5_m_v3_0_d90_sc2_s83_13900.npz
+min5_m_v3_0_d90_sc2_s87_13904.npz
+min5_m_v3_0_d90_sc2_s93_13911.npz
+min5_m_v3_0_d90_sc4_s15_13923.npz
+min5_m_v3_0_d90_sc4_s21_13930.npz
+min5_m_v3_0_d90_sc4_s52_13964.npz
+min5_m_v3_0_d90_sc4_s61_13974.npz
+min5_m_v3_0_d90_sc4_s72_13986.npz
+min5_m_v3_0_d90_sc4_s83_13998.npz
+min5_m_v3_0_d90_sc4_s87_14002.npz
+min5_m_v3_0_d90_sc4_s93_14009.npz
+min5_m_v3_0_d90_sc6_s15_14021.npz
+min5_m_v3_0_d90_sc6_s21_14028.npz
+min5_m_v3_0_d90_sc6_s52_14062.npz
+min5_m_v3_0_d90_sc6_s61_14072.npz
+min5_m_v3_0_d90_sc6_s72_14084.npz
+min5_m_v3_0_d90_sc6_s83_14096.npz
+min5_m_v3_0_d90_sc6_s87_14100.npz
+min5_m_v3_0_d90_sc6_s93_14107.npz
+min5_m_v4_0_d270_sc2_s15_18235.npz
+min5_m_v4_0_d270_sc2_s21_18242.npz
+min5_m_v4_0_d270_sc2_s52_18276.npz
+min5_m_v4_0_d270_sc2_s61_18286.npz
+min5_m_v4_0_d270_sc2_s72_18298.npz
+min5_m_v4_0_d270_sc2_s83_18310.npz
+min5_m_v4_0_d270_sc2_s87_18314.npz
+min5_m_v4_0_d270_sc2_s93_18321.npz
+min5_m_v4_0_d270_sc4_s15_18333.npz
+min5_m_v4_0_d270_sc4_s21_18340.npz
+min5_m_v4_0_d270_sc4_s52_18374.npz
+min5_m_v4_0_d270_sc4_s61_18384.npz
+min5_m_v4_0_d270_sc4_s72_18396.npz
+min5_m_v4_0_d270_sc4_s83_18408.npz
+min5_m_v4_0_d270_sc4_s87_18412.npz
+min5_m_v4_0_d270_sc4_s93_18419.npz
+min5_m_v4_0_d270_sc6_s15_18431.npz
+min5_m_v4_0_d270_sc6_s21_18438.npz
+min5_m_v4_0_d270_sc6_s52_18472.npz
+min5_m_v4_0_d270_sc6_s61_18482.npz
+min5_m_v4_0_d270_sc6_s72_18494.npz
+min5_m_v4_0_d270_sc6_s83_18506.npz
+min5_m_v4_0_d270_sc6_s87_18510.npz
+min5_m_v4_0_d270_sc6_s93_18517.npz
+min5_m_v4_0_d90_sc2_s15_18529.npz
+min5_m_v4_0_d90_sc2_s21_18536.npz
+min5_m_v4_0_d90_sc2_s52_18570.npz
+min5_m_v4_0_d90_sc2_s61_18580.npz
+min5_m_v4_0_d90_sc2_s72_18592.npz
+min5_m_v4_0_d90_sc2_s83_18604.npz
+min5_m_v4_0_d90_sc2_s87_18608.npz
+min5_m_v4_0_d90_sc2_s93_18615.npz
+min5_m_v4_0_d90_sc4_s15_18627.npz
+min5_m_v4_0_d90_sc4_s21_18634.npz
+min5_m_v4_0_d90_sc4_s52_18668.npz
+min5_m_v4_0_d90_sc4_s61_18678.npz
+min5_m_v4_0_d90_sc4_s72_18690.npz
+min5_m_v4_0_d90_sc4_s83_18702.npz
+min5_m_v4_0_d90_sc4_s87_18706.npz
+min5_m_v4_0_d90_sc4_s93_18713.npz
+min5_m_v4_0_d90_sc6_s15_18725.npz
+min5_m_v4_0_d90_sc6_s21_18732.npz
+min5_m_v4_0_d90_sc6_s52_18766.npz
+min5_m_v4_0_d90_sc6_s61_18776.npz
+min5_m_v4_0_d90_sc6_s72_18788.npz
+min5_m_v4_0_d90_sc6_s83_18800.npz
+min5_m_v4_0_d90_sc6_s87_18804.npz
+min5_m_v4_0_d90_sc6_s93_18811.npz
+min5_m_v5_0_d270_sc2_s15_22939.npz
+min5_m_v5_0_d270_sc2_s21_22946.npz
+min5_m_v5_0_d270_sc2_s52_22980.npz
+min5_m_v5_0_d270_sc2_s61_22990.npz
+min5_m_v5_0_d270_sc2_s72_23002.npz
+min5_m_v5_0_d270_sc2_s83_23014.npz
+min5_m_v5_0_d270_sc2_s87_23018.npz
+min5_m_v5_0_d270_sc2_s93_23025.npz
+min5_m_v5_0_d270_sc4_s15_23037.npz
+min5_m_v5_0_d270_sc4_s21_23044.npz
+min5_m_v5_0_d270_sc4_s52_23078.npz
+min5_m_v5_0_d270_sc4_s61_23088.npz
+min5_m_v5_0_d270_sc4_s72_23100.npz
+min5_m_v5_0_d270_sc4_s83_23112.npz
+min5_m_v5_0_d270_sc4_s87_23116.npz
+min5_m_v5_0_d270_sc4_s93_23123.npz
+min5_m_v5_0_d270_sc6_s15_23135.npz
+min5_m_v5_0_d270_sc6_s21_23142.npz
+min5_m_v5_0_d270_sc6_s52_23176.npz
+min5_m_v5_0_d270_sc6_s61_23186.npz
+min5_m_v5_0_d270_sc6_s72_23198.npz
+min5_m_v5_0_d270_sc6_s83_23210.npz
+min5_m_v5_0_d270_sc6_s87_23214.npz
+min5_m_v5_0_d270_sc6_s93_23221.npz
+min5_m_v5_0_d90_sc2_s15_23233.npz
+min5_m_v5_0_d90_sc2_s21_23240.npz
+min5_m_v5_0_d90_sc2_s52_23274.npz
+min5_m_v5_0_d90_sc2_s61_23284.npz
+min5_m_v5_0_d90_sc2_s72_23296.npz
+min5_m_v5_0_d90_sc2_s83_23308.npz
+min5_m_v5_0_d90_sc2_s87_23312.npz
+min5_m_v5_0_d90_sc2_s93_23319.npz
+min5_m_v5_0_d90_sc4_s15_23331.npz
+min5_m_v5_0_d90_sc4_s21_23338.npz
+min5_m_v5_0_d90_sc4_s52_23372.npz
+min5_m_v5_0_d90_sc4_s61_23382.npz
+min5_m_v5_0_d90_sc4_s72_23394.npz
+min5_m_v5_0_d90_sc4_s83_23406.npz
+min5_m_v5_0_d90_sc4_s87_23410.npz
+min5_m_v5_0_d90_sc4_s93_23417.npz
+min5_m_v5_0_d90_sc6_s15_23429.npz
+min5_m_v5_0_d90_sc6_s21_23436.npz
+min5_m_v5_0_d90_sc6_s52_23470.npz
+min5_m_v5_0_d90_sc6_s61_23480.npz
+min5_m_v5_0_d90_sc6_s72_23492.npz
+min5_m_v5_0_d90_sc6_s83_23504.npz
+min5_m_v5_0_d90_sc6_s87_23508.npz
+min5_m_v5_0_d90_sc6_s93_23515.npz

dataset/conditioned_results_v0_5_d45_n40.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b1ad22eb1731f16f69c468666acda2df38d65f9c9b3c19804c6f5a5f50c47ed
+size 2516841027

dataset/min5_m_v1_0_d270_sc2_s10_04118.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21ef8d148671d73c06742cfcf61ffc48d4590fbc70f55f91edab628252160705
+size 2523009

function.py

@@ -0,0 +1,1945 @@
+import pandas as pd
+import numpy as np
+import os
+from scipy import stats
+from scipy.spatial import cKDTree
+from scipy.ndimage import binary_dilation
+from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error
+from skimage.metrics import structural_similarity as ssim
+from tqdm import trange
+
+import matplotlib.pyplot as plt
+import matplotlib.ticker as ticker
+import matplotlib.cm as cm
+import matplotlib
+matplotlib.rcParams['font.sans-serif'] = ['Arial']
+matplotlib.rcParams['font.size'] = 16
+
+
+
+class DataLoader:
+    '''
+    Fucntions: 
+    1. load_predictions: 从 npz 文件加载预测和真实浓度场
+    2. load_metadata: 从 meta txt 文件加载元信息（风速、风向、稳定度、源编号等）
+    3. load_conds_data: 从 pkl 文件加载条件预测数据（如果有）
+    4. log2ppm: 将 log 浓度转换为 ppm 浓度（根据给定的关系）
+    5. get_sample: 根据索引获取单个样本的预测场、真实场、条件预测和元信息
+    '''
+    def __init__(self, pred_npz_path, meta_txt_path, conds_pkl_path):
+        self.pred_npz_path = pred_npz_path
+        self.meta_txt_path = meta_txt_path
+        self.conds_pkl_path = conds_pkl_path
+
+        self.load_predictions()
+        self.meta = self.load_metadata()
+        self.conds_data = self.load_conds_data()
+
+    def load_predictions(self):
+        data = np.load(self.pred_npz_path)
+        self.preds = data['preds'].squeeze(1)
+        self.trues = data['trues'].squeeze(1)
+        _, non_building_mask = PrintMetrics.get_building_area()
+        self.preds = self.preds * non_building_mask[None, :, :]
+        self.trues = self.trues * non_building_mask[None, :, :]
+        return self.preds, self.trues
+
+    def load_metadata(self):
+        df = pd.read_csv(self.meta_txt_path, sep=',', header=None)
+        df.columns = ['npz_colname']
+        pattern = r'v([0-9_]+)_d(\d+)_sc(\d+)_s(\d+)'
+        df[['wind_speed', 'wind_direction', 'sc', 'source_number']] = (
+            df['npz_colname'].str.extract(pattern))
+        df['wind_speed'] = df['wind_speed'].str.replace('_', '.').astype(float)
+        df[['wind_direction', 'sc', 'source_number']] = df[['wind_direction', 'sc',
+                                                            'source_number']].astype(int)
+        return df
+    
+    def load_conds_data(self):
+        conds_data = np.load(self.conds_pkl_path, allow_pickle=True)
+        return conds_data
+    
+    @staticmethod
+    def log2ppm(log_conc):
+        log_conc = np.asarray(log_conc)
+        log_conc = np.minimum(log_conc, 15.0)
+        ppm_conc = np.expm1(log_conc) * (0.7449)
+        return np.maximum(ppm_conc, 0.0)
+    
+    def get_sample(self, idx, in_ppm=True):
+        psi_f = self.preds[idx]
+        psi_t = self.trues[idx]
+        meta = self.meta.iloc[idx]
+        conds_preds = self.conds_data[idx]['conds']['preds']
+        if in_ppm:
+            psi_f = DataLoader.log2ppm(psi_f)
+            psi_t = DataLoader.log2ppm(psi_t)
+            conds_preds = DataLoader.log2ppm(conds_preds)
+        return psi_f, psi_t, conds_preds, meta
+    
+class ObservationModel:
+    '''
+    Functions:
+    1. observation_operator_H: 从浓度场 ψ 中提取点位浓度，使用双线性插值（线性算子）
+    2. observation_operator_H_ens: 对 ensemble 预测场批量应用观测算子，得到每个成员的点位浓度
+    '''
+    @staticmethod # 不依赖实例状态，可以直接通过类调用
+    def observation_operator_H(psi, obs_xy):
+        # 观测算子 M：
+        # 从浓度场 ψ 中提取点位浓度
+        # 使用双线性插值（线性算子）
+        Hh, Ww = psi.shape
+        xs = np.clip(obs_xy[:, 0], 0, Ww - 1 - 1e-6)
+        ys = np.clip(obs_xy[:, 1], 0, Hh - 1 - 1e-6)
+        x0 = np.floor(xs).astype(int)
+        y0 = np.floor(ys).astype(int)
+        x1 = np.clip(x0 + 1, 0, Ww - 1)
+        y1 = np.clip(y0 + 1, 0, Hh - 1)
+        dx = xs - x0
+        dy = ys - y0
+        f00 = psi[y0, x0]
+        f10 = psi[y0, x1]
+        f01 = psi[y1, x0]
+        f11 = psi[y1, x1]
+
+        return (
+            f00 * (1 - dx) * (1 - dy) + 
+            f10 * dx * (1 - dy) + 
+            f01 * (1 - dx) * dy + 
+            f11 * dx * dy
+        )
+    
+    @staticmethod
+    def observation_operator_H_ens(psi_ens, obs_xy):
+        """
+        psi_ens: (N_ens, H, W)
+        obs_xy : (n_obs, 2)
+        return : HX (N_ens, n_obs)
+        """
+        N_ens, Hh, Ww = psi_ens.shape
+        xs = np.clip(obs_xy[:, 0], 0, Ww - 1 - 1e-6)
+        ys = np.clip(obs_xy[:, 1], 0, Hh - 1 - 1e-6)
+        x0 = np.floor(xs).astype(np.int64)
+        y0 = np.floor(ys).astype(np.int64)
+        x1 = np.clip(x0 + 1, 0, Ww - 1)
+        y1 = np.clip(y0 + 1, 0, Hh - 1)
+        dx = xs - x0
+        dy = ys - y0
+        f00 = psi_ens[:, y0, x0]
+        f10 = psi_ens[:, y0, x1]
+        f01 = psi_ens[:, y1, x0]
+        f11 = psi_ens[:, y1, x1]
+
+        HX = (
+            f00 * (1 - dx) * (1 - dy) + 
+            f10 * dx * (1 - dy) + 
+            f01 * (1 - dx) * dy + 
+            f11 * dx * dy
+        )
+        return HX
+    
+class SamplingStrategies:
+    
+    # =========================
+    # (1) Sampling strategies
+    # =========================
+    @staticmethod
+    def sample_random(field_shape, num_points, seed=42):
+        rng = np.random.default_rng(seed)
+        H, W = field_shape
+        _, non_building_mask = PrintMetrics.get_building_area()
+        valid_idx = np.where(non_building_mask.ravel())[0]
+        chosen = rng.choice(valid_idx, size=num_points, replace=False)
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        coords = np.stack([xx.ravel(), yy.ravel()], axis=1)
+        return coords[chosen].astype(float)
+    
+    @staticmethod
+    def sample_uniform(field_shape, num_points, margin=20):
+        H, W = field_shape
+        nx = int(np.ceil(np.sqrt(num_points * W / H)))
+        ny = int(np.ceil(num_points / nx))
+        xs = np.linspace(margin, W - 1 - margin, nx)
+        ys = np.linspace(margin, H - 1 - margin, ny)
+        xx, yy = np.meshgrid(xs, ys)
+        grid_xy = np.stack([xx.ravel(), yy.ravel()], axis=1)
+        _, non_building_mask = PrintMetrics.get_building_area()
+        xi = np.clip(grid_xy[:, 0].astype(int), 0, W - 1)
+        yi = np.clip(grid_xy[:, 1].astype(int), 0, H - 1)
+        valid = non_building_mask[yi, xi] == 1
+        grid_xy = grid_xy[valid]
+        if len(grid_xy) > num_points:
+            idx = np.linspace(0, len(grid_xy) - 1, num_points).astype(int)
+            grid_xy = grid_xy[idx]
+        return grid_xy
+    
+    @staticmethod
+    def two_stage_sampling(
+            true_field,
+            pred_field,
+            num_points,
+            ens_preds_ppm=None,
+            seed=42,
+
+            # ====== 全局控制 ======
+            min_dist=22,
+            n1_ratio=0.65,            # Stage1 比例
+
+            # ====== Stage1 可调参数 =====
+            stage1_grad_power=0.8,    # 梯度权重幂次
+            stage1_value_power=1.2,   # 值权重幂次
+            stage1_center_boost=1.2,  # 是否增强高值区域
+    ):
+        
+        # 内部函数: 基于排斥采样的加权随机选择
+        def repulse_pick(candidate_idx, weights, k, selected_idx):
+            if k <= 0 or len(candidate_idx) == 0:
+                return list(selected_idx)
+            candidate_idx = np.asarray(candidate_idx, dtype=np.int64)
+            weights = np.maximum(np.asarray(weights, dtype=float), 0.0)
+            if weights.sum() <= 0:
+                weights = np.ones_like(weights)
+            weights = weights / weights.sum()
+            overs = min(len(candidate_idx), max(k * 15, 200))
+            cand = rng.choice(candidate_idx, size=overs, replace=False, p=weights)
+            selected = list(selected_idx)
+            for idx in cand:
+                xy = coords[idx]
+                if not selected:
+                    selected.append(idx)
+                    continue
+                sel_xy = coords[np.asarray(selected)]
+                if cKDTree(sel_xy).query(xy, k=1)[0] >= min_dist:
+                    selected.append(idx)
+                if len(selected) >= k + len(selected_idx):
+                    break
+            return selected
+
+        rng = np.random.default_rng(seed)
+        H, W = pred_field.shape
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        coords = np.stack([xx.ravel(), yy.ravel()], axis=1) # 二维坐标网格 (HW, 2) [x,y]
+        v = np.maximum(pred_field, 0.0).ravel()
+        vmax = float(v.max())
+        _, non_building_mask = PrintMetrics.get_building_area()
+        non_building_flat = non_building_mask.ravel().astype(bool)
+
+        if vmax <= 1e-6:
+            valid_idx = np.where(non_building_flat)[0]
+            idx = rng.choice(valid_idx, size=num_points, replace=False)
+            obs_xy = coords[idx]
+            obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+            return obs_xy, obs_val
+
+        n_center_ratio= 1 - n1_ratio
+        n_center = max(2, int(num_points * n_center_ratio))
+        n1 = num_points - n_center
+
+        # 1. 结构修正（LOG）
+        z = np.log1p(np.maximum(pred_field, 0.0))
+        z_flat = z.ravel()
+        gx, gy = np.gradient(z)
+        grad = np.sqrt(gx**2 + gy**2).ravel()
+        nz = z_flat > 1e-6 # 只在非零区域上算分位数，避免大量0把lo/hi压塌
+        z_nz = z_flat[nz]
+        if z_nz.size < 50:
+            support_mask = (grad.reshape(H, W) > np.quantile(grad, 0.90))
+        else:
+            lo = np.quantile(z_nz, 0.70)
+            hi = np.quantile(z_nz, 0.90)
+            core_mask = (z >= lo) & (z <= hi)
+            r_out = 12 # 外扩：把结构带膨胀一��，让采样不只盯着最强边界
+            support_mask = binary_dilation(core_mask, iterations=r_out)
+
+        support_idx = np.where(support_mask.ravel() & non_building_flat)[0]
+        if len(support_idx) < num_points * 2:
+            support_idx = np.where(non_building_flat)[0]
+
+        # Stage1：梯度主导 + 适度保留外圈
+        weights1 = (
+            (grad[support_idx] ** stage1_grad_power) *
+            (z_flat[support_idx] ** stage1_value_power + 1e-6)
+        )
+        if stage1_center_boost > 1.0:
+            weights1 *= (1 + stage1_center_boost * (z_flat[support_idx] / z_flat.max()))
+        selected = repulse_pick(support_idx, weights1, n1, [])
+
+        # Stage 2: 中心峰值区，只取2-3个点
+        peak_idx = np.argmax(z_flat * non_building_flat.astype(float))
+        peak_xy = coords[peak_idx]
+        # print(f"峰值位置: {peak_xy}, z值: {z_flat[peak_idx]:.3f}")
+        selected.append(int(peak_idx)) # 直接把峰值点加进去（1个）
+
+        # 再在峰值极近邻选1-2个，min_dist放松到5保证不重叠
+        if n_center > 1:
+            peak_radius = 10  # 很小的半径，只捕捉最高点附近
+            stage2_idx = np.where(
+                non_building_flat &
+                (np.sqrt((coords[:, 0] - peak_xy[0])**2 +
+                         (coords[:, 1] - peak_xy[1])**2) <= peak_radius)
+            )[0]
+            stage2_idx = np.setdiff1d(stage2_idx, np.array(selected))
+
+            if len(stage2_idx) >= 1:
+                weights2 = z_flat[stage2_idx]
+                weights2 = weights2 / (weights2.sum() + 1e-12)
+                overs = min(len(stage2_idx), max((n_center - 1) * 10, 20))
+                cands = rng.choice(stage2_idx, size=overs, replace=False, p=weights2)
+                for idx in cands:
+                    xy = coords[idx]
+                    if cKDTree(coords[np.array(selected)]).query(xy, k=1)[0] >= 5:
+                        selected.append(int(idx))
+                    if len(selected) >= n_center + len([]):  # 只加到n_center个为止
+                        break
+                    if len(selected) - (num_points - n_center) >= n_center:
+                        break
+        selected = list(dict.fromkeys(selected))
+
+        # 补足剩余点（从Stage1结构带里再补，如果selected不够num_points）
+        if len(selected) < num_points:
+            remain = np.setdiff1d(support_idx, np.array(selected))
+            if len(remain) > 0:
+                w_remain = (
+                    (grad[remain] ** stage1_grad_power) *
+                    (z_flat[remain] ** stage1_value_power + 1e-6)
+                )
+                extra = repulse_pick(remain, w_remain,
+                                     num_points - len(selected), selected)
+                selected = extra
+
+        obs_xy = coords[np.array(selected[:num_points])]
+        obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+
+        return obs_xy, obs_val
+    
+    @staticmethod
+    def two_stage_pro(
+            true_field,
+            pred_field,
+            num_points,
+            ens_preds_ppm=None,
+            seed=42,
+            min_dist=22,
+            n1_ratio=0.65,
+            stage1_grad_power=0.8,
+            stage1_value_power=1.2,
+            stage1_center_boost=1.2,
+    ):
+        import numpy as np
+        from scipy.spatial import cKDTree
+        from scipy.ndimage import binary_dilation
+
+        def repulse_pick(candidate_idx, weights, k, selected_idx, this_min_dist):
+            if k <= 0 or len(candidate_idx) == 0:
+                return list(selected_idx)
+
+            candidate_idx = np.asarray(candidate_idx, dtype=np.int64)
+            weights = np.maximum(np.asarray(weights, dtype=float), 0.0)
+
+            if weights.sum() <= 0:
+                weights = np.ones_like(weights, dtype=float)
+
+            weights = weights / weights.sum()
+
+            overs = min(len(candidate_idx), max(k * 15, 200))
+            cand = rng.choice(candidate_idx, size=overs, replace=False, p=weights)
+
+            selected = list(selected_idx)
+            for idx in cand:
+                idx = int(idx)
+                if idx in selected:
+                    continue
+
+                xy = coords[idx]
+                if not selected:
+                    selected.append(idx)
+                    continue
+
+                sel_xy = coords[np.asarray(selected)]
+                if cKDTree(sel_xy).query(xy, k=1)[0] >= this_min_dist:
+                    selected.append(idx)
+
+                if len(selected) >= k + len(selected_idx):
+                    break
+
+            return selected
+
+        rng = np.random.default_rng(seed)
+        H, W = pred_field.shape
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        coords = np.stack([xx.ravel(), yy.ravel()], axis=1)   # (HW, 2), [x, y]
+
+        v = np.maximum(pred_field, 0.0).ravel()
+        vmax = float(v.max())
+
+        _, non_building_mask = PrintMetrics.get_building_area()
+        non_building_flat = non_building_mask.ravel().astype(bool)
+
+        if vmax <= 1e-6:
+            valid_idx = np.where(non_building_flat)[0]
+            idx = rng.choice(valid_idx, size=min(num_points, len(valid_idx)), replace=False)
+
+            if len(idx) < num_points:
+                raise ValueError(f"Not enough valid non-building points: need {num_points}, got {len(idx)}")
+
+            obs_xy = coords[idx]
+            obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+            return obs_xy, obs_val
+
+        n_center_ratio = 1 - n1_ratio
+        n_center = max(2, int(num_points * n_center_ratio))
+        n1 = num_points - n_center
+
+        z = np.log1p(np.maximum(pred_field, 0.0))
+        z_flat = z.ravel()
+        gx, gy = np.gradient(z)
+        grad = np.sqrt(gx**2 + gy**2).ravel()
+        nz = z_flat > 1e-6
+        z_nz = z_flat[nz]
+        if z_nz.size < 50:
+            support_mask = (grad.reshape(H, W) > np.quantile(grad, 0.90))
+        else:
+            lo = np.quantile(z_nz, 0.70)
+            hi = np.quantile(z_nz, 0.90)
+            core_mask = (z >= lo) & (z <= hi)
+            # r_out = int(0.60 * num_points)
+            r_out = int(np.clip(num_points / 3 + 16 / 3, 12, 24)) # 根据 num_points 动态调整外扩半径，保持在12-24范围内
+            support_mask = binary_dilation(core_mask, iterations=r_out)
+
+        support_idx = np.where(support_mask.ravel() & non_building_flat)[0]
+        if len(support_idx) < num_points * 2:
+            support_idx = np.where(non_building_flat)[0]
+
+        weights1 = (
+            (grad[support_idx] ** stage1_grad_power) *
+            (z_flat[support_idx] ** stage1_value_power + 1e-6)
+        )
+
+        if stage1_center_boost > 1.0:
+            weights1 *= (1 + stage1_center_boost * (z_flat[support_idx] / (z_flat.max() + 1e-12)))
+        selected = repulse_pick(support_idx, weights1, n1, [], min_dist)
+
+        peak_idx = int(np.argmax(z_flat * non_building_flat.astype(float)))
+        peak_xy = coords[peak_idx]
+        selected.append(int(peak_idx))
+
+        if n_center > 1:
+            peak_radius = 10
+            stage2_idx = np.where(
+                non_building_flat &
+                (np.sqrt((coords[:, 0] - peak_xy[0]) ** 2 +
+                        (coords[:, 1] - peak_xy[1]) ** 2) <= peak_radius)
+            )[0]
+            stage2_idx = np.setdiff1d(stage2_idx, np.array(selected))
+
+            if len(stage2_idx) >= 1:
+                weights2 = z_flat[stage2_idx]
+                weights2 = weights2 / (weights2.sum() + 1e-12)
+
+                overs = min(len(stage2_idx), max((n_center - 1) * 10, 20))
+                cands = rng.choice(stage2_idx, size=overs, replace=False, p=weights2)
+
+                for idx in cands:
+                    idx = int(idx)
+                    xy = coords[idx]
+                    if cKDTree(coords[np.array(selected)]).query(xy, k=1)[0] >= 5:
+                        selected.append(idx)
+
+                    if len(selected) >= n_center + len([]):
+                        break
+                    if len(selected) - (num_points - n_center) >= n_center:
+                        break
+
+        selected = list(dict.fromkeys(selected))
+
+        if len(selected) >= num_points:
+            selected = selected[:num_points]
+            obs_xy = coords[np.array(selected)]
+            obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+            return obs_xy, obs_val
+
+        remain = np.setdiff1d(support_idx, np.array(selected))
+        if len(remain) > 0:
+            for d_try in [
+                min_dist,
+                max(1, int(min_dist * 0.7)),
+                max(1, int(min_dist * 0.4)),
+                5,
+                3
+            ]:
+                if len(selected) >= num_points:
+                    break
+
+                remain = np.setdiff1d(support_idx, np.array(selected))
+                if len(remain) == 0:
+                    break
+
+                w_remain = (
+                    (grad[remain] ** stage1_grad_power) *
+                    (z_flat[remain] ** stage1_value_power + 1e-6)
+                )
+
+                selected = repulse_pick(
+                    remain,
+                    w_remain,
+                    num_points - len(selected),
+                    selected,
+                    d_try
+                )
+
+        selected = list(dict.fromkeys(selected))
+
+        if len(selected) < num_points:
+            remain_support = np.setdiff1d(support_idx, np.array(selected))
+
+            if len(remain_support) > 0:
+                need = num_points - len(selected)
+                extra = rng.choice(
+                    remain_support,
+                    size=min(need, len(remain_support)),
+                    replace=False
+                )
+                selected.extend(extra.tolist())
+
+        selected = list(dict.fromkeys(selected))
+
+        if len(selected) < num_points:
+            all_valid = np.where(non_building_flat)[0]
+            remain_all = np.setdiff1d(all_valid, np.array(selected))
+
+            if len(remain_all) > 0:
+                need = num_points - len(selected)
+                extra = rng.choice(
+                    remain_all,
+                    size=min(need, len(remain_all)),
+                    replace=False
+                )
+                selected.extend(extra.tolist())
+
+        selected = list(dict.fromkeys(selected))
+        selected = selected[:num_points]
+        assert len(selected) == num_points, f"Expected {num_points} points, got {len(selected)}"
+
+        obs_xy = coords[np.array(selected)]
+        obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+
+        return obs_xy, obs_val
+
+    @staticmethod
+    def smart_two_pass(
+            enkf,
+            psi_f,
+            conds_preds,
+            true_field,
+            n1,
+            n2,
+            n_rounds=2,
+            phase1_method='two_stage',
+            min_dist_p2=22,
+            under_correct_alpha=1.5,
+            use_localization=False,
+            loc_radius_pixobs=35.0,
+            loc_radius_obsobs=40.0,
+            seed=42,
+            verbose=True,
+    ):
+        """
+        多轮迭代选点 + EnKF 同化。
+
+        每轮流程：
+          Phase 1 — 基于当前先验场选 n1 个点，做 pilot EnKF；
+          Phase 2 — 基于 pilot 残差找欠校正区，再选 n2 个点，做 final EnKF；
+          本轮分析场作为下一轮的先验（psi_f）。
+
+        参数:
+            enkf            : EnKF 实例
+            psi_f           : 初始先验场 (H, W)
+            conds_preds     : 集合预测场 (N_ens, H, W)，协方差来源，全程不变
+            true_field      : 真值场 (H, W)，仅用于观测值提取
+            n1              : 每轮 Phase-1 选点数
+            n2              : 每轮 Phase-2 选点数
+            n_rounds        : 迭代轮数（默认 1，即原始两阶段行为）
+            phase1_method   : Phase-1 采样策略（'two_stage' 或其他 generate 支持的方法）
+            min_dist_p2     : Phase-2 选点与已有点的最小距离（像素）
+            under_correct_alpha : Phase-2 欠校正权重幂次
+            use_localization: 是否使用局地化 EnKF
+            loc_radius_pixobs / loc_radius_obsobs : 局地化半径
+            seed            : 随机种子
+            verbose         : 是否打印中间日志
+
+        返回:
+            psi_a_final     : 最终分析场 (H, W)
+            all_obs_xy      : 所有轮次累计观测坐标 (n_rounds*(n1+n2), 2)
+            all_obs_val     : 所有轮次累计观测值
+            psi_pilot       : 最后一轮的 pilot（Phase-1）分析场
+            obs_xy_p1_last  : 最后一轮 Phase-1 选点坐标
+        """
+        conds_preds = np.asarray(conds_preds)
+        N_ens, H, W = conds_preds.shape
+        rng = np.random.default_rng(seed)
+
+        _, non_building_mask = PrintMetrics.get_building_area()
+        non_building_flat = non_building_mask.ravel().astype(bool)
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        coords = np.stack([xx.ravel(), yy.ravel()], axis=1)
+
+        # 预先构建集合（整个函数中只用这一个 X_f，协方差永远基于原始集合）
+        ens_mean = np.mean(conds_preds, axis=0)
+        X_f_base = conds_preds - ens_mean[None, :, :] + psi_f[None, :, :]
+
+        # 累计所有轮次的观测（跨轮次积累，最终一次性返回）
+        all_obs_xy_list = []
+        all_obs_val_list = []
+
+        # 当前先验：第 1 轮用 psi_f，后续轮次用上一轮的分析场
+        psi_current = psi_f
+        psi_pilot = None
+        obs_xy_p1_last = None
+
+        for round_idx in range(n_rounds):
+            round_seed = seed + round_idx  # 每轮不同种子，避免重复采样
+
+            # ── Phase 1：基于当前先验选点 + pilot EnKF ────────────────
+            if phase1_method == 'two_stage':
+                obs_xy_p1, obs_val_p1 = SamplingStrategies.two_stage_sampling(
+                    true_field=true_field, pred_field=psi_current,
+                    num_points=n1, seed=round_seed)
+            else:
+                obs_xy_p1, obs_val_p1 = SamplingStrategies.generate(
+                    true_field, psi_current, n1, method=phase1_method, seed=round_seed)
+
+            # pilot EnKF 使用当前先验重新中心化的集合
+            ens_mean_cur = np.mean(conds_preds, axis=0)
+            X_f_cur = conds_preds - ens_mean_cur[None, :, :] + psi_current[None, :, :]
+
+            if use_localization:
+                psi_pilot = enkf._enkf_update_localized(
+                    X_f_cur, obs_xy_p1, obs_val_p1,
+                    loc_radius_pixobs, loc_radius_obsobs, round_seed)
+            else:
+                psi_pilot = enkf._enkf_update_standard(X_f_cur, obs_xy_p1, obs_val_p1)
+
+            if verbose:
+                from sklearn.metrics import r2_score as _r2
+                print(f"[SmartEnKF] Round {round_idx+1}/{n_rounds} Phase1: "
+                      f"{n1} 点, pilot R²={_r2(true_field.ravel(), psi_pilot.ravel()):.4f}")
+
+            # ── Phase 2：找欠校正区，补充选点 ──────────────────────────
+            psi_f_flat = psi_current.ravel()
+            psi_pilot_flat = psi_pilot.ravel()
+            correction_map = np.abs(psi_pilot_flat - psi_f_flat)
+
+            nz_vals = psi_f_flat[non_building_flat & (psi_f_flat > 1e-4)]
+            prior_thresh = np.quantile(nz_vals, 0.20) if len(nz_vals) > 20 else 1e-4
+            plume_support = non_building_flat & (psi_f_flat > prior_thresh)
+            cand_idx = np.where(plume_support)[0]
+            if len(cand_idx) < n2 * 3:
+                cand_idx = np.where(non_building_flat & (psi_f_flat > 1e-6))[0]
+
+            prior_cand = psi_f_flat[cand_idx]
+            corr_cand = correction_map[cand_idx]
+            prior_norm = prior_cand / (prior_cand.max() + 1e-12)
+            corr_norm = corr_cand / (corr_cand.max() + 1e-12)
+            under_score = prior_norm * (1.0 - corr_norm + 0.05)
+
+            p1_tree = cKDTree(obs_xy_p1)
+            dist_p1, _ = p1_tree.query(coords[cand_idx], k=1)
+            dist_w = np.tanh(dist_p1 / (min_dist_p2 * 2.5))
+
+            weights_p2 = (under_score ** under_correct_alpha) * (dist_w + 0.05)
+            weights_p2 = np.maximum(weights_p2, 1e-12)
+            weights_p2 /= weights_p2.sum()
+
+            rng_round = np.random.default_rng(round_seed)
+            n_over = min(len(cand_idx), max(n2 * 30, 600))
+            cands = rng_round.choice(cand_idx, size=n_over, replace=False, p=weights_p2)
+
+            selected_p2 = []
+            for cidx in cands:
+                xy = coords[cidx]
+                if p1_tree.query(xy, k=1)[0] < min_dist_p2:
+                    continue
+                if selected_p2:
+                    if cKDTree(coords[np.array(selected_p2)]).query(xy, k=1)[0] < min_dist_p2:
+                        continue
+                selected_p2.append(int(cidx))
+                if len(selected_p2) >= n2:
+                    break
+
+            if len(selected_p2) < n2:
+                remain = np.setdiff1d(cand_idx, np.array(selected_p2, dtype=int))
+                extra = rng_round.choice(remain,
+                                         size=min(n2 - len(selected_p2), len(remain)),
+                                         replace=False)
+                selected_p2.extend(extra.tolist())
+
+            obs_xy_p2 = coords[np.array(selected_p2[:n2])]
+            obs_val_p2 = ObservationModel.observation_operator_H(true_field, obs_xy_p2)
+
+            if verbose:
+                print(f"[SmartEnKF] Round {round_idx+1}/{n_rounds} Phase2: 补充 {n2} 个欠校正区域点")
+
+            # ── Final：本轮全部点 + 当前先验做最终 EnKF ────────────────
+            round_obs_xy = np.vstack([obs_xy_p1, obs_xy_p2])
+            round_obs_val = np.concatenate([obs_val_p1, obs_val_p2])
+
+            if use_localization:
+                psi_a_round = enkf._enkf_update_localized(
+                    X_f_cur, round_obs_xy, round_obs_val,
+                    loc_radius_pixobs, loc_radius_obsobs, round_seed)
+            else:
+                psi_a_round = enkf._enkf_update_standard(X_f_cur, round_obs_xy, round_obs_val)
+
+            if verbose:
+                from sklearn.metrics import r2_score as _r2
+                print(f"[SmartEnKF] Round {round_idx+1}/{n_rounds} Final: "
+                      f"{n1+n2} 点, R²={_r2(true_field.ravel(), psi_a_round.ravel()):.4f}")
+
+            # 累计观测，更新先验进入下一轮
+            all_obs_xy_list.append(round_obs_xy)
+            all_obs_val_list.append(round_obs_val)
+            psi_current = np.maximum(psi_a_round, 0.0)
+            obs_xy_p1_last = obs_xy_p1
+
+        all_obs_xy = np.vstack(all_obs_xy_list)
+        all_obs_val = np.concatenate(all_obs_val_list)
+
+        return (psi_current, all_obs_xy, all_obs_val,
+                np.maximum(psi_pilot, 0.0), obs_xy_p1_last)
+    
+    @staticmethod
+    def generate(true_field, pred_field, num_points, method="uniform", seed=42,
+                 enkf=None, conds_preds=None, **sample_params):
+        field_shape = true_field.shape
+        if method == "random":
+            obs_xy = SamplingStrategies.sample_random(field_shape, num_points, seed)
+        elif method == "uniform":
+            obs_xy = SamplingStrategies.sample_uniform(field_shape, num_points)
+        elif method == "two_stage":
+            obs_xy, _ = SamplingStrategies.two_stage_sampling(
+                true_field,
+                pred_field,
+                num_points,
+                seed=seed,
+                **sample_params
+            )
+        elif method == "two_stage_pro":
+            obs_xy, _ = SamplingStrategies.two_stage_pro(
+                true_field,
+                pred_field,
+                num_points,
+                seed=seed,
+                **sample_params
+            )
+        elif method == "smart_two_pass":
+            if enkf is None or conds_preds is None:
+                raise ValueError(
+                    "method='smart_two_pass' 需要传入 enkf 实例和 conds_preds 集合场。"
+                )
+            # 解析 n1 / n2（支持用 n1_ratio 自动计算）
+            n1_ratio = float(sample_params.pop('n1_ratio', 0.6))
+            n1_default = int(round(num_points * n1_ratio))
+            n1 = int(sample_params.pop('n1', n1_default))
+            if num_points > 1:
+                n1 = max(1, min(n1, num_points - 1))
+            else:
+                n1 = 1
+            n2 = int(sample_params.pop('n2', num_points - n1))
+            return SamplingStrategies.smart_two_pass(
+                enkf=enkf,
+                psi_f=pred_field,
+                conds_preds=conds_preds,
+                true_field=true_field,
+                n1=n1,
+                n2=n2,
+                seed=seed,
+                **sample_params,
+            )
+        else:
+            raise ValueError(f"Unknown observation sampling method: {method}")
+    
+        obs_val = ObservationModel.observation_operator_H(true_field, obs_xy)
+        return obs_xy, obs_val
+    
+
+class EnKF:
+
+    def __init__(
+        self,
+        obs_std_scale=0.08,   # relative observation noise level
+        damping=1.0,
+        jitter=1e-5,
+    ):
+        self.obs_std_scale = obs_std_scale
+        self.damping = damping
+        self.jitter = jitter
+    
+    def standard_enkf(self, psi_f, conds_preds, obs_xy, d_obs):
+        """
+        psi_f: Unet预测的最佳先验场 (H, W)
+        conds_preds: 通过扰动参数生成的集合场 (N_ens, H, W)
+        obs_xy: 监测站坐标 (n_obs, 2)
+        d_obs: 监测站真实浓度 (n_obs,)
+        """
+        conds_preds = np.asarray(conds_preds)
+        N_ens, H, W = conds_preds.shape
+        n_obs = obs_xy.shape[0]
+
+        ens_mean = np.mean(conds_preds, axis=0) # 计算集合均值
+        # 重要：将集合成员的波动叠加到 Unet 预测场 psi_f 上 ,
+        # 确保分析场的统计中心是 Unet 预测的那个场，而不是集合均值（可能有偏差导致更新不好）
+        X_f = conds_preds - ens_mean[None, :, :] + psi_f[None, :, :]
+        X_f_flat = X_f.reshape(N_ens, -1)  # (N_ens, Pixels)
+        HX = ObservationModel.observation_operator_H_ens(X_f, obs_xy)  # (N_ens, n_obs)
+        HX_mean = np.mean(HX, axis=0)
+        X_f_bar = np.mean(X_f_flat, axis=0) # 计算偏差矩阵
+        A_prime = (X_f_flat - X_f_bar[None, :]).T # A_prime (状态偏差): (Pixels, N_ens)
+        Y_prime = (HX - HX_mean).T # Y_prime (观测空间偏差): (n_obs, N_ens)
+        
+        # # 构造观测误差矩阵 R_e
+        # # 基于观测值大小设定自适应噪声 (8% 相对误差)
+        obs_std = self.obs_std_scale * np.maximum(np.abs(d_obs), 1.0) # 先定标准差
+        rng = np.random.default_rng(42) # SVD正交化生成E
+        Z = rng.standard_normal((N_ens, n_obs))
+        U, _, Vt = np.linalg.svd(Z, full_matrices=False)
+        Z = U @ Vt * np.sqrt(N_ens - 1)
+        E = Z * obs_std[None, :]   # (N_ens, n_obs)，E.T即为文献中的E矩阵
+        # 从E计算Re（按照文献公式 Re = EE^T / N-1）
+        E_T = E.T                          # (n_obs, N_ens)，对应文献的E
+        R_e = (E_T @ E_T.T) / (N_ens - 1) # (n_obs, n_obs)
+        R_e += self.jitter * np.eye(n_obs) # 数值稳定项
+        Y_o = d_obs[None, :] + E  # (N_ens, n_obs)
+        
+        # 增益计算与状态更新 (对应公式 3-16, 3-17)
+        # 计算 Pe*H.T 和 H*Pe*H.T 的统计估计值
+        Pe_HT = (A_prime @ Y_prime.T) / (N_ens - 1)
+        H_Pe_HT = (Y_prime @ Y_prime.T) / (N_ens - 1)
+        # 计算集合卡尔曼增益 K_e = Pe*H.T * inverse(H*Pe*H.T + R_e)
+        # 使用 solve 提高数值稳定性
+        K_e = np.linalg.solve((H_Pe_HT + R_e).T, Pe_HT.T).T
+        # 计算创新值 (Innovation): (n_obs, N_ens)
+        # 每个成员根据自己的观测扰动和预测值进行修正
+        innovation = (Y_o - HX).T 
+        # 更新系统状态的集合预测矩阵 X_a
+        # X_a = X_f + K_e * (Y_o - HX_f)
+        X_a_flat = X_f_flat + (self.damping * (K_e @ innovation)).T
+        # 输出最终分析场，取集合均值作为最终结果
+        psi_a_flat = np.mean(X_a_flat, axis=0)
+        psi_a = psi_a_flat.reshape(H, W)
+        # 物理约束：确保浓度不为负数
+        return np.maximum(psi_a, 0.0)
+
+    def enkf_localization(self, psi_f, conds_preds, obs_xy, d_obs,
+                        loc_radius_pixobs=40.0,  # Pixel-Obs localization radius (in pixels)
+                        loc_radius_obsobs=60.0,  # Obs-Obs localization radius (in pixels)
+                        seed=42,
+                        SAVE_DIAGNOSTICS=False,
+        ):
+        conds_preds = np.asarray(conds_preds)
+        N_ens, H, W = conds_preds.shape
+        n_obs = obs_xy.shape[0]
+
+        # ========= 1) prior ensemble centered at psi_f =========
+        ens_mean = np.mean(conds_preds, axis=0)
+        X_f = conds_preds - ens_mean[None, :, :] + psi_f[None, :, :]
+        X_f_flat = X_f.reshape(N_ens, -1)
+
+        HX = ObservationModel.observation_operator_H_ens(X_f, obs_xy)  # 注意：必须用 X_f
+        HX_mean = np.mean(HX, axis=0)
+
+        X_f_bar = np.mean(X_f_flat, axis=0)
+        A_prime = (X_f_flat - X_f_bar[None, :]).T          # (Pixels, N_ens)
+        Y_prime = (HX - HX_mean).T                          # (n_obs, N_ens)
+
+        # # ========= 2) perturbed obs (deterministic-ish, fixed seed) =========
+        obs_std = self.obs_std_scale * np.maximum(np.abs(d_obs), 1.0) # 先定标准差
+        rng = np.random.default_rng(seed) # SVD正交化生成E
+        Z = rng.standard_normal((N_ens, n_obs))
+        U, _, Vt = np.linalg.svd(Z, full_matrices=False)
+        Z = U @ Vt * np.sqrt(N_ens - 1)
+        E = Z * obs_std[None, :]   # (N_ens, n_obs)，E.T即为文献中的E矩阵
+        # 从E计算Re（按照文献公式 Re = EE^T / N-1）
+        E_T = E.T                          # (n_obs, N_ens)，对应文献的E
+        R_e = (E_T @ E_T.T) / (N_ens - 1) # (n_obs, n_obs)
+        R_e += self.jitter * np.eye(n_obs) # 数值稳定项
+        Y_o = d_obs[None, :] + E
+
+        # ========= 3) sample covariances =========
+        Pe_HT = (A_prime @ Y_prime.T) / (N_ens - 1)        # (Pixels, n_obs)
+        H_Pe_HT = (Y_prime @ Y_prime.T) / (N_ens - 1)      # (n_obs, n_obs)
+
+        # ========= 4) localization =========
+        # (a) Pixel-Obs localization: rho_xy (Pixels, n_obs)
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        grid = np.stack([xx.ravel(), yy.ravel()], axis=1)   # (Pixels,2)
+        dx = grid[:, None, 0] - obs_xy[None, :, 0]
+        dy = grid[:, None, 1] - obs_xy[None, :, 1]
+        dist2_xy = dx*dx + dy*dy
+        rho_xy = np.exp(-0.5 * dist2_xy / (loc_radius_pixobs**2))
+
+        # (b) Obs-Obs localization: rho_oo (n_obs, n_obs)
+        dox = obs_xy[:, None, 0] - obs_xy[None, :, 0]
+        doy = obs_xy[:, None, 1] - obs_xy[None, :, 1]
+        dist2_oo = dox*dox + doy*doy
+        rho_oo = np.exp(-0.5 * dist2_oo / (loc_radius_obsobs**2))
+        Pe_HT = Pe_HT * rho_xy
+        H_Pe_HT = H_Pe_HT * rho_oo
+
+        # ========= [诊断] P_e 的谱结构 =========
+        # P_e = A_prime @ A_prime.T / (N_ens-1)，直接分解 A_prime 的奇异值更高效
+        # A_prime shape: (Pixels, N_ens)，SVD给出 P_e 的特征值 = sigma^2
+        U_ens, sigma, Vt_ens = np.linalg.svd(A_prime / np.sqrt(N_ens - 1), full_matrices=False)
+        # sigma shape: (N_ens,)，对应 P_e 的特征值平方根
+        eigenvalues = sigma ** 2  # P_e 的特征值，降序排列
+
+        # --- 指标1：有效秩 r_eff = (Σλ)² / Σλ² 衡量特征值分布均匀程度---
+        # r_eff→1: 近似秩1（能量集中于单一方向）r_eff→N: 各向同性（能量均匀分布）
+        r_eff = (eigenvalues.sum() ** 2) / (eigenvalues ** 2).sum()
+
+        # --- 指标2：主特征值 λ1 = P_e 在主方向上的方差 ---
+        # 只受幅度参数(v, Q)影响，随d单调增大
+        lambda1 = eigenvalues[0]
+        lambda_min = eigenvalues[-2]
+
+        # --- 指标3：方向集中度 λ1/λ2 衡量P_e各向异性程度 ---
+        # 峰值对应最优d配置（d**），超过后集合引入非物理方向
+        ratio_1_2 = eigenvalues[0] / eigenvalues[1] if len(eigenvalues) > 1 else np.inf
+
+        # --- 指标4：主特征向量峰值位置---
+        # 峰值位置随d系统性漂移，随v/Q不变，随n随机漂移
+        u1 = U_ens[:, 0].reshape(H, W)  # u1 = P_e 的第一特征向量，代表集合扰动的主方向
+        u1_peak = np.unravel_index(np.abs(u1).argmax(), u1.shape)
+
+        # ========= 5) Kalman gain =========
+        S = H_Pe_HT + R_e
+        K_e = np.linalg.solve(S.T, Pe_HT.T).T
+
+        innovation = (Y_o - HX).T
+        X_a_flat = X_f_flat + (self.damping * (K_e @ innovation)).T
+
+        psi_a = np.mean(X_a_flat, axis=0).reshape(H, W)
+        psi_a = np.maximum(psi_a, 0.0)
+
+        if SAVE_DIAGNOSTICS:
+            print("=" * 50)
+            print(f"[P_e 谱诊断]")
+            print(f"  指标1 r_eff       = {r_eff:.2f}   # (Σλ)²/Σλ²，建筑影响下界≈2.1")
+            print(f"  指���2 λ1          = {lambda1:.2f} {lambda_min:.2f} # 主方向方差，随d单调增大")
+            print(f"  指标3 λ1/λ2       = {ratio_1_2:.2f}   # 各向异性，d=45°时峰值→最优配置")
+            print(f"  指标4 u1峰值位置   = {u1_peak}  # d变化时系统漂移，v/Q不变")
+            print("=" * 50)
+            diag = {
+                'r_eff':    r_eff,
+                'lambda1':  lambda1,
+                'ratio_1_2': ratio_1_2,
+                'u1_peak_row': u1_peak[0],
+                'u1_peak_col': u1_peak[1],
+            }
+            return psi_a, diag
+        else:
+            return psi_a
+        
+    def _enkf_update_standard(self, X_f, obs_xy, d_obs):
+        """
+        标准 EnKF 更新，直接接受已中心化的集合 X_f (N_ens, H, W)。
+        返回分析场均值 psi_a (H, W)，>=0。
+        """
+        N_ens, H, W = X_f.shape
+        n_obs = obs_xy.shape[0]
+        X_f_flat = X_f.reshape(N_ens, -1)
+ 
+        HX = ObservationModel.observation_operator_H_ens(X_f, obs_xy)
+        HX_mean = np.mean(HX, axis=0)
+        X_f_bar = np.mean(X_f_flat, axis=0)
+        A_prime = (X_f_flat - X_f_bar[None, :]).T   # (Pixels, N_ens)
+        Y_prime = (HX - HX_mean).T                   # (n_obs,  N_ens)
+ 
+        obs_std = self.obs_std_scale * np.maximum(np.abs(d_obs), 1.0)
+        rng = np.random.default_rng(42)
+        Z = rng.standard_normal((N_ens, n_obs))
+        U, _, Vt = np.linalg.svd(Z, full_matrices=False)
+        Z = U @ Vt * np.sqrt(N_ens - 1)
+        E = Z * obs_std[None, :]
+        E_T = E.T
+        R_e = (E_T @ E_T.T) / (N_ens - 1)
+        R_e += self.jitter * np.eye(n_obs)
+        Y_o = d_obs[None, :] + E
+ 
+        Pe_HT   = (A_prime @ Y_prime.T) / (N_ens - 1)
+        H_Pe_HT = (Y_prime @ Y_prime.T) / (N_ens - 1)
+        K_e = np.linalg.solve((H_Pe_HT + R_e).T, Pe_HT.T).T
+        innovation = (Y_o - HX).T
+        X_a_flat = X_f_flat + (self.damping * (K_e @ innovation)).T
+        psi_a = np.mean(X_a_flat, axis=0).reshape(H, W)
+        return np.maximum(psi_a, 0.0)
+ 
+    def _enkf_update_localized(self, X_f, obs_xy, d_obs,
+                                loc_radius_pixobs=35.0,
+                                loc_radius_obsobs=40.0,
+                                seed=42):
+        """
+        局地化 EnKF 更新，直接接受已中心化的集合 X_f (N_ens, H, W)。
+        返回分析场均值 psi_a (H, W)，>=0。
+        """
+        N_ens, H, W = X_f.shape
+        n_obs = obs_xy.shape[0]
+        X_f_flat = X_f.reshape(N_ens, -1)
+ 
+        HX = ObservationModel.observation_operator_H_ens(X_f, obs_xy)
+        HX_mean = np.mean(HX, axis=0)
+        X_f_bar = np.mean(X_f_flat, axis=0)
+        A_prime = (X_f_flat - X_f_bar[None, :]).T
+        Y_prime = (HX - HX_mean).T
+ 
+        obs_std = self.obs_std_scale * np.maximum(np.abs(d_obs), 1.0)
+        rng = np.random.default_rng(seed)
+        Z = rng.standard_normal((N_ens, n_obs))
+        U, _, Vt = np.linalg.svd(Z, full_matrices=False)
+        Z = U @ Vt * np.sqrt(N_ens - 1)
+        E = Z * obs_std[None, :]
+        E_T = E.T
+        R_e = (E_T @ E_T.T) / (N_ens - 1)
+        R_e += self.jitter * np.eye(n_obs)
+        Y_o = d_obs[None, :] + E
+ 
+        Pe_HT   = (A_prime @ Y_prime.T) / (N_ens - 1)
+        H_Pe_HT = (Y_prime @ Y_prime.T) / (N_ens - 1)
+ 
+        yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
+        grid = np.stack([xx.ravel(), yy.ravel()], axis=1)
+        dx = grid[:, None, 0] - obs_xy[None, :, 0]
+        dy = grid[:, None, 1] - obs_xy[None, :, 1]
+        rho_xy = np.exp(-0.5 * (dx*dx + dy*dy) / (loc_radius_pixobs**2))
+        dox = obs_xy[:, None, 0] - obs_xy[None, :, 0]
+        doy = obs_xy[:, None, 1] - obs_xy[None, :, 1]
+        rho_oo = np.exp(-0.5 * (dox*dox + doy*doy) / (loc_radius_obsobs**2))
+ 
+        Pe_HT   = Pe_HT   * rho_xy
+        H_Pe_HT = H_Pe_HT * rho_oo
+ 
+        S = H_Pe_HT + R_e
+        K_e = np.linalg.solve(S.T, Pe_HT.T).T
+        innovation = (Y_o - HX).T
+        X_a_flat = X_f_flat + (self.damping * (K_e @ innovation)).T
+        psi_a = np.mean(X_a_flat, axis=0).reshape(H, W)
+        return np.maximum(psi_a, 0.0)
+        
+        
+class PrintMetrics:
+    @staticmethod
+    def pad_center_crop(arr, center_y, center_x, out_h=256, out_w=256):
+        # Pad and center-crop 2D or 3D array
+        if arr.ndim == 3:
+            C, H, W = arr.shape
+            out = np.zeros((C, out_h, out_w), dtype=arr.dtype)
+        else:
+            H, W = arr.shape
+            out = np.zeros((out_h, out_w), dtype=arr.dtype)
+        y0, x0 = center_y - out_h // 2, center_x - out_w // 2
+        y1, x1 = y0 + out_h, x0 + out_w
+        sy0, sy1 = max(0, y0), min(H, y1)
+        sx0, sx1 = max(0, x0), min(W, x1)
+        dy0, dx0 = sy0 - y0, sx0 - x0
+        dy1, dx1 = dy0 + (sy1 - sy0), dx0 + (sx1 - sx0)
+        if arr.ndim == 3:
+            out[:, dy0:dy1, dx0:dx1] = arr[:, sy0:sy1, sx0:sx1]
+        else:
+            out[dy0:dy1, dx0:dx1] = arr[sy0:sy1, sx0:sx1]
+        return out
+    
+    @staticmethod
+    def get_building_area():
+        # load building data
+        npz_path = '../Gas_unet/Gas_code/dataset_m/5min_m_Data_special/min5_m_v1_0_d270_sc2_s10_04118.npz'
+        data = np.load(npz_path)
+        build_data = data['three_channel_data'][0]
+        non_building_mask = (build_data == 0).astype(np.uint8)
+        center_y, center_x = 498, 538
+        build_data_256 = PrintMetrics.pad_center_crop(build_data, center_y,
+                                                      center_x, 256, 256)
+        non_building_mask = PrintMetrics.pad_center_crop(non_building_mask,
+                                                         center_y, center_x, 256, 256)
+        return build_data_256, non_building_mask
+    
+    @staticmethod
+    def weighted_r2(y_true, y_pred, gamma=1.0, eps=1e-12):
+        """
+        Weighted R2 score emphasizing high-value regions.
+        """
+        y_true = np.asarray(y_true)
+        y_pred = np.asarray(y_pred)
+
+        w = np.maximum(y_true, eps) ** gamma
+        w = w / np.sum(w)
+
+        y_bar = np.sum(w * y_true)
+
+        num = np.sum(w * (y_true - y_pred) ** 2)
+        den = np.sum(w * (y_true - y_bar) ** 2)
+
+        if den < eps:
+            return np.nan
+
+        return 1.0 - num / den
+    
+    
+      
+    @staticmethod
+    def print_metrics(i, wind_speed, wind_direction, sc, source_number,
+                      true_field, pred_field, analysis, obs_xy,
+                      metrics_save_flag=False, metrics_print_flag=True):
+        """
+        Metrics:
+        1) Field-wise (all pixels)
+        2) Plume-aware (true > eps)
+        3) At observations
+        """
+
+        def nmse_metrics(y_true, y_pred):
+            nmse = np.mean((y_true.flatten() - y_pred.flatten())**2) / (np.mean(y_true) * np.mean(y_pred) + 1e-12)
+            return nmse
+        
+        def nmae_metrics(y_true, y_pred):
+            nmae = np.mean(np.abs(y_true.flatten() - y_pred.flatten())) / (np.mean(y_true) + 1e-12)
+            return nmae
+
+        # ========= 保留原始 2D 场 =========
+        _, non_building_mask = PrintMetrics.get_building_area()
+        true_field = np.where(true_field > 0, true_field, 0) * non_building_mask
+        pred_field = np.where(pred_field > 0, pred_field, 0) * non_building_mask
+        analysis   = np.where(analysis   > 0, analysis,   0) * non_building_mask
+        true_flat = true_field.ravel()
+        pred_flat = pred_field.ravel()
+        ana_flat  = analysis.ravel()
+
+        # ===============================
+        # (1) Field-wise (all pixels)
+        # ===============================
+        r2_before = r2_score(true_flat, pred_flat)
+        r2_after  = r2_score(true_flat, ana_flat)
+        mse_before = mean_squared_error(true_flat, pred_flat)
+        mse_after  = mean_squared_error(true_flat, ana_flat)
+        mae_before = mean_absolute_error(true_flat, pred_flat)
+        mae_after  = mean_absolute_error(true_flat, ana_flat)
+        nmse_before = nmse_metrics(true_flat, pred_flat)
+        nmse_after = nmse_metrics(true_flat, ana_flat)
+        nmae_before = nmae_metrics(true_flat, pred_flat)
+        nmae_after = nmae_metrics(true_flat, ana_flat)
+
+        # ===============================
+        # (2) Plume-aware (true > eps)
+        # ===============================
+        plume_mask = true_flat > 1e-6
+        true_p = true_flat[plume_mask]
+        pred_p = pred_flat[plume_mask]
+        ana_p  = ana_flat[plume_mask]
+        r2_plume_before = r2_score(true_p, pred_p)
+        r2_plume_after  = r2_score(true_p, ana_p)
+        mse_plume_before = mean_squared_error(true_p, pred_p)
+        mse_plume_after  = mean_squared_error(true_p, ana_p)
+        mae_plume_before = mean_absolute_error(true_p, pred_p)
+        mae_plume_after  = mean_absolute_error(true_p, ana_p)
+        nmse_plume_before = nmse_metrics(true_p, pred_p)
+        nmse_plume_after = nmse_metrics(true_p, ana_p)
+        nmae_plume_before = nmae_metrics(true_p, pred_p)
+        nmae_plume_after = nmae_metrics(true_p, ana_p)
+
+        # ---- Weighted R2 (plume-aware) ----
+        wr2_plume_before = PrintMetrics.weighted_r2(true_p, pred_p, gamma=1.0)
+        wr2_plume_after  = PrintMetrics.weighted_r2(true_p, ana_p,  gamma=1.0)
+
+
+        # ===============================
+        # (3) At observations
+        # ===============================
+        true_at_obs = ObservationModel.observation_operator_H(true_field, obs_xy)
+        pred_at_obs = ObservationModel.observation_operator_H(pred_field, obs_xy)
+        ana_at_obs  = ObservationModel.observation_operator_H(analysis, obs_xy)
+        r2_obs_before = r2_score(true_at_obs, pred_at_obs)
+        r2_obs_after  = r2_score(true_at_obs, ana_at_obs)
+        mse_obs_before = mean_squared_error(true_at_obs, pred_at_obs)
+        mse_obs_after  = mean_squared_error(true_at_obs, ana_at_obs)
+        mae_obs_before = mean_absolute_error(true_at_obs, pred_at_obs)
+        mae_obs_after  = mean_absolute_error(true_at_obs, ana_at_obs)
+        nmse_obs_before = nmse_metrics(true_at_obs, pred_at_obs)
+        nmse_obs_after = nmse_metrics(true_at_obs, ana_at_obs)
+        nmae_obs_before = nmae_metrics(true_at_obs, pred_at_obs)
+        nmae_obs_after = nmae_metrics(true_at_obs, ana_at_obs)
+
+        if metrics_print_flag:
+            print("=== Assimilation Metrics ===")
+            print("[Field-wise]")
+            print(f"R2  : {r2_before:.4f}->{r2_after:.4f}")
+            print(f"MSE : {mse_before:.4f}->{mse_after:.4f}")
+            print(f"MAE : {mae_before:.4f}->{mae_after:.4f}")
+            print("[Plume-aware]")
+            print(f"R2  : {r2_plume_before:.4f}->{r2_plume_after:.4f}")
+            print(f"MSE : {mse_plume_before:.4f}->{mse_plume_after:.4f}")
+            print(f"MAE : {mae_plume_before:.4f}->{mae_plume_after:.4f}")
+            print(f"W-R2 : {wr2_plume_before:.4f}->{wr2_plume_after:.4f}")
+            print("[At observations]")
+            print(f"R2  : {r2_obs_before:.4f}->{r2_obs_after:.4f}")
+            print(f"MSE : {mse_obs_before:.4f}->{mse_obs_after:.4f}")
+            print(f"MAE : {mae_obs_before:.4f}->{mae_obs_after:.4f}")
+
+        if metrics_save_flag:
+            return {
+                'idx': i,
+                'wind_speed': wind_speed,
+                'wind_direction': wind_direction,
+                'stability_class': sc,
+                'source_number': source_number,
+                "r2_before": r2_before,
+                "r2_after": r2_after,
+                "r2_plume_before": r2_plume_before,
+                "r2_plume_after": r2_plume_after,
+                "w_r2_plume_before": wr2_plume_before,
+                "w_r2_plume_after": wr2_plume_after,
+                "r2_obs_before": r2_obs_before,
+                "r2_obs_after": r2_obs_after,
+                "mse_before": mse_before,
+                "mse_after": mse_after,
+                "mse_plume_before": mse_plume_before,
+                "mse_plume_after": mse_plume_after,
+                "mse_obs_before": mse_obs_before,
+                "mse_obs_after": mse_obs_after,
+                "mae_before": mae_before,
+                "mae_after": mae_after,
+                "mae_plume_before": mae_plume_before,
+                "mae_plume_after": mae_plume_after,
+                "mae_obs_before": mae_obs_before,
+                "mae_obs_after": mae_obs_after,
+                "nmse_before": nmse_before,
+                "nmse_after": nmse_after,
+                "nmse_plume_before": nmse_plume_before,
+                "nmse_plume_after": nmse_plume_after,
+                "nmae_before": nmae_before,
+                "nmae_after": nmae_after,
+                "nmae_plume_before": nmae_plume_before,
+                "nmae_plume_after": nmae_plume_after,
+                "nmse_obs_before": nmse_obs_before,
+                "nmse_obs_after": nmse_obs_after,
+                "nmae_obs_before": nmae_obs_before,
+                "nmae_obs_after": nmae_obs_after,
+            }
+
+class Visualization:
+    def plot_assimilation_with_building(
+        true_field,
+        pred_field,
+        analysis,
+        obs_xy,
+        vmax=10,
+        title_suffix=""
+    ):
+        """
+        - 建筑 mask
+        - 非建筑区浓度
+        - 同化前 / 后对比
+        """
+
+        # ---------- 物理裁剪 + 建筑 mask ----------
+        build_data_256, non_building_mask = PrintMetrics.get_building_area()
+        true_field = np.where(true_field > 0, true_field, 0) * non_building_mask
+        pred_field = np.where(pred_field > 0, pred_field, 0) * non_building_mask
+        analysis   = np.where(analysis   > 0, analysis,   0) * non_building_mask
+
+        # ---------- 画图 ----------
+        fig, axs = plt.subplots(1, 3, figsize=(14, 4), dpi=300)
+        cmap = "inferno"
+        levels = np.linspace(0, vmax, 21)
+
+        im0 = axs[0].contourf(true_field, levels=levels, cmap=cmap, vmin=0, vmax=vmax,
+                            extend='max')
+        axs[0].set_title('True Field' + title_suffix)
+        plt.colorbar(im0, ax=axs[0])
+
+        im1 = axs[1].contourf(pred_field, levels=levels, cmap=cmap, vmin=0, vmax=vmax,
+                            extend='max')
+        axs[1].set_title(r'Prior Prediction Field $\psi^{f}$')
+        plt.colorbar(im1, ax=axs[1])
+
+        im2 = axs[2].contourf(analysis, levels=levels, cmap=cmap, vmin=0, vmax=vmax,
+                            extend='max')
+        axs[2].set_title(r'Analysis $\psi^{a}$')
+        plt.colorbar(im2, ax=axs[2])
+        axs[0].scatter(obs_xy[:, 0], obs_xy[:, 1], c='red', s=15, edgecolors='k')
+        axs[1].scatter(obs_xy[:, 0], obs_xy[:, 1], c='red', s=15, edgecolors='k')
+        axs[2].scatter(obs_xy[:, 0], obs_xy[:, 1], c='red', s=15, edgecolors='k')
+
+        # ---------- 指标 ----------
+        axs[1].text(
+            80, 15,
+            f"$R^2$={r2_score(true_field.ravel(), pred_field.ravel()):.4f}\n"
+            f"$MSE$={mean_squared_error(true_field.ravel(), pred_field.ravel()):.4f}\n"
+            f"$MAE@Obs$={mean_absolute_error(ObservationModel.observation_operator_H(true_field, obs_xy),
+                                            ObservationModel.observation_operator_H(pred_field, obs_xy)):.3f}",
+            color='white'
+        )
+        axs[2].text(
+            80, 15,
+            f"$R^2$={r2_score(true_field.ravel(), analysis.ravel()):.4f}\n"
+            f"$MSE$={mean_squared_error(true_field.ravel(), analysis.ravel()):.4f}\n"
+            f"$MAE@Obs$={mean_absolute_error(ObservationModel.observation_operator_H(true_field, obs_xy),
+                                            ObservationModel.observation_operator_H(analysis, obs_xy)):.3f}",
+            color='white'
+        )
+        plt.tight_layout()
+        plt.show()
+
+    def plot_assimilation_4panel(
+        true_field,
+        pred_field,
+        analysis,
+        obs_xy,
+        obs_val,
+        vmin=0,
+        vmax=10,
+        title_suffix=""
+    ):
+        # ---------- 物理裁剪 + 建筑 mask ----------
+        _, non_building_mask = PrintMetrics.get_building_area()
+        true_field = np.where(true_field > 0, true_field, 0) * non_building_mask
+        pred_field = np.where(pred_field > 0, pred_field, 0) * non_building_mask
+        analysis   = np.where(analysis   > 0, analysis,   0) * non_building_mask
+
+        # ---------- Figure ----------
+        fig, axs = plt.subplots(1, 4, figsize=(18, 4), dpi=300)
+        cmap = "inferno"
+        levels = np.linspace(0, vmax, 21)
+
+        # ---------- (a) True field ----------
+        im0 = axs[0].contourf(true_field, levels=levels, cmap=cmap, vmin=vmin, vmax=vmax,
+                            extend='both')
+        axs[0].set_title("True Field" + title_suffix)
+        plt.colorbar(im0, ax=axs[0])
+
+        # ---------- (b) Prior prediction ----------
+        im1 = axs[1].contourf(pred_field, levels=levels, cmap=cmap, vmin=vmin, vmax=vmax,
+                            extend='both')
+        axs[1].set_title(r"Prior Prediction $\psi^{f}$")
+        plt.colorbar(im1, ax=axs[1])
+
+        # ---------- (c) Observations (points only) ----------
+        sc = axs[2].scatter(
+            obs_xy[:, 0],
+            obs_xy[:, 1],
+            c=obs_val,
+            cmap=cmap,
+            vmin=vmin,
+            vmax=vmax,
+            s=30,
+            edgecolors="k",
+            linewidths=0.4,
+            alpha=0.9
+        )
+        axs[2].set_title("Observations $d_i$")
+        axs[2].set_xlim(0, true_field.shape[1])
+        axs[2].set_ylim(true_field.shape[0], 0)
+        axs[2].set_aspect("equal")
+        axs[2].invert_yaxis()
+        plt.colorbar(sc, ax=axs[2], extend='both')
+
+        # ---------- (d) Analysis field ----------
+        im3 = axs[3].contourf(analysis, levels=levels, cmap=cmap, vmin=vmin, vmax=vmax,
+                            extend='both')
+        axs[3].set_title(r"Analysis $\psi^{a}$")
+        plt.colorbar(im3, ax=axs[3])
+
+        # ---------- Metrics ----------pred_at_obs
+        axs[1].text(
+            0.02, 0.95,
+            f"$R^2$={r2_score(true_field.ravel(), pred_field.ravel()):.4f}\n"
+            f"$MSE$={mean_squared_error(true_field.ravel(), pred_field.ravel()):.4f}",
+            transform=axs[1].transAxes,
+            va="top",
+            color="white"
+        )
+
+        axs[3].text(
+            0.02, 0.95,
+            f"$R^2$={r2_score(true_field.ravel(), analysis.ravel()):.4f}\n"
+            f"$MSE$={mean_squared_error(true_field.ravel(), analysis.ravel()):.4f}",
+            transform=axs[3].transAxes,
+            va="top",
+            color="white"
+        )
+        plt.tight_layout()
+        plt.show()
+
+    def plot_pe_spectrum(all_diags, save_flag=False):
+        C_BLUE = "#488ABA"
+        C_ORANGE = "#e5954e"
+
+        ds = [diag['d'] for diag in all_diags]
+        r_eff_values = [diag['r_eff'] for diag in all_diags]
+        ratio_12_values = [diag['ratio_1_2'] for diag in all_diags]
+
+        fig, ax = plt.subplots(figsize=(6, 3), dpi=300)
+        l1, = ax.plot(ds, r_eff_values,
+                    marker='o', color=C_BLUE, linewidth=1.5,
+                    alpha=0.4,
+                    markersize=6, label=r'$r_{\rm eff}$')
+        ax.set_xlabel(r'Wind direction', labelpad=3)
+        ax.set_ylabel(r'Effective rank $r_{\rm eff}$',
+                    color=C_BLUE, labelpad=4)
+        ax.tick_params(axis='y', colors=C_BLUE)
+        ax.spines['left'].set_color(C_BLUE)
+        ax.set_xticks(ds)
+        ax.set_xticklabels([f'{d}°' for d in ds])
+        ax.set_ylim(1.5, 3)
+
+        # 右轴：λ1/λ2
+        ax2 = ax.twinx()
+        l2, = ax2.plot(ds, ratio_12_values,
+                    marker='s', color=C_ORANGE, linewidth=1.5,
+                    alpha=0.4,
+                    markersize=6, label=r'$\lambda_1/\lambda_2$')
+        ax2.set_ylabel(r'Anisotropy $\lambda_1/\lambda_2$',
+                    color=C_ORANGE, labelpad=4)
+        ax2.tick_params(axis='y', colors=C_ORANGE)
+        ax2.spines['right'].set_color(C_ORANGE)
+        ax2.set_ylim(1.5, 3.5)
+
+        opt_idx = int(np.argmax(ratio_12_values))
+        ax2.axvline(ds[opt_idx], color='grey', linewidth=0.8, linestyle='--', alpha=0.6)
+        ax2.text(ds[opt_idx] + 0.8, 1.58, r'$d^{**}$', color='grey')
+
+        # 统一图例
+        fig.legend(handles=[l1, l2],
+                loc='upper left',
+                bbox_to_anchor=(0.15, 0.95),
+                ncol=1, frameon=False)
+
+        plt.tight_layout()
+        if save_flag:
+            plt.savefig('./figures/test1/reff_ratio.png', dpi=300, bbox_inches='tight',
+                    transparent=True)
+            # plt.savefig('./figures/test1/reff_ratio.svg', dpi=300, bbox_inches='tight', format='svg')
+        plt.show()
+
+    def assimilation_scatter(psi_t_log, psi_f_log, psi_a_log, obs_xy):
+        def log10_formatter(x, pos):
+            return r'$10^{%d}$' % x
+
+        obs_true = np.log10(ObservationModel.observation_operator_H(psi_t_log, obs_xy)+1e-3)
+        obs_prior = np.log10(ObservationModel.observation_operator_H(psi_f_log, obs_xy)+1e-3)
+        obs_analysis = np.log10(ObservationModel.observation_operator_H(psi_a_log, obs_xy)+1e-3)
+
+        fig, ax = plt.subplots(figsize=(5, 4.5), dpi=300)
+        vmin, vamx = -4, 2
+        lim = [vmin, vamx]
+        ax.plot(lim, lim, 'k--', lw=1, label='1:1 line', zorder=1)
+
+        for obs_pred, label, color in zip(
+            [obs_prior, obs_analysis],
+            ['Prior', 'Analysis'],
+            ['steelblue', 'tomato']
+        ):
+            # 散点
+            ax.scatter(obs_true, obs_pred, s=25, alpha=0.6, color=color, zorder=3)
+
+            slope, intercept, r, _, _ = stats.linregress(obs_true, obs_pred)
+            rmse = np.sqrt(np.mean((obs_pred - obs_true) ** 2))
+            x_fit = np.linspace(lim[0], lim[1], 100)
+            ax.plot(x_fit, slope * x_fit + intercept, '-', color=color, lw=1.5,
+                    label=f'{label}r:{r:.2f}', zorder=2)
+
+        ax.set_xlabel('log(True)')
+        ax.set_ylabel('log(Predicted)')
+        ax.legend(loc='lower right', frameon=False)
+        ax.set_xlim(-4, 2)
+        ax.set_ylim(-4, 2)
+        ax.xaxis.set_major_formatter(ticker.FuncFormatter(log10_formatter))
+        ax.yaxis.set_major_formatter(ticker.FuncFormatter(log10_formatter))
+
+        plt.tight_layout()
+        plt.show()
+
+    def none_assimilation_scatter(psi_t_log, psi_f_log, psi_a_log, obs_xy):
+        def sample_independent_points(field, obs_xy, num_points=100, seed=42):
+            H, W = field.shape
+            yy, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
+            all_xy = np.stack([xx.ravel(), yy.ravel()], axis=1)
+            obs_set = set(map(tuple, np.round(obs_xy).astype(int)))
+            obs_mask = np.array([tuple(p) not in obs_set for p in all_xy])
+            _, non_building_mask = PrintMetrics.get_building_area()
+            building_mask = non_building_mask[yy.ravel(), xx.ravel()] == 1
+            num_mask = field.ravel() > 1e-4
+            candidate_xy = all_xy[obs_mask & building_mask & num_mask]
+            rng = np.random.default_rng(seed)
+            idx = rng.choice(len(candidate_xy), num_points, replace=False)
+            return candidate_xy[idx]
+
+        test_xy = sample_independent_points(psi_t_log, obs_xy, 200)
+        obs_true = np.log10(ObservationModel.observation_operator_H(psi_t_log, test_xy) + 1e-6)
+        obs_prior = np.log10(ObservationModel.observation_operator_H(psi_f_log, test_xy) + 1e-6)
+        obs_analysis = np.log10(ObservationModel.observation_operator_H(psi_a_log, test_xy) + 1e-6)
+
+        def log10_formatter(x, pos):
+            return r'$10^{%d}$' % x
+
+        fig, ax = plt.subplots(figsize=(5, 4.5), dpi=300)
+
+        vmin, vmax = -4, 2
+        lim = [vmin, vmax]
+
+        # 1:1 line
+        ax.plot(lim, lim, 'k--', lw=1, label='1:1 line', zorder=1)
+
+        for obs_pred, label, color in zip(
+                [obs_prior, obs_analysis],
+                ['Prior', 'Analysis'],
+                ['steelblue', 'tomato']
+        ):
+
+            # scatter
+            ax.scatter(obs_true, obs_pred,
+                    s=30,
+                    alpha=0.65,
+                    color=color,
+                    zorder=3)
+            slope, intercept, r, _, _ = stats.linregress(obs_true, obs_pred)
+            rmse = np.sqrt(np.mean((obs_pred - obs_true) ** 2))
+            x_fit = np.linspace(lim[0], lim[1], 100)
+            ax.plot(x_fit, slope * x_fit + intercept, '-', color=color, lw=1.5,
+                    label=f'{label}r:{r:.2f}', zorder=2)
+
+        ax.set_xlim(vmin, vmax)
+        ax.set_ylim(vmin, vmax)
+        ax.set_xlabel('log(True)')
+        ax.set_ylabel('log(Predicted)')
+        ax.xaxis.set_major_formatter(ticker.FuncFormatter(log10_formatter))
+        ax.yaxis.set_major_formatter(ticker.FuncFormatter(log10_formatter))
+        ax.legend(loc='lower right', frameon=False)
+        plt.tight_layout()
+        plt.show()
+
+    def methods_comparison(source_idx=25,
+                           num_points = 10,
+                           methods = ['random', 'uniform', 'two_stage']):
+        
+        for method in methods:
+            print(f"\n=== 观测点采样方法: {method} ===")
+            data = np.load(f'./dataset/assim_conds/fields_n{num_points}_{method}_obs1.npz', allow_pickle=True)
+
+            all_fields = data['all_fields']
+            sample_data = all_fields[source_idx]
+            psi_t_log = sample_data['trues_log']
+            psi_f_log = sample_data['preds_log']
+            psi_a_log = sample_data['analysis_log']
+            psi_t_ppm = sample_data['trues_ppm']
+            psi_f_ppm = sample_data['preds_ppm']
+            psi_a_ppm = sample_data['analysis_ppm']
+            obs_xy = sample_data['obs_xy']
+            obs_value_log = sample_data['obs_value_log']
+            obs_value_ppm = sample_data['obs_value_ppm']
+            Visualization.plot_assimilation_with_building(
+                true_field=psi_t_log,
+                pred_field=psi_f_log,
+                analysis=psi_a_log,
+                obs_xy=obs_xy,
+                vmax=10,
+                title_suffix=f" (idx={source_idx})"
+            )
+            Visualization.plot_assimilation_4panel(
+                true_field=psi_t_log,
+                pred_field=psi_f_log,
+                analysis=psi_a_log,
+                obs_xy=obs_xy,
+                obs_val=obs_value_log,
+                vmax=10,
+                title_suffix=f" (idx={source_idx})"
+            )
+            Visualization.plot_assimilation_4panel(
+                true_field=psi_t_ppm,
+                pred_field=psi_f_ppm,
+                analysis=psi_a_ppm,
+                obs_xy=obs_xy,
+                obs_val=obs_value_ppm,
+                vmax=200,
+                title_suffix=f" in PPM SPACE (idx={source_idx})"
+            )
+
+    def plot_n_hist_comparison(obs_tag=1, space_mode="log",
+                               methods=["random", "uniform", "two_stage"],
+                               n_list=[10, 20, 30, 40, 50],
+                               base_dir="./dataset/assim_conds",
+                               plot_mode="after",
+                               target_method="two_stage"):
+        scope_labels = ["overall", "plume", "obs"]
+        metric_keys = {
+            "r2":  ["r2", "r2_plume", "r2_obs"],
+            "nmse": ["nmse", "nmse_plume", "nmse_obs"],
+            "nmae": ["nmae", "nmae_plume", "nmae_obs"],
+        }
+
+        def load_method_df(space, method, num_points):
+            candidate_methods = [method]
+            if method != "two_stage_pro":
+                candidate_methods.append("two_stage_pro")
+            for m in candidate_methods:
+                fp = os.path.join(
+                    base_dir,
+                    f"assimi_{space}_n{num_points}_{m}_obs{obs_tag}.csv"
+                )
+                if os.path.exists(fp):
+                    if m != method:
+                        print(f"[Info] {method} not exsist, back to {fp}")
+                    return pd.read_csv(fp)
+
+            print(f"[Warning] file not found: {candidate_methods}")
+            return None
+
+        def get_metric_value(df, metric_name):
+            before_col = f"{metric_name}_before"
+            after_col = f"{metric_name}_after"
+
+            if before_col not in df.columns or after_col not in df.columns:
+                return np.nan
+
+            before_mean = df[before_col].mean()
+            after_mean = df[after_col].mean()
+
+            if plot_mode == "delta":
+                return after_mean - before_mean
+            return after_mean
+
+        data = {
+            "r2":  np.full((len(n_list), 3), np.nan),
+            "nmse": np.full((len(n_list), 3), np.nan),
+            "nmae": np.full((len(n_list), 3), np.nan),
+        }
+
+
+        for i_n, n in enumerate(n_list):
+            df = load_method_df(space_mode, target_method, n)
+            if df is None:
+                continue
+
+            for metric_type in ["r2", "nmse", "nmae"]:
+                vals = []
+                for mk in metric_keys[metric_type]:
+                    vals.append(get_metric_value(df, mk))
+                data[metric_type][i_n, :] = vals
+
+        # print(f"space_mode = {space_mode}, plot_mode = {plot_mode}, method = {target_method}")
+        # for metric_type in ["r2", "mse", "mae"]:
+        #     print(f"\n{metric_type.upper()}:")
+        #     print(pd.DataFrame(data[metric_type], index=n_list, columns=scope_labels))
+
+        fig, ax1 = plt.subplots(figsize=(12, 6), dpi=300)
+        ax2 = ax1.twinx()
+        x_base = np.arange(len(n_list))
+        scope_offsets = {
+            "overall": -0.24,
+            "plume":   0.00,
+            "obs":     0.24,
+        }
+        bar_w = 0.20
+
+        colors = cm.get_cmap("Blues")
+        scope_colors = {
+            "overall": colors(0.45),
+            "plume": colors(0.65),
+            "obs": colors(0.85),
+            "edge": colors(0.85),
+        }
+        scope_linestyles = {
+            "overall": "-",
+            "plume": "--",
+            "obs": ":",
+        }
+
+        scope_markers_mse = {
+            "overall": "o",
+            "plume": "o",
+            "obs": "o",
+        }
+
+        scope_markers_mae = {
+            "overall": "s",
+            "plume": "s",
+            "obs": "s",
+        }
+
+        # -------------------------
+        # 左轴：R2 柱状图
+        # -------------------------
+        text_r2 = r"$\mathit{R}^2$"
+        for j, scope in enumerate(scope_labels):
+            x = x_base + scope_offsets[scope]
+            y = data["r2"][:, j]
+
+            ax1.bar(
+                x, y,
+                width=bar_w,
+                color=scope_colors[scope],
+                alpha=0.75,
+                label=f"{text_r2}-{scope}",
+                edgecolor=scope_colors["edge"],
+                zorder=2
+            )
+            # 给每根柱子加数值
+            for xi, yi in zip(x, y):
+                if np.isfinite(yi):
+                    ax1.text(
+                        xi, yi + 0.001, f"{yi:.2f}",
+                        ha="center", va="bottom"
+                    )
+
+        ax1.set_ylabel(r"$\mathit{R}^2$")
+        ax1.set_xticks(x_base)
+        ax1.set_xticklabels([f"n={n}" for n in n_list])
+        ax1.grid(axis="y", linestyle="--", alpha=0.25, zorder=0)
+        # ax1.set_ylim(0.6, 1.1)
+
+        if plot_mode == "delta":
+            ax1.axhline(0, color="k", linewidth=1)
+
+        # 在每个 n 下标出 overall / plume / obs
+        y1_min, y1_max = ax1.get_ylim()
+        y_text = y1_min - 0.06 * (y1_max - y1_min)
+        # for i in range(len(n_list)):
+        #     ax1.text(x_base[i] + scope_offsets["overall"], y_text, "overall",
+        #              ha="center", va="top")
+        #     ax1.text(x_base[i] + scope_offsets["plume"], y_text, "plume",
+        #              ha="center", va="top")
+        #     ax1.text(x_base[i] + scope_offsets["obs"], y_text, "obs",
+        #              ha="center", va="top")
+
+        for j, scope in enumerate(scope_labels):
+            x = x_base + scope_offsets[scope]
+
+            ax2.plot(
+                x,
+                data["nmse"][:, j],
+                color=scope_colors[scope],
+                linestyle="--",
+                marker="o",
+                linewidth=1.8,
+                markersize=5,
+                label=f"NMSE-{scope}",
+                markeredgecolor=scope_colors["edge"],
+                zorder=3
+            )
+
+            ax2.plot(
+                x,
+                data["nmae"][:, j],
+                color=scope_colors[scope],
+                linestyle=":",
+                marker="s",
+                linewidth=1.8,
+                markersize=5,
+                markeredgecolor=scope_colors["edge"],
+                label=f"NMAE-{scope}",
+                zorder=3
+            )
+
+        ax2.set_ylabel("NMSE / NMAE")
+        # ax2.set_ylim(0, 0.5)
+
+        if plot_mode == "delta":
+            ax2.axhline(0, color="gray", linewidth=1, alpha=0.6)
+
+        h1, l1 = ax1.get_legend_handles_labels()
+        h2, l2 = ax2.get_legend_handles_labels()
+
+        ax1.legend(
+            h1 + h2,
+            l1 + l2,
+            frameon=False,
+            loc="center left",
+            bbox_to_anchor=(1.15, 0.5)
+        )
+
+        plt.tight_layout()
+        plt.show()
+
+def cal_all_test_resluts(config, conds_pkl_path=None,
+                         data_path=None,
+                         use_localization=True,
+                         save_fields_flag=True,
+                         save_metrics_flag=False):
+    '''
+    使用说明：
+    sample_method_lists = ["random", "uniform", "two_stage"]
+    for method in sample_method_lists:
+        config_test = {
+        "num_points": 20,
+        "sample_method": method,
+        "obs_std_scale": 0.01,
+        "damping": 1,
+        "two_stage_params": {
+            "min_dist": 28,
+            "n1_ratio": 0.6,
+            "stage1_support_frac": 0.2,
+            "stage1_grad_power": 0.8,
+            "stage1_value_power": 1.2,
+            "stage1_center_boost": 1.2,
+            },
+        }
+        cal_all_test_resluts(config_test)
+    '''
+    # ==== 加载数据 ====
+    num_points = config['num_points']
+    sample_method = config['sample_method']  # "random" or "uniform"
+    # Kalman 参数
+    obs_std_scale=config['obs_std_scale']
+    damping=config['damping']
+    sample_method = config["sample_method"]
+    sample_params = {}
+    params_key = f"{sample_method}_params"
+    if params_key in config and isinstance(config[params_key], dict):
+        sample_params = config[params_key]
+
+    if conds_pkl_path is not None:
+        loader = DataLoader(
+            pred_npz_path='./dataset/pre_data/all_test_pred2.npz',
+            meta_txt_path='./dataset/pre_data/combined_test_special.txt',
+            conds_pkl_path=conds_pkl_path
+        )
+    else:
+        loader = DataLoader(
+            pred_npz_path='./dataset/pre_data/all_test_pred2.npz',
+            meta_txt_path='./dataset/pre_data/combined_test_special.txt',
+            conds_pkl_path='./dataset/pre_data/pred_condition/test_results/conditioned_results_v0_5_d45_n40.pkl'
+        )
+    trues, preds = loader.trues, loader.preds
+    # print(f"Total test samples: {len(preds)}")
+    all_metrics_log = []
+    all_metrics_ppm = []
+    all_fields = []
+    enkf = EnKF(obs_std_scale=obs_std_scale, damping=damping)
+    for i in trange(len(preds), desc="Running assimilation"):
+        psi_f_ppm, psi_t_ppm, conds_preds, meta = loader.get_sample(idx=i, in_ppm=True)
+        psi_f_log = np.log1p(np.maximum(psi_f_ppm, 0))
+        psi_t_log = np.log1p(np.maximum(psi_t_ppm, 0))
+        conds_log = np.log1p(np.maximum(conds_preds, 0))
+
+        if sample_method == "smart_two_pass":
+            n1_ratio = float(sample_params.get('n1_ratio', 0.6))
+            n1_default = int(round(num_points * n1_ratio))
+            n1 = int(sample_params.get('n1', n1_default))
+            if num_points > 1:
+                n1 = max(1, min(n1, num_points - 1))
+            else:
+                n1 = 1
+            n2 = num_points - n1
+
+            psi_a_log, obs_xy, all_obs_val_log, _, _ = SamplingStrategies.smart_two_pass(
+                enkf=enkf,
+                psi_f=psi_f_log,
+                conds_preds=conds_log,
+                true_field=psi_t_log,
+                n1=n1,
+                n2=n2,
+                phase1_method=sample_params.get('phase1_method', 'two_stage'),
+                min_dist_p2=sample_params.get('min_dist_p2', 22),
+                under_correct_alpha=sample_params.get('under_correct_alpha', 1.5),
+                use_localization=sample_params.get('use_localization', use_localization),
+                loc_radius_pixobs=sample_params.get('loc_radius_pixobs', 35.0),
+                loc_radius_obsobs=sample_params.get('loc_radius_obsobs', 40.0),
+                seed=42,
+                verbose=sample_params.get('verbose', False),
+            )
+
+            obs_value_log = np.asarray(all_obs_val_log)
+            obs_value_ppm = DataLoader.log2ppm(obs_value_log)
+        else:
+            obs_xy, obs_value_ppm = SamplingStrategies.generate(psi_t_ppm, psi_f_ppm, num_points=num_points,
+                                                  seed=42, method=sample_method,
+                                                  ens_preds_ppm=conds_preds,
+                                                  **sample_params
+                                                  )
+            d_obs_log = np.log1p(np.maximum(obs_value_ppm, 0))  # avoid log(0)
+            if use_localization:
+                psi_a_log = enkf.enkf_localization(psi_f_log, conds_log, obs_xy, d_obs_log,
+                                               loc_radius_pixobs=35.0,
+                                               loc_radius_obsobs=30.0)
+            else:
+                psi_a_log = enkf.standard_enkf(psi_f_log, conds_log, obs_xy, d_obs_log)
+
+            # 计算innovation，判断是否需要同化
+            obs_prior_at_obs = np.log1p(np.maximum(
+                ObservationModel.observation_operator_H(psi_f_ppm, obs_xy), 0
+            ))
+            obs_innovation = np.mean(np.abs(obs_prior_at_obs - d_obs_log))
+            threshold = config.get('innovation_threshold', 0.05)
+
+            if obs_innovation < threshold:
+                psi_a_log = psi_f_log
+            else:
+                psi_a_log = enkf.enkf_localization(
+                    psi_f_log, conds_log, obs_xy, d_obs_log,
+                    loc_radius_pixobs=35.0,
+                    loc_radius_obsobs=30.0
+                )
+            obs_value_log = np.log1p(np.maximum(obs_value_ppm, 0))  # avoid log(0)
+
+        psi_a_ppm = DataLoader.log2ppm(psi_a_log)
+        
+        # 计算指标
+        metrics_log = PrintMetrics.print_metrics(
+            i=i,
+            wind_speed=meta['wind_speed'],
+            wind_direction=meta['wind_direction'],
+            sc=meta['sc'],
+            source_number=meta['source_number'],
+            true_field=psi_t_log,
+            pred_field=psi_f_log,
+            analysis=psi_a_log,
+            obs_xy=obs_xy,
+            metrics_save_flag=True,
+            metrics_print_flag=False
+        )
+        all_metrics_log.append(metrics_log)
+        metrics_ppm = PrintMetrics.print_metrics(
+            i=i,
+            wind_speed=meta['wind_speed'],
+            wind_direction=meta['wind_direction'],
+            sc=meta['sc'],
+            source_number=meta['source_number'],
+            true_field=psi_t_ppm,
+            pred_field=psi_f_ppm,
+            analysis=psi_a_ppm,
+            obs_xy=obs_xy,
+            metrics_save_flag=True,
+            metrics_print_flag=False
+        )
+        all_metrics_ppm.append(metrics_ppm)
+        all_fields.append({
+            "idx": i,
+            "trues_log": psi_t_log,
+            "preds_log": psi_f_log,
+            "analysis_log": psi_a_log,
+            "trues_ppm": psi_t_ppm,
+            "preds_ppm": psi_f_ppm,
+            "analysis_ppm": psi_a_ppm,
+            "obs_xy": obs_xy,
+            "obs_value_log": obs_value_log,
+            "obs_value_ppm": obs_value_ppm,
+        })
+    data_paths = f'./dataset/assim_conds/{data_path}'
+    if not os.path.exists(data_paths):
+        os.makedirs(data_paths)
+    if save_fields_flag:
+        np.savez_compressed(f'./dataset/assim_conds/{data_path}/fields_n{num_points}_{sample_method}_obs{int(obs_std_scale*100)}_damping{damping}.npz',
+                            all_fields=all_fields)
+    all_metrics_df_log = pd.DataFrame(all_metrics_log)
+    all_metrics_df_ppm = pd.DataFrame(all_metrics_ppm)
+    if save_metrics_flag:
+        all_metrics_df_ppm.to_csv(f'./dataset/assim_conds/{data_path}/assimi_ppm_n{num_points}_{sample_method}_obs{int(obs_std_scale*100)}_damping{damping}.csv', index=False)
+        all_metrics_df_log.to_csv(f'./dataset/assim_conds/{data_path}/assimi_log_n{num_points}_{sample_method}_obs{int(obs_std_scale*100)}_damping{damping}.csv', index=False)
+    print("\n=== 平均指标提升 ===")
+    metrics_list = ['r2', 'w_r2_plume', 'r2_plume','mse', 'mae']
+    for metric in metrics_list:
+        before_mean = all_metrics_df_log[f"{metric}_before"].mean()
+        after_mean  = all_metrics_df_log[f"{metric}_after"].mean()
+        delta = after_mean - before_mean
+        print(f'{metric.upper()}: before={before_mean:.4f}, after={after_mean:.4f}, delta={delta:.4f}')
+        before_mean_ppm = all_metrics_df_ppm[f"{metric}_before"].mean()
+        after_mean_ppm  = all_metrics_df_ppm[f"{metric}_after"].mean()
+        delta_ppm = after_mean_ppm - before_mean_ppm
+        print(f'PPM {metric.upper()}: before={before_mean_ppm:.4f}, after={after_mean_ppm:.4f}, delta={delta_ppm:.4f}')