Vectorized stepsampler

[hippke]

• UltraNest version: 3.3.3
• Python version: 3.9.4
• Operating System: Linux

Description

I would like to run a vectorized stepsampler as described in issue 14. I can run the vectorized ReactiveNestedSampler correctly. For my problem, however, a stepsampler is much better suited. The vectorized stepsampler crashes with an error directly after launch.

What I Did

I run this code:

sampler = ReactiveNestedSampler(
            parameters,
            log_likelihood, 
            prior_transform,
            vectorized=True,
            ndraw_min=100,
            ndraw_max=400,
            )
sampler.stepsampler = ultranest.stepsampler.AHARMSampler(nsteps=400)
result = sampler.run()

without the line sampler.stepsampler = ultranest.stepsampler.AHARMSampler(nsteps=400) it runs correctly. With this line it crashes:

starting at [0.37035903 0.26168261 0.74545311 0.13944934 0.86101038 0.0762323
 0.67986505 0.67523488 0.73274571 0.75403661 0.65612488 0.18768482
 0.13987624]
Traceback (most recent call last):
  File "/home/michael//1.py", line 304, in <module>
    result = sampler.run()
  File "/home/michael/miniconda3/envs/py39/lib/python3.9/site-packages/ultranest/integrator.py", line 2170, in run
    for result in self.run_iter(
  File "/home/michael/miniconda3/envs/py39/lib/python3.9/site-packages/ultranest/integrator.py", line 2430, in run_iter
    u, p, L = self._create_point(Lmin=Lmin, ndraw=ndraw, active_u=active_u, active_values=active_values)
  File "/home/michael/miniconda3/envs/py39/lib/python3.9/site-packages/ultranest/integrator.py", line 1697, in _create_point
    u, v, logl, nc = self.stepsampler.__next__(
  File "/home/michael/miniconda3/envs/py39/lib/python3.9/site-packages/ultranest/stepsampler.py", line 1433, in __next__
    assert self.nsteps_done == len(self.history), (self.nsteps_done, len(self.history))
AssertionError: (90, 100)

[JohannesBuchner]

Do other stepsamplers work, such as SliceSampler (e.g., with generate_direction=generate_cube_oriented_direction, as CubeSliceSampler does)?

I think the AHARMSampler is still in development. It is intended to provide superior vectorization.

[hippke]

Thank you for the quick reply! SliceSampler with generate_direction=generate_cube_oriented_direction works, but is very slow (~50x fewer evals/sec compared to ReactiveNestedSampler).

My problem requires ~2e8 evals until convergence with the StepSampler. I can calculate 3e5 models per second on my GPU, so that convergence would be possible in 10 minutes. But the high GPU speed is only possible if I calculate the model in batches of >1e4 (i.e. vectorized). I struggle to feed these high numbers of models into UltraNest.

[JohannesBuchner]

in ReactiveNestedSampler, you can control the batch size minimum and maximum with ndraw_min and ndraw_max, so I guess you could crank up ndraw_min to 10000. This will give you extra evaluations, but if your GPU evals are cheap, it may effectively be faster.

SliceSampler currently does evolve one walker with hit-and-run Monte Carlo or slice sampling, always defining a direction, finding a bracket and then iteratively sampling from this slice and shrinking it until acceptance. This is repeated nsteps times. Since there is a dependence on the previous result, one cannot simply parallelise this for a single walker.

To truly take advantage of GPU parallelism, a sampler should:

• propose a very conservative bracket, and evaluate many shrinkage steps at once, then pick the first successful one to start the next direction (AHARMSampler does this)
• evolve a population of multiple points simultaneously, and return one at each iteration.

If you are motivated, maybe you can implement something to this effect?