diff --git a/save_file.p b/save_file.p
index 6d08d4e..b5e9d4a 100644
Binary files a/save_file.p and b/save_file.p differ
diff --git a/toy_model.py b/toy_model.py
index a6d3513..f775f49 100644
--- a/toy_model.py
+++ b/toy_model.py
@@ -10,23 +10,24 @@ import time, sys, pickle
 ######## CONTROL ########
 
 day = 60*60*24					# define length of day (used for calculating Coriolis as well) (s)
-dt = 60*9						# <----- TIMESTEP (s)
-resolution = 5					# how many degrees between latitude and longitude gridpoints
-nlevels = 20					# how many vertical layers in the atmosphere
-top = 20E3						# top of atmosphere (m)
+resolution = 3					# how many degrees between latitude and longitude gridpoints
+nlevels = 10					# how many vertical layers in the atmosphere
+top = 50E3						# top of atmosphere (m)
 planet_radius = 6.4E6			# define the planet's radius (m)
 insolation = 1370				# TOA radiation from star (W m^-2)
 gravity = 9.81 					# define surface gravity for planet (m s^-2)
+axial_tilt = -23.5				# tilt of rotational axis w.r.t. solar plane
+year = 365*day					# length of year (s)
 
 ###
 
 advection = True 				# if you want to include advection set this to be True
-advection_boundary = 3			# how many gridpoints away from poles to apply advection
+advection_boundary = 4			# how many gridpoints away from poles to apply advection
 
-save = False
+save = True
 load = False
 
-plot = True
+plot = False
 level_plots = False
 
 ###########################
@@ -41,7 +42,7 @@ heights = np.arange(0,top,top/nlevels)
 heights_plot, lat_z_plot = np.meshgrid(lat,heights)
 
 # initialise arrays for various physical fields
-temperature_planet = np.zeros((nlat,nlon)) + 270
+temperature_planet = np.zeros((nlat,nlon)) + 290
 temperature_atmosp = np.zeros((nlat,nlon,nlevels))
 air_pressure = np.zeros_like(temperature_atmosp)
 u = np.zeros_like(temperature_atmosp)
@@ -49,11 +50,6 @@ v = np.zeros_like(temperature_atmosp)
 w = np.zeros_like(temperature_atmosp)
 air_density = np.zeros_like(temperature_atmosp)
 
-# #######################
-upward_radiation = np.zeros(nlevels)
-downward_radiation = np.zeros(nlevels)
-optical_depth = np.zeros(nlevels)
-Q = np.zeros(nlevels)
 # #######################
 
 def profile(a):
@@ -79,19 +75,17 @@ for k in range(nlevels):
 
 ###########################
 
-weight_above = np.interp(x=heights/1E3,xp=standard_height,fp=standard_density)
-top_index = np.argmax(np.array(standard_height) >= top/1E3)
-if standard_height[top_index] == top/1E3:
-	weight_above = np.trapz(np.interp(x=standard_height[top_index:],xp=standard_height,fp=standard_density),standard_height[top_index:])*gravity*1E3
-else:
-	weight_above = np.trapz(np.interp(x=np.insert(standard_height[top_index:], 0, top/1E3),xp=standard_height,fp=standard_density),np.insert(standard_height[top_index:], 0, top/1E3))*gravity*1E3
-weight_above *= 1.1
-# sys.exit()
+# weight_above = np.interp(x=heights/1E3,xp=standard_height,fp=standard_density)
+# top_index = np.argmax(np.array(standard_height) >= top/1E3)
+# if standard_height[top_index] == top/1E3:
+# 	weight_above = np.trapz(np.interp(x=standard_height[top_index:],xp=standard_height,fp=standard_density),standard_height[top_index:])*gravity*1E3
+# else:
+# 	weight_above = np.trapz(np.interp(x=np.insert(standard_height[top_index:], 0, top/1E3),xp=standard_height,fp=standard_density),np.insert(standard_height[top_index:], 0, top/1E3))*gravity*1E3
 
 ###########################
 
