DIST block for PySixDesk

README.md

@@ -86,8 +86,8 @@ By default the jobs will be submitted to HTCondor. If you want to use a differen
       ```python
       myWS.init_study('myStudy')
       ```
-   
-   1. edit the `config.py` file to define the study (e.g. define `.mask` file, add scan parameters, etc...);
+
+   1. edit the `config.py` file to define the parameters you want to scan for your study as well as their range. This also includes adapting the `.mask` file and parametrizing the variables that are to be scanned (use `%` sign to declare these parameters followed by the parameter name defined in `config.py`, e.g. `%SEEDRAN`). Note that PySixDesk will set up a scan making all the possible combinations between defined parameter values (Cartesian product);
 
    1. load definition of study in `config.py` and create/update database:
 
@@ -96,12 +96,14 @@ By default the jobs will be submitted to HTCondor. If you want to use a differen
       myStudy.update_db() # only needed for a new study or when parameters are changed
       ```
 
-   1. prepare and submit pre-processing jobs (e.g. MADX job and sixtrack one turn jobs), and collect results:
+   This initializes a database file. Depending on the type of database you choose in the `config.py` file, it is either an SQLite database stored locally, or, alternatively, a MySQL database hosted by the CERN DBonDemand Service.
+
+   1. prepare and submit pre-processing jobs (i.e. MADX job and sixtrack one turn jobs), and collect results:
 
       ```python
       myStudy.prepare_preprocess_input()
       myStudy.submit(0, 5) # 0 stands for preprocess job, 5 is trial number 
-      myStudy.collect_result(0) # collect results locally
+      myStudy.collect_result(0) # collect results locally once jobs are finished
       ```
 
    1. prepare and submit actual sixtrack jobs, and collect results:
@@ -110,8 +112,10 @@ By default the jobs will be submitted to HTCondor. If you want to use a differen
       myStudy.prepare_sixtrack_input()
       or
       myStudy.prepare_sixtrack_input(True) #True: submit jobs to Boinc
+
       myStudy.submit(1, 5) # 1 stands for sixtrack job, 5 is trial number 
-      myStudy.collect_result(1) # 1 stands for sixtrack job 
+
+      myStudy.collect_result(1) # 1 stands for sixtrack job; collect results once finished
       or
       myStudy.collect_result(1, True) # True: collect results from boinc spool directory
       ```

pysixdesk/lib/machineparams.py

@@ -1,5 +1,3 @@
-from .utils import merge_dicts
-
 # LHC settings
 LHC = {}
 # TODO: fill these in.
@@ -36,8 +34,8 @@
                 tune_x=62.31,
                 tune_y=60.32,
                 )
-LHC['inj'] = merge_dicts(_LHC_DEF, _LHC_INJ)
-LHC['col'] = merge_dicts(_LHC_DEF, _LHC_COL)
+LHC['inj'] = {**_LHC_DEF, **_LHC_INJ}
+LHC['col'] = {**_LHC_DEF, **_LHC_COL}
 
 # HLLHC settings
 HLLHC = {}
@@ -74,5 +72,5 @@
                   tune_x=62.31,
                   tune_y=60.32,
                   )
-HLLHC['inj'] = merge_dicts(_HLLHC_DEF, _HLLHC_INJ)
-HLLHC['col'] = merge_dicts(_HLLHC_DEF, _HLLHC_COL)
+HLLHC['inj'] = {**_HLLHC_DEF, **_HLLHC_INJ}
+HLLHC['col'] = {**_HLLHC_DEF, **_HLLHC_COL}

pysixdesk/lib/preprocess.py

