vide_public/python_source/backend/launchers.py

#+
#   VIDE -- Void IDentification and Examination -- ./python_tools/vide/backend/launchers.py
#   Copyright (C) 2010-2014 Guilhem Lavaux
#   Copyright (C) 2011-2014 P. M. Sutter
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; version 2 of the License.
#
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License along
#   with this program; if not, write to the Free Software Foundation, Inc.,
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#+
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# routines which communicate with individual data analysis portions -
#   they make the main script easier to read

import os
import glob
from . import classes
import numpy as np
import numpy.ma as ma
import os
import shutil
import glob
import subprocess
import sys
from   netCDF4 import Dataset
from backend.classes import *
from backend.cosmologyTools import *
from backend.surveyTools import *
import pickle
import scipy.interpolate as interpolate
import time

NetCDFFile = Dataset
ncFloat = 'f8' # Double precision

class Bunch:
  def __init__(self, **kwds):
    self.__dict__.update(kwds)

#LIGHT_SPEED = 299792.458

# -----------------------------------------------------------------------------
def launchPrep(sample, binPath, workDir=None, inputDataDir=None,
               outputDir=None, figDir=None, logFile=None, useComoving=False,
               continueRun=None):

  startTime = time.time()

  if sample.dataType == "observation":
    sampleName = sample.fullName

    inputParameterFlag = ""
    outputFile = outputDir + "/zobov_slice_" + sampleName

    if sample.dataFile == "":
      datafile = inputDataDir+"/"+sampleName
    else:
      datafile = inputDataDir+"/"+sample.dataFile

    if sample.maskFile == "":
      if sample.nsideForContour == -1:
        sample.nsideForContour = 128

      sample.maskFile = outputDir + "/constructed_mask.fits"
      figureOutMask(datafile, sample.nsideForContour, sample.maskFile)

    if useComoving:
      useComovingFlag = "useComoving"
    else:
      useComovingFlag = ""

    conf="""
      catalog %s
      mask %s
      output %s
      params %s
      zMin %g
      zMax %g
      %s
      %s
      omegaM %g
      nsideForContour %g
      """ % (datafile, sample.maskFile, outputFile,
             outputDir+"/zobov_slice_"+sampleName+".par",
             sample.zBoundary[0], sample.zBoundary[1],
             useComovingFlag, inputParameterFlag, sample.omegaM,
             sample.nsideForContour)

    parmFile = os.getcwd()+"/prep_"+sample.fullName+".par"

    if not (continueRun and jobSuccessful(logFile, "Done!\n")):
      with open(parmFile, mode="wt") as f:
        f.write(conf)

      log = open(logFile, 'w')
      arg1 = "--configFile=%s" % parmFile
      #with open(logFile, 'wt') as log:
      subprocess.call([binPath, arg1], stdout=log, stderr=log)

      # compute mean particle separation
      log.write("\nFlagging edge galaxies.\n")
      (boxVol, nbar) = getSurveyProps(sample)
    
      numTracers = int(open("total_particles.txt", "r").read())
      sample.meanPartSep = (1.*numTracers/boxVol/nbar)**(-1/3.)

      # flag edge galaxies with python routine for now
      galFile = outputDir+"/zobov_slice_"+sampleName
      if os.access("contour_map.fits", os.F_OK):
        os.system("mv %s %s" % ("contour_map.fits", outputDir))
        os.system("mv %s %s" % ("galaxies.txt", outputDir))
      #galFile = outputDir + "galaxies.txt"
      edgeGalFile = outputDir + "/galaxy_edge_flags.txt"
      contourFile = outputDir + "/contour_map.fits"
      findEdgeGalaxies(galFile, sample.maskFile, edgeGalFile, contourFile,
                       sample.zBoundary[0], sample.zBoundary[1], sample.omegaM,
                       useComoving, sample.meanPartSep, 
                       outputDir, log)

      log.write("Done!\n")
      log.close()


      if jobSuccessful(logFile, "Done!\n"):
        endTime = time.time()
        walltime = endTime - startTime
        print("done (%.2fs elapsed)" % walltime)
      else:
        print("FAILED! See log file for details.")
        exit(-1)

    else:
      print("already done!")


    # compute mean particle separation
    # TODO - clean up this duplicate to cover all continueRun options
    (boxVol, nbar) = getSurveyProps(sample)
     
    numTracers = int(open("total_particles.txt", "r").read())
    sample.meanPartSep = (1.*numTracers/boxVol/nbar)**(-1/3.)

    if os.access(parmFile, os.F_OK): os.unlink(parmFile)

    if os.access("contour_map.fits", os.F_OK):
      os.system("mv %s %s" % ("contour_map.fits", outputDir))

    if os.access("comoving_distance.txt", os.F_OK):
      os.system("mv %s %s" % ("comoving_distance.txt", outputDir))

    if os.access("total_particles.txt", os.F_OK):
      os.system("mv %s %s" % ("total_particles.txt", outputDir))

    if os.access("galaxies.txt", os.F_OK):
      os.system("mv %s %s" % ("galaxies.txt", outputDir))
      #os.system("mv %s %s" % ("galaxy_edge_flags.txt", outputDir))

  else: # simulation
    sampleName = sample.fullName

    datafile = inputDataDir+"/"+sample.dataFile

    # compute mean particle separation
    iX = float(sample.mySubvolume[0])
    iY = float(sample.mySubvolume[1])
    xMin = iX/sample.numSubvolumes * sample.boxLen
    yMin = iY/sample.numSubvolumes * sample.boxLen
    xMax = (iX+1)/sample.numSubvolumes * sample.boxLen
    yMax = (iY+1)/sample.numSubvolumes * sample.boxLen
    zMin = sample.zBoundaryMpc[0]
    zMax = sample.zBoundaryMpc[1]

    boxVol = (xMax-xMin)*(yMax-yMin)*(zMax-zMin)
    sample.meanPartSep = (1.*numTracers/boxVol)**(-1/3.)

    # check if the final subsampling is done
    lastSample = sample.subsample.split(', ')[-1]
    doneLine = "Done! %5.2e\n" % float(lastSample)
    if (continueRun and jobSuccessful(logFile, doneLine)):
      print("already done!")
      return

    prevSubSample = -1
    firstSubSample = -1
    for thisSubSample in sample.subsample.split(', '):

      if prevSubSample == -1:
        inputParameterFlag = ""
        outputFile = outputDir+"/zobov_slice_" + sampleName + "_ss" + \
                     thisSubSample
        keepFraction = float(thisSubSample)
        subSampleLine = "subsample %g" % keepFraction
        resubSampleLine = ""
        firstSubSample = keepFraction
      else:
        inputParameterFlag = "inputParameter " + outputDir+"/zobov_slice_"+\
                             sampleName+"_ss"+prevSubSample+".par"
        outputFile = outputDir+"/zobov_slice_" + sampleName + "_ss" + \
                     thisSubSample
        keepFraction = float(thisSubSample)/float(prevSubSample)
        subSampleLine = "subsample %s" % firstSubSample
        resubSampleLine = "resubsample %g" % keepFraction

      includePecVelString = ""
      if sample.usePecVel: includePecVelString = "peculiarVelocities"

      useLightConeString = ""
      if sample.useLightCone: useLightConeString = "cosmo"

      if sample.dataFormat == "multidark" or sample.dataFormat == "random":
        dataFileLine = "multidark " + datafile
      elif sample.dataFormat == "gadget":
        dataFileLine = "gadget " + datafile
      elif sample.dataFormat == "gadget2":
        dataFileLine = "gadget2 " + datafile
      elif sample.dataFormat == "ahf":
        dataFileLine = "gadget " + datafile
      elif sample.dataFormat == "sdf":
        dataFileLine = "sdf " + datafile
      elif sample.dataFormat == "ramses":
        ramsesId = int(os.path.split(datafile)[1][-5:len(datafile)]) # picks out the particle file (should be the output_NNNNN, then extracts the output id "NNNNN" as an integer)
        dataFileLine = "ramses " + datafile + "/"
      else:
        raise ValueError("unknown dataFormat '%s'" % sample.dataFormat)

      iX = float(sample.mySubvolume[0])
      iY = float(sample.mySubvolume[1])

      xMin = iX/sample.numSubvolumes * sample.boxLen
      yMin = iY/sample.numSubvolumes * sample.boxLen
      xMax = (iX+1)/sample.numSubvolumes * sample.boxLen
      yMax = (iY+1)/sample.numSubvolumes * sample.boxLen

      reshiftFlag = ""
      if not sample.shiftSimZ: reshiftFlag = "preReShift"

      if sample.dataFormat == "ramses":
        ramsesIdLine = "ramsesId " + str(ramsesId)
      else:
        ramsesIdLine = ""

      conf="""
      %s
      output %s
      outputParameter %s
      %s
      %s
      gadgetUnit %g
      %s
      rangeX_min %g
      rangeX_max %g
      rangeY_min %g
      rangeY_max %g
      rangeZ_min %g
      rangeZ_max %g
      %s
      %s
      %s
      %s
      """ % (dataFileLine,
             outputFile,
             outputFile+".par",
             includePecVelString,
             useLightConeString,
             sample.dataUnit,
             ramsesIdLine,
             xMin, xMax, yMin, yMax,
             sample.zBoundaryMpc[0], sample.zBoundaryMpc[1],
             subSampleLine,resubSampleLine,inputParameterFlag,reshiftFlag)

      parmFile = os.getcwd()+"/generate_"+sample.fullName+".par"

      with open(parmFile, mode="wt") as f:
        f.write(conf)

      if (prevSubSample == -1):
        cmd = "%s --configFile=%s" % (binPath,parmFile)
        log = open(logFile, 'w')
      else:
        cmd = "%s --configFile=%s" % (binPath,parmFile)
        log = open(logFile, 'a')
      arg1 = "--configFile=%s" % parmFile
 

      subprocess.call(cmd, stdout=log, stderr=log, shell=True)
      log.close()

      # remove intermediate files
      if (prevSubSample != -1):
       os.unlink(outputDir+"/zobov_slice_"+sampleName+"_ss"+prevSubSample+".par")
       os.unlink(outputDir+"/zobov_slice_"+sampleName+"_ss"+prevSubSample)

      doneLine = "Done! %5.2e\n" % keepFraction
      if not jobSuccessful(logFile, doneLine):
        print("FAILED! See log file for details.") ### dies here for now
        exit(-1)

      prevSubSample = thisSubSample

    if jobSuccessful(logFile, doneLine): print("done")

    # place the final subsample
    os.system("mv %s %s" % (outputDir+"/zobov_slice_"+sampleName+"_ss"+\
              prevSubSample, outputDir+"/zobov_slice_"+sampleName))
    os.system("mv %s %s" % (outputDir+"/zobov_slice_"+sampleName+"_ss"+\
              prevSubSample+".par", outputDir+"/zobov_slice_"+sampleName+".par"))

    if os.access("comoving_distance.txt", os.F_OK):
      os.system("mv %s %s" % ("comoving_distance.txt", outputDir))

    if os.access(parmFile, os.F_OK):
      os.unlink(parmFile)

    if os.access("total_particles.txt", os.F_OK):
      os.system("mv %s %s" % ("total_particles.txt", outputDir))

  # save this sample's information
  with open(outputDir+"/sample_info.dat", mode='wb') as output:
    pickle.dump(sample, output, pickle.HIGHEST_PROTOCOL)

  fp = open(outputDir+"/sample_info.txt", 'w')
  fp.write("Sample name: %s\n" % sample.fullName)
  fp.write("Sample nickname: %s\n" % sample.nickName)
  fp.write("Data type: %s\n" % sample.dataType)
  fp.write("Redshift range: %f - %f\n" %(sample.zBoundary[0],sample.zBoundary[1]))
  if (sample.dataType == "simulation"):
    fp.write("Particles placed on lightcone: %g\n" % sample.useLightCone)
    fp.write("Peculiar velocities included: %g\n" % sample.usePecVel)
    if (len(sample.subsample) == 1):
      fp.write("Additional subsampling fraction: %s\n" % sample.subsample)
    else:
      fp.write("Additional subsampling fraction: %s\n" % sample.subsample[-1])
    fp.write("Simulation box length (Mpc/h): %g\n" % sample.boxLen)
    fp.write("Simulation Omega_M: %g\n" % sample.omegaM)
    fp.write("Number of simulation subvolumes: %s\n" % sample.numSubvolumes)
    fp.write("My subvolume index: %s\n" % sample.mySubvolume)
  fp.write("Estimated volume (cubic Mpc/h): %g\n" % boxVol)
  fp.write("Total number of tracers: %d\n" % numTracers)
  fp.write("Estimated mean tracer separation (Mpc/h): %g\n" % sample.meanPartSep)
  fp.write("Minimum void size actually used (Mpc/h): %g\n" % sample.minVoidRadius)
  fp.close()

