nautilus

#!/usr/bin/env python
# This program automates the running of nemo, which otherwise requires tedious
# patch-dependent configuration
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("maps_and_divs", nargs="+", help="map map map ... div div div if not transpose, else interleaved")
parser.add_argument("-T", "--transpose", action="store_true")
parser.add_argument("-d", "--dry-run",   action="store_true")
parser.add_argument("-s", "--serial",    action="store_true")
args = parser.parse_args()
import numpy as np, os, re, sys, shutil, subprocess
from enlib import utils, enmap, mpi
from scipy import optimize

comm = mpi.COMM_WORLD

afreqs = {
	"pa1":["f150"],"pa2":["f150"],"pa3":["f090","f150"],
	"pa4":["f150","f220"],"pa5":["f090","f150"],"pa6":["f090","f150"],
}
sarrs = {
	"s13":["pa1"],"s14":["pa1","pa2"],"s15":["pa1","pa2","pa3"],
	"s16":["pa2","pa3","pa4"],"s17":["pa4","pa5","pa6"],"s18":["pa4","pa5","pa6"],
}
ffreqs = {"f090": 97.0, "f150": 148.0, "f220": 220.}
# This should probably not be hardcoded here, but inferred from the mapdata configA
def get_beam(season, array, freq):
	#beam_fmt = "/home/snaess/project/actpol/depot/sigurdkn/beam/marius/mr3_{season}_{array}_{freq}_nohwp_night_beam_profile_instant_2018022.txt"
	#beam_fmt = "/project/r/rbond/sigurdkn/actpol/mapdata/depot/sigurdkn/beam/marius/mr3_{season}_{array}_{freq}_nohwp_night_beam_profile_instant_2018022.txt"
	#beam_fmt = "/home/r/rbond/sigurdkn/project/actpol/mapdata/depot/mlungu/beams/180805/mr3_{season}_{array}_{freq}_nohwp_night_beam_profile_jitter_180805.txt"
	if int(array[-1:]) < 4:
		return "/home/snaess/project/actpol/depot/sigurdkn/beam/180805/mr3_{season}_{array}_{freq}_nohwp_night_beam_profile_jitter_180805.txt".format(season=season, array=array, freq=freq)
	else:
		return "/home/snaess/project/actpol/depot/sigurdkn/beam/171227/beam_profile_171227_{array}_{freq}_{season}_instant.txt".format(season=season, array=array, freq=freq)

aseas = {}
for season in sarrs:
	for array in sarrs[season]:
		if array not in aseas: aseas[array] = []
		aseas[array].append(season)

nmap = len(args.maps_and_divs)//2
if not args.transpose:
	imaps = args.maps_and_divs[:nmap]
	idivs = args.maps_and_divs[nmap:]
else:
	imaps = args.maps_and_divs[0::2]
	idivs = args.maps_and_divs[1::2]

