Python中用循环传递前一条件结果并保存输出至CSV

我想用Python的while循环替代下面的多分支if语句，但不知道怎么正确实现：

if t == 0:
    za = ds.isel(time=t)
    abc = get_grids_inside_rad(127.2, 15.7)
    print(abc)
if t == 1:
    za = ds.isel(time=t)
    ghi = get_grids_inside_rad(abc.lon,abc.lat)
    print(ghi)
if t == 2:
    za = ds.isel(time=t)
    jkl = get_grids_inside_rad(ghi.lon,ghi.lat)
    print(jkl)
if t == 3:
    za = ds.isel(time=t)
    mno = get_grids_inside_rad(jkl.lon,jkl.lat)
    print(mno)
if t == 4:
    za = ds.isel(time=t)
    pqr = get_grids_inside_rad(mno.lon,ghi.lat)
    print(pqr)

我的需求：

• 将前一个时间步的输出作为下一个循环的输入；
• 将最终所有输出保存到CSV文件。

get_grids_inside_rad(127.2, 15.7)是一个读取经纬度的函数。当前只能输出第一个时间步的结果，无法打印第二个及之后的内容，我期望最终输出包含4组lon、lat、msl数据（len(ds.time)等于4）。

以下是完整代码：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import xarray as xr
import cartopy.crs as ccrs
import cartopy.mpl.ticker as cticker
import matplotlib.ticker as mticker
from cartopy.util import add_cyclic_point
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import metpy.calc as mpcalc
import csv
import sys


ds = xr.open_dataset('prmsl_koppu_oct16.nc')

def haversine(lon1, lat1, lon2, lat2):
    # convert decimal degrees to radians
    lon1 = np.deg2rad(lon1)
    lon2 = np.deg2rad(lon2)
    lat1 = np.deg2rad(lat1)
    lat2 = np.deg2rad(lat2)

    # haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    r = 6371
    return c * r
    
def get_grids_inside_rad(ctr_lon, ctr_lat):
    lon = np.arange(0, 360, 1.25)
    lat = np.arange(-90, 90, 1.25)
                        # get coordinates of all points on the grid
    grid_lon, grid_lat = np.meshgrid(lon, lat)
    dists_in_km = haversine(grid_lon, grid_lat, ctr_lon, ctr_lat)
    dists_in_deg = dists_in_km / 111
                         # find nearby points
    thr = 2.0
    coordinates = []
    for i in range(grid_lon.shape[0]):
      for j in range(grid_lon.shape[1]):
            this_lon = grid_lon[i, j]
            this_lat = grid_lat[i, j]
            dist = dists_in_deg[i, j]
            if dist <= thr:
                df = pd.DataFrame({"lon":[this_lon],"lat":[this_lat],"dist":[dist]})
                coordinates.append(df)
                        
    xyz = pd.concat(coordinates)
    xyz.reset_index(drop=True, inplace=True)
    xyz.index = range(1, len(xyz)+1, 1)
      
    lat = list(xyz.iloc[:, 1])
    lon = list(xyz.iloc[:, 0])
      
    filtered_rad=za.sel(lat=lat, lon=lon, method='nearest')
      
    yyy=filtered_rad.sortby('lon')
    df = yyy.to_dataframe().reset_index().drop_duplicates()
    df = df[df["msl"]==df["msl"].min()]
    return df
    

t=0
while t < len(ds.time):
     if t == 0:
         za = ds.isel(time=t)
         abc = get_grids_inside_rad(127.2, 15.7)
         print(abc)
     if t == 1:
         za = ds.isel(time=t)
         ghi = get_grids_inside_rad(abc.lon,abc.lat)
         print(ghi)
     else: 
         break
 t += 1

ds结构说明：包含time、lat、lon维度，变量为msl（海平面气压）。

原代码的问题 #