# -----------------------------------------------------------------------------
def launchZobov(sample, binPath, outputDir=None, logDir=None, continueRun=None,
                 numZobovDivisions=None, numZobovThreads=None, 
                 mergingThreshold=0.2):

  startTime = time.time()

  sampleName = sample.fullName

  datafile = outputDir+"zobov_slice_"+sampleName

  logFile = logDir+"/zobov_"+sampleName+".out"

  vozScript = "./scr_"+sampleName

  if os.access(vozScript, os.F_OK):
    os.unlink(vozScript)

  maxDen = mergingThreshold
  if sample.dataType == "simulation" and numZobovDivisions == 1:
    print("  WARNING! You are using a single ZOBOV division with a simulation. Periodic boundaries will not be respected!")

  if not (continueRun and jobSuccessful(logFile, "Done!\n")):
    for fileName in glob.glob(outputDir+"/part._"+sampleName+".*"):
      os.unlink(fileName)

    if os.access(outputDir+"/adj_"+sampleName+".dat", os.F_OK):
      os.unlink(outputDir+"adj_"+sampleName+".dat")

    if os.access(outputDir+"/voidDesc_"+sampleName+".out", os.F_OK):
      os.unlink(outputDir+"/voidDesc_"+sampleName+".out")

    cmd = [binPath+"/vozinit", datafile, "0.1", "1.0", str(numZobovDivisions), \
                      "_"+sampleName, str(numZobovThreads), \
                      binPath, outputDir]
    log = open(logFile, 'w')
    subprocess.call(cmd, stdout=log, stderr=log)
    log.close()

    cmd = ["./%s" % vozScript]
    log = open(logFile, 'a')
    subprocess.call(cmd, stdout=log, stderr=log)
    log.close()

    # re-weight the volumes based on selection function
    if sample.dataType == "observation" and \
      sample.selFunFile != None:

      # load volumes
      volFile = outputDir+"/vol_"+sampleName+".dat"
      with open(volFile, mode="rb") as File:
        numPartTot = np.fromfile(File, dtype=np.int32,count=1)
        vols = np.fromfile(File, dtype=np.float32,count=numPartTot)

      # load redshifts
      partFile = outputDir+"/zobov_slice_"+sample.fullName
      with open(partFile, mode="rb") as File:
        chk = np.fromfile(File, dtype=np.int32,count=1)
        Np = np.fromfile(File, dtype=np.int32,count=1)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # skipping through the file to get to the correct location
        # x
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # y
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # z
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # RA
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # Dec
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

        # redshift
        chk = np.fromfile(File, dtype=np.int32,count=1)
        redshifts = np.fromfile(File, dtype=np.float32,count=Np)
        chk = np.fromfile(File, dtype=np.int32,count=1)

      # build selection function interpolation
      selfuncData = np.genfromtxt(sample.selFunFile)
      selfunc = interpolate.interp1d(selfuncData[:,0], selfuncData[:,1],
                                     kind='cubic', bounds_error=False,
                                     fill_value=1.0)
      # re-weight and write
      for i in range(len(vols)):
        vols[i] *= selfunc(redshifts[i])

      sample.hasWeightedVolumes = True
      volFile = outputDir+"/vol_weighted_"+sampleName+".dat"
      with open(volFile, mode='wb') as File:
        numPartTot.astype(np.int32).tofile(File)
        vols.astype(np.float32).tofile(File)

      volFileToUse = outputDir+"/vol_weighted_"+sampleName+".dat"
    else:
      volFileToUse = outputDir+"/vol_"+sampleName+".dat"

    # Zobov places guard particles at the domain walls; adjacenies to these
    # guard particles are not tracked in the final tesselation.
    # Therefore, if a galaxy marked as "interior" is missing an adjancency, that
    # means it erroneously connected to a guard particle.
    #
    # Re-weight the volumes of any edge galaxies to prevent watershed
    # from spilling outside of survey region.
    if sample.dataType == "observation":

      log = open(logFile, 'a')
      log.write("\nDouble-checking all edge flags.\n")

      # read in the edge flag information
      edgeFile = outputDir+"/galaxy_edge_flags.txt"
      edgeFlags = np.loadtxt(edgeFile, dtype=np.int32)

      # load in and count adjacenies
      log.write("  Loading adjancies...\n")
      part = []
      adjFile = outputDir+"adj_"+sampleName+".dat"
      with open(adjFile, mode="rb") as File:
        numPart = np.fromfile(File, dtype=np.int32,count=1)[0]

        # i'm tired and can't think of a better way to do this right now
        for p in range(numPart):
          part.append(Bunch(nadjs = 0, adjs = []))

        nadjPerPart = np.fromfile(File, dtype=np.int32,count=numPart)
        for p in range(numPart):
          nin = np.fromfile(File, dtype=np.int32, count=1)[0]
          for n in range(nin):
            pAdj = np.fromfile(File, dtype=np.int32, count=1)[0]
            if (p < pAdj):
              part[p].adjs.append(pAdj)
              part[pAdj].adjs.append(p)
 
      for p in range(numPart):
        nHave = len(part[p].adjs)
        nExpected = nadjPerPart[p]

        # an interior galaxy slipped through the boundary
        if edgeFlags[p] == 0 and nHave != nExpected:
          log.write("  Found an unflagged galaxy: %d. Correcting.\n" % (p))
          edgeFlags[p] = 4

      # re-write the new edge flag information
      np.savetxt(edgeFile, edgeFlags, fmt="%d")
      log.write("  Done double-checking. New flags saved.\n")

      # read in the appropriate volume file
      log.write("Re-weighting edge galaxy volumes...\n")
      with open(volFileToUse, mode="rb") as File:
        numPartTot = np.fromfile(File, dtype=np.int32,count=1)[0]
        vols = np.fromfile(File, dtype=np.float32, count=numPartTot)
      
      # set edge galaxy volumes to nearly 0 (implying very high density)
      vols[ edgeFlags>0 ] = 1.e-10

      volFile = outputDir+"/vol_weighted_"+sampleName+".dat"
      log.write("Saving new volume file to %s.\n" % (volFile))
      with open(volFile, mode='wb') as File:
        numPartTot.astype(np.int32).tofile(File)
        vols.astype(np.float32).tofile(File)

      sample.hasWeightedVolumes = True
      volFileToUse = outputDir+"/vol_weighted_"+sampleName+".dat"
      log.write("Done re-weighting.\n\n")

      log.close()
    else:
      volFileToUse = outputDir+"/vol_"+sampleName+".dat"

    cmd = [binPath+"/jozov2", \
           outputDir+"/adj_"+sampleName+".dat", \
           volFileToUse, \
           outputDir+"/voidPart_"+sampleName+".dat", \
           outputDir+"/voidZone_"+sampleName+".dat", \
           outputDir+"/voidDesc_"+sampleName+".out", \
           str(maxDen)
           ]
    log = open(logFile, 'a')
    subprocess.call(cmd, stdout=log, stderr=log)
    log.close()

    # don't need the subbox files
    for fileName in glob.glob(outputDir+"/part._"+sampleName+".*"):
      os.unlink(fileName)

    if jobSuccessful(logFile, "Done!\n"):
      endTime = time.time()
      walltime = endTime - startTime
      print("done (%.2fs elapsed)" % walltime)
    else:
      print("FAILED! See log file for details.")
      exit(-1)

  else:
    print("already done!")

  if os.access(vozScript, os.F_OK):
    os.unlink(vozScript)

