如何在Python的for循环中从函数获取数组（Matlab转Python）

我是Python新手，希望编写一个接收n个输入、返回m个输出的函数。我可以在Matlab中实现这样的函数，例如：

[a,b,c,d] = my_function(x,y,z,h)

其中(x,y,z,h)是多维数组。

具体来说，我定义了以下数组：

A=np.zeros((3,150,150))
B=np.zeros((3,150,150))
C=np.zeros((3,150,150))
D=np.zeros((3,150,150))
p=zero(3,1)  # 这里应该是np.zeros((3,1))

同时有如下for循环：

x.shape=(1540,1838)
step=10
jj=0
for j in range(0,1838,step):
     jj=jj+1 #tot= 184

     ii=0
     for i in range(0,1540,step):
         ii=ii+1 #total 154
         pos=x(i,j)  # 这里Python里应该用x[i,j]

我希望在该循环内执行自定义函数，得到类似如下的输出：

a=A(:,ii,jj)
b=B(:,:,ii,jj)
c=C(:,:,ii,jj)
d=D(:,ii,jj)

请问如何用Python风格实现这个功能？

一、多输入多输出函数的Python实现 #

Python里做这事很简单，不用复杂语法，直接在函数末尾用return返回多个值就行，调用时直接把返回值解包到多个变量，逻辑和Matlab差不多：

import numpy as np

def my_function(x, y, z, h):
    # 这里写你的计算逻辑，根据输入生成需要的a、b、c、d
    a = np.random.rand(3)  # 示例输出，对应你要的维度
    b = np.random.rand(3, 150)
    c = np.random.rand(3, 150)
    d = np.random.rand(3)
    return a, b, c, d

# 调用函数，直接接收多个返回值
a, b, c, d = my_function(input_x, input_y, input_z, input_h)

二、修正代码里的基础错误 #

1. numpy数组创建：p=zero(3,1)是错的，得改成p = np.zeros((3, 1))——要调用numpy的zeros函数，而且维度得用元组传。
2. 数组索引：Python里numpy数组用方括号[]索引，不是Matlab的圆括号()，所以pos=x(i,j)要改成pos = x[i, j]。

三、更Pythonic的循环写法 #

手动加jj、ii计数器太麻烦，用enumerate函数直接生成带计数的循环变量就行，还能指定start=1让计数从1开始，和你原来的逻辑保持一致：

x = np.random.rand(1540, 1838)  # 生成示例x数组
step = 10

# 外层循环：jj从1开始计数
for jj, j in enumerate(range(0, 1838, step), start=1):
    # 内层循环：ii从1开始计数
    for ii, i in enumerate(range(0, 1540, step), start=1):
        pos = x[i, j]
        # 在这里调用自定义函数，传入需要的参数
        a, b, c, d = my_function(pos, ...)
        
        # 转换Matlab的索引到Python：注意Python是0开头的索引，所以要减1
        # Matlab的A(:,ii,jj) → Python的A[:, ii-1, jj-1]
        a_arr = A[:, ii-1, jj-1]
        # 你写的B(:,:,ii,jj)可能是笔误，因为B定义的是三维数组(3,150,150)，按需求调整索引
        b_arr = B[:, ii-1, jj-1]  # 取所有通道、第ii行、第jj列
        c_arr = C[:, ii-1, jj-1]
        d_arr = D[:, ii-1, jj-1]

四、重点：Matlab和Numpy的索引差异 #

Matlab的索引是从1开始的，而Python的numpy数组是从0开始的，所以你原来的ii、jj从1计数，访问numpy数组时必须减1，才能对应到正确的位置。比如Matlab里的A(:,1,1)，对应numpy里的A[:,0,0]。

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