Commit 5e73c082 authored by Viktor Reshniak's avatar Viktor Reshniak
Browse files

add MLELoss class

parent 5a6f5b82
Loading
Loading
Loading
Loading
+78 −26
Original line number Diff line number Diff line
@@ -1147,32 +1147,84 @@ class Histogram(object):
			result = least_squares(residual, #jac=jacobian_residual,
				x0=params_init, bounds=[params_lbnd,params_ubnd], method='trf', verbose=0, max_nfev=1000)
		elif loss=='mle':
			def peak_loss(params):
			gaussian_peak = Gaussian(ndims=3, covariance_parameterization=covariance_parameterization)
			class MLELoss(object):
				def __call__(self, params):
					self.func_params = np.asarray(params).copy()

					# background
					bkgr = params[0]**2

					# peak
				g,dg,d2g = gaussian_mixture(params[nbkgr:],fit_points,npeaks=1,covariance_parameterization=covariance_parameterization,return_gradient=True,return_hessian=True)
					g = gaussian_peak(params[nbkgr:], fit_points)

					# fit
					self.fit = bkgr + g

					return (self.fit-fit_data*np.log(self.fit)).sum()

				def gradient(self, params, *args, **kwargs):
					self.grad_params = np.asarray(params).copy()
					if not np.all(self.grad_params==self.func_params):  g = self.__call__(params)

					self.dloss_dfit = 1 - fit_data/self.fit

					dloss_dbkgr = [2*params[0]*self.dloss_dfit.sum()] #,0,0,0]
					dloss_dpeak = gaussian_peak.gradient(params[nbkgr:], fit_points, dloss_dfit=self.dloss_dfit)

				fit = bkgr + g.reshape(fit_data.shape)
					return np.concatenate((dloss_dbkgr,dloss_dpeak))

				# gradient of the loss
				dloss_dfit  = (1-fit_data/fit)
				dloss_dbkgr = [2*params[0]*dloss_dfit[detector_mask].sum(),0,0,0]
				dloss_dg    = (dloss_dfit.reshape((-1,1))*dg)[detector_mask.ravel(),:].sum(axis=0)
				dloss = np.concatenate((dloss_dbkgr,dloss_dg))
				def hessian(self, params, *args, **kwargs):
					self.hess_params = np.asarray(params).copy()
					if not np.all(self.hess_params==self.func_params):  g = self.__call__(params)
					if not np.all(self.hess_params==self.grad_params): dg = self.gradient(params)

				# hessian of the loss
				# d2loss_db02 = 2*(dloss_dfit.reshape((-1,1,1))[detector_mask.ravel(),...].sum(axis=0) + params[0] * (fit_data/fit**2).reshape((-1,1,1))[detector_mask.ravel(),...].sum(axis=0))
				d2loss_db02 = 2*(dloss_dfit.ravel()[detector_mask.ravel()].sum() + params[0] * (fit_data/fit**2).ravel()[detector_mask.ravel()].sum())
				d2loss_dg2  = (dloss_dfit.reshape((-1,1,1))*d2g + (fit_data/fit**2).reshape((-1,1,1))*dg[...,np.newaxis]*dg[:,np.newaxis,:])[detector_mask.ravel(),...].sum(axis=0)
					d2loss_dfit2 = fit_data/self.fit**2

				d2loss_dbkgr2 = np.array([[d2loss_db02,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,0]])
					d2loss_dbkgr2 = 2 * self.dloss_dfit.sum() + (2*params[0])**2 * d2loss_dfit2.sum()
					d2loss_dpeak2 = gaussian_peak.hessian(params[nbkgr:],fit_points,dloss_dfit=self.dloss_dfit,d2loss_dfit2=d2loss_dfit2)

					d2loss = np.zeros((len(params),len(params)))
					d2loss[:nbkgr,:nbkgr] = d2loss_dbkgr2
					d2loss[nbkgr:,nbkgr:] = d2loss_dpeak2

					return d2loss
			# @MemoizeJacHess
			# def peak_loss(params):
			# 	# background
			# 	bkgr = params[0]**2

			# 	# peak
			# 	g,dg,d2g = gaussian_mixture(params[nbkgr:],fit_points,npeaks=1,covariance_parameterization=covariance_parameterization,return_gradient=True,return_hessian=True)
			# 	# g = gaussian_mixture(params[nbkgr:],fit_points,npeaks=1,covariance_parameterization=covariance_parameterization,return_gradient=False,return_hessian=False)
			# 	print(g.shape,dg[0].shape)
			# 	exit()
			# 	fit = bkgr + g

			# 	# derivative wrt fit
			# 	dloss_dfit   = 1 - fit_data/fit
			# 	d2loss_dfit2 = fit_data/fit**2

			# 	# gradient of the loss
			# 	dloss_dbkgr = [2*params[0]*dloss_dfit.sum()]#,0,0,0]
			# 	# dloss_dpeak = np.concatenate([ (dloss_dfit.reshape((-1,1))*dgi).sum(axis=0) for dgi in dg ])
			# 	# dloss = np.concatenate((dloss_dbkgr,dloss_dpeak))
			# 	# dloss_dpeak = (dloss_dfit.reshape((-1,1))*dg).sum(axis=0)
			# 	dloss_dpeak = np.einsum('k,ki->i',dloss_dfit,dg)
			# 	dloss = np.concatenate((dloss_dbkgr,dloss_dpeak))
			# 	print(dloss)
			# 	exit()

				d2loss = np.block([[d2loss_dbkgr2, np.zeros((nbkgr,len(params)-nbkgr))],[np.zeros((len(params)-nbkgr,nbkgr)), d2loss_dg2]])
			# 	# hessian of the loss
			# 	d2loss = np.zeros((len(params),len(params)))
			# 	d2loss_dbkgr2 = 2 * dloss_dfit.sum() + (2*params[0])**2 * d2loss_dfit2.sum()
			# 	d2loss_dpeak2 = (dloss_dfit.reshape((-1,1,1))*d2g + (d2loss_dfit2.reshape((-1,1))*dg)[...,np.newaxis]*dg[:,np.newaxis,:]).sum(axis=0)
			# 	d2loss[:nbkgr,:nbkgr] = d2loss_dbkgr2
			# 	d2loss[nbkgr:,nbkgr:] = d2loss_dpeak2

				return (fit-fit_data*np.log(fit))[detector_mask].sum(), dloss, d2loss
			peak_loss = MemoizeJacHess(peak_loss)
			# 	return (fit-fit_data*np.log(fit)).sum(), dloss, d2loss #None, None
			# peak_loss = MemoizeJacHess(peak_loss)
			peak_loss = MLELoss()
			result = minimize(peak_loss,
				jac  = peak_loss.derivative, #grad,
				hess = peak_loss.hessian, #hess,