From 0d8bc41b1d5fbaf83f567f0aa15f24b414fac6ef Mon Sep 17 00:00:00 2001
From: Minas Karamanis <minaskar@gmail.com>
Date: Thu, 11 Jul 2024 15:13:30 +0200
Subject: [PATCH] n_steps control for pocoMC

---
 cosmosis/samplers/poco/poco_sampler.py | 4 ++++
 cosmosis/samplers/poco/sampler.yaml    | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/cosmosis/samplers/poco/poco_sampler.py b/cosmosis/samplers/poco/poco_sampler.py
index 74735557..95f31c6f 100755
--- a/cosmosis/samplers/poco/poco_sampler.py
+++ b/cosmosis/samplers/poco/poco_sampler.py
@@ -41,6 +41,8 @@ def config(self):
             self.flow = self.read_ini("flow", str, "nsf6")
             self.precondition = self.read_ini("precondition", bool, True)
             self.dynamic = self.read_ini("dynamic", bool, True)
+            self.n_steps = seld.read_ini("n_steps", int, len(pipeline.varied_params))
+            self.n_max_steps = seld.read_ini("n_max_steps", int, 10*len(pipeline.varied_params))
             seed = self.read_ini("seed", int, 0)
             if seed == 0:
                 seed = None
@@ -79,6 +81,8 @@ def config(self):
                 flow=self.flow,
                 precondition=self.precondition,
                 dynamic=self.dynamic,
+                n_steps=self.n_steps,
+                n_max_steps=self.n_max_steps,
                 output_dir=self.output_dir,
                 pool=self.pool,
                 blobs_dtype=float,
diff --git a/cosmosis/samplers/poco/sampler.yaml b/cosmosis/samplers/poco/sampler.yaml
index 6f244cfa..26f875ed 100644
--- a/cosmosis/samplers/poco/sampler.yaml
+++ b/cosmosis/samplers/poco/sampler.yaml
@@ -35,6 +35,8 @@ params:
     dynamic: (bool, default=True) Whether to use dynamic particle allocation.
     n_total: (integer, default=4096) Total number of effective samples to generate. Higher values will increase the accuracy of the posterior estimate (recommended ~10000 for publication quality corner plots).
     n_evidence: (integer, default=4096) Number of samples to generate for the Bayesian evidence calculation.
+    n_steps: (integer, default=None) Number of MCMC steps. This will be automatically adjusted based on the sampling efficiency. Default is `n_steps=n_dim`. Larger values will provide more robust results at higher computational cost.
+    n_max_steps: (integer, default=None) Maximum number of MCMC steps. Default is `n_max_steps=10*n_dim`. This should be greater than `n_steps`.
     save_every: (int, default=None) Save the state of the sampler every `save_every` iterations. If not provided, the sampler will not save intermediate states.
     seed: (int; default=0) A random seed for reproducibility. Default is to generate automatically.
     progress: (bool, default=True) Whether to show a progress bar during sampling.