mj-will / nessai Goto Github PK

By default when using a reparmeterisations dictionary, nessai assumes that any parameter that hasn't been specified and doesn't have a known default shouldn't be reparameterised, that is NullReparameterisation is used.

There may be cases where you have various parameters that don't need a specific reparameterisation but using no reparameterisation might cause issues. For example, if the prior is very narrow.

Describe the solution you'd like
It would be helpful to have a default fallback reparameterisation that can be used when no reparameterisation is specified. By default, this could be None to preserve the current behaviour.

Describe the bug

The test suite fails when corner 2.2.2 is included in the test environment:

$ python -m pytest tests/test_plot.py -x
======================================== test session starts =========================================
platform linux -- Python 3.11.3, pytest-7.3.0, pluggy-1.0.0
rootdir: /home/duncan/git/nessai
configfile: pyproject.toml
plugins: cov-4.0.0
collected 75 items

tests/test_plot.py .......................................................F

============================================== FAILURES ==============================================
__________________________________________ test_corner_plot __________________________________________

live_points = array([(-4.65345777,  0.96432746, nan, nan, 0),
       (-2.21440385,  4.9681729 , nan, nan, 0),
       ( 1.67006296,  ... -3.00372179, nan, nan, 0)],
      dtype=[('x', '<f8'), ('y', '<f8'), ('logP', '<f8'), ('logL', '<f8'), ('it', '<i4')])

    def test_corner_plot(live_points):
        """Test the corner plot."""
>       fig = plot.corner_plot(live_points)

tests/test_plot.py:446:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
nessai/plot.py:58: in wrapper
    return func(*args, **kwargs)
