code: """
  Classes and functions for templates.
  """
  
  from __future__ import absolute_import, division, print_function
  
  import sys
  from glob import glob
  import os
  import traceback
  
  import numpy as np
  from astropy.io import fits
  
  from .utils import native_endian, elapsed, transmission_Lyman
  
  from .rebin import rebin_template, trapz_rebin
  
  
  class Template(object):
      """A spectral Template PCA object.
  
      The template data is read from a redrock-format template file.
      Alternatively, the data can be specified in the constructor.
  
      Args:
          filename (str): the path to the template file, either absolute or
              relative to the RR_TEMPLATE_DIR environment variable.
  
      """
      def __init__(self, filename=None, spectype=None, redshifts=None,
          wave=None, flux=None, subtype=None):
  
          if filename is not None:
              fx = None
              if os.path.exists(filename):
                  fx = fits.open(filename, memmap=False)
              else:
                  xfilename = os.path.join(os.getenv('RR_TEMPLATE_DIR'), filename)
                  if os.path.exists(xfilename):
                      fx = fits.open(xfilename, memmap=False)
                  else:
                      raise IOError('unable to find '+filename)
  
              hdr = fx['BASIS_VECTORS'].header
              if 'VERSION' in hdr:
                  self._version = hdr['VERSION']
              else:
                  self._version = 'unknown'
  
              self.wave = np.asarray(hdr['CRVAL1'] + \
                  hdr['CDELT1']*np.arange(hdr['NAXIS1']), dtype=np.float64)
              if 'LOGLAM' in hdr and hdr['LOGLAM'] != 0:
                  self.wave = 10**self.wave
  
              self.flux = np.asarray(native_endian(fx['BASIS_VECTORS'].data),
                  dtype=np.float64)
  
              self._redshifts = None
  
              ## find out if redshift info is present in the file
              old_style_templates = True
              try:
                  self._redshifts = native_endian(fx['REDSHIFTS'].data)
                  old_style_templates = False
              except KeyError:
                  pass
  
              fx.close()
  
              self._rrtype = hdr['RRTYPE'].strip().upper()
              if old_style_templates:
                  if self._rrtype == 'GALAXY':
                      # redshifts = 10**np.arange(np.log10(1+0.005),
                      # np.log10(1+2.0), 1.5e-4) - 1
                      self._redshifts = 10**np.arange(np.log10(1-0.005),
                          np.log10(1+1.7), 3e-4) - 1
                  elif self._rrtype == 'STAR':
                      self._redshifts = np.arange(-0.002, 0.00201, 4e-5)
                  elif self._rrtype == 'QSO':
                      self._redshifts = 10**np.arange(np.log10(1+0.05),
                          np.log10(1+6.0), 5e-4) - 1
                  else:
                      raise ValueError("Unknown redshift range to use for "
                          "template type {}".format(self._rrtype))
                  zmin = self._redshifts[0]
                  zmax = self._redshifts[-1]
                  print("DEBUG: Using default redshift range {:.4f}-{:.4f} for "
                      "{}".format(zmin, zmax, os.path.basename(filename)))
              else:
                  zmin = self._redshifts[0]
                  zmax = self._redshifts[-1]
                  print("DEBUG: Using redshift range {:.4f}-{:.4f} for "
                      "{}".format(zmin, zmax, os.path.basename(filename)))
  
              self._subtype = None
              if 'RRSUBTYP' in hdr:
                  self._subtype = hdr['RRSUBTYP'].strip().upper()
              else:
                  self._subtype = ''
  
          else:
              self._rrtype = spectype
              self._redshifts = redshifts
              self.wave = wave
              self.flux = flux
              self._subtype = subtype
  
          self._nbasis = self.flux.shape[0]
          self._nwave = self.flux.shape[1]
  
  
      @property
      def nbasis(self):
          return self._nbasis
  
      @property
      def nwave(self):
          return self._nwave
  
      @property
      def template_type(self):
          return self._rrtype
  
      @property
      def sub_type(self):
          return self._subtype
  
      @property
      def full_type(self):
          """Return formatted type:subtype string.
          """
          if self._subtype != '':
              return '{}:::{}'.format(self._rrtype, self._subtype)
          else:
              return self._rrtype
  
      @property
      def redshifts(self):
          return self._redshifts
  
  
      def eval(self, coeff, wave, z):
          """Return template for given coefficients, wavelengths, and redshift
  
          Args:
              coeff : array of coefficients length self.nbasis
              wave : wavelengths at which to evaluate template flux
              z : redshift at which to evaluate template flux
  
          Returns:
              template flux array
  
          Notes:
              A single factor of (1+z)^-1 is applied to the resampled flux
              to conserve integrated flux after redshifting.
  
          """
          assert len(coeff) == self.nbasis
          flux = self.flux.T.dot(coeff).T / (1+z)
          return trapz_rebin(self.wave*(1+z), flux, wave)
  
  
  
  
  def find_templates(template_dir=None):
      """Return list of redrock-\*.fits template files
  
      Search directories in this order, returning results from first one found:
          - template_dir
          - $RR_TEMPLATE_DIR
          - <redrock_code>/templates/
  
      Args:
          template_dir (str): optional directory containing the templates.
  
      Returns:
          list: a list of template files.
  
      """
      if template_dir is None:
          if 'RR_TEMPLATE_DIR' in os.environ:
              template_dir = os.environ['RR_TEMPLATE_DIR']
          else:
              thisdir = os.path.dirname(__file__)
              tempdir = os.path.join(os.path.abspath(thisdir), 'templates')
              if os.path.exists(tempdir):
                  template_dir = tempdir
  
      if template_dir is None:
          raise IOError("ERROR: can't find template_dir, $RR_TEMPLATE_DIR, or {rrcode}/templates/")
      else:
          print('DEBUG: Read templates from {}'.format(template_dir) )
  
      return sorted(glob(os.path.join(template_dir, 'rrtemplate-*.fits')))
  
  
  class DistTemplatePiece(object):
      """One piece of the distributed template data.
  
      This is a simple container for storing interpolated templates for a set of
      redshift values.  It is used for communicating the interpolated templates
      between processes.
  
      In the MPI case, each process will store at most two of these
      simultaneously.  This is the data that is computed on a single process and
      passed between processes.
  
      Args:
          index (int): the chunk index of this piece- this corresponds to
              the process rank that originally computed this piece.
          redshifts (array): the redshift range contained in this piece.
          data (list): a list of dictionaries, one for each redshift, and
              each containing the 2D interpolated template values for all
              "wavehash" keys.
  
      """
      def __init__(self, index, redshifts, data):
          self.index = index
          self.redshifts = redshifts
          self.data = data
  
  
  def _mp_rebin_template(template, dwave, zlist, qout):
      """Function for multiprocessing version of rebinning.
      """
      try:
          results = dict()
          for z in zlist:
              binned = rebin_template(template, z, dwave)
              results[z] = binned
          qout.put(results)
      except:
          exc_type, exc_value, exc_traceback = sys.exc_info()
          lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
          lines = [ "MP rebin: {}".format(x) for x in lines ]
          print("".join(lines))
          sys.stdout.flush()
      return
  
  
  class DistTemplate(object):
      """Distributed template data interpolated to all redshifts.
  
      For a given template, the redshifts are distributed among the
      processes in the communicator.  Then each process will rebin the
      template to those redshifts for the wavelength grids specified by
      dwave.
  
      Args:
          template (Template): the template to distribute
          dwave (dict): the keys are the "wavehash" and the values
              are a 1D array containing the wavelength grid.
          mp_procs (int): if not using MPI, restrict the number of
              multiprocesses to this.
          comm (mpi4py.MPI.Comm): (optional) the MPI communicator.
  
      """
      def __init__(self, template, dwave, mp_procs=1, comm=None):
          self._comm = comm
          self._template = template
          self._dwave = dwave
  
          self._comm_rank = 0
          self._comm_size = 1
          if self._comm is not None:
              self._comm_rank = self._comm.rank
              self._comm_size = self._comm.size
  
          self._distredshifts = np.array_split(self._template.redshifts,
              self._comm_size)
  
          myz = self._distredshifts[self._comm_rank]
          nz = len(myz)
  
          data = list()
  
          # In the case of not using MPI (comm == None), one process is rebinning
          # all the templates.  In that scenario, use multiprocessing
          # workers to do the rebinning.
  
          if self._comm is not None:
              # MPI case- compute our local redshifts
              for z in myz:
                  binned = rebin_template(self._template, z, self._dwave)
                  data.append(binned)
          else:
              # We don't have MPI, so use multiprocessing
              import multiprocessing as mp
  
              qout = mp.Queue()
              work = np.array_split(myz, mp_procs)
              procs = list()
              for i in range(mp_procs):
                  p = mp.Process(target=_mp_rebin_template,
                      args=(self._template, self._dwave, work[i], qout))
                  procs.append(p)
                  p.start()
  
              # Extract the output into a single list
              results = dict()
              for i in range(mp_procs):
                  res = qout.get()
                  results.update(res)
              for z in myz:
                  data.append(results[z])
  
          # Correct spectra for Lyman-series
          for i, z in enumerate(myz):
              for k in list(self._dwave.keys()):
                  T = transmission_Lyman(z,self._dwave[k])
                  for vect in range(data[i][k].shape[1]):
                      data[i][k][:,vect] *= T
  
          self._piece = DistTemplatePiece(self._comm_rank, myz, data)
  
  
      @property
      def comm(self):
          return self._comm
  
      @property
      def template(self):
          return self._template
  
      @property
      def local(self):
          return self._piece
  
  
      def cycle(self):
          """Pass our piece of data to the next process.
  
          If we have returned to our original data, then return True, otherwise
          return False.
  
          Args:
              Nothing
  
          Returns (bool):
              Whether we have finished (True) else False.
  
          """
          # If we are not using MPI, this function is a no-op, so just return.
          if self._comm is None:
              return True
  
          rank = self._comm_rank
          nproc = self._comm_size
  
          to_proc = rank + 1
          if to_proc >= nproc:
              to_proc = 0
  
          from_proc = rank - 1
          if from_proc < 0:
              from_proc = nproc - 1
  
          # Send our data and get a request handle for later checking.
  
          req = self._comm.isend(self._piece, to_proc)
  
          # Receive our data
  
          incoming = self._comm.recv(source=from_proc)
  
          # Wait for send to finishself._comm_rank = self._comm.rank
  
          req.wait()
  
          # Now replace our local piece with the new one
  
          self._piece = incoming
  
          # Are we done?
  
          done = False
          if self._piece.index == rank:
              done = True
  
          return done
  
  
  def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1):
      """Read and distribute templates from disk.
  
      This reads one or more template files from disk and distributes them among
      an MPI communicator.  Each process will locally store interpolated data
      for a redshift slice of each template.  For a single redshift, the template
      is interpolated to the wavelength grids specified by "dwave".
  
      As an example, imagine 3 templates with independent redshift ranges.  Also
      imagine that the communicator has 2 processes.  This function would return
      a list of 3 DistTemplate objects.  Within each of those objects, the 2
      processes store the interpolated data for a subset of the redshift range:
  
      DistTemplate #1:  zmin1 <---- p0 ----> | <---- p1 ----> zmax1
      DistTemplate #2:  zmin2 <-- p0 --> | <-- p1 --> zmax2
      DistTemplate #3:  zmin3 <--- p0 ---> | <--- p1 ---> zmax3
  
      Args:
          dwave (dict): the dictionary of wavelength grids.  Keys are the
              "wavehash" and values are an array of wavelengths.
          templates (str or None): if None, find all templates from the
              redrock template directory.  If a path to a file is specified,
              load that single template.  If a path to a directory is given,
              load all templates in that directory.
          comm (mpi4py.MPI.Comm): (optional) the MPI communicator.
          mp_procs (int): if not using MPI, restrict the number of
              multiprocesses to this.
  
      Returns:
          list: a list of DistTemplate objects.
  
      """
      timer = elapsed(None, "", comm=comm)
  
      template_files = None
  
      if (comm is None) or (comm.rank == 0):
          # Only one process needs to do this
          if templates is not None:
              if os.path.isfile(templates):
                  # we are using just a single file
                  template_files = [ templates ]
              elif os.path.isdir(templates):
                  # this is a template dir
                  template_files = find_templates(template_dir=templates)
              else:
                  print("{} is neither a file nor a directory"\
                      .format(templates))
                  sys.stdout.flush()
                  if comm is not None:
                      comm.Abort()
          else:
              template_files = find_templates()
  
      if comm is not None:
          template_files = comm.bcast(template_files, root=0)
  
      template_data = list()
      if (comm is None) or (comm.rank == 0):
          for t in template_files:
              template_data.append(Template(filename=t))
  
      if comm is not None:
          template_data = comm.bcast(template_data, root=0)
  
      timer = elapsed(timer, "Read and broadcast of {} templates"\
          .format(len(template_files)), comm=comm)
  
      # Compute the interpolated templates in a distributed way with every
      # process generating a slice of the redshift range.
  
      dtemplates = list()
      for t in template_data:
          dtemplates.append(DistTemplate(t, dwave, mp_procs=mp_procs, comm=comm))
  
      timer = elapsed(timer, "Rebinning templates", comm=comm)
  
      return dtemplates
  
apis: numpy.arange numpy.array_split numpy.log10