注水算法解功率分配问题-Python

问题描述：

分配总功率$P$给$N$个发射机，使总接收通信速率最大。

问题建模为：
$$\begin{array}{rl}& \underset{\mathbf{p}}{\min }-\sum _i^{N}Wlo{g}_2(1+\frac{h_i{p}_i}{{\sigma }^2})\\ & \mathrm{s}.\mathrm{t}.\sum _i^N{p}_i=P\\ & p_i\ge 0\end{array}$$
其中${\sigma }^2$为接收机的噪声功率，$W$为信道带宽，$P$为发射机的总功率，$p_i$为分给第$i$个接收机、信道增益为$h_i$的发射功率。

对等式约束引入一个乘子$v\in \mathbb{R}$，对不等式约束引入乘子$\lambda \in {\mathbb{R}}^N$，得到Lagrange函数$L$:
$$\begin{array}{r}L(\mathbf{p},\lambda ,v)=-\sum _i^{N}Wlo{g}_2(1+\frac{h_i{p}_i}{{\sigma }^2})-\sum _i^N\lambda p_i+v(\sum _i^N{p}_i-P)\end{array}$$

令${\mathbf{p}}^{\ast }$和$({\lambda }^{\ast },v^{\ast })$分别为原问题和对偶问题的某对最优解，则得到如下KKT条件：

$$\begin{array}{rl}& \sum _i^N{p}_i^{\ast }=P\\ & p_i^{\ast }\ge 0,i=1,...,N\\ & {\lambda }_i^{\ast }\ge 0,i=1,...,N\\ & {\lambda }_i^{\ast }{p}_i^{\ast }=0,i=1,...,N\\ & -\frac{W{\log }_{2}e}{{\sigma }^2/h_i+p_i^{\ast }}-{\lambda }_i^{\ast }+v^{\ast }=0i=1,...,N\end{array}$$

