## DEV Community is a community of 621,286 amazing developers

We're a place where coders share, stay up-to-date and grow their careers.

# Advent of Code 2020: Day 17 using 3D/4D Convolution in TensorFlow in Python

Yuan Gao Updated on ・4 min read

Running through these quickly! Today's is another short one because the solution looks a lot like Day 12's solution

# The Challenge Part 1

This challenge is another cellular automata question, asking you to run a cellular automata for a few cycles. In fact, it's an easier one than Day 12, but involves doing it in 3D.

Fortunately for us, our solution to Day 12 was done in TensorFlow, and extending it to 3D is easy:

## 3D cellular automata in TensorFlow

Copy-pasting most of the Day 12 solution, and adjusting to add an extra dimension. I'm manually defining a fixed space for the automata, and then dumping the initial state somewhere near the middle of it, though it should be relatively easy to let it expand outwards.

``````import numpy as np
import tensorflow as tf

bool_data = data.view('U1') == '#'
bool_data = bool_data.reshape((1, data.size, -1))

zinit, yinit, xinit = bool_data.shape
zdim, ydim, xdim = 40,40,40

padding = tf.constant([[1, 1], [1, 1], [1, 1]])
neighbors_np = np.ones([3, 3, 3])
neighbors_np[1, 1, 1] = 0
conv_filt = tf.reshape(tf.cast(neighbors_np, tf.float32),  [3, 3, 3, 1, 1])

ystart = xstart = (xdim-xinit)//2
yend = xend = xstart+xinit
zstart = zdim//2
zend = zstart+1
init_data_np = np.zeros([zdim, ydim, xdim], dtype=np.bool)
init_data_np[zstart:zend, ystart:yend, xstart:xend] = bool_data
init_data = tf.convert_to_tensor(init_data_np)

@tf.function
def generate(this_gen):
padded = tf.reshape(padded, [1, zdim+2, ydim+2, xdim+2, 1]) # need for tf.nn.convolution

neighbors = tf.reshape(convolved, [zdim, xdim, ydim])

three_neighbors = tf.math.equal(neighbors, 3)
two_or_three_neighbors = tf.math.logical_or(tf.math.equal(neighbors, 2), three_neighbors)

next_gen = tf.math.logical_or(this_gen, three_neighbors)
next_gen = tf.math.logical_and(next_gen, two_or_three_neighbors)
return next_gen

generation = init_data
for _ in range(6):
generation = generate(generation)

print("total", tf.math.reduce_sum(tf.cast(generation, tf.int32)))
``````

This is the full solution, the meat of it is the `generate()` function which is almost exactly the same as last time, but implementing the specific cellular automata rules for this problem.

# The Challenge Part 2

The second part of the challege is exactly the same as the first, but extending the automata to run in 4 dimensions.

## 4D Convolution in TensorFlow

This presents a slight challenge to us because TensorFlow doesn't usually do 4D convolution. Fortunately someone has implemented a method to do it, but this is for TensorFlow 1, while we're using the newer TensorFlow 2. I took some time to adapt the code for our use-case

The full code is therefore:

``````import numpy as np
import tensorflow as tf

bool_data = data.view('U1') == '#'
bool_data = bool_data.reshape((1, 1, data.size, -1))

winit, zinit, yinit, xinit = bool_data.shape
wdim, xdim, ydim, zdim = 40, 40, 40, 40

padding = tf.constant([[1, 1], [1, 1], [1, 1], [1, 1]])
neighbors_np = np.ones([3, 3, 3, 3])
neighbors_np[1, 1, 1, 1] = 0
conv_filt = tf.reshape(tf.cast(neighbors_np, tf.float32),  [3, 3, 3, 3, 1, 1])

init_data_np = np.zeros([wdim, zdim, xdim, ydim], dtype=np.bool)
ystart = xstart = (xdim-xinit)//2
yend = xend = xstart+xinit
wstart = zstart = zdim//2
wend = zend = zstart+1
init_data_np[wstart:wend, zstart:zend, ystart:yend, xstart:xend] = bool_data
init_data = tf.convert_to_tensor(init_data_np)

@tf.function
def conv4d(data, conv_filt):
# input, kernel, and output sizes
(b, wi, zi, yi, xi, c) = data.shape.as_list()
(wk, zk, yk, xk, ik, ok) = conv_filt.shape.as_list()

# output size and tensor
wo = wi - wk + 1
results = [ None ]*wo

# convolve each kernel frame i with each input frame j
for i in range(wk):
for j in range(wi):

# add results to this output frame
out_frame = j - (i - wk//2) - (wi - wo)//2
if out_frame < 0 or out_frame >= wo:
continue

# convolve input frame j with kernel frame i
frame_conv3d = tf.nn.convolution(tf.reshape(data[:,:,j,:,:], (b, zi, yi, xi, c)), conv_filt[:,:,:,i])

if results[out_frame] is None:
results[out_frame] = frame_conv3d
else:
results[out_frame] += frame_conv3d

return tf.stack(results, axis=2)

@tf.function
def generate(this_gen):
padded = tf.reshape(padded, [1, wdim+2, zdim+2, ydim+2, xdim+2, 1]) # need for tf.nn.convolution

neighbors = tf.reshape(convolved, [wdim, zdim, xdim, ydim])

three_neighbors = tf.math.equal(neighbors, 3)
two_or_three_neighbors = tf.math.logical_or(tf.math.equal(neighbors, 2), three_neighbors)

next_gen = tf.math.logical_or(this_gen, three_neighbors)
next_gen = tf.math.logical_and(next_gen, two_or_three_neighbors)
return next_gen

generation = init_data
for _ in range(6):
generation = generate(generation)

print("total", tf.math.reduce_sum(tf.cast(generation, tf.int32)))
``````

Thanks to the power of TensorFlow, we didn't really need to think too much about the implementation of this cellular automata: instead of thinking about iterating over all the dimensions and worrying that we have off-by-one errors, we can just paint it in broad strokes by doing tensor operations.

Onwards!