@@ -190,7 +190,13 @@ def sixtrack_prep_job(self, fort_cfg, source_prefix=None,
             madx_fc3 = source_prefix / madx_fc3
 
         # concatenate
-        utils.concatenate_files([dest, madx_fc3], output_file)
+        # utils.concatenate_files([dest, madx_fc3], output_file)
+        utils.sandwich(dest,
+                       output_file,
+                       path_prefix=source_prefix,
+                       logger=self._logger)
+        # with open(output_file, 'r') as fp:
+        #     print(''.join(fp.readlines()))
         # if not source.samefile(output_file):
         #     utils.diff(source, output_file, logger=self._logger)
 
@@ -338,6 +344,8 @@ def madx_prep(self, output_file='madx_in'):
         # show diff
         # utils.diff(self.madx_cfg["mask_file"], output_file,
         #            logger=self._logger)
+        # with open(output_file, 'r') as fp:
+        #     print(''.join(fp.readlines()))
 
     def madx_run(self, mask):
         """Runs madx.
@@ -361,6 +369,8 @@ def madx_run(self, mask):
             raise Exception(content)
         else:
             self._logger.info("MADX has completed properly!")
+        # with open('madx_stdout', 'r') as fp:
+        #     print(''.join(fp.readlines()))
 
     def madx_job(self):
         """Controls madx job execution.
@@ -495,5 +505,5 @@ def write_oneturnresult(self):
         where = f"task_id={job.task_id}"
         job_table['status'] = 'incomplete'
         job_table['mtime'] = int(time.time() * 1E7)
-        job.db.update('preprocess_wu'. job_table, where)
+        job.db.update('preprocess_wu', job_table, where)
         raise e

pysixdesk/lib/sixtrack.py

@@ -284,8 +284,34 @@ def sixtrack_prep_job(self, fort_cfg, source_prefix=None,
         if source_prefix is not None:
             madx_fc3 = source_prefix / madx_fc3
 
+        utils.sandwich(dest,
+                       output_file,
+                       path_prefix=source_prefix,
+                       logger=self._logger)
+        # with open(output_file, 'r') as fp:
+        #     lines = fp.readlines()
+        #     print(''.join(lines))
+
+
+        # Ugly hack to generate the particle distribution for use with DIST block
+        if self.fort_cfg['dist_type'] == 'polar':
+            amps = json.loads(fort_cfg['phase_space_var1'])
+            angle = fort_cfg['phase_space_var2']
+            num_particle_pairs = int(fort_cfg['nss'] / 2.)
+            os.system(f'python init_polar_dist.py {amps[0]} {amps[1]} ' +
+                      f'{angle} {num_particle_pairs}')
+        elif self.fort_cfg['dist_type'] == 'cartesian':
+            amps_x = json.loads(fort_cfg['phase_space_var1'])
+            amps_y = json.loads(fort_cfg['phase_space_var2'])
+            num_particle_pairs = int(fort_cfg["nss"] / 2.)
+            os.system(
+                f'python init_cartesian_dist.py {amps_x[0]} {amps_x[1]} ' +
+                f'{amps_y[0]} {amps_y[1]} {num_particle_pairs}')
+        else:
+            raise ValueError('Unknown dist_type')
+
         # concatenate
-        utils.concatenate_files([dest, madx_fc3], output_file)
+        # utils.concatenate_files([dest, madx_fc3], output_file)
 
     def sixtrack_run(self, output_file):
         """Runs sixtrack.
@@ -438,7 +464,8 @@ def boinc_submit(self, job_name):
         """
         # zip all the input files, e.g. fort.3 fort.2 fort.8 fort.16
         input_zip = job_name + '.zip'
-        inputs = ['fort.2', 'fort.3', 'fort.8', 'fort.16'] + self.cr_inputs
+        inputs = ['fort.2', 'fort.3', 'fort.8', 'fort.16', 'fort.34'] + self.cr_inputs
+        inputs += ['input_dist.txt', 'fort.3.aper']
         with zipfile.ZipFile(input_zip, 'w', zipfile.ZIP_DEFLATED) as ziph:
             for infile in inputs:
                 if infile in os.listdir('.'):

pysixdesk/lib/study.py

@@ -83,7 +83,7 @@ def _defaults(self):
             'fc.34': 'fort.34'}
 
         self.oneturn_sixtrack_input['input'] = dict(self.madx_output)
-        self.sixtrack_input['temp'] = 'fort.3'
+        self.sixtrack_input['fort_file'] = 'fort.3'
         self.sixtrack_output = ['fort.10']
 
         self.db_info['db_type'] = 'sql'
@@ -248,6 +248,7 @@ def _update_db_params(self):
         madx_keys = self.params.madx.keys()
         six_keys = (list(self.params.sixtrack.keys()) +
                     list(self.params.phasespace.keys()))
+
         # preprocess_wu already in the table
         prep_entries = self.db.select('preprocess_wu', madx_keys)
         prep_wu_done = set(prep_entries)
@@ -394,6 +395,9 @@ def _prep_sixtrack_cfg(self):
         self.sixtrack_config['sixtrack']['output_files'] = json.dumps(self.sixtrack_output)
         self.sixtrack_config['sixtrack']['test_turn'] = json.dumps(self.env['test_turn'])
         self.sixtrack_config['six_results'] = self.tables['six_results']
+        # TODO: UGLY HACK TO ADD THE aperture_losses FILE...
+        # self.sixtrack_config['aperture_losses'] = self.tables['aperture_losses']
+
         if self.collimation:
             self.sixtrack_config['aperture_losses'] = self.tables['aperture_losses']
             self.sixtrack_config['collimation_losses'] = self.tables['collimation_losses']
@@ -445,17 +449,18 @@ def _run_calcs(self):
                 self._logger.info('Queued update of sixtack_wu/wu_id:'
                                   f'{row["wu_id"]} with {update_cols}.')
 
-        # update everything at once
-        self._logger.info(f'Updating {len(calc_out_to_be_updated)} rows of '
-                          'sixtrack_wu.')
-        # turn list of dicts into dict of lists
-        calc_out_to_be_updated = {k: [dic[k] for dic in calc_out_to_be_updated]
-                                  for k in calc_out_to_be_updated[0]}
-        where_to_be_updated = {k: [dic[k] for dic in where_to_be_updated]
-                               for k in where_to_be_updated[0]}
-        self.db.updatem('sixtrack_wu',
-                        calc_out_to_be_updated,
-                        where=where_to_be_updated)
+        if calc_out_to_be_updated:
+            # update everything at once
+            self._logger.info(f'Updating {len(calc_out_to_be_updated)} rows of '
+                              'sixtrack_wu.')
+            # turn list of dicts into dict of lists
+            calc_out_to_be_updated = {k: [dic[k] for dic in calc_out_to_be_updated]
+                                      for k in calc_out_to_be_updated[0]}
+            where_to_be_updated = {k: [dic[k] for dic in where_to_be_updated]
+                                   for k in where_to_be_updated[0]}
+            self.db.updatem('sixtrack_wu',
+                            calc_out_to_be_updated,
+                            where=where_to_be_updated)
 
     def update_db(self, db_check=False):
         '''Update the database whith the user-defined parameters'''
@@ -545,11 +550,12 @@ def query(index):
                 print(query_list)
                 for i in wus:
                     print(i)
+            return results
 
         if job == 0 or job == 2:
-            query(0)
+            return query(0)
         if job == 1 or job == 2:
-            query(1)
+            return query(1)
 
     def submit(self, typ, trials=5, *args, **kwargs):
         '''Sumbit the preporcess or sixtrack jobs to htctondor.
@@ -627,6 +633,8 @@ def collect_result(self, typ, boinc=False):
             info_sec['boinc_results'] = self.env['boinc_results']
             info_sec['boinc'] = boinc
             info_sec['st_pre'] = self.st_pre
+            # TODO: UGLY HACK TO PARSE THE APERTURE_LOSSES FILE...
+            # config['aperture_losses'] = self.tables['aperture_losses']
 
             info_sec['outs'] = Table.result_table(self.sixtrack_output)
             if self.collimation:
@@ -670,7 +678,9 @@ def prepare_sixtrack_input(self, resubmit=False, boinc=False, groupby=None,
             action = 'resubmit'
         else:
             constraints = "status='incomplete' and preprocess_id in %s" % str(
-                preprocess_outs[0])
+                preprocess_outs[0]).replace(',)', ')')  # if there is a single
+            # job, the comma of the tuple (1,) breaks the sql query because
+            # there is no proper sql sanitization...
             action = 'submit'
         names = self.tables['sixtrack_wu'].keys()
         outputs = self.db.select('sixtrack_wu',

pysixdesk/lib/study_params.py

@@ -1,6 +1,7 @@
 import re
 import logging
 
+from math import pi
 from pathlib import Path
 from collections import OrderedDict
 from itertools import product
@@ -9,7 +10,7 @@
 
 from . import machineparams
 from .constants import PROTON_MASS
-from .utils import PYSIXDESK_ABSPATH, merge_dicts
+from .utils import PYSIXDESK_ABSPATH, linspace
 
 
 class StudyParams:
@@ -46,7 +47,7 @@ def __init__(self,
         # comment regexp
         self._reg_comment = re.compile(r'^(\s?!|\s?/).*', re.MULTILINE)
         # placeholder pattern regexp
-        self._reg = re.compile(r'%([a-zA-Z0-9_]+/?)')
+        self._reg = re.compile(r'%(?!FILE|%)([a-zA-Z0-9_]+/?)')
         self.fort_path = fort_path
         self.mask_path = mask_path
         # initialize empty calculation queue
@@ -85,18 +86,17 @@ def __init__(self,
                             "writebins": 1,
                             }
         self.machine_defaults = machine_defaults
-        self.defaults = merge_dicts(self.f3_defaults, self.machine_defaults)
+        self.defaults = {**self.f3_defaults, **self.machine_defaults}
         # phasespace params
         # TODO: find sensible defaults for the phasespace parameters.
-        amp = [8, 10, 12]  # The amplitude
-        self.phasespace = {"amp": list(zip(amp, amp[1:])),
-                           "kang": list(range(1, 1 + 1)),
-                           "kmax": 5,
+        self.phasespace = {"phase_space_var1": [],
+                           "phase_space_var2": [],
                            }
 
         self.madx = self.find_patterns(self.mask_path)
         self.sixtrack = self.find_patterns(self.fort_path,
                                            mandatory=['chrom_eps', 'CHROM'])
+        self.sixtrack['dist_type'] = 'None'
 
     @property
     def _sixtrack_only(self):
@@ -117,6 +117,10 @@ def oneturn(self):
         sixtrack['toggle_coll/'] = '/'
         return sixtrack
 
+    @staticmethod
+    def da_angles(start=0, end=pi/2, n=7):
+            return linspace(start, end, n + 2)[1: -1]  # exclusive
+
     def keys(self):
         """Gets the keys of 'self.madx', 'self.sixtrack' and 'self.phasespace'.
 
@@ -180,7 +184,7 @@ def find_patterns(self, file_path, keep_none=True, mandatory=None):
 
     @staticmethod
     def _combinations_prep(**kwargs):
-        '''Sanitizes the paramter values.
+        '''Sanitizes the parameter values.
 
         Args:
             **kwargs: Parameter name = parameter value.

pysixdesk/lib/utils.py

@@ -59,6 +59,58 @@ def download_output(filenames, dest, zp=True):
                 shutil.copy(filename, dest)
 
 
+def sandwich(in_file, out_file, path_prefix='', logger=None):
+    '''
+    Looks for any patterns matching '^%FILE:.*' in in_file, then replaces the
+    match with the content of the file following the match.
+    A path prefix can be specified to look for matched file in another
+    directory. If the matched file is not found, comment out the pattern.
+    Example:
+        contents of in_file:
+            aaaaaaaaa
+            %FILE:insert.txt
+            aaaaaaaaa
+        contents of insert.txt:
+            bbbbbbbbb
+            bbbbbbbbb
+        writes to out_file:
+            aaaaaaaaa
+            bbbbbbbbb
+            bbbbbbbbb
+            aaaaaaaaa
+    TODO: maybe it's best to return the the sandwiched lines and not write the
+    file, so that more parsing can happen in memory without reopening and
+    writing files. Also might be best to take as input the contents itself for
+    the same reason.
+    '''
+
+    if logger is not None:
+        display = logger.warning
+    else:
+        display = print
+
+    with open(in_file, 'r') as f:
+        in_lines = f.read()
+
+    reg = re.compile(r'^%FILE:.*', re.MULTILINE)
+    for m in re.finditer(reg, in_lines):
+        m_str = m.group()
+        try:
+            if path_prefix is None:
+                path_prefix = ''
+            fname = os.path.join(path_prefix, m_str.split(':')[1].lstrip())
+            with open(fname, 'r') as f:
+                file_lines = f.read()
+            in_lines = re.sub(f'{m_str}', file_lines, in_lines)
+        except FileNotFoundError as e:
+            display(e)
+            display(f'Commenting out {m_str} for {out_file}')
+            in_lines = re.sub(f'{m_str}', f'/{m_str}', in_lines)
+
+    with open(out_file, 'w') as out:
+        out.write(in_lines)
+
+
 def check_fort3_block(fort3, block):
     '''Check the existence of the given block in fort.3'''
 
@@ -238,15 +290,13 @@ def condor_logger(name):
     return logger
 
 
-def merge_dicts(x, y):
-    """Merges two dicts.
-
-    Returns:
-        dict: Merged dict
-    """
-    z = x.copy()
-    z.update(y)
-    return z
+def linspace(a, b, n):
+    '''Numpyless linear spacing function.
+    '''
+    if n < 2:
+        return a
+    diff = (float(b) - a)/(n - 1)
+    return [diff * i + a for i in range(n)]
 
 
 class ProgressBar(object):