Cythonised all the things!

Used cython to improve speed by factor of 10, fixed bug in advection, decamped functions to claude libraries

claude_low_level_library.py

@@ -0,0 +1,84 @@
+# claude low level library
+
+import numpy as np
+cimport numpy as np
+
+ctypedef np.float64_t DTYPE_f
+sigma = 5.67E-8
+
+# define various useful differential functions:
+# gradient of scalar field a in the local x direction at point i,j
+def scalar_gradient_x(a,dx,nlon,i,j,k):
+	return (a[i,(j+1)%nlon,k]-a[i,(j-1)%nlon,k])/dx[i]
+
+def scalar_gradient_x_2D(a,dx,nlon,i,j)
+	return (a[i,(j+1)%nlon]-a[i,(j-1)%nlon])/dx[i]
+
+# gradient of scalar field a in the local y direction at point i,j
+def scalar_gradient_y(a,dy,nlat,i,j,k):
+	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
+
+def scalar_gradient_y_2d(a,dy,nlat,i,j)
+	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
+
+def scalar_gradient_z(a,dz,i,j,k):
+	output = np.zeros_like(a)
+	nlevels = len(dz)
+	if output.ndim == 1:
+		if k == 0:
+			return (a[k+1]-a[k])/dz[k]
+		elif k == nlevels-1:
+			return (a[k]-a[k-1])/dz[k]
+		else:
+			return (a[k+1]-a[k-1])/(2*dz[k])
+	else:
+		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])
+
+def surface_optical_depth(lat):
+	return 4 + np.cos(lat*np.pi/90)*2.5/2
+
+def thermal_radiation(a):
+	return sigma*(a**4)
+
+# power incident on (lat,lon) at time t
+def solar(insolation, lat, lon, t, day, year, axial_tilt):
+	sun_longitude = -t % day
+	sun_longitude *= 360/day
+	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 profile(a):
+	return np.mean(np.mean(a,axis=0),axis=0)

claude_low_level_library.pyx

@@ -0,0 +1,89 @@
+# claude low level library
+
+import numpy as np
+cimport numpy as np
+cimport cython
+
+ctypedef np.float64_t DTYPE_f
+cdef float inv_180 = np.pi/180
+cdef float inv_90 = np.pi/90
+
+# define various useful differential functions:
+# gradient of scalar field a in the local x direction at point i,j
+def scalar_gradient_x(np.ndarray a,np.ndarray dx,np.int_t nlon,np.int_t i,np.int_t j,np.int_t k):
+	return (a[i,(j+1)%nlon,k]-a[i,(j-1)%nlon,k])/dx[i]
+
+def scalar_gradient_x_2D(a,dx,nlon,i,j):
+	return (a[i,(j+1)%nlon]-a[i,(j-1)%nlon])/dx[i]
+
+# gradient of scalar field a in the local y direction at point i,j
+def scalar_gradient_y(np.ndarray a,DTYPE_f dy,np.int_t nlat,np.int_t i,np.int_t j,np.int_t k):
+	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
+
+def scalar_gradient_y_2d(a,dy,nlat,i,j):
+	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
+
+def scalar_gradient_z(a,dz,i,j,k):
+	cdef np.ndarray output = np.zeros_like(a)
+	cdef np.int_t nlevels = len(dz)
+	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])
+
+def scalar_gradient_z_1D(np.ndarray a,np.ndarray dz,np.int_t i,np.int_t j,np.int_t k):
+	cdef np.ndarray output = np.zeros_like(a)
+	cdef np.int_t nlevels = len(dz)
+	if k == 0:
+		return (a[k+1]-a[k])/dz[k]
+	elif k == nlevels-1:
+		return (a[k]-a[k-1])/dz[k]
+	else:
+		return (a[k+1]-a[k-1])/(2*dz[k])
+
+def surface_optical_depth(DTYPE_f lat):
+	return 4 + np.cos(lat*inv_90)*2.5
+
+def thermal_radiation(DTYPE_f a):
+	cdef DTYPE_f sigma = 5.67E-8
+	return sigma*(a**4)
+
+# power incident on (lat,lon) at time t
+def solar(DTYPE_f insolation,DTYPE_f  lat,DTYPE_f lon,np.int_t t,DTYPE_f  day,DTYPE_f  year,DTYPE_f  axial_tilt):
+	cdef float sun_longitude = -t % day
+	cdef float sun_latitude = axial_tilt*np.cos(t*2*np.pi/year)
+	cdef float value = insolation*np.cos((lat-sun_latitude)*inv_180)
+	cdef float lon_diff, cos_lon
+
+	if value < 0:	
+		return 0
+	else:
+		sun_longitude *= 360/day
+		lon_diff = lon-sun_longitude
+		cos_lon = np.cos(lon_diff*inv_180) 
+		value *= cos_lon
+		
+		if value < 0:
+			if lat + sun_latitude > 90:
+				return insolation*np.cos((lat+sun_latitude)*inv_180)*cos_lon
+			elif lat + sun_latitude < -90:
+				return insolation*np.cos((lat+sun_latitude)*inv_180)*cos_lon
+			else:
+				return 0
+		else:
+			return value
+
+def profile(a):
+	return np.mean(np.mean(a,axis=0),axis=0)