最优解之间满足：
$$p_i^{\ast }=\{\begin{array}{ll}\frac{W{\log }_{2}e}{v^{\ast }}-\frac{{\sigma }^2}{h_i}& v^{\ast }<\frac{W{\log }_{2}e}{{\sigma }^2/h_i}\\ 0& v^{\ast }\ge \frac{W{\log }_{2}e}{{\sigma }^2/h_i}\end{array}$$

或者，更简洁地， p i ∗ = max ⁡ { 0 , W log ⁡ 2 e v ∗ − σ 2 h i } p_i^*=\max \{0,\frac{W\log_2e}{v^*}-\frac{\sigma^2}{h_i}\} pi∗​=max{0,v∗Wlog2​e​−hi​σ2​}。将$p_i^{\ast }$带入条件${\sum }_i^N{p}_i^{\ast }=P$得到：
∑ i N max ⁡ { 0 , W log ⁡ 2 e v ∗ − σ 2 h i } = P \sum_i^N \max \{0,\frac{W\log_2e}{v^*}-\frac{\sigma^2}{h_i}\}=P i∑N​max{0,v∗Wlog2​e​−hi​σ2​}=P
方程左端是$\frac{W{\log }_{2}e}{v^{\ast }}$的分段线性增函数，分割点为${\sigma }^2/h_i$，因此上述方程有唯一确定的解。

上述解决问题的方法称为注水。这是因为，可以将${\sigma }^2/h_i$看做第$i$片区域的水平线，然后对整个区域注水，使其具有深度$\frac{W{\log }_{2}e}{v^{\ast }}$，所需的总水量为 ∑ i N max ⁡ { 0 , W log ⁡ 2 e v ∗ − σ 2 h i } \sum_i^N \max \{0,\frac{W\log_2e}{v^*}-\frac{\sigma^2}{h_i}\} ∑iN​max{0,v∗Wlog2​e​−hi​σ2​}，不断注水，直至总水量为$P$。第$i$个区域的水位深度即为最优$p_i^{\ast }$。

上述解需要确定水位$\frac{{\log }_{2}e}{v^{\ast }}$，编程上的思想为：

• 以最差信道的高度$\frac{{\log }_{2}e}{{\sigma }^2/h_{worst}}$作为水位，注水（分配功率）
• 若注水到该水位时功率有剩余，则均分剩余功率；
• 若水量不足，则移除该最差信道（即不分配功率），重复以上操作。

Python代码：

# -*- coding: utf-8 -*-
"""
2021.01.29 注水法进行功率分配
注水法代码参考 https://pyphysim.readthedocs.io/en/latest/_modules/pyphysim/comm/waterfilling.html
"""

import numpy as np

def waterfilling(Channels, TotalPower, NoisePower):
    """ 注水算法进行功率分配
        Channels: 信道增益
        TotalPower: 待分配的总发射功率
        NoisePower: 接收端的噪声功率
    
    Returns:
        Powers: optimum powers (分配的功率)
        mu: water level (水位)
    """
    ### 降序排列信道增益
    Channels_SortIndexes = np.argsort(Channels)[::-1]
    Channels_Sorted = Channels[Channels_SortIndexes]
    """
    计算接触最差信道的水位，对这个最差的信道分配零功率。
    此后，按此水位为每个信道分配功率，
        如果功率之和少于总功率，则均分剩余功率给各个信道（增加水位）；
        如果功率之和多于总功率，则移除最坏信道，重复操作
    """
    N_Channels = Channels.size ## 总信道数
    N_RemovedChannels = 0  ## 移除的信道数
    ## 按最差信道计算最低水位
    WaterLevel = NoisePower / (Channels_Sorted[N_Channels-N_RemovedChannels-1])
    Powers = WaterLevel - (NoisePower /Channels_Sorted[np.arange(0, N_Channels - N_RemovedChannels)])
    
    ## 当功率之和多于总功率时，移除最坏信道，直至总功率够分
    while (sum(Powers)>TotalPower) and (N_RemovedChannels<N_Channels):
        N_RemovedChannels += 1
        WaterLevel = NoisePower / (Channels_Sorted[N_Channels-N_RemovedChannels-1])
        Powers = WaterLevel - (NoisePower /Channels_Sorted[np.arange(0, N_Channels - N_RemovedChannels)])
    
    ## 将剩余的功率均分给各个(剩余的)信道
    Power_Remine = TotalPower-np.sum(Powers)
    Powers_Opt_Temp = Powers + Power_Remine/(N_Channels - N_RemovedChannels)
    
    ## 将功率分配情况按原信道顺序排列
    Powers_Opt = np.zeros([N_Channels])
    Powers_Opt[Channels_SortIndexes[np.arange(0, N_Channels-N_RemovedChannels)]] = Powers_Opt_Temp
    
    WaterLevel = Powers_Opt_Temp[0] + NoisePower / Channels_Sorted[0]
    
    return Powers_Opt, WaterLevel

if __name__ == '__main__':
    """  测试代码
    """
    power_Noise = 1e-8
    channels = np.random.random(10)*1e-10
    alpha = power_Noise/np.array(channels)
    power_Tx = 1000
    powers,waterlevel = waterfilling(np.array(channels), power_Tx, power_Noise)
    print(powers,waterlevel)

    import matplotlib
    import matplotlib.pylab as plt
#    %matplotlib inline
    
    matplotlib.rcParams.update({'font.size': 14})
    buckets = channels.size
    axis = np.arange(0.5,buckets+1.5,1)
    index = axis+0.5
    X = powers.copy()
    Y = alpha+ X
    
    # to include the last data point as a step, we need to repeat it
    A = np.concatenate((alpha,[alpha[-1]]))
    X = np.concatenate((X,[X[-1]]))
    Y = np.concatenate((Y,[Y[-1]]))
    
    plt.xticks(index)
    plt.xlim(0.5,buckets+0.5)
#    plt.ylim(0,1.5)
    plt.step(axis,A,where='post',label =r'$\sigma^2/h_i$',lw=2)
    plt.step(axis,Y,where='post',label=r'$\sigma^2/h_i + p_i$',lw=2)
    plt.legend(loc='upper left')
    plt.xlabel('Bucket Number')
    plt.ylabel('Power Level')
    plt.title('Water Filling Solution')
    plt.show()