src/make_kmz.py

   1 from __future__ import absolute_import
   2 from __future__ import print_function
   3 import simplekml as kml
   4 import os.path as osp
   5 import json
   6 import sys
   7 import logging
   8 from utils import update_nested_dict, load_sys_cfg
   9 from six.moves.urllib.parse import urljoin
  10 import posixpath as pxp
  11 import requests
  12 from six.moves import map
  13 from six.moves import zip
  14
  15 sys_cfg = load_sys_cfg()
  16 sys_cfg.sims_path = 'fdds/simulations'
  17 sys_cfg.sims_url_path = 'simulations'
  18
  19 def make_kmz(job_id, steps, mode, only_vars):
  20     """
  21     Create KMZ file from visualization stored in wrfxweb.
  22     :param job_id: string, the name of job directory
  23     :param steps: string '1,1,1,3' takes every 3rd frame in domain 4, etc. Default: all 1
  24     :param mode: string, 'inc', ' to include image files (default), 'ref' to use links only
  25     :param only_var: list of strings variables to include or None to include  all
  26     """
  27
  28     logging.info('make_kmz: job_id=%s' % job_id)
  29     job_path = osp.join(osp.abspath(sys_cfg.sims_path),job_id)
  30     url_prefix = pxp.join(sys_cfg.url_root,sys_cfg.sims_url_path,job_id)
  31     logging.debug('make_kmz: job_path %s' % job_path)
  32     logging.debug('make_kmz: url_prefix %s' % url_prefix)
  33
  34
  35     if mode == '' or mode == "inc":
  36         kmz_filename = job_id + '_inc.kmz'
  37         href_prefix = osp.abspath(job_path)
  38         href_join = osp.join
  39         logging.debug('make_kmz: kmz file will include images from %s' % href_prefix)
  40         logging.info('make_kmz: kmz file will include images')
  41     elif mode == "ref":
  42         kmz_filename = job_id + '_ref.kmz'
  43         href_prefix = url_prefix
  44         href_join = pxp.join
  45         logging.debug('make_kmz: kmz file will link to images from %s' % href_prefix)
  46         logging.info('make_kmz: kmz file will link to images')
  47     else:
  48         logging.error('make_kmz: arg 2 must be "inc" or "ref" or omitted')
  49         exit(1)
  50
  51     # read the catalog and the manifest
  52     cat_path = osp.join(job_path,'catalog.json')
  53     cat = json.load(open(cat_path))
  54     if job_id not in cat:
  55          logging.error('job id %s not in the catalog' % job_id)
  56          sys.exit(1)
  57     cat_entry = cat[job_id]
  58     mf = json.load(open(osp.join(sys_cfg.sims_path,cat_entry['manifest_path'])))
  59
  60     mdomain = max(list(map(int, list(mf.keys()))))
  61     if steps=='':
  62         step=[1]
  63     else:
  64         step=list(map(int, steps.split(',')))
  65     if len(step) == 1:
  66         step = step*mdomain
  67     elif len(step) != mdomain:
  68         logging.error('make_kmz: steps needed for all up to max domain number = %s' % mdomain)
  69         sys.exit(1)
  70
  71     description = cat_entry['description']
  72     logging.info('make_kmz: job description: %s' % description)
  73
  74     # transpose var and time in manifest, output to frame
  75     frame={}
  76     for domain in mf:
  77         for ts_esmf in mf[domain]:
  78             for var in mf[domain][ts_esmf]:
  79                 if only_vars is None or var in only_vars:
  80                     update_nested_dict(frame,{domain:{var:{ts_esmf:mf[domain][ts_esmf][var]}}})
  81
  82     doc = kml.Kml(name=description)
  83
  84     for domain in sorted(frame):
  85         domain_folder = doc.newfolder(name = domain)
  86         istep = step[int(domain)-1]
  87         logging.info('make_kmz: processing domain %s step %s' % (domain, istep))
  88         for var in frame[domain]:
  89             var_folder = domain_folder.newfolder(name = var)
  90             ts_esmf = sorted(frame[domain][var].keys())
  91             ts_esmf = ts_esmf[1::istep]
  92             ts_esmf1 = ts_esmf[1:]
  93             ts_esmf1.append(None)
  94             for ts_esmf,ts_esmf1 in zip(ts_esmf,ts_esmf1):
  95                 ts_folder = var_folder.newfolder(name = ts_esmf)
  96                 ts_folder.timespan.begin = ts_esmf.replace('_','T')+'Z'
  97                 if ts_esmf1 is not None:
  98                      ts_folder.timespan.end = ts_esmf1.replace('_','T')+'Z'
  99                 frame_data = frame[domain][var][ts_esmf]
 100                 raster_path = frame_data['raster']
 101                 coords = frame_data['coords']
 102                 if 'colorbar' in frame_data:
 103                     # add colorbar to KMZ
 104                     cb_path = frame_data['colorbar']
 105                     cbo = ts_folder.newscreenoverlay(name='colorbar')
 106                     cbo.overlayxy = kml.OverlayXY(x=0,y=1,xunits=kml.Units.fraction,yunits=kml.Units.fraction)
 107                     cbo.screenxy = kml.ScreenXY(x=0.02,y=0.95,xunits=kml.Units.fraction,yunits=kml.Units.fraction)
 108                     cbo.size = kml.Size(x=150,y=300,xunits=kml.Units.pixel,yunits=kml.Units.pixel)
 109                     cbo.color = kml.Color.rgb(255,255,255,a=150)
 110                     cbo.visibility = 0
 111                     cbo.icon.href = href_join(href_prefix,cb_path)
 112
 113                 # add ground overlay
 114                 ground = ts_folder.newgroundoverlay(name=var,color='80ffffff')
 115                 ground.gxlatlonquad.coords = coords
 116                 ground.visibility = 0
 117                 ground.icon.href = href_join(href_prefix,raster_path)
 118
 119     # build output file
 120     kmz_path = osp.join(job_path,kmz_filename)
 121     logging.info('make_kmz: creating file %s' % kmz_path)
 122     doc.savekmz(kmz_path)
 123     url = pxp.join(url_prefix,kmz_filename)
 124     logging.info('make_kmz: file created at %s' % url)
 125     try:
 126         r = requests.get(url, stream=True)
 127         content_size = int(r.headers['Content-Length'])
 128         logging.info('make_kmz: file size is %s' % content_size)
 129         cat[job_id]['kml_url']=url
 130         cat[job_id]['kml_size']=content_size
 131         json.dump(cat, open(cat_path,'w'), indent=4, separators=(',', ': '))
 132     except Exception as e:
 133         logging.warning('make_kmz: accessing the file over the web failed with exception %s' % e)
 134
 135
 136
 137
 138
 139
 140 if __name__ == '__main__':
 141
 142     sys.argv
 143
 144     if len(sys.argv) < 2:
 145         print('usage: make_kmz.sh job_id steps mode only_vars variable1 variable2 ...')
 146         print('job_id: the name of job directory in ' + sys_cfg.sims_path)
 147         print("steps: '1,1,1,3' takes every 3rd frame in domain 4, etc. Default: all 1")
 148         print("mode: inc to include image files (default), ref to use links only")
 149         print('variable (optional): variables to include; if absent all will be included')
 150         sys.exit(1)
 151
 152     logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 153
 154     job_id = sys.argv[1]
 155
 156     if len(sys.argv) >= 3:
 157         steps = sys.argv[2]
 158     else:
 159         steps = ''
 160
 161     if len(sys.argv) >= 4:
 162         mode = sys.argv[3]
 163     else:
 164         mode = "inc"
 165
 166     if len(sys.argv) >= 5:
 167         only_vars = sys.argv[4:]
 168     else:
 169         only_vars = None
 170
 171     make_kmz(job_id, steps, mode, only_vars)
 172