Refs #8801 Add support for partial DOS.

Code/Mantid/Framework/PythonInterface/plugins/algorithms/WorkflowAlgorithms/DensityOfStates.py

@@ -47,8 +47,11 @@ def PyInit(self):
 			self.declareProperty(name='ZeroThreshold', defaultValue=3.0, 
 				doc='Ignore frequencies below the this threshold. Default is 3.0')
 
+			self.declareProperty(StringArrayProperty('Ions', Direction.Input), 
+				doc="List of Ions to use to calculate partial density of states. If left blank, total density of states will be calculated")
+
 			self.declareProperty(MatrixWorkspaceProperty('OutputWorkspace', '', Direction.Output), 
-			  doc="Name to give the output workspace.")
+				doc="Name to give the output workspace.")
 
 			#regex pattern for a floating point number
 			self._float_regex = '\-?(?:\d+\.?\d*|\d*\.?\d+)'
@@ -60,32 +63,42 @@ def PyExec(self):
 			self._get_properties()
 
 			file_name = self.getPropertyValue('File')
-			frequencies, ir_intensities, raman_intensities, weights = self._read_data_from_file(file_name)
+			file_data = self._read_data_from_file(file_name)
+			frequencies, ir_intensities, raman_intensities, weights = file_data[:4]
+
+			prog_reporter = Progress(self,0.0,1.0,1)
+			if self._calc_partial:
+				eigenvectors = file_data[4]
+				prog_reporter.report("Calculating partial density of states")
+				self._compute_partial(frequencies, eigenvectors, weights)
+				mtd[self._ws_name].setYUnit('(D/A)^2/amu')
+				mtd[self._ws_name].setYUnitLabel('Intensity')
 
-			if self._spec_type == 'DOS':
+			elif self._spec_type == 'DOS':
+				prog_reporter.report("Calculating density of states")
 				self._compute_DOS(frequencies, np.ones_like(frequencies), weights)
-				#set y units on output workspace
 				mtd[self._ws_name].setYUnit('(D/A)^2/amu')
 				mtd[self._ws_name].setYUnitLabel('Intensity')
 
 			elif self._spec_type == 'IR_Active':
 				if ir_intensities.size == 0:
 					raise ValueError("Could not load any IR intensities from file.")
 
+				prog_reporter.report("Calculating IR intensities")
 				self._compute_DOS(frequencies, ir_intensities, weights)
-				#set y units on output workspace
 				mtd[self._ws_name].setYUnit('(D/A)^2/amu')
 				mtd[self._ws_name].setYUnitLabel('Intensity')
 
 			elif self._spec_type == 'Raman_Active':
 				if raman_intensities.size == 0:
 					raise ValueError("Could not load any Raman intensities from file.")
 
+				prog_reporter.report("Calculating Raman intensities")
 				self._compute_raman(frequencies, raman_intensities, weights)
-				#set y units on output workspace
 				mtd[self._ws_name].setYUnit('A^4')
 				mtd[self._ws_name].setYUnitLabel('Intensity')
 
+
 			self.setProperty("OutputWorkspace", self._ws_name)
 
 #----------------------------------------------------------------------------------------
@@ -102,6 +115,8 @@ def _get_properties(self):
 			self._peak_width = self.getProperty('PeakWidth').value
 			self._scale = self.getProperty('Scale').value
 			self._zero_threshold = self.getProperty('ZeroThreshold').value
+			self._ions = self.getProperty('Ions').value
+			self._calc_partial = (len(self._ions) > 0)
 
 #----------------------------------------------------------------------------------------
 
@@ -138,6 +153,43 @@ def _draw_peaks(self, hist, peaks):
 
 			return dos
 
+#----------------------------------------------------------------------------------------
+
+		def _compute_partial(self, frequencies, eigenvectors, weights):
+			"""
+			Compute partial Density Of States.
+
+			This uses the eigenvectors in a .phonon file to calculate
+			the partial density of states.
+
+			@param frequencies - frequencies read from file
+			@param eigenvectors - eigenvectors read from file
+			@param weights - weights for each frequency block
+			"""
+
+			intensities = []
+			for block_vectors in eigenvectors:
+				block_intensities = []
+				for mode in xrange(self._num_branches):
+					#only select vectors for the ions we're interested in
+					lower, upper = mode*self._num_ions, (mode+1)*self._num_ions
+					vectors = block_vectors[lower:upper] 
+					vectors = vectors[self._partial_ion_numbers]
+
+					#compute intensity
+					exponent = np.empty(vectors.shape)
+					exponent.fill(2)
+					vectors = np.power(vectors, exponent)
+					total = np.sum(vectors)
+
+					block_intensities.append(total)
+
+				intensities += block_intensities
+
+			intensities = np.asarray(intensities)
+			self._compute_DOS(frequencies, intensities, weights)
+
+
 #----------------------------------------------------------------------------------------
 
 		def _compute_DOS(self, frequencies, intensities, weights):
@@ -245,12 +297,12 @@ def _read_data_from_file(self, file_name):
 			elif ext == '.castep':
 				file_data = self._parse_castep_file(file_name)
 
-			frequencies, ir_intensities, raman_intensities, weights = file_data
+			frequencies = file_data[0]
 
 			if ( frequencies.size == 0 ):
 				raise ValueError("Failed to load any frequencies from file.")
 
