view_restart.py

#!/usr/bin/env python


# Tobey Carman, June 2014
# Spatial Ecology Lab

# Quick prototype script for making images to "see" dvm-dos-tem restart files.
# Not meant for any sort of scientific analysis. Simply a way to take a quick
# glance at a netCDF restart file that is generated by the model and get an
# idea of the variables and how things are setup.

import argparse
import textwrap
import itertools

import matplotlib.pyplot as plt
import numpy as np
import netCDF4 as nc
import matplotlib.gridspec as gridspec


import collections

import logging


# turns out there is no need to use this - just get a numpy array straight
# from a netcdf variable, slice, and then reshape as needed.
def flat_gen(x):
  '''Generator to flatten arbitrarily deep nested list...'''

  # Basically copied from here:
  # http://stackoverflow.com/questions/16176742/python-3-replacement-for-deprecated-compiler-ast-flatten-function

  # might be dangerous if we have even deeper nested structures - too deep of
  # recursion?? But for now it should work..

  # Note that I encountered some problems testing this in the embedded IPython
  # shell. See here:
  # https://github.com/ipython/ipython/issues/62

  def iselement(e):
    return not(isinstance(e, collections.Iterable) and not isinstance(e, str))
  for el in x:
    if iselement(el):
      yield el
    else:
      for sub in flat_gen(el): yield sub

# try it out:
#print(list(flat_gen(["junk",["nested stuff"],[],[[[],['deep']]]])))

print "Finding all variables that are in terms of each 3D dimension group..."

def plot_3d_groups(dataset, dimension_group, gs, row=0):
  #from IPython import embed; embed()
  for i, dim_grp in enumerate(dimension_group):
    logging.info("Axes {0:}  {1:}".format(i, dim_grp))
    vars = [var for var in dataset.variables if dataset.variables[var].dimensions == dim_grp]
    
    ax = plt.subplot(gs[row,i])

    s = dataset.variables[vars[0]].shape
    flatdata = np.vstack( (dataset.variables[vars[0]][:].reshape((s[0]*s[1], s[2]))) )
    for j, v in enumerate(vars):
      logging.info("  {0:} {1:}".format(v, dataset.variables[v].shape))
      
      s = dataset.variables[v].shape
      flatdata = np.vstack( (flatdata, dataset.variables[v][:].reshape((s[0]*s[1], s[2]))) )
  
    ax.set_xticklabels(())
    ax.set_xticks(())
    ax.set_xticklabels(())
    
    ax.set_title( "(%s)" % ('A' if (row==0) else 'B') )
    ax.set_yticks(())
    ax.set_yticklabels(())
    #ax.set_yticks( range(0, len(vars)) )
    ax.set_ylabel("%s, %s, %s" % (dim_grp[0], dim_grp[1], dim_grp[2]))
    
    logging.info("(rows, cols) of image being plotted: %s" % ["%s"%i for i in flatdata.shape])
    
    #from IPython import embed; embed()
    
    plt.imshow(dataset.variables[v][:].reshape((s[0]*s[1], s[2])), aspect=1.0, interpolation='nearest')



def plot_2d_groups(dataset, dimension_group, gs, col=0):
  
  logging.info("Finding all variables that are in terms of each 2D dimension group...")
  for i, dim_grp in enumerate(dimension_group):
    logging.info("Axes {0:}  {1:}".format(i, dim_grp))
    vars = [var for var in dataset.variables if dataset.variables[var].dimensions == dim_grp]