claude_setup.py

@@ -0,0 +1,8 @@
+from setuptools import setup
+from Cython.Build import cythonize
+import numpy
+
+setup(
+	include_dirs=[numpy.get_include()],
+    ext_modules = cythonize("*.pyx")
+)

claude_top_level_library.py

@@ -0,0 +1,101 @@
+# claude_top_level_library
+
+import claude_low_level_library as low_level
+import numpy as np
+
+# laplacian of scalar field a
+def laplacian(a):
+	output = np.zeros_like(a)
+	if output.ndim == 2:
+		for i in np.arange(1,nlat-1):
+			for j in range(nlon):
+				output[i,j] = (scalar_gradient_x_2D(a,i,(j+1)%nlon) - scalar_gradient_x_2D(a,i,(j-1)%nlon))/dx[i] + (scalar_gradient_y_2D(a,i+1,j) - scalar_gradient_y_2D(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,u,v,dx,dy):
+	output = np.zeros_like(a)
+	nlat, nlon, nlevels = output.shape[:]
+
+	au = a*u
+	av = a*v
+
+	for i in range(nlat):
+		for j in range(nlon):
+			for k in range(nlevels):
+				output[i,j,k] = low_level.scalar_gradient_x(au,dx,nlon,i,j,k) + low_level.scalar_gradient_y(av,dy,nlat,i,j,k) #+ 0.1*scalar_gradient_z(a*w,i,j,k)
+	return output
+
+def radiation_calculation(temperature_world, temperature_atmos, air_pressure, air_density, heat_capacity_earth, albedo, insolation, lat, lon, heights, dz, t, dt, day, year, axial_tilt):
+	# calculate change in temperature of ground and atmosphere due to radiative imbalance
+	nlat, nlon, nlevels = temperature_atmos.shape[:]	
+
+	upward_radiation = np.zeros(nlevels)
+	downward_radiation = np.zeros(nlevels)
+	optical_depth = np.zeros(nlevels)
+	Q = np.zeros(nlevels)
+
+	for i in range(nlat):
+		for j in range(nlon):
+			# calculate optical depth
+			pressure_profile = air_pressure[i,j,:]
+			density_profile = air_density[i,j,:]
+			fl = 0.1
+			optical_depth = low_level.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] = low_level.thermal_radiation(temperature_world[i,j])
+			for k in np.arange(1,nlevels):
+				upward_radiation[k] = (upward_radiation[k-1] - (optical_depth[k]-optical_depth[k-1])*(low_level.thermal_radiation(temperature_atmos[i,j,k])))/(1+optical_depth[k-1]-optical_depth[k])
+
+			# calculate downward longwave flux, bc is zero at TOA (in model)
+			downward_radiation[-1] = 0
+			for k in np.arange(0,nlevels-1)[::-1]:
+				downward_radiation[k] = (downward_radiation[k+1] - low_level.thermal_radiation(temperature_atmos[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] = -low_level.scalar_gradient_z_1D(upward_radiation-downward_radiation,dz,0,0,k)/(1E3*density_profile[k])
+				# make sure model does not have a higher top than 50km!!
+				# approximate SW heating of ozone
+				if heights[k] > 20E3:
+					Q[k] += low_level.solar(5,lat[i],lon[j],t,day, year, axial_tilt)*((((heights[k]-20E3)/1E3)**2)/(30**2))/(24*60*60)
+
+			temperature_atmos[i,j,:] += Q*dt
+
+			# update surface temperature with shortwave radiation flux
+			temperature_world[i,j] += dt*((1-albedo[i,j])*(low_level.solar(insolation,lat[i],lon[j],t, day, year, axial_tilt) + downward_radiation[0]) - upward_radiation[0])/heat_capacity_earth[i,j] 
+	
+	return temperature_world, temperature_atmos
+
+def velocity_calculation(u,v,air_pressure,old_pressure,air_density,coriolis,gravity,dx,dy,dt):
+	# 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(u)
+
+	nlat,nlon,nlevels = air_pressure.shape[:]
+
+	# calculate acceleration of atmosphere using primitive equations on beta-plane
+	for i in np.arange(1,nlat-1):
+		for j in range(nlon):
+			for k in range(nlevels):
+				u_temp[i,j,k] += dt*( -u[i,j,k]*low_level.scalar_gradient_x(u,dx,nlon,i,j,k) - v[i,j,k]*low_level.scalar_gradient_y(u,dy,nlat,i,j,k) + coriolis[i]*v[i,j,k] - low_level.scalar_gradient_x(air_pressure,dx,nlon,i,j,k)/air_density[i,j,k] )
+				v_temp[i,j,k] += dt*( -u[i,j,k]*low_level.scalar_gradient_x(v,dx,nlon,i,j,k) - v[i,j,k]*low_level.scalar_gradient_y(v,dy,nlat,i,j,k) - coriolis[i]*u[i,j,k] - low_level.scalar_gradient_y(air_pressure,dy,nlat,i,j,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)
+
+	u += u_temp
+	v += v_temp
+	w = w_temp
+
+	# approximate friction
+	u *= 0.95
+	v *= 0.95
+
+	return u,v,w

claude_top_level_library.pyx

@@ -0,0 +1,110 @@
+# claude_top_level_library
+
+import claude_low_level_library as low_level
+import numpy as np
+cimport numpy as np
+cimport cython
+
+ctypedef np.float64_t DTYPE_f
+
+# laplacian of scalar field a
+# def laplacian(a):
+# 	output = np.zeros_like(a)
+# 	if output.ndim == 2:
+# 		for i in np.arange(1,nlat-1):
+# 			for j in range(nlon):
+# 				output[i,j] = (scalar_gradient_x_2D(a,i,(j+1)%nlon) - scalar_gradient_x_2D(a,i,(j-1)%nlon))/dx[i] + (scalar_gradient_y_2D(a,i+1,j) - scalar_gradient_y_2D(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,u,v,dx,dy):
+	output = np.zeros_like(a)
+	nlat, nlon, nlevels = output.shape[:]
+
+	au = a*u
+	av = a*v
+
+	for i in range(nlat):
+		for j in range(nlon):
+			for k in range(nlevels):
+				output[i,j,k] = low_level.scalar_gradient_x(au,dx,nlon,i,j,k) + low_level.scalar_gradient_y(av,dy,nlat,i,j,k) #+ 0.1*scalar_gradient_z(a*w,i,j,k)
+	return output
+
+def radiation_calculation(np.ndarray temperature_world, np.ndarray temperature_atmos, np.ndarray air_pressure, np.ndarray air_density, np.ndarray heat_capacity_earth, np.ndarray albedo, DTYPE_f insolation, np.ndarray lat, np.ndarray lon, np.ndarray heights, np.ndarray dz, np.int_t t, np.int_t dt, DTYPE_f day, DTYPE_f year, DTYPE_f axial_tilt):
+	# calculate change in temperature of ground and atmosphere due to radiative imbalance
+	cdef np.int_t nlat,nlon,nlevels,i,j
+	cdef DTYPE_f fl = 0.1
+	cdef np.ndarray upward_radiation,downward_radiation,optical_depth,Q, pressure_profile, density_profile
+
+	nlat = temperature_atmos.shape[0]	
+	nlon = temperature_atmos.shape[1]
+	nlevels = temperature_atmos.shape[2]
+
+	upward_radiation = np.zeros(nlevels)
+	downward_radiation = np.zeros(nlevels)
+	optical_depth = np.zeros(nlevels)
+	Q = np.zeros(nlevels)
+
+	for i in range(nlat):
+		for j in range(nlon):
+			# calculate optical depth
+			pressure_profile = air_pressure[i,j,:]
+			density_profile = air_density[i,j,:]
+			optical_depth = low_level.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] = low_level.thermal_radiation(temperature_world[i,j])
+			for k in np.arange(1,nlevels):
+				upward_radiation[k] = (upward_radiation[k-1] - (optical_depth[k]-optical_depth[k-1])*(low_level.thermal_radiation(temperature_atmos[i,j,k])))/(1+optical_depth[k-1]-optical_depth[k])
+
+			# calculate downward longwave flux, bc is zero at TOA (in model)
+			downward_radiation[-1] = 0
+			for k in np.arange(0,nlevels-1)[::-1]:
+				downward_radiation[k] = (downward_radiation[k+1] - low_level.thermal_radiation(temperature_atmos[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] = -low_level.scalar_gradient_z_1D(upward_radiation-downward_radiation,dz,0,0,k)/(1E3*density_profile[k])
+				# make sure model does not have a higher top than 50km!!
+				# approximate SW heating of ozone
+				if heights[k] > 20E3:
+					Q[k] += low_level.solar(5,lat[i],lon[j],t,day, year, axial_tilt)*((((heights[k]-20E3)/1E3)**2)/(30**2))/(24*60*60)
+
+			temperature_atmos[i,j,:] += Q*dt
+
+			# update surface temperature with shortwave radiation flux
+			temperature_world[i,j] += dt*((1-albedo[i,j])*(low_level.solar(insolation,lat[i],lon[j],t, day, year, axial_tilt) + downward_radiation[0]) - upward_radiation[0])/heat_capacity_earth[i,j] 
+	
+	return temperature_world, temperature_atmos
+
+def velocity_calculation(u,v,air_pressure,old_pressure,air_density,coriolis,gravity,dx,dy,dt):
+	# 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(u)
+
+	nlat,nlon,nlevels = air_pressure.shape[:]
+
+	# calculate acceleration of atmosphere using primitive equations on beta-plane
+	for i in np.arange(1,nlat-1):
+		for j in range(nlon):
+			for k in range(nlevels):
+				u_temp[i,j,k] += dt*( -u[i,j,k]*low_level.scalar_gradient_x(u,dx,nlon,i,j,k) - v[i,j,k]*low_level.scalar_gradient_y(u,dy,nlat,i,j,k) + coriolis[i]*v[i,j,k] - low_level.scalar_gradient_x(air_pressure,dx,nlon,i,j,k)/air_density[i,j,k] )
+				v_temp[i,j,k] += dt*( -u[i,j,k]*low_level.scalar_gradient_x(v,dx,nlon,i,j,k) - v[i,j,k]*low_level.scalar_gradient_y(v,dy,nlat,i,j,k) - coriolis[i]*u[i,j,k] - low_level.scalar_gradient_y(air_pressure,dy,nlat,i,j,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)
+
+	u += u_temp
+	v += v_temp
+	w = w_temp
+
+	# approximate friction
+	u *= 0.95
+	v *= 0.95
+
+	return u,v,w

toy_model.py

@@ -6,6 +6,8 @@
 import numpy as np 
 import matplotlib.pyplot as plt
 import time, sys, pickle
+import claude_low_level_library as low_level
+import claude_top_level_library as top_level
 
 ######## CONTROL ########
 
@@ -19,16 +21,21 @@ 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)
 
+dt_spinup = 60*137
+dt_main = 60*9
+spinup_length = 0*day
+
 ###
 
 advection = True 				# if you want to include advection set this to be True
-advection_boundary = 4			# how many gridpoints away from poles to apply advection
+advection_boundary = 8			# how many gridpoints away from poles to apply advection
 
-save = True
+save = False 					# save current state to file?
-load = False
+load = False  					# load initial state from file?
 
-plot = False
+plot = True 					# display plots of output?
-level_plots = False
+level_plots = True 				# display plots of output on vertical levels?
+nplots = 5						# how many levels you want to see plots of (evenly distributed through column)
 
 ###########################
 
@@ -42,19 +49,16 @@ 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)) + 290
+temperature_world = np.zeros((nlat,nlon)) + 290
-temperature_atmosp = np.zeros((nlat,nlon,nlevels))
+temperature_atmos = np.zeros((nlat,nlon,nlevels))
-air_pressure = np.zeros_like(temperature_atmosp)
+air_pressure = np.zeros_like(temperature_atmos)
-u = np.zeros_like(temperature_atmosp)
+u = np.zeros_like(temperature_atmos)
-v = np.zeros_like(temperature_atmosp)
+v = np.zeros_like(temperature_atmos)
-w = np.zeros_like(temperature_atmosp)
+w = np.zeros_like(temperature_atmos)
-air_density = np.zeros_like(temperature_atmosp)
+air_density = np.zeros_like(temperature_atmos)
 
 # #######################
 
-def profile(a):
-	return np.mean(np.mean(a,axis=0),axis=0)
-
 # read temperature and density in from standard atmosphere
 f = open("standard_atmosphere.txt", "r")
 standard_height = []
@@ -71,7 +75,7 @@ density_profile = np.interp(x=heights/1E3,xp=standard_height,fp=standard_density
 temp_profile = np.interp(x=heights/1E3,xp=standard_height,fp=standard_temp)
 for k in range(nlevels):
 	air_density[:,:,k] = density_profile[k]
-	temperature_atmosp[:,:,k] = temp_profile[k]
+	temperature_atmos[:,:,k] = temp_profile[k]
 
 ###########################
 
@@ -84,8 +88,8 @@ for k in range(nlevels):
 
 ###########################
 
-albedo = np.zeros_like(temperature_planet) + 0.2
+albedo = np.zeros_like(temperature_world) + 0.2
-heat_capacity_earth = np.zeros_like(temperature_planet) + 1E6
+heat_capacity_earth = np.zeros_like(temperature_world) + 1E6
 
 albedo_variance = 0.001
 for i in range(nlat):
@@ -96,7 +100,7 @@ specific_gas = 287
 thermal_diffusivity_roc = 1.5E-6
 sigma = 5.67E-8
 
-air_pressure = air_density*specific_gas*temperature_atmosp
+air_pressure = air_density*specific_gas*temperature_atmos
 
 # define planet size and various geometric constants
 circumference = 2*np.pi*planet_radius
@@ -115,136 +119,32 @@ 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,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,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:
-		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
-
-def scalar_gradient_z(a,i,j,k):
-	output = np.zeros_like(a)
-	if output.ndim == 1:
-		if k == 0:
-			return (a[k+1]-a[k])/dz[k]
-		elif k == nlevels-1:
-			return (a[k]-a[k-1])/dz[k]
-		else:
-			return (a[k+1]-a[k-1])/(2*dz[k])
-	else:
-		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)
-	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(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) #+ 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
-	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 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 ####################
 
 if plot:
 	# set up plot
 	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[0].contourf(lon_plot, lat_plot, temperature_world, cmap='seismic')
-	ax[1].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,0], cmap='seismic')
+	ax[1].contourf(lon_plot, lat_plot, temperature_atmos[:,:,0], cmap='seismic')
 	plt.subplots_adjust(left=0.1, right=0.75)
 	ax[0].set_title('Ground temperature')
 	ax[1].set_title('Atmosphere temperature')