-			return frequencies, ir_intensities, raman_intensities, weights
+			return file_data
 
 #----------------------------------------------------------------------------------------
 
@@ -287,8 +339,38 @@ def _parse_phonon_file_header(self, f_handle):
 					self._num_ions = int(line.strip().split()[-1])
 				elif 'Number of branches' in line:
 					self._num_branches = int(line.strip().split()[-1])
+				elif self._calc_partial and 'Fractional Co-ordinates' in line:
+					#we're calculating partial density of states
+					ion_dict = dict( (ion, []) for ion in self._ions)
+
+					if self._num_ions == None:
+						raise IOError("Failed to parse file. Invalid file header.")
+
+					#extract the mode number for each of the ions we're interested in 
+					for _ in xrange(self._num_ions):
+						line = f_handle.readline()
+						line_data = line.strip().split()
+						ion = line_data[4]
+
+						if ion in self._ions:
+							mode = int(line_data[0])-1 #-1 to convert to zero based indexing
+							ion_dict[ion].append(mode)
+
+					self._partial_ion_numbers = []
+					for ion, ion_nums in ion_dict.items():
+						if len(ion_nums) == 0:
+							logger.warning("Could not find any ions of type %s" % ion)		
+						self._partial_ion_numbers += ion_nums
+
+					self._partial_ion_numbers = sorted(self._partial_ion_numbers)
+					self._partial_ion_numbers = np.asarray(self._partial_ion_numbers)
+
+					if self._partial_ion_numbers.size == 0:
+						raise ValueError("Could not find any of the specified ions")
 
-				if 'END header' in line: 
+				if 'END header' in line:
+					if self._num_ions == None or self._num_branches == None:
+						raise IOError("Failed to parse file. Invalid file header.")
 					return
 
 #----------------------------------------------------------------------------------------
@@ -299,11 +381,32 @@ def _parse_phonon_freq_block(self, f_handle):
 
 				@param f_handle - handle to the file.
 				"""
+				prog_reporter = Progress(self,0.0,1.0,1)
 				for _ in xrange( self._num_branches):
 					line = f_handle.readline()
 					line_data = line.strip().split()[1:]
 					line_data = map(float, line_data)
 					yield line_data
+
+				prog_reporter.report("Reading frequencies.")
+
+#----------------------------------------------------------------------------------------
+		def _parse_phonon_eigenvectors(self, f_handle):
+			vectors = []
+			prog_reporter = Progress(self,0.0,1.0,self._num_branches*self._num_ions)
+			for _ in xrange(self._num_ions*self._num_branches):
+				line = f_handle.readline()
+
+				if not line:
+					raise IOError("Could not parse file. Invalid file format.")
+
+				line_data = line.strip().split()
+				vector_componets = line_data[2::2]
+				vector_componets = map(float, vector_componets)
+				vectors.append(vector_componets)
+				prog_reporter.report("Reading eigenvectors.")
+
+			return np.asarray(vectors)
 
 #----------------------------------------------------------------------------------------
 
@@ -323,6 +426,7 @@ def _parse_phonon_file(self, file_name):
 			block_count = 0
 
 			frequencies, ir_intensities, raman_intensities, weights = [], [], [], []
+			eigenvectors = []
 			data_lists = (frequencies, ir_intensities, raman_intensities)
 			with open(file_name, 'rU') as f_handle:
 				self._parse_phonon_file_header(f_handle)
@@ -345,20 +449,26 @@ def _parse_phonon_file(self, file_name):
 							for data_list, item in zip(data_lists, line_data):
 								data_list.append(item)
 
-					#skip over eigenvectors
 					vector_match = eigenvectors_regex.match(line)
 					if vector_match:
-						for _ in xrange(self._num_ions*self._num_branches):
-							line = f_handle.readline()
-							if not line:
-								raise IOError("Could not parse file. Invalid file format.")
-
+						if self._calc_partial:
+							#parse eigenvectors for partial dos
+							vectors = self._parse_phonon_eigenvectors(f_handle)
+							eigenvectors.append(vectors)
+						else:
+							#skip over eigenvectors
+							for _ in xrange(self._num_ions*self._num_branches):
+								line = f_handle.readline()
+								if not line:
+									raise IOError("Bad file format. Uexpectedly reached end of file.")
+
 			frequencies = np.asarray(frequencies) 
 			ir_intensities = np.asarray(ir_intensities)
+			eigenvectors = np.asarray(eigenvectors)
 			raman_intensities = np.asarray(raman_intensities)
 			warray = np.repeat(weights, self._num_branches)
 
-			return frequencies, ir_intensities, raman_intensities, warray
+			return frequencies, ir_intensities, raman_intensities, warray, eigenvectors
 
 #----------------------------------------------------------------------------------------
 
@@ -393,6 +503,7 @@ def _parse_castep_freq_block(self, f_handle):
 
 			@param f_handle - handle to the file.
 			"""
+			prog_reporter = Progress(self,0.0,1.0,1)
 			for _ in xrange(self._num_branches):
 				line = f_handle.readline()
 				line_data = line.strip().split()[1:-1]
@@ -411,6 +522,8 @@ def _parse_castep_freq_block(self, f_handle):
 				line_data = map(float, line_data)
 				yield line_data
 
+			prog_reporter.report("Reading frequencies.")
+
 #----------------------------------------------------------------------------------------
 
 		def _find_castep_freq_block(self, f_handle, data_regex):