Source code for pylhc_submitter.constants.autosix

"""
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) ---
[docs] @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."""
[docs] @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 ---
[docs] 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
[docs] 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"