+	cbar_ax = f.add_axes([0.85, 0.15, 0.05, 0.7])
+
 	# allow for live updating as calculations take place
 
 	if level_plots:
-		g, bx = plt.subplots(nlevels,figsize=(9,7),sharex=True)
+
+		level_divisions = int(np.floor(nlevels/nplots))
+		level_plots_levels = range(nlevels)[::level_divisions]
+
+		g, bx = plt.subplots(nplots,figsize=(9,8),sharex=True)
 		g.canvas.set_window_title('CLAuDE atmospheric levels')
-		for k in range(nlevels):
+		for k, z in zip(range(nplots), level_plots_levels):	
-			bx[k].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,k], cmap='seismic')
+			z += 1
-			bx[k].set_title(str(heights[k])+' km')
+			bx[k].contourf(lon_plot, lat_plot, temperature_atmos[:,:,z], cmap='seismic')
+			bx[k].set_title(str(heights[z])+' km')
 			bx[k].set_ylabel('Latitude')
 		bx[-1].set_xlabel('Longitude')
 
@@ -256,116 +156,67 @@ t = 0
 
 if load:
 	# load in previous save file
-	temperature_atmosp,temperature_planet,u,v,w,t,air_density,albedo = pickle.load(open("save_file.p","rb"))
+	temperature_atmos,temperature_world,u,v,w,t,air_density,albedo = pickle.load(open("save_file.p","rb"))
 
 while True:
 
 	initial_time = time.time()
 
