content: """
  Module taking care of the forward modelling: theoretical dispersion
  curve given a 1D crustal model of velocities and densities.
  Uses the binaries of the Computer Programs in Seismology, with
  must be installed in *COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR*
  """
  
  import numpy as np
  import matplotlib.pyplot as plt
  import os
  import shutil
  import itertools as it
  from easyprocess import EasyProcess
  import tempfile
  import pickle
  
  # getting the dir of the binaries of the Computer Programs in Seismology
  
  # import CONFIG class initalised in ./configs/tmp_config.pickle
  config_pickle = 'configs/tmp_config.pickle'
  f = open(name=config_pickle, mode='rb')
  CONFIG = pickle.load(f)
  f.close()
      
  # import variables from initialised CONFIG class.
  COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR=CONFIG.COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR
  
  # default header of the model file:
  # isotropic, 1D, flat Earth with layers of constant velocity
  MODEL_HEADER = """MODEL.01
  TEST
  ISOTROPIC
  KGS
  FLAT EARTH
  1-D
  CONSTANT VELOCITY
  LINE08
  LINE09
  LINE10
  LINE11
  H   VP  VS RHO  QP  QS ETAP ETAS FREFP FREFS"""
  
  
  class VsModel:
      """
      Class holding a layered model of Vs function of depth,
      with Vp/Vs and rho/Vs ratio fixed.
      """
      def __init__(self, vs, dz, ratio_vp_vs, ratio_rho_vs, name='',
                   store_vg_at_periods=None):
          """
          Initializes model with layers' Vs (vs), layers' thickness (dz),
          and layers' ratio Vp/Vs and rho/Vs (ratio_vp_vs, ratio_rho_vs).
          """
          # checking shapes
          nlayers = np.size(vs)
          if np.size(dz) != nlayers - 1:
              raise Exception("Size of dz should be nb of layers minus 1")
          if not np.size(ratio_vp_vs) in [1, nlayers]:
              raise Exception("Size of ratio_vp_vs should be nb of layers or 1")
          if not np.size(ratio_rho_vs) in [1, nlayers]:
              raise Exception("Size of ratio_rho_vs should be nb of layers or 1")
  
          self.name = name
          self.vs = np.array(vs)
          self.dz = np.array(dz)
          self.ratio_vp_vs = np.array(ratio_vp_vs)
          self.ratio_rho_vs = np.array(ratio_rho_vs)
  
          # storing vg model at selected periods if required
          self.stored_vgperiods = store_vg_at_periods
          if not store_vg_at_periods is None:
              self.stored_vg = self.vg_model(store_vg_at_periods)
          else:
              self.stored_vg = None
  
      def misfit_to_vg(self, periods, vg, sigmavg, squared=True,
                       use_storedvg=True, storevg=False):
          """
          Misfit of modelled vg to observed vg
  
                [vg_model - vg]**2
          = Sum ------------------  over periods
                  2 x sigmavg**2
          """
          # using stored vg model if required and available, else re-calculating it
          if use_storedvg and np.all(periods == self.stored_vgperiods):
              vg_model = self.stored_vg
          else:
              vg_model = self.vg_model(periods, store=storevg)
  
          misfit = np.sum(((vg_model - vg) / sigmavg)**2) / 2.0
          if squared:
              misfit = np.sqrt(misfit)
          return misfit
  
      def vg_model(self, periods, store=False):
          """
          Modelled group velocities, vg, function of period
          """
          vs = self.vs
          vp = self.ratio_vp_vs * self.vs
          rho = self.ratio_rho_vs * self.vs
          dz = np.r_[self.dz, 0]  # we append a fake thickness
          vg = Rayleigh_group_velocities(periods, dz=dz, vp=vp, vs=vs, rho=rho)
          if store:
              # storing group velocities if required
              self.stored_vgperiods = periods
              self.stored_vg = vg
          return vg
  
      def get_vs_at(self, z):
          """
          Returns Vs ad depth(s) *z*
          """
          indices = np.searchsorted(np.r_[0, self.dz.cumsum()], z, side='right') - 1
          if np.any(indices) < 0:
              raise Exception("Depth out of range")
          return self.vs[indices]
  
      def plot(self, periods, obsvgarrays=None, fig=None, color='r'):
          """
          Plots modelled and observed group velocity function of period (top)
          and the model itself, i.e. Vs vs depth (bottom)
          """
          if not fig:
              fig = plt.figure(figsize=(6.5, 10), tight_layout=True)
              axlist = [fig.add_subplot(211), fig.add_subplot(212)]
              legend = True
          else:
              axlist = fig.get_axes()
              legend = False  # no need to add legend to existing fig
  
          # 1st subplot: group velocity vs period
          ax = axlist[0]
          self.plot_vg(periods, obsvgarrays=obsvgarrays, ax=ax, legend=legend, color=color)
          ax.set_title(self.name)
  
          # 2nd subplot: Vs vs depth
          ax = axlist[1]
          self.plot_model(ax=ax, color=color)
  
          fig.canvas.draw()
          fig.show()
          return fig
  
      def plot_vg(self, periods, obsvgarrays=None, ax=None, legend=True, color='r'):
          """
          Plots modelled and observed group velocity function of period
          """
          # creating figure if not given as input
          fig = None
          if not ax:
              fig = plt.figure()
              ax = fig.add_subplot(111)
  
          vg_model = self.vg_model(periods)
          ax.plot(periods, vg_model, lw=1.5, color=color, label=self.name)
          if obsvgarrays:
              for i, vgarray in enumerate(obsvgarrays):
                  label = 'Observed dispersion curves' if not i else None
                  ax.plot(periods, vgarray, lw=0.5, color='k', label=label)
          ax.set_xlabel('Period (sec)')
          ax.set_ylabel('Group velocity (km/s)')
          if legend:
              ax.legend(loc='best', fontsize=11, framealpha=0.8)
          ax.grid(True)
  
          if fig:
              fig.show()
  
      def plot_model(self, ax=None, color='r', format_axes=True):
          """
          Plots the model, i.e. Vs vs depth
          """
          # creating figure if not given as input
          fig = None
          if not ax:
              fig = plt.figure()
              ax = fig.add_subplot(111)
  
          x = list(it.chain.from_iterable([[v, v] for v in self.vs]))
          y = [0.0] + list(it.chain.from_iterable([[z, z] for z in np.cumsum(self.dz)])) + \
              [self.dz.sum() + 15]
          ax.plot(x, y, lw=1.5, color=color)
  
          if format_axes:
              ax.set_ylim(sorted(ax.get_ylim(), reverse=True))
              ax.set_xlabel('Vs (km/s)')
              ax.set_ylabel('Depth (km)')
              ax.grid(True)
  
          if fig:
              fig.show()
  
  
  def Rayleigh_group_velocities(periods, dz, vp, vs, rho, verbose=False):
      """
      Returns the array of Rayleigh wave group velocities at
      selected periods, from the 1-D layered Earth model
      contained in *dz* (thicknesses), *vp* (P wave velocities),
      *vs* (S wave velocities) and *rho* (densities).
  
      The Computer Programs in Seismology, located in dir
      *COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR*,
      are used for the computation.
      """
      if not COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR:
          raise Exception("Please provide the dir of the Computer Programs in Seismology")
  
      # making and moving to temporary dir
      current_dir = os.getcwd()
      tmp_dir = tempfile.mkdtemp()
      os.chdir(tmp_dir)
  
      # preparing input files
      if verbose:
          print 'Preparing model and periods files'
      create_model_file('model', dz, vp, vs, rho)
      f = open('periods', 'w')
      f.write('\n'.join([str(p) for p in periods]))
      f.close()
  
      # preparing model
      if verbose:
          print "Calling sprep96"
      cmd = os.path.join(COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR, 'sprep96')
      # Rayleigh wave, fundamental mode
      p = EasyProcess('"{}" -M model -PARR periods -NMOD 1 -R'.format(cmd)).call()
      if verbose:
          print p.stdout
  
      # phase dispersion curve
      if verbose:
          print "Calling sdisp96"
      cmd = os.path.join(COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR, 'sdisp96')
      p = EasyProcess('"{}" -v'.format(cmd)).call()
      if verbose:
          print p.stdout
  
      # group dispersion curve
      if verbose:
          print "Calling sregn96"
      cmd = os.path.join(COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR, 'sregn96')
      p = EasyProcess('"{}"'.format(cmd)).call()
      if verbose:
          print p.stdout
  
      # exporting group velocities (-U) of Rayleigh waves (-R) in ascii file
      if verbose:
          print "Calling sdpegn96"
      cmd = os.path.join(COMPUTER_PROGRAMS_IN_SEISMOLOGY_DIR, 'sdpegn96')
      p = EasyProcess('"{}" -R -S -U -XLOG -PER -ASC'.format(cmd)).call()
      if verbose:
          print p.stdout
  
      # loading group velocities from 6th column of ascii file
      vg = np.loadtxt('SREGN.ASC', skiprows=1, usecols=(5,))
  
      # removing temp dir
      os.chdir(current_dir)
      shutil.rmtree(tmp_dir)
  
      return vg
  
  
  def create_model_file(path, dz, vp, vs, rho):
      """
      Writing the 1D model to ascci file, to be used as input
      by the Computer Programs in Seismology
      """
      qp = np.zeros_like(dz)
      qs = np.zeros_like(dz)
      etap = np.zeros_like(dz)
      etas = np.zeros_like(dz)
      frefp = np.ones_like(dz)
      frefs = np.ones_like(dz)
  
      f = open(path, mode='w')
      f.write(MODEL_HEADER)
  
      a = np.vstack((dz, vp, vs, rho, qp, qs, etap, etas, frefp, frefs))
      for col in a.T:
          f.write('\n')
          col.tofile(f, sep=' ')
  
      f.close()