pipeline.py 5.39 KB
Newer Older
1
2
# -*- coding: utf-8 -*-

3
4
5
6
7
import os
import numpy as np

from mpi4py import MPI

8
9
10
from keepers import Loggable

from likelihoods import Likelihood
11
from magnetic_fields import MagneticFieldFactory
12
13
from observers import Observer
from priors import Prior
14
15
16
17
18
from imagine.pymultinest import pymultinest

comm = MPI.COMM_WORLD
size = comm.size
rank = comm.rank
19
20
21


class Pipeline(Loggable, object):
22
23
24
25
26
27
28
29
30
31
32
33
34
    """
    The pipeline
    - posses all the building blocks: magnetic_field, observer,
        likelihood and prior.
    - if multiple log-likelihoods and log-priors are given: sum the result
    - coordinates the repeated observation in order to compute an ensemble
    - controls which parameters of the magnetic field are tested
        (active parameters)


    """
    def __init__(self, magnetic_field_factory, observer, likelihood, prior,
                 active_variables=[], ensemble_size=1):
35
        self.logger.debug("Setting up pipeline.")
36
        self.magnetic_field_factory = magnetic_field_factory
37
38
39
        self.observer = observer
        self.likelihood = likelihood
        self.prior = prior
40
        self.active_variables = active_variables
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        self.ensemble_size = ensemble_size

    @property
    def observer(self):
        return self._observer

    @observer.setter
    def observer(self, observer):
        if not isinstance(observer, Observer):
            raise TypeError("observer must be an instance of Observer-class.")
        self.logger.debug("Setting observer.")
        self._observer = observer

    @property
    def likelihood(self):
        return self._likelihood

    @likelihood.setter
    def likelihood(self, likelihood):
60
        self.logger.debug("Setting likelihood.")
61
62
63
64
65
66
        self._likelihood = ()
        for l in likelihood:
            if not isinstance(l, Likelihood):
                raise TypeError(
                    "likelihood must be an instance of Likelihood-class.")
            self._likelihood += (l,)
67
68
69
70
71
72
73

    @property
    def prior(self):
        return self._prior

    @prior.setter
    def prior(self, prior):
74
        self.logger.debug("Setting prior.")
75
76
77
78
79
80
        self._prior = ()
        for p in prior:
            if not isinstance(p, Prior):
                raise TypeError(
                    "prior must be an instance of Prior-class.")
            self._prior += (p,)
81
82

    @property
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
    def magnetic_field_factory(self):
        return self._magnetic_field_factory

    @magnetic_field_factory.setter
    def magnetic_field_factory(self, magnetic_field_factory):
        if not isinstance(magnetic_field_factory, MagneticFieldFactory):
            raise TypeError(
                "magnetic_field_factory must be an instance of the "
                "MagneticFieldFactory-class.")
        self.logger.debug("Setting magnetic_field_factory.")
        self._magnetic_field_factory = magnetic_field_factory

    @property
    def active_variables(self):
        return self._active_variables

    @active_variables.setter
    def active_variables(self, active_variables):
        if not isinstance(active_variables, list):
            raise TypeError(
                    "active_variables must be a list.")
        self.logger.debug("Resetting active_variables to %s" %
                          str(active_variables))
        new_active = []
        for av in active_variables:
            new_active += [str(av)]
        self._active_variables = new_active
110
111
112
113
114
115
116
117
118
119
120
121
122
123

    @property
    def ensemble_size(self):
        return self._ensemble_size

    @ensemble_size.setter
    def ensemble_size(self, ensemble_size):

        ensemble_size = int(ensemble_size)
        if ensemble_size <= 0:
            raise ValueError("ensemble_size must be positive!")
        self.logger.debug("Setting ensemble size to %i." % ensemble_size)
        self._ensemble_size = ensemble_size

124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
    def _multinest_likelihood(self, cube, ndim, nparams):
        cube_content = np.empty(ndim)
        for i in xrange(ndim):
            cube_content[i] = cube[i]
        if rank != 0:
            raise RuntimeError("_multinest_likelihood must only be called on "
                               "rank==0.")
        for i in xrange(1, size):
            comm.send(cube_content, dest=i)

        return self._core_likelihood(cube_content)

    def _listen_for_likelihood_calls(self):
        cube = comm.recv(obj=None, source=0)
        self._core_likelihood(cube)

    def _core_likelihood(self, cube):
        # translate cube to variables
        variables = {}
        for i, av in enumerate(self.active_variables):
            variables[av] = cube[i]

        # create magnetic field
        b_field = self.magnetic_field_factory(variables=variables,
                                              ensemble_size=self.ensemle_size)
        # create observables
        observables = self.observer(b_field)

        # add up individual log-likelihood terms
        likelihood = 0
        for like in self.likelihood:
            likelihood += like(observables)

        return likelihood

159
    def __call__(self, variables):
160
161
162
163
164
165
166
167
168
169
170
171

        if rank == 0:
            # kickstart pymultinest
            if not os.path.exists("chains"):
                os.mkdir("chains")
            pymultinest.run(self._multinest_likelihood,
                            self.prior,
                            len(self.active_variables),
                            verbose=True)
        else:
            # let all other nodes listen for likelihood evaluations
            self._listen_for_likelihood_calls()