Closed wbalmer closed 8 months ago
Thanks a lot William! Great that you are trying to implement this!
In FitModel
, there was a similar issue from what I remember when adding run_ultranest
. I had to make the embedded function for the likelihood an actual method of the class. Within run_multinest
and run_ultranest
, I then added a local function for the likelihood that calls the class method. This should indeed be possible here as well (I think), although I am not sure how straightforward. I can have a look so the code structure can be used for your implementation.
Would really help indeed if the opacity memory could be shared between cores, but I am not sufficiently experienced with multiprocessing to understand what is the best way to do that... Maybe Paul and Evert have put some thought into that?
I have restructured the AtmosphericRetrieval
class in the last commit on the main
branch (see b2adbcbd623f3951cdcd06bfe0653492af661134).
In run_retrieval
you can run the Dynesty sampler, which can use either prior_transform()
and lnlike_func()
, or lnprior_dynesty()
and lnlike_dynesty()
. The latter in case anything Dynesty-specific needs to be implemented with the prior/lnlike otherwise these can be removed.
Would this help with your Dynesty implementation?
I am not sure if lnlike_func
/lnlike_dynesty
can be pickled or if it needs to be moved outside of the class. In that last scenario it would be a bit more work to implement because the attributes should be converted into parameters. But hopefully it works like this π€.
I have set it up a bit differently! There is now the setup_retrieval
method, and then separately run_multinest
/run_dynesty
. The latter will have only MultiNest/Dynesty parameters, so perhaps the kwargs_multinest
is not needed. Sorry about that, but we can leave it for now and decide later if this new setup works.
Wow! Awesome stuff. I will pull soon and try out later this week ^_^ thanks for entertaining this implementation Tomas!
Hi Tomas - I've implemented an initial run_dynesty
on my fork's branch add_run_dynesty that is based on the backtracks
implementation. That branch is up to date with commit 3d150aa.
I've tested the various permutations (dynamic bool, n_pool None/int, mpi_pool bool, resume bool) of the multiprocessing locally to confirm that they at least begin sampling and save the intermediate products. I have a few lines at the end that take the dynesty
output results object and saves a retrieval_post_equal_weights.dat
file that is equivalent to the output from pymultinest
.
I'm going to try to test this on my cluster next, I couldn't do more than a few processes with mpi or a multiprocess Pool on my local machine due to memory constraints.
Let me know what you think, and if you'd like me to open a PR for that branch.
Awesome! Good to hear that it is working already!
Would be great if you could create a PR for that π
Thanks again for implementing the Dynesty support in AtmosphericRetrieval
! Since I had been making some updates in FitModel
, I have adapted run_dynesty
in there, which was quite straightforward (see 1f02b5a).
Looks great! I had a commit I forgot to push to do something similar, though I hadn't yet documented it. I'll test on my end with some more AF Lep b fits. I think we can close the issue now :).
Hi Tomas, I just wanted to open this issue to keep track of our discussion about dynesty support. This isn't very urgent at all, but it would go a long way to help me be more productive with the retrievals I plan to do in the near future (for reasons I describe below). I hope to start with adding dynesty support for the retrievals and then we can likely take the lessons learned there and add support for the forward modeling (I think that is less important since the nuances of nested sampling aren't really a potential limiting factor in doing the grid interpolation/model comparison).
Some hand wavy reasons for implementing dynesty:
MultiNest
) and dynamic nested sampling (effectively reproducesUltraNest
).The issue with a clean port, e.g. adding a
run_dynesty
function that is effectively the same asrun_multinest
except it passes the loglike and prior functions to the dynesty interface right now is that withinspecies.fit.retrieval.run_multinest
theloglike_func
andprior_func
are local objects, functions defined insiderun_multinest
and not class attributes or global objects. In this way, when attempting to set a poolwith dynesty.pool.Pool(npool, loglike_func, prior_func) as pool:
, you get an error likeAttributeError: Can't pickle local object 'AtmosphericRetrieval.run_dynesty.<locals>.loglike_func'
and when trying to initialize dynesty on one core
dsampler = dynesty.NestedSampler(loglike_func, prior_func,len(self.parameters))
, you get an error likeTypeError: AtmosphericRetrieval.run_dynesty.<locals>.prior_func() missing 2 required positional arguments: 'n_dim' and 'n_param'
I've tried to hack something together to show these errors on a dummy branch of my fork, dynesty_branch_2.
I've pseudo-coded / parallel implemented a bigger overhaul in a separate file
species.fit.dretrieval.py
on a branch of my fork dynesty-retrieval which is my first attempt to solve this by making loglike_func and prior_func attributes of the AtmosphericRetrieval class so that they can be interpreted by dynesty. This of course involves making the variables defined withinrun_multinest
that loglike and prior rely on into class attributes as well, which became something of a mess - I don't intend that code to be the actual strategy of implementation, more like its a sketch book π .If there's a different way to do this that I've overlooked, however, that doesn't involve such a big change to the structure of the AtmosphericRetrieval class, I'd prefer that. Hence why I've done most of this in pseudo code and in a separate file to start tinkering with. I was able to run a full retrieval on some test data using multiprocessing locally using dynesty by calling AtmosphericRetrieval with
dretrieval.py
, but I haven't been able to get a mpi implementation to work.I'm also trying to think through if there is any memory saving way to set up the retrieval across multiple cores, that doesn't involve having to load the opacity tables for each core..., that has been an issue on my cluster. That is definitely out of the scope of this issue but may be possible in the future with the schwimmbad/dynesty architecture this issue could help set up? I am thinking of the main/worker structure (like the mpi-demo.py in schwimmbad) that loads the AtmosphericRetrieval and RadTrans object once, but is able to distribute the calc_flux to workers in the pool. There is probably a reason I am not thinking of that this is not feasible.