mix_climate

#!/usr/bin/env python

from __future__ import print_function
from argparse import ArgumentParser

import flo_utils as fu
import json
from math import floor, ceil
import netCDF4 as nc
import numpy as np
import os
import re
import subprocess as sp
import sys



def combine_files (factors, output, options):
  config = json.load(open(options.config))
  filesets = config["filesets"]
  varnames = list(filesets.values())[0].keys()
  for varname in varnames:
    if varname == "smb":
      for (n, fileset) in enumerate(factors.keys()):
        if filesets[fileset].get(varname, False):
          f = factors[fileset]
          if f > 1:
            f=1
          elif f < 0 :
            f = 0
          tens_l = "%01d0"%(floor(f*10))
          tens_h = "%01d0"%(ceil(f*10))
          remains_h = f-floor(f)
          remains_l = 1-remains_h
          path_l = re.sub("<TENS>", tens_l, filesets[fileset][varname])
          path_h = re.sub("<TENS>", tens_h, filesets[fileset][varname])
          arguments = ["cdo"]
          arguments = arguments + ["-add", "-mulc,%s"%(remains_l), path_l ]
          arguments = arguments + ["-mulc,%s"%(remains_h), path_h ]
    else:
      arguments = ["cdo"]
      for (n, fileset) in enumerate(factors.keys()):
        if (n < len(factors.keys())-1):
          arguments = arguments + ["-add", "-mulc,%s"%(factors[fileset]), filesets[fileset][varname] ]
        else:
          arguments = arguments + ["-mulc,%s"%(factors[fileset]), filesets[fileset][varname] ]
    arguments.append ("%s_%s.nc"%(options.output, varname))
    fu.cerr (fu.qo (arguments))




def get_files_and_factors(options):
  config = json.load(open(options.config))
  ts_file = nc.Dataset(config["ts_file"])
  ts_vars = [ x for x in ts_file.variables.keys() if x not in ["time"]]
  ts_data = {ts: ts_file.variables[ts] for ts in ts_vars}
  time = ts_file.variables["time"]
  factors = {ts: np.interp(options.year, time[:], ts_data[ts][:]) for ts in ts_vars}
  fu.cerr ("mixing ratios:")
  fu.cerr ("\n".join([ "%s : %s"%(k, factors[k]) for k in factors.keys()]))
  return (factors) # dictionary filename : factor


def parse_args():
  parser = ArgumentParser()
  parser.description = "Combine climate files to create intermediate climate states"
#  parser.add_argument ("startyear", help='First year to use in average',  type=int)
#  parser.add_argument ("endyear", help='Last year to use in average',  type=int)
  #  parser.add_argument ("run", help='run to average')
  parser.add_argument("-c", "--config",
                      help='''Config file (required)''', required = True)
  parser.add_argument("-y", "--year",
                      help='''year to use for obtaining coefficients''', type=float, required = True)
  parser.add_argument("-o", "--output",
                      help='''output file name''', required=True)

  parser.add_argument("-v", "--verbose",
                    help='''Be verbose''', action="store_true")

  options = parser.parse_args()
  if options.verbose:
      fu.cerr(dir(options))
  return options

def main():
  options = parse_args()
  factors = get_files_and_factors(options)
  combine_files(factors, options.output, options)

if __name__ == "__main__":
    main()