diff --git a/save_file.p b/save_file.p
new file mode 100644
index 0000000..b780f7e
Binary files /dev/null and b/save_file.p differ
diff --git a/toy_model.py b/toy_model.py
index ea11e76..41301b3 100644
--- a/toy_model.py
+++ b/toy_model.py
@@ -7,10 +7,12 @@ import numpy as np
 import matplotlib.pyplot as plt
 import time, sys, pickle
 
-### CONTROL
+### CONTROL ###
+
 day = 60*60*24					# define length of day (used for calculating Coriolis as well) (s)
 dt = 60*9						# <----- TIMESTEP (s)
 resolution = 3					# how many degrees between latitude and longitude gridpoints
+nlevels = 4						# how many vertical layers in the atmosphere
 planet_radius = 6.4E6			# define the planet's radius (m)
 insolation = 1370				# TOA radiation from star (W m^-2)
 
@@ -18,7 +20,7 @@ advection = True 				# if you want to include advection set this to be True (cur
 advection_boundary = 7			# how many gridpoints away from poles to apply advection
 
 save = False
-load = True
+load = False
 
 ###########################
 
@@ -28,14 +30,16 @@ lon = np.arange(0,360,resolution)
 nlat = len(lat)
 nlon = len(lon)
 lon_plot, lat_plot = np.meshgrid(lon, lat)
+heights = np.arange(0,100E3,100E3/nlevels)
 
 # initialise arrays for various physical fields
 temperature_planet = np.zeros((nlat,nlon)) + 270
-temperature_atmosp = np.zeros_like(temperature_planet) + 270
-air_pressure = np.zeros_like(temperature_planet)
-u = np.zeros_like(temperature_planet)
-v = np.zeros_like(temperature_planet)
-air_density = np.zeros_like(temperature_planet) + 1.3
+temperature_atmosp = np.zeros((nlat,nlon,nlevels)) + 270
+air_pressure = np.zeros_like(temperature_atmosp)
+u = np.zeros_like(temperature_atmosp)
+v = np.zeros_like(temperature_atmosp)
+w = np.zeros_like(temperature_atmosp)
+air_density = np.zeros_like(temperature_atmosp) + 1.3
 
 albedo = np.zeros_like(temperature_planet) + 0.5
 heat_capacity_earth = np.zeros_like(temperature_planet) + 1E7
@@ -54,7 +58,10 @@ for i in range(nlat):
 # heat_capacity_earth[2:30,85:110] = 1E6
 
 # define physical constants
-epsilon = 0.75
+epsilon = np.zeros(nlevels)
+epsilon[0] = 0.75
+for i in np.arange(1,nlevels):	
+	epsilon[i] = epsilon[i-1]*0.5
 heat_capacity_atmos = 1E7
 
 specific_gas = 287
@@ -75,35 +82,67 @@ angular_speed = 2*np.pi/day
 for i in range(nlat):
 	dx[i] = dy*np.cos(lat[i]*np.pi/180)
 	coriolis[i] = angular_speed*np.sin(lat[i]*np.pi/180)
+dz = np.zeros(nlevels)
+for k in range(nlevels-1): dz[k] = heights[k+1] - heights[k]
+dz[-1] = dz[-2]
 
 ###################### FUNCTIONS ######################
 
 # define various useful differential functions:
 # gradient of scalar field a in the local x direction at point i,j
-def scalar_gradient_x(a,i,j):
-	return (a[i,(j+1)%nlon]-a[i,(j-1)%nlon])/dx[i]
+def scalar_gradient_x(a,i,j,k=999):
+	if k == 999:
+		return (a[i,(j+1)%nlon]-a[i,(j-1)%nlon])/dx[i]
+	else:
+		return (a[i,(j+1)%nlon,k]-a[i,(j-1)%nlon,k])/dx[i]
 
 # gradient of scalar field a in the local y direction at point i,j
-def scalar_gradient_y(a,i,j):
-	if i == 0 or i == nlat-1:
-		return 0
+def scalar_gradient_y(a,i,j,k=999):
+	if k == 999:
+		if i == 0:
+			return 2*(a[i+1,j]-a[i,j])/dy
+		elif i == nlat-1:
+			return 2*(a[i,j]-a[i-1,j])/dy
+		else:
+			return (a[i+1,j]-a[i-1,j])/dy
 	else:
-		return (a[i+1,j]-a[i-1,j])/dy
+		if i == 0:
+			return 2*(a[i+1,j,k]-a[i,j,k])/dy
+		elif i == nlat-1:
+			return 2*(a[i,j,k]-a[i-1,j,k])/dy
+		else:
+			return (a[i+1,j,k]-a[i-1,j,k])/dy
 