-albedo = np.zeros_like(temperature_planet) + 0.5
-heat_capacity_earth = np.zeros_like(temperature_planet) + 1E7
+albedo = np.zeros_like(temperature_planet) + 0.2
+heat_capacity_earth = np.zeros_like(temperature_planet) + 1E6
 
 albedo_variance = 0.001
 for i in range(nlat):
@@ -186,48 +180,57 @@ def divergence_with_scalar(a):
 	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)
+				output[i,j] = scalar_gradient_x(a*u,i,j,k) + scalar_gradient_y(a*v,i,j,k) #+ 0.1*scalar_gradient_z(a*w,i,j,k)
 	return output	
 
 # power incident on (lat,lon) at time t
 def solar(insolation, lat, lon, t):
 	sun_longitude = -t % day
 	sun_longitude *= 360/day
-	value = insolation*np.cos(lat*np.pi/180)*np.cos((lon-sun_longitude)*np.pi/180)
-	if value < 0:	return 0
-	else:	return value
+	sun_latitude = axial_tilt*np.cos(t*2*np.pi/year)
+
+	value = insolation*np.cos((lat-sun_latitude)*np.pi/180)
+
+	if value < 0:	
+		return 0
+	else:
+
+		lon_diff = lon-sun_longitude
+		value *= np.cos(lon_diff*np.pi/180)
+		
+		if value < 0:
+			if lat + sun_latitude > 90:
+				return insolation*np.cos((lat+sun_latitude)*np.pi/180)*np.cos(lon_diff*np.pi/180)
+			elif lat + sun_latitude < -90:
+				return insolation*np.cos((lat+sun_latitude)*np.pi/180)*np.cos(lon_diff*np.pi/180)
+			else:
+				return 0
+		else:
+			return value
 
 def surface_optical_depth(lat):
-	return 3.75 + np.cos(lat*np.pi/90)*4.5/2
+	return 4 + np.cos(lat*np.pi/90)*2.5/2
+
+def smooth(a,window):
+	output = np.zeros_like(a)
+	diff = int(np.floor(window/2))
+	for i in range(len(output)):
+		if i-diff < 0:
+			output[i] = np.mean(a[i:i+diff+1])
+		elif i+diff+1 > len(output):
+			output[i] = np.mean(a[i-diff:i])
+		else:
+			output[i] = np.mean(a[i-diff:i+diff+1])
+	return output
+
+def thermal_radiation(a):
+	return sigma*(a**4)
 
 #################### SHOW TIME ####################
 
-# pressure_profile = profile(air_pressure)
-# density_profile = profile(air_density)
-# temperature_profile = profile(temperature_atmosp)
-
-# print(dz,density_profile,gravity,pressure_profile/1E2,temperature_profile)
-
-# plt.plot(pressure_profile/1E2,heights/1E3,color='blue',label='Model atmosphere')
-# temp = np.zeros_like(pressure_profile)
-# temp[0] = pressure_profile[0]
-# for k in np.arange(1,nlevels): temp[k] = temp[k-1] - dz[k]*density_profile[k]*gravity
-# plt.plot(temp/1E2,heights/1E3,color='red',label='Implied hydrostatic balance')
-# plt.xlabel('Pressure (hPa)')
-# plt.ylabel('Height (km)')
-# plt.legend()
-# plt.show()
-
-# plt.plot(temp/pressure_profile)
-# plt.show()
-
-# sys.exit()
-
-#################################################
-
 if plot:
 	# set up plot
-	f, ax = plt.subplots(2,figsize=(9,7))
+	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[:,:,0], cmap='seismic')
@@ -237,10 +240,9 @@ if plot:
 	# allow for live updating as calculations take place
 
 	if level_plots:
-		plot_subsample = [0,5,10,15]
-		g, bx = plt.subplots(len(plot_subsample),figsize=(9,7),sharex=True)
+		g, bx = plt.subplots(nlevels,figsize=(9,7),sharex=True)
 		g.canvas.set_window_title('CLAuDE atmospheric levels')
-		for k in plot_subsample:
+		for k in range(nlevels):
 			bx[k].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,k], cmap='seismic')
 			bx[k].set_title(str(heights[k])+' km')
 			bx[k].set_ylabel('Latitude')
@@ -260,8 +262,8 @@ while True:
 
 	initial_time = time.time()
 
