mcextract.py 2.96 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12
import yaml
import numpy as np
import itertools

'''This module can be used to easily extract Monte Carlo results from the *.results.yml file produced by the loadleveller library.'''

class Observable:
    def __init__(self, num_tasks):
        self.rebinning_bin_length = np.zeros(num_tasks)
        self.rebinning_bin_count = np.zeros(num_tasks)
        self.autocorrelation_time = np.zeros(num_tasks)+np.nan

Lukas Weber's avatar
Lukas Weber committed
13 14
        self.mean = [None for i in range(num_tasks)]
        self.error = [None for i in range(num_tasks)]
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37

class MCArchive:
    def __init__(self, filename):
        with open(filename, 'r') as f:
            doc = yaml.load(f)

        param_names = set(sum([list(task['parameters'].keys()) for task in doc], []))
        observable_names = set(sum([list(task['results'].keys()) for task in doc], []))
        self.num_tasks = len(doc)
        
        self.parameters = dict(zip(param_names, [[None for _ in range(self.num_tasks)] for _ in param_names]))
        self.observables = dict(zip(observable_names, [Observable(self.num_tasks) for _ in observable_names]))

        for i, task in enumerate(doc):
            for param, value in task['parameters'].items():
                self.parameters[param][i] = value

            for obs, value in task['results'].items():
                o = self.observables[obs]
                o.rebinning_bin_length[i] = value['rebinning_bin_length']
                o.rebinning_bin_count[i] = value['rebinning_bin_count']
                o.autocorrelation_time[i] = value['autocorrelation_time']

Lukas Weber's avatar
Lukas Weber committed
38 39
                o.mean[i] = value['mean']
                o.error[i] = value['error']
40

Lukas Weber's avatar
Lukas Weber committed
41 42
    def filter_mask(self, filter):
        if not filter:
43 44
            return [True for _ in range(self.num_tasks)]

Lukas Weber's avatar
Lukas Weber committed
45
        return [all(self.parameters[key][i] == val for key, val in filter.items()) for i in range(self.num_tasks)]
46

Lukas Weber's avatar
Lukas Weber committed
47 48
    def get_parameter(self, name, unique=False, filter={}):
        selection = list(itertools.compress(self.parameters[name], self.filter_mask(filter)))
49 50 51 52 53
        
        if unique:
            return list(sorted(set(selection)))
        return selection

Lukas Weber's avatar
Lukas Weber committed
54
    def get_observable(self, name, filter={}):
55 56 57
        orig = self.observables[name]
        
        selection = Observable(0)
Lukas Weber's avatar
Lukas Weber committed
58
        mask = self.filter_mask(filter)
59 60 61
        selection.rebinning_bin_count = orig.rebinning_bin_count[mask]
        selection.rebinning_bin_length = orig.rebinning_bin_length[mask]
        selection.autocorrelation_time = orig.autocorrelation_time[mask]
Lukas Weber's avatar
Lukas Weber committed
62 63 64 65 66 67 68
        
        selection.mean = [m for i, m in enumerate(orig.mean) if mask[i]]
        selection.error = [m for i, m in enumerate(orig.error) if mask[i]]

        if all(len(m) == len(selection.mean[0]) for m in selection.mean):
            selection.mean = np.array(selection.mean)
            selection.error = np.array(selection.error)
69

Lukas Weber's avatar
Lukas Weber committed
70
        if selection.mean.shape[1] == 1:
71 72 73
            selection.mean = selection.mean.flatten()
            selection.error = selection.error.flatten()

74
        return selection