pylhc_submitter/constants/autosix.py
"""
Autosix
-------
Collections of constants and paths used in autosix.
:module: constants.autosix
:author: jdilly
"""
from dataclasses import dataclass, fields, MISSING
from pathlib import Path
from typing import Union
import logging
from pylhc_submitter.constants.external_paths import MADX_BIN, SIXDESK_UTILS
from pylhc_submitter.constants.general import PMASS
LOG = logging.getLogger(__name__)
# Program Paths ----------------------------------------------------------------
UTILITIES_DIR: Path = Path("utilities")
BASH_DIR: Path = UTILITIES_DIR / "bash"
SETENV_SH: Path = BASH_DIR / "set_env.sh"
MAD_TO_SIXTRACK_SH: Path = BASH_DIR / "mad6t.sh"
RUNSIX_SH: Path = BASH_DIR / "run_six.sh"
RUNSTATUS_SH: Path = BASH_DIR / "run_status"
DOT_PROFILE: Path = BASH_DIR / "dot_profile"
SIXDB: Path = UTILITIES_DIR / "externals" / "SixDeskDB" / "sixdb"
SIXDESKLOCKFILE: Path = "sixdesklock"
# Constants and Requirements ---------------------------------------------------
HEADER_BASEDIR = "BASEDIR"
# AutoSix Environment (also defines defaults) ---
@dataclass
class AutoSixEnvironment:
mask_text: str
working_directory: Path
executable: Union[str, Path] = MADX_BIN
python2: Union[str, Path] = None
python3: Union[str, Path] = "python3"
da_turnstep: int = 100
sixdesk_directory: Path = SIXDESK_UTILS
unlock: bool = False
max_stage: 'Stage' = None
ssh: str = None
stop_workspace_init: bool = False
apply_mad6t_hacks: bool = False
resubmit: bool = False
max_materialize: int = None
# Sixenv ---
class _SetDynamicallyType:
pass
SetDynamically = _SetDynamicallyType()
"""Sentinel value defining a parameter that is set later dynamically."""
@dataclass
class SixDeskEnvironment:
TURNS: int
AMPMIN: int
AMPMAX: int
AMPSTEP: int
ANGLES: int
JOBNAME: str = SetDynamically
WORKSPACE: str = SetDynamically
BASEDIR: str = SetDynamically
SCRATCHDIR: str = SetDynamically
TURNSPOWER: int = SetDynamically
RESUBMISSION: int = 0 # 0: never, 1: if fort.10, 2: always
PLATFORM: str = "HTCondor"
LOGLEVEL: int = 0 # 0: basic + errors , 1: + info, >=2: + debug
FIRSTSEED: int = 1
LASTSEED: int = 60
RUNTYPE: str = "col" # 'col' or 'inj'
NPAIRS: int = 30 # 1-32 particle pairs
EMITTANCE: float = 3.75 # normalized emittance
DIMENSIONS: int = 6 # Phase-Space dimensions
WRITEBINS: int = 500
ENERGY: float = None
GAMMA: float = SetDynamically # calculated from ENERGY, needs to be here so asdict() converts it
def __post_init__(self):
# Check Runtype, Energy and Gamma ---
energy_map = {'inj': 450_000., 'col': 7000_000.}
if self.RUNTYPE not in energy_map.keys():
raise ValueError(f"RUNTYPE needs to be one of : {list(energy_map.keys())}")
if self.ENERGY is None:
energy = energy_map[self.RUNTYPE]
LOG.debug(f"Energy for '{self.RUNTYPE}' defaults to {energy}")
self.ENERGY = energy
if self.ENERGY >= 6500_000 and self.RUNTYPE != list(energy_map.keys())[1]:
LOG.warning(f"Runtype is {self.RUNTYPE}, yet energy is set to {self.ENERGY}. Are you sure?")
self.GAMMA = self.ENERGY / 1000 / PMASS
# Check Seeds ---
if any(getattr(self, key) is None for key in SEED_KEYS): # set by user to None
for key in SEED_KEYS:
setattr(self, key, 0)
# the following checks are limits of SixDesk in 2020
# and might be fixed upstream in the future
if self.AMPMAX < self.AMPMIN:
raise ValueError("Given AMPMAX is smaller than AMPMIN.")
if (self.AMPMAX - self.AMPMIN) % self.AMPSTEP:
raise ValueError("The amplitude range need to be dividable by the amplitude steps!")
if not self.ANGLES % 2:
raise ValueError("The number of angles needs to be an uneven one.")
SIXENV_REQUIRED = [f.name for f in fields(SixDeskEnvironment) if (f.default is MISSING)] # required by user
SIXENV_OPTIONAL = [f.name for f in fields(SixDeskEnvironment) if (f.default is not MISSING) and (f.default is not SetDynamically)]
SEED_KEYS = ["FIRSTSEED", "LASTSEED"]
# SixDB and Postprocess ---
SIXTRACK_OUTPUT_FILES = "fort.10.gz", "Sixout.zip"
SIXTRACK_INPUT_CHECK_FILES = "JOB_NOT_YET_STARTED", "JOB_NOT_YET_COMPLETED"
HEADER_NTOTAL, HEADER_INFO, HEADER_HINT = "NTOTAL", "INFO", "HINT"
MEAN, STD, MIN, MAX, N = "MEAN", "STD", "MIN", "MAX", "N"
SEED, ANGLE, ALOST1, ALOST2, AMP = "SEED", "ANGLE", "ALOST1", "ALOST2", "A"
# Errors ---
class StageSkip(Exception):
""" Indicates that the stage was not completed or skipped entirely.
This can be due to an error or on purpose (e.g. user interaction before
restart). """
pass
class StageStop(Exception):
""" A signal sent at the end of a Stage indicating, that it has succeeded
and that any iteration should be stopped after this Stage as the jobs have
been submitted and the user needs to wait for them to finish. """
pass
# Workspace Paths --------------------------------------------------------------
# Input ---
def get_workspace_path(jobname: str, basedir: Path) -> Path:
return basedir / f"workspace-{jobname}"
def get_scratch_path(basedir: Path) -> Path:
return basedir / f"scratch-0"
def get_sixjobs_path(jobname: str, basedir: Path) -> Path:
return get_workspace_path(jobname, basedir) / "sixjobs"
def get_sixdeskenv_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "sixdeskenv"
def get_sysenv_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "sysenv"
def get_masks_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "mask"
def get_track_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "track"
def get_database_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / f"{jobname}.db"
def get_sixtrack_input_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "sixtrack_input"
def get_mad6t_mask_path(jobname: str, basedir: Path) -> Path:
return get_sixtrack_input_path(jobname, basedir) / "mad6t.sh"
def get_mad6t1_mask_path(jobname: str, basedir: Path) -> Path:
return get_sixtrack_input_path(jobname, basedir) / "mad6t1.sh"
# Output ---
def get_autosix_results_path(jobname: str, basedir: Path) -> Path:
return get_sixjobs_path(jobname, basedir) / "autosix_output"
def get_stagefile_path(jobname: str, basedir: Path) -> Path:
return get_autosix_results_path(jobname, basedir) / "stages_completed.txt"
def get_tfs_da_path(jobname: str, basedir: Path) -> Path:
return get_autosix_results_path(jobname, basedir) / f"{jobname}_da.tfs"
def get_tfs_da_seed_stats_path(jobname: str, basedir: Path) -> Path:
return get_autosix_results_path(jobname, basedir) / f"{jobname}_da_per_seed.tfs"
def get_tfs_da_angle_stats_path(jobname: str, basedir: Path) -> Path:
return get_autosix_results_path(jobname, basedir) / f"{jobname}_da_per_angle.tfs"