1. else: break会导致t=0时执行完第一个if后直接跳出循环，永远无法进入后续时间步的逻辑；
2. 硬编码t==0、t==1的判断无法适配4个时间步的循环需求；
3. 没有收集所有时间步的结果，无法实现保存到CSV的需求；
4. get_grids_inside_rad依赖全局变量za，代码耦合性高。

修正后的代码 #

核心思路：用变量保存上一步的经纬度结果，循环中依次传入下一个时间步，同时收集所有结果最后写入CSV。

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import xarray as xr
import cartopy.crs as ccrs
import cartopy.mpl.ticker as cticker
import matplotlib.ticker as mticker
from cartopy.util import add_cyclic_point
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import metpy.calc as mpcalc
import csv
import sys


ds = xr.open_dataset('prmsl_koppu_oct16.nc')

def haversine(lon1, lat1, lon2, lat2):
    # 角度转弧度
    lon1 = np.deg2rad(lon1)
    lon2 = np.deg2rad(lon2)
    lat1 = np.deg2rad(lat1)
    lat2 = np.deg2rad(lat2)

    # 半正矢公式计算距离
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    r = 6371
    return c * r
    
def get_grids_inside_rad(ctr_lon, ctr_lat, za):
    # 将za作为参数传入，避免依赖全局变量
    lon = np.arange(0, 360, 1.25)
    lat = np.arange(-90, 90, 1.25)
    # 生成网格坐标
    grid_lon, grid_lat = np.meshgrid(lon, lat)
    dists_in_km = haversine(grid_lon, grid_lat, ctr_lon, ctr_lat)
    dists_in_deg = dists_in_km / 111
    # 筛选阈值内的点
    thr = 2.0
    coordinates = []
    for i in range(grid_lon.shape[0]):
      for j in range(grid_lon.shape[1]):
            this_lon = grid_lon[i, j]
            this_lat = grid_lat[i, j]
            dist = dists_in_deg[i, j]
            if dist <= thr:
                df = pd.DataFrame({"lon":[this_lon],"lat":[this_lat],"dist":[dist]})
                coordinates.append(df)
                        
    xyz = pd.concat(coordinates)
    xyz.reset_index(drop=True, inplace=True)
    xyz.index = range(1, len(xyz)+1, 1)
      
    lat_list = list(xyz.iloc[:, 1])
    lon_list = list(xyz.iloc[:, 0])
      
    filtered_rad=za.sel(lat=lat_list, lon=lon_list, method='nearest')
      
    yyy=filtered_rad.sortby('lon')
    df = yyy.to_dataframe().reset_index().drop_duplicates()
    df = df[df["msl"]==df["msl"].min()]
    # 添加时间步标识列
    df['time_step'] = t
    return df


# 初始化结果收集列表
all_results = []
t = 0
# 初始经纬度
current_lon, current_lat = 127.2, 15.7

while t < len(ds.time):
    za = ds.isel(time=t)
    # 调用函数，传入当前经纬度和当前时间步的数据
    step_result = get_grids_inside_rad(current_lon, current_lat, za)
    print(step_result)
    all_results.append(step_result)
    
    # 更新下一轮循环的经纬度
    current_lon = step_result['lon'].iloc[0]
    current_lat = step_result['lat'].iloc[0]
    
    t += 1

# 合并所有结果并保存到CSV
final_df = pd.concat(all_results, ignore_index=True)
final_df.to_csv('time_step_results.csv', index=False)
print("所有结果已保存到time_step_results.csv")

关键修改点 #

1. 移除else: break语句，确保循环能遍历所有时间步；
2. 将za作为参数传入get_grids_inside_rad，消除全局变量依赖；
3. 用current_lon和current_lat保存上一步的结果，作为下一次循环的输入；
4. 新增all_results列表收集每个时间步的数据，最后合并成DataFrame保存为CSV；
5. 给每个时间步的结果添加time_step列，方便区分不同时间的数据。

内容的提问来源于stack exchange，提问作者Lyndz