Merge pull request #38 from OMS-NetZero/ensemblegen

joint lognormal TCR/ECS ensemble generation now available

.gitignore

@@ -18,6 +18,7 @@
 
 # Ipython checkpoints
 .ipynb_checkpoints/*
+notebooks/.ipynb_checkpoints/*
 
 # autosaved python version of index.ipynb
 index.py

fair/forcing/aerosols.py

@@ -4,7 +4,6 @@
 import os
 import sys
 import numpy as np
-from scipy.interpolate import Rbf
 from ..constants import molwt
 from ..RCPs.rcp45 import Emissions as r45e
 

fair/tools/ensemble.py

@@ -0,0 +1,124 @@
+import numpy as np
+import scipy.stats as st
+import os
+from functools import reduce
+
+def mvlognorm(data, n=1000, seed=None, correlated=True):
+    """Returns joint lognormal random variables.
+    
+    Inputs:
+        data: (m, p) array of 'observations' where m is the number
+            of observations and p is the number of parameters to
+            estimate.
+        
+    Keywords:
+        n: number of 
+        seed: random seed for variable generation
+        correlated: logical. If True, assume random variables are
+            correlated, and calculate the correlation coefficient.
+        
+    Outputs:
+        (n, p) array of simulated joint lognormal random variables.
+        
+    Method: This function takes the input data (untransformed) and 
+    calculates the mean, standard deviation and correlation coefficient
+    of the data, transforms it to a lognormal, and returns simulated
+    data based on the input 'observations'.
+    
+    It is based on the MethylCapSig R package by Deepak N. Ayyala et al.,
+    https://CRAN.R-project.org/package=MethylCapSig
+    """
+
+    p = data.shape[1]
+    mu = np.mean(data, axis=0)
+    sigma = np.std(data, axis=0)
+    if correlated:
+        corr = np.corrcoef(np.log(data), rowvar=False)
+    else:
+        corr = np.eye((p))
+
+    alpha = np.log(mu) - 0.5*np.log(1.0 + sigma/mu**2)
+    beta = np.diag(np.log(1.0 + sigma/(mu**2)))
+
+    delta = reduce(np.matmul, [np.sqrt(beta), corr, np.sqrt(beta)])
+    delta_eigenvalues, delta_eigenvectors = np.linalg.eig(delta)
+    root_delta = reduce(np.matmul, 
+      [delta_eigenvectors, np.diag(np.sqrt(delta_eigenvalues)), delta_eigenvectors.T])
+
+    out = st.norm.rvs(size=(n,p), random_state=seed)
+    for i in range(n):
+        out[i,:] = np.exp(alpha[:] + np.matmul(root_delta, out[i,:]))
+
+    return out
+
+
+def tcrecs_generate(tcrecs_in='cmip5', dist='lognorm', n=1000, correlated=True,
+                    strip_ecs_lt_tcr=True,
+                    seed=None):
+    """Generates a distribution of TCR and ECS.
+    
+    Inputs:
+        tcrecs_in: either 'cmip5' for pre-shipped CMIP5 TCR (transient climate
+        response) and ECS (equilibrium climate sensitivity) values, or a 
+        2-column array of TCR and ECS values to sample from.
+        
+    Keywords:
+        dist: string. Distribution to use when constructing the
+            joint distribution. Accepted values are norm and 
+            lognorm (default).
+        n: number of samples to generate. Default 1000.
+        correlated: logical. If True (default), assume ECS and TCR
+            inputs are correlated. The function calculates the
+            correlation coefficient automatically.
+        strip_ecs_lt_tcr: logical. If True (default), remove values
+            where ECS < TCR in place (but still return n samples.)
+        seed: random seed for generating variables.
+
+    Output:
+        (n, 2) array of sampled ECS, TCR pairs."""
+
+    if type(tcrecs_in) is str and tcrecs_in=='cmip5':
+        filepath = os.path.join(os.path.dirname(__file__),
+          'tcrecs/cmip5tcrecs.csv')
+        tcrecs_in = np.loadtxt(filepath, delimiter=',', skiprows=3)
+    try:
+        assert(type(tcrecs_in) is np.ndarray)
+        assert(tcrecs_in.shape[1] == 2)
+    except AssertionError:
+        raise ValueError('tcrecs_in should "cmip5" or an array of shape (n, 2)')
+
+    dist = dist.lower()
+
+    def _genvar(tcrecs, dist, n, seed, correlated):
+        if dist=='lognorm':
+            out = mvlognorm(tcrecs, n=n, seed=seed, correlated=correlated)
+        elif dist=='norm':
+            mu = np.mean(tcrecs, axis=0)
+            sigma = np.std(tcrecs, axis=0)
+            if correlated:
+                cov = np.cov(tcrecs, rowvar=False)
+            else:
+                cov = np.diag(np.var(tcrecs, axis=0))
+            out = st.multivariate_normal.rvs(mu, cov, size=n, random_state=seed)
+        else:
+            raise ValueError('dist should be "norm" or "lognorm"')
+        return out
+
+    tcrecs_out = _genvar(tcrecs_in, dist, n, seed, correlated)
+
+    if strip_ecs_lt_tcr:
+        tcrecs_out = np.delete(tcrecs_out, 
+          np.where(tcrecs_out[:,0] > tcrecs_out[:,1]), axis=0)
+        nreq = n - len(tcrecs_out[:,0])
+        while nreq>0:
+            # required otherwise we are repeating values
+            # this is still deterministic for constant ensemble size
+            if seed is not None:
+                seed = seed + nreq
+            new = _genvar(tcrecs_out, dist, n-len(tcrecs_out[:,0]),
+              seed, correlated)
+            tcrecs_out = np.append(tcrecs_out, new, axis=0)
+            tcrecs_out = np.delete(tcrecs_out, np.where(tcrecs_out[:,0] > 
+              tcrecs_out[:,1]), axis=0)
+            nreq = n - len(tcrecs_out[:,0])
+    return tcrecs_out

fair/tools/tcrecs/cmip5tcrecs.csv

@@ -0,0 +1,26 @@
+# ECS and TCR from CMIP5 climate models
+# Source: Forster et al., 2013, JGRA, DOI https://doi.org/10.1002/jgrd.50174
+tcr,ecs
+2,3.83
+1.7,2.82
+2.1,2.87
+2.4,3.69
+1.8,2.89
+2.1,3.25
+1.8,4.08
+2.4,4.17
+2,3.97
+1.1,2.39
+1.3,2.44
+1.7,2.31
+1.5,2.11
+2.5,4.59
+1.3,2.08
+2,4.13
+1.5,2.61
+1.5,2.72
+2.2,4.67
+2,3.63
+2,3.45
+1.6,2.6
+1.4,2.8

tests/tcrecs/tcrecs_output.txt

@@ -0,0 +1,100 @@
+1.451538196588102902e+00 2.469551245165491959e+00
+1.297573956030297238e+00 2.261243882526820403e+00
+9.692874404794418197e-01 2.158150409724756358e+00
+2.334330482845776711e+00 3.854521109121046063e+00
+1.113379754906709262e+00 2.247321290101159441e+00
+1.425949410376278870e+00 2.115064806223128979e+00
+1.218132271209856121e+00 2.626338885548963553e+00
+2.240018805638363641e+00 3.686812627038122514e+00
+9.964933666846257365e-01 2.374794188829417063e+00
+2.076890436804943008e+00 4.868241433193782441e+00
+1.804163200298524394e+00 3.444817284145909042e+00
+3.008985012029306017e+00 5.163866444729952043e+00
+1.132265849296000271e+00 2.329519982756000918e+00
+1.683723132963531599e+00 3.737505809553693936e+00
+1.207867816287714557e+00 2.669705110561091832e+00
+1.116113657162634665e+00 2.359215966706684764e+00
+1.375323789401841079e+00 2.687646700174538772e+00
+2.320474834878655468e+00 4.403537949018546449e+00
+2.097879239714476096e+00 2.674937320248055528e+00
+2.192818650168901318e+00 4.275872560087226937e+00
+2.037596173320204151e+00 3.916090669356882437e+00
+1.296953350358093893e+00 3.049329514639044447e+00
+1.288011404878481203e+00 3.017570454340144259e+00
+3.310337903006294269e+00 4.836514375225618778e+00
+1.931791866855064432e+00 3.323556453910847530e+00
+1.639858329544262361e+00 3.347357242456931115e+00
+2.286098073422712318e+00 3.074216938739542382e+00
+3.263746334708250707e+00 4.343232397467771122e+00
+1.100781748513349978e+00 1.945423918162827004e+00
+1.556662136870744728e+00 2.995974825466610092e+00
+1.624228871195223967e+00 3.145926444859620918e+00
+2.384785393119428942e+00 6.665786223479530648e+00
+3.531270776396765232e+00 5.414418793806882846e+00
+2.717451332540535613e+00 4.031173348880250806e+00
+1.939022222243606874e+00 2.915906952122077289e+00
+1.215940238750810698e+00 2.581455414999639153e+00
+1.987245521422330796e+00 2.965659909972387265e+00
+2.127680271879808238e+00 4.222347216197413822e+00
+2.307099357126805916e+00 3.242730397085637151e+00
+1.715285602384358565e+00 3.602734403472037439e+00
+1.699498120169044890e+00 2.612998310865429108e+00
+2.135293940332133378e+00 4.422936854804563112e+00
+1.729228789023881552e+00 3.094314351160178500e+00
+1.292024615393371656e+00 2.255100031839395225e+00
+1.042139951493962435e+00 2.416084424396452768e+00
+1.024008011101124938e+00 2.522822597325784511e+00
+1.619401536921436069e+00 3.298248326414765952e+00
+1.697711425825302545e+00 2.625813559511153361e+00
+2.336063775158583322e+00 4.635439967075168255e+00
+2.463940776335298288e+00 3.724986338679467046e+00
+2.772817167397374050e+00 5.432422503419477522e+00
+1.712565374431410214e+00 3.196098096563476876e+00
+1.186944183534079711e+00 2.817093130404066947e+00
+1.588557790513311252e+00 2.431925579100456591e+00
+1.078101470883030322e+00 3.553182424469676803e+00
+1.035028963215754105e+00 2.580930491345172495e+00
+2.164008392805430336e+00 3.962506754929142172e+00
+1.626271251591156153e+00 3.304900713542970436e+00
+2.377184434239144828e+00 3.906982339569422535e+00
+1.545023447726212584e+00 2.353890886832713747e+00
+2.645593845544210510e+00 3.927925677003198057e+00
+2.183236298008408571e+00 3.887655323016329323e+00
+1.382967507447103506e+00 2.055886523564756097e+00
+6.757589337461044066e-01 1.959614280912242990e+00
+1.662868555096991940e+00 3.017208741803167449e+00
+1.213101792810705293e+00 2.195561501738052534e+00
+1.381503902478510071e+00 3.127131607142050296e+00
+1.759743256403233813e+00 3.842818279912441515e+00
+2.515020834360082613e+00 3.192709672366153217e+00
+2.087944403707748364e+00 4.421002138112082314e+00
+2.888969005075551877e+00 5.066268623714294428e+00
+2.820310642735099460e+00 4.296721105594568435e+00
+1.559151809878610573e+00 2.744498264258125264e+00
+1.998314298175857484e+00 3.377067146824034083e+00
+1.654080102865330915e+00 3.236626211702226197e+00
+1.309223475250872903e+00 2.464408289117991480e+00
+3.170041447517683242e+00 4.378560915470894876e+00
+1.717132757071669902e+00 3.362592903955825108e+00
+9.702685103821293566e-01 1.953683320995563077e+00
+2.906225375642801367e+00 4.295897565324628964e+00
+1.775141992480178121e+00 3.922507536555935204e+00
+3.545461371194654987e+00 4.986659920734763496e+00
+1.691623925394495798e+00 3.026751507285943887e+00
+2.961873788693437870e+00 4.053609064264886541e+00
+1.740252882932598943e+00 3.468605439337008267e+00
+1.281542281436591990e+00 2.736135933783638396e+00
+1.595366759105774968e+00 2.556198156043188785e+00
+2.932421830273996388e+00 3.451723996254647986e+00
+1.251245729123137140e+00 2.083971792901944564e+00
+1.346614368267833317e+00 2.750423391648577631e+00
+1.312483183878368154e+00 3.821929439724604194e+00
+2.448780430572317535e+00 4.174750064642457126e+00
+3.717528648138180625e+00 4.815857734052825023e+00
+1.377949246965856744e+00 2.736635126278957220e+00
+1.674266844936252108e+00 2.502958841170364490e+00
+1.122242291973260198e+00 2.442624901048536135e+00
+1.319800402547825824e+00 3.573987875511071000e+00
+2.008710870466959619e+00 3.008138906562877413e+00
+3.346478291922688975e+00 4.700443343584020717e+00
+2.647250704391764575e+00 4.580744339593612047e+00

tests/test_fair.py

@@ -4,7 +4,7 @@
 import os
 import numpy as np
 from fair.RCPs import rcp3pd, rcp45, rcp6, rcp85
-from fair.tools import magicc, steady
+from fair.tools import magicc, steady, ensemble
 from fair.ancil import natural, cmip5_annex2_forcing
 from fair.constants import molwt
 from fair.forcing.ghg import myhre
@@ -416,3 +416,14 @@ def test_rcp_aliases():
     assert np.allclose(C, C_expected)
     assert np.allclose(F, F_expected)
     assert np.allclose(T, T_expected)
+
+
+def test_ensemble():
+    tcrecs = ensemble.tcrecs_generate('cmip5', dist='lognorm', n=100,
+      seed=40000)
+    # check first 100 lines
+    filepath = os.path.join(os.path.dirname(__file__), 
+      'tcrecs/tcrecs_output.txt')
+    tcrecs_compare = np.loadtxt(filepath)
+    assert np.allclose(tcrecs,tcrecs_compare)
+