from __future__ import division, print_function, absolute_import
import numpy as np
import warnings
from climlab import constants as const
from climlab.radiation.radiation import _Radiation_LW
from climlab.domain import Field, Axis, domain
from .utils import _prepare_general_arguments
from .utils import _climlab_to_rrtm, _rrtm_to_climlab
# These values will get overridden by reading from Fortran extension
nbndlw = 1; ngptlw = 1
try:
# The compiled fortran extension module
from climlab_rrtmg import rrtmg_lw as _rrtmg_lw
nbndlw = int(_rrtmg_lw.parrrtm.nbndlw)
ngptlw = int(_rrtmg_lw.parrrtm.ngptlw)
# Initialize absorption data
_rrtmg_lw.climlab_rrtmg_lw_ini(const.cp)
except:
warnings.warn('Cannot import and initialize compiled Fortran extension, RRTMG_LW module will not be functional.')
# Longwave spectral band limits (wavenumbers in cm^-1)
# Data copied from rrtmg_lw_v4.85/gcm_model/src/rrtmg_lw_init.f90
wavenum_bounds = np.array([ 10., 350., 500., 630., 700., 820.,
980.,1080.,1180.,1390.,1480.,1800.,
2080.,2250.,2380.,2600.,3250.])
wavenum_delta = np.diff(wavenum_bounds)
wavenum_ax = Axis(axis_type='abstract', bounds=wavenum_bounds)
[docs]
class RRTMG_LW(_Radiation_LW):
def __init__(self,
# GENERAL, used in both SW and LW
icld = 1, # Cloud overlap method, 0: Clear only, 1: Random, 2, Maximum/random] 3: Maximum
irng = 1, # more monte carlo stuff
idrv = 0, # whether to also calculate the derivative of flux with respect to surface temp
permuteseed = 300,
inflglw = 2,
iceflglw = 1,
liqflglw = 1,
tauc = 0., # in-cloud optical depth
tauaer = 0., # Aerosol optical depth at mid-point of LW spectral bands
return_spectral_olr = False, # Whether or not to return OLR averaged over each band
**kwargs):
super(RRTMG_LW, self).__init__(**kwargs)
# define INPUTS
self.add_input('icld', icld)
self.add_input('irng', irng)
self.add_input('idrv', idrv)
self.add_input('permuteseed', permuteseed)
self.add_input('inflglw', inflglw)
self.add_input('iceflglw', iceflglw)
self.add_input('liqflglw', liqflglw)
self.add_input('tauc', tauc)
self.add_input('tauaer', self._spectral_field(tauaer))
self.add_input('return_spectral_olr', return_spectral_olr)
# Spectrally-decomposed OLR
if self.return_spectral_olr:
# Adjust output flag
self._ispec = 1 # Spectral OLR output flag, 0: only calculate total fluxes, 1: also return spectral OLR
# set up appropriately sized Field object to store the spectral OLR diagnostic
spectral_axes = {**self.OLR.domain.axes, 'wavenumber': wavenum_ax}
spectral_domain = domain._Domain(axes=spectral_axes)
# HACK need to reorder axes in the domain object
### I think we want to change this. Should be consistent with dimension ordering for tauaer and other spectrally resolved fields
shape = list(self.OLR.shape)
shape.append(wavenum_ax.num_points)
spectral_domain.shape = tuple(shape)
spectral_domain.axis_index = {**self.OLR.domain.axis_index, 'wavenumber': len(shape)-1}
# This ensures that the spectral dimension (length: nbndlw) is appended after existing grid dimensions
blank_field = Field((self.OLR[...,np.newaxis] * wavenum_delta), domain=spectral_domain)
self.add_diagnostic('OLR_spectral', blank_field)
else:
self._ispec = 0, # Spectral OLR output flag, 0: only calculate total fluxes, 1: also return spectral OLR
[docs]
def _spectral_field(self, field):
full_spectral_axes = {**self.Tatm.domain.axes, 'wavenumber': wavenum_ax}
full_spectral_domain = domain._Domain(axes=full_spectral_axes)
return Field(field *
np.repeat(np.ones_like(self.Tatm[np.newaxis, ...]), nbndlw, axis=0),
domain=full_spectral_domain)
[docs]
def _prepare_lw_arguments(self):
# scalar integer arguments
icld = self.icld
ispec = self._ispec
irng = self.irng
idrv = self.idrv
permuteseed = self.permuteseed
inflglw = self.inflglw
iceflglw = self.iceflglw
liqflglw = self.liqflglw
(ncol, nlay, play, plev, tlay, tlev, tsfc,
h2ovmr, o3vmr, co2vmr, ch4vmr, n2ovmr, o2vmr, cfc11vmr,
cfc12vmr, cfc12vmr, cfc22vmr, ccl4vmr,
cldfrac, ciwp, clwp, relq, reic) = _prepare_general_arguments(self)
# surface emissivity
emis = self.emissivity * np.ones((ncol,nbndlw))
# These arrays have an extra dimension for number of bands
# in-cloud optical depth [nbndlw,ncol,nlay]
tauc = _climlab_to_rrtm(self.tauc * np.ones_like(self.Tatm))
# broadcast to get [nbndlw,ncol,nlay]
tauc = tauc * np.ones([nbndlw,ncol,nlay])
tauaer = _climlab_to_rrtm(self.tauaer, spectral_axis=True)
# # Aerosol optical depth at mid-point of LW spectral bands, needs to be [ncol,nlay,nbndlw]
# tauaer = self.tauaer[..., ::-1] # flip pressure order
# if len(self.Tatm.shape)==1: # (num_lev)
# # Need to append an extra dimension for singleton horizontal ncol
# tauaer = tauaer[..., np.newaxis, :] # [nbndlw, ncol, nlay]
# # transpose to get [ncol,nlay,nbndlw]
# tauaer = np.transpose(tauaer, (1,2,0))
args = [ncol, nlay, icld, ispec, permuteseed, irng, idrv, const.cp,
play, plev, tlay, tlev, tsfc,
h2ovmr, o3vmr, co2vmr, ch4vmr, n2ovmr, o2vmr,
cfc11vmr, cfc12vmr, cfc22vmr, ccl4vmr, emis,
inflglw, iceflglw, liqflglw,
cldfrac, ciwp, clwp, reic, relq, tauc, tauaer,]
return args
[docs]
def _compute_heating_rates(self):
'''Prepare arguments and call the RRTGM_LW driver to calculate
radiative fluxes and heating rates'''
(ncol, nlay, icld, ispec, permuteseed, irng, idrv, cp,
play, plev, tlay, tlev, tsfc,
h2ovmr, o3vmr, co2vmr, ch4vmr, n2ovmr, o2vmr,
cfc11vmr, cfc12vmr, cfc22vmr, ccl4vmr, emis,
inflglw, iceflglw, liqflglw,
cldfrac, ciwp, clwp, reic, relq, tauc, tauaer,) = self._prepare_lw_arguments()
if icld == 0: # clear-sky only
cldfmcl = np.zeros((ngptlw,ncol,nlay))
ciwpmcl = np.zeros((ngptlw,ncol,nlay))
clwpmcl = np.zeros((ngptlw,ncol,nlay))
reicmcl = np.zeros((ncol,nlay))
relqmcl = np.zeros((ncol,nlay))
taucmcl = np.zeros((ngptlw,ncol,nlay))
else:
# Call the Monte Carlo Independent Column Approximation (McICA, Pincus et al., JC, 2003)
(cldfmcl, ciwpmcl, clwpmcl, reicmcl, relqmcl, taucmcl) = \
_rrtmg_lw.climlab_mcica_subcol_lw(
ncol, nlay, icld,
permuteseed, irng, play,
cldfrac, ciwp, clwp, reic, relq, tauc)
# Call the RRTMG_LW driver to compute radiative fluxes
(olr_sr, uflx, dflx, hr, uflxc, dflxc, hrc, duflx_dt, duflxc_dt) = \
_rrtmg_lw.climlab_rrtmg_lw(ncol, nlay, icld, ispec, idrv,
play, plev, tlay, tlev, tsfc,
h2ovmr, o3vmr, co2vmr, ch4vmr, n2ovmr, o2vmr,
cfc11vmr, cfc12vmr, cfc22vmr, ccl4vmr, emis,
inflglw, iceflglw, liqflglw, cldfmcl,
taucmcl, ciwpmcl, clwpmcl, reicmcl, relqmcl,
tauaer)
# Output is all (ncol,nlay+1) or (ncol,nlay)
self.LW_flux_up = _rrtm_to_climlab(uflx) + 0.*self.LW_flux_up
self.LW_flux_down = _rrtm_to_climlab(dflx) + 0.*self.LW_flux_down
self.LW_flux_up_clr = _rrtm_to_climlab(uflxc) + 0.*self.LW_flux_up_clr
self.LW_flux_down_clr = _rrtm_to_climlab(dflxc) + 0.*self.LW_flux_down_clr
# Compute quantities derived from fluxes, including OLR
self._compute_LW_flux_diagnostics()
# Except for spectrally-decomposed TOA flux, olr_sr (ncol, nbndlw)
if self.return_spectral_olr:
# Need to deal with broadcasting for two different cases: single column and latitude axis
# case single column: self.OLR is (1,), self.OLR_spectral is (1, nbndlw), olr_sr is (1,nbndlw)
# squeeze olr_sr down to (nbndlw,)
# then use np.squeeze(olr_sr)[..., np.newaxis, :] to get back to (1, nbndlw)
# case latitude axis: self.OLR is (num_lat,1), self.OLR_spectral is (num_lat, 1, nbndlw), olr_sr is (num_lat, nbndlw)
# np.squeeze(olr_sr) has no effect in this case
# add the newaxis because the domain has a size-1 depth axis ---> (num_lat, 1, nbndlw)
self.OLR_spectral = np.squeeze(olr_sr)[...,np.newaxis,:] + 0.*self.OLR_spectral
# calculate heating rates from flux divergence
LWheating_Wm2 = np.array(np.diff(self.LW_flux_net, axis=-1)) + 0.*self.Tatm
LWheating_clr_Wm2 = np.array(np.diff(self.LW_flux_net_clr, axis=-1)) + 0.*self.Tatm
self.heating_rate['Ts'] = np.array(-self.LW_flux_net[..., -1, np.newaxis]) + 0.*self.Ts
self.heating_rate['Tatm'] = LWheating_Wm2
# Convert to K / day
Catm = self.Tatm.domain.heat_capacity
self.TdotLW = LWheating_Wm2 / Catm * const.seconds_per_day
self.TdotLW_clr = LWheating_clr_Wm2 / Catm * const.seconds_per_day
