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

3 4 5
import os
import numpy as np

6 7
from scipy import optimize

8 9
from mpi4py import MPI

10 11
from keepers import Loggable

12 13 14 15
from imagine.likelihoods import Likelihood
from imagine.magnetic_fields import MagneticFieldFactory
from imagine.observers import Observer
from imagine.priors import Prior
16
from imagine import pymultinest
17
from imagine.sample import Sample
18 19 20 21

comm = MPI.COMM_WORLD
size = comm.size
rank = comm.rank
22

23 24 25
WORK_TAG = 0
DIE_TAG = 1

26 27

class Pipeline(Loggable, object):
28 29 30 31 32 33 34 35 36 37 38 39
    """
    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,
40
                 active_variables=[], ensemble_size=1,
41
                 pymultinest_parameters={}, sample_callback=None):
42
        self.logger.debug("Setting up pipeline.")
43
        self.magnetic_field_factory = magnetic_field_factory
44 45 46
        self.observer = observer
        self.likelihood = likelihood
        self.prior = prior
47
        self.active_variables = active_variables
48
        self.ensemble_size = ensemble_size
49
        self.likelihood_rescaler = 1.
50

51
        # setting defaults for pymultinest
52 53 54 55 56
        self.pymultinest_parameters = {'verbose': True,
                                       'n_iter_before_update': 1,
                                       'n_live_points': 100}
        self.pymultinest_parameters.update(pymultinest_parameters)

57 58
        self.sample_callback = sample_callback

59
        self.fixed_random_seed = None
Theo Steininger's avatar
fix  
Theo Steininger committed
60
        self.likelihood_threshold = 0.
61
        self.check_threshold = False
62

63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79
    @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):
80
        self.logger.debug("Setting likelihood.")
81
        self._likelihood = ()
Theo Steininger's avatar
Theo Steininger committed
82
        if not (isinstance(likelihood, list) or
83 84
                isinstance(likelihood, tuple)):
            likelihood = [likelihood]
85 86 87 88 89
        for l in likelihood:
            if not isinstance(l, Likelihood):
                raise TypeError(
                    "likelihood must be an instance of Likelihood-class.")
            self._likelihood += (l,)
90 91 92 93 94 95 96

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

    @prior.setter
    def prior(self, prior):
97
        self.logger.debug("Setting prior.")
98 99 100 101
        if not isinstance(prior, Prior):
            raise TypeError(
                "prior must be an instance of Prior-class.")
        self._prior = prior
102 103

    @property
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 129 130
    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
131 132 133 134 135 136 137 138 139 140 141 142 143 144

    @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

145 146 147 148
    def _multinest_likelihood(self, cube, ndim, nparams):
        cube_content = np.empty(ndim)
        for i in xrange(ndim):
            cube_content[i] = cube[i]
149 150 151 152

        # heuristic for minimizers:
        # if a parameter value from outside of the cube is requested, return
        # the worst possible likelihood value
153
        if np.any(cube_content > 1.) or np.any(cube_content < 0.):
154
            self.logger.info('Cube %s requested. Returned most negative '
Theo Steininger's avatar
Theo Steininger committed
155
                             'possible number.' % cube_content)
156 157
            return np.nan_to_num(-np.inf)

158 159 160
        if rank != 0:
            raise RuntimeError("_multinest_likelihood must only be called on "
                               "rank==0.")
Theo Steininger's avatar
Theo Steininger committed
161 162
        error_count = 0
        while error_count < 5:
Theo Steininger's avatar
Theo Steininger committed
163 164 165
            for i in xrange(1, size):
                comm.send(cube_content, dest=i, tag=WORK_TAG)
            self.logger.debug("Sent multinest-cube to nodes with rank > 0.")
Theo Steininger's avatar
Theo Steininger committed
166
            likelihood = self._core_likelihood(cube_content)
167 168
            if likelihood < self.likelihood_threshold or \
               not self.check_threshold:
Theo Steininger's avatar
Theo Steininger committed
169 170 171 172 173
                break
            else:
                self.logger.error("Positive log-likelihood value encountered!"
                                  "Redoing calculation.")
        return likelihood
174 175

    def _listen_for_likelihood_calls(self):
176 177 178 179 180 181 182 183
        status = MPI.Status()
        while True:
            cube = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
            if status == DIE_TAG:
                self.logger.debug("Received DIE_TAG from rank 0.")
                break
            self.logger.debug("Received cube from rank 0.")
            self._core_likelihood(cube)
184 185

    def _core_likelihood(self, cube):
186 187
        self.logger.debug("Beginning Likelihood-calculation for %s." %
                          str(cube))
188 189 190 191 192 193
        # translate cube to variables
        variables = {}
        for i, av in enumerate(self.active_variables):
            variables[av] = cube[i]

        # create magnetic field
194 195
        self.logger.debug("Creating magnetic field.")
        b_field = self.magnetic_field_factory.generate(
196 197 198
                                          variables=variables,
                                          ensemble_size=self.ensemble_size,
                                          random_seed=self.fixed_random_seed)
199

200
        # create observables
201
        self.logger.debug("Creating observables.")
202 203 204
        observables = self.observer(b_field)

        # add up individual log-likelihood terms
205
        self.logger.debug("Evaluating likelihood(s).")
206 207
        likelihood = ()
        total_likelihood = 0
208
        for like in self.likelihood:
209 210 211
            current_likelihood = like(observables)
            likelihood += (current_likelihood, )
            total_likelihood += current_likelihood
212

Theo Steininger's avatar
Theo Steininger committed
213
        self.logger.info("Evaluated likelihood: %f for %s" %
214 215 216 217 218 219 220 221 222 223 224
                         (total_likelihood, str(cube)))

        if self.sample_callback is not None:
            self.logger.debug("Creating sample-object.")
            sample = Sample(variables=variables,
                            magnetic_field=b_field,
                            observables=observables,
                            likelihood=likelihood,
                            total_likelihood=total_likelihood)
            self.sample_callback(sample)

225
        return total_likelihood * self.likelihood_rescaler
226

227
    def __call__(self):
228 229 230

        if rank == 0:
            # kickstart pymultinest
Theo Steininger's avatar
Theo Steininger committed
231
            self.logger.info("Starting pymultinest.")
232 233 234 235 236
            if not os.path.exists("chains"):
                os.mkdir("chains")
            pymultinest.run(self._multinest_likelihood,
                            self.prior,
                            len(self.active_variables),
237
                            **self.pymultinest_parameters)
238 239 240 241
            self.logger.info("pymultinest finished.")
            for i in xrange(1, size):
                self.logger.debug("Sending DIE_TAG to rank %i." % i)
                comm.send(None, dest=i, tag=DIE_TAG)
242 243 244
        else:
            # let all other nodes listen for likelihood evaluations
            self._listen_for_likelihood_calls()
245

246 247
    def find_minimum(self, starting_guess=None, method='Nelder-Mead',
                     **kwargs):
248 249 250 251 252 253
        if starting_guess is None:
            starting_guess = np.zeros(len(self.active_variables)) + 0.5

        if rank == 0:
            # kickstart pymultinest
            self.logger.info("Starting minimizer.")
Theo Steininger's avatar
Theo Steininger committed
254
            call_func = lambda z: -self._multinest_likelihood(
255 256 257
                                                 z,
                                                 len(self.active_variables),
                                                 len(self.active_variables))
Theo Steininger's avatar
Theo Steininger committed
258
            minimum = optimize.minimize(fun=call_func,
259 260 261
                                        x0=starting_guess,
                                        method=method,
                                        **kwargs)
262 263 264 265 266 267 268 269 270 271
            self.logger.info("Minimizer finished.")
            for i in xrange(1, size):
                self.logger.debug("Sending DIE_TAG to rank %i." % i)
                comm.send(None, dest=i, tag=DIE_TAG)
        else:
            minimum = None
            # let all other nodes listen for likelihood evaluations
            self._listen_for_likelihood_calls()
        minimum = comm.bcast(minimum, root=0)
        return minimum