mcextract.py 3.2 KB
Newer Older
1 2 3 4
import yaml
import numpy as np
import itertools

Lukas Weber's avatar
Lukas Weber committed
5 6 7 8 9
try:
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader
    
10 11 12 13 14 15 16 17
'''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
18 19
        self.mean = [np.array([np.nan]) for i in range(num_tasks)]
        self.error = [np.array([np.nan]) for i in range(num_tasks)]
20 21 22 23

class MCArchive:
    def __init__(self, filename):
        with open(filename, 'r') as f:
Lukas Weber's avatar
Lukas Weber committed
24
            doc = yaml.load(f, Loader=SafeLoader)
25 26 27 28 29 30 31 32 33 34 35 36 37 38

        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]
39 40 41
                o.rebinning_bin_length[i] = int(value.get('rebinning_bin_length',0))
                o.rebinning_bin_count[i] = int(value.get('rebinning_bin_count',0))
                o.autocorrelation_time[i] = float(value.get('autocorrelation_time',0))
42

43 44
                o.mean[i] = np.array(value['mean'], dtype=float)
                o.error[i] = np.array(value['error'], dtype=float)
45

Lukas Weber's avatar
Lukas Weber committed
46 47
    def filter_mask(self, filter):
        if not filter:
48 49
            return [True for _ in range(self.num_tasks)]

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

Lukas Weber's avatar
Lukas Weber committed
52 53
    def get_parameter(self, name, unique=False, filter={}):
        selection = list(itertools.compress(self.parameters[name], self.filter_mask(filter)))
54 55 56 57 58
        
        if unique:
            return list(sorted(set(selection)))
        return selection

Lukas Weber's avatar
Lukas Weber committed
59
    def get_observable(self, name, filter={}):
60 61 62
        orig = self.observables[name]
        
        selection = Observable(0)
Lukas Weber's avatar
Lukas Weber committed
63
        mask = self.filter_mask(filter)
64 65 66
        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
67 68 69 70 71 72 73
        
        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)
74

Lukas Weber's avatar
Lukas Weber committed
75 76 77
            if selection.mean.shape[1] == 1:
                selection.mean = selection.mean.flatten()
                selection.error = selection.error.flatten()
78

79
        return selection