From f7593442942915085f979f87f315b7a625ace228 Mon Sep 17 00:00:00 2001
From: jacopok <jacopok@gmail.com>
Date: Wed, 27 Nov 2024 11:35:55 +0100
Subject: [PATCH] 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):