def testhessdiag(self):
fun = lambda x: x[0] + x[1]**2 + x[2]**3
Hfun = nd.Hessdiag(fun)
hd = Hfun([1, 2, 3])
htrue = [0., 2., 18.]
for (hi, hit) in zip(hd, htrue):
self.assertAlmostEqual(hi, hit)
def _calculate_univariate_uncertainty(parametrization, alpha, totalhet, num_snps, num_samples):
NCKoef = 0.319 # this koef gives proportion of causal variants that explain 90% of heritability.
# it is specific to BGMG with single gaussian, with MAF specific model
funcs = [('pi', lambda x: x._pi),
('nc', lambda x: x._pi * num_snps),
('nc@p9', lambda x: x._pi * num_snps * NCKoef),
('sig2_beta', lambda x: x._sig2_beta),
('sig2_zero', lambda x: x._sig2_zero),
('h2', lambda x: x._sig2_beta * x._pi * totalhet)]
stats = [('mean', lambda x: np.mean(x)),
('median', lambda x: np.median(x)),
('std', lambda x: np.std(x)),
('lower', lambda x: np.percentile(x, 100.0 * ( alpha/2))),
('upper', lambda x: np.percentile(x, 100.0 * (1-alpha/2)))]
hessian = _hessian_robust(nd.Hessian(parametrization._calc_cost)(parametrization._init_vec),
nd.Hessdiag(parametrization._calc_cost)(parametrization._init_vec))
x_sample = np.random.multivariate_normal(parametrization._init_vec, np.linalg.inv(hessian), num_samples)
sample = [parametrization._vec_to_params(x) for x in x_sample]
result = {}
for func_name, func in funcs:
result[func_name] = {'point_estimate': func(parametrization._vec_to_params(parametrization._init_vec))}
param_vector = [func(s) for s in sample]
for stat_name, stat in stats:
result[func_name][stat_name] = stat(param_vector)
return result, sample
def _calculate_bivariate_uncertainty(parametrization, ci_samples, alpha, totalhet, num_snps, num_samples):
NCKoef = 0.319 # this koef gives proportion of causal variants that explain 90% of heritability.
# it is specific to BGMG with single gaussian, with MAF specific model
funcs = [('sig2_zero_T1', lambda x: x._sig2_zero[0]),
('sig2_zero_T2', lambda x: x._sig2_zero[1]),
('sig2_beta_T1', lambda x: x._sig2_beta[0]),
('sig2_beta_T2', lambda x: x._sig2_beta[1]),
('h2_T1', lambda x: x._sig2_beta[0] * (x._pi[0] + x._pi[2]) * totalhet),
('h2_T2', lambda x: x._sig2_beta[1] * (x._pi[1] + x._pi[2]) * totalhet),
('rho_zero', lambda x: x._rho_zero),
('rho_beta', lambda x: x._rho_beta),
('rg', lambda x: x._rho_beta * x._pi[2] / np.sqrt((x._pi[0] + x._pi[2]) * (x._pi[1] + x._pi[2]))),
('pi1', lambda x: x._pi[0]),
('pi2', lambda x: x._pi[1]),
('pi12', lambda x: x._pi[2]),
('pi1u', lambda x: x._pi[0] + x._pi[2]),
('pi2u', lambda x: x._pi[1] + x._pi[2]),
('nc1', lambda x: num_snps * x._pi[0]),
('nc2', lambda x: num_snps * x._pi[1]),
('nc12', lambda x: num_snps * x._pi[2]),
('nc1u', lambda x: num_snps * (x._pi[0] + x._pi[2])),
('nc2u', lambda x: num_snps * (x._pi[1] + x._pi[2])),
('nc1@p9', lambda x: NCKoef * num_snps * x._pi[0]),
('nc2@p9', lambda x: NCKoef * num_snps * x._pi[1]),
('nc12@p9', lambda x: NCKoef * num_snps * x._pi[2]),
('nc1u@p9', lambda x: NCKoef * num_snps * (x._pi[0] + x._pi[2])),
('nc2u@p9', lambda x: NCKoef * num_snps * (x._pi[1] + x._pi[2])),
('totalpi', lambda x: np.sum(x._pi)),
('totalnc', lambda x: num_snps * np.sum(x._pi)),
('totalnc@p9', lambda x: NCKoef * num_snps * np.sum(x._pi)),
('pi1_over_totalpi', lambda x: x._pi[0] / np.sum(x._pi)),
('pi2_over_totalpi', lambda x: x._pi[1] / np.sum(x._pi)),
('pi12_over_totalpi', lambda x: x._pi[2] / np.sum(x._pi)),
('pi1_over_pi1u', lambda x: x._pi[0] / (x._pi[0] + x._pi[2])),
('pi2_over_pi2u', lambda x: x._pi[1] / (x._pi[1] + x._pi[2])),
('pi12_over_pi1u', lambda x: x._pi[2] / (x._pi[0] + x._pi[2])),
('pi12_over_pi2u', lambda x: x._pi[2] / (x._pi[1] + x._pi[2])),
('pi1u_over_pi2u', lambda x: (x._pi[0] + x._pi[2]) / (x._pi[1] + x._pi[2])),
('pi2u_over_pi1u', lambda x: (x._pi[1] + x._pi[2]) / (x._pi[0] + x._pi[2]))]
stats = [('mean', lambda x: np.mean(x)),
('median', lambda x: np.median(x)),
('std', lambda x: np.std(x)),
('lower', lambda x: np.percentile(x, 100.0 * ( alpha/2))),
('upper', lambda x: np.percentile(x, 100.0 * (1-alpha/2)))]
hessian = _hessian_robust(nd.Hessian(parametrization._calc_cost)(parametrization._init_vec),
nd.Hessdiag(parametrization._calc_cost)(parametrization._init_vec))
x_sample = np.random.multivariate_normal(parametrization._init_vec, np.linalg.inv(hessian), num_samples)
sample = [parametrization._vec_to_params(x, params1=ci_s1, params2=ci_s2) for ci_s1, ci_s2, x in zip(ci_samples[0], ci_samples[1], x_sample)]
result = {}
for func_name, func in funcs:
result[func_name] = {'point_estimate': func(parametrization._vec_to_params(parametrization._init_vec))}
param_vector = [func(s) for s in sample]
for stat_name, stat in stats:
result[func_name][stat_name] = stat(param_vector)
return result, sample
def laplacian(self):
"""
The laplacian for a scalar field is defined as the
divergence of the of the gradiant of the scalar
field, and is written in cartesian coordinates
Lap(f) = D^2_x f + D^2_y f + D^2_z f + ...
"""
return sum(nd.Hessdiag(self.f)(self.x))
def compute_pseudopot_frequencies(self, r):
'''
This is only valid if xp, yp, zp is the trapping position. Return frequency (i.e. omega/(2*pi))
'''
hessdiag = nd.Hessdiag(self.compute_pseudopot,
step=1e-6)(r) / (self.__scale**2)
'''
Now d2Udx2 has units of J/m^2. Then w = sqrt(d2Udx2/(mass)) has units of angular frequency
'''
return np.sqrt(qe * abs(hessdiag) / self.m) / (2 * np.pi), hessdiag > 0
def numerical_hessian(func: Callable | None,
params: Iterable[zfit.Parameter],
hessian=None) -> tf.Tensor:
"""Calculate numerically the hessian matrix of func with respect to `params`.
Args:
func: Function without arguments that depends on `params`
params: Parameters that `func` implicitly depends on and with respect to which the
derivatives will be taken.
Returns:
Hessian matrix
"""
from ..core.parameter import assign_values
params = convert_to_container(params)
def wrapped_func(param_values):
assign_values(params, param_values)
value = func()
if hasattr(value, "numpy"):
value = value.numpy()
return value
param_vals = znp.stack(params)
original_vals = [param.value() for param in params]
if hessian == "diag":
hesse_func = numdifftools.Hessdiag(
wrapped_func,
order=2,
# TODO: maybe add step to remove numerical problems?
base_step=1e-4,
)
else:
hesse_func = numdifftools.Hessian(
wrapped_func,
order=2,
base_step=1e-4,
)
if tf.executing_eagerly():
computed_hessian = convert_to_tensor(hesse_func(param_vals))
else:
computed_hessian = tf.numpy_function(hesse_func,
inp=[param_vals],
Tout=tf.float64)
n_params = param_vals.shape[0]
if hessian == "diag":
computed_hessian.set_shape((n_params, ))
else:
computed_hessian.set_shape((n_params, n_params))
assign_values(params, original_vals)
return computed_hessian
def get_Hessian(self, One_Dimension=True):
if One_Dimension:
f = nd.Derivative(partial(self.hmmtract.objective_1D, False), n=1)
ff = nd.Derivative(partial(self.hmmtract.objective_1D, False), n=2)
grad = -f(self.hmmtract.x[1:])
hess = -ff(self.hmmtract.x[1:])
else:
f = nd.Derivative(self._loglikelihood, n=1)
ff = nd.Hessdiag(self._loglikelihood)
grad = -f(self.x)
hess = -ff(self.x)
return grad, hess
def numerical_hessian(func: Callable,
params: Iterable["zfit.Parameter"],
hessian=None) -> tf.Tensor:
"""Calculate numerically the hessian matrix of func with respect to `params`.
Args:
func: Function without arguments that depends on `params`
params: Parameters that `func` implicitly depends on and with respect to which the
derivatives will be taken.
Returns:
Hessian matrix
"""
params = convert_to_container(params)
def wrapped_func(param_values):
for param, value in zip(params, param_values):
param.assign(value)
return func().numpy()
param_vals = tf.stack(params)
original_vals = [param.read_value() for param in params]
if hessian == 'diag':
hesse_func = numdifftools.Hessdiag(
wrapped_func,
# TODO: maybe add step to remove numerical problems?
# step=1e-4
)
else:
hesse_func = numdifftools.Hessian(wrapped_func,
# base_step=1e-4
)
computed_hessian = tf.py_function(hesse_func,
inp=[param_vals],
Tout=tf.float64)
n_params = param_vals.shape[0]
if hessian == 'diag':
computed_hessian.set_shape((n_params, ))
else:
computed_hessian.set_shape((n_params, n_params))
for param, val in zip(params, original_vals):
param.set_value(val)
return computed_hessian
def compute_pseudopot_frequencies(self, r):
'''
This is only valid if xp, yp, zp is the trapping position. Return frequency (i.e. 2*pi*omega)
'''
ev_to_joule = 1.60217657e-19
m = 6.64215568e-26 # 40 amu in kg
hessdiag = nd.Hessdiag(self.compute_pseudopot)(r)
d2Udx2 = ev_to_joule * hessdiag[0]
d2Udy2 = ev_to_joule * hessdiag[1]
d2Udz2 = ev_to_joule * hessdiag[2]
'''
Now d2Udx2 has units of J/m^2. Then w = sqrt(d2Udx2/(mass)) has units of angular frequency
'''
fx = np.sqrt(abs(d2Udx2) / m) / (2 * np.pi)
fy = np.sqrt(abs(d2Udy2) / m) / (2 * np.pi)
fz = np.sqrt(abs(d2Udz2) / m) / (2 * np.pi)
return [fx, fy, fz]
def _calculate_univariate_uncertainty(parametrization, alpha, totalhet,
num_snps, num_samples):
funcs, stats = _calculate_univariate_uncertainty_funcs(
alpha, totalhet, num_snps)
hessian = _hessian_robust(
nd.Hessian(parametrization._calc_cost)(parametrization._init_vec),
nd.Hessdiag(parametrization._calc_cost)(parametrization._init_vec))
x_sample = np.random.multivariate_normal(parametrization._init_vec,
np.linalg.inv(hessian),
num_samples)
sample = [parametrization._vec_to_params(x) for x in x_sample]
result = {}
for func_name, func in funcs:
result[func_name] = {
'point_estimate':
func(parametrization._vec_to_params(parametrization._init_vec))
}
param_vector = [func(s) for s in sample]
for stat_name, stat in stats:
result[func_name][stat_name] = stat(param_vector)
return result, sample
def laplacian(f, x):
"""Laplacian of scalar field f evaluated at x"""
hes = nd.Hessdiag(f)(x)
return sum(hes)
def get_dense_nuts_step(
start=None,
adaptation_window=101,
doubling=True,
initial_weight=10,
use_hessian=False,
use_hessian_diag=False,
hessian_regularization=1e-8,
model=None,
**kwargs,
):
"""Get a NUTS step function with a dense mass matrix
The entries in the mass matrix will be tuned based on the sample
covariances during tuning. All extra arguments are passed directly to
``pymc3.NUTS``.
Args:
start (dict, optional): A starting point in parameter space. If not
provided, the model's ``test_point`` is used.
adaptation_window (int, optional): The (initial) size of the window
used for sample covariance estimation.
doubling (bool, optional): If ``True`` (default) the adaptation window
is doubled each time the matrix is updated.
"""
model = modelcontext(model)
if not all_continuous(model.vars):
raise ValueError("NUTS can only be used for models with only "
"continuous variables.")
if start is None:
start = model.test_point
mean = model.dict_to_array(start)
if use_hessian or use_hessian_diag:
try:
import numdifftools as nd
except ImportError:
raise ImportError(
"The 'numdifftools' package is required for Hessian "
"computations")
logger.info("Numerically estimating Hessian matrix")
if use_hessian_diag:
hess = nd.Hessdiag(model.logp_array)(mean)
var = np.diag(-1.0 / hess)
else:
hess = nd.Hessian(model.logp_array)(mean)
var = -np.linalg.inv(hess)
factor = 1
success = False
while not success:
var[np.diag_indices_from(var)] += factor * hessian_regularization
try:
np.linalg.cholesky(var)
except np.linalg.LinAlgError:
factor *= 2
else:
success = True
else:
var = np.eye(len(mean))
potential = QuadPotentialDenseAdapt(
model.ndim,
mean,
var,
initial_weight,
adaptation_window=adaptation_window,
doubling=doubling,
)
return pm.NUTS(potential=potential, model=model, **kwargs)
def hess(self, s, p=0):
''' Return hessian diagonal approximation. nd.Hessian takes long time. In testing so far
Hessdiag is an OOM faster and works just as good if not better.
'''
return np.diag(
nd.Hessdiag(lambda x: self.u(x, 0))(s.reshape(len(self))))
def OutlierHMC(pintpsr, outdir='.', Nsamples=20000, Nburnin=1000):
"""Run full Hamiltonian Monte Carlo outlier analysis for given pulsar
:param pintpsr: enterprise PintPulsar object
:param outdir: Desired output directory for chains and
plots, default is current working directory
:param Nsamples: Number of samples to generate with HMC, default is 20000
:param Nburnin: Number of samples for HMC burn-in phase, default is 1000
"""
# Extract pulsar name and load Interval Likelihood object
psr = pintpsr.name
likob = itvl.Interval(pintpsr)
# Now get an approximate maximum of the posterior
def func(x):
ll, _ = likob.full_loglikelihood_grad(x)
return -np.inf if np.isnan(ll) else ll
def jac(x):
_, j = likob.full_loglikelihood_grad(x)
return j
# Ensure full path to outdir exists
os.makedirs(outdir, exist_ok=True)
# Compute the approximate max and save to a pickle file
endpfile = f'{outdir}/{psr}-endp.pickle'
if not os.path.isfile(endpfile):
endp = likob.pstart
for iter in range(3):
res = so.minimize(lambda x: -func(x),
endp,
jac=lambda x: -jac(x),
hess=None,
method='L-BFGS-B', options={'disp': True})
endp = res['x']
pickle.dump(endp,open(endpfile,'wb'))
else:
endp = pickle.load(open(endpfile,'rb'))
# Next to whiten the likelihood, compute the Hessian of the posterior
nhyperpars = likob.ptadict[likob.pname + '_outlierprob'] + 1
hessfile = f'{outdir}/{psr}-fullhessian.pickle'
if not os.path.isfile(hessfile):
reslice = np.arange(0,nhyperpars)
def partfunc(x):
p = np.copy(endp)
p[reslice] = x
return likob.full_loglikelihood_grad(p)[0]
ndhessdiag = nd.Hessdiag(func)
ndparthess = nd.Hessian(partfunc)
# Create a good-enough approximation for the Hessian
nhdiag = ndhessdiag(endp)
nhpart = ndparthess(endp[reslice])
fullhessian = np.diag(nhdiag)
fullhessian[:nhyperpars,:nhyperpars] = nhpart
pickle.dump(fullhessian,open(hessfile,'wb'))
else:
fullhessian = pickle.load(open(hessfile,'rb'))
# Whiten the likelihood and starting parameter vector
wl = itvl.whitenedLikelihood(likob, endp, -fullhessian)
likob.pstart = endp
wlps = wl.forward(endp)
# Set up and run HMC sampler
chaindir = f'{outdir}/chains_{psr}'
if not os.path.exists(chaindir):
os.makedirs(chaindir)
chainfile = chaindir + '/samples.txt'
if not os.path.isfile(chainfile) or len(open(chainfile,'r').readlines()) < Nsamples-1:
# Run NUTS for 20000 samples, with a burn-in of
# 1000 samples (target acceptance = 0.6)
samples, lnprob, epsilon = nuts6(wl.loglikelihood_grad, Nsamples, Nburnin,
wlps, 0.6,
verbose=True,
outFile=chainfile,
pickleFile=chaindir + '/save')
#------------POST PROCESSING------------
# Undo all the coordinate transformations and save samples to file
parsfile = f'{outdir}/{psr}-pars.npy'
if not os.path.isfile(parsfile):
samples = np.loadtxt(chaindir + '/samples.txt')
fullsamp = wl.backward(samples[:,:-2])
funnelsamp = likob.backward(fullsamp)
pars = likob.multi_full_backward(funnelsamp)
np.save(parsfile,pars)
else:
pars = np.load(parsfile)
# Corner plot of the hyperparameter posteriors
parnames = list(likob.ptadict.keys())
if not os.path.isfile(f'{outdir}/{psr}-corner.pdf'):
corner.corner(pars[:,:nhyperpars], labels=parnames[:nhyperpars], show_titles=True);
plt.savefig(f'{outdir}/{psr}-corner.pdf')
# Array of outlier probabilities
pobsfile = f'{outdir}/{psr}-pobs.npy'
if not os.path.isfile(pobsfile):
nsamples = len(pars)
nobs = len(likob.Nvec)
# basic likelihood
lo = likob
outps = np.zeros((nsamples,nobs),'d')
sigma = np.zeros((nsamples,nobs),'d')
for i,p in enumerate(pars):
outps[i,:], sigma[i,:] = poutlier(p,lo)
out = np.zeros((nsamples,nobs,2),'d')
out[:,:,0], out[:,:,1] = outps, sigma
np.save(pobsfile,out)
else:
out = np.load(pobsfile)
outps, sigma = out[:,:,0], out[:,:,1]
avgps = np.mean(outps,axis=0)
medps = np.median(outps,axis=0)
# Residual plot with outliers highlighted
spd = 86400.0 # seconds per day
residualplot = f'{outdir}/{psr}-residuals.pdf'
outliers = medps > 0.1
if not os.path.isfile(residualplot):
nout = np.sum(outliers)
nbig = nout
print("Big: {}".format(nbig))
if nout == 0:
outliers = medps > 5e-4
nout = np.sum(outliers)
print("Plotted: {}".format(nout))
plt.figure(figsize=(15,6))
psrobj = likob.psr
# convert toas to mjds
toas = psrobj.toas/spd
# red noise at the starting fit point
_, _ = likob.full_loglikelihood_grad(endp)
rednoise = psrobj.residuals - likob.detresiduals
# plot tim-file residuals (I think)
plt.errorbar(toas,psrobj.residuals,yerr=psrobj.toaerrs,fmt='.',alpha=0.3)
# red noise
# plt.plot(toas,rednoise,'r-')
# possible outliers
plt.errorbar(toas[outliers],psrobj.residuals[outliers],yerr=psrobj.toaerrs[outliers],fmt='rx')
plt.savefig(residualplot)
# Text file with exact indices of outlying TOAs and their
# outlier probabilities
outlier_indices = f'{outdir}/outliers.txt'
with open(outlier_indices, 'w') as f:
for ii, elem in enumerate(outliers):
if elem:
f.write('TOA Index {}: Outlier Probability {}\n'.format(likob.psr.desort[ii], medps[ii]))
return medps[likob.psr.desort]
