Advent of Code 2021 in pure TensorFlow - day 11

The Day 11 problem has lots in common with Day 9. In fact, will re-use some computer vision concepts like the pixel neighborhood, and we’ll be able to solve both parts in pure TensorFlow by using only a tf.queue as a support data structure.

Day 11: Dumbo Octopus

You can click on the title above to read the full text of the puzzle. The TLDR version is:

Our dataset is a grid of numbers. Every number represents an energy level of an octopus. On every time step, every energy level increases by 1 unit, and when the energy level goes beyond 9, the octopus flashes. When an octopus flashes it increases the energy level of all (in every direction) the octopus in the neighborhood. This may trigger a cascade of flashes. Finally, any octopus that flashed during this time step, rests its energy level to 0.

For example:

Before any steps:

After step 1:

After step 2:

An octopus is highlighted when it flashed during the given time step. Part one asks us to find out how many total flashes are there after 100 steps?

Design phase: part one

The problem requires modeling the increment on every time step, and the propagation of the increment to the neighborhood when an energy level goes beyond 9. Thus, considering the input dataset as a 2D grayscale image, we can treat every energy level as a pixel. Thus, we need:

  1. Define a function that given a pixel returns the 8-neighborhood coordinates. Differently from the solution presented in day 9: design phase, where the problem required us to only look at the 4-neighborhood, this time we need to consider also the diagonal pixels. Moreover, we can’t simplify in the same way the problem padding the input image, since every update may trigger a cascade of updates that may involve the padding pixel introduced, and we don’t want it.
  2. Given that every flash may trigger a cascade of flash, we need to increment by 1 every value in the neighborhood of a flashing pixel and keep track of all the neighbors that exceeded the value of 9. Pushing the coordinates into a queue, and repeating the process for every pixel in the queue until the propagation is not complete.

Input pipeline

We create a object for reading the text file line-by-line as usual. Since we want to work with all the pixels at once, we convert the dataset in a 2D tf.Tensor - this is our octopus population.

