2_yr_spinup_eqpolediff_T42.py

import climt
import sympl
from datetime import timedelta

import os
os.system("rm 2_yr_spinup_eqpolediff_T42_37vlvl_3damped.nc")

### -----------------------
# 10th Dec 2019
# Script to initialise a simple dry gcm from rest and spin up for one year, 
# saving all fields.
# Include 300k-240k eq pole temp diff
"""TODO: Update with shortwave radiation and climt.SlabSurface(), for an aquaplanet!!"""
### -----------------------

# Set timestep
timestep = timedelta(minutes=20)

# Convection
convection = climt.EmanuelConvection()

# Radiation: only call every 1 hr
radiation_timestep = timedelta(minutes=60)
radiation = sympl.UpdateFrequencyWrapper(climt.GrayLongwaveRadiation(), 
                                         radiation_timestep)
# Simple physics (boundary layer): use tendencies to work better in spectral space
simple_physics = sympl.TimeDifferencingWrapper(climt.SimplePhysics())

# Set up grid (3degx3deg, 15 v levels)
grid = climt.get_grid(nx=128, ny=64, nz=37, x_name='lon', y_name='lat')

dycore = climt.GFSDynamicalCore(
        [simple_physics, radiation,
         convection], number_of_damped_levels=3
    )

state = climt.get_default_state([dycore], grid_state=grid)

# Set initial/boundary conditions
latitudes = state[('latitude')].values
longitudes = state[('longitude')].values
surface_shape = latitudes.shape

temperature_equator = 300
temperature_pole = 240

temperature_profile = temperature_equator - (
    (temperature_equator - temperature_pole)*(
        np.sin(np.radians(latitudes))**2))

state[('surface_temperature')] = sympl.DataArray(
    temperature_profile*np.ones(surface_shape),
    dims=['lat', 'lon'], attrs={'units': 'degK'})

# Initialise saving object
fields_to_store = ['air_temperature', 'air_pressure', 'eastward_wind',
                   'northward_wind', 'surface_pressure', 'atmosphere_relative_vorticity',
                   'specific_humidity', 'surface_temperature',
                   'latitude', 'longitude', 'atmosphere_convective_available_potential_energy',
                   'convective_precipitation_rate', 'stratiform_precipitation_rate'
                   'convective_heating_rate']

netcdf_monitor = sympl.NetCDFMonitor('2_yr_spinup_eqpolediff_T42_37vlvl_3damped.nc',
                                    store_names=fields_to_store,
                                    write_on_store=True)

# Step forward
diagnostics, new_state = dycore(state, timestep)
state.update(diagnostics)
state.update(new_state)
state['time'] += timestep

#%%time
for step in range(2*24*3*365):
    diagnostics, new_state = dycore(state, timestep)
    state.update(new_state)
    state['time'] += timestep
    if step%(24*3) == 0:
        print(state['time'])
        netcdf_monitor.store(state)
print("Finished at ", state['time'])