python求解势函数和流函数的简单办法

木叶村洗剪吹

参考沙颖凯 2014-3-2的帖子 https://bbs.06climate.com/forum. ... =%CA%C6%BA%AF%CA%FD
lz把matlab代码转化成了经常使用的python，发现大佬用的差分格式是五点的空间差分，算法是松弛迭代法，求解泊松方程得到得势函数，在求解时，将边值设置成0视为合理。
虽然改进后得迭代法已经比原来快了一倍，但还是很慢啊……有木有
在等待了N个小时之后，lz意识到泊松方程是线性的，且我们用的差分格式决定了我们的线性系数矩阵是稀疏的，lz尝试了用python线性求解器，非常快速地解决了这个问题。
这里省略了散度场的求解，将散度场作为输入量，求解线性方程直接得到势函数
同理，在求解流函数时，只要把散度场替换成涡度场，求解线性方程也能直接得到流函数
实现了流函数势函数自由~
————————————————————————————————————
欢迎交流讨论~~~

def dx_atmos(longitude, latitude):
    R = 6378137.0                                  # 赤道半径，单位为米
    dtheta = np.gradient(longitude,axis=1)                    # 计算经度的梯度
    dtheta=np.array(dtheta)
    dx = dtheta  * np.cos(np.radians(latitude)) * (np.pi / 180) * R  # 转换为米
    return dx
def dy_atmos(latitude):
    R = 6356752.3                          # 经圈半径，单位为米
    dtheta = np.gradient(latitude,axis=0)             # 计算纬度的梯度
    dtheta = np.array(dtheta)
    dy = dtheta * (np.pi / 180) * R           # 转换为米
    return dy

def grad_atmos(longitude, latitude, H):
    dx = dx_atmos(longitude, latitude)        # 计算经向网格间距
    dy = dy_atmos(latitude)                    # 计算纬向网格间距

    # 计算场 H 的梯度
    grad_y, grad_x = np.gradient(H)           # 计算 H 在经度和纬度方向的梯度

    # 归一化梯度，计算 H 在经度和纬度方向的实际梯度
    grad_x = grad_x / dx                       # 将经度方向的梯度进行归一化
    grad_y = grad_y / dy                       # 将纬度方向的梯度进行归一化

    return grad_x, grad_y

def index(i, j, ny):
    return i * ny + j

def chi_potential(latitude, longitude, divh):
    """
    ps：为了让结果更可靠
    计算速度势和辐散风分量。
    返回：
    chi : np.ndarray
        速度势数组。
    Uchi : np.ndarray
        辐散风 U 分量。
    Vchi : np.ndarray
        辐散风 V 分量。
    """   
    M  = latitude.shape[0]                      # 获取经度和纬度的维度
    N  = latitude.shape[1]
    #divh = divh_atmos(longitude, latitude, u, v)  # 计算水平散度

    dx2 = dx_atmos(longitude, latitude) ** 2   # 计算纬向梯度的平方
    dy2 = dy_atmos(latitude) ** 2               # 计算经向梯度的平方
    #print(dx2.shape)
    #print(f"dy2:{dy2.mean()}")
    #print(f"div:{divh.mean()}")
    #divh*=1e7

    # 2. 构建离散化矩阵
    n_interior = (M - 2) * (N - 2)  # 内部网格点数
    A = sp.lil_matrix((n_interior, n_interior))  # 创建稀疏矩阵
    #print(f"A.shape:{A.shape}")
    # 填充矩阵 A
    for i in range(1,M - 1):
        for j in range(1, N - 1):
            idx = index(i - 1 , j - 1, N - 2)
            #print(f"i:{i} j:{j}")
            A[idx, idx] = 2*((1/dx2[i,j])+(1/dy2[i,j]))  # 对角线
            if i>1:
                A[idx, index(i - 1 - 1, j - 1, N-2)] = -1/dy2[i-1,j]   # 上侧点
            if i<M-2:
                A[idx, index(i + 1 - 1, j - 1, N-2)] = -1/dy2[i+1,j]   # 下侧点
            if j>1:
                A[idx, index(i - 1, j - 1 - 1, N-2)] = -1/dx2[i,j-1]   # 左侧点
            if j<N-2:
                A[idx, index(i - 1, j + 1 - 1, N-2)] = -1/dx2[i,j+1]   # 右侧点

    # 3. 定义右边项
    b = np.zeros(n_interior)

    # 这里定义一个示例源项
    for i in range(1, M - 1):
        for j in range(1, N - 1):
            b[index(i - 1, j - 1 , N - 2)] = divh[i,j]
    # 4. 使用 pyAMG 求解线性系统
    ml = pyamg.ruge_stuben_solver(A)  # 创建多重网格求解器
    chi_interior = ml.solve(b, tol=1e-10)  # 求解

    # 结果转换为二维数组格式用于可视化
    chi = np.zeros((M, N))  # 初始化整个解向量
    chi[1:-1, 1:-1] = chi_interior.reshape((M - 2, N - 2))  # 将内部解放入完整解中

    #print(f"chi:{chi.mean()}")
    #print(f"chi[10,10]:{chi[10,10]}")

    # 计算辐散风分量
    DchiDx, DchiDy = grad_atmos(longitude, latitude, chi)
    Uchi = -DchiDx  # 辐散风的 U 分量
    Vchi = -DchiDy  # 辐散风的 V 分量

    # 将边界值设为 NaN
    chi[0, :] = np.nan; Uchi[0, :] = np.nan; Vchi[0, :] = np.nan
    chi[-1, :] = np.nan; Uchi[-1, :] = np.nan; Vchi[-1, :] = np.nan
    chi[:, 0] = np.nan; Uchi[:, 0] = np.nan; Vchi[:, 0] = np.nan
    chi[:, -1] = np.nan; Uchi[:, -1] = np.nan; Vchi[:, -1] = np.nan

    return chi, Uchi, Vchi

SunJiaming

楼主可以参考下windspharm这个库，也很好用

木叶村洗剪吹

SunJiaming 发表于 2024-11-26 22:08
楼主可以参考下windspharm这个库，也很好用

windspharm真的有人装上了吗，听说能满足各种气象上计算需求，就是安装了两次都卡在import某个包上了（叹气

一大碗年糕

windspharm+1哈哈

SunJiaming

木叶村洗剪吹 发表于 2024-11-26 22:18
windspharm真的有人装上了吗，听说能满足各种气象上计算需求，就是安装了两次都卡在import某个包上了（叹 ...

我是在conda里建了个虚拟环境单独装的，这个库貌似对python的版本有要求

GXYYYYY

木叶村洗剪吹 发表于 2024-11-26 22:18
windspharm真的有人装上了吗，听说能满足各种气象上计算需求，就是安装了两次都卡在import某个包上了（叹 ...

windspharm对python和numpy的版本有要求，可以参考这篇帖子下的回复http://bbs.06climate.com/forum.p ... ighlight=windspharm

木叶村洗剪吹

GXYYYYY 发表于 2024-11-27 10:29
windspharm对python和numpy的版本有要求，可以参考这篇帖子下的回复http://bbs.06climate.com/for ...

十分感谢！！

linktv

木叶村洗剪吹 发表于 2024-11-27 17:07
十分感谢！！

如果其他方法都没解决，记得把python版本降到3.9及以下，windspharm没法在3.10以后版本上运行，就是装上了也没用

木叶村洗剪吹

linktv 发表于 2024-12-2 11:16
如果其他方法都没解决，记得把python版本降到3.9及以下，windspharm没法在3.10以后版本上运行，就是装上 ...

哈哈哈在上面的帖子里眼熟大佬了（！谢谢！已经装上了！