-	if t < 30*day:
+	if t < spinup_length:
-		dt = 60*137
+		dt = dt_spinup
 		velocity = False
 	else:
-		dt = 60*9
+		dt = dt_main
 		velocity = True
 
 	# print current time in simulation to command line
 	print("+++ t = " + str(round(t/day,2)) + " days +++", end='\r')
-	print('T: ',round(temperature_planet.max()-273.15,1),' - ',round(temperature_planet.min()-273.15,1),' C')
+	print('T: ',round(temperature_world.max()-273.15,1),' - ',round(temperature_world.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
+	before_radiation = time.time()
-	for i in range(nlat):
+	temperature_world, temperature_atmos = top_level.radiation_calculation(temperature_world, temperature_atmos, air_pressure, air_density, heat_capacity_earth, albedo, insolation, lat, lon, heights, dz, t, dt, day, year, axial_tilt)
-		for j in range(nlon):
+	time_taken = float(round(time.time() - before_radiation,3))
-
+	# print('Radiation: ',str(time_taken),'s')
-			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] = 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])*(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[-1] = 0
-			for k in np.arange(0,nlevels-1)[::-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])
-				# 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) + 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
+	air_pressure = air_density*specific_gas*temperature_atmos
 
 	if velocity:
 
-		# introduce temporary arrays to update velocity in the atmosphere
+		before_velocity = time.time()
-		u_temp = np.zeros_like(u)
+		u,v,w = top_level.velocity_calculation(u,v,air_pressure,old_pressure,air_density,coriolis,gravity,dx,dy,dt)
-		v_temp = np.zeros_like(v)
+		time_taken = float(round(time.time() - before_velocity,3))
-		w_temp = np.zeros_like(w)
+		# print('Velocity: ',str(time_taken),'s')
-
-		# calculate acceleration of atmosphere using primitive equations on beta-plane
-		for i in np.arange(1,nlat-1):
-			for j in range(nlon):
-				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] += -(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 += w_temp
-
-		# approximate friction
-		u *= 0.95
-		v *= 0.95
 
+		before_advection = time.time()
 		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))
