Convert SPE files into NXSPE files

spe2nxspe.py

#!/usr/bin/env python

'''
Created on Nov 24, 2009

@author: Stuart Campbell - SNS, Oak Ridge National Laboratory
'''

import nxs
import optparse
import numpy
import math
import os

VERSION = "1.1"
PROGRAM = "spe2nxspe.py"


def writeProgramDetails(nxfile):
    nxfile.makedata("definition",'char',[5])
    nxfile.opendata("definition")
    nxfile.putattr("version", VERSION)
    nxfile.putdata("NXSPE")
    nxfile.closedata()
    
    nxfile.makedata("program_name",'char',[len(PROGRAM)])
    nxfile.opendata("program_name")
    nxfile.putattr("version", VERSION)
    nxfile.putdata(PROGRAM)
    nxfile.closedata()
      

if __name__ == '__main__':
   
    info = []
    info.append("This program will convert ASCII SPE/PHX files into HDF5 ")
    info.append("based NXSPE files.\n")
    info.append("You have the option of storing a number of additional ")
    info.append("parameters in the NXSPE as well.  These can then be")
    info.append("read automatically into programs such as mslice. \n")
    info.append("Stuart Campbell ([email protected])")

#    parser = optparse.OptionParser("usage: %prog [options] <runnumber>",
 #                                  None, optparse.Option, VERSION, 'error',
  #                                 " ".join(info))
 
    parser = optparse.OptionParser(usage="usage: %prog [options] ", version=VERSION, 
	description=" ".join(info))
        
    parser.add_option("-i", "--inst" , dest="inst",
		      help="The instrument name (e.g. CNCS)")
    
    parser.add_option("-c", "--cvinfo" , dest="cvinfo",
		      help="The location of the cvinfo file")
    
    parser.add_option("-r", "--run-number", dest="runnumber",
		     help="The run number.")
    
    parser.add_option("-s", "--spe", dest="spe",
		      help="The SPE filename")
	
    parser.add_option("-p", "--phx", dest="phx",
		      help="The PHX filename")
	
    parser.add_option("-o", "--output", dest="output",
		      help="The output filename for the resulting NXSPE file.")
    
    parser.add_option("--seblock", dest="seblock",
		     help="The name of the motor as stored in the cvinfo file (case-sensitive).")
    
    parser.add_option("-a", "--angle", dest="rotation",
		     help="The value of the sample angle (degrees).")
    
    parser.add_option("-t", "--temperature", dest="temperature",
		     help="The temperature of the sample (K).")
    
    parser.add_option("--psi", dest="psi",
		     help="The value of psi (as used in mslice plotting software). PSI=ANGLE+OFFSET")
    
    parser.add_option("-e", "--ei", dest="energy",
		     help="The value of the fixed energy (meV).")
	
    parser.add_option("--lambda-scaling", dest="lambdascaling",
		     help="Has ki/kf scaling been applied (1=yes,0=no).")
    
    parser.add_option("-q", "--quiet", action="store_true", dest="quiet",
      help="Disable print statements")
	
    parser.add_option("-f", "--force", action="store_true", dest="force",
      help="Force overwriting of any files.")

    parser.add_option("", "--grab", action="store_true", dest="grab",
                      help="Switch for OIC operation to copy input data "\
                      +"files to working directory.")
    parser.set_defaults(grab=False)
    

    (options, args) = parser.parse_args()
    
    if options.spe is None:
	print "ERROR: Please specify an SPE filename."
    else:
	# Check that the specified SPE file exists and is readable
	spefilename = options.spe
	#print "Using SPE file : ", spefilename

    if options.phx is None:
	print "ERROR: Please specify an PHX filename."
    else:
	# Check that the specified PHX file exists and is readable
	phxfilename = options.phx
	
    # Instrument Name
    if options.inst is None:
	instrument = "UNKNOWN"
    else:
	instrument = options.inst
    
    if (instrument == "SEQ"):
	long_instrument = "SEQUOIA"
    else:
	long_instrument = instrument
    
    # Run Number
    if options.runnumber is None:
	runnumber = "UNKNOWN"
    else:
	runnumber = options.runnumber
    
    # Determine the output filename.
    if options.output is None:
        outputfilename = spefilename.replace(".spe", ".nxspe")
    else:
        outputfilename = options.output
	#print "Output Filename:", outputfilename

    # motor name
    if options.seblock is None:
	motor_name = None
    else:
	motor_name = options.seblock

    # Hack to get around DGSreduction bug if we can't find the cvinfo file.
    if options.rotation is not None:
	if options.rotation.startswith("--"):
		options.rotation = None

    # rotation angle
    if options.rotation is None:
	rotation_angle = None
    else:
	rotation_angle = options.rotation
    
    # Psi angle
    if options.psi is None:
	psi_angle = None
    else:
	psi_angle = options.psi

    # Temperature
    if options.temperature is None:
	sample_temperature = None
    else:
	sample_temperature = options.temperature
	
    # Incident Energy
    if options.energy is None:
	parser.error("You need to specify the energy.")
    else:
	incident_energy = options.energy
	
    print outputfilename
   
    if options.lambdascaling is None:
	lambda_scaling = 0
    else:
        lambda_scaling = options.lambdascaling

    # OIC Operation: grab input files (plus PAR) and copy to working directory
    if options.grab:
        parfilename = spefilename.replace(".spe", ".par")
        import shutil
        shutil.copy(spefilename, '.')
        shutil.copy(phxfilename, '.')
        shutil.copy(parfilename, '.')

        dloc = os.path.dirname(spefilename)
        shutil.copy(os.path.join(dloc, '*.stdout.log'), '.')
        
        # Remove locations for input files
        spefilename = os.path.basename(spefilename)
        phxfilename = os.path.basename(phxfilename)
         
    # Check to see if the output file already exists
    if os.path.isfile(outputfilename):
        if options.force:
            if not options.quiet:
                print "Overwriting %s" % outputfilename
		os.remove(outputfilename)
            else:
                raise IOError('Output file %s already exists.  Please use -f or --force to overwrite.' 
                              % outputfilename)			

