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 bash:
`source /sl/cmd3/cc8-64/Cmd3Off/tune.cmd3_runs_scripts.sh`

This script is inspired by listoffdata.py script (/sl/cmd3/cc7-64/Cmd3Off/scripts/listoffdata.py)
"""

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)
    """
    
    points, seasons = [], []
    # Get seasons (they are sorted by date)
    for e in EnergyPoint.objects.filter(name__startswith="Season"):
        info = [e.name[7:], e.startrun, e.endrun]
        seasons.append(info)
    
    # Get energy points (they are sorted by date) and match corresponing seasons
    season_idx = 0
    for e in EnergyPoint.objects.filter(startrun__gte=startrun, endrun__lte=stoprun):
        while((e.startrun < seasons[season_idx][1]) or (e.endrun > seasons[season_idx][2])):
            season_idx+=1
        info = (e.name, e.startrun, e.endrun, e.energy, e.starttime, e.endtime, seasons[season_idx][0])
        if not(e.name.startswith("Season")):
            points.append(info)
    return points
    
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()