pipeline.py 10.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2017 Max-Planck-Society
#
# IMAGINE is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.
18

19 20
import os
import numpy as np
21
from scipy import optimize
22
from mpi4py import MPI
23 24
from keepers import Loggable

25 26 27 28
from imagine.likelihoods import Likelihood
from imagine.magnetic_fields import MagneticFieldFactory
from imagine.observers import Observer
from imagine.priors import Prior
29
from imagine import pymultinest
30
from imagine.sample import Sample
31 32 33 34

comm = MPI.COMM_WORLD
size = comm.size
rank = comm.rank
35

36 37 38
WORK_TAG = 0
DIE_TAG = 1

39
class Pipeline(Loggable, object):
40 41 42 43 44 45 46 47 48 49 50 51
    """
    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,
52
                 active_variables=[], ensemble_size=1,
53
                 pymultinest_parameters={}, sample_callback=None):
54
        self.logger.debug("Setting up pipeline.")
55
        self.magnetic_field_factory = magnetic_field_factory
56 57 58
        self.observer = observer
        self.likelihood = likelihood
        self.prior = prior
59
        self.active_variables = active_variables
60
        self.ensemble_size = ensemble_size
Jiaxin Wang's avatar
Jiaxin Wang committed
61
        # rescaling total likelihood in _core_likelihood
62
        self.likelihood_rescaler = 1.
63
        # setting defaults for pymultinest
Jiaxin Wang's avatar
Jiaxin Wang committed
64 65 66
        self.pymultinest_parameters = {"verbose": True,
                                       "n_iter_before_update": 1,
                                       "n_live_points": 100}
67
        self.pymultinest_parameters.update(pymultinest_parameters)
68
        self.sample_callback = sample_callback
Jiaxin Wang's avatar
Jiaxin Wang committed
69
        # seed for generating random field
70
        self.fixed_random_seed = None
Jiaxin Wang's avatar
Jiaxin Wang committed
71
        # checking likelihood threshold in _multinest_likelihood
72
        self.check_threshold = False
Jiaxin Wang's avatar
Jiaxin Wang committed
73
        self.likelihood_threshold = 0.
74

75 76 77 78 79 80
    @property
    def observer(self):
        return self._observer

    @observer.setter
    def observer(self, observer):
Jiaxin Wang's avatar
Jiaxin Wang committed
81
        if not isinstance(observer, Observer):#{
82
            raise TypeError("observer must be an instance of Observer-class.")
Jiaxin Wang's avatar
Jiaxin Wang committed
83
        #}
84 85 86 87 88 89 90 91 92
        self.logger.debug("Setting observer.")
        self._observer = observer

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

    @likelihood.setter
    def likelihood(self, likelihood):
93
        self.logger.debug("Setting likelihood.")
94
        self._likelihood = ()
Jiaxin Wang's avatar
Jiaxin Wang committed
95
        if not (isinstance(likelihood, list) or isinstance(likelihood, tuple)):#{
96
            likelihood = [likelihood]
Jiaxin Wang's avatar
Jiaxin Wang committed
97 98 99 100 101
        #}
        for l in likelihood:#{
            if not isinstance(l, Likelihood):#{
                raise TypeError("likelihood must be an instance of Likelihood-class.")
            #}
102
            self._likelihood += (l,)
Jiaxin Wang's avatar
Jiaxin Wang committed
103
        #}
104 105 106 107 108 109 110

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

    @prior.setter
    def prior(self, prior):
111
        self.logger.debug("Setting prior.")
Jiaxin Wang's avatar
Jiaxin Wang committed
112 113 114
        if not isinstance(prior, Prior):#{
            raise TypeError("prior must be an instance of Prior-class.")
        #}
115
        self._prior = prior
116 117

    @property
118 119 120 121 122
    def magnetic_field_factory(self):
        return self._magnetic_field_factory

    @magnetic_field_factory.setter
    def magnetic_field_factory(self, magnetic_field_factory):
Jiaxin Wang's avatar
Jiaxin Wang committed
123 124 125
        if not isinstance(magnetic_field_factory, MagneticFieldFactory):#{
            raise TypeError("magnetic_field_factory must be an instance of the MagneticFieldFactory-class.")
        #}
126 127 128 129 130 131 132 133 134
        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):
Jiaxin Wang's avatar
Jiaxin Wang committed
135 136 137 138
        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))
139
        new_active = []
Jiaxin Wang's avatar
Jiaxin Wang committed
140
        for av in active_variables:#{
141
            new_active += [str(av)]
Jiaxin Wang's avatar
Jiaxin Wang committed
142
        #}
143
        self._active_variables = new_active
144 145 146 147 148 149 150 151

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

    @ensemble_size.setter
    def ensemble_size(self, ensemble_size):
        ensemble_size = int(ensemble_size)
Jiaxin Wang's avatar
Jiaxin Wang committed
152
        if ensemble_size <= 0:#{
153
            raise ValueError("ensemble_size must be positive!")
Jiaxin Wang's avatar
Jiaxin Wang committed
154
        #}
155 156 157
        self.logger.debug("Setting ensemble size to %i." % ensemble_size)
        self._ensemble_size = ensemble_size

158 159 160 161
    def _multinest_likelihood(self, cube, ndim, nparams):
        cube_content = np.empty(ndim)
        for i in xrange(ndim):
            cube_content[i] = cube[i]
162 163 164
        # heuristic for minimizers:
        # if a parameter value from outside of the cube is requested, return
        # the worst possible likelihood value
Jiaxin Wang's avatar
Jiaxin Wang committed
165 166
        if np.any(cube_content > 1.) or np.any(cube_content < 0.):#{
            self.logger.info("Cube %s requested. Returned most negative possible number." % cube_content)
167
            return np.nan_to_num(-np.inf)
Jiaxin Wang's avatar
Jiaxin Wang committed
168 169 170 171
        #}
        if rank != 0:#{
            raise RuntimeError("_multinest_likelihood must only be called on rank==0.")
        #}
Theo Steininger's avatar
Theo Steininger committed
172 173
        error_count = 0
        while error_count < 5:
Jiaxin Wang's avatar
Jiaxin Wang committed
174
            for i in xrange(1, size):#{
Theo Steininger's avatar
Theo Steininger committed
175
                comm.send(cube_content, dest=i, tag=WORK_TAG)
Jiaxin Wang's avatar
Jiaxin Wang committed
176
            #}
Theo Steininger's avatar
Theo Steininger committed
177
            self.logger.debug("Sent multinest-cube to nodes with rank > 0.")
Theo Steininger's avatar
Theo Steininger committed
178
            likelihood = self._core_likelihood(cube_content)
Jiaxin Wang's avatar
Jiaxin Wang committed
179
            if likelihood < self.likelihood_threshold or not self.check_threshold:
Theo Steininger's avatar
Theo Steininger committed
180 181
                break
            else:
Jiaxin Wang's avatar
Jiaxin Wang committed
182 183
                self.logger.error("Positive log-likelihood value encountered! Redoing calculation.")
        #}
Theo Steininger's avatar
Theo Steininger committed
184
        return likelihood
185

Jiaxin Wang's avatar
Jiaxin Wang committed
186
    # on slave nodes
187
    def _listen_for_likelihood_calls(self):
188
        status = MPI.Status()
Jiaxin Wang's avatar
Jiaxin Wang committed
189
        while True:#{
190
            cube = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
Jiaxin Wang's avatar
Jiaxin Wang committed
191
            if status == DIE_TAG:#{
192 193
                self.logger.debug("Received DIE_TAG from rank 0.")
                break
Jiaxin Wang's avatar
Jiaxin Wang committed
194
            #}
195 196
            self.logger.debug("Received cube from rank 0.")
            self._core_likelihood(cube)
Jiaxin Wang's avatar
Jiaxin Wang committed
197
        #}
198 199

    def _core_likelihood(self, cube):
Jiaxin Wang's avatar
Jiaxin Wang committed
200
        self.logger.debug("Beginning Likelihood-calculation for %s." % str(cube))
201 202
        # translate cube to variables
        variables = {}
Jiaxin Wang's avatar
Jiaxin Wang committed
203
        for i, av in enumerate(self.active_variables):#{
204
            variables[av] = cube[i]
Jiaxin Wang's avatar
Jiaxin Wang committed
205
        #}
206
        # create magnetic field
207 208
        self.logger.debug("Creating magnetic field.")
        b_field = self.magnetic_field_factory.generate(
209 210 211
                                          variables=variables,
                                          ensemble_size=self.ensemble_size,
                                          random_seed=self.fixed_random_seed)
212
        # create observables
213
        self.logger.debug("Creating observables.")
214 215
        observables = self.observer(b_field)
        # add up individual log-likelihood terms
216
        self.logger.debug("Evaluating likelihood(s).")
217 218
        likelihood = ()
        total_likelihood = 0
Jiaxin Wang's avatar
Jiaxin Wang committed
219
        for like in self.likelihood:#{
220 221 222
            current_likelihood = like(observables)
            likelihood += (current_likelihood, )
            total_likelihood += current_likelihood
Jiaxin Wang's avatar
Jiaxin Wang committed
223 224
        #}
        self.logger.info("Evaluated likelihood: %f for %s" % (total_likelihood, str(cube)))
225

Jiaxin Wang's avatar
Jiaxin Wang committed
226
        if self.sample_callback is not None:#{
227 228 229 230 231 232 233
            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)
Jiaxin Wang's avatar
Jiaxin Wang committed
234
        #}
235
        return total_likelihood * self.likelihood_rescaler
236

237
    def __call__(self):
Jiaxin Wang's avatar
Jiaxin Wang committed
238
        if rank == 0:#{
239
            # kickstart pymultinest
Theo Steininger's avatar
Theo Steininger committed
240
            self.logger.info("Starting pymultinest.")
Jiaxin Wang's avatar
Jiaxin Wang committed
241
            if not os.path.exists("chains"):#{
242
                os.mkdir("chains")
Jiaxin Wang's avatar
Jiaxin Wang committed
243 244
            #}
            # run pymultinest
245 246 247
            pymultinest.run(self._multinest_likelihood,
                            self.prior,
                            len(self.active_variables),
248
                            **self.pymultinest_parameters)
249
            self.logger.info("pymultinest finished.")
Jiaxin Wang's avatar
Jiaxin Wang committed
250
            for i in xrange(1, size):#{
251 252
                self.logger.debug("Sending DIE_TAG to rank %i." % i)
                comm.send(None, dest=i, tag=DIE_TAG)
Jiaxin Wang's avatar
Jiaxin Wang committed
253 254
            #}
        #}
255 256 257
        else:
            # let all other nodes listen for likelihood evaluations
            self._listen_for_likelihood_calls()
258

Jiaxin Wang's avatar
Jiaxin Wang committed
259 260
    def find_minimum(self, starting_guess=None, method="Nelder-Mead", **kwargs):
        if starting_guess is None:#{
261
            starting_guess = np.zeros(len(self.active_variables)) + 0.5
Jiaxin Wang's avatar
Jiaxin Wang committed
262 263
        #}
        if rank == 0:#{
264 265
            # kickstart pymultinest
            self.logger.info("Starting minimizer.")
Theo Steininger's avatar
Theo Steininger committed
266
            call_func = lambda z: -self._multinest_likelihood(
267 268 269
                                                 z,
                                                 len(self.active_variables),
                                                 len(self.active_variables))
Theo Steininger's avatar
Theo Steininger committed
270
            minimum = optimize.minimize(fun=call_func,
271 272 273
                                        x0=starting_guess,
                                        method=method,
                                        **kwargs)
274
            self.logger.info("Minimizer finished.")
Jiaxin Wang's avatar
Jiaxin Wang committed
275
            for i in xrange(1, size):#{
276 277
                self.logger.debug("Sending DIE_TAG to rank %i." % i)
                comm.send(None, dest=i, tag=DIE_TAG)
Jiaxin Wang's avatar
Jiaxin Wang committed
278 279 280
            #}
        #}
        else:#{
281 282 283
            minimum = None
            # let all other nodes listen for likelihood evaluations
            self._listen_for_likelihood_calls()
Jiaxin Wang's avatar
Jiaxin Wang committed
284
        #}
285 286
        minimum = comm.bcast(minimum, root=0)
        return minimum