-	if t < 14*day:
-		dt = 60*47
+	if t < 30*day:
+		dt = 60*137
 		velocity = False
 	else:
 		dt = 60*9
@@ -269,45 +271,51 @@ while True:
 
 	# print current time in simulation to command line
 	print("+++ t = " + str(round(t/day,2)) + " days +++", end='\r')
-	print('U:',u.max(),u.min(),'V: ',v.max(),v.min(),'W: ',w.max(),w.min())
-	# print(np.mean(np.mean(air_density,axis=0),axis=0))
+	print('T: ',round(temperature_planet.max()-273.15,1),' - ',round(temperature_planet.min()-273.15,1),' C')
+	print('U: ',round(u.max(),2),' - ',round(u.min(),2),' V: ',round(v.max(),2),' - ',round(v.min(),2),' W: ',round(w.max(),2),' - ',round(w.min(),2))
+	# print(profile(air_density))
+	# print(profile(air_pressure)/100)
 
 	# calculate change in temperature of ground and atmosphere due to radiative imbalance
 	for i in range(nlat):
 		for j in range(nlon):
 
+			upward_radiation = np.zeros(nlevels)
+			downward_radiation = np.zeros(nlevels)
+			optical_depth = np.zeros(nlevels)
+			Q = np.zeros(nlevels)
+
 			# calculate optical depth
 			pressure_profile = air_pressure[i,j,:]
 			density_profile = air_density[i,j,:]
 			fl = 0.1
 			optical_depth = surface_optical_depth(lat[i])*(fl*(pressure_profile/pressure_profile[0]) + (1-fl)*(pressure_profile/pressure_profile[0])**4)
-
+			
 			# calculate upward longwave flux, bc is thermal radiation at surface
-			upward_radiation[0] = sigma*temperature_planet[i,j]**4
+			upward_radiation[0] = thermal_radiation(temperature_planet[i,j])
 			for k in np.arange(1,nlevels):
-				upward_radiation[k] = upward_radiation[k-1] + (optical_depth[k]-optical_depth[k-1])*(upward_radiation[k-1] - sigma*temperature_atmosp[i,j,k]**4)
+				upward_radiation[k] = (upward_radiation[k-1] - (optical_depth[k]-optical_depth[k-1])*(thermal_radiation(temperature_atmosp[i,j,k])))/(1+optical_depth[k-1]-optical_depth[k])
 
 			# calculate downward longwave flux, bc is zero at TOA (in model)
-			downward_radiation[nlevels-1] = 0
+			downward_radiation[-1] = 0
 			for k in np.arange(0,nlevels-1)[::-1]:
-				if k == 0:
-					downward_radiation[k] = downward_radiation[k+1] + (optical_depth[k+1]-optical_depth[k])*(sigma*temperature_atmosp[i,j,k]**4 - downward_radiation[k+1])
-				else:
-					downward_radiation[k] = downward_radiation[k+1] + (optical_depth[k]-optical_depth[k-1])*(sigma*temperature_atmosp[i,j,k]**4 - downward_radiation[k+1])
+				downward_radiation[k] = (downward_radiation[k+1] - thermal_radiation(temperature_atmosp[i,j,k])*(optical_depth[k+1]-optical_depth[k]))/(1 + optical_depth[k] - optical_depth[k+1])
+			
 			# gradient of difference provides heating at each level
 			for k in np.arange(nlevels):
 				Q[k] = -scalar_gradient_z(upward_radiation-downward_radiation,0,0,k)/(1E3*density_profile[k])
-			temperature_atmosp[i,j,:] += Q
+				# make sure model does not have a higher top than 50km!!
+				# approximate SW heating of ozone
+				if heights[k] > 20E3:
+					Q[k] += solar(5,lat[i],lon[j],t)*((((heights[k]-20E3)/1E3)**2)/(30**2))/(24*60*60)
+
+			temperature_atmosp[i,j,:] += Q*dt
 
 			# update surface temperature with shortwave radiation flux
