LETKF.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Modify the following variables before each experiment\n",
    "# hcovlocal_scale   L_h\n",
    "# covinflate1       RTPS\n",
    "# exptname          output file name\n",
    "# filename_climo    input file 64/256   3hr/12hr\n",
    "# filename_truth    input file\n",
    "# nobs              number of observation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ARCTAN\n",
    "\n",
    "from __future__ import print_function\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from netCDF4 import Dataset\n",
    "import sys, time, os\n",
    "from sqgturb import SQG, rfft2, irfft2, cartdist,enkf_update,gaspcohn, bulk_ensrf\n",
    "\n",
    "\n",
    "def submit():\n",
    "    # horizontal covariance localization length scale in meters.\n",
    "    hcovlocal_scale = 4500  #float(sys.argv[1])\n",
    "    hcovlocal_scale = hcovlocal_scale * 1000\n",
    "    covinflate1 = 0.9\n",
    "    covinflate2 = -1\n",
    "\n",
    "    exptname = os.getenv('exptname','EnKF_N64_12hrly_256_fixed_arctan_{}_{}'.format(hcovlocal_scale/1000 ,covinflate1)) \n",
    "    threads = int(os.getenv('OMP_NUM_THREADS','1'))\n",
    "\n",
    "    diff_efold = None # use diffusion from climo file\n",
    "\n",
    "    profile = False # turn on profiling?\n",
    "\n",
    "    use_letkf = True  #False use LETKF\n",
    "    global_enkf = False #False # global EnSRF solve\n",
    "    read_restart = False\n",
    "    # if savedata not None, netcdf filename will be defined by env var 'exptname'\n",
    "    # if savedata = 'restart', only last time is saved (so expt can be restarted)\n",
    "    #savedata = True \n",
    "    #savedata = 'restart'\n",
    "    savedata = None\n",
    "    #nassim = 101 \n",
    "    #nassim_spinup = 1\n",
    "    nassim = 100 # assimilation times to run\n",
    "    nassim_spinup = 100\n",
    "\n",
    "    direct_insertion = False \n",
    "    if direct_insertion: print('# direct insertion!')\n",
    "\n",
    "    nanals = 20 # ensemble members  20\n",
    "\n",
    "    oberrstdev = 0.01 # ob error standard deviation in K\n",
    "\n",
    "    # nature run created using sqg_run.py.\n",
    "    filename_climo = 'sqg_N64_12hrly.nc' # file name for forecast model climo \n",
    "    # perfect model\n",
    "    filename_truth = 'sqg_N64_12hrly.nc' # file name for nature run to draw obs \n",
    "\n",
    "\n",
    "    print('# filename_modelclimo=%s' % filename_climo)\n",
    "    print('# filename_truth=%s' % filename_truth)\n",
    "\n",
    "    # fix random seed for reproducibility.\n",
    "    rsobs = np.random.RandomState(42) # fixed seed for observations\n",
    "    rsics = np.random.RandomState() # varying seed for initial conditions\n",
    "    rsjump = np.random.RandomState(888) # fixed seed for observations\n",
    "    \n",
    "\n",
    "    # get model info\n",
    "    nc_climo = Dataset(filename_climo)\n",
    "    # parameter used to scale PV to temperature units.\n",
    "    scalefact = nc_climo.f*nc_climo.theta0/nc_climo.g\n",
    "    # initialize qg model instances for each ensemble member.\n",
    "    x = nc_climo.variables['x'][:]\n",
    "    y = nc_climo.variables['y'][:]\n",
    "    x, y = np.meshgrid(x, y)\n",
    "    nx = len(x); ny = len(y)\n",
    "    dt = nc_climo.dt\n",
    "    if diff_efold == None: diff_efold=nc_climo.diff_efold\n",
    "    pvens = np.empty((nanals,2,ny,nx),np.float32)\n",
    "    if not read_restart:\n",
    "        pv_climo = nc_climo.variables['pv']\n",
    "        indxran = rsics.choice(pv_climo.shape[0],size=nanals,replace=False)\n",
    "    else:\n",
    "        ncinit = Dataset('%s_restart.nc' % exptname, mode='r', format='NETCDF4_CLASSIC')\n",
    "        ncinit.set_auto_mask(False)\n",
    "        pvens[:] = ncinit.variables['pv_b'][-1,...]/scalefact\n",
    "        tstart = ncinit.variables['t'][-1]\n",
    "        #for nanal in range(nanals):\n",
    "        #    print(nanal, pvens[nanal].min(), pvens[nanal].max())\n",
    "    # get OMP_NUM_THREADS (threads to use) from environment.\n",
    "    models = []\n",
    "    for nanal in range(nanals):\n",
    "        if not read_restart:\n",
    "            pvens[nanal] = pv_climo[indxran[nanal]]\n",
    "            #print(nanal, pvens[nanal].min(), pvens[nanal].max())\n",
    "        pvens[nanal] = pv_climo[0] + np.random.normal(0,1000,size=(2,ny,nx))\n",
    "        models.append(\\\n",
    "        SQG(pvens[nanal],\n",
    "        nsq=nc_climo.nsq,f=nc_climo.f,dt=dt,U=nc_climo.U,H=nc_climo.H,\\\n",
    "        r=nc_climo.r,tdiab=nc_climo.tdiab,symmetric=nc_climo.symmetric,\\\n",
    "        diff_order=nc_climo.diff_order,diff_efold=diff_efold,threads=threads))\n",
    "    if read_restart: ncinit.close()\n",
    "\n",
    "    # vertical localization scale\n",
    "    Lr = np.sqrt(models[0].nsq)*models[0].H/models[0].f\n",
    "    vcovlocal_fact = gaspcohn(np.array(Lr/hcovlocal_scale))\n",
    "    #vcovlocal_fact = 0.0 # no increment at opposite boundary\n",
    "    #vcovlocal_fact = 1.0 # no vertical localization\n",
    "\n",
    "    print('# use_letkf=%s global_enkf=%s' % (use_letkf,global_enkf))\n",
    "    print(\"# hcovlocal=%g vcovlocal=%s diff_efold=%s covinf1=%s covinf2=%s nanals=%s\" %\\\n",
    "        (hcovlocal_scale/1000.,vcovlocal_fact,diff_efold,covinflate1,covinflate2,nanals))\n",
    "\n",
    "    # if nobs > 0, each ob time nobs ob locations are randomly sampled (without\n",
    "    # replacement) from the model grid\n",
    "    # if nobs < 0, fixed network of every Nth grid point used (N = -nobs)\n",
    "    nobs = -256   # number of obs to assimilate (randomly distributed)\n",
    "    #nobs = -1 # fixed network, every -nobs grid points. nobs=-1 obs at all pts.\n",
    "\n",
    "    # nature run\n",
    "    nc_truth = Dataset(filename_truth)\n",
    "    pv_truth = nc_truth.variables['pv']\n",
    "    # set up arrays for obs and localization function\n",
    "    if nobs < 0:\n",
    "        nskip = -nobs\n",
    "        if (nx*ny)%nobs != 0:\n",
    "            raise ValueError('nx*ny must be divisible by nobs')\n",
    "        nobs = (nx*ny)//nskip**2\n",
    "        print('# fixed network nobs = %s' % nobs)\n",
    "        fixed = True\n",
    "        nobs = nskip\n",
    "    else:\n",
    "        fixed = False\n",
    "        print('# random network nobs = %s' % nobs)\n",
    "    #if nobs == nx*ny//2: fixed=True # used fixed network for obs every other grid point\n",
    "    oberrvar = oberrstdev**2*np.ones(nobs,float)\n",
    "    pvob = np.empty((2,nobs),float)\n",
    "    covlocal = np.empty((ny,nx),float)\n",
    "    covlocal_tmp = np.empty((nobs,nx*ny),float)\n",
    "    xens = np.empty((nanals,2,nx*ny),float)\n",
    "    if not use_letkf:\n",
    "        obcovlocal = np.empty((nobs,nobs),float)\n",
    "    else:\n",
    "        obcovlocal = None\n",
    "\n",
    "    if global_enkf: # model-space localization matrix\n",
    "        n = 0\n",
    "        covlocal_modelspace = np.empty((nx*ny,nx*ny),float)\n",
    "        x1 = x.reshape(nx*ny); y1 = y.reshape(nx*ny)\n",
    "        for n in range(nx*ny):\n",
    "            dist = cartdist(x1[n],y1[n],x1,y1,nc_climo.L,nc_climo.L)\n",
    "            covlocal_modelspace[n,:] = gaspcohn(dist/hcovlocal_scale)\n",
    "\n",
    "    obtimes = nc_truth.variables['t'][:]\n",
    "    if read_restart:\n",
    "        timeslist = obtimes.tolist()\n",
    "        ntstart = timeslist.index(tstart)\n",
    "        print('# restarting from %s.nc ntstart = %s' % (exptname,ntstart))\n",
    "    else:\n",
    "        ntstart = 0\n",
    "    assim_interval = obtimes[1]-obtimes[0]\n",
    "    assim_timesteps = int(np.round(assim_interval/models[0].dt))\n",
    "    print('# assim interval = %s secs (%s time steps)' % (assim_interval,assim_timesteps))\n",
    "    print('# ntime,pverr_a,pvsprd_a,pverr_b,pvsprd_b,obinc_b,osprd_b,obinc_a,obsprd_a,omaomb/oberr,obbias_b,inflation,tr(P^a)/tr(P^b)')\n",
    "\n",
    "    # initialize model clock\n",
    "    for nanal in range(nanals):\n",
    "        models[nanal].t = obtimes[ntstart]\n",
    "        models[nanal].timesteps = assim_timesteps\n",
    "\n",
    "    # initialize output file.\n",
    "    if savedata is not None:\n",
    "        nc = Dataset('%s.nc' % exptname, mode='w', format='NETCDF4_CLASSIC')\n",
    "        nc.r = models[0].r\n",
    "        nc.f = models[0].f\n",
    "        nc.U = models[0].U\n",
    "        nc.L = models[0].L\n",
    "        nc.H = models[0].H\n",
    "        nc.nanals = nanals\n",
    "        nc.hcovlocal_scale = hcovlocal_scale\n",
    "        nc.vcovlocal_fact = vcovlocal_fact\n",
    "        nc.oberrstdev = oberrstdev\n",
    "        nc.g = nc_climo.g; nc.theta0 = nc_climo.theta0\n",
    "        nc.nsq = models[0].nsq\n",
    "        nc.tdiab = models[0].tdiab\n",
    "        nc.dt = models[0].dt\n",
    "        nc.diff_efold = models[0].diff_efold\n",
    "        nc.diff_order = models[0].diff_order\n",
    "        nc.filename_climo = filename_climo\n",
    "        nc.filename_truth = filename_truth\n",
    "        nc.symmetric = models[0].symmetric\n",
    "        xdim = nc.createDimension('x',models[0].N)\n",
    "        ydim = nc.createDimension('y',models[0].N)\n",
    "        z = nc.createDimension('z',2)\n",
    "        t = nc.createDimension('t',None)\n",
    "        obs = nc.createDimension('obs',nobs)\n",
    "        ens = nc.createDimension('ens',nanals)\n",
    "        pv_t =\\\n",
    "        nc.createVariable('pv_t',np.float32,('t','z','y','x'),zlib=True)\n",
    "        pv_b =\\\n",
    "        nc.createVariable('pv_b',np.float32,('t','ens','z','y','x'),zlib=True)\n",
    "        pv_a =\\\n",
    "        nc.createVariable('pv_a',np.float32,('t','ens','z','y','x'),zlib=True)\n",
    "        pv_a.units = 'K'\n",
    "        pv_b.units = 'K'\n",
    "        inf = nc.createVariable('inflation',np.float32,('t','z','y','x'),zlib=True)\n",
    "        pv_obs = nc.createVariable('obs',np.float32,('t','obs'))\n",
    "        x_obs = nc.createVariable('x_obs',np.float32,('t','obs'))\n",
    "        y_obs = nc.createVariable('y_obs',np.float32,('t','obs'))\n",
    "        # eady pv scaled by g/(f*theta0) so du/dz = d(pv)/dy\n",
    "        xvar = nc.createVariable('x',np.float32,('x',))\n",
    "        xvar.units = 'meters'\n",
    "        yvar = nc.createVariable('y',np.float32,('y',))\n",
    "        yvar.units = 'meters'\n",
    "        zvar = nc.createVariable('z',np.float32,('z',))\n",
    "        zvar.units = 'meters'\n",
    "        tvar = nc.createVariable('t',np.float32,('t',))\n",
    "        tvar.units = 'seconds'\n",
    "        ensvar = nc.createVariable('ens',np.int32,('ens',))\n",
    "        ensvar.units = 'dimensionless'\n",
    "        xvar[:] = np.arange(0,models[0].L,models[0].L/models[0].N)\n",
    "        yvar[:] = np.arange(0,models[0].L,models[0].L/models[0].N)\n",
    "        zvar[0] = 0; zvar[1] = models[0].H\n",
    "        ensvar[:] = np.arange(1,nanals+1)\n",
    "\n",
    "    # initialize kinetic energy error/spread spectra\n",
    "    kespec_errmean = None; kespec_sprdmean = None\n",
    "\n",
    "    ncount = 0\n",
    "    nanals2 = 4 # ensemble members used for kespec spread\n",
    "\n",
    "    rmse = np.zeros(nassim)\n",
    "    #Jump\n",
    "    percentage  = 0.1\n",
    "    #print(\"percentage of noise steps\", percentage)\n",
    "    noiselevel=np.array((900,600))\n",
    "    noise_idx = np.array((1,1 - percentage*15/35 ,1-percentage))\n",
    "    jump_dist = rsjump.uniform(0,1,300)\n",
    "    jump = np.where(jump_dist > 1-percentage)[0] +1000\n",
    "    #print(jump)\n",
    "    # Initialize an empty dictionary\n",
    "    jump_idx = {}\n",
    "    for i in range(len(noiselevel)):\n",
    "        value = noiselevel[i]\n",
    "        keys = np.where((jump_dist <= noise_idx[i]) & (jump_dist > noise_idx[i+1]))[0].tolist()\n",
    "        for key in keys:\n",
    "            jump_idx[key] = value\n",
    "    #print(jump_idx)\n",
    "\n",
    "    indxob = np.sort(rsobs.choice(nx*ny,nobs,replace=False))\n",
    "    indx_unob = np.setdiff1d(np.arange(nx*ny), indxob)\n",
    "    obs_save = np.zeros((100,nobs))\n",
    "\n",
    "    for ntime in range(nassim): #nassim\n",
    "\n",
    "        # check model clock\n",
    "        if models[0].t != obtimes[ntime+ntstart]:\n",
    "            raise ValueError('model/ob time mismatch %s vs %s' %\\\n",
    "            (models[0].t, obtimes[ntime+ntstart]))\n",
    "\n",
    "        t1 = time.time()\n",
    "        if not fixed:\n",
    "            # randomly choose points from model grid\n",
    "            if nobs == nx*ny:\n",
    "                indxob = np.arange(nx*ny)\n",
    "            else:\n",
    "                indxob = np.sort(rsobs.choice(nx*ny,nobs,replace=False))\n",
    "        else:\n",
    "                if ntime == 0:\n",
    "                    print(\"not Random\")\n",
    "        if (ntime == jump).any():\n",
    "            jumpnoise = rsjump.normal(0,jump_idx[ntime],size=(1,2,ny,nx)) \n",
    "            jumpnoise_reshape =jumpnoise.reshape(2,ny*nx)  \n",
    "        for k in range(2):\n",
    "            # surface temp obs\n",
    "            if (ntime == jump).any():\n",
    "                pvob[k] = np.arctan(scalefact*(pv_truth[ntime+ntstart,k,:,:].ravel()[indxob] +  jumpnoise_reshape[k,indxob]))\n",
    "            else:\n",
    "                pvob[k] = np.arctan(scalefact*pv_truth[ntime+ntstart,k,:,:].ravel()[indxob])\n",
    "            pvob[k] += rsobs.normal(scale=oberrstdev,size=nobs) # add ob errors\n",
    "        xob = x.ravel()[indxob]\n",
    "        yob = y.ravel()[indxob]\n",
    "        # compute covariance localization function for each ob\n",
    "        if not fixed or ntime == 0:\n",
    "            for nob in range(nobs):\n",
    "                dist = cartdist(xob[nob],yob[nob],x,y,nc_climo.L,nc_climo.L)\n",
    "                covlocal = gaspcohn(dist/hcovlocal_scale)\n",
    "                covlocal_tmp[nob] = covlocal.ravel()\n",
    "                dist = cartdist(xob[nob],yob[nob],xob,yob,nc_climo.L,nc_climo.L)\n",
    "                if not use_letkf: obcovlocal[nob] = gaspcohn(dist/hcovlocal_scale)\n",
    "\n",
    "        # first-guess spread (need later to compute inflation factor)\n",
    "        fsprd = ((pvens - pvens.mean(axis=0))**2).sum(axis=0)/(nanals-1)\n",
    "\n",
    "        # compute forward operator.\n",
    "        # hxens is ensemble in observation space.\n",
    "        hxens = np.empty((nanals,2,nobs),float)\n",
    "        for nanal in range(nanals):\n",
    "            for k in range(2):\n",
    "                hxens[nanal,k,...] = np.arctan(scalefact*pvens[nanal,k,...].ravel()[indxob]) # surface pv obs\n",
    "        hxensmean_b = hxens.mean(axis=0)\n",
    "        obsprd = ((hxens-hxensmean_b)**2).sum(axis=0)/(nanals-1)\n",
    "        # innov stats for background\n",
    "        obfits = pvob - hxensmean_b\n",
    "        obfits_b = (obfits**2).mean()\n",
    "        obbias_b = obfits.mean()\n",
    "        obsprd_b = obsprd.mean()\n",
    "        pvensmean_b = pvens.mean(axis=0).copy()\n",
    "        if (ntime == jump).any():\n",
    "            pverr_b = (scalefact*(pvensmean_b-pv_truth[ntime+ntstart]+jumpnoise))**2\n",
    "        else:\n",
    "            pverr_b = (scalefact*(pvensmean_b-pv_truth[ntime+ntstart]))**2\n",
    "        pvsprd_b = ((scalefact*(pvensmean_b-pvens))**2).sum(axis=0)/(nanals-1)\n",
    "\n",
    "        if savedata is not None:\n",
    "            if savedata == 'restart' and ntime != nassim-1:\n",
    "                pass\n",
    "            else:\n",
    "                pv_t[ntime] = pv_truth[ntime+ntstart]\n",
    "                pv_b[ntime,:,:,:] = scalefact*pvens\n",
    "                #pv_obs[ntime] = pvob\n",
    "                #x_obs[ntime] = xob\n",
    "                #y_obs[ntime] = yob\n",
    "\n",
    "        # EnKF update\n",
    "        # create 1d state vector.\n",
    "        xens = pvens.reshape(nanals,2,nx*ny)\n",
    "        # update state vector.\n",
    "        if direct_insertion and nobs == nx*ny:\n",
    "            print('wrongwrong!!!!!')\n",
    "            for nanal in range(nanals):\n",
    "                xens[nanal] =\\\n",
    "                pv_truth[ntime+ntstart].reshape(2,nx*ny) + \\\n",
    "                rsobs.normal(scale=oberrstdev,size=(2,nx*ny))/scalefact\n",
    "            xens = xens - xens.mean(axis=0) + \\\n",
    "            pv_truth[ntime+ntstart].reshape(2,nx*ny) + \\\n",
    "            rsobs.normal(scale=oberrstdev,size=(2,nx*ny))/scalefact\n",
    "        else:\n",
    "            # hxens,pvob are in PV units, xens is not \n",
    "            if global_enkf and not use_letkf:\n",
    "                xens = bulk_ensrf(xens,indxob,pvob,oberrvar,covlocal_modelspace,vcovlocal_fact,scalefact)\n",
    "            else:\n",
    "                xens =\\\n",
    "                enkf_update(xens,hxens,pvob,oberrvar,covlocal_tmp,vcovlocal_fact,obcovlocal=obcovlocal)\n",
    "        # back to 3d state vector\n",
    "        pvens = xens.reshape((nanals,2,ny,nx))\n",
    "        t2 = time.time()\n",
    "        #print('cpu time for EnKF update',t2-t1)\n",
    "\n",
    "        # forward operator on posterior ensemble.\n",
    "        for nanal in range(nanals):\n",
    "            for k in range(2):\n",
    "                hxens[nanal,k,...] = np.arctan(scalefact*pvens[nanal,k,...].ravel()[indxob]) # surface pv obs\n",
    "\n",
    "        # ob space diagnostics\n",
    "        hxensmean_a = hxens.mean(axis=0)\n",
    "        obsprd_a = (((hxens-hxensmean_a)**2).sum(axis=0)/(nanals-1)).mean()\n",
    "        # expected value is HPaHT (obsprd_a).\n",
    "        obinc_a = ((hxensmean_a-hxensmean_b)*(pvob-hxensmean_a)).mean()\n",
    "        # expected value is HPbHT (obsprd_b).\n",
    "        obinc_b = ((hxensmean_a-hxensmean_b)*(pvob-hxensmean_b)).mean()\n",
    "        # expected value R (oberrvar).\n",
    "        omaomb = ((pvob-hxensmean_a)*(pvob-hxensmean_b)).mean()\n",
    "\n",
    "        # posterior multiplicative inflation.\n",
    "        pvensmean_a = pvens.mean(axis=0)\n",
    "        pvprime = pvens-pvensmean_a\n",
    "        asprd = (pvprime**2).sum(axis=0)/(nanals-1)\n",
    "        asprd_over_fsprd = asprd.mean()/fsprd.mean()\n",
    "        if covinflate2 < 0:\n",
    "            # relaxation to prior stdev (Whitaker & Hamill 2012)\n",
    "            asprd = np.sqrt(asprd); fsprd = np.sqrt(fsprd)\n",
    "            inflation_factor = 1.+covinflate1*(fsprd-asprd)/asprd\n",
    "        else:\n",
    "            # Hodyss et al 2016 inflation (covinflate1=covinflate2=1 works well in perfect\n",
    "            # model, linear gaussian scenario)\n",
    "            # inflation = asprd + (asprd/fsprd)**2((fsprd/nanals)+2*inc**2/(nanals-1))\n",
    "            inc = pvensmean_a - pvensmean_b\n",
    "            inflation_factor = covinflate1*asprd + \\\n",
    "            (asprd/fsprd)**2*((fsprd/nanals) + covinflate2*(2.*inc**2/(nanals-1)))\n",
    "            inflation_factor = np.sqrt(inflation_factor/asprd)\n",
    "        pvprime = pvprime*inflation_factor\n",
    "        pvens = pvprime + pvensmean_a\n",
    "\n",
    "        # print out analysis error, spread and innov stats for background\n",
    "        if (ntime == jump).any():\n",
    "            print('jumppp')\n",
    "            pverr_a = (scalefact*(pvensmean_a-pv_truth[ntime+ntstart]-jumpnoise))**2\n",
    "        else:\n",
    "            pverr_a = (scalefact*(pvensmean_a-pv_truth[ntime+ntstart]))**2\n",
    "        pvsprd_a = ((scalefact*(pvensmean_a-pvens))**2).sum(axis=0)/(nanals-1)\n",
    "        print(\"%s %g %g %g %g %g %g %g %g %g %g %g %g\" %\\\n",
    "        (ntime+ntstart,np.sqrt(pverr_a.mean()),np.sqrt(pvsprd_a.mean()),\\\n",
    "        np.sqrt(pverr_b.mean()),np.sqrt(pvsprd_b.mean()),\\\n",
    "        obinc_b,obsprd_b,obinc_a,obsprd_a,omaomb/oberrvar.mean(),obbias_b,inflation_factor.mean(),asprd_over_fsprd))\n",
    "        \n",
    "        rmse[ntime] = np.sqrt(pverr_a.mean())\n",
    "        \n",
    "        #np.save(runing_rmse,rmse)\n",
    "        # save data.\n",
    "        if savedata is not None:\n",
    "            if savedata == 'restart' and ntime != nassim-1:\n",
    "                pass\n",
    "            else:\n",
    "                pv_a[ntime,:,:,:] = scalefact*pvens\n",
    "                tvar[ntime] = obtimes[ntime+ntstart]\n",
    "                inf[ntime] = inflation_factor\n",
    "                nc.sync()\n",
    "\n",
    "        # run forecast ensemble to next analysis time\n",
    "        t1 = time.time()\n",
    "        for nanal in range(nanals):\n",
    "            pvens[nanal] = models[nanal].advance(pvens[nanal])\n",
    "        t2 = time.time()\n",
    "        #print('cpu time for ens forecast',t2-t1)\n",
    "\n",
    "\n",
    "        # compute spectra of error and spread\n",
    "        if ntime >= nassim_spinup:\n",
    "            pvfcstmean = pvens.mean(axis=0)\n",
    "            pverrspec = scalefact*rfft2(pvfcstmean - pv_truth[ntime+ntstart])\n",
    "            psispec = models[0].invert(pverrspec)\n",
    "            psispec = psispec/(models[0].N*np.sqrt(2.))\n",
    "            kespec = (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "            if kespec_errmean is None:\n",
    "                kespec_errmean =\\\n",
    "                (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "            else:\n",
    "                kespec_errmean = kespec_errmean + kespec\n",
    "            for nanal in range(nanals2):\n",
    "                pvsprdspec = scalefact*rfft2(pvens[nanal] - pvfcstmean)\n",
    "                psispec = models[0].invert(pvsprdspec)\n",
    "                psispec = psispec/(models[0].N*np.sqrt(2.))\n",
    "                kespec = (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "                if kespec_sprdmean is None:\n",
    "                    kespec_sprdmean =\\\n",
    "                    (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real/nanals2\n",
    "                else:\n",
    "                    kespec_sprdmean = kespec_sprdmean+kespec/nanals2\n",
    "            ncount += 1\n",
    "\n",
    "    if savedata: nc.close()\n",
    "\n",
    "    #np.save(runing_rmse,rmse)\n",
    "\n",
    "    return None\n",
    "\n",
    "\n",
    "## Run main code\n",
    "if __name__ == \"__main__\":\n",
    "    submit()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Linear\n",
    "\n",
    "from __future__ import print_function\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from netCDF4 import Dataset\n",
    "import sys, time, os\n",
    "from sqgturb import SQG, rfft2, irfft2, cartdist,enkf_update,gaspcohn, bulk_ensrf\n",
    "\n",
    "\n",
    "def submit():\n",
    "    # horizontal covariance localization length scale in meters.\n",
    "    hcovlocal_scale = 1500  #float(sys.argv[1])\n",
    "    hcovlocal_scale = hcovlocal_scale * 1000\n",
    "    covinflate1 = 0.5\n",
    "    covinflate2 = -1\n",
    "\n",
    "    exptname = os.getenv('exptname','EnKF_N64_3hrly_256_fixed_linear_{}_{}'.format(hcovlocal_scale/1000 ,covinflate1)) \n",
    "    threads = int(os.getenv('OMP_NUM_THREADS','1'))\n",
    "\n",
    "    diff_efold = None # use diffusion from climo file\n",
    "\n",
    "    profile = False # turn on profiling?\n",
    "\n",
    "    use_letkf = True  #False use LETKF\n",
    "    global_enkf = False #False # global EnSRF solve\n",
    "    read_restart = False\n",
    "    # if savedata not None, netcdf filename will be defined by env var 'exptname'\n",
    "    # if savedata = 'restart', only last time is saved (so expt can be restarted)\n",
    "    #savedata = True \n",
    "    #savedata = 'restart'\n",
    "    savedata = None\n",
    "    #nassim = 101 \n",
    "    #nassim_spinup = 1\n",
    "    nassim = 100 # assimilation times to run\n",
    "    nassim_spinup = 100\n",
    "\n",
    "    direct_insertion = False \n",
    "    if direct_insertion: print('# direct insertion!')\n",
    "\n",
    "    nanals = 20 # ensemble members  20\n",
    "\n",
    "    oberrstdev = 1 # ob error standard deviation in K\n",
    "\n",
    "    # nature run created using sqg_run.py.\n",
    "    filename_climo = 'sqg_N64_3hrly.nc' # file name for forecast model climo \n",
    "    # perfect model\n",
    "    filename_truth = 'sqg_N64_3hrly.nc' # file name for nature run to draw obs \n",
    "\n",
    "\n",
    "    print('# filename_modelclimo=%s' % filename_climo)\n",
    "    print('# filename_truth=%s' % filename_truth)\n",
    "\n",
    "    # fix random seed for reproducibility.\n",
    "    rsobs = np.random.RandomState(42) # fixed seed for observations\n",
    "    rsics = np.random.RandomState() # varying seed for initial conditions\n",
    "    rsjump = np.random.RandomState(888) # fixed seed for observations\n",
    "    \n",
    "\n",
    "    # get model info\n",
    "    nc_climo = Dataset(filename_climo)\n",
    "    # parameter used to scale PV to temperature units.\n",
    "    scalefact = nc_climo.f*nc_climo.theta0/nc_climo.g\n",
    "    # initialize qg model instances for each ensemble member.\n",
    "    x = nc_climo.variables['x'][:]\n",
    "    y = nc_climo.variables['y'][:]\n",
    "    x, y = np.meshgrid(x, y)\n",
    "    nx = len(x); ny = len(y)\n",
    "    dt = nc_climo.dt\n",
    "    if diff_efold == None: diff_efold=nc_climo.diff_efold\n",
    "    pvens = np.empty((nanals,2,ny,nx),np.float32)\n",
    "    if not read_restart:\n",
    "        pv_climo = nc_climo.variables['pv']\n",
    "        indxran = rsics.choice(pv_climo.shape[0],size=nanals,replace=False)\n",
    "    else:\n",
    "        ncinit = Dataset('%s_restart.nc' % exptname, mode='r', format='NETCDF4_CLASSIC')\n",
    "        ncinit.set_auto_mask(False)\n",
    "        pvens[:] = ncinit.variables['pv_b'][-1,...]/scalefact\n",
    "        tstart = ncinit.variables['t'][-1]\n",
    "        #for nanal in range(nanals):\n",
    "        #    print(nanal, pvens[nanal].min(), pvens[nanal].max())\n",
    "    # get OMP_NUM_THREADS (threads to use) from environment.\n",
    "    models = []\n",
    "    for nanal in range(nanals):\n",
    "        if not read_restart:\n",
    "            pvens[nanal] = pv_climo[indxran[nanal]]\n",
    "            #print(nanal, pvens[nanal].min(), pvens[nanal].max())\n",
    "        pvens[nanal] = pv_climo[0] + np.random.normal(0,1000,size=(2,ny,nx))\n",
    "        models.append(\\\n",
    "        SQG(pvens[nanal],\n",
    "        nsq=nc_climo.nsq,f=nc_climo.f,dt=dt,U=nc_climo.U,H=nc_climo.H,\\\n",
    "        r=nc_climo.r,tdiab=nc_climo.tdiab,symmetric=nc_climo.symmetric,\\\n",
    "        diff_order=nc_climo.diff_order,diff_efold=diff_efold,threads=threads))\n",
    "    if read_restart: ncinit.close()\n",
    "\n",
    "    # vertical localization scale\n",
    "    Lr = np.sqrt(models[0].nsq)*models[0].H/models[0].f\n",
    "    vcovlocal_fact = gaspcohn(np.array(Lr/hcovlocal_scale))\n",
    "    #vcovlocal_fact = 0.0 # no increment at opposite boundary\n",
    "    #vcovlocal_fact = 1.0 # no vertical localization\n",
    "\n",
    "    print('# use_letkf=%s global_enkf=%s' % (use_letkf,global_enkf))\n",
    "    print(\"# hcovlocal=%g vcovlocal=%s diff_efold=%s covinf1=%s covinf2=%s nanals=%s\" %\\\n",
    "        (hcovlocal_scale/1000.,vcovlocal_fact,diff_efold,covinflate1,covinflate2,nanals))\n",
    "\n",
    "    # if nobs > 0, each ob time nobs ob locations are randomly sampled (without\n",
    "    # replacement) from the model grid\n",
    "    # if nobs < 0, fixed network of every Nth grid point used (N = -nobs)\n",
    "    nobs = -256         # number of obs to assimilate (randomly distributed)\n",
    "    #nobs = -1 # fixed network, every -nobs grid points. nobs=-1 obs at all pts.\n",
    "\n",
    "    # nature run\n",
    "    nc_truth = Dataset(filename_truth)\n",
    "    pv_truth = nc_truth.variables['pv']\n",
    "    # set up arrays for obs and localization function\n",
    "    if nobs < 0:\n",
    "        nskip = -nobs\n",
    "        if (nx*ny)%nobs != 0:\n",
    "            raise ValueError('nx*ny must be divisible by nobs')\n",
    "        nobs = (nx*ny)//nskip**2\n",
    "        print('# fixed network nobs = %s' % nobs)\n",
    "        fixed = True\n",
    "        nobs = nskip\n",
    "    else:\n",
    "        fixed = False\n",
    "        print('# random network nobs = %s' % nobs)\n",
    "    if nobs == nx*ny//2: fixed=True # used fixed network for obs every other grid point\n",
    "    oberrvar = oberrstdev**2*np.ones(nobs,float)\n",
    "    pvob = np.empty((2,nobs),float)\n",
    "    covlocal = np.empty((ny,nx),float)\n",
    "    covlocal_tmp = np.empty((nobs,nx*ny),float)\n",
    "    xens = np.empty((nanals,2,nx*ny),float)\n",
    "    if not use_letkf:\n",
    "        obcovlocal = np.empty((nobs,nobs),float)\n",
    "    else:\n",
    "        obcovlocal = None\n",
    "\n",
    "    if global_enkf: # model-space localization matrix\n",
    "        n = 0\n",
    "        covlocal_modelspace = np.empty((nx*ny,nx*ny),float)\n",
    "        x1 = x.reshape(nx*ny); y1 = y.reshape(nx*ny)\n",
    "        for n in range(nx*ny):\n",
    "            dist = cartdist(x1[n],y1[n],x1,y1,nc_climo.L,nc_climo.L)\n",
    "            covlocal_modelspace[n,:] = gaspcohn(dist/hcovlocal_scale)\n",
    "\n",
    "    obtimes = nc_truth.variables['t'][:]\n",
    "    if read_restart:\n",
    "        timeslist = obtimes.tolist()\n",
    "        ntstart = timeslist.index(tstart)\n",
    "        print('# restarting from %s.nc ntstart = %s' % (exptname,ntstart))\n",
    "    else:\n",
    "        ntstart = 0\n",
    "    assim_interval = obtimes[1]-obtimes[0]\n",
    "    assim_timesteps = int(np.round(assim_interval/models[0].dt))\n",
    "    print('# assim interval = %s secs (%s time steps)' % (assim_interval,assim_timesteps))\n",
    "    print('# ntime,pverr_a,pvsprd_a,pverr_b,pvsprd_b,obinc_b,osprd_b,obinc_a,obsprd_a,omaomb/oberr,obbias_b,inflation,tr(P^a)/tr(P^b)')\n",
    "\n",
    "    # initialize model clock\n",
    "    for nanal in range(nanals):\n",
    "        models[nanal].t = obtimes[ntstart]\n",
    "        models[nanal].timesteps = assim_timesteps\n",
    "\n",
    "    # initialize output file.\n",
    "    if savedata is not None:\n",
    "        nc = Dataset('%s.nc' % exptname, mode='w', format='NETCDF4_CLASSIC')\n",
    "        nc.r = models[0].r\n",
    "        nc.f = models[0].f\n",
    "        nc.U = models[0].U\n",
    "        nc.L = models[0].L\n",
    "        nc.H = models[0].H\n",
    "        nc.nanals = nanals\n",
    "        nc.hcovlocal_scale = hcovlocal_scale\n",
    "        nc.vcovlocal_fact = vcovlocal_fact\n",
    "        nc.oberrstdev = oberrstdev\n",
    "        nc.g = nc_climo.g; nc.theta0 = nc_climo.theta0\n",
    "        nc.nsq = models[0].nsq\n",
    "        nc.tdiab = models[0].tdiab\n",
    "        nc.dt = models[0].dt\n",
    "        nc.diff_efold = models[0].diff_efold\n",
    "        nc.diff_order = models[0].diff_order\n",
    "        nc.filename_climo = filename_climo\n",
    "        nc.filename_truth = filename_truth\n",
    "        nc.symmetric = models[0].symmetric\n",
    "        xdim = nc.createDimension('x',models[0].N)\n",
    "        ydim = nc.createDimension('y',models[0].N)\n",
    "        z = nc.createDimension('z',2)\n",
    "        t = nc.createDimension('t',None)\n",
    "        obs = nc.createDimension('obs',nobs)\n",
    "        ens = nc.createDimension('ens',nanals)\n",
    "        pv_t =\\\n",
    "        nc.createVariable('pv_t',np.float32,('t','z','y','x'),zlib=True)\n",
    "        pv_b =\\\n",
    "        nc.createVariable('pv_b',np.float32,('t','ens','z','y','x'),zlib=True)\n",
    "        pv_a =\\\n",
    "        nc.createVariable('pv_a',np.float32,('t','ens','z','y','x'),zlib=True)\n",
    "        pv_a.units = 'K'\n",
    "        pv_b.units = 'K'\n",
    "        inf = nc.createVariable('inflation',np.float32,('t','z','y','x'),zlib=True)\n",
    "        pv_obs = nc.createVariable('obs',np.float32,('t','obs'))\n",
    "        x_obs = nc.createVariable('x_obs',np.float32,('t','obs'))\n",
    "        y_obs = nc.createVariable('y_obs',np.float32,('t','obs'))\n",
    "        # eady pv scaled by g/(f*theta0) so du/dz = d(pv)/dy\n",
    "        xvar = nc.createVariable('x',np.float32,('x',))\n",
    "        xvar.units = 'meters'\n",
    "        yvar = nc.createVariable('y',np.float32,('y',))\n",
    "        yvar.units = 'meters'\n",
    "        zvar = nc.createVariable('z',np.float32,('z',))\n",
    "        zvar.units = 'meters'\n",
    "        tvar = nc.createVariable('t',np.float32,('t',))\n",
    "        tvar.units = 'seconds'\n",
    "        ensvar = nc.createVariable('ens',np.int32,('ens',))\n",
    "        ensvar.units = 'dimensionless'\n",
    "        xvar[:] = np.arange(0,models[0].L,models[0].L/models[0].N)\n",
    "        yvar[:] = np.arange(0,models[0].L,models[0].L/models[0].N)\n",
    "        zvar[0] = 0; zvar[1] = models[0].H\n",
    "        ensvar[:] = np.arange(1,nanals+1)\n",
    "\n",
    "    # initialize kinetic energy error/spread spectra\n",
    "    kespec_errmean = None; kespec_sprdmean = None\n",
    "\n",
    "    ncount = 0\n",
    "    nanals2 = 4 # ensemble members used for kespec spread\n",
    "\n",
    "    rmse = np.zeros(nassim)\n",
    "    #Jump\n",
    "    percentage  = 0.1\n",
    "    #print(\"percentage of noise steps\", percentage)\n",
    "    noiselevel=np.array((900,600))\n",
    "    noise_idx = np.array((1,1 - percentage*15/35 ,1-percentage))\n",
    "    jump_dist = rsjump.uniform(0,1,300)\n",
    "    jump = np.where(jump_dist > 1-percentage)[0] +1000\n",
    "    #print(jump)\n",
    "    # Initialize an empty dictionary\n",
    "    jump_idx = {}\n",
    "    for i in range(len(noiselevel)):\n",
    "        value = noiselevel[i]\n",
    "        keys = np.where((jump_dist <= noise_idx[i]) & (jump_dist > noise_idx[i+1]))[0].tolist()\n",
    "        for key in keys:\n",
    "            jump_idx[key] = value\n",
    "    #print(jump_idx)\n",
    "\n",
    "    indxob = np.sort(rsobs.choice(nx*ny,nobs,replace=False))\n",
    "    indx_unob = np.setdiff1d(np.arange(nx*ny), indxob)\n",
    "    obs_save = np.zeros((100,nobs))\n",
    "\n",
    "    for ntime in range(nassim): #nassim\n",
    "\n",
    "        # check model clock\n",
    "        if models[0].t != obtimes[ntime+ntstart]:\n",
    "            raise ValueError('model/ob time mismatch %s vs %s' %\\\n",
    "            (models[0].t, obtimes[ntime+ntstart]))\n",
    "\n",
    "        t1 = time.time()\n",
    "        if not fixed:\n",
    "            # randomly choose points from model grid\n",
    "            if nobs == nx*ny:\n",
    "                indxob = np.arange(nx*ny)\n",
    "            else:\n",
    "                indxob = np.sort(rsobs.choice(nx*ny,nobs,replace=False))\n",
    "        else:\n",
    "                if ntime == 0:\n",
    "                    print(\"not RRR\")\n",
    "        if (ntime == jump).any():\n",
    "            jumpnoise = rsjump.normal(0,jump_idx[ntime],size=(1,2,ny,nx)) \n",
    "            jumpnoise_reshape =jumpnoise.reshape(2,ny*nx)  \n",
    "        for k in range(2):\n",
    "            # surface temp obs\n",
    "            if (ntime == jump).any():\n",
    "                pvob[k] = scalefact*(pv_truth[ntime+ntstart,k,:,:].ravel()[indxob] +  jumpnoise_reshape[k])\n",
    "            else:\n",
    "                pvob[k] = scalefact*pv_truth[ntime+ntstart,k,:,:].ravel()[indxob]\n",
    "            pvob[k] += rsobs.normal(scale=oberrstdev,size=nobs) # add ob errors\n",
    "        xob = x.ravel()[indxob]\n",
    "        yob = y.ravel()[indxob]\n",
    "        # compute covariance localization function for each ob\n",
    "        if not fixed or ntime == 0:\n",
    "            for nob in range(nobs):\n",
    "                dist = cartdist(xob[nob],yob[nob],x,y,nc_climo.L,nc_climo.L)\n",
    "                covlocal = gaspcohn(dist/hcovlocal_scale)\n",
    "                covlocal_tmp[nob] = covlocal.ravel()\n",
    "                dist = cartdist(xob[nob],yob[nob],xob,yob,nc_climo.L,nc_climo.L)\n",
    "                if not use_letkf: obcovlocal[nob] = gaspcohn(dist/hcovlocal_scale)\n",
    "\n",
    "        # first-guess spread (need later to compute inflation factor)\n",
    "        fsprd = ((pvens - pvens.mean(axis=0))**2).sum(axis=0)/(nanals-1)\n",
    "\n",
    "        # compute forward operator.\n",
    "        # hxens is ensemble in observation space.\n",
    "        hxens = np.empty((nanals,2,nobs),float)\n",
    "        for nanal in range(nanals):\n",
    "            for k in range(2):\n",
    "                hxens[nanal,k,...] = scalefact*pvens[nanal,k,...].ravel()[indxob] # surface pv obs\n",
    "        hxensmean_b = hxens.mean(axis=0)\n",
    "        obsprd = ((hxens-hxensmean_b)**2).sum(axis=0)/(nanals-1)\n",
    "        # innov stats for background\n",
    "        obfits = pvob - hxensmean_b\n",
    "        obfits_b = (obfits**2).mean()\n",
    "        obbias_b = obfits.mean()\n",
    "        obsprd_b = obsprd.mean()\n",
    "        pvensmean_b = pvens.mean(axis=0).copy()\n",
    "        if (ntime == jump).any():\n",
    "            pverr_b = (scalefact*(pvensmean_b-pv_truth[ntime+ntstart]+jumpnoise))**2\n",
    "        else:\n",
    "            pverr_b = (scalefact*(pvensmean_b-pv_truth[ntime+ntstart]))**2\n",
    "        pvsprd_b = ((scalefact*(pvensmean_b-pvens))**2).sum(axis=0)/(nanals-1)\n",
    "\n",
    "        if savedata is not None:\n",
    "            if savedata == 'restart' and ntime != nassim-1:\n",
    "                pass\n",
    "            else:\n",
    "                pv_t[ntime] = pv_truth[ntime+ntstart]\n",
    "                pv_b[ntime,:,:,:] = scalefact*pvens\n",
    "                #pv_obs[ntime] = pvob\n",
    "                #x_obs[ntime] = xob\n",
    "                #y_obs[ntime] = yob\n",
    "\n",
    "        # EnKF update\n",
    "        # create 1d state vector.\n",
    "        xens = pvens.reshape(nanals,2,nx*ny)\n",
    "        # update state vector.\n",
    "        if direct_insertion and nobs == nx*ny:\n",
    "            print('wrongwrong!!!!!')\n",
    "            for nanal in range(nanals):\n",
    "                xens[nanal] =\\\n",
    "                pv_truth[ntime+ntstart].reshape(2,nx*ny) + \\\n",
    "                rsobs.normal(scale=oberrstdev,size=(2,nx*ny))/scalefact\n",
    "            xens = xens - xens.mean(axis=0) + \\\n",
    "            pv_truth[ntime+ntstart].reshape(2,nx*ny) + \\\n",
    "            rsobs.normal(scale=oberrstdev,size=(2,nx*ny))/scalefact\n",
    "        else:\n",
    "            # hxens,pvob are in PV units, xens is not \n",
    "            if global_enkf and not use_letkf:\n",
    "                xens = bulk_ensrf(xens,indxob,pvob,oberrvar,covlocal_modelspace,vcovlocal_fact,scalefact)\n",
    "            else:\n",
    "                xens =\\\n",
    "                enkf_update(xens,hxens,pvob,oberrvar,covlocal_tmp,vcovlocal_fact,obcovlocal=obcovlocal)\n",
    "        # back to 3d state vector\n",
    "        pvens = xens.reshape((nanals,2,ny,nx))\n",
    "        t2 = time.time()\n",
    "        #print('cpu time for EnKF update',t2-t1)\n",
    "\n",
    "        # forward operator on posterior ensemble.\n",
    "        for nanal in range(nanals):\n",
    "            for k in range(2):\n",
    "                hxens[nanal,k,...] = scalefact*pvens[nanal,k,...].ravel()[indxob] # surface pv obs\n",
    "\n",
    "        # ob space diagnostics\n",
    "        hxensmean_a = hxens.mean(axis=0)\n",
    "        obsprd_a = (((hxens-hxensmean_a)**2).sum(axis=0)/(nanals-1)).mean()\n",
    "        # expected value is HPaHT (obsprd_a).\n",
    "        obinc_a = ((hxensmean_a-hxensmean_b)*(pvob-hxensmean_a)).mean()\n",
    "        # expected value is HPbHT (obsprd_b).\n",
    "        obinc_b = ((hxensmean_a-hxensmean_b)*(pvob-hxensmean_b)).mean()\n",
    "        # expected value R (oberrvar).\n",
    "        omaomb = ((pvob-hxensmean_a)*(pvob-hxensmean_b)).mean()\n",
    "\n",
    "        # posterior multiplicative inflation.\n",
    "        pvensmean_a = pvens.mean(axis=0)\n",
    "        pvprime = pvens-pvensmean_a\n",
    "        asprd = (pvprime**2).sum(axis=0)/(nanals-1)\n",
    "        asprd_over_fsprd = asprd.mean()/fsprd.mean()\n",
    "        if covinflate2 < 0:\n",
    "            # relaxation to prior stdev (Whitaker & Hamill 2012)\n",
    "            asprd = np.sqrt(asprd); fsprd = np.sqrt(fsprd)\n",
    "            inflation_factor = 1.+covinflate1*(fsprd-asprd)/asprd\n",
    "        else:\n",
    "            # Hodyss et al 2016 inflation (covinflate1=covinflate2=1 works well in perfect\n",
    "            # model, linear gaussian scenario)\n",
    "            # inflation = asprd + (asprd/fsprd)**2((fsprd/nanals)+2*inc**2/(nanals-1))\n",
    "            inc = pvensmean_a - pvensmean_b\n",
    "            inflation_factor = covinflate1*asprd + \\\n",
    "            (asprd/fsprd)**2*((fsprd/nanals) + covinflate2*(2.*inc**2/(nanals-1)))\n",
    "            inflation_factor = np.sqrt(inflation_factor/asprd)\n",
    "        pvprime = pvprime*inflation_factor\n",
    "        pvens = pvprime + pvensmean_a\n",
    "\n",
    "        # print out analysis error, spread and innov stats for background\n",
    "        if (ntime == jump).any():\n",
    "            print('jumppp')\n",
    "            pverr_a = (scalefact*(pvensmean_a-pv_truth[ntime+ntstart]-jumpnoise))**2\n",
    "        else:\n",
    "            pverr_a = (scalefact*(pvensmean_a-pv_truth[ntime+ntstart]))**2\n",
    "        pvsprd_a = ((scalefact*(pvensmean_a-pvens))**2).sum(axis=0)/(nanals-1)\n",
    "        print(\"%s %g %g %g %g %g %g %g %g %g %g %g %g\" %\\\n",
    "        (ntime+ntstart,np.sqrt(pverr_a.mean()),np.sqrt(pvsprd_a.mean()),\\\n",
    "        np.sqrt(pverr_b.mean()),np.sqrt(pvsprd_b.mean()),\\\n",
    "        obinc_b,obsprd_b,obinc_a,obsprd_a,omaomb/oberrvar.mean(),obbias_b,inflation_factor.mean(),asprd_over_fsprd))\n",
    "        \n",
    "        rmse[ntime] = np.sqrt(pverr_a.mean())\n",
    "        \n",
    "        #np.save(runing_rmse,rmse)\n",
    "        # save data.\n",
    "        if savedata is not None:\n",
    "            if savedata == 'restart' and ntime != nassim-1:\n",
    "                pass\n",
    "            else:\n",
    "                pv_a[ntime,:,:,:] = scalefact*pvens\n",
    "                tvar[ntime] = obtimes[ntime+ntstart]\n",
    "                inf[ntime] = inflation_factor\n",
    "                nc.sync()\n",
    "\n",
    "        # run forecast ensemble to next analysis time\n",
    "        t1 = time.time()\n",
    "        for nanal in range(nanals):\n",
    "            pvens[nanal] = models[nanal].advance(pvens[nanal])\n",
    "        t2 = time.time()\n",
    "        #print('cpu time for ens forecast',t2-t1)\n",
    "\n",
    "\n",
    "        # compute spectra of error and spread\n",
    "        if ntime >= nassim_spinup:\n",
    "            pvfcstmean = pvens.mean(axis=0)\n",
    "            pverrspec = scalefact*rfft2(pvfcstmean - pv_truth[ntime+ntstart])\n",
    "            psispec = models[0].invert(pverrspec)\n",
    "            psispec = psispec/(models[0].N*np.sqrt(2.))\n",
    "            kespec = (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "            if kespec_errmean is None:\n",
    "                kespec_errmean =\\\n",
    "                (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "            else:\n",
    "                kespec_errmean = kespec_errmean + kespec\n",
    "            for nanal in range(nanals2):\n",
    "                pvsprdspec = scalefact*rfft2(pvens[nanal] - pvfcstmean)\n",
    "                psispec = models[0].invert(pvsprdspec)\n",
    "                psispec = psispec/(models[0].N*np.sqrt(2.))\n",
    "                kespec = (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real\n",
    "                if kespec_sprdmean is None:\n",
    "                    kespec_sprdmean =\\\n",
    "                    (models[0].ksqlsq*(psispec*np.conjugate(psispec))).real/nanals2\n",
    "                else:\n",
    "                    kespec_sprdmean = kespec_sprdmean+kespec/nanals2\n",
    "            ncount += 1\n",
    "\n",
    "    if savedata: nc.close()\n",
    "\n",
    "    #np.save(runing_rmse,rmse)\n",
    "\n",
    "    return None\n",
    "\n",
    "\n",
    "## Run main code\n",
    "if __name__ == \"__main__\":\n",
    "    submit()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}