ptemcee plot_mean_log_posterior memory leak with negative infinities

Hi there,

While using ptemcee, I noticed a memory leak and after some investigation narrowed it down to plot_mean_log_posterior called when checkpointing each iteration of the sampler. In particular, it's using large amounts of memory (>= 100 GB) when trying to save this figure, when there are negative "infinities" (~ -1E305) in the mean_log_posterior array.

Running on ldas-pcdev2, I was able to reproduce this issue with the following standalone script, and reading the same array from meanlogposterior.npy; when I run this script, I get a segfault.

import numpy as np

def plot_mean_log_posterior(mean_log_posterior, outdir, label):
    import matplotlib.pyplot as plt

    ntemps, nsteps = mean_log_posterior.shape
    ymax = np.max(mean_log_posterior)
    ymin = np.min(mean_log_posterior[:, -100:])
    ymax += 0.1 * (ymax - ymin)
    ymin -= 0.1 * (ymax - ymin)

    fig, ax = plt.subplots()
    idxs = np.arange(nsteps)
    ax.plot(idxs, mean_log_posterior.T)
    ax.set(xlabel="Iteration", ylabel=r"$\langle\mathrm{log-posterior}\rangle$",
           ylim=(ymin, ymax))
    fig.tight_layout()
    fig.savefig("{}/{}_checkpoint_meanlogposterior.png".format(outdir, label))
    plt.close(fig)

if __name__ == "__main__":
    path = "/path/to/meanlogposterior.npy"
    mean_log_posterior = np.load(path)
    plot_mean_log_posterior(mean_log_posterior, ".", "test")

We were able to correct this issue by filtering out the negative "infinities", like so:

def plot_mean_log_posterior(mean_log_posterior, outdir, label):
    import matplotlib.pyplot as plt

    # new code
    mean_log_posterior[ mean_log_posterior < -1E100 ] = np.nan

    ntemps, nsteps = mean_log_posterior.shape
    ymax = np.nanmax(mean_log_posterior) # modified
    ymin = np.nanmin(mean_log_posterior[:, -100:]) # modified
    ymax += 0.1 * (ymax - ymin)
    ymin -= 0.1 * (ymax - ymin)

    fig, ax = plt.subplots()
    idxs = np.arange(nsteps)
    ax.plot(idxs, mean_log_posterior.T)
    ax.set(xlabel="Iteration", ylabel=r"$\langle\mathrm{log-posterior}\rangle$",
           ylim=(ymin, ymax))
    fig.tight_layout()
    fig.savefig("{}/{}_checkpoint_meanlogposterior.png".format(outdir, label))
    plt.close(fig)

However, there may also be better solutions to this problem! For reference, I am using a (lightly-modified) bilby 1.1.5, numpy 1.21.1, and matplotlib 3.4.2