/Framework/PythonInterface/plugins/algorithms/LoadNMoldyn4Ascii.py
Python | 309 lines | 210 code | 49 blank | 50 comment | 32 complexity | 9178e889ba64e552f075475538f171c5 MD5 | raw file
- #pylint: disable=no-init
- from __future__ import (absolute_import, division, print_function)
- from mantid.simpleapi import *
- from mantid.kernel import *
- from mantid.api import *
- import numpy as np
- import scipy.constants as sc
- import ast
- import fnmatch
- import re
- import os
- #------------------------------------------------------------------------------
- VARIABLE_REGEX = re.compile(r'#\s+variable name:\s+(.*)')
- TYPE_REGEX = re.compile(r'#\s+type:\s+([A-z]+)')
- AXIS_REGEX = re.compile(r'#\s+axis:\s+([A-z]+)\|([A-z]+)')
- UNIT_REGEX = re.compile(r'#\s+units:\s+(.*)')
- SLICE_1D_HEADER_REGEX = re.compile(r'#slice:\[([0-9]+)[A-z]*\]')
- SLICE_2D_HEADER_REGEX = re.compile(r'#slice:\[([0-9]+)[A-z]*,\s+([0-9]+)[A-z]*\]')
- #------------------------------------------------------------------------------
- class LoadNMoldyn4Ascii(PythonAlgorithm):
- _axis_cache = None
- _data_directory = None
- #------------------------------------------------------------------------------
- def category(self):
- return 'Inelastic\\DataHandling;Simulation'
- def summary(self):
- return 'Imports functions from .dat files output by nMOLDYN 4.'
- #------------------------------------------------------------------------------
- def PyInit(self):
- self.declareProperty(FileProperty('Directory', '',
- action=FileAction.Directory),
- doc='Path to directory containg .dat files')
- self.declareProperty(StringArrayProperty('Functions'),
- doc='Names of functions to attempt to load from file')
- self.declareProperty(WorkspaceProperty('OutputWorkspace', '',
- direction=Direction.Output),
- doc='Output workspace name')
- #------------------------------------------------------------------------------
- def validateInputs(self):
- issues = dict()
- if len(self.getProperty('Functions').value) == 0:
- issues['Functions'] = 'Must specify at least one function to load'
- return issues
- #------------------------------------------------------------------------------
- def PyExec(self):
- self._axis_cache = {}
- self._data_directory = self.getPropertyValue('Directory')
- # Convert the simplified function names to the actual file names
- data_directory_files = [os.path.splitext(f)[0] for f in fnmatch.filter(os.listdir(self._data_directory), '*.dat')]
- logger.debug('All data files: {0}'.format(data_directory_files))
- functions_input = [x.strip().replace(',', '') for x in self.getProperty('Functions').value]
- functions = [f for f in data_directory_files if f.replace(',', '') in functions_input]
- logger.debug('Functions to load: {0}'.format(functions))
- loaded_function_workspaces = []
- for func_name in functions:
- try:
- # Load the intensity data
- function = self._load_function(func_name)
- # Load (or retrieve) the axis data
- v_axis = self._load_axis(function[2][0])
- x_axis = self._load_axis(function[2][1])
- # Perform axis unit conversions
- v_axis = self._axis_conversion(*v_axis)
- x_axis = self._axis_conversion(*x_axis)
- # Create the workspace for function
- create_workspace = AlgorithmManager.Instance().create('CreateWorkspace')
- create_workspace.initialize()
- create_workspace.setLogging(False)
- create_workspace.setProperty('OutputWorkspace', func_name)
- create_workspace.setProperty('DataX', x_axis[0])
- create_workspace.setProperty('DataY', function[0])
- create_workspace.setProperty('NSpec', v_axis[0].size)
- create_workspace.setProperty('UnitX', x_axis[1])
- create_workspace.setProperty('YUnitLabel', function[1])
- create_workspace.setProperty('VerticalAxisValues', v_axis[0])
- create_workspace.setProperty('VerticalAxisUnit', v_axis[1])
- create_workspace.setProperty('WorkspaceTitle', func_name)
- create_workspace.execute()
- loaded_function_workspaces.append(func_name)
- except ValueError as rerr:
- logger.warning('Failed to load function {0}. Error was: {1}'.format(func_name, str(rerr)))
- # Process the loaded workspaces
- out_ws_name = self.getPropertyValue('OutputWorkspace')
- if len(loaded_function_workspaces) == 0:
- raise RuntimeError('Failed to load any functions for data')
- GroupWorkspaces(InputWorkspaces=loaded_function_workspaces,
- OutputWorkspace=out_ws_name)
- # Set the output workspace
- self.setProperty('OutputWorkspace', out_ws_name)
- #------------------------------------------------------------------------------
- def _load_function(self, function_name):
- """
- Loads a function from the data directory.
- @param function_name Name of the function to load
- @return Tuple of (Numpy array of data, unit, (v axis name, x axis name))
- @exception ValueError If function is not found
- """
- function_filename = os.path.join(self._data_directory, '{0}.dat'.format(function_name))
- if not os.path.isfile(function_filename):
- raise ValueError('File for function "{0}" not found'.format(function_name))
- data = None
- axis = (None, None)
- unit = None
- with open(function_filename, 'rU') as f_handle:
- while True:
- line = f_handle.readline()
- if not line:
- break
- # Ignore empty lines
- if len(line[0]) == 0:
- pass
- # Parse header lines
- elif line[0] == '#':
- variable_match = VARIABLE_REGEX.match(line)
- if variable_match and variable_match.group(1) != function_name:
- raise ValueError('Function name differs from file name')
- axis_match = AXIS_REGEX.match(line)
- if axis_match:
- axis = (axis_match.group(1), axis_match.group(2))
- unit_match = UNIT_REGEX.match(line)
- if unit_match:
- unit = unit_match.group(1)
- slice_match = SLICE_2D_HEADER_REGEX.match(line)
- if slice_match:
- dimensions = (int(slice_match.group(1)), int(slice_match.group(2)))
- # Now parse the data
- data = self._load_2d_slice(f_handle, dimensions)
- return (data, unit, axis)
- #------------------------------------------------------------------------------
- def _load_axis(self, axis_name):
- """
- Loads an axis by name from the data directory.
- @param axis_name Name of axis to load
- @return Tuple of (Numpy array of data, unit, name)
- @exception ValueError If axis is not found
- """
- if axis_name in self._axis_cache:
- return self._axis_cache[axis_name]
- axis_filename = os.path.join(self._data_directory, '{0}.dat'.format(axis_name))
- if not os.path.isfile(axis_filename):
- raise ValueError('File for axis "{0}" not found'.format(axis_name))
- data = None
- unit = None
- with open(axis_filename, 'rU') as f_handle:
- while True:
- line = f_handle.readline()
- if not line:
- break
- # Ignore empty lines
- if len(line[0]) == 0:
- pass
- # Parse header lines
- elif line[0] == '#':
- variable_match = VARIABLE_REGEX.match(line)
- if variable_match and variable_match.group(1) != axis_name:
- raise ValueError('Axis name differs from file name')
- unit_match = UNIT_REGEX.match(line)
- if unit_match:
- unit = unit_match.group(1)
- slice_match = SLICE_1D_HEADER_REGEX.match(line)
- if slice_match:
- length = int(slice_match.group(1))
- # Now parse the data
- data = self._load_1d_slice(f_handle, length)
- return (data, unit, axis_name)
- #------------------------------------------------------------------------------
- def _load_1d_slice(self, f_handle, length):
- """
- Loads a 1D slice from the open file.
- @param f_handle Handle to the open file with the iterator at the slice header
- @param length Length of data
- @return Numpy array of length [length]
- """
- data = np.ndarray(shape=(length), dtype=float)
- for idx in range(length):
- line = f_handle.readline()
- # End of file or empty line (either way end of data)
- if not line or len(line) == 0:
- break
- data[idx] = ast.literal_eval(line)
- return data
- #------------------------------------------------------------------------------
- def _load_2d_slice(self, f_handle, dimensions):
- """
- Loads a 2D slice from the open file.
- @param f_handle Handle to the open file with the iterator at the slice header
- @param dimensions Tuple containing dimensions (rows/vertical axis, cols/x axis)
- @return Numpy array of shape [dimensions]
- """
- data = np.ndarray(shape=dimensions, dtype=float)
- for v_idx in range(dimensions[0]):
- line = f_handle.readline()
- # End of file or empty line (either way end of data)
- if not line or len(line) == 0:
- break
- values = [ast.literal_eval(s) for s in line.split()]
- data[v_idx] = np.array(values)
- return data
- #------------------------------------------------------------------------------
- def _axis_conversion(self, data, unit, name):
- """
- Converts an axis to a Mantid axis type (possibly performing a unit
- conversion).
- @param data The axis data as Numpy array
- @param unit The axis unit as read from the file
- @param name The axis name as read from the file
- @return Tuple containing updated axis details
- """
- logger.debug('Axis for conversion: name={0}, unit={1}'.format(name, unit))
- # Q (nm**-1) to Q (Angstrom**-1)
- if name.lower() == 'q' and unit.lower() == 'inv_nm':
- logger.information('Axis {0} will be converted to Q in Angstrom**-1'.format(name))
- unit = 'MomentumTransfer'
- data /= sc.nano # nm to m
- data *= sc.angstrom # m to Angstrom
- # Frequency (THz) to Energy (meV)
- elif name.lower() == 'frequency' and unit.lower() == 'thz':
- logger.information('Axis {0} will be converted to energy in meV'.format(name))
- unit = 'Energy'
- data *= sc.tera # THz to Hz
- data *= sc.value('Planck constant in eV s') # Hz to eV
- data /= sc.milli # eV to meV
- # Time (ps) to TOF (s)
- elif name.lower() == 'time' and unit.lower() == 'ps':
- logger.information('Axis {0} will be converted to time in microsecond'.format(name))
- unit = 'TOF'
- data *= sc.micro # ps to us
- # No conversion
- else:
- unit = 'Empty'
- return (data, unit, name)
- #------------------------------------------------------------------------------
- AlgorithmFactory.subscribe(LoadNMoldyn4Ascii)