forked from AmbaPant/mantid
-
Notifications
You must be signed in to change notification settings - Fork 1
/
BivariateGaussian.py
251 lines (219 loc) · 9.8 KB
/
BivariateGaussian.py
1
2
3
4
5
6
7
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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
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
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
# Mantid Repository : https://github.com/mantidproject/mantid
#
# Copyright © 2018 ISIS Rutherford Appleton Laboratory UKRI,
# NScD Oak Ridge National Laboratory, European Spallation Source,
# Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
# SPDX - License - Identifier: GPL - 3.0 +
import numpy as np
from mantid.api import IFunction1D, FunctionFactory
def bivariate_normal(X, Y, sigmax=1.0, sigmay=1.0,
mux=0.0, muy=0.0, sigmaxy=0.0):
"""
Bivariate Gaussian distribution for equal shape *X*, *Y*.
See `bivariate normal
<http://mathworld.wolfram.com/BivariateNormalDistribution.html>`_
at mathworld.
This was copied from https://matplotlib.org/2.2.4/_modules/matplotlib/mlab.html#bivariate_normal
"""
Xmu = X - mux
Ymu = Y - muy
rho = sigmaxy / (sigmax * sigmay)
z = Xmu ** 2 / sigmax ** 2 + Ymu ** 2 / sigmay ** 2 - 2 * rho * Xmu * Ymu / (sigmax * sigmay)
denom = 2 * np.pi * sigmax * sigmay * np.sqrt(1 - rho ** 2)
return np.exp(-z / (2 * (1 - rho ** 2))) / denom
class BivariateGaussian(IFunction1D):
"""
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", 1.0) # Amplitude
self.declareParameter("MuX", 0.0) # Mean along the x direction
self.declareParameter("MuY", 0.0) # Mean along the y direction
self.declareParameter("SigX", 0.05) # sigma along the x direction
self.declareParameter("SigY", 0.05) # sigma along the y direction
self.declareParameter("SigP", 0.0) # 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'))
if nX*nY*2 != t.size:
raise ValueError(f"Input array cannot be resized as a {nX} x {nY} x 2 matrix, please modify "
"attributes nX and 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
zRet += Bg
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, penalty=None):
"""
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.addConstraints("{:4.4e} < {:s} < {:4.4e}".format(
boundsDict[param][1], param, boundsDict[param][0]))
if penalty is not None:
self.setConstraintPenaltyFactor(param, penalty)
except ValueError:
raise UserWarning("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
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)