hammurapy.py 10.1 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 tempfile
21
import subprocess
Theo Steininger's avatar
Theo Steininger committed
22
import xml.etree.ElementTree as et
23

24
import numpy as np
25

26
from d2o import distributed_data_object
27

28 29
from nifty import HPSpace

30
from imagine.observers.observer import Observer
31 32 33
from .observable_mixins import ObservableMixin
from .model_mixins import MagneticFieldModel
from imagine.observables import Observable
34 35


36
class Hammurapy(Observer):
37 38 39
    def __init__(self, hammurabi_executable, magnetic_field_model, observables,
                 input_directory='./input', working_directory_base='.',
                 nside=64):
40
        self.hammurabi_executable = os.path.abspath(hammurabi_executable)
41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57

        if not isinstance(magnetic_field_model, MagneticFieldModel):
            raise TypeError("magnetic_field_model must be an instance of the "
                            "MagneticField class.")
        self.magnetic_field_model = magnetic_field_model

        if not isinstance(observables, list):
            if isinstance(observables, tuple):
                observables = list(observables)
            else:
                observables = [observables]
        for obs in observables:
            if not isinstance(obs, ObservableMixin):
                raise TypeError("observables must be an instance of the "
                                "ObservableMixin class.")
        self.observables = observables

Theo Steininger's avatar
Theo Steininger committed
58
        self.input_directory = os.path.abspath(input_directory)
59
        self.working_directory_base = os.path.abspath(working_directory_base)
60

61
        self.nside = int(nside)
62
        self._hpSpace = HPSpace(nside=self.nside)
63

64 65
        self.last_call_log = ""

66
    @property
67
    def magnetic_field_class(self):
68
        return self.magnetic_field_model.magnetic_field_class
69 70

    def _make_temp_folder(self):
71
        prefix = os.path.join(self.working_directory_base, 'temp_hammurabi_')
72 73
        return tempfile.mkdtemp(prefix=prefix)

74 75 76 77 78 79 80 81 82 83 84
    def _remove_folder(self, path):
        # Try multiple times in order to overcome 'Directory not empty' errors
        # Hopefully open file handles get closed in the meantime
        n = 0
        while n < 10:
            temp_process = subprocess.Popen(['rm', '-rf', path],
                                            stderr=subprocess.PIPE)
            # wait until process is finished
            errlog = temp_process.communicate()[1]
            # check if there were some errors
            if errlog == '':
Theo Steininger's avatar
Theo Steininger committed
85
                self.logger.debug("Successfully removed temporary folder.")
86 87 88
                break
        else:
            self.logger.warning('Could not delete %s' % path)
89 90

    def _call_hammurabi(self, path):
Theo Steininger's avatar
Theo Steininger committed
91 92 93 94
        temp_process = subprocess.Popen(
                                [self.hammurabi_executable, 'parameters.xml'],
                                stdout=subprocess.PIPE,
                                cwd=path)
95 96
        self.last_call_log = temp_process.communicate()[0]

97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
    def _initialize_observable_dict(self, magnetic_field):
        observable_dict = {}
        for observable in self.observables:
            observable_dict = self._initialize_observable_dict_helper(
                    observable_dict, magnetic_field, observable)
        return observable_dict

    def _initialize_observable_dict_helper(self, observable_dict,
                                           magnetic_field, observable):
        ensemble_space = magnetic_field.domain[0]
        for component in observable.component_names:
            # It is important to initialize the Observables with an explicit
            # value. Otherwise the d2o will not instantaneuosly be created
            # (c.f. lazy object creation).
            observable_dict[component] = Observable(
                                    val=0,
                                    domain=(ensemble_space, self._hpSpace),
                                    distribution_strategy='equal')
        return observable_dict
116