population = tf.convert_to_tensor(
        .map(lambda string: tf.strings.to_number(string, out_type=tf.int64))

We can now start implementing our TensorFlow program for solving part 1. We can start by the definition of the _neighs function.

The 8-neighborhood

The first point of the design phase explains why we need to take care of the 8-neighborhood without padding the input. Hence, our neighborhood will contain 8 pixels if we are inside the image. It will contain 3 or 5 pixels of we are on a corner or on the side of the image, respectively.

We need to take care of these conditions manually.

def _neighs(grid: tf.Tensor, center: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    y, x = center[0], center[1]

    shape = tf.shape(grid, tf.int64) - 1

    if tf.logical_and(tf.less(y, 1), tf.less(x, 1)):  # 0,0
        mask = tf.constant([(1, 0), (0, 1), (1, 1)])
    elif tf.logical_and(tf.equal(y, shape[0]), tf.equal(x, shape[1])):  # h,w
        mask = tf.constant([(-1, 0), (0, -1), (-1, -1)])
    elif tf.logical_and(tf.less(y, 1), tf.equal(x, shape[1])):  # top right
        mask = tf.constant([(0, -1), (1, 0), (1, -1)])
    elif tf.logical_and(tf.less(x, 1), tf.equal(y, shape[0])):  # bottom left
        mask = tf.constant([(-1, 0), (-1, 1), (0, 1)])
    elif tf.less(x, 1):  # left
        mask = tf.constant([(1, 0), (-1, 0), (-1, 1), (0, 1), (1, 1)])
    elif tf.equal(x, shape[0]):  # right
        mask = tf.constant([(-1, 0), (1, 0), (0, -1), (-1, -1), (1, -1)])
    elif tf.less(y, 1):  # top
        mask = tf.constant([(0, -1), (0, 1), (1, 0), (1, -1), (1, 1)])
    elif tf.equal(y, shape[1]):  # bottom
        mask = tf.constant([(0, -1), (0, 1), (-1, 0), (-1, -1), (-1, 1)])
    else:  # generic
        mask = tf.constant(
            [(-1, 0), (0, -1), (1, 0), (0, 1), (-1, -1), (1, 1), (-1, 1), (1, -1)]

    coords = center + tf.cast(mask, tf.int64)
    neighborhood = tf.gather_nd(grid, coords)
    return neighborhood, coords

The _neihgs function accepts the grid (2D tf.Tensor) and the coordinate of a pixel, that’s the center of the neighborhood.

Depending on the center coordinate respect to the grid, the function builds a different mask used to create the pixel coordinate and gather the value in that coordinates.

Note that TensorFlow coordinates, by convention, are in y,x format and not in x,y format. We are now ready for implementing the algorithm.

Simulating the population behavior

The puzzle asks us to simulate 100 steps, and count how many flashes happened during them. Thus, we need to completely simulate the octopus population behavior over time.

We already have our data organized in a grid, and from the second point of the design phase, we know that the unique data structure needed for model the propagation is a tf.queue.

def __init__(self, population, steps):

    self._steps = steps
    self._population = tf.Variable(population, dtype=tf.int64)
    self._counter = tf.Variable(0, dtype=tf.int64)

    self._zero = tf.constant(0, dtype=tf.int64)
    self._one = tf.constant(1, dtype=tf.int64)
    self._nine = tf.constant(9, tf.int64)
    self._ten = tf.constant(10, dtype=tf.int64)

    self._queue = tf.queue.FIFOQueue(-1, [tf.int64])

    self._flashmap = tf.Variable(tf.zeros_like(self._population))

The constructor requires the population and the number of steps to simulate. The population input variable is then used to initialize a tf.Variable that we’ll use to store the population status after every time step.

The _counter variable will be our puzzle output, and the _flashmap is a helper variable with the same shape of the population that we use to keep track of the flashing pixels during every time step - so we can know what pixel flashed and keep track of them.

Other than the _queue used for modeling the flash propagation, we also define some constant with the dtype tf.int64 that we’ll later use inside the TensorFlow program. This is a good practice to follow since we avoid the creation of several different (and useless) constants in the graph, and we’ll always refer to the same constants instead.

Modeling the population behavior is straightforward - it’s just a matter of converting in TensorFlow code the puzzle instructions.

def __call__(self):
    for step in tf.range(self._steps):
        # First, the energy level of each octopus increases by 1.

        # Then, any octopus with an energy level greater than 9 flashes.
        flashing_coords = tf.where(tf.greater(self._population, self._nine))

        # This increases the energy level of all adjacent octopuses by 1, including octopuses that are diagonally adjacent.
        # If this causes an octopus to have an energy level greater than 9, it also flashes.
        # This process continues as long as new octopuses keep having their energy level increased beyond 9.
        # (An octopus can only flash at most once per step.)
        while tf.greater(self._queue.size(), 0):
            p = self._queue.dequeue()
            if tf.greater(self._flashmap[p[0], p[1]], 0):
            self._flashmap.scatter_nd_update([p], [1])

            _, neighs_coords = self._neighs(self._population, p)
            updates = tf.repeat(
                tf.shape(neighs_coords, tf.int64)[0],
            self._population.scatter_nd_add(neighs_coords, updates)
            flashing_coords = tf.where(tf.greater(self._population, self._nine))

        # Finally, any octopus that flashed during this step has its energy level set to 0, as it used all of its energy to flash.
        indices = tf.where(tf.equal(self._flashmap, self._one))
        if tf.greater(tf.size(indices), 0):
            shape = tf.shape(indices, tf.int64)
            updates = tf.repeat(self._zero, shape[0])
            self._population.scatter_nd_update(indices, updates)


        # tf.print(step, self._population, summarize=-1)
    return self._counter

The code is the precise implementation of the puzzle instructions, the comments are part of the original puzzle text and after every comment, there’s the equivalent TensorFlow implementation.


Here we go!

steps = tf.constant(100, tf.int64)
flash_counter = FlashCounter(population, steps)
tf.print("Part one: ", flash_counter())

Part one is solved! Let’s see what part two is about.

Design phase: part 2

The propagation of the flashes causes a nice phenomenon: the octopuses are synchronizing! For example

After step 194:

After step 195:

Part two asks us to determine the first step during which all the octopuses flash.

Part two implementation

Our algorithm already models the population behavior, hence we already have implemented all we need to solve part two.

In fact, part two solution is identical to part 1 solution with the for loop replaced with a while. In part one, we have a fixed number of steps, instead this time we re-use the counter variable to keep track of how many steps we performed until the while condition holds.

def find_sync_step(self):
    # use count as step
    while tf.logical_not(tf.reduce_all(tf.equal(self._population, self._zero))):
         # First, the energy level of each octopus increases by 1.
         # [ very same code of previous solution, omitted!]
         # Full version:

Since we re-used the same internal status of the FlashCounter class, we have two options. Create a new object for invoking the find_sync_step or speeding up the computation knowing that our search starts from step 100 used to solve part 1.

tf.print("Part two: ", steps + flash_counter.find_sync_step())

Here we go! Challenge 11 is solved in pure TensorFlow pretty easily.


You can see the complete solution in folder 11 on the dedicated Github repository:

Solving this problem has, once again, demonstrated how TensorFlow can be used to solve any kind of programming problem, and how we can think of TensorFlow as a different programming language, especially if we are able to write pure TensorFlow programs (hence tf.function graph-converted functions).

If you missed the articles about the previous days’ solutions, here’s a handy list:

The next article will be about my solution to Day 12 problem. In that article I’ll show how to work with graphs in the “traditional” way: no neural networks, only an adjacency matrix, and a search algorithm implemented using recursion. However, tf.function can’t be used to write recursive functions :( hence I end up writing a pure TensorFlow eager solution.

For any feedback or comment, please use the Disqus form below - thanks!

Don't you want to miss the next article? Do you want to be kept updated?
Subscribe to the newsletter!

Related Posts

Advent of Code 2021 in pure TensorFlow - day 12

Day 12 problem projects us the world of graphs. TensorFlow can be used to work on graphs pretty easily since a graph can be represented as an adjacency matrix, and thus, we can have a tf.Tensor containing our graph. However, the "natural" way of exploring a graph is using recursion, and as we'll see in this article, this prevents us to solve the problem using a pure TensorFlow program, but we have to work only in eager mode.

Advent of Code 2021 in pure TensorFlow - day 10

The day 10 challenge projects us in the world of syntax checkers and autocomplete tools. In this article, we'll see how TensorFlow can be used as a generic programming language for implementing a toy syntax checker and autocomplete.

Advent of Code 2021 in pure TensorFlow - day 9

The day 9 challenge can be seen as a computer vision problem. TensorFlow contains some computer vision utilities that we'll use - like the image gradient - but it's not a complete framework for computer vision (like OpenCV). Anyway, the framework offers primitive data types like tf.TensorArray and tf.queue that we can use for implementing a flood-fill algorithm in pure TensorFlow and solve the problem.

Advent of Code 2021 in pure TensorFlow - day 8

The day 8 challenge is, so far, the most boring challenge faced 😅. Designing a TensorFlow program - hence reasoning in graph mode - would have been too complicated since the solution requires lots of conditional branches. A known AutoGraph limitation forbids variables to be defined in only one branch of a TensorFlow conditional if the variable is used afterward. That's why the solution is in pure TensorFlow eager.

Advent of Code 2021 in pure TensorFlow - day 7

The day 7 challenge is easily solvable with the help of the TensorFlow ragged tensors. In this article, we'll solve the puzzle while learning what ragged tensors are and how to use them.

Advent of Code 2021 in pure TensorFlow - day 6

The day 6 challenge has been the first one that obliged me to completely redesign for part 2 the solution I developed for part 1. For this reason, in this article, we'll see two different approaches to the problem. The former will be computationally inefficient but will completely model the problem, hence it will be easy to understand. The latter, instead, will be completely different and it will focus on the puzzle goal instead of the complete modeling.

Advent of Code 2021 in pure TensorFlow - day 5

The day 5 challenge is easily solvable in pure TensorFlow thanks to its support for various distance functions and the power of the tf.math package. The problem only requires some basic math knowledge to be completely solved - and a little bit of computer vision experience doesn't hurt.

Advent of Code 2021 in pure TensorFlow - day 4

Using tensors for representing and manipulating data is very convenient. This representation allows changing shape, organizing, and applying generic transformations to the data. TensorFlow - by design - executes all the data manipulation in parallel whenever possible. The day 4 challenge is a nice showcase of how choosing the correct data representation can easily simplify a problem.

Advent of Code 2021 in pure TensorFlow - day 3

A Solution to the AoC day 3 puzzle in pure TensorFlow. This challenge allows us to explore the TensorArray data type and find their limitations when used inside a static-graph context. We'll also use a tf.function experimental (but very useful) feature for avoiding useless retraces and reusing the same graph with tensors of different shapes.

Advent of Code 2021 in pure TensorFlow - day 2

A Solution to the AoC day 2 puzzle in pure TensorFlow. How to use Enums in TensorFlow programs and the limitations of tf.Tensor used for type annotation