import osimport reimport globimport numpy as npimport rasteriofrom scipy.stats import t, ffrom tqdm import tqdmbase_dir = r"E:\非线性Granger因果"ndvi_dir = os.path.join(base_dir, "NDVI")pre_dir = os.path.join(base_dir, "PRE")tmp_dir = os.path.join(base_dir, "TMP")rad_dir = os.path.join(base_dir, "RAD")out_dir = os.path.join(base_dir, "相关分析结果")pearson_dir = os.path.join(out_dir, "皮尔逊相关")partial_dir = os.path.join(out_dir, "偏相关")multiple_dir = os.path.join(out_dir, "复相关")for d in [pearson_dir, partial_dir, multiple_dir]:    os.makedirs(d, exist_ok=True)start_year = 2010end_year = 2024alpha = 0.05min_valid_n = 30nodata_value = -9999.0def get_ym(path):    name = os.path.basename(path)    m = re.search(r"(\d{4})_(\d{2})", name)    if m is None:        return None    y = int(m.group(1))    mo = int(m.group(2))    if start_year <= y <= end_year and 1 <= mo <= 12:        return f"{y:04d}_{mo:02d}"    return Nonedef collect_files(folder):    files = glob.glob(os.path.join(folder, "*.tif"))    d = {}    for fp in files:        ym = get_ym(fp)        if ym is not None:            d[ym] = fp    return dndvi_dict = collect_files(ndvi_dir)pre_dict = collect_files(pre_dir)tmp_dict = collect_files(tmp_dir)rad_dict = collect_files(rad_dir)common_ym = sorted(set(ndvi_dict) & set(pre_dict) & set(tmp_dict) & set(rad_dict))if len(common_ym) == 0:    raise RuntimeError("没有找到四类数据共同存在的年月文件")print(f"共同年月数量：{len(common_ym)}")print(f"起始年月：{common_ym[0]}")print(f"结束年月：{common_ym[-1]}")ndvi_files = [ndvi_dict[x] for x in common_ym]pre_files = [pre_dict[x] for x in common_ym]tmp_files = [tmp_dict[x] for x in common_ym]rad_files = [rad_dict[x] for x in common_ym]def read_stack(files):    arrs = []    profile = None    src_nodata = None    for fp in tqdm(files):        with rasterio.open(fp) as src:            a = src.read(1).astype(np.float32)            if profile is None:                profile = src.profile.copy()                src_nodata = src.nodata            if src.nodata is not None:                a[a == src.nodata] = np.nan            a[~np.isfinite(a)] = np.nan            arrs.append(a)    return np.stack(arrs, axis=0), profile, src_nodataprint("读取NDVI")ndvi, profile, src_nodata = read_stack(ndvi_files)print("读取降水")pre, _, _ = read_stack(pre_files)print("读取气温")tmp, _, _ = read_stack(tmp_files)print("读取辐射")rad, _, _ = read_stack(rad_files)n_time, rows, cols = ndvi.shaper_ndvi_pre = np.full((rows, cols), np.nan, dtype=np.float32)p_ndvi_pre = np.full((rows, cols), np.nan, dtype=np.float32)sig_ndvi_pre = np.zeros((rows, cols), dtype=np.uint8)r_ndvi_tmp = np.full((rows, cols), np.nan, dtype=np.float32)p_ndvi_tmp = np.full((rows, cols), np.nan, dtype=np.float32)sig_ndvi_tmp = np.zeros((rows, cols), dtype=np.uint8)r_ndvi_rad = np.full((rows, cols), np.nan, dtype=np.float32)p_ndvi_rad = np.full((rows, cols), np.nan, dtype=np.float32)sig_ndvi_rad = np.zeros((rows, cols), dtype=np.uint8)pr_ndvi_pre = np.full((rows, cols), np.nan, dtype=np.float32)pp_ndvi_pre = np.full((rows, cols), np.nan, dtype=np.float32)psig_ndvi_pre = np.zeros((rows, cols), dtype=np.uint8)pr_ndvi_tmp = np.full((rows, cols), np.nan, dtype=np.float32)pp_ndvi_tmp = np.full((rows, cols), np.nan, dtype=np.float32)psig_ndvi_tmp = np.zeros((rows, cols), dtype=np.uint8)pr_ndvi_rad = np.full((rows, cols), np.nan, dtype=np.float32)pp_ndvi_rad = np.full((rows, cols), np.nan, dtype=np.float32)psig_ndvi_rad = np.zeros((rows, cols), dtype=np.uint8)multiple_R = np.full((rows, cols), np.nan, dtype=np.float32)multiple_R2 = np.full((rows, cols), np.nan, dtype=np.float32)multiple_p = np.full((rows, cols), np.nan, dtype=np.float32)multiple_sig = np.zeros((rows, cols), dtype=np.uint8)def pearson_pvalue(r, n):    if n <= 2 or not np.isfinite(r) or abs(r) >= 1:        if np.isfinite(r) and abs(r) == 1 and n > 2:            return 0.0        return np.nan    tv = r * np.sqrt((n - 2) / max(1e-12, 1 - r * r))    return 2 * t.sf(abs(tv), n - 2)def partial_pvalue(r, n, k):    df = n - k - 2    if df <= 0 or not np.isfinite(r) or abs(r) >= 1:        if np.isfinite(r) and abs(r) == 1 and df > 0:            return 0.0        return np.nan    tv = r * np.sqrt(df / max(1e-12, 1 - r * r))    return 2 * t.sf(abs(tv), df)def safe_corr(a, b):    if np.std(a) == 0 or np.std(b) == 0:        return np.nan    return np.corrcoef(a, b)[0, 1]def save_float(arr, path):    p = profile.copy()    p.update(dtype=rasterio.float32, count=1, nodata=nodata_value, compress="lzw")    out = arr.astype(np.float32)    out[~np.isfinite(out)] = nodata_value    with rasterio.open(path, "w", **p) as dst:        dst.write(out, 1)def save_uint8(arr, path):    p = profile.copy()    p.update(dtype=rasterio.uint8, count=1, nodata=255, compress="lzw")    out = arr.astype(np.uint8)    with rasterio.open(path, "w", **p) as dst:        dst.write(out, 1)print("开始逐像元计算")for i in tqdm(range(rows)):    for j in range(cols):        y0 = ndvi[:, i, j]        x10 = pre[:, i, j]        x20 = tmp[:, i, j]        x30 = rad[:, i, j]        mask = np.isfinite(y0) & np.isfinite(x10) & np.isfinite(x20) & np.isfinite(x30)        n = int(mask.sum())        if n < min_valid_n:            continue        y = y0[mask]        x1 = x10[mask]        x2 = x20[mask]        x3 = x30[mask]        if np.std(y) == 0 or np.std(x1) == 0 or np.std(x2) == 0 or np.std(x3) == 0:            continue        r01 = safe_corr(y, x1)        r02 = safe_corr(y, x2)        r03 = safe_corr(y, x3)        r12 = safe_corr(x1, x2)        r13 = safe_corr(x1, x3)        r23 = safe_corr(x2, x3)        r_ndvi_pre[i, j] = r01        r_ndvi_tmp[i, j] = r02        r_ndvi_rad[i, j] = r03        p01 = pearson_pvalue(r01, n)        p02 = pearson_pvalue(r02, n)        p03 = pearson_pvalue(r03, n)        p_ndvi_pre[i, j] = p01        p_ndvi_tmp[i, j] = p02        p_ndvi_rad[i, j] = p03        if np.isfinite(p01) and p01 < alpha:            sig_ndvi_pre[i, j] = 1        if np.isfinite(p02) and p02 < alpha:            sig_ndvi_tmp[i, j] = 1        if np.isfinite(p03) and p03 < alpha:            sig_ndvi_rad[i, j] = 1        corr = np.array([            [1.0, r01, r02, r03],            [r01, 1.0, r12, r13],            [r02, r12, 1.0, r23],            [r03, r13, r23, 1.0]        ], dtype=np.float64)        try:            inv_corr = np.linalg.inv(corr)            pr01 = -inv_corr[0, 1] / np.sqrt(inv_corr[0, 0] * inv_corr[1, 1])            pr02 = -inv_corr[0, 2] / np.sqrt(inv_corr[0, 0] * inv_corr[2, 2])            pr03 = -inv_corr[0, 3] / np.sqrt(inv_corr[0, 0] * inv_corr[3, 3])            pr_ndvi_pre[i, j] = pr01            pr_ndvi_tmp[i, j] = pr02            pr_ndvi_rad[i, j] = pr03            pp01 = partial_pvalue(pr01, n, 2)            pp02 = partial_pvalue(pr02, n, 2)            pp03 = partial_pvalue(pr03, n, 2)            pp_ndvi_pre[i, j] = pp01            pp_ndvi_tmp[i, j] = pp02            pp_ndvi_rad[i, j] = pp03            if np.isfinite(pp01) and pp01 < alpha:                psig_ndvi_pre[i, j] = 1            if np.isfinite(pp02) and pp02 < alpha:                psig_ndvi_tmp[i, j] = 1            if np.isfinite(pp03) and pp03 < alpha:                psig_ndvi_rad[i, j] = 1        except Exception:            pass        try:            r_yx = np.array([r01, r02, r03], dtype=np.float64)            r_xx = np.array([                [1.0, r12, r13],                [r12, 1.0, r23],                [r13, r23, 1.0]            ], dtype=np.float64)            r2 = float(r_yx @ np.linalg.inv(r_xx) @ r_yx.T)            r2 = min(max(r2, 0.0), 1.0)            rv = np.sqrt(r2)            multiple_R[i, j] = rv            multiple_R2[i, j] = r2            k = 3            df1 = k            df2 = n - k - 1            if df2 > 0 and r2 < 1:                fv = (r2 / k) / ((1 - r2) / df2)                pv = 1 - f.cdf(fv, df1, df2)            elif df2 > 0 and r2 == 1:                pv = 0.0            else:                pv = np.nan            multiple_p[i, j] = pv            if np.isfinite(pv) and pv < alpha:                multiple_sig[i, j] = 1        except Exception:            passprint("保存皮尔逊相关结果")save_float(r_ndvi_pre, os.path.join(pearson_dir, "NDVI_降水_皮尔逊相关系数.tif"))save_float(p_ndvi_pre, os.path.join(pearson_dir, "NDVI_降水_皮尔逊相关P值.tif"))save_uint8(sig_ndvi_pre, os.path.join(pearson_dir, "NDVI_降水_皮尔逊显著性.tif"))save_float(r_ndvi_tmp, os.path.join(pearson_dir, "NDVI_气温_皮尔逊相关系数.tif"))save_float(p_ndvi_tmp, os.path.join(pearson_dir, "NDVI_气温_皮尔逊相关P值.tif"))save_uint8(sig_ndvi_tmp, os.path.join(pearson_dir, "NDVI_气温_皮尔逊显著性.tif"))save_float(r_ndvi_rad, os.path.join(pearson_dir, "NDVI_辐射_皮尔逊相关系数.tif"))save_float(p_ndvi_rad, os.path.join(pearson_dir, "NDVI_辐射_皮尔逊相关P值.tif"))save_uint8(sig_ndvi_rad, os.path.join(pearson_dir, "NDVI_辐射_皮尔逊显著性.tif"))print("保存偏相关结果")save_float(pr_ndvi_pre, os.path.join(partial_dir, "NDVI_降水_控制气温辐射_偏相关系数.tif"))save_float(pp_ndvi_pre, os.path.join(partial_dir, "NDVI_降水_控制气温辐射_偏相关P值.tif"))save_uint8(psig_ndvi_pre, os.path.join(partial_dir, "NDVI_降水_控制气温辐射_偏相关显著性.tif"))save_float(pr_ndvi_tmp, os.path.join(partial_dir, "NDVI_气温_控制降水辐射_偏相关系数.tif"))save_float(pp_ndvi_tmp, os.path.join(partial_dir, "NDVI_气温_控制降水辐射_偏相关P值.tif"))save_uint8(psig_ndvi_tmp, os.path.join(partial_dir, "NDVI_气温_控制降水辐射_偏相关显著性.tif"))save_float(pr_ndvi_rad, os.path.join(partial_dir, "NDVI_辐射_控制降水气温_偏相关系数.tif"))save_float(pp_ndvi_rad, os.path.join(partial_dir, "NDVI_辐射_控制降水气温_偏相关P值.tif"))save_uint8(psig_ndvi_rad, os.path.join(partial_dir, "NDVI_辐射_控制降水气温_偏相关显著性.tif"))print("保存复相关结果")save_float(multiple_R, os.path.join(multiple_dir, "NDVI_降水气温辐射_复相关系数R.tif"))save_float(multiple_R2, os.path.join(multiple_dir, "NDVI_降水气温辐射_决定系数R2.tif"))save_float(multiple_p, os.path.join(multiple_dir, "NDVI_降水气温辐射_复相关P值.tif"))save_uint8(multiple_sig, os.path.join(multiple_dir, "NDVI_降水气温辐射_复相关显著性.tif"))print("全部完成")print(f"结果输出路径：{out_dir}")