diff --git a/work/week12/starting_point/generate_input.py b/work/week12/starting_point/generate_input.py
index cf23df2..f5e266f 100644
--- a/work/week12/starting_point/generate_input.py
+++ b/work/week12/starting_point/generate_input.py
@@ -1,127 +1,131 @@
 import os
 import numpy as np
 import math
 import argparse
 import matplotlib.pyplot as plt
 
 
 def isSquare(n):
     ''' Check if an integer value is a perfect square or not. '''
     x = n // 2
     seen = set([x])
     while x * x != n:
         x = (x + (n // x)) // 2
         if x in seen:
             return False
         seen.add(x)
     return True
 
 
 def plotDistributions(heat2d, temp2d):
     ''' Plot 2D heat and temperature distributions. '''
     fig, axes = plt.subplots(1, 2, figsize=(12, 5))
 
     ax = axes[0]
     ax.set_title('Heat source distribution')
     ax.set_xlabel('Y')
     ax.set_ylabel('X')
-    ax.imshow(heat2d.real, extent=[-1, 1, -1, 1])
+    ax.imshow(heat2d, extent=[-1, 1, -1, 1])
 
     ax = axes[1]
     ax.set_title('Temperature distribution')
     ax.set_xlabel('Y')
     ax.set_ylabel('X')
-    ax.imshow(temp2d.real, extent=[-1, 1, -1, 1])
+    ax.imshow(temp2d, extent=[-1, 1, -1, 1])
 
     return fig
 
 
 def main():
 
     parser = argparse.ArgumentParser()
     parser.add_argument('-n', '--number', type=int, default=32**2,
                         help='Number of particles')
     parser.add_argument('-o', '--outfile', type=str, default='step-00000.csv',
                         help='Name of output file')
     parser.add_argument('-H', '--heatdist', type=str, default='null',
                         help='Heat distribution type ("null", "punctual" or "circular")')
     parser.add_argument('-R', '--radius', type=float, default=1 / 3,
                         help='Heat source radius (if heatdist is set to "circular")')
     parser.add_argument('-p', '--plot', default=False, action='store_true',
                         help='Plot heat and temperature distribution')
     parser.add_argument('--hr', type=float, default=3,
                         help='Specific heat rate')
 
     args = parser.parse_args()
 
     # Get number of particles and square root of it
     N = int(args.number)
     if not(isSquare(N)):
         raise ValueError('Number of particles must be square')
     sqrtN = int(math.sqrt(N))
 
     # Set initial positions on a 2D rectilinear grid (z = 0)
     pos1D = np.linspace(-1, 1, sqrtN)
     pos2D = np.array(np.meshgrid(pos1D, pos1D)).T
 
     # Set heat volumetric source
     hv2d = np.zeros((sqrtN, sqrtN))
     if args.heatdist == 'null':
         pass  # do nothing
     elif args.heatdist == 'punctual':
         hv2d[int(sqrtN / 4), :] = -sqrtN
         hv2d[int(3 * sqrtN / 4), :] = sqrtN
     elif args.heatdist == 'circular':
         R = args.radius
         if R < 0:
             raise ValueError('radius must be positive')
         for i in range(sqrtN):
             for j in range(sqrtN):
                 if (pos1D[i]**2 + pos1D[j]**2 < R**2):
                     hv2d[i, j] = N
     else:
         raise ValueError('heatdist can be "null", "punctual" or "circular"')
 
     # Compute frequency modes in the Fourier space
     L = 2  # normalization factor
     q = np.fft.fftfreq(sqrtN) * 2 * np.pi / L * sqrtN
     q2d = np.einsum('i,j->ij', np.ones(sqrtN), q)
 
     # Compute Laplacian of frequency modes
     laplacian_q = q2d**2 + q2d.T**2
     laplacian_q[0, 0] = 1.
 
     # Compute 2D Fourier transform of heat distribution
     hv2d_hat = np.fft.fft2(hv2d)
     hv2d_hat[0, 0] = 0.  # setting origin to zero
 
     # Set temperature distribution to respect equilibrium for given heat distribution
     theta2d = np.fft.ifft2(hv2d_hat / (args.hr * laplacian_q))
 
+    # Extract real component of both matrices
+    theta2d = theta2d.real
+    hv2d = hv2d.real
+
     # Reshape matrices into serialized vectors
     positions = np.hstack((pos2D.reshape(N, 2), np.zeros((N, 1))))
     hv = np.reshape(hv2d, (hv2d.size, 1))
     theta = np.reshape(theta2d, (theta2d.size, 1))
 
     # Generate serialized vectors for other, constant quantities
     hr = args.hr * np.ones((N, 1))
     velocities = np.zeros((N, 3))
     forces = np.zeros((N, 3))
     masses = np.ones((N, 3))
 
     # Save data into csv file
     file_data = np.hstack((positions, velocities, forces, masses, theta, hr, hv))
     outdir = os.path.join(os.getcwd(), 'build', 'dumps')
     if not os.path.isdir(outdir):
         os.makedirs(outdir)
     filepath = os.path.join(os.getcwd(), 'build', 'dumps', args.outfile)
     np.savetxt(filepath, file_data, delimiter=" ")
 
     # Plot the heat and temperature fields
     if args.plot:
         plotDistributions(hv2d, theta2d)
         plt.show()
 
 
 if __name__ == '__main__':
     main()