Add support for nessai sampler in pycbc inference

[mj-will]

This PR adds basic support for the nested sampler nessai in pycbc inference.

This implementation is intended as a starting point that can be refined and improved upon, so there are some limitations (see below). Below I mention a few points that I think are important and/or may need some discussion/tweaks.

Changes

• Add pycbc/inference/io/nessai.py which includes two classes: the main interface and the nessai model object
• Add pycbc/inference/sampler/nessai.py (based on the equivalent for dynesty)
• Add nessai to the dictionary of samplers in pycbc/inference/sampler/__init__.py

Config files

I wanted to avoid needing to provide types for every single keyword argument nessai supports, so I opted to use ast.literal_eval to evaluate the arguments. I can understand wanting to avoid literal evaluations, so I can change this if folks have suggestions.

Requirements

I made some changes to nessai to enable easier integration with pycbc, so using nessai requires at least nessai>=0.10

Limitations

There are a few places the implementation could be improved:

• I have not included support for i-nessai; it requires some additional methods in model class nessai uses
• Resuming from checkpoints works, but as it stands, there is no means to automatically checkpoint the sampler. nessai supports this internally, but writes to its own file. I'm keen to hear peoples suggestions about how to handle this.
• Some of GW specific features in nessai are not supported, this may require more changes to nessai

Testing

I haven't added any unit or integration tests directly, but I can do if that is desired. Instead, I've been testing the implementation using the scripts in this repo: https://github.com/mj-will/nessai-pycbc-testing and the runs I've been doing as a part of other projects.

[ahnitz]

@mj-will Is it possible for nessai to provide a callback with whatever data nessai needs to checkpoint? And similar for resuming? Where possible, we'd like to keep the checkpoint information with our results files. This makes it easier to manage on a heterogeneous cluster where we want a higher level workflow to understand the checkpoint files.

[ahnitz]

This direction can be implemented pretty straightforwardly, see https://github.com/gwastro/pycbc/blob/master/pycbc/inference/sampler/base_cube.py#L82

[mj-will]

Thanks, I'll look at adding that. Unfortunately, the current implementation requires both directions, though it's something I'm hoping to change in a future release.

[ahnitz]

Ok, I don't think that's unreasonable, but would require some changes to the distribution package to support.

[mj-will]

I'd prefer to leave this not implemented and then revisit both methods together when either the distribution package supports the other direction or nessai doesn't require both. It is unused in the standard sampler, so also wouldn't be testing by running the example. What do you think?

[mj-will]

@ahnitz So it already supports something like that for resuming, where you can give it an object rather than a resume file path (see the function below and self.resume_data.

When you say a callback, would that be a function the user provides to nessai that it can then call internally to checkpoint?

pycbc/pycbc/inference/sampler/nessai.py

Line 225 in 4169605

def resume_from_checkpoint(self):

)

[ahnitz]

@cdcapano Any suggestions for the things that @mj-will raises?

[ahnitz]

@ahnitz So it already supports something like that for resuming, where you can give it an object rather than a resume file path (see the function below and self.resume_data.

When you say a callback, would that be a function the user provides to nessai that it can then call internally to checkpoint?

pycbc/pycbc/inference/sampler/nessai.py

Line 225 in 4169605

def resume_from_checkpoint(self):

)

Yes, exactly. I assume the issue is that you don't have any method of forcing a return to the calling code so that it can do a checkpoint. If that is a technical choice (which may be reasonable depending on the code structure), the alternative would be to take a function and call it with the data you'd have checkpointed at the times you'd normally checkpoint. That still takes away options related to checkpoint frequency, but depending on the library implementation might be the only straightforward solution.

[mj-will]

@ahnitz So it already supports something like that for resuming, where you can give it an object rather than a resume file path (see the function below and self.resume_data.
When you say a callback, would that be a function the user provides to nessai that it can then call internally to checkpoint?

pycbc/pycbc/inference/sampler/nessai.py

Line 225 in 4169605

def resume_from_checkpoint(self):

)

Yes, exactly. I assume the issue is that you don't have any method of forcing a return to the calling code so that it can do a checkpoint. If that is a technical choice (which may be reasonable depending on the code structure), the alternative would be to take a function and call it with the data you'd have checkpointed at the times you'd normally checkpoint. That still takes away options related to checkpoint frequency, but depending on the library implementation might be the only straightforward solution.

Yes, that's exactly the issue. Adding support for a checkpointing callback is definitely the more straightforward of the two options, so I'll add that to the roadmap for future nessai releases. Thanks for the suggestion!

[mj-will]

From the discussion on the parameter estimation dev call:

• Address code climate issues
• Add nessai to https://github.com/gwastro/pycbc/tree/master/examples/inference/samplers (this may require making the model 2d)
  • Check plot still makes sense
• Update nessai version will new release
• Add nessai to https://github.com/gwastro/pycbc/blob/master/companion.txt
• Add a comment about the literal_eval

[mj-will]

Here's an example of what the plot from the samplers example will look like (I've only included a subset of samplers):

[mj-will]

@ahnitz I think I've now addressed the all the relevant code climate issues and what remains are issues with either existing code, code that it's directly copied from existing samplers or code that cannot be changed because it is dictated by how nessai is written, e.g. the use of x and N.

[ahnitz]

@mj-will Looks good