# -----------------------------------------------------------------------------
def launchPrune(sample, binPath,
                summaryFile=None, logFile=None, outputDir=None,
                continueRun=None, useComoving=False, mergingThreshold=0.2):

  startTime = time.time()

  sampleName = sample.fullName

  numVoids = sum(1 for line in \
                 open(outputDir+"/voidDesc_"+sampleName+".out"))
  numVoids -= 2

  maxDen = mergingThreshold
  if sample.dataType == "observation":
    totalPart = open(outputDir+"/total_particles.txt", "r").read()
    observationLine = " --isObservation"
  else:
    observationLine = ""

  periodicLine = " --periodic='" + getPeriodic(sample) + "'"

  if sample.hasWeightedVolumes:
    volFileLine = " --partVol=" + outputDir+"/vol_weighted_"+\
                  str(sampleName)+".dat"
  else:
    volFileLine = " --partVol=" + outputDir+"/vol_"+str(sampleName)+".dat"

  if useComoving:
    useComovingFlag = " --useComoving"
  else:
    useComovingFlag = ""

  if sample.minVoidRadius == -1:
    minRadius = -1
    for line in open(outputDir+"/sample_info.txt"):
      if "Estimated mean tracer separation" in line:
        minRadius = float(line.split()[5])
        break
    if minRadius == -1:
      print("Could not grab mean tracer separation!?")
      exit(-1)
  else:
    minRadius = sample.minVoidRadius

  if not (continueRun and (jobSuccessful(logFile, "NetCDF: Not a valid ID\n") \
          or jobSuccessful(logFile, "Done!\n"))):
    cmd = binPath
    cmd += " --partFile=" + outputDir+"/zobov_slice_"+str(sampleName)
    cmd += " --voidDesc=" + outputDir+"/voidDesc_"+str(sampleName)+".out"
    cmd += " --void2Zone="+outputDir+"/voidZone_"+str(sampleName)+".dat"
    cmd += " --zone2Part=" + outputDir+"/voidPart_"+str(sampleName)+".dat"
    cmd += " --partAdj=" + outputDir+"/adj_"+str(sampleName)+".dat"
    cmd += " --partEdge=" + outputDir+"galaxy_edge_flags.txt"
    cmd += " --extraInfo=" + outputDir+"/zobov_slice_"+str(sampleName)+".par"
    cmd += " --maxCentralDen=" + str(maxDen)
    cmd += " --rMin=" + str(minRadius)
    cmd += " --numVoids=" + str(numVoids)
    cmd += observationLine
    cmd += periodicLine
    cmd += volFileLine
    cmd += useComovingFlag
    cmd += " --omegaM=" + str(sample.omegaM)
    cmd += " --outputDir=" + outputDir
    cmd += " --sampleName=" + str(sampleName)
    log = open(logFile, 'w')
    subprocess.call(cmd, stdout=log, stderr=log, shell=True)
    log.close()

    if jobSuccessful(logFile, "NetCDF: Not a valid ID\n") or \
       jobSuccessful(logFile, "Done!\n"):
      endTime = time.time()
      walltime = endTime - startTime
      print("done (%.2fs elapsed)" % walltime)
    else:
      print("FAILED! See log file for details.")
      #exit(-1)

  else:
    print("already done!")