Newer
Older
from __future__ import (absolute_import, division, print_function)
import numpy as np
from mantid.api import IFunction1D, FunctionFactory
from matplotlib.mlab import bivariate_normal
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
"""
BivariateGaussian implements a bivariate gaussian (BivariateGaussian) in Mantid (M) as a 1D function. This is done so that it can be
fit in a straightforward fashion using Mantid's Fit() function. To achieve this, we use the flattened
version of the 2D profile and fit it as a 1D function. It is built on matplotlib.mlab.bivariate_normal, which
is available on SNS analysis machines.
To make it compatible with fitting, X, Y, and E must all be the same shape. This is possible if we input
twice and reconstruct the BivariateGaussian in the function.
If h is an n*n 2d profile we're trying to fitt and th, ph are the n*1 arrays containing the x and y coordinates,
we would fit it as follows:
TH, PH = np.meshgrid(th, ph,indexing='ij') #Get 2D version
m = BivariateGaussian.BivariateGaussian()
m.init()
m['A'] = 1.0
# ... #Set initial parameters
m['nX'] = len(th) #Length needed for reconstruction
m['nY'] = len(ph) #Length needed for reconstruction
pos = np.empty(TH.shape + (2,))
pos[:,:,0] = TH
pos[:,:,1] = PH
h2 = m.function2D(pos)
H = np.empty(h.shape + (2,))
H[:,:,0] = h
H[:,:,1] = h
bvgWS = CreateWorkspace(OutputWorkspace='bvgWS',DataX=pos.ravel(),DataY=H.ravel(),dataE=np.sqrt(H.ravel()))
fitResults = Fit(Function=m, InputWorkspace='bvgWS', Output='bvgfit')
BS - May 18 2018
"""
def init(self):
self.declareParameter("A") #Amplitude
self.declareParameter("MuX") #Mean along the x direction
self.declareParameter("MuY") #Mean along the y direction
self.declareParameter("SigX") #sigma along the x direction
self.declareParameter("SigY") #sigma along the y direction
self.declareParameter("SigP") #interaction term rho
self.declareParameter("Bg") #constant BG terms
self.declareAttribute("nX", 50) #used for reconstructing 2d profile
self.declareAttribute("nY", 50) #used for reconstruction 2d profile
self.addConstraints("0 < A") #Require amplitude to be positive
self.addConstraints("0 < SigX") #standard deviations must be positive
self.addConstraints("0 < SigY") #standard deviations must be positive
self.addConstraints("0 < Bg") #standard deviations must be positive
def function1D(self, t):
"""
function1D returns the flattened version of the function.
Input, t, may be in one of two forms:
1) a 1D array (e.g. pos.ravel() from the example).
2) a 3D array (e.g. pos from the example)
Output
If input is of type 1, a 1D array matching the size of the
input is returned. This allows fitting to take place.
If input is of type 2, a 1D array matching the size of
pos.ravel() is returned.
"""
if t.ndim == 1:
nX = int(self.getAttributeValue('nX'))
nY = int(self.getAttributeValue('nY'))
pos = t.reshape(nX, nY, 2)
elif t.ndim == 3:
pos = t
X = pos[...,0]
Y = pos[...,1]
A = self.getParamValue(0)
MuX = self.getParamValue(1)
MuY = self.getParamValue(2)
SigX = self.getParamValue(3)
SigY = self.getParamValue(4)
SigP = self.getParamValue(5)
Bg = self.getParamValue(6)
SigXY = SigX*SigY*SigP
Z = A*bivariate_normal(X,Y, sigmax=SigX, sigmay=SigY,
mux=MuX,muy=MuY,sigmaxy=SigXY)
if t.ndim == 1:
zRet = np.empty(Z.shape+(2,))
zRet[:,:,0] = Z
zRet[:,:,1] = Z
elif t.ndim == 3:
zRet = Z
return zRet.ravel()
def getMu(self):
MuX = self.getParamValue(1)
MuY = self.getParamValue(2)
return np.array([MuX, MuY])
def getSigma(self):
SigX = self.getParamValue(3)
SigY = self.getParamValue(4)
SigP = self.getParamValue(5)
return np.array([[SigX**2,SigX*SigY*SigP],[SigX*SigY*SigP, SigY**2]])
def getMuSigma(self):
return self.getMu(), self.getSigma()
def setConstraints(self, boundsDict):
"""
setConstraints sets fitting constraints for the mbvg function.
Intput:
boundsDict: a dictionary object where each key is a parameter as a string
('A', 'MuX', 'MuY', 'SigX', 'SigY', 'SigP', 'Bg') and the value is
an array with the lower bound and upper bound
"""
for param in boundsDict.keys():
try:
if boundsDict[param][0] < boundsDict[param][1]:
constraintString = "{:4.4e} < {:s} < {:4.4e}".format(boundsDict[param][0], param, boundsDict[param][1])
self.addConstraints(constraintString)
else:
self.log().information('Setting constraints on mbvg; reversing bounds')
self.addConstraints("{:4.4e} < A < {:4.4e}".format(boundsDict[param][1], boundsDict[param][0]))
except ValueError:
self.log().warning("Cannot set parameter {:s} for mbvg. Valid choices are " +
"('A', 'MuX', 'MuY', 'SigX', 'SigY', 'SigP', 'Bg')".format(param))
def function2D(self, t):
"""
function2D returns the 2D version of the BivariateGaussian.
Input may be in two forms:
1) 1D array (e.g. pos.ravel()). This will be reshaped into an nX*nY*2 array, so
it must contain nX*nY*2 elements.
2) 3D array of size A*B*2. A and B are arbitrary integers.
Output:
a 2D array either size nX*nY (intput type 1) or A*B (input type 2) with intensities
of the BivariateGaussian.
"""
if t.ndim == 1:
nX = int(self.getAttributeValue('nX'))
nY = int(self.getAttributeValue('nY'))
pos = t.reshape(nX, nY, 2)
elif t.ndim == 3:
pos = t
X = pos[...,0]
Y = pos[...,1]
A = self.getParamValue(0)
MuX = self.getParamValue(1)
MuY = self.getParamValue(2)
SigX = self.getParamValue(3)
SigY = self.getParamValue(4)
SigP = self.getParamValue(5)
Bg = self.getParamValue(6)
SigXY = SigX*SigY*SigP
Z = A*bivariate_normal(X,Y, sigmax=SigX, sigmay=SigY,
mux=MuX,muy=MuY,sigmaxy=SigXY)
Z += Bg
return Z
def function3D(self, t):
if t.ndim == 4:
pos = t[:,:,:,1:]
X = pos[...,0]
Y = pos[...,1]
A = self.getParamValue(0)
MuX = self.getParamValue(1)
MuY = self.getParamValue(2)
SigX = self.getParamValue(3)
SigY = self.getParamValue(4)
SigP = self.getParamValue(5)
Bg = self.getParamValue(6)
SigXY = SigX*SigY*SigP
Z = A*bivariate_normal(X,Y, sigmax=SigX, sigmay=SigY,
mux=MuX,muy=MuY,sigmaxy=SigXY)
Z += Bg
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
return Z
# Evaluate the function for a differnt set of paremeters (trialc)
def function1DDiffParams(self, xvals, trialc):
#First, grab the original parameters and set to trialc
c = np.zeros(self.numParams())
for i in range(self.numParams()):
c[i] = self.getParamValue(i)
self.setParameter(i, trialc[i])
#Get the trial values
f_trial = self.function1D(xvals)
#Now return to the orignial
for i in range(self.numParams()):
self.setParameter(i, c[i])
return f_trial
# Construction the Jacobian (df) for the function
def functionDeriv1D(self, xvals, jacobian, eps=1.e-3):
f_int = self.function1D(xvals)
#Fetch parameters into array c
c = np.zeros(self.numParams())
for i in range(self.numParams()):
c[i] = self.getParamValue(i)
nc = np.prod(np.shape(c))
for k in range(nc):
dc = np.zeros(nc)
if k == 1 or k == 2:
epsUse = 1.e-3
else:
epsUse = eps
dc[k] = max(epsUse,epsUse*c[k])
f_new = self.function1DDiffParams(xvals,c+dc)
for i,dF in enumerate(f_new-f_int):
jacobian.set(i,k,dF/dc[k])
FunctionFactory.subscribe(BivariateGaussian)