+			# atmosp_addition = dt*(thermal_diffusivity_air*laplacian(temperature_atmos))
-			atmosp_addition = dt*divergence_with_scalar(temperature_atmosp)
+
-			temperature_atmosp[advection_boundary:-advection_boundary,:,:] -= atmosp_addition[advection_boundary:-advection_boundary,:,:]
+			atmosp_addition = dt*top_level.divergence_with_scalar(temperature_atmos,u,v,dx,dy)
-			temperature_atmosp[advection_boundary-1,:,:] -= 0.5*atmosp_addition[advection_boundary-1,:,:]
+			temperature_atmos[advection_boundary:-advection_boundary,:,:] -= atmosp_addition[advection_boundary:-advection_boundary,:,:]
-			temperature_atmosp[-advection_boundary,:,:] -= 0.5*atmosp_addition[-advection_boundary,:,:]
+			temperature_atmos[advection_boundary-1,:,:] -= 0.5*atmosp_addition[advection_boundary-1,:,:]
+			temperature_atmos[-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)
+			# density_addition = dt*divergence_with_scalar(air_density)
-			air_density[advection_boundary:-advection_boundary,:,:] -= density_addition[advection_boundary:-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-1),:,:] -= 0.5*density_addition[advection_boundary-1,:,:]
-			air_density[-advection_boundary,:,:] -= 0.5*density_addition[-advection_boundary,:,:]
+			# air_density[-advection_boundary,:,:] -= 0.5*density_addition[-advection_boundary,:,:]
 
