make-results-files.py

#!/usr/bin/env python3

"""
Reduce event-by-event observables (multiplicity, Qn vectors) to batch
quantitities (<mult>, flow cumulants) for each design point and centrality.

Creates files <IC name>_<observable>.dat.
"""

import glob
import sys
import subprocess

import numpy as np
from hic import flow


def process_file(fn):
    """
    Read a file generated by urqmd-observables and calculate batch quantities.

    fn : filename
        Must be gzipped and contain a line

            Nch_mid Nch_flow Q2re Q2im Q3re ...

        for each event.

    returns : [<Nch_flow>, v_2{2}, v_3{2}]

    """
    with subprocess.Popen(('zcat', fn), stdout=subprocess.PIPE) as proc:
        obs = np.array([l.split() for l in proc.stdout], dtype=np.float64)

    # Nch_mid = obs[:, 0]
    Nch_flow = obs[:, 1]
    mult = [Nch_flow.mean()]

    q2, q3 = obs.view(complex)[:, 1:3].T
    vnk = flow.FlowCumulant(Nch_flow, {2: q2, 3: q3})
    k = 2
    cumulants = [vnk.flow(n, k, imaginary='zero') for n in (2, 3)]

    return mult + cumulants


def main():
    try:
        path = sys.argv[1].rstrip('/')
    except IndexError:
        print('usage: {} IC_observables_folder'.format(sys.argv[0]))
        exit()
    ic = path.split('/')[-1]

    cent_strs = tuple(sorted(g.split('/')[-1].rstrip('.gz')
                             for g in glob.iglob(path + '/0/*')))
    cent_mids = tuple(.5*sum(int(i) for i in cstr.split('-'))
                      for cstr in cent_strs)

    with open('{}_cent.dat'.format(ic), 'w') as centfile:
        print(*cent_mids, file=centfile)

    npoints = len(glob.glob(path + '/*'))

    with \
            open('{}_mult.dat'.format(ic), 'w') as multfile, \
            open('{}_v2-2.dat'.format(ic), 'w') as v2file, \
            open('{}_v3-2.dat'.format(ic), 'w') as v3file:
        for n in range(npoints):
            print(ic, n)
            mult, v2, v3 = zip(
                *(process_file('{}/{}/{}.gz'.format(path, n, cstr))
                  for cstr in cent_strs)
            )
            print(*mult, file=multfile)
            print(*v2, file=v2file)
            print(*v3, file=v3file)


if __name__ == "__main__":
    main()