117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137
    def _write_parameter_xml(self, magnetic_field, local_ensemble_index,
                             working_directory):
        # load the default xml
        try:
            default_parameters_xml = os.path.join(self.input_directory,
                                                  'default_parameters.xml')
            tree = et.parse(default_parameters_xml)
        except IOError:
            import imagine
            module_path = os.path.split(
                            imagine.observers.hammurapy.__file__)[0]
            default_parameters_xml = os.path.join(
                                module_path, 'input/default_parameters.xml')
            tree = et.parse(default_parameters_xml)

        root = tree.getroot()

        # modify the default_parameters.xml
        custom_parameters = [
                ['./Grid/Integration/shell/auto/nside_min', 'value',
                 self.nside],
138
                ['./Grid/Integration/nside_sim', 'value', self.nside],
139 140 141
#                ['./Interface/fe_grid', 'read', '1'],
#                ['./Interface/fe_grid', 'filename',
#                 os.path.join(self.input_directory, 'fe_grid.bin')]
142 143 144 145 146 147
                             ]

        # access the magnetic-field's random-eed d2o directly, since we
        # know that the distribution strategy is the same for the
        # randam samples and the magnetic field itself
        random_seed = magnetic_field.random_seed.data[local_ensemble_index]
Theo Steininger's avatar
Theo Steininger committed
148 149
        self.logger.debug("Local ensemble index %i uses random seed %i" %
                          (local_ensemble_index, random_seed))
150
        custom_parameters += [['./MagneticField/Random', 'seed',
151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
                               random_seed]]

        for key, value in magnetic_field.parameters.iteritems():
            desc = magnetic_field.descriptor_lookup[key]
            custom_parameters += [desc + [value]]

        # set up grid parameters
        grid_space = magnetic_field.domain[1]
        lx, ly, lz = \
            np.array(grid_space.shape)*np.array(grid_space.distances)/2.
        nx, ny, nz = grid_space.shape
        custom_parameters += [['./Grid/Box/x_min', 'value', -lx],
                              ['./Grid/Box/x_max', 'value', lx],
                              ['./Grid/Box/y_min', 'value', -ly],
                              ['./Grid/Box/y_max', 'value', ly],
                              ['./Grid/Box/z_min', 'value', -lz],
                              ['./Grid/Box/z_max', 'value', lz],
                              ['./Grid/Box/nx', 'value', nx],
                              ['./Grid/Box/ny', 'value', ny],
                              ['./Grid/Box/nz', 'value', nz]]

        for parameter in custom_parameters:
            root.find(parameter[0]).set(parameter[1], str(parameter[2]))

        self.magnetic_field_model.update_parameter_xml(root)

        for observable in self.observables:
            observable.update_parameter_xml(root)

        parameters_file_path = os.path.join(working_directory,
                                            'parameters.xml')
        tree.write(parameters_file_path)

184
    def _fill_observable_dict(self, observable_dict, working_directory,
185 186 187 188 189 190 191
                              local_ensemble_index):
        for observable in self.observables:
            observable.fill_observable_dict(
                            observable_dict=observable_dict,
                            working_directory=working_directory,
                            local_ensemble_index=local_ensemble_index,
                            nside=self.nside)
192 193 194 195
        return observable_dict

    def __call__(self, magnetic_field):

196 197 198 199
        if not isinstance(magnetic_field, self.magnetic_field_class):
            raise ValueError("Given magnetic field is not a subclass of" +
                             " %s" % str(self.magnetic_field_class))

200 201
        observable_dict = self._initialize_observable_dict(
                                            magnetic_field=magnetic_field)
202 203 204 205 206

        # iterate over ensemble and put result into result_observable
        # get the local shape by creating a dummy d2o
        ensemble_number = magnetic_field.shape[0]
        dummy = distributed_data_object(global_shape=(ensemble_number,),
207 208
                                        distribution_strategy='equal',
                                        dtype=np.float)
209 210 211

        local_length = dummy.distributor.local_length
        for local_ensemble_index in xrange(local_length):
Theo Steininger's avatar
Theo Steininger committed
212 213
            self.logger.debug("Processing local_ensemble_index %i." %
                              local_ensemble_index)
214 215 216
            # create a temporary folder
            working_directory = self._make_temp_folder()

217
            self._write_parameter_xml(
218
                                    magnetic_field=magnetic_field,
219 220
                                    local_ensemble_index=local_ensemble_index,
                                    working_directory=working_directory)
221

222
            # call hammurabi
223
            self._call_hammurabi(working_directory)
224 225 226 227 228 229 230 231 232 233 234

            # if hammurabi failed, _fill_observable_dict will fail
            try:
                self._fill_observable_dict(observable_dict,
                                           working_directory,
                                           local_ensemble_index)
            except:
                self.logger.critical("Hammurabi failed! Last call log:\n" +
                                     self.last_call_log)
                raise
            finally:
235
                self._remove_folder(working_directory)
236 237

        return observable_dict