-			temperature_planet += dt*(thermal_diffusivity_roc*laplacian(temperature_planet))
+			# temperature_world += dt*(thermal_diffusivity_roc*laplacian(temperature_world))
 			
+			time_taken = float(round(time.time() - before_advection,3))
+			# print('Advection: ',str(time_taken),'s')
+	
+	before_plot = time.time()
 	if plot:
 		# update plot
-		test = ax[0].contourf(lon_plot, lat_plot, temperature_planet, cmap='seismic')
+		test = ax[0].contourf(lon_plot, lat_plot, temperature_world, 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()))
@@ -374,7 +225,7 @@ while True:
 		ax[0].axhline(y=0,color='black',alpha=0.3)
 		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].contourf(heights_plot, lat_z_plot, np.transpose(np.mean(temperature_atmos,axis=1)), cmap='seismic')
 		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}$')
@@ -383,30 +234,31 @@ while True:
 		ax[1].set_ylabel('Height (m)')
 		ax[1].set_xlabel('Latitude')
 
-		cbar_ax = f.add_axes([0.85, 0.15, 0.05, 0.7])
 		f.colorbar(test, cax=cbar_ax)
 		cbar_ax.set_title('Temperature (K)')
 		f.suptitle( 'Time ' + str(round(24*t/day,2)) + ' hours' )
 		
 		if level_plots:
