# Water-filling algorithm explained and implemented in Python

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.