#    spefilename = '/Users/Shared/Data/SEQ_1131.spe'
#    phxfilename = '/Users/Shared/Data/SEQ_1131.phx'
#    spefilename = '/Users/Shared/Data/CNCS_5162.spe'
#    phxfilename = '/Users/Shared/Data/CNCS_5162.phx'
#    outputfilename = '/Users/Shared/Data/crap.nxs'
    
    # Open the SPE file
    spe_handle = open(spefilename, "r")
    
    # Open the PHX file
    phx_file = open(phxfilename, "r")
    
    # Get the number of energy and angle bins from the SPE file
    (spe_angles, spe_energies) = spe_handle.readline().split()
    # ...and convert them to integers
    energies = int(spe_energies)
    angles = int(spe_angles)
    
    # Now lets check that these values match the PHX file...
    phx_angles = int(phx_file.readline())
    
    if (phx_angles != angles):
	raise "SPE and PHX files do not match"
    
    # Let's read the rest of the PHX file
    theta, psi, dtheta, dpsi = numpy.loadtxt(phx_file, usecols=(2,3,4,5), unpack=True)
    
    # "### Phi Grid" - ignore
    spe_handle.readline()
        
    angle_lines = int(math.ceil((float(spe_angles) + 1.0) / 8.0))
    for i in range(angle_lines):
        fields = spe_handle.readline().split()
        if i == 0:
            angle_bins = numpy.float32(fields)
        else:
            angle_bins = numpy.concatenate((angle_bins, numpy.float32(fields)))    
        
    print angle_bins.__len__()
    
    # "### Energy Grid" - ignore
    buffer = spe_handle.readline()
    
    # Create an array to hold the energy bins
    energy_bins = numpy.zeros(energies+1) 
    n = 0

    energy_lines = int(math.ceil((float(spe_energies) + 1.0) / 8.0))
    for i in range(energy_lines):
	line = spe_handle.readline()
	for j in range(0,len(line)-10,10):
	    energy_bins[n] = numpy.float(line[j:j+10])
	    n += 1
	    
    #print energy_bins.__len__()  

    if not options.quiet:
        print "Total Number of Angles = %s, Energy bins = %s." % (angles, energies)

    # Let's create the NeXus file
    file = nxs.open(outputfilename, 'w5')
    # Construct the name of the NXentry
    if (instrument == "UNKNOWN") or (runnumber == "UNKNOWN"):
	entry_name = "entry"
    else:
	entry_name = instrument + "_" + runnumber
	
    # Make the NXentry
    file.makegroup(entry_name,"NXentry")
    file.opengroup(entry_name,"NXentry")
    
    # Write some static config info...
    writeProgramDetails(file)
    
    note = "This is a file that has been converted from an SPE file"   
   
    # Create an NXSPE_info
    file.makegroup("NXSPE_info", "NXcollection")
    file.opengroup("NXSPE_info", "NXcollection")
  
    if incident_energy is not None:
	file.makedata("fixed_energy", "float64", [1])
	file.opendata("fixed_energy")
	file.putdata(incident_energy)
	file.putattr('units', 'meV')
	file.putattr('mode', 'direct')
	file.closedata()

    if psi_angle is not None:
		file.makedata("psi", "float64", [1])
		file.opendata("psi")
		file.putdata(psi_angle)
		file.putattr('units', 'degrees')
		file.closedata()    	

    if lambda_scaling is not None: 
    	file.makedata("ki_over_kf_scaling", "int8", [1])
    	file.opendata("ki_over_kf_scaling") 
	file.putdata(lambda_scaling)
    	file.closedata() 

    # Close NXSPE_info
    file.closegroup()

    # Make the NXinstrument
    file.makegroup("instrument", "NXinstrument")
    file.opengroup("instrument", "NXinstrument")
    
    # Instrument name
    file.makedata("name",'char',[len(long_instrument)])
    file.opendata("name")
    file.putdata(long_instrument)
    file.putattr("short_name", instrument)
    file.closedata()
    
    
    
    # Make the NXchopper for the Fermi
    file.makegroup("fermi", "NXfermi_chopper")
    file.opengroup("fermi", "NXfermi_chopper")
    
    # TODO: Get the Ei and put it here - but what is it called! 
    if incident_energy is not None:
	file.makedata("energy", "float64", [1])
	file.opendata("energy")
	file.putdata(incident_energy)
	file.putattr('units', 'meV')
	file.closedata()
    
    # NXchopper (fermi)
    file.closegroup()
    
    # NXinstrument
    file.closegroup()
       
    # Make the NXsample
    file.makegroup("sample", "NXsample")
    file.opengroup("sample", "NXsample")
    
    # TODO: put the angles here!
    if rotation_angle is not None:
		file.makedata("rotation_angle", "float64", [1])
		file.opendata("rotation_angle")
		file.putdata(rotation_angle)
		file.putattr('units', 'degrees')
		file.closedata()
    
    # Sample Temperature
    if sample_temperature is not None:
		file.makedata("temperature", "float64", [1])
		file.opendata("temperature")
		file.putdata(sample_temperature)
		file.putattr("units", "K")
		file.closedata()

    # Motor Name
    if motor_name is not None:
	file.makedata("seblock",'char',[len(motor_name)])
    	file.opendata("seblock")
    	file.putdata(motor_name)
    	file.closedata()
    
    # NXsample
    file.closegroup()
       
    # Make an NXdata
    file.makegroup("data","NXdata")
    file.opengroup("data","NXdata")

    # Write the 2theta (polar) angle
    file.makedata('polar','float64',[angles])
    file.opendata('polar')
    file.putdata(theta)
    file.closedata()

    # Write the 2theta (polar) angle widths
    file.makedata('polar_width','float64',[angles])
    file.opendata('polar_width')
    file.putdata(dtheta)
    file.closedata()

    # Write the psi (azimuthal) angle
    file.makedata('azimuthal','float64',[angles])
    file.opendata('azimuthal')
    file.putdata(psi)
    file.closedata()

    # Write the psi (azimuthal) angle widths
    file.makedata('azimuthal_width','float64',[angles])
    file.opendata('azimuthal_width')
    file.putdata(dpsi)
    file.closedata()

    # Write the energy bins
    file.makedata('energy','float64',[energies+1])
    file.opendata('energy')
    file.putattr('units', 'meV')
    file.putdata(energy_bins)
    file.closedata()

    # Just for the moment - for the sake of speed - just read 50 angles
    #angles = 50

    # Lets create the data and errors arrays
    #file.makedata('data', 'float64', [angles, energies])
    #file.makedata('error', 'float64', [angles, energies])
    file.compmakedata('data', 'float64', [angles, energies], 'lzw', [1,energies])
    file.compmakedata('error', 'float64', [angles, energies], 'lzw', [1,energies])

    # How many lines does the data take up ?
    data_lines = int(math.ceil((float(spe_energies)) / 8.0))

    # Now lets loop over the angle bins
    for angle in range(angles):
	# '### S(Phi,w)'
	buffer = spe_handle.readline()
	# Let's create an array to hold the data
	data = numpy.zeros(energies)
	# Data point counter
	nd = 0
        # Now lets read the data
	for i in range(data_lines):
	    line = spe_handle.readline()
	    for j in range(0,len(line)-10,10):
		data[nd] = numpy.float(line[j:j+10])
		nd += 1
	
	file.opendata('data')
	file.putslab(data,[angle,0], [1,energies])
	file.closedata()
	
	# '### Errors'
	buffer = spe_handle.readline()
	# Let's create an array to hold the errors
	error = numpy.zeros(energies)
	# Error point counter
	ne = 0
        # Now lets read the data
	for i in range(data_lines):
	    line = spe_handle.readline()
	    for j in range(0,len(line)-10,10):
		error[ne] = numpy.float(line[j:j+10])
		ne += 1
	
	file.opendata('error')
	file.putslab(error, [angle,0], [1,energies])
	file.closedata()
	
    # Close NXdata
    file.closegroup()

    # Close the NeXus File
    file.close()
	
    # Close the PHX file
    phx_file.close()
	
    # Close the SPE file
    spe_handle.close()