text: 
  import sys
  import re
  from westtools import WESTTool, WESTDataReader
  import numpy, h5py, operator, time
  import westpa
  from westpa import h5io
  
  from west import Segment
  from west.states import InitialState
  from west.data_manager import (weight_dtype, n_iter_dtype, seg_id_dtype, utime_dtype, vstr_dtype, 
                                 istate_type_dtype, istate_status_dtype)
  
  class Trace:
      '''A class representing a trace of a certain trajectory segment back to its origin.'''
      
      def __init__(self, summary, endpoint_type, basis_state, initial_state, data_manager = None):
          self.summary = summary
          self.endpoint_type = endpoint_type
          self.basis_state = basis_state
          self.initial_state = initial_state
          self.data_manager = data_manager or westpa.rc.get_data_manager()
          
          # A mapping from aux file names to open h5py.File objects, to minimize time
          
          self._auxfiles = {}
              
      def __len__(self):
          try:
              return len(self.summary)
          except TypeError:
              return 0
          
      def __getitem__(self, sl):
          return self.summary[sl]
      
      def __iter__(self):
          return iter(self.summary)
              
      @classmethod
      def from_data_manager(cls, n_iter, seg_id, data_manager = None):
          '''Construct and return a trajectory trace whose last segment is identified
          by ``seg_id`` in the iteration number ``n_iter``.'''
          
          data_manager = data_manager or westpa.rc.get_data_manager()
          
          # These values are used later on
          endpoint_type = None
          pcoord_dtype = None
          pcoord_pt_shape = None
          
          seginfo = []
          parent_id = seg_id
          
          while n_iter > 0 and parent_id >= 0:
              seg_id = parent_id
              iter_group = data_manager.get_iter_group(n_iter)
              pcoord_ds = iter_group['pcoord']
              seg_index = iter_group['seg_index']
              n_segs = pcoord_ds.shape[0]
              pcoord_len = pcoord_ds.shape[1]
              
              assert seg_id < n_segs
              
              indexrow = seg_index[seg_id]
              final_pcoord = pcoord_ds[seg_id, pcoord_len-1]
              weight = indexrow['weight']
              cputime = indexrow['cputime']
              walltime = indexrow['walltime']
              
              try:
                  parent_id = int(indexrow['parent_id'])
              except IndexError:
                  # old HDF5 version
                  parent_id = int(iter_group['parents'][indexrow['parents_offset']])
                  
              if endpoint_type is None:
                  endpoint_type = indexrow['endpoint_type']
                  pcoord_pt_shape = pcoord_ds.shape[2:]
                  pcoord_dtype = pcoord_ds.dtype
                  
              seginfo.append((n_iter, seg_id, weight, walltime, cputime, final_pcoord))
              
              del iter_group, pcoord_ds, seg_index
              n_iter -= 1
              
          # loop terminates with parent_id set to the identifier of the initial state, 
          # seg_id set to the identifier of the first segment in the trajectory, and
          # n_iter set to one less than the iteration of the first segment
          first_iter = n_iter + 1
          first_seg_id = seg_id
          first_parent_id = parent_id
  
          # Initial segment (for fetching initial state)
          first_segment = Segment(n_iter=first_iter, seg_id=first_seg_id, parent_id=first_parent_id)
          
          seginfo.reverse()
          
          summary_dtype = numpy.dtype([('n_iter', n_iter_dtype),
                                       ('seg_id', seg_id_dtype),
                                       ('weight', weight_dtype),
                                       ('walltime', utime_dtype),
                                       ('cputime', utime_dtype),
                                       ('final_pcoord', pcoord_dtype, pcoord_pt_shape),
                                       ])
          
          summary = numpy.array(seginfo, dtype=summary_dtype)
          
          try:
              initial_state = data_manager.get_segment_initial_states([first_segment], first_iter)[0]
          except KeyError:
              # old HDF5 version
              assert parent_id < 0
              istate_pcoord = data_manager.get_iter_group(first_iter)['pcoord'][first_seg_id,0]
              istate_id = -(first_parent_id+1)
              basis_state = None
              initial_state = InitialState(istate_id, None, iter_created=0, pcoord=istate_pcoord)
              
          else:
              basis_state = data_manager.get_basis_states(first_iter)[initial_state.basis_state_id]
              
          return cls(summary, endpoint_type, basis_state, initial_state, data_manager)
              
      def get_segment_data_slice(self, datafile, dsname, n_iter, seg_id, slice_=None, index_data=None,
                                 iter_prec=None):
          '''Return the data from the dataset named ``dsname`` within the given ``datafile`` (an open
          h5py.File object) for the given iteration and segment. By default, it is assumed that the
          dataset is stored in the iteration group for iteration ``n_iter``, but if ``index_data``
          is provided, it must be an iterable (preferably a simple array) of (n_iter,seg_id) pairs,
          and the index in the ``index_data`` iterable of the matching n_iter/seg_id pair is used as
          the index of the data to retrieve.
          
          If an optional ``slice_`` is provided, then the given slicing tuple is appended to that
          used to retrieve the segment-specific data (i.e. it can be used to pluck a subset of the
          data that would otherwise be returned).
          '''
  
          if slice_ is None:
              slice_ = numpy.s_[...] 
                      
          if index_data is not None:
              dataset = datafile[dsname]
  
              for i, (i_n_iter,i_seg_id) in enumerate(index_data):
                  if (i_n_iter,i_seg_id) == (n_iter,seg_id):
                      break
              else:
                  raise KeyError((n_iter,seg_id))
              
              itpl = (i,) + slice_
              return dataset[itpl]
          else:
              if not iter_prec:
                  iter_prec = datafile.attrs.get('west_iter_prec', self.data_manager.default_iter_prec)
              igname_tail = 'iter_{:0{iter_prec:d}d}'.format(int(n_iter),iter_prec=int(iter_prec))
              try:
                  iter_group = datafile['/iterations/' + igname_tail]
              except KeyError:
                  iter_group = datafile[igname_tail]
              
              dataset = iter_group[dsname]
              itpl = (seg_id,) + slice_
  
              return dataset[itpl]
  
      def trace_timepoint_dataset(self, dsname, slice_=None, auxfile=None,index_ds=None):
          '''Return a trace along this trajectory over a dataset which is layed out as [seg_id][timepoint][...].
          Overlapping values at segment boundaries are accounted for.  Returns (data_trace, weight), where 
          data_trace is a time series of the dataset along this trajectory, and weight is the corresponding
          trajectory weight at each time point.
          
          If ``auxfile`` is given, then load the dataset from the given HDF5 file, which must be 
          layed out the same way as the main HDF5 file (e.g. iterations arranged as
          iterations/iter_*).
          
          If index_ds is given, instead of reading data per-iteration from iter_* groups, then the
          given index_ds is used as an index of n_iter,seg_id pairs into ``dsname``. In this case,
          the target data set need not exist on a per-iteration basis inside iter_* groups.
          
          If ``slice_`` is given, then *further* slice the data returned from the HDF5 dataset. This can
          minimize I/O if it is known (and specified) that only a subset of the data along the
          trajectory is needed.
          '''
          
          # Figure out where to look for the dataset
          if isinstance(auxfile, str):
              datafile = h5py.File(auxfile, 'r')
              close_datafile = True
          elif auxfile is not None:
              datafile = auxfile
              close_datafile = False
          else:
              datafile = self.data_manager.we_h5file
              close_datafile = False
              
          iter_prec = self.data_manager.iter_prec
          get_data_slice = self.get_segment_data_slice
              
          # Load the index if we use it
          if index_ds is not None:
              if isinstance(index_ds,str):
                  index_ds = datafile[index_ds]
              index_data = index_ds[...]
          else:
              index_data = None
              
          # Be sure to retrieve the time series
          if not slice_:
              first_sl = numpy.index_exp[:, ...]
              other_sl = numpy.index_exp[1:,...]
          else:
              first_sl = numpy.index_exp[:] + slice_
              other_sl = numpy.index_exp[1:] + slice_ 
          
          # Retrieve the first segment's data
          first_n_iter, first_seg_id = self.summary[0]['n_iter'], self.summary[0]['seg_id']
          first_iter_data = get_data_slice(datafile, dsname, first_n_iter, first_seg_id, first_sl, index_data, iter_prec)
  
          n_segs = len(self)
          n_points_per_seg = len(first_iter_data)
          
          length = n_points_per_seg + (n_segs-1)*(n_points_per_seg-1)
          tracedata = numpy.empty((length,) + first_iter_data.shape[1:], dtype=first_iter_data.dtype)
          traceweight = numpy.empty((length,), weight_dtype)
          
          # Store first segment data
          tracedata[0:n_points_per_seg] = first_iter_data
          traceweight[0:n_points_per_seg] = self.summary[0]['weight']
          del first_iter_data
          
          # Store remainder of data
          
          for iseg, summary_item in enumerate(self.summary[1:]):
              n_iter = summary_item['n_iter']
              seg_id = summary_item['seg_id']
              weight = summary_item['weight']
              
              offset = n_points_per_seg + iseg*(n_points_per_seg-1)
              length = n_points_per_seg - 1
              seg_data = get_data_slice(datafile, dsname, n_iter, seg_id, other_sl, index_data, iter_prec)
              
              tracedata[offset:offset+length] = seg_data 
              traceweight[offset:offset+length] = weight
              del seg_data
          
          if close_datafile:
              datafile.close()
          
          return tracedata, traceweight
  
      """
      # This is disabled until there is a real use for it; the following code is 
      # outdated
      def trace_perseg_dataset(self, dsname):
          '''Return a trace along this trajectory over a dataset which is layed out as [seg_id][...].
          Returns (data_trace, weight), where  data_trace is a time series of the dataset along this
          trajectory, and weight is the corresponding trajectory weight at each time point.'''
  
          first_n_iter, first_seg_id = self.summary[0]['n_iter'], self.summary[0]['seg_id']
          first_iter_group = self.data_manager.get_iter_group(first_n_iter)
          first_iter_ds = first_iter_group[dsname]
          n_segs = len(self)
          tracedata = numpy.empty((n_segs,) + first_iter_ds.shape[1:], dtype=first_iter_ds.dtype)
          traceweight = numpy.empty((n_segs,), weight_dtype)
          tracedata[0] = first_iter_ds[first_seg_id]
          traceweight[0] = self.summary[0]['weight']
          for isegm1, summary_item in enumerate(self.summary[1:]):
              iseg = isegm1 + 1
              n_iter = summary_item['n_iter']
              seg_id = summary_item['seg_id']            
              iter_group = self.data_manager.get_iter_group(n_iter)
              seg_data = iter_group[dsname][seg_id]
              tracedata[iseg] = seg_data
              traceweight[iseg] = summary_item['weight']
              del seg_data
              
          return tracedata, traceweight
      """
          
  class WTraceTool(WESTTool):
      prog='w_trace'
      description = '''\
  Trace individual WEST trajectories and emit (or calculate) quantities along the
  trajectory.
  
  Trajectories are specified as N_ITER:SEG_ID pairs. Each segment is traced back
  to its initial point, and then various quantities (notably n_iter and seg_id)
  are printed in order from initial point up until the given segment in the given
  iteration.
  
  Output is stored in several files, all named according to the pattern given by
  the -o/--output-pattern parameter. The default output pattern is "traj_%d_%d",
  where the printf-style format codes are replaced by the iteration number and
  segment ID of the terminal segment of the trajectory being traced.
  
  Individual datasets can be selected for writing using the -d/--dataset option
  (which may be specified more than once). The simplest form is ``-d dsname``,
  which causes data from dataset ``dsname`` along the trace to be stored to
  HDF5.  The dataset is assumed to be stored on a per-iteration basis, with
  the first dimension corresponding to seg_id and the second dimension
  corresponding to time within the segment.  Further options are specified
  as comma-separated key=value pairs after the data set name, as in
  
      -d dsname,alias=newname,index=idsname,file=otherfile.h5,slice=[100,...]
      
  The following options for datasets are supported:
  
      alias=newname
          When writing this data to HDF5 or text files, use ``newname``
          instead of ``dsname`` to identify the dataset. This is mostly of
          use in conjunction with the ``slice`` option in order, e.g., to
          retrieve two different slices of a dataset and store then with
          different names for future use.
  
      index=idsname
          The dataset is not stored on a per-iteration basis for all
          segments, but instead is stored as a single dataset whose
          first dimension indexes n_iter/seg_id pairs. The index to
          these n_iter/seg_id pairs is ``idsname``.
      
      file=otherfile.h5
          Instead of reading data from the main WEST HDF5 file (usually
          ``west.h5``), read data from ``otherfile.h5``.
          
      slice=[100,...]
          Retrieve only the given slice from the dataset. This can be
          used to pick a subset of interest to minimize I/O.
          
  -------------------------------------------------------------------------------
  '''
  
  
      pcoord_formats = {'u8': '%20d',
                        'i8': '%20d',
                        'u4': '%10d',
                        'i4': '%11d',
                        'u2': '%5d',
                        'i2': '%6d',
                        'f4': '%14.7g',
                        'f8': '%023.15g'}
      
      def __init__(self):
          super(WTraceTool,self).__init__()
          
          self.data_reader = WESTDataReader()
          #self.h5storage = HDF5Storage()
          self.output_file = None
          self.output_pattern = None
          self.endpoints = None
          self.datasets = []
  
          
      # Interface for command-line tools
      def add_args(self, parser):
          self.data_reader.add_args(parser)
          #self.h5storage.add_args(parser)
          parser.add_argument('-d', '--dataset', dest='datasets',
                              #this breaks argparse (see http://bugs.python.org/issue11874) 
                              #metavar='DSNAME[,alias=ALIAS][,index=INDEX][,file=FILE][,slice=SLICE]',
                              metavar='DSNAME',
                              action='append',
                              help='''Include the dataset named DSNAME in trace output. An extended form like
                              DSNAME[,alias=ALIAS][,index=INDEX][,file=FILE][,slice=SLICE] will
                              obtain the dataset from the given FILE instead of the main WEST HDF5 file,
                              slice it by SLICE, call it ALIAS in output, and/or access per-segment data by a n_iter,seg_id
                              INDEX instead of a seg_id indexed dataset in the group for n_iter.''')
          parser.add_argument('endpoints',  metavar='N_ITER:SEG_ID', nargs='+',
                              help='''Trace trajectory ending (or at least alive at) N_ITER:SEG_ID.''')
          
          #tgroup = parser.add_argument_group('trace options')
          ogroup = parser.add_argument_group('output options')
          ogroup.add_argument('--output-pattern', default='traj_%d_%d',
                              help='''Write per-trajectory data to output files/HDF5 groups whose names begin with OUTPUT_PATTERN,
                                   which must contain two printf-style format flags which will be replaced with the iteration number
                                   and segment ID of the terminal segment of the trajectory being traced.
                                   (Default: %(default)s.)''')
          ogroup.add_argument('-o', '--output', default='trajs.h5',
                              help='Store intermediate data and analysis results to OUTPUT (default: %(default)s).')
          
      
      def process_args(self, args):
          self.data_reader.process_args(args)
          #self.h5storage.process_args(args)
          self.endpoints = [list(map(int,endpoint.split(':'))) for endpoint in args.endpoints]
          self.output_pattern = args.output_pattern
          
          for dsstr in args.datasets or []:
              self.datasets.append(self.parse_dataset_string(dsstr))        
          
          #self.h5storage.open_analysis_h5file()
          self.output_file = h5py.File(args.output)
          
      def parse_dataset_string(self, dsstr):
          dsinfo = {}
  
          r = re.compile(r',(?=[^\]]*(?:\[|$))')
          fields = r.split(dsstr)
  
          dsinfo['dsname'] = fields[0]
  
          for field in (field.strip() for field in fields[1:]):
              k,v = field.split('=')
              k = k.lower()
              if k in ('alias', 'file', 'index'):
                  dsinfo[k] = v
              elif k == 'slice':
                  try:
                      dsinfo['slice'] = eval('numpy.index_exp' + v)
                  except SyntaxError:
                      raise SyntaxError('invalid index expression {!r}'.format(v))
              else:
                  raise ValueError('invalid dataset option {!r}'.format(k))
              
          return dsinfo
      
      def go(self):
          self.data_reader.open('r')
          
          #Create a new 'trajectories' group if this is the first trace
          try:
              trajs_group = h5io.create_hdf5_group(self.output_file, 'trajectories', replace=False, creating_program=self.prog)
          except ValueError:
              trajs_group = self.output_file['trajectories']
          
          for n_iter, seg_id in self.endpoints:
              trajname = self.output_pattern % (n_iter,seg_id)
              trajgroup = trajs_group.create_group(trajname)
  
              trace = Trace.from_data_manager(n_iter,seg_id, self.data_reader.data_manager)
              
              with open(trajname + '_trace.txt', 'wt') as trace_output:
                  self.emit_trace_text(trace, trace_output)
                  
              self.emit_trace_h5(trace, trajgroup)
              
              aux_h5files = {}
              for dsinfo in self.datasets:
                  dsname = dsinfo['dsname']
                  filename = dsinfo.get('file')
                  if filename:
                      try:
                          aux_h5file = aux_h5files[filename]
                      except KeyError:
                          aux_h5file = aux_h5files[filename] = h5py.File(filename, 'r')
                  else:
                      aux_h5file = None
                      
                  slice_ = dsinfo.get('slice')
                  alias = dsinfo.get('alias', dsname)
                  index = dsinfo.get('index')
                  
                  data, weights = trace.trace_timepoint_dataset(dsname, auxfile=aux_h5file, slice_=slice_,index_ds=index)
                  
                  # Save data to HDF5
                  try:
                      del trajgroup[alias]
                  except KeyError:
                      pass
                  trajgroup[alias] = data
                  
                  # All weight vectors will be the same length, so only store in HDF5 once
                  if not ('weights' in trajgroup and trajgroup['weights'].shape == weights.shape):
                      try:
                          del trajgroup['weights']
                      except KeyError:
                          pass    
                      trajgroup['weights'] = weights
                              
      def emit_trace_h5(self, trace, output_group):
          for dsname in ('basis_state', 'initial_state', 'segments'):
              try:
                  del output_group[dsname]
              except KeyError:
                  pass
          
          if trace.basis_state:
              output_group['basis_state'] = trace.basis_state.as_numpy_record()
          output_group['initial_state'] = trace.initial_state.as_numpy_record()
          output_group['segments'] = trace.summary
              
      def emit_trace_text(self, trace, output_file):
          '''Dump summary information about each segment in the given trace to the given output_file,
          which must be opened for writing in text mode.  Output columns are separated by at least
          one space.'''
          
          if not trace:
              return
                  
          pcoord_ndim = trace[0]['final_pcoord'].shape[0]
          lastseg = trace[-1]
          len_n_iter = max(6, len(str(lastseg['n_iter'])))
          len_seg_id = max(6, max(len(str(seg_id)) for seg_id in trace['seg_id']))
          seg_pattern = '    '.join(['{n_iter:{len_n_iter}d}',
                                     '{seg_id:{len_seg_id}d}',
                                     '{weight:22.17e}',
                                     '{walltime:10.6g}',
                                     '{cputime:10.6g}',
                                     '{pcoord_str:s}'
                                     ]) + '\n'
                                    
          
          output_file.write('''\
  # Trace of trajectory ending in n_iter:seg_id {n_iter:d}:{seg_id:d} (endpoint type {endpoint_type_text:s})   
  # column  0: iteration (0 => initial state)
  # column  1: seg_id (or initial state ID)
  # column  2: weight
  # column  3: wallclock time (s)
  # column  4: CPU time (s)
  '''.format(n_iter = int(lastseg['n_iter']), 
             seg_id = int(lastseg['seg_id']), 
             endpoint_type_text = Segment.endpoint_type_names[trace.endpoint_type]))
          
          
          if pcoord_ndim == 1:
              output_file.write('''\
  # column  5: final progress coordinate value            
  ''')
          else:
              fpcbegin = 5
              fpcend = fpcbegin + pcoord_ndim - 1
              output_file.write('''\
  # columns {fpcbegin:d} -- {fpcend:d}: final progress coordinate value
  '''.format(fpcbegin=fpcbegin,fpcend=fpcend))
          
          
          pcoord_formats = self.pcoord_formats
          
          # Output row for initial state
          initial_state = trace.initial_state
          pcoord_str = '    '.join(pcoord_formats.get(pcfield.dtype.str[1:], '%s') % pcfield 
                                   for pcfield in initial_state.pcoord)
          output_file.write(seg_pattern.format(n_iter=0, seg_id=initial_state.state_id,
                                               weight=0.0, walltime=0, cputime=0, pcoord_str=pcoord_str,
                                               len_n_iter=len_n_iter,len_seg_id=len_seg_id))
          
          # Output rows for segments
          for segment in trace:            
              pcoord_str = '    '.join(pcoord_formats.get(pcfield.dtype.str[1:], '%s') % pcfield 
                                       for pcfield in segment['final_pcoord'])
              output_file.write(seg_pattern.format(n_iter = int(segment['n_iter']), 
                                                   seg_id = int(segment['seg_id']),
                                                   weight = float(segment['weight']),
                                                   walltime = float(segment['walltime']),
                                                   cputime = float(segment['cputime']), 
                                                   pcoord_str=pcoord_str,
                                                   len_n_iter=len_n_iter,
                                                   len_seg_id=len_seg_id))
          
      
          
  
  if __name__ == '__main__':
      WTraceTool().main()
      
  