Ticket #8391: IndirectBayes.py

File IndirectBayes.py, 31.7 KB (added by Samuel Jackson, 7 years ago)

Updated version of the script

Line 
1# Bayes routines
2# Fortran programs use fixed length arrays whereas Python has variable lenght lists
3# Input : the Python list is padded to Fortrans length using procedure PadArray
4# Output : the Fortran numpy array is sliced to Python length using dataY = yout[:ny]
5#
6from IndirectImport import *
7if is_supported_f2py_platform():
8    QLr     = import_f2py("QLres")
9    QLd     = import_f2py("QLdata")
10    Qse     = import_f2py("QLse")
11    Que     = import_f2py("Quest")
12    resnorm = import_f2py("ResNorm")
13else:
14    unsupported_message()
15
16from mantid.simpleapi import *
17from mantid import config, logger, mtd
18from IndirectCommon import *
19import sys, platform, math, os.path, numpy as np
20mp = import_mantidplot()
21
22def CalcErange(inWS,ns,erange,nbin):
23        rscl = 1.0
24        array_len = 4096                           # length of array in Fortran
25        N,X,Y,E = GetXYE(inWS,ns,array_len)        # get data
26        Xin = mtd[inWS].readX(0)
27        ndat = len(Xin)-1                                                       # get no. points from length of x array
28        imin = 0
29        for i in range(0,ndat):
30                if Xin[i] <= erange[0]:
31                        imin = i
32        imin += 1
33        imax = ndat
34        for i in range(ndat,0,-1):
35                if Xin[i] >= erange[1]:
36                        imax = i
37        nbin = int(nbin)
38        bnorm=1.0/nbin
39        N=0
40        Xout = []
41        for i in range(imin,imax,nbin):
42                N += 1
43                XXD = 0.0
44                for j in range(0,nbin):
45                        XXD = XXD+Xin[i+j]
46                Xout.append(bnorm*XXD)
47        nout = [N, imin+1, imax]
48        Xdat=PadArray(Xout,array_len)
49        if nout[0] == 0:
50                error = 'Erange - calculated points is Zero'                   
51                logger.notice('ERROR *** ' + error)
52                sys.exit(error)
53        if nout[1] == 0:
54                error = 'Erange - calculated Imin is Zero'                     
55                logger.notice('ERROR *** ' + error)
56                sys.exit(error)
57        if nout[2] == 0:
58                error = 'Erange - calculated Imax is Zero'                     
59                logger.notice('ERROR *** ' + error)
60                sys.exit(error)
61        return nout,bnorm,Xdat,X,Y,E
62
63def GetXYE(inWS,n,array_len):
64        Xin = mtd[inWS].readX(n)
65        N = len(Xin)-1                                                  # get no. points from length of x array
66        Yin = mtd[inWS].readY(n)
67        Ein = mtd[inWS].readE(n)
68        X=PadArray(Xin,array_len)
69        Y=PadArray(Yin,array_len)
70        E=PadArray(Ein,array_len)
71        return N,X,Y,E
72
73def GetResNorm(resnormWS,ngrp):
74        if ngrp == 0:                                # read values from WS
75                dtnorm = mtd[resnormWS+'_Intensity'].readY(0)
76                xscale = mtd[resnormWS+'_Stretch'].readY(0)
77        else:                                        # constant values
78                dtnorm = []
79                xscale = []
80                for m in range(0,ngrp):
81                        dtnorm.append(1.0)
82                        xscale.append(1.0)
83        dtn=PadArray(dtnorm,51)                      # pad for Fortran call
84        xsc=PadArray(xscale,51)
85        return dtn,xsc
86       
87def ReadNormFile(readRes,resnormWS,nsam,Verbose):            # get norm & scale values
88        if readRes:                   # use ResNorm file option=o_res
89                Xin = mtd[resnormWS+'_Intensity'].readX(0)
90                nrm = len(Xin)                                          # no. points from length of x array
91                if nrm == 0:                           
92                        error = 'ResNorm file has no Intensity points'                 
93                        logger.notice('ERROR *** ' + error)
94                        sys.exit(error)
95                Xin = mtd[resnormWS+'_Stretch'].readX(0)                                        # no. points from length of x array
96                if len(Xin) == 0:                               
97                        error = 'ResNorm file has no xscale points'
98                        logger.notice('ERROR *** ' + error)
99                        sys.exit(error)
100                if nrm != nsam:                         # check that no. groups are the same
101                        error = 'ResNorm groups (' +str(nrm) + ') not = Sample (' +str(nsam) +')'                       
102                        logger.notice('ERROR *** ' + error)
103                        sys.exit(error)
104                else:
105                        dtn,xsc = GetResNorm(resnormWS,0)
106        else:
107                # do not use ResNorm file
108                dtn,xsc = GetResNorm(resnormWS,nsam)
109        return dtn,xsc
110
111#Reads in a width ASCII file
112def ReadWidthFile(readWidth,widthFile,numSampleGroups,Verbose):
113        widthY = []
114        widthE = []
115
116        if readWidth:
117
118                if Verbose:     
119                        logger.notice('Width file is ' + widthFile)
120
121                # read ascii based width file
122                try:
123                        wfPath = FileFinder.getFullPath(widthFile)
124                        handle = open(wfPath, 'r')
125                        asc = []
126
127                        for line in handle:
128                                line = line.rstrip()
129                                asc.append(line)
130                        handle.close()
131
132                except Exception, e:
133                        error = 'Failed to read width file'     
134                        logger.notice('ERROR *** ' + error)
135                        sys.exit(error)
136
137                numLines = len(asc)
138               
139                if numLines == 0:
140                        error = 'No groups in width file'       
141                        logger.notice('ERROR *** ' + error)
142                        sys.exit(error)
143               
144                if numLines != numSampleGroups:                         # check that no. groups are the same
145                        error = 'Width groups (' +str(numLines) + ') not = Sample (' +str(numSampleGroups) +')' 
146                        logger.notice('ERROR *** ' + error)
147                        sys.exit(error)
148
149        else: 
150                # no file: just use constant values
151                widthY = np.zeros(numSampleGroups)
152                widthE = np.zeros(numSampleGroups)
153
154        # pad for Fortran call
155        widthY = PadArray(widthY,51)
156        widthE = PadArray(widthE,51)
157
158        return widthY, widthE
159
160# QLines programs
161def QLRun(program,samWS,resWS,resnormWS,erange,nbins,Fit,wfile,Loop,Verbose,Plot,Save):
162        StartTime(program)
163
164        #expand fit options
165        elastic, background, width, resnorm = Fit
166       
167        #convert true/false to 1/0 for fortran
168        o_el = 1 if elastic else 0
169        o_w1 = 1 if width else 0
170        o_res = 1 if resnorm else 0
171
172        #fortran code uses background choices defined using the following numbers
173        if background == 'Sloping':
174                o_bgd = 2
175        elif background == 'Flat':
176                o_bgd = 1
177        elif background == 'Zero':
178                o_bgd = 0
179
180        fitOp = [o_el, o_bgd, o_w1, o_res]
181
182        workdir = getDefaultWorkingDirectory()
183
184        facility = config['default.facility']
185        array_len = 4096                                                   # length of array in Fortran
186        CheckXrange(erange,'Energy')
187
188        nbin,nrbin = nbins[0], nbins[1]
189
190        if Verbose:
191                logger.notice('Sample is ' + samWS)
192                logger.notice('Resolution is ' + resWS)
193
194        if facility == 'ISIS':
195                CheckAnalysers(samWS,resWS,Verbose)
196                efix = getEfixed(samWS)
197                theta,Q = GetThetaQ(samWS)
198
199        nsam,ntc = CheckHistZero(samWS)
200
201        totalNoSam = nsam
202       
203        #check if we're performing a sequential fit
204        if Loop != True:
205                nsam = 1
206
207        nres,ntr = CheckHistZero(resWS)
208       
209        if program == 'QL':
210                if nres == 1:
211                        prog = 'QLr'                                            # res file
212                else:
213                        prog = 'QLd'                                            # data file
214                        CheckHistSame(samWS,'Sample',resWS,'Resolution')
215        elif program == 'QSe':
216                if nres == 1:
217                        prog = 'QSe'                                            # res file
218                else:
219                        error = 'Stretched Exp ONLY works with RES file'
220                        sys.exit(error)
221
222        if Verbose:
223                logger.notice('Version is ' +prog)
224                logger.notice(' Number of spectra = '+str(nsam))
225                logger.notice(' Erange : '+str(erange[0])+' to '+str(erange[1]))
226
227        Wy,We = ReadWidthFile(width,wfile,totalNoSam,Verbose)
228        dtn,xsc = ReadNormFile(resnorm,resnormWS,totalNoSam,Verbose)
229
230        fname = samWS[:-4] + '_'+ prog
231        probWS = fname + '_Prob'
232        fitWS = fname + '_Fit'
233        datWS = fname + '_Data'
234        wrks=os.path.join(workdir, samWS[:-4])
235        if Verbose:
236                logger.notice(' lptfile : '+wrks+'_'+prog+'.lpt')
237        lwrk=len(wrks)
238        wrks.ljust(140,' ')
239        wrkr=resWS
240        wrkr.ljust(140,' ')
241        wrk = [wrks, wrkr]
242
243        # initialise probability list
244        if program == 'QL':
245                prob0 = []
246                prob1 = []
247                prob2 = []
248        xQ = np.array([Q[0]])
249        for m in range(1,nsam):
250                xQ = np.append(xQ,Q[m])
251        xProb = xQ
252        xProb = np.append(xProb,xQ)
253        xProb = np.append(xProb,xQ)
254        eProb = np.zeros(3*nsam)
255
256        group = ''
257        for m in range(0,nsam):
258                if Verbose:
259                        logger.notice('Group ' +str(m)+ ' at angle '+ str(theta[m]))
260                nsp = m+1
261                nout,bnorm,Xdat,Xv,Yv,Ev = CalcErange(samWS,m,erange,nbin)
262                Ndat = nout[0]
263                Imin = nout[1]
264                Imax = nout[2]
265                if prog == 'QLd':
266                        mm = m
267                else:
268                        mm = 0
269                Nb,Xb,Yb,Eb = GetXYE(resWS,mm,array_len)         # get resolution data
270                numb = [nsam, nsp, ntc, Ndat, nbin, Imin, Imax, Nb, nrbin]
271                rscl = 1.0
272                reals = [efix, theta[m], rscl, bnorm]
273                if prog == 'QLr':
274                        nd,xout,yout,eout,yfit,yprob=QLr.qlres(numb,Xv,Yv,Ev,reals,fitOp,
275                                                                                                   Xdat,Xb,Yb,Wy,We,dtn,xsc,
276                                                                                                   wrks,wrkr,lwrk)
277                        message = ' Log(prob) : '+str(yprob[0])+' '+str(yprob[1])+' '+str(yprob[2])+' '+str(yprob[3])
278                        if Verbose:
279                                logger.notice(message)
280                if prog == 'QLd':
281                        nd,xout,yout,eout,yfit,yprob=QLd.qldata(numb,Xv,Yv,Ev,reals,fitOp,
282                                                                                                        Xdat,Xb,Yb,Eb,Wy,We,
283                                                                                                        wrks,wrkr,lwrk)
284                        message = ' Log(prob) : '+str(yprob[0])+' '+str(yprob[1])+' '+str(yprob[2])+' '+str(yprob[3])
285                        if Verbose:
286                                logger.notice(message)
287                if prog == 'QSe':
288                        nd,xout,yout,eout,yfit,yprob=Qse.qlstexp(numb,Xv,Yv,Ev,reals,fitOp,
289                                                                                                        Xdat,Xb,Yb,Wy,We,dtn,xsc,
290                                                                                                        wrks,wrkr,lwrk)
291                dataX = xout[:nd]
292                dataX = np.append(dataX,2*xout[nd-1]-xout[nd-2])
293                yfit_list = np.split(yfit[:4*nd],4)
294                dataF0 = yfit_list[0]
295                dataF1 = yfit_list[1]
296                if program == 'QL':
297                        dataF2 = yfit_list[2]
298                        dataF3 = yfit_list[3]
299                dataG = np.zeros(nd)
300                datX = dataX
301                datY = yout[:nd]
302                datE = eout[:nd]
303                datX = np.append(datX,dataX)
304                datY = np.append(datY,dataF1[:nd])
305                datE = np.append(datE,dataG)
306                res1 = dataF1[:nd] - yout[:nd]
307                datX = np.append(datX,dataX)
308                datY = np.append(datY,res1)
309                datE = np.append(datE,dataG)
310                nsp = 3
311                names = 'data,fit.1,diff.1'
312                res_plot = [0, 1, 2]
313                if program == 'QL':
314                        datX = np.append(datX,dataX)
315                        datY = np.append(datY,dataF2[:nd])
316                        datE = np.append(datE,dataG)
317                        res2 = dataF2[:nd] - yout[:nd]
318                        datX = np.append(datX,dataX)
319                        datY = np.append(datY,res2)
320                        datE = np.append(datE,dataG)
321                        nsp += 2
322                        names += ',fit.2,diff.2'
323                        res_plot.append(4)
324                        prob0.append(yprob[0])
325                        prob1.append(yprob[1])
326                        prob2.append(yprob[2])
327
328                # create result workspace
329                fitWS = fname+'_Result'
330                fout = fitWS +'_'+ str(m)
331
332                CreateWorkspace(OutputWorkspace=fout, DataX=datX, DataY=datY, DataE=datE,
333                        Nspec=nsp, UnitX='DeltaE', VerticalAxisUnit='Text', VerticalAxisValues=names)
334               
335                # append workspace to list of results
336                group += fout + ','
337
338       
339        GroupWorkspaces(InputWorkspaces=group,OutputWorkspace=fitWS)
340
341        if program == 'QL':
342                yPr0 = np.array([prob0[0]])
343                yPr1 = np.array([prob1[0]])
344                yPr2 = np.array([prob2[0]])
345                for m in range(1,nsam):
346                        yPr0 = np.append(yPr0,prob0[m])
347                        yPr1 = np.append(yPr1,prob1[m])
348                        yPr2 = np.append(yPr2,prob2[m])
349                yProb = yPr0
350                yProb = np.append(yProb,yPr1)
351                yProb = np.append(yProb,yPr2)
352                CreateWorkspace(OutputWorkspace=probWS, DataX=xProb, DataY=yProb, DataE=eProb,
353                        Nspec=3, UnitX='MomentumTransfer')
354                outWS = C2Fw(samWS[:-4],fname)
355                if (Plot != 'None'):
356                        QLPlotQL(fname,Plot,res_plot,Loop)
357        if program == 'QSe':
358                outWS = C2Se(fname)
359                if (Plot != 'None'):
360                        QLPlotQSe(fname,Plot,res_plot,Loop)
361
362        #Add some sample logs to the output workspace
363        CopyLogs(InputWorkspace=samWS, OutputWorkspace=outWS)
364        AddSampleLog(Workspace=outWS, LogName="res_file", LogType="String", LogText=resWS)
365        AddSampleLog(Workspace=outWS, LogName="fit_program", LogType="String", LogText=prog)
366        AddSampleLog(Workspace=outWS, LogName="background", LogType="String", LogText=str(background))
367        AddSampleLog(Workspace=outWS, LogName="elastic_peak", LogType="String", LogText=str(elastic))
368        AddSampleLog(Workspace=outWS, LogName="energy_min", LogType="Number", LogText=str(erange[0]))
369        AddSampleLog(Workspace=outWS, LogName="energy_max", LogType="Number", LogText=str(erange[1]))
370        AddSampleLog(Workspace=outWS, LogName="sample_binning", LogType="Number", LogText=str(nbin))
371        AddSampleLog(Workspace=outWS, LogName="resolution_binning", LogType="Number", LogText=str(nrbin))
372        CopyLogs(InputWorkspace=samWS, OutputWorkspace=fitWS)
373        AddSampleLog(Workspace=fitWS, LogName="res_file", LogType="String", LogText=resWS)
374        AddSampleLog(Workspace=fitWS, LogName="fit_program", LogType="String", LogText=prog)
375        AddSampleLog(Workspace=fitWS, LogName="background", LogType="String", LogText=str(background))
376        AddSampleLog(Workspace=fitWS, LogName="elastic_peak", LogType="String", LogText=str(elastic))
377        AddSampleLog(Workspace=fitWS, LogName="energy_min", LogType="Number", LogText=str(erange[0]))
378        AddSampleLog(Workspace=fitWS, LogName="energy_max", LogType="Number", LogText=str(erange[1]))
379        AddSampleLog(Workspace=fitWS, LogName="sample_binning", LogType="Number", LogText=str(nbin))
380        AddSampleLog(Workspace=fitWS, LogName="resolution_binning", LogType="Number", LogText=str(nrbin))
381
382        AddSampleLog(Workspace=outWS, LogName="resnorm", LogType="String", LogText=str(resnorm))
383        AddSampleLog(Workspace=fitWS, LogName="resnorm", LogType="String", LogText=str(resnorm))
384        if resnorm:
385                AddSampleLog(Workspace=outWS, LogName="resnorm_file", LogType="String", LogText=resnormWS)
386                AddSampleLog(Workspace=fitWS, LogName="resnorm_file", LogType="String", LogText=resnormWS)
387       
388        AddSampleLog(Workspace=outWS, LogName="width", LogType="String", LogText=str(width))
389        AddSampleLog(Workspace=fitWS, LogName="width", LogType="String", LogText=str(width))
390       
391        if width:
392                AddSampleLog(Workspace=outWS, LogName="width_file", LogType="String", LogText=wfile)
393                AddSampleLog(Workspace=fitWS, LogName="width_file", LogType="String", LogText=wfile)
394
395        if Save:
396                fit_path = os.path.join(workdir,fitWS+'.nxs')
397                SaveNexusProcessed(InputWorkspace=fitWS, Filename=fit_path)
398                out_path = os.path.join(workdir, outWS+'.nxs')                                  # path name for nxs file
399                SaveNexusProcessed(InputWorkspace=outWS, Filename=out_path)
400                if Verbose:
401                        logger.notice('Output fit file created : ' + fit_path)
402                        logger.notice('Output paramter file created : ' + out_path)
403        EndTime(program)
404
405
406def LorBlock(a,first,nl):                                 #read Ascii block of Integers
407        line1 = a[first]
408        first += 1
409        val = ExtractFloat(a[first])               #Q,AMAX,HWHM,BSCL,GSCL
410        Q = val[0]
411        AMAX = val[1]
412        HWHM = val[2]
413        BSCL = val[3]
414        GSCL = val[4]
415        first += 1
416        val = ExtractFloat(a[first])               #A0,A1,A2,A4
417        int0 = [AMAX*val[0]]
418        bgd1 = BSCL*val[1]
419        bgd2 = BSCL*val[2]
420        zero = GSCL*val[3]
421        first += 1
422        val = ExtractFloat(a[first])                #AI,FWHM first peak
423        fw = [2.*HWHM*val[1]]
424        int = [AMAX*val[0]]
425        if nl >= 2:
426                first += 1
427                val = ExtractFloat(a[first])            #AI,FWHM second peak
428                fw.append(2.*HWHM*val[1])
429                int.append(AMAX*val[0])
430        if nl == 3:
431                first += 1
432                val = ExtractFloat(a[first])            #AI,FWHM third peak
433                fw.append(2.*HWHM*val[1])
434                int.append(AMAX*val[0])
435        first += 1
436        val = ExtractFloat(a[first])                 #SIG0
437        int0.append(val[0])
438        first += 1
439        val = ExtractFloat(a[first])                  #SIGIK
440        int.append(AMAX*math.sqrt(math.fabs(val[0])+1.0e-20))
441        first += 1
442        val = ExtractFloat(a[first])                  #SIGFK
443        fw.append(2.0*HWHM*math.sqrt(math.fabs(val[0])+1.0e-20))
444        if nl >= 2:                                      # second peak
445                first += 1
446                val = ExtractFloat(a[first])                  #SIGIK
447                int.append(AMAX*math.sqrt(math.fabs(val[0])+1.0e-20))
448                first += 1
449                val = ExtractFloat(a[first])                  #SIGFK
450                fw.append(2.0*HWHM*math.sqrt(math.fabs(val[0])+1.0e-20))
451        if nl == 3:                                       # third peak
452                first += 1
453                val = ExtractFloat(a[first])                  #SIGIK
454                int.append(AMAX*math.sqrt(math.fabs(val[0])+1.0e-20))
455                first += 1
456                val = ExtractFloat(a[first])                  #SIGFK
457                fw.append(2.0*HWHM*math.sqrt(math.fabs(val[0])+1.0e-20))
458        first += 1
459        return first,Q,int0,fw,int                                      #values as list
460
461def C2Fw(prog,sname):
462        workdir = config['defaultsave.directory']
463        outWS = sname+'_Workspace'
464        Vaxis = []
465
466        dataX = np.array([])
467        dataY = np.array([])
468        dataE = np.array([])
469
470        nhist = 0
471        for nl in range(1,4):
472                file = sname + '.ql' +str(nl)
473                handle = open(os.path.join(workdir, file), 'r')
474                asc = []
475                for line in handle:
476                        line = line.rstrip()
477                        asc.append(line)
478                handle.close()
479                lasc = len(asc)
480                var = asc[3].split()                                                    #split line on spaces
481                nspec = var[0]
482                ndat = var[1]
483                var = ExtractInt(asc[6])
484                first = 7
485                Xout = []
486
487                ns = int(nspec)
488
489                YData = [[] for i in range(6)]
490                EData = [[] for i in range(6)]
491
492                for m in range(0,ns):
493                        first,Q,i0,fw,it = LorBlock(asc,first,nl)
494                        Xout.append(Q)
495
496                        for i in range(0,nl):
497                                #collect amplitude and width data
498                                YData[i*2].append(fw[i])
499                                YData[i*2+1].append(it[i])
500                                EData[i*2].append(fw[nl+i])
501                                EData[i*2+1].append(it[nl+i])
502
503                nhist += nl*2
504               
505                for i in range(0,nl):
506                        #append amplitude
507                        dataX = np.append(dataX, np.array(Xout))
508                        dataY = np.append(dataY, np.array(YData[i*2+1]))
509                        dataE = np.append(dataE, np.array(EData[i*2+1]))
510                        Vaxis.append('ampl.'+str(nl)+'.'+str(i+1))
511
512                        #append width
513                        dataX = np.append(dataX, np.array(Xout))
514                        dataY = np.append(dataY, np.array(YData[i*2]))
515                        dataE = np.append(dataE, np.array(EData[i*2]))
516                        Vaxis.append('width.'+str(nl)+'.'+str(i+1))
517
518        CreateWorkspace(OutputWorkspace=outWS, DataX=dataX, DataY=dataY, DataE=dataE, Nspec=nhist,
519                UnitX='MomentumTransfer', VerticalAxisUnit='Text', VerticalAxisValues=Vaxis, YUnitLabel='')
520
521        return outWS
522
523def SeBlock(a,first):                                 #read Ascii block of Integers
524        line1 = a[first]
525        first += 1
526        val = ExtractFloat(a[first])               #Q,AMAX,HWHM
527        Q = val[0]
528        AMAX = val[1]
529        HWHM = val[2]
530        first += 1
531        val = ExtractFloat(a[first])               #A0
532        int0 = [AMAX*val[0]]
533        first += 1
534        val = ExtractFloat(a[first])                #AI,FWHM first peak
535        fw = [2.*HWHM*val[1]]
536        int = [AMAX*val[0]]
537        first += 1
538        val = ExtractFloat(a[first])                 #SIG0
539        int0.append(val[0])
540        first += 1
541        val = ExtractFloat(a[first])                  #SIG3K
542        int.append(AMAX*math.sqrt(math.fabs(val[0])+1.0e-20))
543        first += 1
544        val = ExtractFloat(a[first])                  #SIG1K
545        fw.append(2.0*HWHM*math.sqrt(math.fabs(val[0])+1.0e-20))
546        first += 1
547        be = ExtractFloat(a[first])                  #EXPBET
548        first += 1
549        val = ExtractFloat(a[first])                  #SIG2K
550        be.append(math.sqrt(math.fabs(val[0])+1.0e-20))
551        first += 1
552        return first,Q,int0,fw,int,be                                      #values as list
553       
554def C2Se(sname):
555        workdir = config['defaultsave.directory']
556        prog = 'QSe'
557        outWS = sname+'_Workspace'
558        handle = open(os.path.join(workdir, sname+'.qse'), 'r')
559        asc = []
560        for line in handle:
561                line = line.rstrip()
562                asc.append(line)
563        handle.close()
564        lasc = len(asc)
565        var = asc[3].split()                                                    #split line on spaces
566        nspec = var[0]
567        ndat = var[1]
568        var = ExtractInt(asc[6])
569        first = 7
570        Xout = []
571        Yf = []
572        Ef = []
573        Yi = []
574        Ei = []
575        Yb = []
576        Eb = []
577        ns = int(nspec)
578
579        dataX = np.array([])
580        dataY = np.array([])
581        dataE = np.array([])
582
583        for m in range(0,ns):
584                first,Q,int0,fw,it,be = SeBlock(asc,first)
585                Xout.append(Q)
586                Yf.append(fw[0])
587                Ef.append(fw[1])
588                Yi.append(it[0])
589                Ei.append(it[1])
590                Yb.append(be[0])
591                Eb.append(be[1])
592        Vaxis = []
593
594        dataX = np.append(dataX,np.array(Xout))
595        dataY = np.append(dataY,np.array(Yi))
596        dataE = np.append(dataE,np.array(Ei))
597        nhist = 1
598        Vaxis.append('ampl')
599
600        dataX = np.append(dataX, np.array(Xout))
601        dataY = np.append(dataY, np.array(Yf))
602        dataE = np.append(dataE, np.array(Ef))
603        nhist += 1
604        Vaxis.append('width')
605
606        dataX = np.append(dataX,np.array(Xout))
607        dataY = np.append(dataY,np.array(Yb))
608        dataE = np.append(dataE,np.array(Eb))
609        nhist += 1
610        Vaxis.append('beta')
611
612        logger.notice('Vaxis=' + str(Vaxis))
613        CreateWorkspace(OutputWorkspace=outWS, DataX=dataX, DataY=dataY, DataE=dataE, Nspec=nhist,
614                UnitX='MomentumTransfer', VerticalAxisUnit='Text', VerticalAxisValues=Vaxis, YUnitLabel='')
615        return outWS
616
617def QLPlotQL(inputWS,Plot,res_plot,Loop):
618        if Loop:
619                if (Plot == 'Prob' or Plot == 'All'):
620                        pWS = inputWS+'_Prob'
621                        p_plot=mp.plotSpectrum(pWS,[1,2],False)
622
623                if (Plot == 'FwHm' or Plot == 'All'):
624                        ilist = [1,3,5]
625                        i_plot=mp.plotSpectrum(inputWS+'_Workspace',ilist,True)
626                        i_layer = i_plot.activeLayer()
627                        i_layer.setAxisTitle(mp.Layer.Left,'Amplitude')
628                if (Plot == 'Intensity' or Plot == 'All'):
629                        wlist = [0,2,4]
630                        w_plot=mp.plotSpectrum(inputWS+'_Workspace',wlist,True)
631                        w_layer = w_plot.activeLayer()
632                        w_layer.setAxisTitle(mp.Layer.Left,'Full width half maximum (meV)')
633        if (Plot == 'Fit' or Plot == 'All'):
634                fWS = inputWS+'_Result_0'
635                f_plot=mp.plotSpectrum(fWS,res_plot,False)
636
637def QLPlotQSe(inputWS,Plot,res_plot,Loop):
638        if Loop:
639                if (Plot == 'FwHm' or Plot == 'All'):
640                        i_plot=mp.plotSpectrum(inputWS+'_Workspace',1,True)
641                        i_layer = i_plot.activeLayer()
642                        i_layer.setAxisTitle(mp.Layer.Left,'Amplitude')
643                if (Plot == 'Intensity' or Plot == 'All'):
644                        w_plot=mp.plotSpectrum(inputWS+'_Workspace',0,True)
645                        w_layer = w_plot.activeLayer()
646                        w_layer.setAxisTitle(mp.Layer.Left,'Full width half maximum (meV)')
647                if (Plot == 'Beta' or Plot == 'All'):
648                        b_plot=mp.plotSpectrum(inputWS+'_Workspace',2,True)
649                        b_layer = b_plot.activeLayer()
650                        b_layer.setAxisTitle(mp.Layer.Left,'Beta')
651        if (Plot == 'Fit' or Plot == 'All'):
652                fWS = inputWS+'_Result_0'
653                f_plot=mp.plotSpectrum(fWS,res_plot,False)
654
655# Quest programs
656def CheckBetSig(nbs):
657        Nsig = int(nbs[1])
658        if Nsig == 0:
659                error = 'Number of sigma points is Zero'                       
660                logger.notice('ERROR *** ' + error)
661                sys.exit(error)
662        if Nsig > 200:
663                error = 'Max number of sigma points is 200'                     
664                logger.notice('ERROR *** ' + error)
665                sys.exit(error)
666        Nbet = int(nbs[0])
667        if Nbet == 0:
668                error = 'Number of beta points is Zero'                 
669                logger.notice('ERROR *** ' + error)
670                sys.exit(error)
671        if Nbet > 200:
672                error = 'Max number of beta points is 200'                     
673                logger.notice('ERROR *** ' + error)
674                sys.exit(error)
675        return Nbet,Nsig
676
677def QuestRun(samWS,resWS,nbs,erange,nbins,Fit,Loop,Verbose,Plot,Save):
678        StartTime('Quest')
679        #expand fit options
680        elastic, background, width, resnorm = Fit
681       
682        #convert true/false to 1/0 for fortran
683        o_el = 1 if elastic else 0
684        o_w1 = 1 if width else 0
685        o_res = 1 if resnorm else 0
686
687        #fortran code uses background choices defined using the following numbers
688        if background == 'Sloping':
689                o_bgd = 2
690        elif background == 'Flat':
691                o_bgd = 1
692        elif background == 'Zero':
693                o_bgd = 0
694
695        fitOp = [o_el, o_bgd, o_w1, o_res]
696
697        workdir = getDefaultWorkingDirectory()
698
699        array_len = 4096                           # length of array in Fortran
700        CheckXrange(erange,'Energy')
701        nbin,nrbin = nbins[0],nbins[1]
702        if Verbose:
703                logger.notice('Sample is ' + samWS)
704                logger.notice('Resolution is ' + resWS)
705        CheckAnalysers(samWS,resWS,Verbose)
706        nsam,ntc = CheckHistZero(samWS)
707       
708        if Loop != True:
709                nsam = 1
710
711        efix = getEfixed(samWS)
712        theta,Q = GetThetaQ(samWS)
713        nres,ntr = CheckHistZero(resWS)
714        if nres == 1:
715                prog = 'Qst'                        # res file
716        else:
717                error = 'Stretched Exp ONLY works with RES file'                       
718                logger.notice('ERROR *** ' + error)
719                sys.exit(error)
720        if Verbose:
721                logger.notice(' Number of spectra = '+str(nsam))
722                logger.notice(' Erange : '+str(erange[0])+' to '+str(erange[1]))
723
724        fname = samWS[:-4] + '_'+ prog
725        wrks=os.path.join(workdir, samWS[:-4])
726        if Verbose:
727                logger.notice(' lptfile : ' + wrks +'_Qst.lpt')
728        lwrk=len(wrks)
729        wrks.ljust(140,' ')
730        wrkr=resWS
731        wrkr.ljust(140,' ')
732        wrk = [wrks, wrkr]
733        Nbet,Nsig = nbs[0], nbs[1]
734        eBet0 = np.zeros(Nbet)                  # set errors to zero
735        eSig0 = np.zeros(Nsig)                  # set errors to zero
736        rscl = 1.0
737        Qaxis = ''
738        for m in range(0,nsam):
739                if Verbose:
740                        logger.notice('Group ' +str(m)+ ' at angle '+ str(theta[m]))
741                nsp = m+1
742                nout,bnorm,Xdat,Xv,Yv,Ev = CalcErange(samWS,m,erange,nbin)
743                Ndat = nout[0]
744                Imin = nout[1]
745                Imax = nout[2]
746                Nb,Xb,Yb,Eb = GetXYE(resWS,0,array_len)
747                numb = [nsam, nsp, ntc, Ndat, nbin, Imin, Imax, Nb, nrbin, Nbet, Nsig]
748                reals = [efix, theta[m], rscl, bnorm]
749                xsout,ysout,xbout,ybout,zpout=Que.quest(numb,Xv,Yv,Ev,reals,fitOp,
750                                                                                        Xdat,Xb,Yb,wrks,wrkr,lwrk)
751                dataXs = xsout[:Nsig]               # reduce from fixed Fortran array
752                dataYs = ysout[:Nsig]
753                dataXb = xbout[:Nbet]
754                dataYb = ybout[:Nbet]
755                zpWS = fname + '_Zp' +str(m)
756                if (m > 0):
757                        Qaxis += ','
758                Qaxis += str(Q[m])
759                for n in range(0,Nsig):
760                        yfit_list = np.split(zpout[:Nsig*Nbet],Nsig)
761                        dataYzp = yfit_list[n]
762                        if n == 0:
763                                CreateWorkspace(OutputWorkspace=zpWS, DataX=xbout[:Nbet], DataY=dataYzp[:Nbet], DataE=eBet0,
764                                        Nspec=1, UnitX='MomentumTransfer')
765                        else:
766                                CreateWorkspace(OutputWorkspace='__Zpt', DataX=xbout[:Nbet], DataY=dataYzp[:Nbet], DataE=eBet0,
767                                        Nspec=1, UnitX='MomentumTransfer')
768                                ConjoinWorkspaces(InputWorkspace1=zpWS, InputWorkspace2='__Zpt', CheckOverlapping=False)                               
769                if m == 0:
770                        xSig = dataXs
771                        ySig = dataYs
772                        eSig = eSig0           
773                        xBet = dataXb
774                        yBet = dataYb
775                        eBet = eBet0           
776                        groupZ = zpWS
777                else:
778                        xSig = np.append(xSig,dataXs)
779                        ySig = np.append(ySig,dataYs)
780                        eSig = np.append(eSig,eSig0)
781                        xBet = np.append(xBet,dataXb)
782                        yBet = np.append(yBet,dataYb)
783                        eBet = np.append(eBet,eBet0)
784                        groupZ = groupZ +','+ zpWS
785        CreateWorkspace(OutputWorkspace=fname+'_Sigma', DataX=xSig, DataY=ySig, DataE=eSig,
786                Nspec=nsam, UnitX='', VerticalAxisUnit='MomentumTransfer', VerticalAxisValues=Qaxis)
787        CreateWorkspace(OutputWorkspace=fname+'_Beta', DataX=xBet, DataY=yBet, DataE=eBet,
788                Nspec=nsam, UnitX='', VerticalAxisUnit='MomentumTransfer', VerticalAxisValues=Qaxis)
789        group = fname + '_Sigma,'+ fname + '_Beta'
790        GroupWorkspaces(InputWorkspaces=group,OutputWorkspace=fname+'_Fit')     
791        GroupWorkspaces(InputWorkspaces=groupZ,OutputWorkspace=fname+'_Contour')
792
793        CopyLogs(InputWorkspace=samWS, OutputWorkspace=fname+'_Fit')
794        AddSampleLog(Workspace=fname+'_Fit', LogName="res_file", LogType="String", LogText=resWS)
795        AddSampleLog(Workspace=fname+'_Fit', LogName="background", LogType="String", LogText=str(background))
796        AddSampleLog(Workspace=fname+'_Fit', LogName="elastic_peak", LogType="String", LogText=str(elastic))
797        AddSampleLog(Workspace=fname+'_Fit', LogName="energy_min", LogType="Number", LogText=str(erange[0]))
798        AddSampleLog(Workspace=fname+'_Fit', LogName="energy_max", LogType="Number", LogText=str(erange[1]))
799        AddSampleLog(Workspace=fname+'_Fit', LogName="sample_binning", LogType="Number", LogText=str(nbin))
800        AddSampleLog(Workspace=fname+'_Fit', LogName="sigma", LogType="Number", LogText=str(Nsig))
801        AddSampleLog(Workspace=fname+'_Fit', LogName="beta", LogType="Number", LogText=str(Nbet))
802        CopyLogs(InputWorkspace=samWS, OutputWorkspace=fname+'_Contour')
803        AddSampleLog(Workspace=fname+'_Contour', LogName="res_file", LogType="String", LogText=resWS)
804        AddSampleLog(Workspace=fname+'_Contour', LogName="background", LogType="String", LogText=str(background))
805        AddSampleLog(Workspace=fname+'_Contour', LogName="elastic_peak", LogType="String", LogText=str(elastic))
806        AddSampleLog(Workspace=fname+'_Contour', LogName="energy_min", LogType="Number", LogText=str(erange[0]))
807        AddSampleLog(Workspace=fname+'_Contour', LogName="energy_max", LogType="Number", LogText=str(erange[1]))
808        AddSampleLog(Workspace=fname+'_Contour', LogName="sample_binning", LogType="Number", LogText=str(nbin))
809        AddSampleLog(Workspace=fname+'_Contour', LogName="sigma", LogType="Number", LogText=str(Nsig))
810        AddSampleLog(Workspace=fname+'_Contour', LogName="beta", LogType="Number", LogText=str(Nbet))
811
812        if Save:
813                fpath = os.path.join(workdir,fname+'_Fit.nxs')
814                SaveNexusProcessed(InputWorkspace=fname+'_Fit', Filename=fpath)
815                cpath = os.path.join(workdir,fname+'_Contour.nxs')
816                SaveNexusProcessed(InputWorkspace=fname+'_Contour', Filename=cpath)
817                if Verbose:
818                        logger.notice('Output file for Fit : ' + fpath)
819                        logger.notice('Output file for Contours : ' + cpath)
820        if (Plot != 'None'):
821                QuestPlot(fname,Plot)
822        EndTime('Quest')
823
824def QuestPlot(inputWS,Plot):
825        if (Plot == 'Sigma'):
826                sig_plot=mp.plotSpectrum(inputWS+'_Sigma',0,True)
827        if (Plot == 'Beta'):
828                beta_plot = mp.plotSpectrum(inputWS+'_Beta',0,True)
829        if(Plot == 'All'):
830                mp.plotSpectrum([inputWS+'_Sigma',inputWS+'_Beta'], 0, True)
831
832# ResNorm programs
833def ResNormRun(vname,rname,erange,nbin,Verbose=False,Plot='None',Save=False):
834        StartTime('ResNorm')
835       
836        workdir = getDefaultWorkingDirectory()
837
838        array_len = 4096                                    # length of Fortran array
839        CheckXrange(erange,'Energy')
840        CheckAnalysers(vname,rname,Verbose)
841        nvan,ntc = CheckHistZero(vname)
842        theta,Q = GetThetaQ(vname)
843        efix = getEfixed(vname)
844        nres,ntr = CheckHistZero(rname)
845        nout,bnorm,Xdat,Xv,Yv,Ev = CalcErange(vname,0,erange,nbin)
846        Ndat = nout[0]
847        Imin = nout[1]
848        Imax = nout[2]
849        wrks=os.path.join(workdir, vname[:-4])
850        if Verbose:
851                logger.notice(' Number of spectra = '+str(nvan))
852                logger.notice(' lptfile : ' + wrks +'_resnrm.lpt')
853        lwrk=len(wrks)
854        wrks.ljust(140,' ')                              # pad for fixed Fortran length
855        wrkr=rname
856        wrkr.ljust(140,' ')
857        Nb,Xb,Yb,Eb = GetXYE(rname,0,array_len)
858        rscl = 1.0
859        xPar = np.array([theta[0]])
860        for m in range(1,nvan):
861                xPar = np.append(xPar,theta[m])
862        ePar = np.zeros(nvan)
863        fname = vname[:-4]
864        for m in range(0,nvan):
865                if Verbose:
866                        logger.notice('Group ' +str(m)+ ' at angle '+ str(theta[m]))
867                ntc,Xv,Yv,Ev = GetXYE(vname,m,array_len)
868                nsp = m+1
869                numb = [nvan, nsp, ntc, Ndat, nbin, Imin, Imax, Nb]
870                reals = [efix, theta[0], rscl, bnorm]
871                nd,xout,yout,eout,yfit,pfit=resnorm.resnorm(numb,Xv,Yv,Ev,reals,
872                                                                        Xdat,Xb,Yb,wrks,wrkr,lwrk)
873                if Verbose:
874                        message = ' Fit paras : '+str(pfit[0])+' '+str(pfit[1])
875                        logger.notice(message)
876                dataX = xout[:nd]
877                dataX = np.append(dataX,2*xout[nd-1]-xout[nd-2])
878                if m == 0:
879                        yPar1 = np.array([pfit[0]])
880                        yPar2 = np.array([pfit[1]])
881                        CreateWorkspace(OutputWorkspace='Data', DataX=dataX, DataY=yout[:nd], DataE=eout[:nd],
882                                NSpec=1, UnitX='DeltaE')
883                        CreateWorkspace(OutputWorkspace='Fit', DataX=dataX, DataY=yfit[:nd], DataE=np.zeros(nd),
884                                NSpec=1, UnitX='DeltaE')
885                else:
886                        yPar1 = np.append(yPar1,pfit[0])
887                        yPar2 = np.append(yPar2,pfit[1])
888                        CreateWorkspace(OutputWorkspace='__datmp', DataX=dataX, DataY=yout[:nd], DataE=eout[:nd],
889                                NSpec=1, UnitX='DeltaE')
890                        ConjoinWorkspaces(InputWorkspace1='Data', InputWorkspace2='__datmp', CheckOverlapping=False)                           
891                        CreateWorkspace(OutputWorkspace='__f1tmp', DataX=dataX, DataY=yfit[:nd], DataE=np.zeros(nd),
892                                NSpec=1, UnitX='DeltaE')
893                        ConjoinWorkspaces(InputWorkspace1='Fit', InputWorkspace2='__f1tmp', CheckOverlapping=False)                             
894        CreateWorkspace(OutputWorkspace=fname+'_ResNorm_Intensity', DataX=xPar, DataY=yPar1, DataE=xPar,
895                NSpec=1, UnitX='MomentumTransfer')
896        CreateWorkspace(OutputWorkspace=fname+'_ResNorm_Stretch', DataX=xPar, DataY=yPar2, DataE=xPar,
897                NSpec=1, UnitX='MomentumTransfer')
898        group = fname + '_ResNorm_Intensity,'+ fname + '_ResNorm_Stretch'
899        GroupWorkspaces(InputWorkspaces=group,OutputWorkspace=fname+'_ResNorm')
900        GroupWorkspaces(InputWorkspaces='Data,Fit',OutputWorkspace=fname+'_ResNorm_Fit')
901        CopyLogs(InputWorkspace=vname, OutputWorkspace=fname+'_ResNorm')
902        AddSampleLog(Workspace=fname+'_ResNorm', LogName="energy_min", LogType="Number", LogText=str(erange[0]))
903        AddSampleLog(Workspace=fname+'_ResNorm', LogName="energy_max", LogType="Number", LogText=str(erange[1]))
904        AddSampleLog(Workspace=fname+'_ResNorm', LogName="van_binning", LogType="Number", LogText=str(nbin))
905        CopyLogs(InputWorkspace=vname, OutputWorkspace=fname+'_ResNorm_Fit')
906        AddSampleLog(Workspace=fname+'_ResNorm_Fit', LogName="energy_min", LogType="Number", LogText=str(erange[0]))
907        AddSampleLog(Workspace=fname+'_ResNorm_Fit', LogName="energy_max", LogType="Number", LogText=str(erange[1]))
908        AddSampleLog(Workspace=fname+'_ResNorm_Fit', LogName="van_binning", LogType="Number", LogText=str(nbin))
909        if Save:
910                par_path = os.path.join(workdir,fname+'_ResNorm.nxs')
911                SaveNexusProcessed(InputWorkspace=fname+'_ResNorm', Filename=par_path)
912                fit_path = os.path.join(workdir,fname+'_ResNorm_Fit.nxs')
913                SaveNexusProcessed(InputWorkspace=fname+'_ResNorm_Fit', Filename=fit_path)
914                if Verbose:
915                        logger.notice('Parameter file created : ' + par_path)
916                        logger.notice('Fit file created : ' + fit_path)
917        if (Plot != 'None'):
918                ResNormPlot(fname,Plot)
919        EndTime('ResNorm')
920
921def ResNormPlot(inputWS,Plot):
922        if (Plot == 'Intensity' or Plot == 'All'):
923                iWS = inputWS + '_ResNorm_Intensity'
924                i_plot=mp.plotSpectrum(iWS,0,False)
925        if (Plot == 'Stretch' or Plot == 'All'):
926                sWS = inputWS + '_ResNorm_Stretch'
927                s_plot=mp.plotSpectrum(sWS,0,False)
928        if (Plot == 'Fit' or Plot == 'All'):
929                fWS = inputWS + '_ResNorm_Fit'
930                f_plot=mp.plotSpectrum(fWS,0,False)