sources/elabelizer.py

"""
The script is for corresponding the elabels from database and compton average measurements
run it on cc-8 with the command:
`python3 elabelizer.py`

If this script is not working properly, try to setup environment firstly with the following command in shell:
`source /sl/cmd3/cc8-64/Cmd3Off/tune.cmd3_runs_scripts.sh`
"""

import os 
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'settings_read_all')


import django
django.setup()

from cmdweb.apps.cmd3off.energypoints.models import EnergyPoint, Param, ParamData
from urllib.request import urlopen

    
def retrieve_elabels(startrun, stoprun):
    """Retrieves elabels from database
    
    Parameters:
    -----------
    startrun : int
        first run for retrieving
    stoprun : int
        last run for retrieving
    
    Returns:
    --------
    List[Tuple[str, int, int, float, datetime, datetime]
        list of elabels and information about them
        (elabel, firstrun, lastrun, nominal energy, starttime, stoptime, season)
    """
    elabels = []
    season = ""
    for e in EnergyPoint.objects.all():
        if (e.name).startswith('Season'):
            season = e.name[7:]
        if e.startrun >= startrun and e.endrun <= stoprun:
            elabel_info = (e.name, e.startrun, e.endrun, e.energy, e.starttime, e.endtime, season)
            if (e.name).startswith('Season'):
                continue
            elabels.append(elabel_info)
    return elabels
    
def get_available_compton_seasons():
    tables_url = 'https://cmd.inp.nsk.su/~compton/gitlist/compton_tables/raw/dev/tables/'
    with urlopen(tables_url) as f:
        data = f.read().decode('utf-8')
        datarows = list(map(lambda x: x[:-4], filter(lambda x: x.endswith('.csv'), data.split('\n'))))
    return datarows

def retrieve_comptons(seasons=None):
    """Retrieves compton measurements from tables
    
    Parameters:
    -----------
    seasons: Optional[List[str]]
        list of seasons for retrieving
        (default is None, it means retrieving all available seasons)
    
    Returns:
    --------
    List[Tuple[str, float, int, int, float]]
        list of compton means measurements
        (season, energy_point, first_run, last_run, mean_energy)
        first_run and last_run can be different from the corresponding energy points from elabels database
    """
    
    available_seasons = get_available_compton_seasons()
    if seasons is not None:
        for s in seasons:
            if s not in available_seasons:
                raise ValueError(f"Season {s} is not found. List of available seasons: {available_seasons}")
    else:
        seasons = available_seasons
        
    def parse_compton_row(row, season):
        items = row.split(',')
        return (season, float(items[0]), int(items[1]), int(items[2]), float(items[3]))
    
    def retrieve_compton_one_season(season):
        with urlopen(f'https://cmd.inp.nsk.su/~compton/gitlist/compton_tables/raw/dev/tables/{season}.csv') as f:
            r = f.read().decode('utf-8').strip()
            rows = list(map(lambda x: parse_compton_row(x, season), r.split('\n')[1:]))
        return rows
                    
    datarows = []
    for s in seasons:
        seasonrows = retrieve_compton_one_season(s)
        datarows.extend(seasonrows)
    
    return datarows
        
def combine_compton_elabel(epoints, comptonpoints):
    """Combines compton energy points and elabels from database
    
    Parameters:
    -----------
    epoints : List[Tuple[...]]
        list of energy points
    comptonpoints : List[Tuple[...]]
        list of compton points
    
    Returns:
    --------
    combpoints : List[Tuple[...]]
        list of combined points (outer join)
    """
    
    epoints_sorted = sorted(epoints, key=lambda x: x[1])
    cpoints_sorted = sorted(comptonpoints, key=lambda x: x[2])
    
    combined_data = []
    
    eidx, cidx = 0, 0
    while (eidx < len(epoints_sorted)) and (cidx < len(cpoints_sorted)):
        erow = epoints_sorted[eidx]
        crow = cpoints_sorted[cidx]
        
        cstart, cstop, cenergy, cseason = crow[2], crow[3], crow[4], crow[0]
        estart, estop, elabel, eseason = erow[1], erow[2], erow[0], erow[6]
        
        if (cstart >= estart) and (cstop <= estop):
            combrow = (eseason, elabel, cenergy)
            eidx += 1
            cidx += 1
            #print(combrow)
        elif (cstart > estop):
            combrow = (eseason, elabel, None)
            eidx += 1
        elif (estart > cstop):
            combrow = (cseason, None, cenergy)
            cidx += 1
        else:
            raise Exception("Something wrong")
        
        #print(combrow)
        combined_data.append(combrow)
    
    for i in range(eidx, len(epoints_sorted)):
        erow = epoints_sorted[i]
        elabel = erow[0]
        combrow = (None, elabel, None)
        combined_data.append(combrow)
        
    for i in range(cidx, len(cpoints_sorted)):
        crow = cpoints_sorted[i]
        cstart, cstop, cenergy, cseason = crow[2], crow[3], crow[4], crow[0]
        combrow = (cseason, None, cenergy)
        combined_data.append(combrow)
    
    return combined_data


def elabelize():
    RUNLIMITS = (17405, 200000)
    SEASONS = None
    epoints = retrieve_elabels(*RUNLIMITS)
    comptonpoints = retrieve_comptons(SEASONS)
    
    combdata = combine_compton_elabel(epoints, comptonpoints)
    return combdata
    
def main():
    combdata = elabelize()
    for c in combdata:
        season, elabel, energy = c
        if season is None:
            season = ''
        if elabel is None:
            elabel = ''
        print(f'{season:>15s}, {elabel:>13s}, {energy}')
    print('The number of not found points in compton is', len(list(filter(lambda x: x[2] is None, combdata))))
    
if __name__=="__main__":
    main()