Wireless

Water-filling algorithm explained and implemented in Python

Jake

Aug 10, 2021 • 5 min read

In this post, I'm going to show you how the water-filling algorithm works, and how to implement it in Python. For a more in-depth overview of the water-filling algorithm, check our post showing the algorithm's behavior and derivation.

Water-filling is a generic method for allocating transmission power over (equalizing) communications channels without inter-channel interference (ICI). That is, the channels are assumed to be in a technical sense orthogonal. This can mean OFDMA, TDMA, ZF-MIMO, or any other means of removing the ICI by design.

We already had a look at the same problem in our previous post, where we solved the problem using a generic convex optimization algorithm from CVXPY. Here, we derive a custom algorithm for solving the problem more commonly known as water-filling.

Let's first define what we mean by channel capacity in this context. We can write the throughput for \(N\) orthogonal channels by the sum of their individual rates. That is,

\[ \sum_{n=1}^N \text{log}_2 (1 + \text{SNR}_n) \]

To get to the capacity of the said system, we need to set some constraints on the SNR. We will use the sum power budget (\(P\)), which limits the available power across the channels as

\[ \sum_{n=1}^N p_n = P \]

Further, we must define the \(\text{SNR}_n\) in terms of the \(p_n\) as

\[ \text{SNR}(p_n) = {g_n p_n \over N_0} \]

where \(g_n\) denotes the channel gain for the \(n^\text{th}\) channel, \(p_n\) is the corresponding power allocation and \(N_0\) is a fixed noise-level determined by the given SNR.

Now, we can write the capacity as

\[ C = \sum_{n=1}^N \text{log}_2 (1 + \text{SNR}(p^*_n)) \]

where \(p^*_n\) is the sum-rate maximizing power allocation for channel \(n\). Now, we have the basic premise to come up with an algorithm that finds the optimal power allocation, namely, the water-filling algorithm.

Algorithm derivation

To see, where the water filling algorithm actually comes from, we formulate the capacity-achieving power allocation as a throughput maximization problem, where we try to maximize the sum throughput over \(N\) channels with varying channel gains while being limited by a power budget. This problem can be mathematically formulated as

\[ \begin{array}{rl} \underset{p_n}{\text{max.}} & \displaystyle \sum_{n=1}^N \text{log}_2 (1 + {g_n p_n \over N_0}) \\ \text{s.t.} & \displaystyle \sum_{n=1}^N p_n = P \\ & \displaystyle p_n \geq 0 \end{array} \]

This is a convex problem, so any solution satisfying the Karush-Kuhn-Tucker (KKT) conditions can be expected to be optimal.

We start with the stationary condition. The derivative of the Lagrangian w.r.t. all \(p_n\) has to equal zero

\[ \begin{array}{rl} \Delta_{p_n} (-\sum_{n=1}^N \text{log}_2 (1 + {g_n p_n \over N_0}) + \alpha(\sum_{n=1}^N p_n - P) - \sum_{n=1}^N\lambda_n p_n) & \displaystyle = 0 \\ \Leftrightarrow -{g_n \over g_n p_n + N_0} + \alpha - \lambda_n & \displaystyle = 0 \\ \Leftrightarrow {g_n \over g_n p_n + N_0} & \displaystyle = \alpha - \lambda_n \\ \Leftrightarrow {g_n \over {\alpha - \lambda_n} } & \displaystyle = (g_n p_n + N_0) \\ \Leftrightarrow {g_n \over {\alpha - \lambda_n} } - N_0 & \displaystyle = g_n p_n \\ \Leftrightarrow {1 \over {\alpha - \lambda_n} } - {N_0 \over g_n} & \displaystyle = p_n \end{array} \]

This gives as the solution for channel-specific power allocation \[ p_n = \Big({1 \over {\alpha - \lambda_n} } - {N_0 \over g_n}\Big)^{+} \hspace2ex (1) \] We consider only the positive solutions due to the primal constraint \(p_n \geq 0\).

Having the channel-specific dual variables \(\lambda_n\) appearing in the solution makes solving the problem quite complex. Luckily, there is a way to get rid of them. Taking a look at the complementary slackness KKT condition \[\lambda_n p_n = 0\]

This states that either \(\lambda_n = 0\) or \(p_n = 0\), which means that if \(p_n > 0\), then \(\lambda_n = 0\) and the power allocation reduces to \[ p_n = \Big({1 \over {\alpha} } - {N_0 \over g_n}\Big)^{+} \hspace2ex (1) \]

Now, we can search for such \(\alpha\) that satisfies the power budget allocation, and we get the optimal power allocation.

Water-filling in Python

We will use the same environment setup as in our previous article, where we solve the same problem by using CVXPY.

import numpy as np

# Number of channels 
N = 10
N0 = 1 # Normalized noise level
SNR_dB = 10 # The signal to noise ratio in dB
P = 10**(SNR_dB/10) # Sum power budget defined via the SNR

# The channel specific gains drawn from Gaussian distribution
g = np.abs(np.random.randn(N, 1))

We can see that adjusting \(\alpha\) acts as a scaling coefficient for the power allocation. Thus, increasing \(\alpha\) always lowers power allocation and vice versa. Keeping this in mind, we will implement a bisection search for the power allocation also known as the water-filling, where we search for the water level (power scaling), by bisecting the search range for \(\alpha\).

To begin with, introduce the initial bounds for \(\alpha\). Here, zero is the absolute lower bound, and the higher bound is set by a sufficiently large value.

# Bisection search for alpha

alpha_low = min(N0/g) # Initial low
alpha_high = (P + np.sum(N0/g)) / N # Initial high

The stop threshold is used to set the point when the algorithm has converged. Our test condition is \( | \alpha_\text{low} - \alpha_\text{high} | < \epsilon_\text{stop} \)

stop_threshold = 1e-5 # Stop threshold

The complete bisection loop is given by

# Iterate while low/high bounds are further than stop_threshold
while(np.abs(alpha_low - alpha_high) > threshold):
    alpha = (alpha_low + alpha_high) / 2 # Test value in the middle of low/high

    # Solve the power allocation
    p = 1/alpha - N0/g 
    p[p < 0] = 0 # Consider only positive power allocation
    
    # Test sum-power constraints
    if (np.sum(p) > P): # Exceeds power limit => lower the upper bound
        alpha_low = alpha
    else: # Less than power limit => increase the lower bound
        alpha_high = alpha
        
# Print the achievable rate in nats/s
print(np.sum(np.log(1 + g*p/N0)))

You can test the solution against the generic version from here. Both should provide identical solutions given the same parameters.

We were able to write a custom algorithm for solving the channel capacity-achieving power allocation with a couple of lines of code. Just by having a look at the KKT conditions and identifying the solution. This gives us a good place to start for more complex systems. For example, in an upcoming post, we will utilize the same water-filling algorithm to explore the capacity of the MIMO channel.

The complete source can be found on Github.

Please, leave your comments and suggestions in the video comments section.

Algorithm derivation

Water-filling in Python

Further reading