hammurapy.py 11.6 KB
Newer Older
1
2
3
4
# -*- coding: utf-8 -*-

import os
import tempfile
5
import subprocess
Theo Steininger's avatar
Theo Steininger committed
6
import xml.etree.ElementTree as et
7

8
import numpy as np
9

10
from d2o import distributed_data_object
11

12
13
from nifty import HPSpace

14
from imagine.observers.observer import Observer
15
16
17
from .observable_mixins import ObservableMixin
from .model_mixins import MagneticFieldModel
from imagine.observables import Observable
18
19


20
class Hammurapy(Observer):
21
22
23
    def __init__(self, hammurabi_executable, magnetic_field_model, observables,
                 input_directory='./input', working_directory_base='.',
                 nside=64):
24
        self.hammurabi_executable = os.path.abspath(hammurabi_executable)
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

        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
42
        self.input_directory = os.path.abspath(input_directory)
43
        self.working_directory_base = os.path.abspath(working_directory_base)
44

45
        self.nside = int(nside)
46
        self._hpSpace = HPSpace(nside=self.nside)
47

48
49
        self.last_call_log = ""

50
    @property
51
    def magnetic_field_class(self):
52
        return self.magnetic_field_model.magnetic_field_class
53
54

    def _make_temp_folder(self):
55
        prefix = os.path.join(self.working_directory_base, 'temp_hammurabi_')
56
57
        return tempfile.mkdtemp(prefix=prefix)

58
59
60
61
62
63
64
65
66
67
68
    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
69
                self.logger.debug("Successfully removed temporary folder.")
70
71
72
                break
        else:
            self.logger.warning('Could not delete %s' % path)
73
74

    def _call_hammurabi(self, path):
Theo Steininger's avatar
Theo Steininger committed
75
76
77
78
        temp_process = subprocess.Popen(
                                [self.hammurabi_executable, 'parameters.xml'],
                                stdout=subprocess.PIPE,
                                cwd=path)
79
80
        self.last_call_log = temp_process.communicate()[0]

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    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
100

101
    def _build_parameter_dict(self, parameter_dict, magnetic_field,
Theo Steininger's avatar
Theo Steininger committed
102
103
                              local_ensemble_index):

104
105
        parameter_dict.update(self.magnetic_field_model.parameter_dict)

Theo Steininger's avatar
Theo Steininger committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
        parameter_dict.update(
                {('./Interface/fe_grid', 'read'): '1',
                 ('./Interface/fe_grid', 'filename'):
                     os.path.join(self.input_directory, 'fe_grid.bin'),
                 })
        # access the magnetic-field's random-seed 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]
        parameter_dict.update(
                {('./Galaxy/MagneticField/Random', 'seed'): random_seed})

        for key, value in magnetic_field.parameters.iteritems():
            large_key = magnetic_field.descriptor_lookup[key]
            parameter_dict[large_key] = value

122
123
124
        grid_space = magnetic_field.domain[1]
        lx, ly, lz = np.array(grid_space.shape)*np.array(grid_space.distances)
        nx, ny, nz = grid_space.shape
Theo Steininger's avatar
Theo Steininger committed
125
126
127
128
129
130
131
132
133

        parameter_dict.update({('./Grid/Box/lx', 'value'): lx,
                               ('./Grid/Box/ly', 'value'): ly,
                               ('./Grid/Box/lz', 'value'): lz,
                               ('./Grid/Box/nx', 'value'): nx,
                               ('./Grid/Box/ny', 'value'): ny,
                               ('./Grid/Box/nz', 'value'): nz})
        parameter_dict.update(
                {('./Grid/Integration/nside', 'value'): self.nside})
134

135
    def _write_parameter_file(self, working_directory):
Theo Steininger's avatar
Theo Steininger committed
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
        # 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()
        for key, value in parameter_dict.iteritems():
            root.find(key[0]).set(key[1], str(value))
152
153

        parameters_file_path = os.path.join(working_directory,
Theo Steininger's avatar
Theo Steininger committed
154
155
                                            'parameters.xml')
        tree.write(parameters_file_path)
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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
    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],
                ['./Interface/fe_grid', 'read', '1'],
                ['./Interface/fe_grid', 'filename',
                 os.path.join(self.input_directory, 'fe_grid.bin')]
                             ]

        # 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]
        custom_parameters += [['./Galaxy/MagneticField/Random', 'seed',
                               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)

221
    def _fill_observable_dict(self, observable_dict, working_directory,
222
223
224
225
226
227
228
                              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)
229
230
231
232
        return observable_dict

    def __call__(self, magnetic_field):

233
234
235
236
        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))

237
238
        observable_dict = self._initialize_observable_dict(
                                            magnetic_field=magnetic_field)
239
240
241
242
243

        # 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,),
244
245
                                        distribution_strategy='equal',
                                        dtype=np.float)
246
247
248

        local_length = dummy.distributor.local_length
        for local_ensemble_index in xrange(local_length):
Theo Steininger's avatar
Theo Steininger committed
249
250
            self.logger.debug("Processing local_ensemble_index %i." %
                              local_ensemble_index)
251
252
253
            # create a temporary folder
            working_directory = self._make_temp_folder()

254
            self._write_parameter_xml(
255
                                    magnetic_field=magnetic_field,
256
257
                                    local_ensemble_index=local_ensemble_index,
                                    working_directory=working_directory)
258

259
            # call hammurabi
260
            self._call_hammurabi(working_directory)
261
262
263
264
265
266
267
268
269
270
271

            # 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:
272
273
                pass
                #self._remove_folder(working_directory)
274
275

        return observable_dict