#    if len(plt.gcf().axes) > 1:
#      ax = plt.subplot(gs[i,col], sharex=plt.gcf().axes[0])
#    else:
#      ax = plt.subplot(gs[i,col])

    ax = plt.subplot(gs[i,col])

    flatdata = list(itertools.chain.from_iterable( dataset.variables[vars[0]] ))
    for j, v in enumerate(vars):
      logging.info( "  {0:} {1:}".format(v, dataset.variables[v].shape) )
      if not j == 0:
        flatdata = np.vstack( (flatdata, list(itertools.chain.from_iterable( dataset.variables[vars[j]]))) )
      else:
        pass # already added the first var's data ouside loop
  
    #ax.set_title("??")
    
    ax.set_xlabel("%s, %s" % (dim_grp[0], dim_grp[1]))
    ax.set_xticks(())
    ax.set_xticklabels(())
    
    ax.set_ylabel("%i" % i)
    ax.set_yticks( range(0, len(vars)) )
    ax.set_yticklabels(())
    
    logging.info("(rows, cols) of image being plotted: %s" % ["%s"%i for i in flatdata.shape])
    #from IPython import embed; embed()
    plt.imshow(flatdata, interpolation='nearest', aspect=1.0) # 4 -> height is 4x the width
    if len(plt.gcf().axes) > 1:
      plt.imshow(flatdata, interpolation='nearest', aspect=1.0) # 4 -> height is 4x the width




def main(fileA, compareFile):
  logging.basicConfig(level=logging.DEBUG)
  logging.info("Loading dataset(s)")
  dsA = nc.Dataset(fileA)
  if compareFile:
    dsB = nc.Dataset(compareFile)


  logging.info("Finding set of 'dimension groups' that covers all variable in the file")
  # make a list of dimesions for every variable in the file
  # reduce to a set of 'dimension groups' that covers all variables in each file
  dim_grpA = set( [dsA.variables[v].dimensions for v in dsA.variables] )

  if compareFile:
    dim_grpB = set( [dsB.variables[v].dimensions for v in dsB.variables] )
    if dim_grpA != dim_grpB:
      logging.warn(")The two files don't seem to have the same variables/dimensions!")

  logging.info("Separating dimension groups into 2D group and 3D group")
  dim_grpA_2D = [g for g in dim_grpA if len(g) == 2]
  dim_grpA_3D = [g for g in dim_grpA if len(g) == 3]

  if compareFile:
    dim_grpB_2D = [g for g in dim_grpB if len(g) == 2]
    dim_grpB_3D = [g for g in dim_grpB if len(g) == 3]
    if (dim_grpA_2D != dim_grpB_2D) or (dim_grpA_3D != dim_grpB_3D):
      logging.warn("Problem with not matching between two files!")

  logging.info("2D dimension groups: %s" % dim_grpA_2D)
  logging.info("3D dimension groups: %s" % dim_grpA_3D)

  # Set up figure for 2D image plots
  gs = gridspec.GridSpec( len(dim_grpA_2D), 1 )
  if compareFile:
    gs = gridspec.GridSpec( len(dim_grpB_2D), 2 )

  fig = plt.figure()
  fig.suptitle("2D variables for %s" % fileA)
  if compareFile:
    fig.suptitle(textwrap.dedent('''2D variables for
        (A) %s
        (B) %s''' % (fileA, compareFile)))

  plot_2d_groups(dsA, dim_grpA_2D, gs, col=0)

  if compareFile:
    plot_2d_groups(dsB, dim_grpB_2D, gs, col=1)

  plt.show()



  gs = gridspec.GridSpec( 1, len(dim_grpA_3D) )
  if compareFile:
    gs = gridspec.GridSpec( 2, len(dim_grpB_3D) )

  fig = plt.figure()
  fig.suptitle("3D variables for %s" % fileA)
  if compareFile:
    fig.suptitle(textwrap.dedent('''3D variables for
      (A) %s
      (B) %s''' % (fileA, compareFile)))

  plot_3d_groups(dsA, dim_grpA_3D, gs, row=0)
  if compareFile:
    plot_3d_groups(dsB, dim_grpB_3D, gs, row=1)

  plt.show()




if __name__ == '__main__':
  
  parser = argparse.ArgumentParser(

    formatter_class=argparse.RawDescriptionHelpFormatter,
      
      description=textwrap.dedent('''\
        Show some restart files as images.
        '''),
        
      epilog=textwrap.dedent('''''')
  )
  
  parser.add_argument('file', default='',
      help="A file to show")

  parser.add_argument('--compare', default=None,
      help='path to a NetCDF file (B) to plot for comparison')

  print "Parsing command line arguments..."
  args = parser.parse_args()

  main(args.file, args.compare)