PageRenderTime 1ms CodeModel.GetById 31ms app.highlight 12ms RepoModel.GetById 0ms app.codeStats 0ms

/Framework/PythonInterface/plugins/algorithms/LoadNMoldyn4Ascii.py

https://github.com/wdzhou/mantid
Python | 309 lines | 210 code | 49 blank | 50 comment | 36 complexity | 9178e889ba64e552f075475538f171c5 MD5 | raw file
  1#pylint: disable=no-init
  2
  3from __future__ import (absolute_import, division, print_function)
  4from mantid.simpleapi import *
  5from mantid.kernel import *
  6from mantid.api import *
  7
  8import numpy as np
  9import scipy.constants as sc
 10import ast
 11import fnmatch
 12import re
 13import os
 14
 15#------------------------------------------------------------------------------
 16
 17VARIABLE_REGEX = re.compile(r'#\s+variable name:\s+(.*)')
 18TYPE_REGEX = re.compile(r'#\s+type:\s+([A-z]+)')
 19AXIS_REGEX = re.compile(r'#\s+axis:\s+([A-z]+)\|([A-z]+)')
 20UNIT_REGEX = re.compile(r'#\s+units:\s+(.*)')
 21
 22SLICE_1D_HEADER_REGEX = re.compile(r'#slice:\[([0-9]+)[A-z]*\]')
 23SLICE_2D_HEADER_REGEX = re.compile(r'#slice:\[([0-9]+)[A-z]*,\s+([0-9]+)[A-z]*\]')
 24
 25#------------------------------------------------------------------------------
 26
 27
 28class LoadNMoldyn4Ascii(PythonAlgorithm):
 29
 30    _axis_cache = None
 31    _data_directory = None
 32
 33#------------------------------------------------------------------------------
 34
 35    def category(self):
 36        return 'Inelastic\\DataHandling;Simulation'
 37
 38    def summary(self):
 39        return 'Imports functions from .dat files output by nMOLDYN 4.'
 40
 41#------------------------------------------------------------------------------
 42
 43    def PyInit(self):
 44        self.declareProperty(FileProperty('Directory', '',
 45                                          action=FileAction.Directory),
 46                             doc='Path to directory containg .dat files')
 47
 48        self.declareProperty(StringArrayProperty('Functions'),
 49                             doc='Names of functions to attempt to load from file')
 50
 51        self.declareProperty(WorkspaceProperty('OutputWorkspace', '',
 52                                               direction=Direction.Output),
 53                             doc='Output workspace name')
 54
 55#------------------------------------------------------------------------------
 56
 57    def validateInputs(self):
 58        issues = dict()
 59
 60        if len(self.getProperty('Functions').value) == 0:
 61            issues['Functions'] = 'Must specify at least one function to load'
 62
 63        return issues
 64
 65#------------------------------------------------------------------------------
 66
 67    def PyExec(self):
 68        self._axis_cache = {}
 69        self._data_directory = self.getPropertyValue('Directory')
 70
 71        # Convert the simplified function names to the actual file names
 72        data_directory_files = [os.path.splitext(f)[0] for f in fnmatch.filter(os.listdir(self._data_directory), '*.dat')]
 73        logger.debug('All data files: {0}'.format(data_directory_files))
 74        functions_input = [x.strip().replace(',', '') for x in self.getProperty('Functions').value]
 75        functions = [f for f in data_directory_files if f.replace(',', '') in functions_input]
 76        logger.debug('Functions to load: {0}'.format(functions))
 77
 78        loaded_function_workspaces = []
 79        for func_name in functions:
 80            try:
 81                # Load the intensity data
 82                function = self._load_function(func_name)
 83                # Load (or retrieve) the axis data
 84                v_axis = self._load_axis(function[2][0])
 85                x_axis = self._load_axis(function[2][1])
 86
 87                # Perform axis unit conversions
 88                v_axis = self._axis_conversion(*v_axis)
 89                x_axis = self._axis_conversion(*x_axis)
 90
 91                # Create the workspace for function
 92                create_workspace = AlgorithmManager.Instance().create('CreateWorkspace')
 93                create_workspace.initialize()
 94                create_workspace.setLogging(False)
 95                create_workspace.setProperty('OutputWorkspace', func_name)
 96                create_workspace.setProperty('DataX', x_axis[0])
 97                create_workspace.setProperty('DataY', function[0])
 98                create_workspace.setProperty('NSpec', v_axis[0].size)
 99                create_workspace.setProperty('UnitX', x_axis[1])
