Source code for hepi.nllfast.run

import os
import pkgutil
import re
import shutil
import stat
from string import Template
import subprocess
from typing import List
import warnings
from hepi.input import Input, is_gluino, is_slepton, is_squark, is_weakino
from hepi.results import Result
from hepi.run import RunParam, Runner
from hepi.util import namehash
from smpl import debug
import pyslha
from uncertainties import ufloat



class NLLfastRunner(Runner):
	#TODO treat stop sbot separately
	#TODO separate nll and nnll

	def _get_nf_proc(self,p:Input):
		d = pyslha.read(self.get_output_dir() + p.slha)
		mg = d.blocks["MASS"][1000021]
		ms = 0.
		for r in range(1000001,1000007):
			ms += d.blocks["MASS"][r]
		for r in range(2000001,2000007):
			ms += d.blocks["MASS"][r]
		ms /= 12
		if p.has_gluino() and p.has_squark():
			return 'sg','cteq',ms,mg
		if is_gluino(p.particle1) and is_gluino(p.particle2):
			return 'gg','cteq',ms,mg
		if is_squark(p.particle1) and is_squark(p.particle2):
			if p.particle1 > 0 and p.particle2 >0:
				return 'ss','cteq',ms,mg
			elif (p.particle1 > 0 and p.particle2 <0 )or(p.particle1 < 0 and p.particle2 >0 ) :
				s = p.particle1 if p.particle1 > p.particle2 else p.particle2
				b = p.particle2 if p.particle1 > p.particle2 else p.particle1
				return 'sb','cteq',ms,mg
		return "UNKNOWN_PROCESS_OR_UNIMPLEMENTED_PROCESS"

	def _get_nf_input(self, p:Input) -> dict:
		# TODO return masses of squark and gluino
		d = {}
		#d["ps_inlo"] = int(p.order)
		d["nf_final_state_in"],d["nf_pdf"],d["nf_squark_mass"],d["nf_gluino_mass"] = self._get_ps_proc(p)
		return d

	def _check_input(self,p: Input)-> bool:
		"""Checks input parameter for compatibility with Prospino"""
		if p.mu_f != 1. or  p.mu_r != 1.:
			warnings.warn("NLL-fast does not support varying the scales manually.")
			return False
		if p.pdf_lo != "cteq6l1"  or  p.pdf_nlo != "cteq66":
			warnings.warn("NLL-fast does not support all pdfs (CTEQ6L1 and CTEQ66 allowed defaults).")
			return False
		return True

	def _is_valid(self, file: str, p: Input, d) -> bool:
		return super()._is_valid(file, p, d)

	def _parse_file(self, file: str) -> Result:
		#TODO parse result
		pass
		ret=[]
		with open(file) as output:
			for line in output:
				if line.startswith("nn") or line.startswith("ng") or line.startswith("ns") or line.startswith("sg") or line.startswith("ll") or line.startswith("gg") or line.startswith("ss") or line.startswith("sb"): # TODO generalize
					for s in line[2:].split():
						ret.append(float(s))
		return Result(ufloat(ret[8],ret[8]*ret[9]),ufloat(ret[10],ret[10]*ret[11]) if ret[10]!= 0. else None,None)	

	def _prepare(self, p: Input, **kwargs) -> RunParam:
		rp = super()._prepare(p,**kwargs)
		if not rp.skip:
			d = p.__dict__
			data = pkgutil.get_data(__name__, "prospino_main.f90_template").decode('utf-8')
			src = Template(data)
			#compute dependent pieces for template
			for k,v in self._get_nf_input(p).items():
				d[k]=v
			result = src.substitute(d)
			#open(rp.in_file, "w").write(result)
			open(rp.out_file, "w").write(result + "\n\n")
			#rdir = self.get_output_dir() + rp.name + ".rdir"
			#if os.path.exists(rdir) and os.path.isdir(rdir):
			#	shutil.rmtree(rdir)
			#shutil.copytree(self.get_path(),rdir)
			#open(rdir  +"/prospino_main.f90", "w").write(result)
			## compile
			#subprocess.Popen("cd " + rdir + " && make", shell=True,stdout=subprocess.DEVNULL).wait()

			open(rp.execute, "w").write("#!/bin/sh\n"+
				"{path}/nll-fast {proc} {pdf} {sq} {gl}> {out}".format(
				path=self.get_path(),out=rp.out_file , proc = d["nf_final_state_in"], sq = d["nf_squark_mass"], gl = d["nf_gluino_mass"], pdf = d["nf_pdf"]))
			st = os.stat(rp.execute)
			os.chmod(rp.execute, st.st_mode | stat.S_IEXEC)


			sname = d['slha']
			with open(self.get_output_dir() + sname, 'r') as f:
				open(rp.out_file, "a").write(f.read() + "\n\n")
		return rp



# Legacy 
default_nllfast_runner = NLLfastRunner("~/git/nll-fast/")
"""Default Prospino Runner to provide backward compatibility"""
run = default_nllfast_runner.run
set_path = default_nllfast_runner.set_path
get_path = default_nllfast_runner.get_path