将AWK多分支条件计算逻辑转换为Python（Pandas）实现

原始数据 #

data = {'f_geno': ["AA", "AA", "AA", "BB", "BB", "BB", "AB", "AB", "AB"],
        'ch_geno': ["AA", "AB", "BB", "AA", "AB", "BB", "AA", "BB", "AB"],
        'freq_A': [0.50, 0.46, 0.49, 0.57, 0.55, 0.44, 0.37, 0.66, 0.46],
        'freq_B': [0.50, 0.54, 0.51, 0.43, 0.45, 0.56, 0.63, 0.34, 0.54]
        }

原始AWK脚本 #

BEGIN {
    FS = OFS = ","
}
{
    if ($1 == "AA" && $2 == "AA") {
        $5 = (1 / $3)
    } else if ($1 == "AA" && $2 == "AB") {
        $5 = (0.5 / $3)
    } else if ($1 == "AA" && $2 == "BB") {
        $5 = (0.001)
    } else if ($1 == "BB" && $2 == "AA") {
        $5 = (0.001)
    } else if ($1 == "BB" && $2 == "AB") {
        $5 = (0.5 / $4)
    } else if ($1 == "BB" && $2 == "BB") {
        $5 = (1 / $4) 
    } else if ($1 == "AB" && $2 == "AA") {
        $5 = (0.5 / $3)
    } else if ($1 == "AB" && $2 == "BB") {
        $5 = (0.5 / $4)  
    } else {
        $5 = (($3 + $4) / (4 * $3 * $4))
    }
    
    print 
}

Pandas实现方案 #

使用numpy.select批量处理条件计算，效率远高于逐行apply，适合大数据量场景：

import pandas as pd
import numpy as np

# 创建DataFrame
df = pd.DataFrame(data)

# 定义条件列表与对应计算结果
conditions = [
    (df['f_geno'] == 'AA') & (df['ch_geno'] == 'AA'),
    (df['f_geno'] == 'AA') & (df['ch_geno'] == 'AB'),
    (df['f_geno'] == 'AA') & (df['ch_geno'] == 'BB'),
    (df['f_geno'] == 'BB') & (df['ch_geno'] == 'AA'),
    (df['f_geno'] == 'BB') & (df['ch_geno'] == 'AB'),
    (df['f_geno'] == 'BB') & (df['ch_geno'] == 'BB'),
    (df['f_geno'] == 'AB') & (df['ch_geno'] == 'AA'),
    (df['f_geno'] == 'AB') & (df['ch_geno'] == 'BB')
]

results = [
    1 / df['freq_A'],
    0.5 / df['freq_A'],
    0.001,
    0.001,
    0.5 / df['freq_B'],
    1 / df['freq_B'],
    0.5 / df['freq_A'],
    0.5 / df['freq_B']
]

# 应用条件计算，default对应AWK的else分支
df['calc_result'] = np.select(conditions, results, default=(df['freq_A'] + df['freq_B']) / (4 * df['freq_A'] * df['freq_B']))

# 输出结果
print(df)

运行结果 #

得到的DataFrame包含新计算列calc_result，与AWK脚本输出完全一致：

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