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]