pythonABC/data_manipulation.py at master · Neojume/pythonABC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import os.path
import os
import pickle

__all__ = ['ABCData',
           'get_filename',
           'save',
           'load']

class ABCData(object):

    '''
    Storage class for abc-data.
    '''

    def __init__(self, algorithm):
        '''
        Create a data-container for this combination.
        '''

        self.algorithm = type(algorithm).__name__
        self.alg_args = algorithm.get_parameters()
        self.problem = type(algorithm.problem).__name__

        self.num_data = 0
        self.list_of_samples = []
        self.list_of_accepts = []
        self.list_of_sim_calls = []
        self.list_of_sim_locs = []

    def add_datum(self, algorithm):
        '''
        Add data to the database.
        '''
        self.list_of_samples.append(algorithm.samples)
        self.list_of_sim_calls.append(algorithm.sim_calls)
        try:
            self.list_of_accepts.append(algorithm.accepted)
        except AttributeError:
            # Algorithm has no accepted attribute
            pass
        try:
            self.list_of_sim_locs.append(algorithm.xs)
        except AttributeError:
            pass

        self.num_data += 1


def get_filename(algorithm):
    '''
    Creates a unique filename for this algorithm.

    Arguments
    ---------
    algorithm : instance of an `ABC_Algorithm`
        The algorithm instance to generate a name for.
    '''

    return str(algorithm) + '.abc'


def load(algorithm, data_dir=False):
    '''
    Loads the results for the given algorithm with given parameters for
    the given problem. If there are no results for these combinations None is
    returned.

    Parameters
    ----------
    algorithm :
        An instance of an ABC algorithm
    data_dir : boolean
        Whether there is a directory `data` where the file is stored

    Returns
    -------
    data : ABCData or None
        The loaded data. Or None if there is no data.
    '''

    filename = get_filename(algorithm)

    if data_dir:
        path = os.path.join(os.getcwd(), 'data', filename)
    else:
        path = os.path.join(os.getcwd(), filename)

    if os.path.isfile(path):
        with open(path, 'rb') as f:
            data = pickle.load(f)
        return data
    else:
        return None


def save(algorithm, data_dir=False):
    '''
    Saves the results for the given algorithm with given parameters for
    the given problem.

    Parameters
    ----------
    algorithm :
        An instance of an ABC algorithm
    data_dir : boolean
        Whether there is a directory `data` where the file is stored
    '''

    filename = get_filename(algorithm)

    if data_dir:
        path = os.path.join(os.getcwd(), 'data', filename)
    else:
        path = os.path.join(os.getcwd(), filename)

    if os.path.isfile(path):
        # If file exists open and append to the existing database
        with open(path, 'rb') as f:
            data = pickle.load(f)
    else:
        # Otherwise create a new database
        data = ABCData(algorithm)

    data.add_datum(algorithm)

    with open(path, 'wb') as f:
        pickle.dump(data, f)