100                create_workspace.setProperty('YUnitLabel', function[1])
101                create_workspace.setProperty('VerticalAxisValues', v_axis[0])
102                create_workspace.setProperty('VerticalAxisUnit', v_axis[1])
103                create_workspace.setProperty('WorkspaceTitle', func_name)
104                create_workspace.execute()
105
106                loaded_function_workspaces.append(func_name)
107
108            except ValueError as rerr:
109                logger.warning('Failed to load function {0}. Error was: {1}'.format(func_name, str(rerr)))
110
111        # Process the loaded workspaces
112        out_ws_name = self.getPropertyValue('OutputWorkspace')
113        if len(loaded_function_workspaces) == 0:
114            raise RuntimeError('Failed to load any functions for data')
115        GroupWorkspaces(InputWorkspaces=loaded_function_workspaces,
116                        OutputWorkspace=out_ws_name)
117
118        # Set the output workspace
119        self.setProperty('OutputWorkspace', out_ws_name)
120
121#------------------------------------------------------------------------------
122
123    def _load_function(self, function_name):
124        """
125        Loads a function from the data directory.
126
127        @param function_name Name of the function to load
128        @return Tuple of (Numpy array of data, unit, (v axis name, x axis name))
129        @exception ValueError If function is not found
130        """
131        function_filename = os.path.join(self._data_directory, '{0}.dat'.format(function_name))
132        if not os.path.isfile(function_filename):
133            raise ValueError('File for function "{0}" not found'.format(function_name))
134
135        data = None
136        axis = (None, None)
137        unit = None
138
139        with open(function_filename, 'rU') as f_handle:
140            while True:
141                line = f_handle.readline()
142                if not line:
143                    break
144
145                # Ignore empty lines
146                if len(line[0]) == 0:
147                    pass
148
149                # Parse header lines
150                elif line[0] == '#':
151                    variable_match = VARIABLE_REGEX.match(line)
152                    if variable_match and variable_match.group(1) != function_name:
153                        raise ValueError('Function name differs from file name')
154
155                    axis_match = AXIS_REGEX.match(line)
156                    if axis_match:
157                        axis = (axis_match.group(1), axis_match.group(2))
158
159                    unit_match = UNIT_REGEX.match(line)
160                    if unit_match:
161                        unit = unit_match.group(1)
162
163                    slice_match = SLICE_2D_HEADER_REGEX.match(line)
164                    if slice_match:
165                        dimensions = (int(slice_match.group(1)), int(slice_match.group(2)))
166                        # Now parse the data
167                        data = self._load_2d_slice(f_handle, dimensions)
168
169        return (data, unit, axis)
170
171#------------------------------------------------------------------------------
172
173    def _load_axis(self, axis_name):
174        """
175        Loads an axis by name from the data directory.
176
177        @param axis_name Name of axis to load
178        @return Tuple of (Numpy array of data, unit, name)
179        @exception ValueError If axis is not found
180        """
181        if axis_name in self._axis_cache:
182            return self._axis_cache[axis_name]
183
184        axis_filename = os.path.join(self._data_directory, '{0}.dat'.format(axis_name))
185        if not os.path.isfile(axis_filename):
186            raise ValueError('File for axis "{0}" not found'.format(axis_name))
187
188        data = None
189        unit = None
190
191        with open(axis_filename, 'rU') as f_handle:
192            while True:
193                line = f_handle.readline()
194                if not line:
195                    break
196
197                # Ignore empty lines
198                if len(line[0]) == 0:
199                    pass
200
201                # Parse header lines
202                elif line[0] == '#':
203                    variable_match = VARIABLE_REGEX.match(line)
204                    if variable_match and variable_match.group(1) != axis_name:
205                        raise ValueError('Axis name differs from file name')
206
207                    unit_match = UNIT_REGEX.match(line)
208                    if unit_match:
209                        unit = unit_match.group(1)
210
211                    slice_match = SLICE_1D_HEADER_REGEX.match(line)
212                    if slice_match:
213                        length = int(slice_match.group(1))
214                        # Now parse the data
215                        data = self._load_1d_slice(f_handle, length)
216
217        return (data, unit, axis_name)
218
219#------------------------------------------------------------------------------
220
221    def _load_1d_slice(self, f_handle, length):
222        """
223        Loads a 1D slice from the open file.
224
225        @param f_handle Handle to the open file with the iterator at the slice header
226        @param length Length of data
227        @return Numpy array of length [length]
228        """
229        data = np.ndarray(shape=(length), dtype=float)
230
231        for idx in range(length):
232            line = f_handle.readline()
233
234            # End of file or empty line (either way end of data)
235            if not line or len(line) == 0:
236                break
237
238            data[idx] = ast.literal_eval(line)
239
240        return data
241
242#------------------------------------------------------------------------------
243
244    def _load_2d_slice(self, f_handle, dimensions):
245        """
246        Loads a 2D slice from the open file.
247
248        @param f_handle Handle to the open file with the iterator at the slice header
249        @param dimensions Tuple containing dimensions (rows/vertical axis, cols/x axis)
250        @return Numpy array of shape [dimensions]
251        """
252        data = np.ndarray(shape=dimensions, dtype=float)
253
254        for v_idx in range(dimensions[0]):
255            line = f_handle.readline()
256
257            # End of file or empty line (either way end of data)
258            if not line or len(line) == 0:
259                break
260
261            values = [ast.literal_eval(s) for s in line.split()]
262            data[v_idx] = np.array(values)
263
264        return data
265
266#------------------------------------------------------------------------------
267
268    def _axis_conversion(self, data, unit, name):
269        """
270        Converts an axis to a Mantid axis type (possibly performing a unit
271        conversion).
272
273        @param data The axis data as Numpy array
274        @param unit The axis unit as read from the file
275        @param name The axis name as read from the file
276        @return Tuple containing updated axis details
277        """
278        logger.debug('Axis for conversion: name={0}, unit={1}'.format(name, unit))
279
280        # Q (nm**-1) to Q (Angstrom**-1)
281        if name.lower() == 'q' and unit.lower() == 'inv_nm':
282            logger.information('Axis {0} will be converted to Q in Angstrom**-1'.format(name))
283            unit = 'MomentumTransfer'
284            data /= sc.nano # nm to m
285            data *= sc.angstrom # m to Angstrom
286
287        # Frequency (THz) to Energy (meV)
288        elif name.lower() == 'frequency' and unit.lower() == 'thz':
289            logger.information('Axis {0} will be converted to energy in meV'.format(name))
290            unit = 'Energy'
291            data *= sc.tera # THz to Hz
292            data *= sc.value('Planck constant in eV s') # Hz to eV
293            data /= sc.milli # eV to meV
294
295        # Time (ps) to TOF (s)
296        elif name.lower() == 'time' and unit.lower() == 'ps':
297            logger.information('Axis {0} will be converted to time in microsecond'.format(name))
298            unit = 'TOF'
299            data *= sc.micro # ps to us
300
301        # No conversion
302        else:
303            unit = 'Empty'
304
305        return (data, unit, name)
306
307#------------------------------------------------------------------------------
308
309AlgorithmFactory.subscribe(LoadNMoldyn4Ascii)