-# laplacian of scalar field a in the local x direction
+def scalar_gradient_z(a,i,j,k):
+	if k == 0:
+		return (a[i,j,k+1]-a[i,j,k])/dz[k]
+	elif k == nlevels-1:
+		return (a[i,j,k]-a[i,j,k-1])/dz[k]
+	else:
+		return (a[i,j,k+1]-a[i,j,k-1])/(2*dz[k])
+
+# laplacian of scalar field a
 def laplacian(a):
 	output = np.zeros_like(a)
-	for i in np.arange(1,len(a[:,0])-1):
-		for j in range(len(a[0,:])):
-			output[i,j] = (scalar_gradient_x(a,i,(j+1)%nlon) - scalar_gradient_x(a,i,(j-1)%nlon)/dx[i]) + (scalar_gradient_y(a,i+1,j) - scalar_gradient_y(a,i-1,j))/dy
-	return output
+	if output.ndim == 2:
+		for i in np.arange(1,nlat-1):
+			for j in range(nlon):
+				output[i,j] = (scalar_gradient_x(a,i,(j+1)%nlon) - scalar_gradient_x(a,i,(j-1)%nlon)/dx[i]) + (scalar_gradient_y(a,i+1,j) - scalar_gradient_y(a,i-1,j))/dy
+		return output
+	if output.ndim == 3:
+		for i in np.arange(1,nlat-1):
+			for j in range(nlon):
+				for k in range(nlevels-1):
+					output[i,j,k] = (scalar_gradient_x(a,i,(j+1)%nlon,k) - scalar_gradient_x(a,i,(j-1)%nlon,k)/dx[i]) + (scalar_gradient_y(a,i+1,j,k) - scalar_gradient_y(a,i-1,j,k))/dy + (scalar_gradient_z(a,i,j,k+1)-scalar_gradient_z(a,i,j,k-1))/2*dz[k]
+		return output
 
 # divergence of (a*u) where a is a scalar field and u is the atmospheric velocity field
 def divergence_with_scalar(a):
 	output = np.zeros_like(a)
-	for i in range(len(a[:,0])):
-		for j in range(len(a[0,:])):
-			output[i,j] = scalar_gradient_x(a*u,i,j) + scalar_gradient_y(a*v,i,j)
+	for i in range(nlat):
+		for j in range(nlon):
+			for k in range(nlevels):
+				output[i,j] = scalar_gradient_x(a*u,i,j,k) + scalar_gradient_y(a*v,i,j,k) + scalar_gradient_z(a*w,i,j,k)
 	return output	
 
 # power incident on (lat,lon) at time t
@@ -120,7 +159,7 @@ def solar(insolation, lat, lon, t):
 f, ax = plt.subplots(2,figsize=(9,9))
 f.canvas.set_window_title('CLAuDE')
 ax[0].contourf(lon_plot, lat_plot, temperature_planet, cmap='seismic')
-ax[1].contourf(lon_plot, lat_plot, temperature_atmosp, cmap='seismic')
+ax[1].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,0], cmap='seismic')
 plt.subplots_adjust(left=0.1, right=0.75)
 ax[0].set_title('Ground temperature')
 ax[1].set_title('Atmosphere temperature')
@@ -153,8 +192,14 @@ while True:
 	# calculate change in temperature of ground and atmosphere due to radiative imbalance
 	for i in range(nlat):
 		for j in range(nlon):
-			temperature_planet[i,j] += dt*((1-albedo[i,j])*solar(insolation,lat[i],lon[j],t) + epsilon*sigma*temperature_atmosp[i,j]**4 - sigma*temperature_planet[i,j]**4)/heat_capacity_earth[i,j]
-			temperature_atmosp[i,j] += dt*(epsilon*sigma*temperature_planet[i,j]**4 - 2*epsilon*sigma*temperature_atmosp[i,j]**4)/heat_capacity_atmos
+			temperature_planet[i,j] += dt*((1-albedo[i,j])*solar(insolation,lat[i],lon[j],t) + epsilon[0]*sigma*temperature_atmosp[i,j,0]**4 - sigma*temperature_planet[i,j]**4)/heat_capacity_earth[i,j]
+			for k in range(nlevels):
+				if k == 0:
+					temperature_atmosp[i,j,k] += dt*epsilon[k]*sigma*(temperature_planet[i,j]**4 + epsilon[k+1]*temperature_atmosp[i,j,k+1]**4 - 2*temperature_atmosp[i,j,k]**4)/(air_density[i,j,k]*heat_capacity_atmos)
+				elif k == nlevels-1:
+					temperature_atmosp[i,j,k] += dt*epsilon[k]*sigma*(epsilon[k-1]*temperature_atmosp[i,j,k-1]**4 - 2*temperature_atmosp[i,j,k]**4)/(air_density[i,j,k]*heat_capacity_atmos)
+				else:
+					temperature_atmosp[i,j,k] += dt*epsilon[k]*sigma*(epsilon[k+1]*temperature_atmosp[i,j,k+1]**4 + epsilon[k-1]*temperature_atmosp[i,j,k-1]**4 - 2*temperature_atmosp[i,j,k]**4)/(air_density[i,j,k]*heat_capacity_atmos)
 
 	# update air pressure
 	air_pressure = air_density*specific_gas*temperature_atmosp