for ind in range(comm.rank, len(imaps), comm.size):
	imapfile, idivfile = imaps[ind], idivs[ind]
	print imapfile
	# Parse the file name looking for array and season information. We need all this
	# just to infer which beam to use!
	season, array, freq = None, None, None
	toks = ".".join(os.path.basename(imapfile).split(".")[:-1]).split("_")
	for tok in toks:
		m = re.match(r"(pa\d)", tok)
		if m and not array:  array  = m.group(1)
		m = re.match(r"(s\d\d)", tok)
		if m and not season: season = m.group(1)
		m = re.match(r"\+?(f\d\d\d)", tok)
		if m and not freq:   freq   = m.group(1)
	if array is None:
		print "Could not infer array from '%s'. Skipping" % imapfile
		continue
	if season is None:
		if array in aseas: season = aseas[array][0]
		else:
			print "Could not infer season from '%s'. Skipping" % imapfile
			continue
	if freq is None:
		if array in afreqs: freq = afreqs[array][0]
		else:
			print "Could not infer freq from '%s'. Skipping" % imapfile
			continue
	beam = get_beam(season, array, freq)

	def find_largest_rectangle(mask, p0):
		def calc_chisq(x):
			box = utils.nint(x.reshape(2,2))
			if np.any(box[1] < box[0]):      return np.inf
			if np.any(box[0] < 0):           return np.inf
			if np.any(box[1] >= mask.shape): return np.inf
			h, w   = box[1]-box[0]
			area   = h*w
			bad    = np.sum(~mask[box[0,0]:box[1,0],box[0,1]:box[1,1]])
			chisq = -area + 1e6*bad
			#print "A %4d %4d %4d %4d %15.7e" % (x[0],x[1],x[2],x[3],chisq)
			return chisq
		p0 = np.array(p0)
		x0 = np.array([p0,p0+1],float).reshape(-1)
		# Powell sometimes gets stuck here. This seems like a case where simplex is
		# better.
		x  = optimize.fmin(calc_chisq, x0, disp=False)
		return utils.nint(x.reshape(2,2))

	# Automatically find the reference region by finding the biggest
	# rectangle that has no holes. We base this on a low-resolution
	# version of the map.
	shape, wcs = enmap.read_map_geometry(idivfile)
	down= 1
	div = enmap.read_fits(idivfile, sel=(Ellipsis, slice(None,None,down), slice(None, None, down))).preflat[0]
	area= div.area()

	if area < np.pi:
		# Not the huge advanced act patch, which must be handled separately
		p0  = np.array(div.shape[-2:])//2
		ref = 0
		for i in range(3):
			ref = np.median(div[div>ref/8])
		mask = div>ref/5
		pixbox = find_largest_rectangle(mask, p0)
		box    = div.pix2sky(pixbox.T).T/utils.degree
		# Shrink box to whole degrees to make nemo happy
		box[0] = np.ceil(box[0])
		box[1] = np.floor(box[1])
		dec1, ra1, dec2, ra2 = box.reshape(-1)

		# If the box is too tall, switch to dec slice mode
		if dec2-dec1 > 30:
			dec1 = "'numDecSteps'"
			dec2 = 8
		
		print ra1, ra2, dec1, dec2

		# Estimate memory usage. Assume double precision and 16 copies
		mem = shape[-2]*shape[-1]*2*24*8/1024.**3
		print "estimated memory: %.1fG" % mem

		# Generate the nemo parameter string
		nemo_params = r"""# Nautilus auto-generated parameters for nemo
unfilteredMaps: [{
  mapFileName: "%(map)s",
  weightsFileName: "%(div)s",
  obsFreqGHz: %(freq)s, units: 'uK',
  beamFileName: "%(beam)s",
  addNoise: null,
  pointSourceRemoval: null,
  pointSourceMask: null,
  surveyMask: null
}]

# Detection options
thresholdSigma: 3.5
minObjPix: 1
rejectBorder: 0 #would be nice to calibrate this.
objIdent: 'nemo-'
longNames: False
catalogCuts: ['SNR > 3.5']
findCenterOfMass: True
useInterpolator: True

useMPI: False

# Photometry options
photometryOptions: {photFilter: 'Beam'}

# tileDeck options - cut-up each map into smaller sections, store in a multi-extension .fits file
makeTileDeck: False
tileOverlapDeg: 1.0

mapFilters: [{
    label: "Beam", class: "BeamRealSpaceMatchedFilter",
    params: {noiseParams: {method: "max(dataMap,CMB)",
      matchedFilterClass: 'BeamMatchedFilter',
      RADecSection:  [%(ra1)s, %(ra2)s, %(dec1)s, %(dec2)s],
      kernelMaxArcmin: 7.,
      symmetrize: False,
      noiseGridArcmin: 20.,
      saveHighPassMap: False,
      saveRMSMap: False},
      bckSub: True,
      outputUnits: 'uK',
    },
  }]

# Set this to True to generate a sky sim (with noise), run all the filters over it, and measure contamination
# Set numSkySims to number required - we need to average over many as results vary a fair bit
estimateContaminationFromSkySim: False
numSkySims: 1

# Set this to True to estimate contamination by running cluster finder over inverted maps
# This is sensitive to how well point source masking is done
estimateContaminationFromInvertedMaps: False
""" % {
		"map":os.path.abspath(imapfile), "div":os.path.abspath(idivfile), "freq":ffreqs[freq],
		"dec1":dec1, "ra1":ra1, "dec2":dec2, "ra2":ra2, "beam":beam}
		ntask = 1

	else:
		mem = 100.0
		nemo_params = r"""
# Nautilus auto-generated parameters for nemo
unfilteredMaps: [{
  mapFileName: "%(map)s",
  weightsFileName: "%(div)s",
  obsFreqGHz: %(freq)s, units: 'uK',
  beamFileName: "%(beam)s",
  addNoise: null,
  pointSourceRemoval: null,
  pointSourceMask: null,
  surveyMask: null,
}]

# MPI?
useMPI: True

# tileDeck options - cut-up each map into smaller sections, store in a multi-extension .fits file
makeTileDeck: True
tileOverlapDeg: 1.0

# User-defined tiles
# These will automatically be expanded by tileOverlapDeg, i.e., don't need to handle overlaps here
tileDefinitions: [
  {extName: '0_0',  RADecSection: [122.7, 103.8, -61., -51.]},
  {extName: '0_1',  RADecSection: [103.8, 66.6, -61., -51.]},
  {extName: '0_2',  RADecSection: [66.6, 14.3, -61., -51.]},
  {extName: '0_3',  RADecSection: [14.3, 321., -61., -51.]},
  {extName: '0_4',  RADecSection: [321., 258., -61., -51.]},
  {extName: '1_0',  RADecSection: [117., 67.3, -51., -41.]},
  {extName: '1_1',  RADecSection: [67.3, 13.3, -51., -41.]},
  {extName: '1_2',  RADecSection: [13.3, 334.2, -51., -41.]},
  {extName: '1_3',  RADecSection: [334.2, 300., -51., -41.]},
  {extName: '1_4',  RADecSection: [300., 264., -51., -41.]},
  {extName: '2_0',  RADecSection: [110., 68.2, -41., -31.]},
  {extName: '2_1',  RADecSection: [68.2, 15.5, -41., -31.]},
  {extName: '2_2',  RADecSection: [15.5, 334.2, -41., -31.]},
  {extName: '2_3',  RADecSection: [334.2, 300., -41., -31.]},
  {extName: '2_4',  RADecSection: [300., 266., -41., -31.]},
  {extName: '3_0',  RADecSection: [110., 45., -31., -21.]},
  {extName: '3_1',  RADecSection: [45., 345., -31., -21.]},
  {extName: '3_2',  RADecSection: [345., 276., -31., -21.]},
  {extName: '4_0',  RADecSection: [102., 45., -21., -11.]},
  {extName: '4_1',  RADecSection: [45., 345., -21., -11.]},
  {extName: '4_2',  RADecSection: [345., 281., -21., -11.]},
  {extName: '5_0',  RADecSection: [86., 45., -11., 5.5]},
  {extName: '5_1',  RADecSection: [45., 345., -11., 5.5]},
  {extName: '5_2',  RADecSection: [345., 291., -11., 5.5]},
  {extName: '5A_0', RADecSection: [179.9, 105., -6., 5.5]},
  {extName: '5A_1', RADecSection: [262., 217., -6., 5.5]},
  {extName: '5A_2', RADecSection: [217., 180.1, -6., 5.5]},
  {extName: '6_0',  RADecSection: [179.9, 100., 5.5, 22.]},
  {extName: '6_1',  RADecSection: [81., 6., 5.5, 22.]},
  {extName: '6_2',  RADecSection: [6., 298., 5.5, 22.]},
  {extName: '6_3',  RADecSection: [269., 217., 5.5, 22.]},
  {extName: '6_4',  RADecSection: [217., 180.1, 5.5, 22.]},
]
# Corresponding regions in tiles to use for noise part of matched filter
# IF these are modified, tileDeck files will need to be re-made (delete them and rerun nemo)
# Format for each entry: extName: [RAMin, RAMax, decMin, decMax]
tileNoiseRegions: {
  '0_0': [111., 103., -59., -52.],
  '0_1': [97., 73., -59., -52.],
  '0_2': [61., 45., -59., -52.],
  '0_3': [5., 350., -59., -52.],
  '0_4': [318., 300., -59., -52.],
  '1_0': [107., 90., -49., -42.],
  '1_1': [34., 16., -49., -42.],
  '1_2': [10., 352., -49., -42.],
  '1_3': [332., 315., -49., -42.],
  '1_4': [298., 280., -49., -42.],
  '2_0': [104., 88., -40., -32.],
  '2_1': [34., 16., -40., -32.],
  '2_2': [10., 352., -41., -33.],
  '2_3': [332., 315., -40., -32.],
  '2_4': [298., 280., -40., -32.],
  '3_0': [85., 60., -30., -22.],
  '3_1': [34., 16., -30., -22.],
  '3_2': [320., 305., -30., -22.],
  '4_0': [84., 55., -20., -12.],
  '4_1': [30., 10., -20., -12.],
  '4_2': [315., 295., -20., -12.],
  '5_0': [77., 57., -5., 4.],
  '5_1': [32.3, 38.2, -7.5, -2.5],
  '5_2': [320., 300., -7., 2.],
  '5A_0': [170., 140., -4., 4.],
  '5A_1': [234., 218., -4., 2.],
  '5A_2': [216., 200., -4., 2.],
  '6_0': [150., 130., 10., 18.],
  '6_1': [30., 10., 10., 18.],
  '6_2': [357., 340., 7., 16.],
  '6_3': [236., 218., 7., 16.],
  '6_4': [215., 195., 7., 16.],
}

# Detection options
thresholdSigma: 3.5
minObjPix: 1
rejectBorder: 0 #would be nice to calibrate this.
objIdent: 'nemo-'
longNames: False
catalogCuts: ['SNR > 3.5']
findCenterOfMass: True
useInterpolator: True

# Photometry options
photometryOptions: {photFilter: 'Beam'}

mapFilters: [{
  label: "Beam", class: "BeamRealSpaceMatchedFilter",
  params: {noiseParams: {method: "max(dataMap,CMB)",
      matchedFilterClass: 'BeamMatchedFilter',
      RADecSection:  'tileNoiseRegions',
      kernelMaxArcmin: 7.,
      symmetrize: False,
      noiseGridArcmin: 20.,
      saveHighPassMap: False,
      saveRMSMap: False},
    bckSub: True,
    outputUnits: 'uK'},
  }]

# Set this to True to generate a sky sim (with noise), run all the filters over it, and measure contamination
# Set numSkySims to number required - we need to average over many as results vary a fair bit
estimateContaminationFromSkySim: False
numSkySims: 1

# Set this to True to estimate contamination by running cluster finder over inverted maps
# This is sensitive to how well point source masking is done
estimateContaminationFromInvertedMaps: False
""" % {
		"map":os.path.abspath(imapfile), "div":os.path.abspath(idivfile), "freq":ffreqs[freq], "beam":beam}
		ntask = 10

	# Set up nemo work directory
	workdir = os.getcwd() + "/" + imapfile + ".work"
	#shutil.rmtree(workdir, ignore_errors=True)
	utils.mkdir(workdir)
	with open(workdir + "/nemo.yml", "w") as f:
		f.write(nemo_params)
	ofile = ".".join(imapfile.split(".")[:-1]) + "_catalog.txt"
	runfile = workdir + "/batch.txt"

	# Build a batch script. Nemo is not parallel, so this will just ask for
	# one core per node
	batch = r"""#!/bin/bash
#SBATCH --nodes 1 --ntasks-per-node=%(ntask)s --cpus-per-task=1 --time=4:00:00
#SBATCH --job-name %(name)s
cd "%(wdir)s"
OMP_NUM_THREADS=1 mpirun -n %(ntask)s nemo nemo.yml
cp nemo/nemo_optimalCatalog.csv "%(ofile)s"
""" % {"name": "_".join(["nemo"]+toks), "wdir": workdir, "ofile": os.path.abspath(ofile),
		"mem": mem*1024, "ntask":ntask}

	with open(runfile, "w") as f:
		f.write(batch)
	
	if not args.dry_run:
		if not args.serial:
			subprocess.call(["sbatch",runfile])
		else:
			subprocess.call(["bash",runfile])