如何考虑不同的输入大小优化代码？

Question

我想以最有效的方式计算蒸汽属性，考虑标量、向量和矩阵作为两个参数的输入选项。我担心的是，我必须使用 if 块相对于输入（标量、向量或矩阵）的大小，使代码变得很长。我是一个简单的机械工程师，对 python 很陌生，任何关于如何优化代码的帮助都非常感谢。这是代码：

from iapws.iapws97 import _Region4
import numpy as np

def h_x(P,x): 
    ''' spec enthalpy in liquid, steam and wet (two-phase flow) regions
    P - pressure in bar
    x - drayness steam fraction [-]
    h - specific heat of wet region returned [kJ/kW]
    '''

    mm = len(np.shape(x))
    if mm == 0:
        h_  = _Region4(P/10,0)['h'] 
        h__ = _Region4(P/10,1)['h'] 
        # return h_ + x * (h__ - h_)
        return h_ + x * (h__ - h_)        
    elif mm == 1:
        return np.array([ _Region4(i/10,0)['h'] + j * ( _Region4(i/10,1)['h'] - _Region4(i/10,0)['h'] ) for i,j in zip(P,x) ])    
    elif mm == 2:
        mmm,nnn = x.shape
        h     = np.ndarray(shape=(mmm,nnn)) #(mm,nn)
        for i in range(mmm):
            for j in range(nnn):
                h_  = _Region4(P[i,j]/10,0)['h']
                h__ = _Region4(P[i,j]/10,1)['h']
                h[i,j]   = h_ + x[i,j] * (h__ - h_)
        return h        
    else:
        print('h_x input must be scalar, vector or 2D matrix!')

# code testing
P = np.array([[.0234,.0193,0.244],[.0244,.0185,0.254]])
x = np.array([[.812,.782,.620],[.912,.882,.820]])
h_x(P,x)

Answer 1

你真的只做一种计算，但用两种不同的方式。您可以使用内置的map 将其拉出并将其应用于函数作为输入的任何内容。如果失败了，那么你有一个单一的（不可迭代的）值，你可以直接应用你的计算。

# Define a dummy func to make code work
def _Region4(a, b):
    return {'h': a + 3 * b}

import numpy as np

def calculate(P, x):
    '''
    Spec enthalpy in liquid, steam and wet (two-phase flow) regions
    P - pressure in bar
    x - drayness steam fraction [-]
    h - specific heat of wet region returned [kJ/kW]
    '''
    h_  = _Region4(P/10,0)['h'] 
    h__ = _Region4(P/10,1)['h'] 
    return h_ + x * (h__ - h_) 

def h_x(P,x):
    shape = np.shape(x)
    dimensions = len(shape)
    # Check for wrong input
    if dimensions > 2:
        raise ValueError('h_x input must be scalar, vector or 2D matrix!')
    # Try a general mapping
    try:
        return np.array(list(map(calculate, P.flat, x.flat))).reshape(shape)
    # if it fails then you got a pair of scalar
    except AttributeError:
        return calculate(P, x)

P = np.array([[.023,.23,.05],[.023,.23,.05]])
x = np.array([[.92,.98,.99],[.92,.98,.99]])
print(h_x(P, x))
"""
Out:
    [[2.7623 2.963  2.975 ]
    [2.7623 2.963  2.975 ]]
"""
print(h_x(3, 4))
"""
Out:
    12.3
"""