nessai/plot.py:634: in corner_plot
    fig = corner.corner(
../../opt/mambaforge/envs/nessai/lib/python3.11/site-packages/corner/corner.py:280: in corner
    return arviz_corner(
../../opt/mambaforge/envs/nessai/lib/python3.11/site-packages/corner/arviz_corner.py:136: in arviz_corner
    fig = corner_impl(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

xs = array([[-4.65345777,  0.96432746],
       [-2.21440385,  4.9681729 ],
       [ 1.67006296,  3.63502087],
       [ 1.55... 0.03564489],
       [ 3.90688656, -2.11997105],
       [-4.1205496 ,  4.93055757],
       [ 0.96344951, -3.00372179]])
bins = 32, range = None, axes_scale = 'linear', weights = None, color = '#02979d', hist_bin_factor = 1
smooth = 0.9, smooth1d = None, labels = array(['x', 'y'], dtype='<U4'), label_kwargs = {}
titles = array(['x', 'y'], dtype='<U4'), show_titles = True, title_fmt = '.2f', title_kwargs = {}
truths = None, truth_color = '#f5b754', scale_hist = False, quantiles = [0.16, 0.84]
title_quantiles = [0.16, 0.84], verbose = False, fig = None, max_n_ticks = 5, top_ticks = False
use_math_text = False, reverse = False, labelpad = 0.0, hist_kwargs = {'density': True}
hist2d_kwargs = {'fill_contours': True, 'levels': (0.3934693402873666, 0.8646647167633873, 0.9888910034617577), 'plot_datapoints': True, 'plot_density': True}

    def corner_impl(
        xs,
        bins=20,
        range=None,
        axes_scale="linear",
        weights=None,
        color=None,
        hist_bin_factor=1,
        smooth=None,
        smooth1d=None,
        labels=None,
        label_kwargs=None,
        titles=None,
        show_titles=False,
        title_fmt=".2f",
        title_kwargs=None,
        truths=None,
        truth_color="#4682b4",
        scale_hist=False,
        quantiles=None,
        title_quantiles=None,
        verbose=False,
        fig=None,
        max_n_ticks=5,
        top_ticks=False,
        use_math_text=False,
        reverse=False,
        labelpad=0.0,
        hist_kwargs=None,
        **hist2d_kwargs,
    ):
        if quantiles is None:
            quantiles = []
        if title_kwargs is None:
            title_kwargs = dict()
        if label_kwargs is None:
            label_kwargs = dict()

        # If no separate titles are set, copy the axis labels
        if titles is None:
            titles = labels

        # deal with title quantiles so they much quantiles unless desired otherwise
        if title_quantiles is None:
            if len(quantiles) > 0:
                title_quantiles = quantiles
            else:
                # a default for when quantiles not supplied.
                title_quantiles = [0.16, 0.5, 0.84]

        if show_titles and len(title_quantiles) != 3:
>           raise ValueError(
                "'title_quantiles' must contain exactly three values; "
                "pass a length-3 list or array using the 'title_quantiles' argument"
            )
E           ValueError: 'title_quantiles' must contain exactly three values; pass a length-3 list or array using the 'title_quantiles' argument

../../opt/mambaforge/envs/nessai/lib/python3.11/site-packages/corner/core.py:83: ValueError

I believe the incompatibility was introduced by dfm/corner.py#193.

To Reproduce

mamba create -n nessai python=3.11 corner=2.2.2 nessai=0.8.0 pytest pytest-cov
conda activate nessai
python -m pytest tests/test_plot.py -x

Expected behaviour

The tests pass

Screenshots

Additional context
Add any other context about the problem here.

Describe the bug
AttributeError: module 'jax' has no attribute 'Array'

To Reproduce
`
from jax.config import config

config.update("jax_enable_x64", True)

import pylab as plt
import tensorflow_probability.substrates.jax as tfp
from jax import random, numpy as jnp
from jax import vmap

import jaxns
from jaxns import ExactNestedSampler
from jaxns import Model
from jaxns import PriorModelGen, Prior
from jaxns import TerminationCondition
from jaxns import analytic_log_evidence

tfpd = tfp.distributions

#%%
def log_likelihood(x):
X = x[np.newaxis, ...]
return -1 * model.heterodyne_minusLogLikelihood(X)

def prior_model() -> PriorModelGen:
x = yield Prior(tfpd.Uniform(low=model.lower_bound, high=model.upper_bound), name='x')
return x

model = Model(prior_model=prior_model,
log_likelihood=log_likelihood)

ns = exact_ns = ExactNestedSampler(model=model, num_live_points=200, num_parallel_samplers=1,
max_samples=1e4)

termination_reason, state = exact_ns(random.PRNGKey(42),
term_cond=TerminationCondition(live_evidence_frac=1e-4))
results = exact_ns.to_results(state, termination_reason)

`

Expected behaviour
Crash occurs before sampler begins

Screenshots

Additional context
Python 3.10.4
JAX 0.3.15

The unit and integration tests are failing because of the latest release of numpy (v1.25).

I installed the version locally and test a few examples and they all seem fine, so it may be something specific to the tests.

Describe the bug
In version 0.7.1 of nessai, setting the precision for the log-likelihood doesn't change the value that then gets used when running the sampler.

To Reproduce

from nessai import config

config.LOGL_DTYPE = "f4"

from nessai.livepoint import get_dtype

print(get_dtype(["x"]))

gives

[('x', '<f8'), ('logP', '<f8'), ('logL', '<f8'), ('it', '<i4')]

whereas you would expect logL to be f4.

Expected behaviour
The data type for logL should change when config.LOGL_DTYPE is changed.

Additional context
This likely affects all of the variables that are in lists in nessai.config.

This can be circumvented using

from nessai import config

config.LOGL_DTYPE = "f4"
config.CORE_PARAMETERS_DTYPE = ["f8", "f4", "i4"]

from nessai.livepoint import reset_extra_live_points_parameters

reset_extra_live_points_parameters()

but this is not how it is intended to work.

The change in v0.9.0 to support regex for adding reparameterisations has lead to a bug where parameters are added multiple times in the reparameterisations.

For example, if the model has names ["x_1", "x_10", "x_11"] then when adding the reparameterisations the names list has duplicates: ['x1', 'x10', 'x11', 'x10', 'x11']. This occurs because "x_1" partially matches "x_10".

Some of the unit tests are failing with errors related to structured arrays.

Compared to tests that recently passed, the main difference seems to be the version of numpy (1.22.4 vs 1.23.0). This new release contains changes to how structured arrays are compared as explained here.

We should investigate this and see if this breaks any existing code in nessai.

Example

np.testing.assert_array_equal(
            x_prime[['logP', 'logL']], x_out[['logL', 'logP']])
E       TypeError: Cannot compare structured arrays unless they have a common dtype.  I.e. `np.result_type(arr1, arr2)` must be defined.

Possible solutions

For testing, we should consider using assert_structured_arrays_equal which was introduced in #170

Add the final insertion indices p-value to the results file.

Describe the solution you'd like
The result.json should contain key along the lines of final_p_value.

Ruff (https://github.com/charliermarsh/ruff) is a rust based alternative to flake8. It seems to be must faster, so we could consider switching to it.

I recently found about NESSAI and I'm trying to implement some of my analysis of radial velocity time-series (for exoplanets) to test the performance compared to dynesty, as I have a few problems which will hugely benefit from a 2x reduction in computing time.

To give some context, I'm working with the data of this article (https://ui.adsabs.harvard.edu/abs/2022A%26A...658A.115F/abstract). So basically 2 time-series fitted simultaneously, with Gaussian processes and a custom likelihood function (written below).

In my first try it is running much slower than dinesty (parallel). My test case takes around 3x longer (in a 4-cores/8-threads laptop). I am also getting a few warnings I have trouble linking to configuration parameters. The final distributions (and logZ) is basically the same.

I have trouble figuring out if I am doing something wrong with NESSAI, or if this is just a problem at which dynesty is very good. I'm also not sure the parallelization is scaling correctly. The CPU load seems quite low (compared to dynesty).

I'm using the sampler as:

cpu = 8

fs = FlowSampler(TestModel(), 
                    output=output, 
                    resume=False, 
                    seed=1235,
                    nlive=live,
                    max_threads=cpu,
                    n_pool=cpu-1,
                    max_poolsize_scale=100)

The model is written at the end of the message.

I'm getting this warning nonstop:
02-24 15:16 nessai.proposal.flowproposal WARNING : Rejection sampling accepted less than 1 percent of samples!

Could the default configuration of the flow proposal have issues with this problem? Or maybe that it's not the right problem for it? I'm also not sure the paralellization is scaling correctly. The CPU load is quite low (compared to dynesty).

Here's the model, in case there's something that shouldn't be written the way it is. The priors are a mix of gaussian and uniform, with distributions created with scity.stats.

global bjd_rv,rv,erv,bjd_fwhm,fwhm,efwhm  #This is the data
global E18R, E19R, E21R  #Some chunks with independent properties
global E18F, E19F, E21F   #Same for the secondary time-series

class TestModel(Model):
    def __init__(self):
        self.names = labels
        self.bounds = dict_labels

    def log_prior(self, x):

        log_p = np.log(self.in_bounds(x))

        for n in self.names:
            log_p += prior_dict[n].logpdf(x[n])
        return log_p

    def log_likelihood(self, x):

        A1   = np.exp(x["LN_A1_RV"])
        A1_2 = np.exp(x["LN_A2_RV"])
        A2   = np.exp(x["LN_A1_FW"])
        A2_2 = np.exp(x["LN_A2_FW"])
        PROT = x["PROT"]
        TS1  = np.exp(x["LN_TS1"])
        TS2  = np.exp(x["LN_TS2"])

        zero_rv = rv.copy()
        zero_rv[E18R] = x['V0_E18_RV']
        zero_rv[E19R] = x['V0_E19_RV']
        zero_rv[E21R] = x['V0_E21_RV']

        res_rv = rv - zero_rv

        erv1 = erv.copy()
        erv1[E18R] = np.sqrt(erv[E18R]*erv[E18R] + np.exp(2*x['LJ_E18_RV']))
        erv1[E19R] = np.sqrt(erv[E19R]*erv[E19R] + np.exp(2*x['LJ_E19_RV']))
        erv1[E21R] = np.sqrt(erv[E21R]*erv[E21R] + np.exp(2*x['LJ_E21_RV']))



        K1 = terms2.SHOTerm(sigma=A1, rho=PROT, tau=TS1) + \
                terms2.SHOTerm(sigma=A1_2,rho=0.5 * PROT, tau=TS2)
        GP1 = celerite2.GaussianProcess(K1)
        GP1.compute(bjd_rv, yerr=erv1)
        lnlike_rv = GP1.log_likelihood(res_rv)



        zero_fw = fwhm.copy()
        zero_fw[E18F] = x['V0_E18_FW']
        zero_fw[E19F] = x['V0_E19_FW']
        zero_fw[E21F] = x['V0_E21_FW']

        res_fw = fwhm - zero_fw

        efwhm1 = efwhm.copy()
        efwhm1[E18F] = np.sqrt(efwhm[E18F]*efwhm[E18F] + np.exp(2*x['LJ_E18_FW']))
        efwhm1[E19F] = np.sqrt(efwhm[E19F]*efwhm[E19F] + np.exp(2*x['LJ_E19_FW']))
        efwhm1[E21F] = np.sqrt(efwhm[E21F]*efwhm[E21F] + np.exp(2*x['LJ_E21_FW']))

        K2 = terms2.SHOTerm(sigma=A2, rho=PROT, tau=TS1) + \
                terms2.SHOTerm(sigma=A2_2, rho=0.5 * PROT, tau=TS2)
        GP2 = celerite2.GaussianProcess(K2)
        GP2.compute(bjd_fwhm, yerr=efwhm1)
        lnlike_fw = GP2.log_likelihood(res_fw)

        return lnlike_rv + lnlike_fw

A change introduced in v0.10.0 has significantly slowed down bilby analyses, I suspect this is due to the changes made to support non-vectorised prior functions.

I'm looking into the cause of this and plan to make bug fix release to address this. In the meantime, I recommend users stick with v0.9.1.

The multiprocess library uses dill instead of pickle for serialisation so it might be a useful optional alternative to the standard multiprocessing library.

Currently, the main limitation is the lack of support for ABC but this is being addressed in uqfoundation/dill#450 so we should consider trying it out once there's a release that contains this fix.

Some of the default settings in nessai are a bit out-of-date and more optimal settings could be used instead. I'm creating this issue to track the settings that could/should be updated.

General sampling settings

• Set maximum_uninformed such that by default it's a few times the number of live points

FlowProposal settings

• Set update_poolsize=True by default

Flow settings

• Increase Batch size
• Maybe increase patience?

GW Reparameterisations

• Change angle-sine to use default reparam

A lot of defaults were changed in #126 and a lot of examples have extra configuration that is probably no longer needed. We should update examples to reflect the new defaults.

Various examples use np.log(self.in_bounds(x)) in the prior. This results in np.log(True) which returns float16, this can lead to rounding issues when computing the prior.

To-Do

• Add a check in verify_model to warn the user if the log_prior returns float16
• Update examples to cast the value to float64.

See for example: https://readthedocs.org/projects/nessai/builds/22109909/

Describe the bug
If the user specifies the sky-ra-dec reparameterisation but sets the prior on RA to [-pi, pi], then the resulting reparameterisation is not valid. This is because it is assumed the prior on RA is uniform on [0, 2pi] but this is not checked.

To Reproduce
Run full_gw_example.py and set prior on RA to uniform on [-pi, pi]

verify_rescaling then returns:

RuntimeError: Duplicate samples must map to same input values. Check the rescaling and inverse rescaling functions for ra.

Expected behaviour
This should be checked and either an error raised or it accounted for in the reparameterisation.

nessai is now included in the bilby master branch and the documentation should be updated to reflect this.

• Update README
• Update installation instructions in the docs

Describe the bug
The example eggbox.py returns an error instead of finishing normally

Traceback (most recent call last):
  File "eggbox.py", line 68, in <module>
    fs.run()
  File "/home/michaelwilliams/git_repos/nessai/nessai/flowsampler.py", line 140, in run
    self.ns.nested_sampling_loop()
  File "/home/michaelwilliams/git_repos/nessai/nessai/nestedsampler.py", line 1185, in nested_sampling_loop
    self.consume_sample()
  File "/home/michaelwilliams/git_repos/nessai/nessai/nestedsampler.py", line 609, in consume_sample
    c, proposed = next(self.yield_sample(worst))
  File "/home/michaelwilliams/git_repos/nessai/nessai/nestedsampler.py", line 560, in yield_sample
    newparam = self.proposal.draw(copy(oldparam))
  File "/home/michaelwilliams/git_repos/nessai/nessai/proposal/flowproposal.py", line 1537, in draw
    self.populate(worst_point, N=self.poolsize)
  File "/home/michaelwilliams/git_repos/nessai/nessai/proposal/flowproposal.py", line 1449, in populate
    z_samples[accepted:(accepted+n), ...] = z[:n]
UnboundLocalError: local variable 'z_samples' referenced before assignment

To Reproduce
Run eggbox.py in examples/: https://github.com/mj-will/nessai/blob/main/examples/eggbox.py

Expected behaviour
The example should run and finish sampling.

Currently, the docs are not tested in any of the GitHub Actions, we should consider adding a workflow to ensure the docs build correctly.

Possible solutions

• https://github.com/marketplace/actions/sphinx-build

Describe the bug
If prior_sampling=True and the nested sampler is being run through FlowSampler and an error is raised because the nested sampler only returns samples, not the expected evidence + samples.

Error

Traceback (most recent call last):
  File "/home/michael/git_repos/nessai/examples/sample_prior.py", line 82, in <module>
    fp.run()
  File "/home/michael/git_repos/nessai/nessai/flowsampler.py", line 116, in run
    self.logZ, self.nested_samples = \
ValueError: too many values to unpack (expected 2)

To Reproduce

fp = FlowSampler(GaussianModel(), prior_sampling=True)
pp.run()

Expected behaviour
The sampler should draw the prior samples and those should be saved by FlowSampler

Describe the bug

When specifying that no reparameterisation should be used for a parameter, an error is raised due to the prior bounds being passed to NullReparameterisation.

For example:

reparameterisations={                                         
    'x': {'reparameterisation': 'default'},                              
    'y': 'none'                                               
}    

raises

Traceback (most recent call last):
  File "reparameterisations_example.py", line 73, in <module>
    fp.run()
  File "/home/michaelwilliams/git_repos/nessai/nessai/flowsampler.py", line 110, in run
    self.ns.initialise()
  File "/home/michaelwilliams/git_repos/nessai/nessai/nestedsampler.py", line 669, in initialise
    self._flow_proposal.initialise()
  File "/home/michaelwilliams/git_repos/nessai/nessai/proposal/flowproposal.py", line 430, in initialise
    self.set_rescaling()
  File "/home/michaelwilliams/git_repos/nessai/nessai/proposal/flowproposal.py", line 624, in set_rescaling
    self.configure_reparameterisations(self.reparameterisations)
  File "/home/michaelwilliams/git_repos/nessai/nessai/proposal/flowproposal.py", line 573, in configure_reparameterisations
    r = rc(prior_bounds=prior_bounds, **default_config)
TypeError: __init__() got an unexpected keyword argument 'prior_bounds'

Expected behaviour

An instance of NullReparameterisation should be added to the reparameterisation.

We should add a workflow that tests compatibility with bilby.

This workflow should probably be allowed to fail given that there will be cases where we'll need to update the bilby API

In some scenarios, there are extra parameters that are sampled but are not of interest, e.g. calibration in GW analyses. It would be helpful to be able to exclude them from the trace plot.

Describe the bug
When trying to re-run a sampler in which resume was not explicitly set to False (e.g. in a jupyter notebook),
the following error appears:

Traceback (most recent call last):
  File "/home/rodrigo/gravity/git/nessai/examples/2d_gaussian.py", line 63, in <module>
    fs = FlowSampler(GaussianModel(), output=output, seed=1234)
  File "/home/rodrigo/opt/miniconda3/envs/cw-fun/lib/python3.9/site-packages/nessai/flowsampler.py", line 76, in __init__
    kwargs['flow_config'], weights_file)
KeyError: 'flow_config'

To Reproduce
This can be easily reproduced by doing the following in the 2d_gaussian.py example:

fs = FlowSampler(GaussianModel(), output=output, seed=1234)
fs.run()
fs = FlowSampler(GaussianModel(), output=output, seed=1234)

The slow integration tests are being skipped because some of the normal integration tests are failing and being rerun.

See for example https://github.com/mj-will/nessai/actions/runs/2035671470

Cause
In pytest-integration slow integration tests are skipped in normal integration tests fail as described here:

If a (quick) integration tests fails, slow integration tests are not run.

So it seems that tests that are rerun and pass still count as failing.

Describe the bug

Trace plot causes crash with the error ValueError: Image size of 1200x69120 pixels is too large. It must be less than 2^16 in each direction.

To Reproduce
I have 92 parameters. I would imagine trying a 100-D gaussian would reproduce.

Expected behaviour
No crash. If it's not possible to plot then catch the exception?

Screenshots

Traceback (most recent call last):
  File "/home/john/work/students/ann/fast_tutorial_mylikelihood.py", line 303, in <module>
    result = bilby.run_sampler(
  File "/home/john/src/bilby/bilby/core/sampler/__init__.py", line 234, in run_sampler
    result = sampler.run_sampler()
  File "/home/john/src/bilby/bilby/core/sampler/base_sampler.py", line 97, in wrapped
    output = method(self, *args, **kwargs)
  File "/home/john/src/bilby/bilby/core/sampler/nessai.py", line 247, in run_sampler
    self.fs.run(**run_kwargs)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/flowsampler.py", line 240, in run
    self.logZ, self.nested_samples = self.ns.nested_sampling_loop()
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/samplers/nestedsampler.py", line 1249, in nested_sampling_loop
    self.update_state()
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/samplers/nestedsampler.py", line 1163, in update_state
    self.plot_trace(
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/samplers/nestedsampler.py", line 1094, in plot_trace
    fig = plot_trace(
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/plot.py", line 58, in wrapper
    return func(*args, **kwargs)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/nessai/plot.py", line 501, in plot_trace
    fig.savefig(filename, bbox_inches="tight")
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/figure.py", line 3285, in savefig
    self.canvas.print_figure(fname, **kwargs)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backend_bases.py", line 2308, in print_figure
    renderer = _get_renderer(
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backend_bases.py", line 1559, in _get_renderer
    print_method(io.BytesIO())
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backend_bases.py", line 2204, in <lambda>
    print_method = functools.wraps(meth)(lambda *args, **kwargs: meth(
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/_api/deprecation.py", line 410, in wrapper
    return func(*inner_args, **inner_kwargs)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backends/backend_agg.py", line 517, in print_png
    self._print_pil(filename_or_obj, "png", pil_kwargs, metadata)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backends/backend_agg.py", line 463, in _print_pil
    FigureCanvasAgg.draw(self)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backends/backend_agg.py", line 399, in draw
    self.renderer = self.get_renderer()
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/_api/deprecation.py", line 384, in wrapper
    return func(*inner_args, **inner_kwargs)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backends/backend_agg.py", line 416, in get_renderer
    self.renderer = RendererAgg(w, h, self.figure.dpi)
  File "/home/john/.conda/envs/ann_msci/lib/python3.10/site-packages/matplotlib/backends/backend_agg.py", line 84, in __init__
    self._renderer = _RendererAgg(int(width), int(height), dpi)
ValueError: Image size of 1200x69120 pixels is too large. It must be less than 2^16 in each direction.

Additional context
Add any other context about the problem here.

NumPy has already dropped support for Python 3.8, we should do the same.

This would also be a good chance to include Python 3.11 in the CI.

Describe the bug
When trying to us a multiprocessing pool from ray for the pool argument introduced in #148 an AttributeError is raised:

Traceback (most recent call last):
  File "/home/michael/projects/nessai-dev/ray/test_ray.py", line 40, in <module>
    fs = FlowSampler(
  File "/home/michael/git_repos/nessai/nessai/flowsampler.py", line 109, in __init__
    self.ns = NestedSampler(model, output=self.output,
  File "/home/michael/git_repos/nessai/nessai/nestedsampler.py", line 174, in __init__
    self.model.configure_pool(pool=pool, n_pool=n_pool)
  File "/home/michael/git_repos/nessai/nessai/model.py", line 213, in configure_pool
    self.n_pool = self.pool._processes
AttributeError: 'Pool' object has no attribute '_processes'

To Reproduce

#!/usr/bin/env python

import numpy as np
from scipy.stats import norm
from ray.util.multiprocessing import Pool

from nessai.flowsampler import FlowSampler
from nessai.model import Model
from nessai.utils import setup_logger
from nessai.utils.multiprocessing import initialise_pool_variables

output = './outdir/test_ray/'
logger = setup_logger(output=output)


class GaussianModel(Model):
    def __init__(self):
        self.names = ['x', 'y']
        self.bounds = {'x': [-10, 10], 'y': [-10, 10]}

    def log_prior(self, x):
        log_p = np.log(self.in_bounds(x))
        for n in self.names:
            log_p -= np.log(self.bounds[n][1] - self.bounds[n][0])
        return log_p

    def log_likelihood(self, x):
        log_l = np.zeros(x.size)
        for n in self.names:
            log_l += norm.logpdf(x[n])
        return log_l


model = GaussianModel()
pool = Pool(
    initializer=initialise_pool_variables,
    initargs=(model,)
)

fs = FlowSampler(
    model,
    output=output,
    resume=False,
    seed=1234,
    pool=pool,
)

fs.run()

Expected behaviour
Sampler should run and use the ray multiprocessing pool.

Additional context

• Version: main@18fdccb3773a48a3404c11874866874634900597
• ray multiprocessing docs

Currently, if logging to the same directory, the logger appends to the log file instead of overwriting the existing one. Generally, this isn't a problem but there are cases when it might be more helpful to overwrite the log.

We can probably do this by adding an option to setup_logger that changes the mode as described in the logging documentation for FileHandler or adding an option to pass kwargs.

We should update the version of black in the requirements and pre-commits to use the 2024 style.

In some cases, multiple parameters will also use the same reparameterisation and specifying a list is impractical because, for example, the settings are in an ini file.

We should support the use of regex patterns when specifying reparameterisations.

The test with ray recently started failing in the CI, locally the test still passes. We should investigate the cause of this.

Describe the bug
corner_plot returns a ValueError if an ndarray is passed for the truths kwarg. A list object is fine.

To Reproduce

from nessai.plot import corner_plot
from nessai.livepoint import dict_to_live_points
import matplotlib.pyplot as plt
import numpy as np

truth = np.array([
    0.,
    0.,
])

samples = {
    "x0":np.random.normal(size=10000),
    "x1":np.random.normal(size=10000)
}

result = dict_to_live_points(samples)

corner_plot(
    result,
    exclude=["logL","logP","it"],
    truths = truth,
)

Expected behaviour
A corner plot should appear.

Screenshots
N/A

Additional context
nessai version: 0.11.0
python version: 3.12.1

Describe the bug

The sampler crashes when running bilby_example.py if gwpy is installed. This error is not caused by code in nessai.

The following error is produced:

Traceback (most recent call last):
  File "bilby_example.py", line 55, in <module>
    result = bilby.run_sampler(outdir=outdir, label=label, resume=True, plot=True,
  File "/nfshome/store03/users/michael.williams/git_repos/bilby/bilby/core/sampler/__init__.py", line 185, in run_sampler
    result = sampler.run_sampler()
  File "/nfshome/store03/users/michael.williams/git_repos/bilby/bilby/core/sampler/nessai.py", line 142, in run_sampler
    out.run(save=True, plot=self.plot)
  File "/nfshome/store03/users/michael.williams/git_repos/nessai/nessai/flowsampler.py", line 128, in run
    plot.plot_live_points(self.posterior_samples,
  File "/nfshome/store03/users/michael.williams/git_repos/nessai/nessai/plot.py", line 56, in plot_live_points
    fig.map_offdiag(sns.scatterplot, hue=hue,
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/axisgrid.py", line 1311, in map_offdiag
    self.map_lower(func, **kwargs)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/axisgrid.py", line 1281, in map_lower
    self._map_bivariate(func, indices, **kwargs)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/axisgrid.py", line 1465, in _map_bivariate
    self._plot_bivariate(x_var, y_var, ax, func, **kws)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/axisgrid.py", line 1505, in _plot_bivariate
    func(x=x, y=y, **kwargs)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/_decorators.py", line 46, in inner_f
    return f(**kwargs)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/relational.py", line 820, in scatterplot
    p.plot(ax, kwargs)
  File "/home/michael.williams/.conda/envs/nessai/lib/python3.8/site-packages/seaborn/relational.py", line 651, in plot
    points = ax.scatter(*args, **kws)
TypeError: scatter() takes from 3 to 4 positional arguments but 5 were given

To Reproduce

import numpy as np
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import gwpy

x = np.random.randn(100, 2)
df = pd.DataFrame(data=x)

fig = sns.PairGrid(df, corner=True, diag_sharey=False)
fig.map_diag(plt.hist)
fig.map_offdiag(sns.scatterplot)

Expected behaviour
The sampler should produce a posterior plot and finish finalising the sampling process.

Additional context
The error has been reported to the gwpy team gwpy/gwpy#1311 and in the meantime, the sampler can run with plot=False when using bilby or without gwpy installed to avoid this error.

I'm documenting this error here for anyone that encounters it and is unsure what causes it.

An inference mode was added to PyTorch in version 1.9.0. It seems to be an even more extreme version of torch.no_grad() and promises to be even faster.

To quote the documentation here:

Inference mode is the extreme version of no-grad mode. Just like in no-grad mode, computations in inference mode are not recorded in the backward graph, but enabling inference mode will allow PyTorch to speed up your model even more.

We should consider testing inference mode in the population stage of nessai since this requires repeated calls to the normalising flow and is the main bottleneck at the moment.

Useful links

PyTorch 1.9.0 release notes
Inference mode documentation
Explanation of different autograd mechanics

Currently, nessai only supports the fork start method for multiprocessing, since it is the only methods that does not require pickling the model. We should consider support other start methods.

As it stands, there is nothing in the model that prevents it from being pickled, so we should test the other methods and see if enabling them is feasible. There are however cases where the user may define the model in such a way that it is not pickleable.

Python 3.12 changes how some of mock functions and classes such that one cannot call:

assert mock.called_once_with(...)

We should remove all instances of this in the test suite.