-			temperature_planet[i,j] += dt*((1-albedo[i,j])*solar(insolation,lat[i],lon[j],t) + scalar_gradient_z(downward_radiation,0,0,0) - sigma*temperature_planet[i,j]**4)/heat_capacity_earth[i,j]
-
-			# plt.plot(downward_radiation,heights,label='downward',color='blue')
-			# plt.plot(upward_radiation,heights,label='upward',color='red')
-			# plt.legend()
-			# plt.show()
+			temperature_planet[i,j] += dt*((1-albedo[i,j])*(solar(insolation,lat[i],lon[j],t) + downward_radiation[0]) - upward_radiation[0])/heat_capacity_earth[i,j] 
 
 	# update air pressure
+	old_pressure = np.copy(air_pressure)
 	air_pressure = air_density*specific_gas*temperature_atmosp
 
 	if velocity:
@@ -323,16 +331,21 @@ while True:
 				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] += dt*( (air_pressure[i,j,k] - np.trapz(y=air_density[i,j,k:]*gravity,dx=dz[(k+1):]) - weight_above)/(1E5*dz[k]*air_density[i,j,k]) )
+					w_temp[i,j,k] += -(air_pressure[i,j,k]-old_pressure[i,j,k])/(dt*air_density[i,j,k]*gravity)
+
+				# plt.plot(w_temp[i,j,:],heights)
+				# plt.title('Vertical acceleration')
+				# plt.show()
+				# sys.exit()
+
 
 		u += u_temp
 		v += v_temp
-		w[:,:,1:-2] += w_temp[:,:,1:-2]
+		w += w_temp
 
-		# approximate surface friction
-		u[:,:,0] *= 0.99
-		v[:,:,0] *= 0.99
+		# approximate friction
+		u *= 0.95
+		v *= 0.95
 
 		if advection:
 			# allow for thermal advection in the atmosphere, and heat diffusion in the atmosphere and the ground
@@ -353,6 +366,7 @@ while True:
 	if plot:
 		# update plot
 		test = ax[0].contourf(lon_plot, lat_plot, temperature_planet, cmap='seismic')
+		ax[0].streamplot(lon_plot, lat_plot, u[:,:,0], v[:,:,0], color='white',density=1)
 		ax[0].set_title('$\it{Ground} \quad \it{temperature}$')
 		ax[0].set_xlim((lon.min(),lon.max()))
 		ax[0].set_ylim((lat.min(),lat.max()))
@@ -361,8 +375,8 @@ while True:
 		ax[0].set_xlabel('Longitude')
 
 		ax[1].contourf(heights_plot, lat_z_plot, np.transpose(np.mean(temperature_atmosp,axis=1)), cmap='seismic')
-		ax[1].contour(heights_plot,lat_z_plot, np.transpose(np.mean(u,axis=1)), colors='white')
-		ax[1].streamplot(heights_plot, lat_z_plot, np.transpose(np.mean(v,axis=1)),np.transpose(np.mean(w,axis=1)),color='black',density=0.75)
+		ax[1].contour(heights_plot,lat_z_plot, np.transpose(np.mean(u,axis=1)), colors='white',levels=20,linewidths=1,alpha=0.8)
+		ax[1].streamplot(heights_plot, lat_z_plot, np.transpose(np.mean(v,axis=1)),np.transpose(np.mean(10*w,axis=1)),color='black',density=0.75)
 		ax[1].set_title('$\it{Atmospheric} \quad \it{temperature}$')
 		ax[1].set_xlim((-90,90))
 		ax[1].set_ylim((0,heights.max()))
@@ -375,10 +389,10 @@ while True:
 		f.suptitle( 'Time ' + str(round(24*t/day,2)) + ' hours' )
 		
 		if level_plots:
-			for k in plot_subsample:
+			for k in range(nlevels):
 				index = nlevels-1-k
 				test = bx[index].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,k], cmap='seismic')
-				bx[index].streamplot(lon_plot, lat_plot, u[:,:,k], v[:,:,k],density=0.75,color='black')
+				bx[index].streamplot(lon_plot, lat_plot, u[:,:,k], v[:,:,k],density=0.5,color='black')
 				# g.colorbar(test,cax=bx[nlevels-1-k])
 				bx[index].set_title(str(heights[k]/1E3)+' km')
 				bx[index].set_ylabel('Latitude')