-			for k in range(nlevels):
+			for k, z in zip(range(nplots), level_plots_levels):	
-				index = nlevels-1-k
+				z += 1
-				test = bx[index].contourf(lon_plot, lat_plot, temperature_atmosp[:,:,k], cmap='seismic')
+				bx[k].contourf(lon_plot, lat_plot, temperature_atmos[:,:,z], cmap='seismic')
-				bx[index].streamplot(lon_plot, lat_plot, u[:,:,k], v[:,:,k],density=0.5,color='black')
+				bx[k].streamplot(lon_plot, lat_plot, u[:,:,z], v[:,:,z], color='white',density=1)
-				# g.colorbar(test,cax=bx[nlevels-1-k])
+				bx[k].set_title(str(round(heights[z]/1E3))+' km')
-				bx[index].set_title(str(heights[k]/1E3)+' km')
+				bx[k].set_ylabel('Latitude')
-				bx[index].set_ylabel('Latitude')
+				bx[k].set_xlim((lon.min(),lon.max()))
-				bx[index].set_xlim((lon.min(),lon.max()))
+				bx[k].set_ylim((lat.min(),lat.max()))				
-				bx[index].set_ylim((lat.min(),lat.max()))
 			bx[-1].set_xlabel('Longitude')
 
 		plt.pause(0.01)
 		
 		ax[0].cla()
 		ax[1].cla()
+
 		if level_plots:
-			for k in range(nlevels):
+			for k in range(nplots):
 				bx[k].cla()
+	time_taken = float(round(time.time() - before_plot,3))
+	# print('Plotting: ',str(time_taken),'s')
 
 	# advance time by one timestep
 	t += dt
@@ -416,6 +268,6 @@ while True:
 	print('Time: ',str(time_taken),'s')
 
 	if save:
-		pickle.dump((temperature_atmosp,temperature_planet,u,v,w,t,air_density,albedo), open("save_file.p","wb"))
+		pickle.dump((temperature_atmos,temperature_world,u,v,w,t,air_density,albedo), open("save_file.p","wb"))
 
 	# sys.exit()