Ticket #9086: WrappedReduction.py

File WrappedReduction.py, 14.4 KB (added by Owen Arnold, 7 years ago)
Line 
1import numpy
2import sys, os
3import dis, inspect, opcode
4import matplotlib.pyplot as plt
5from mantid.simpleapi import *
6
7############################################################################################
8def lineno():
9    """
10    call signature(s)::
11    lineno()
12
13    Returns the current line number in our program.
14   
15    No Arguments.
16
17
18    Working example
19    >>> print "This is the line number ",lineno(),"\n"
20   
21    """
22    return inspect.currentframe().f_back.f_lineno
23
24def decompile(code_object):
25        '''     taken from http://thermalnoise.wordpress.com/2007/12/30/exploring-python-bytecode/
26
27        decompile extracts dissasembly information from the byte code and stores it in a
28                list for further use.
29       
30        call signature(s)::
31                instructions=decompile(f.f_code)
32
33        Required arguments:
34        =========   =====================================================================
35        f.f_code    A  bytecode object ectracted with inspect.currentframe()
36                    or anyother mechanism that returns byte code.   
37       
38        Optional keyword arguments: NONE       
39
40        Outputs:
41        =========   =====================================================================
42        instructions  a list of offsets, op_codes, names, arguments, argument_type,
43                        argument_value which can be deconstructed to find out various things
44                        about a function call.
45
46        Examples:
47       
48        f = inspect.currentframe().f_back.f_back
49        i = f.f_lasti  # index of the last attempted instruction in byte code
50        ins=decompile(f.f_code)
51        pretty_print(ins)
52       
53
54        '''
55        code = code_object.co_code
56        variables = code_object.co_cellvars + code_object.co_freevars
57        instructions = []
58        n = len(code)
59        i = 0
60        e = 0
61        while i < n:
62                i_offset = i
63                i_opcode = ord(code[i])
64                i = i + 1
65                if i_opcode >= opcode.HAVE_ARGUMENT:
66                        i_argument = ord(code[i]) + (ord(code[i+1]) << (4*2)) + e
67                        i = i +2
68                        if i_opcode == opcode.EXTENDED_ARG:
69                                e = iarg << 16
70                        else:
71                                e = 0
72                        if i_opcode in opcode.hasconst:
73                                i_arg_value = repr(code_object.co_consts[i_argument])
74                                i_arg_type = 'CONSTANT'
75                        elif i_opcode in opcode.hasname:
76                                i_arg_value = code_object.co_names[i_argument]
77                                i_arg_type = 'GLOBAL VARIABLE'
78                        elif i_opcode in opcode.hasjrel:
79                                i_arg_value = repr(i + i_argument)
80                                i_arg_type = 'RELATIVE JUMP'
81                        elif i_opcode in opcode.haslocal:
82                                i_arg_value = code_object.co_varnames[i_argument]
83                                i_arg_type = 'LOCAL VARIABLE'
84                        elif i_opcode in opcode.hascompare:
85                                i_arg_value = opcode.cmp_op[i_argument]
86                                i_arg_type = 'COMPARE OPERATOR'
87                        elif i_opcode in opcode.hasfree:
88                                i_arg_value = variables[i_argument]
89                                i_arg_type = 'FREE VARIABLE'
90                        else:
91                                i_arg_value = i_argument
92                                i_arg_type = 'OTHER'
93                else:
94                        i_argument = None
95                        i_arg_value = None
96                        i_arg_type = None
97                instructions.append( (i_offset, i_opcode, opcode.opname[i_opcode], i_argument, i_arg_type, i_arg_value) )
98        return instructions
99 
100# Print the byte code in a human readable format
101def pretty_print(instructions):
102        print '%5s %-20s %3s  %5s  %-20s  %s' %  ('OFFSET', 'INSTRUCTION', 'OPCODE', 'ARG', 'TYPE', 'VALUE')
103        for (offset, op, name, argument, argtype, argvalue) in instructions:
104                print '%5d  %-20s (%3d)  ' % (offset, name, op),
105                if argument != None:
106                        print '%5d  %-20s  (%s)' % (argument, argtype, argvalue),
107                print
108
109
110def lhs(output='names'):
111        '''     
112        call signature(s)::
113        NamesofOutPutsExpected    = lhs(output='names')   # This is the default
114        NumberOfOutputsExpected  = lhs('number')
115        Number, Names                    = lhs('both')
116
117        Return how many values the caller is expecting or the names of the output variables on the left of the = in a list or both.
118
119        Required arguments:     NONE
120       
121        Optional keyword arguments: NONE       
122
123
124        Outputs:
125        =========   =====================================================================
126        numReturns      Number of return values on expected on the left of the equal sign.
127
128        Examples:
129
130        This function is not designed for cammand line use.  Using in a function can
131        follow the form below.
132       
133        def MyFunction():
134                """Documentation for MyFunction
135                """
136                ReturnVarNumber, ReturnVarNames=lhs('both')
137                ### Do something useful
138                single = 1, double_1 = 1, double_2=2
139               
140                if ReturnVarNumber == 0:
141                        # process like a void return
142                        return
143                elif  ReturnVarNumber == 1:
144                        # return a single output
145                        return single
146                elif ReturnVarNumber == 2:
147                        # retrun a pair
148                        return double_1, double_2
149                       
150        x     = MyFunction() # returns x=1
151        x, y = MyFunction() # returns x=1, y=2
152
153        '''
154        f = inspect.currentframe().f_back.f_back
155        i = f.f_lasti  # index of the last attempted instruction in byte code
156        ins=decompile(f.f_code)
157        #pretty_print(ins)
158
159        CallFunctionLocation={}
160        first=False; StartIndex=0; StartOffset=0
161        # we must list all of the operators that behave like a function call in byte-code
162        OperatorNames=set(['CALL_FUNCTION','UNARY_POSITIVE','UNARY_NEGATIVE','UNARY_NOT','UNARY_CONVERT','UNARY_INVERT','GET_ITER', 'BINARY_POWER','BINARY_MULTIPLY','BINARY_DIVIDE', 'BINARY_FLOOR_DIVIDE', 'BINARY_TRUE_DIVIDE', 'BINARY_MODULO','BINARY_ADD','BINARY_SUBTRACT','BINARY_SUBSCR','BINARY_LSHIFT','BINARY_RSHIFT','BINARY_AND','BINARY_XOR','BINARY_OR'])
163
164        for index in range(len(ins)):
165                (offset, op, name, argument, argtype, argvalue) = ins[index]
166                if name in OperatorNames:
167                        if not first:
168                                CallFunctionLocation[StartOffset] = (StartIndex,index)
169                        StartIndex=index
170                        StartOffset = offset
171
172        (offset, op, name, argument, argtype, argvalue) = ins[-1]
173        CallFunctionLocation[StartOffset]=(StartIndex,len(ins)-1) # append the index of the last entry to form the last boundary
174
175        #print CallFunctionLocation
176        #pretty_print( ins[CallFunctionLocation[i][0]:CallFunctionLocation[i][1]] )
177        # In our case i should always be the offset of a Call_Function instruction. We can use this to baracket
178        # the bit which we are interested in
179       
180        OutputVariableNames=[]
181        (offset, op, name, argument, argtype, argvalue) = ins[CallFunctionLocation[i][0] + 1] 
182        if name == 'POP_TOP':  # no Return Values
183                pass
184                #return OutputVariableNames
185        if name == 'STORE_FAST' or name == 'STORE_NAME': # One Return Value
186                OutputVariableNames.append(argvalue)
187        if name == 'UNPACK_SEQUENCE': # Many Return Values, One equal sign
188                for index in range(argvalue):
189                        (offset_, op_, name_, argument_, argtype_, argvalue_) = ins[CallFunctionLocation[i][0] + 1 + 1 +index] 
190                        OutputVariableNames.append(argvalue_)
191        maxReturns = len(OutputVariableNames)
192        if name == 'DUP_TOP': # Many Return Values, Many equal signs
193                # The output here should be a multi-dim list which mimics the variable unpacking sequence.
194                # For instance a,b=c,d=f() => [ ['a','b'] , ['c','d'] ]
195                #              a,b=c=d=f() => [ ['a','b'] , 'c','d' ]  So on and so forth.
196
197                # put this in a loop and stack the results in an array.
198                count = 0; maxReturns = 0 # Must count the maxReturns ourselves in this case
199                while count < len(ins[CallFunctionLocation[i][0] :CallFunctionLocation[i][1]]):
200                        (offset_, op_, name_, argument_, argtype_, argvalue_) = ins[CallFunctionLocation[i][0]+count] 
201                        #print 'i= ',i,'count = ', count, 'maxReturns = ',maxReturns
202                        if name_ == 'UNPACK_SEQUENCE': # Many Return Values, One equal sign
203                                hold=[]
204                                #print 'argvalue_ = ', argvalue_, 'count = ',count
205                                if argvalue_ > maxReturns:
206                                        maxReturns=argvalue_
207                                for index in range(argvalue_):
208                                        (_offset_, _op_, _name_, _argument_, _argtype_, _argvalue_) = ins[CallFunctionLocation[i][0] + count+1+index] 
209                                        hold.append(_argvalue_)
210                                count = count + argvalue_
211                                OutputVariableNames.append(hold)
212                        # Need to now skip the entries we just appended with the for loop.
213                        if name_ == 'STORE_FAST' or name_ == 'STORE_NAME': # One Return Value
214                                if 1 > maxReturns:
215                                        maxReturns = 1
216                                OutputVariableNames.append(argvalue_)
217                        count = count + 1
218
219        # Now that OutputVariableNames is filled with the right stuff we need to output the correct thing. Either the maximum number of
220        # variables to unpack in the case of multiple ='s or just the length of the array or just the naames of the variables.         
221       
222        if output== 'names':
223                return OutputVariableNames
224        elif output == 'number':
225                return maxReturns
226        elif output == 'both':
227                return (maxReturns,OutputVariableNames)
228       
229        return 0 # Should never get to here
230
231############################################################################################
232def centerbins(xvals):
233        ''' for a given numpy array return the bin centers
234        '''
235        NewXvals=( xvals + numpy.roll(xvals,-1) )/2 # calculate the bin center
236        return numpy.delete(NewXvals,-1) # remove the last element which is junk
237       
238
239def CorrectMonitor(TheWorkspaceToCorrect):
240        ''' Correct the monitors so that I/Io would return a flat line if the detector was in the direct beam
241        '''
242        WorkspaceToCorrect=CloneWorkspace(TheWorkspaceToCorrect)
243        CorrectedOutput=0
244        if WorkspaceToCorrect.isMultiPeriod():
245                instrument=WorkspaceToCorrect[0].getInstrument()
246        else:
247                instrument=WorkspaceToCorrect.getInstrument()
248
249        CorrectionType=instrument.getStringParameter('correction')[0]
250        if CorrectionType == 'polynomial':
251                CorrectedOutput=PolynomialCorrection(InputWorkspace=WorkspaceToCorrect, Coefficients=instrument.getStringParameter('polystring')[0], Operation='Multiply') #Operation='Divide'
252        elif CorrectionType == 'exponential':
253                CorrectedOutput=ExponentialCorrection(InputWorkspace=WorkspaceToCorrect,C0=instrument.getNumberParameter('C0')[0], C1=instrument.getNumberParameter('C1')[0], Operation='Multiply')#Operation='Divide'
254
255        return CorrectedOutput
256
257def CorrectPolarization(TheDataIn):
258        '''
259        This function takes a two period data set or four period data set and corrects for polarization efficiencies.
260        '''
261        # Make a copy of the data to work on locally
262        DataIn=CloneWorkspace(TheDataIn)
263        # Check the the workspace is in wavelength
264       
265        if DataIn.isMultiPeriod():
266                instrument=DataIn[0].getInstrument()
267        else:
268                instrument=DataIn.getInstrument()
269
270        # Calculate the required polynomials
271        Ones=DataIn[1] * 0.0 + 1.0 # The polynomial correction function operates on a workspace of ones. 
272                                   # In the case of multidetector data we expect the monitors have been stripped away.   
273        rho    = PolynomialCorrection(InputWorkspace=Ones,Coefficients=instrument.getStringParameter('crho')[0])
274        Pp     = PolynomialCorrection(InputWorkspace=Ones,Coefficients=instrument.getStringParameter('cPp')[0])
275        if(len(DataIn) == 2):# R+ and R-
276                # split out the spin-states and name them according to the journal article
277                Ip=DataIn[0];  Ia=DataIn[1]
278                #print "Ip,Ia ",Ip, Ia
279                D = Pp * ( 1.0 + rho) 
280                nIp = ( Ip * (rho * Pp + 1.0) + Ia * (Pp - 1.0) ) / D
281                nIa = ( Ip * (rho * Pp - 1.0) + Ia * (Pp + 1.0) ) / D
282                #print "nIp,nIa ",nIp, nIa
283                DataOut=GroupWorkspaces([nIp, nIa])
284                return DataOut
285
286        elif(len(DataIn)== 4):# R++ , R+- , R-+ , R--
287               
288                alpha  = PolynomialCorrection(InputWorkspace=Ones,Coefficients=instrument.getStringParameter('calpha')[0])
289                Ap     = PolynomialCorrection(InputWorkspace=Ones,Coefficients=instrument.getStringParameter('cAp')[0])
290
291                # split out the spin-states and name them according to the journal article
292                Ipp=DataIn[0];  Iaa=DataIn[1]; Ipa=DataIn[2]; Iap=DataIn[3]
293               
294                ############## Precalculate parts of the correction equation
295                A0 = Iaa * Pp * Ap + Ap * Ipa * rho * Pp + Ap * Iap * Pp * alpha + Ipp * Ap * alpha * rho * Pp
296                A1 = Pp * Iaa
297                A2 = Pp * Iap
298                A3 = Ap * Iaa
299                A4 = Ap * Ipa
300                A5 = Ap * alpha * Ipp
301                A6 = Ap * alpha * Iap
302                A7 = Pp * rho   * Ipp
303                A8 = Pp * rho   * Ipa
304
305                D = Pp * Ap *( 1.0 + rho + alpha + rho * alpha ) 
306
307                ############## Build the corrections    # Note the mathematical symmetry
308                nIpp = (A0 - A1 + A2 - A3 + A4 + A5 - A6 + A7 - A8 + Ipp + Iaa - Ipa - Iap) / D
309                nIaa = (A0 + A1 - A2 + A3 - A4 - A5 + A6 - A7 + A8 + Ipp + Iaa - Ipa - Iap) / D
310                nIpa = (A0 - A1 + A2 + A3 - A4 - A5 + A6 + A7 - A8 - Ipp - Iaa + Ipa + Iap) / D
311                nIap = (A0 + A1 - A2 - A3 + A4 + A5 - A6 - A7 + A8 - Ipp - Iaa + Ipa + Iap) / D
312
313                DataOut=GroupWorkspaces([nIpp, nIaa, nIpa, nIap])
314                return DataOut
315        else: 
316                print len(DataIn), ' This function only operates on 2 or 4 period data' 
317                return DataIn; # Should scream with an error here.  We can only correct data which has 2 or 4 spin-states
318
319def ref(RunNumbers='8169+8171+8173+8175+8177+8179+8181', theta=1.0):
320        '''Basic shorthand for the default point detector reduction.
321       
322          Example:
323       
324         >>>  rqz, rLambda , ThetaOut= ref(RunNumbers='8169+8171+8173+8175+8177+8179+8181', theta=1.0)
325        '''
326        ReturnVarNumber, ReturnVarNames=lhs('both')
327        dd=Load(RunNumbers); 
328        trans=CreateTransmissionWorkspaceAuto(FirstTransmissionRun=dd,I0MonitorIndex='2',ProcessingInstructions='2')
329        #Now correct for the monitor to detector efficiency and optionally correct the monitors to deal with the polarization corrections too
330        trans=CorrectMonitor(trans)
331        (_,rl,ThetaOut)=ReflectometryReductionOneAuto(InputWorkspace=dd,I0MonitorIndex='2',ProcessingInstructions='3',ThetaIn=theta,FirstTransmissionRun=trans,StrictSpectrumChecking=False)
332       
333        RcorrectedVsLambda = CorrectPolarization(rl)
334        CloneWorkspace(InputWorkspace=RcorrectedVsLambda,OutputWorkspace=ReturnVarNames[1])
335       
336        RcorrectedVsQz = RcorrectedVsLambda.convertUnits('MomentumTransfer')
337        CloneWorkspace(InputWorkspace=RcorrectedVsQz,OutputWorkspace=ReturnVarNames[0])
338       
339        return mtd[ReturnVarNames[0]], mtd[ReturnVarNames[1]] , ThetaOut
340
341def refq(RunNumbers='8169+8171+8173+8175+8177+8179+8181', theta=1.0):
342        '''Outputs I/Io vs qz.  See ref for more documentation.
343       
344          Example:
345          >>> xx=refq(RunNumbers='8169+8171+8173+8175+8177+8179+8181', theta=1.0)
346        '''
347        rq,rl,ThetaOut=ref(RunNumbers, theta)
348        # Ensure that the desired output is coppied and named correctly.  This is required since Mantid is not completely pass by value.
349        ReturnVarNumber, ReturnVarNames=lhs('both')
350        CloneWorkspace(InputWorkspace=rq,OutputWorkspace=ReturnVarNames[ReturnVarNumber-1])
351        return mtd[ReturnVarNames[ReturnVarNumber-1]]
352
353def refl(RunNumbers='8169+8171+8173+8175+8177+8179+8181', theta=1.0):
354        '''Outputs I/Io vs lambda. See ref for more documentation.
355       
356          Example: 
357          >>> xx = refl(RunNumbers='8169+8171+8173+8175+8177+8179+8181')
358        '''
359        rq,rl,ThetaOut=ref(RunNumbers, theta)
360        # Ensure that the desired output is coppied and named correctly.  This is required since Mantid is not completely pass by value.
361        ReturnVarNumber, ReturnVarNames=lhs('both')
362        CloneWorkspace(InputWorkspace=rq,OutputWorkspace=ReturnVarNames[ReturnVarNumber-1])
363        return mtd[ReturnVarNames[ReturnVarNumber-1]]
364