@@ -164,15 +209,19 @@ while True:
 		# introduce temporary arrays to update velocity in the atmosphere
 		u_temp = np.zeros_like(u)
 		v_temp = np.zeros_like(v)
+		w_temp = np.zeros_like(w)
 
 		# calculate acceleration of atmosphere using primitive equations on beta-plane
 		for i in np.arange(1,nlat-1):
 			for j in range(nlon):
-				u_temp[i,j] += dt*( -u[i,j]*scalar_gradient_x(u,i,j) - v[i,j]*scalar_gradient_y(u,i,j) + coriolis[i]*v[i,j] - scalar_gradient_x(air_pressure,i,j)/air_density[i,j] )
-				v_temp[i,j] += dt*( -u[i,j]*scalar_gradient_x(v,i,j) - v[i,j]*scalar_gradient_y(v,i,j) - coriolis[i]*u[i,j] - scalar_gradient_y(air_pressure,i,j)/air_density[i,j] )
+				for k in range(nlevels):
+					u_temp[i,j,k] += dt*( -u[i,j,k]*scalar_gradient_x(u,i,j,k) - v[i,j,k]*scalar_gradient_y(u,i,j,k) + coriolis[i]*v[i,j,k] - scalar_gradient_x(air_pressure,i,j,k)/air_density[i,j,k] )
+					v_temp[i,j,k] += dt*( -u[i,j,k]*scalar_gradient_x(v,i,j,k) - v[i,j,k]*scalar_gradient_y(v,i,j,k) - coriolis[i]*u[i,j,k] - scalar_gradient_y(air_pressure,i,j,k)/air_density[i,j,k] )
+					w_temp[i,j,k] += 0
 
 		u += u_temp
 		v += v_temp
+		w += w_temp
 
 		u *= 0.99
 		v *= 0.99
@@ -180,15 +229,15 @@ while True:
 		if advection:
 			# allow for thermal advection in the atmosphere, and heat diffusion in the atmosphere and the ground
 			atmosp_addition = dt*(thermal_diffusivity_air*laplacian(temperature_atmosp) + divergence_with_scalar(temperature_atmosp))
-			temperature_atmosp[advection_boundary:-advection_boundary,:] -= atmosp_addition[advection_boundary:-advection_boundary,:]
-			temperature_atmosp[advection_boundary-1,:] -= 0.5*atmosp_addition[advection_boundary-1,:]
-			temperature_atmosp[-advection_boundary,:] -= 0.5*atmosp_addition[-advection_boundary,:]
+			temperature_atmosp[advection_boundary:-advection_boundary,:,:] -= atmosp_addition[advection_boundary:-advection_boundary,:,:]
+			temperature_atmosp[advection_boundary-1,:,:] -= 0.5*atmosp_addition[advection_boundary-1,:,:]
+			temperature_atmosp[-advection_boundary,:,:] -= 0.5*atmosp_addition[-advection_boundary,:,:]
 
 			# as density is now variable, allow for mass advection
 			density_addition = dt*divergence_with_scalar(air_density)
-			air_density[advection_boundary:-advection_boundary,:] -= density_addition[advection_boundary:-advection_boundary,:]
-			air_density[(advection_boundary-1),:] -= 0.5*density_addition[advection_boundary-1,:]
-			air_density[-advection_boundary,:] -= 0.5*density_addition[-advection_boundary,:]
+			air_density[advection_boundary:-advection_boundary,:,:] -= density_addition[advection_boundary:-advection_boundary,:,:]
+			air_density[(advection_boundary-1),:,:] -= 0.5*density_addition[advection_boundary-1,:,:]
+			air_density[-advection_boundary,:,:] -= 0.5*density_addition[-advection_boundary,:,:]
 
 			temperature_planet += dt*(thermal_diffusivity_roc*laplacian(temperature_planet))
 	
@@ -196,8 +245,8 @@ while True:
 	test = ax[0].contourf(lon_plot, lat_plot, temperature_planet, cmap='seismic')
 	ax[0].set_title('$\it{Ground} \quad \it{temperature}$')
 
-	ax[1].contourf(lon_plot, lat_plot, temperature_atmosp, cmap='seismic')
-	ax[1].streamplot(lon_plot,lat_plot,u,v,density=0.75,color='black')
+	ax[1].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,0], cmap='seismic')
+	ax[1].streamplot(lon_plot,lat_plot,u[:,:,0],v[:,:,0],density=0.75,color='black')
 	ax[1].set_title('$\it{Atmospheric} \quad \it{temperature}$')
 	for i in ax: 
 		i.set_xlim((lon.min(),lon.max()))