From f7593442942915085f979f87f315b7a625ace228 Mon Sep 17 00:00:00 2001 From: jacopok Date: Wed, 27 Nov 2024 11:35:55 +0100 Subject: [PATCH 1/4] first implementation of the spline representation of the coordinate change --- ultranest/hotstart.py | 60 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 49 insertions(+), 11 deletions(-) diff --git a/ultranest/hotstart.py b/ultranest/hotstart.py index a7e83e6c..0b27c05c 100644 --- a/ultranest/hotstart.py +++ b/ultranest/hotstart.py @@ -14,6 +14,7 @@ from .utils import resample_equal, vectorize +from scipy.interpolate import PchipInterpolator def get_auxiliary_problem(loglike, transform, ctr, invcov, enlargement_factor, df=1): """Return a new loglike and transform based on an auxiliary distribution. @@ -343,6 +344,45 @@ def compute_quantile_intervals_refined(steps, upoints, uweights, logsteps_max=20 return ulos, uhis, uinterpspace +def get_aux_splines(upoints, uweights, delta_v): + """ Compute spline representations of the auxiliary transforms. + """ + + splines = [] + derivatives = [] + for i in range(upoints.shape[1]): + these_upoints = upoints[:, i] + sort = np.argsort(these_upoints) + cdf = np.cumsum(uweights[sort]) + + u0 = np.min(these_upoints) + u1 = np.max(these_upoints) + + for _ in range(5): + u_range = np.linspace(u0, u1, num=20) + cdf_interp = np.interp(u_range, these_upoints[sort], cdf) + pdf_interp = np.gradient(cdf_interp, u_range) + slope_outside = (1-(u1 - u0)) / (1 - delta_v) + + u0 = u_range[np.argmin(np.where(pdf_interp > slope_outside, u_range, 1))] + u1 = u_range[np.argmax(np.where(pdf_interp > slope_outside, u_range, 0))] + + if np.all(pdf_interp >= slope_outside): + break + + v_0 = u0 / slope_outside + v_1 = 1 - (1-u1) / slope_outside + + v_values = np.linspace(0, 1, num=10000) + v_to_interp = np.concatenate(([0], v_0 + (v_1-v_0)*cdf_interp, [1])) + u_to_interp = np.concatenate(([0], u_range, [1])) + + interpolator = PchipInterpolator(v_to_interp, u_to_interp) + splines.append(interpolator) + derivatives.append(interpolator.derivative()) + + return splines, derivatives + def get_auxiliary_contbox_parameterization( param_names, loglike, transform, upoints, uweights, vectorized=False, ): @@ -412,23 +452,23 @@ def get_auxiliary_contbox_parameterization( mask = np.logical_and(upoints > 0, upoints < 1).all(axis=1) assert np.all(mask), ( 'upoints must be between 0 and 1, have:', upoints[~mask,:]) - steps = 10**-(1.0 * np.arange(1, 8, 2)) + # steps = 10**-(1.0 * np.arange(1, 8, 2)) + nsamples, ndim = upoints.shape - assert nsamples > 10 - ulos, uhis, uinterpspace = compute_quantile_intervals_refined(steps, upoints, uweights) aux_param_names = param_names + ['aux_logweight'] + delta_v = 0.5**(1/ndim) + splines, derivatives = get_aux_splines(upoints, uweights, delta_v) + def aux_transform(u): ndim2, = u.shape assert ndim2 == ndim + 1 umod = np.empty(ndim) log_aux_volume_factors = 0 for i in range(ndim): - ulo_here = np.interp(u[-1], uinterpspace, ulos[:,i]) - uhi_here = np.interp(u[-1], uinterpspace, uhis[:,i]) - umod[i] = ulo_here + (uhi_here - ulo_here) * u[i] - log_aux_volume_factors += np.log(uhi_here - ulo_here) + umod[i] = splines[i](u[i]) + log_aux_volume_factors += np.log(derivatives[i](u[i])) return np.append(transform(umod), log_aux_volume_factors) def aux_transform_vectorized(u): @@ -437,10 +477,8 @@ def aux_transform_vectorized(u): umod = np.empty((nsamples, ndim2 - 1)) log_aux_volume_factors = np.zeros((nsamples, 1)) for i in range(ndim): - ulo_here = np.interp(u[:,-1], uinterpspace, ulos[:,i]) - uhi_here = np.interp(u[:,-1], uinterpspace, uhis[:,i]) - umod[:,i] = ulo_here + (uhi_here - ulo_here) * u[:,i] - log_aux_volume_factors[:,0] += np.log(uhi_here - ulo_here) + umod[:, i] = splines[i](u[:,i]) + log_aux_volume_factors[:,0] += np.log(derivatives[i](u[:, i])) return np.hstack((transform(umod), log_aux_volume_factors)) def aux_loglikelihood(x): From 67eca853959977bdd1cf366df2f3cb3c31bcc05f Mon Sep 17 00:00:00 2001 From: jacopok Date: Wed, 27 Nov 2024 12:31:55 +0100 Subject: [PATCH 2/4] remove unused variable --- ultranest/hotstart.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ultranest/hotstart.py b/ultranest/hotstart.py index 0b27c05c..b25787f9 100644 --- a/ultranest/hotstart.py +++ b/ultranest/hotstart.py @@ -373,7 +373,6 @@ def get_aux_splines(upoints, uweights, delta_v): v_0 = u0 / slope_outside v_1 = 1 - (1-u1) / slope_outside - v_values = np.linspace(0, 1, num=10000) v_to_interp = np.concatenate(([0], v_0 + (v_1-v_0)*cdf_interp, [1])) u_to_interp = np.concatenate(([0], u_range, [1])) From 3425332fe0ecb70771785cd73b4b3d91e6e35f10 Mon Sep 17 00:00:00 2001 From: jacopok Date: Wed, 27 Nov 2024 12:33:29 +0100 Subject: [PATCH 3/4] Fix whitespace --- ultranest/hotstart.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/ultranest/hotstart.py b/ultranest/hotstart.py index b25787f9..0a897a24 100644 --- a/ultranest/hotstart.py +++ b/ultranest/hotstart.py @@ -16,6 +16,7 @@ from scipy.interpolate import PchipInterpolator + def get_auxiliary_problem(loglike, transform, ctr, invcov, enlargement_factor, df=1): """Return a new loglike and transform based on an auxiliary distribution. @@ -345,9 +346,9 @@ def compute_quantile_intervals_refined(steps, upoints, uweights, logsteps_max=20 def get_aux_splines(upoints, uweights, delta_v): - """ Compute spline representations of the auxiliary transforms. + """ Compute spline representations of the auxiliary transforms. """ - + splines = [] derivatives = [] for i in range(upoints.shape[1]): @@ -357,19 +358,19 @@ def get_aux_splines(upoints, uweights, delta_v): u0 = np.min(these_upoints) u1 = np.max(these_upoints) - + for _ in range(5): u_range = np.linspace(u0, u1, num=20) cdf_interp = np.interp(u_range, these_upoints[sort], cdf) pdf_interp = np.gradient(cdf_interp, u_range) slope_outside = (1-(u1 - u0)) / (1 - delta_v) - + u0 = u_range[np.argmin(np.where(pdf_interp > slope_outside, u_range, 1))] u1 = u_range[np.argmax(np.where(pdf_interp > slope_outside, u_range, 0))] if np.all(pdf_interp >= slope_outside): break - + v_0 = u0 / slope_outside v_1 = 1 - (1-u1) / slope_outside @@ -382,6 +383,7 @@ def get_aux_splines(upoints, uweights, delta_v): return splines, derivatives + def get_auxiliary_contbox_parameterization( param_names, loglike, transform, upoints, uweights, vectorized=False, ): @@ -452,14 +454,14 @@ def get_auxiliary_contbox_parameterization( assert np.all(mask), ( 'upoints must be between 0 and 1, have:', upoints[~mask,:]) # steps = 10**-(1.0 * np.arange(1, 8, 2)) - + nsamples, ndim = upoints.shape aux_param_names = param_names + ['aux_logweight'] delta_v = 0.5**(1/ndim) splines, derivatives = get_aux_splines(upoints, uweights, delta_v) - + def aux_transform(u): ndim2, = u.shape assert ndim2 == ndim + 1 From a8bf72b646ce0d564115d7ee866d70a739e71e2c Mon Sep 17 00:00:00 2001 From: jacopok Date: Thu, 12 Dec 2024 12:14:33 +0100 Subject: [PATCH 4/4] KDE implementation of auxiliary transformation --- ultranest/hotstart.py | 145 +++++++++++++++++++++++++++--------------- 1 file changed, 92 insertions(+), 53 deletions(-) diff --git a/ultranest/hotstart.py b/ultranest/hotstart.py index 0a897a24..31011de6 100644 --- a/ultranest/hotstart.py +++ b/ultranest/hotstart.py @@ -14,7 +14,9 @@ from .utils import resample_equal, vectorize -from scipy.interpolate import PchipInterpolator +from scipy.interpolate import PchipInterpolator, make_interp_spline +from scipy.integrate import trapezoid +from getdist import MCSamples def get_auxiliary_problem(loglike, transform, ctr, invcov, enlargement_factor, df=1): @@ -344,48 +346,77 @@ def compute_quantile_intervals_refined(steps, upoints, uweights, logsteps_max=20 return ulos, uhis, uinterpspace +def conservative_prior_guess(all_x, all_p, single_dim_fraction): + """ Given a function p:[0, 1] -> R+, + represented by some samples (`all_p`) corresponding + to the points `all_x`, return another normalized function + which is the sum of p and a constant, weighed by `single_dim_fraction` + and `(1-single_dim_fraction)` respectively. + """ + + integral = trapezoid(all_p, x=all_x) + + return all_p * single_dim_fraction + np.ones_like(all_p) * (1 - single_dim_fraction) * integral -def get_aux_splines(upoints, uweights, delta_v): +def get_aux_splines(upoints, uweights, volume_contribution=0.5, beta=1.): """ Compute spline representations of the auxiliary transforms. + + upoints: shape (n_points, dim) + uweights: shape (n_points,) + volume_contribution: fraction of total prior mass which will + be taken up by the "guess" distribution represented by upoints + (default: 0.5) + beta: temperature for the prior guess (default: 1) + """ - splines = [] - derivatives = [] - for i in range(upoints.shape[1]): - these_upoints = upoints[:, i] - sort = np.argsort(these_upoints) - cdf = np.cumsum(uweights[sort]) - - u0 = np.min(these_upoints) - u1 = np.max(these_upoints) - - for _ in range(5): - u_range = np.linspace(u0, u1, num=20) - cdf_interp = np.interp(u_range, these_upoints[sort], cdf) - pdf_interp = np.gradient(cdf_interp, u_range) - slope_outside = (1-(u1 - u0)) / (1 - delta_v) - - u0 = u_range[np.argmin(np.where(pdf_interp > slope_outside, u_range, 1))] - u1 = u_range[np.argmax(np.where(pdf_interp > slope_outside, u_range, 0))] - - if np.all(pdf_interp >= slope_outside): - break - - v_0 = u0 / slope_outside - v_1 = 1 - (1-u1) / slope_outside - - v_to_interp = np.concatenate(([0], v_0 + (v_1-v_0)*cdf_interp, [1])) - u_to_interp = np.concatenate(([0], u_range, [1])) - - interpolator = PchipInterpolator(v_to_interp, u_to_interp) - splines.append(interpolator) - derivatives.append(interpolator.derivative()) - - return splines, derivatives + + n_points, dim = upoints.shape + assert uweights.shape == (n_points, ) + single_dim_fraction = volume_contribution**(1/dim) + + mcsamples = MCSamples( + samples=upoints, + weights=uweights, + names=[str(i) for i in range(upoints.shape[1])], + ranges={str(i):[0, 1] for i in range(upoints.shape[1])} + ) + + densities, inverse_cumulatives = [], [] + + for i in range(dim): + density_interior = mcsamples.get1DDensity(str(i)) + + interior_x = density_interior.x + + all_x = np.sort(np.unique(np.concatenate(( + np.linspace(0, 1, num=1024), + interior_x + )))) + + all_p = np.maximum(density_interior(all_x), 0) + + all_p_thresholded = conservative_prior_guess(all_x, all_p**beta, single_dim_fraction) + + interpolator = PchipInterpolator(all_x, all_p_thresholded) + cumulative = interpolator.antiderivative() + integral = cumulative(1) - cumulative(0) + cumulative_points = cumulative(all_x) / integral + + unique_cumulative, unique_idx = np.unique(cumulative_points, return_index=True) + + inverse_cumulative = make_interp_spline(unique_cumulative, all_x[unique_idx], k=1) + inverse_cumulatives.append(inverse_cumulative) + + cumulative = make_interp_spline(all_x, cumulative_points, k=1) + densities.append(cumulative.derivative()) + + return densities, inverse_cumulatives def get_auxiliary_contbox_parameterization( - param_names, loglike, transform, upoints, uweights, vectorized=False, + param_names, loglike, transform, upoints, uweights=None, vectorized=False, + volume_contribution=0.5, beta=1., ): """Return a new loglike and transform based on an auxiliary distribution. @@ -397,16 +428,13 @@ def get_auxiliary_contbox_parameterization( This is achieved by deforming the prior space, and undoing that transformation by correction weights in the likelihood. A additional parameter, "aux_logweight", is added at the end, - which contains the correction weight. You can ignore it. - - The auxiliary distribution used for transformation/weighting is - factorized. Each axis considers the ECDF of the auxiliary samples, - and segments it into quantile segments. Within each segment, - the parameter edges in u-space are linearly interpolated. - To see the interpolation quantiles for each axis, use:: - - steps = 10**-(1.0 * np.arange(1, 8, 2)) - ulos, uhis, uinterpspace = compute_quantile_intervals_refined(steps, upoints, uweights) + which contains the correction weight. + If this parameter is negative for most of the sampling, + it means that the samples are being taken from the "expanded" region, + therefore the hotstart is working as expected. + If it is positive, the auxiliary prior guess was likely mis-specified; + the sampler may still give correct results, but it will take + as long as a regular nested sampling run (or longer). Parameters ------------ @@ -422,6 +450,15 @@ def get_auxiliary_contbox_parameterization( Weights of samples (needs to sum of 1) vectorized: bool whether the loglike & transform functions are vectorized + volume_contribution: float + fraction of total prior mass which will + be taken up by the "guess" distribution (default: 0.5), + should be between 0 and 1 + beta: float + inverse temperature for the prior guess (default: 1). + The guess will be rescaled as prior**beta: higher beta + will make the prior guess more informative and peaked, + lower beta will make it flatter and more conservative. Returns --------- @@ -453,14 +490,16 @@ def get_auxiliary_contbox_parameterization( mask = np.logical_and(upoints > 0, upoints < 1).all(axis=1) assert np.all(mask), ( 'upoints must be between 0 and 1, have:', upoints[~mask,:]) - # steps = 10**-(1.0 * np.arange(1, 8, 2)) nsamples, ndim = upoints.shape + if uweights is None: + uweights = np.ones(nsamples) / nsamples + assert uweights.shape == (nsamples,) + aux_param_names = param_names + ['aux_logweight'] - delta_v = 0.5**(1/ndim) - splines, derivatives = get_aux_splines(upoints, uweights, delta_v) + densities, inv_cumulatives = get_aux_splines(upoints, uweights, volume_contribution, beta) def aux_transform(u): ndim2, = u.shape @@ -468,8 +507,8 @@ def aux_transform(u): umod = np.empty(ndim) log_aux_volume_factors = 0 for i in range(ndim): - umod[i] = splines[i](u[i]) - log_aux_volume_factors += np.log(derivatives[i](u[i])) + umod[i] = inv_cumulatives[i](u[i]) + log_aux_volume_factors -= np.log(densities[i](umod[i])) return np.append(transform(umod), log_aux_volume_factors) def aux_transform_vectorized(u): @@ -478,8 +517,8 @@ def aux_transform_vectorized(u): umod = np.empty((nsamples, ndim2 - 1)) log_aux_volume_factors = np.zeros((nsamples, 1)) for i in range(ndim): - umod[:, i] = splines[i](u[:,i]) - log_aux_volume_factors[:,0] += np.log(derivatives[i](u[:, i])) + umod[:, i] = inv_cumulatives[i](u[:,i]) + log_aux_volume_factors[:,0] -= np.log(densities[i](umod[:, i])) return np.hstack((transform(umod), log_aux_volume_factors)) def aux_loglikelihood(x):