text:  #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  """
  Created on Fri Aug 14 19:11:57 2020
  
  @author: twguest
  """
  import sys
  import os
  import numpy as np
  
  from felpy.utils.os_utils import mkdir_p
  from felpy.utils.job_utils import JobScheduler
   
  <<<<<<< HEAD
  from wpg.wavefront import Wavefront
  from labwork.about import logs, dCache 
  
  =======
  from felpy.model.core.wavefront import Wavefront
  >>>>>>> 108cfb9b6fc97d3841ee1db54862523eee5b184e
  
  indir = dCache + "NanoKB-Pulse/source/"
  outdir = dCache + "NanoKB-Pulse/source_data/"
  
  mkdir_p(outdir)
  tmp_dir = outdir + "/tmp/"
  mkdir_p(tmp_dir)
  
  intensity_dir = tmp_dir + "/integrated_intensity/"
  mkdir_p(intensity_dir)
  
  complex_dir = tmp_dir + "/complex_wavefield/"
  mkdir_p(complex_dir)
  
  
  
  def launch():
      """
      This part launches the jobs that run in main 
      """
      
      cwd = os.getcwd()
      script = os.path.basename(__file__)
      
      
      js = JobScheduler(cwd + "/" + script, logDir = logs,
                        jobName = "extract_intensity_", partition = 'exfel', nodes = 4, jobType = 'array',
                        jobArray = os.listdir(indir))
      
      js.run(test = False)
      
      
  
  def extract_intensity(fname):
      
      wfr = Wavefront()
      wfr.load_hdf5(indir + fname)
      
     
      np.save(intensity_dir + fname, wfr.get_intensity().sum(-1))
      np.save(complex_dir + fname, wfr.as_complex_array().sum(-1))
      
  def compile_intensity_data():
      """
      note, it would be nice if this could read from squeue
      """
  
      
      np.save(outdir + "integrated_intensity" ,np.dstack(np.load(intensity_dir + ii) for ii in os.listdir(intensity_dir)))
      np.save(outdir + "complex_wavefield" ,np.dstack(np.load(intensity_dir + ii) for ii in os.listdir(intensity_dir)))
      
  
  if __name__ == '__main__':
      from tqdm import tqdm
      
      mkdir_p(outdir)
  
      try:
          extract_intensity(sys.argv[1])
      except:
          for fname in tqdm(os.listdir(indir)):
              print(fname)
              extract_intensity(fname)
  
      #compile_intensity_data()