diff --git a/CHANGELOG.md b/CHANGELOG.md index f4c4d114a..711f6b9a0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Features +- [#570](https://github.com/pybop-team/PyBOP/pull/570) - Updates the contour and surface plots, adds mixed chain effective sample size computation, x0 to optim.log - [#566](https://github.com/pybop-team/PyBOP/pull/566) - Adds `UnitHyperCube` transformation class, fixes incorrect application of gradient transformation. - [#569](https://github.com/pybop-team/PyBOP/pull/569) - Adds parameter specific learning rate functionality to GradientDescent optimiser. - [#282](https://github.com/pybop-team/PyBOP/issues/282) - Restructures the examples directory. diff --git a/examples/scripts/comparison_examples/spm_XNES.py b/examples/scripts/comparison_examples/spm_XNES.py index 09d0ab395..90d5aa1d0 100644 --- a/examples/scripts/comparison_examples/spm_XNES.py +++ b/examples/scripts/comparison_examples/spm_XNES.py @@ -52,3 +52,6 @@ # Plot the cost landscape with optimisation path pybop.plot.surface(optim) + +# Plot contour +pybop.plot.contour(optim) diff --git a/pybop/optimisers/base_optimiser.py b/pybop/optimisers/base_optimiser.py index d008f5eeb..0a790accf 100644 --- a/pybop/optimisers/base_optimiser.py +++ b/pybop/optimisers/base_optimiser.py @@ -61,11 +61,11 @@ def __init__( ): # First set attributes to default values self.parameters = Parameters() - self.x0 = None + self.x0 = optimiser_kwargs.get("x0", []) + self.log = dict(x=[], x_best=[], cost=[], cost_best=[], x0=[]) self.bounds = None self.sigma0 = 0.02 self.verbose = True - self.log = dict(x=[], x_best=[], cost=[], cost_best=[]) self.minimising = True self._transformation = None self._needs_sensitivities = False @@ -86,7 +86,6 @@ def __init__( else: try: - self.x0 = optimiser_kwargs.get("x0", []) cost_test = cost(self.x0) warnings.warn( "The cost is not an instance of pybop.BaseCost, but let's continue " @@ -132,6 +131,7 @@ def set_base_options(self): """ # Set initial values, if x0 is None, initial values are unmodified. self.parameters.update(initial_values=self.unset_options.pop("x0", None)) + self.log_update(x0=self.parameters.reset_initial_value()) self.x0 = self.parameters.reset_initial_value(apply_transform=True) # Set default bounds (for all or no parameters) @@ -197,7 +197,7 @@ def _run(self): """ raise NotImplementedError - def log_update(self, x=None, x_best=None, cost=None, cost_best=None): + def log_update(self, x=None, x_best=None, cost=None, cost_best=None, x0=None): """ Update the log with new values. @@ -234,9 +234,8 @@ def apply_transformation(values): self.log["x"].extend(x) if x_best is not None: - x_best = convert_to_list(x_best) - x_best = apply_transformation(x_best) - self.log["x_best"].extend([x_best]) + x_best = apply_transformation([x_best]) + self.log["x_best"].extend(x_best) if cost is not None: cost = convert_to_list(cost) @@ -246,6 +245,9 @@ def apply_transformation(values): cost_best = convert_to_list(cost_best) self.log["cost_best"].extend(cost_best) + if x0 is not None: + self.log["x0"].extend(x0) + def name(self): """ Returns the name of the optimiser, to be overwritten by child classes. diff --git a/pybop/plot/contour.py b/pybop/plot/contour.py index 922627028..bab365ceb 100644 --- a/pybop/plot/contour.py +++ b/pybop/plot/contour.py @@ -36,6 +36,8 @@ def contour( bounds : numpy.ndarray, optional A 2x2 array specifying the [min, max] bounds for each parameter. If None, uses `cost.parameters.get_bounds_for_plotly`. + apply_transform : bool, optional + Converts contour landscape into transformed landscape that was used by the optimiser. steps : int, optional The number of grid points to divide the parameter space into along each dimension (default: 10). show : bool, optional @@ -156,11 +158,12 @@ def contour( layout_options = dict( title="Cost Landscape", title_x=0.5, - title_y=0.9, + title_y=0.905, width=600, height=600, xaxis=dict(range=bounds[0], showexponent="last", exponentformat="e"), yaxis=dict(range=bounds[1], showexponent="last", exponentformat="e"), + legend=dict(orientation="h", yanchor="bottom", y=1, xanchor="right", x=1), ) if hasattr(cost, "parameters"): name = cost.parameters.keys() @@ -170,7 +173,8 @@ def contour( # Create contour plot and update the layout fig = go.Figure( - data=[go.Contour(x=x, y=y, z=costs, connectgaps=True)], layout=layout + data=[go.Contour(x=x, y=y, z=costs, colorscale="Viridis", connectgaps=True)], + layout=layout, ) if plot_optim: @@ -184,7 +188,8 @@ def contour( mode="markers", marker=dict( color=[i / len(optim_trace) for i in range(len(optim_trace))], - colorscale="YlOrBr", + colorscale="Greys", + size=8, showscale=False, ), showlegend=False, @@ -192,21 +197,40 @@ def contour( ) # Plot the initial guess - if optim.x0 is not None: + if optim.log["x0"] is not None: fig.add_trace( go.Scatter( - x=[optim.x0[0]], - y=[optim.x0[1]], + x=[optim.log["x0"][0]], + y=[optim.log["x0"][1]], mode="markers", - marker_symbol="circle", + marker_symbol="x", marker=dict( - color="mediumspringgreen", - line_color="mediumspringgreen", + color="white", + line_color="black", line_width=1, size=14, showscale=False, ), - showlegend=False, + name="Initial values", + ) + ) + + # Plot optimised value + if optim.log["x_best"] is not None: + fig.add_trace( + go.Scatter( + x=[optim.log["x_best"][-1][0]], + y=[optim.log["x_best"][-1][1]], + mode="markers", + marker_symbol="cross", + marker=dict( + color="black", + line_color="white", + line_width=1, + size=14, + showscale=False, + ), + name="Final values", ) ) diff --git a/pybop/plot/voronoi.py b/pybop/plot/voronoi.py index 4c3e2483a..4dad41223 100644 --- a/pybop/plot/voronoi.py +++ b/pybop/plot/voronoi.py @@ -293,9 +293,9 @@ def surface( y=y_optim, mode="markers", marker=dict( + color=[i / len(x_optim) for i in range(len(x_optim))], + colorscale="Greys", size=8, - symbol="x-thin", - line=dict(color="white", width=1), showscale=False, ), text=[f"f={val:.2f}" for val in f], @@ -308,31 +308,52 @@ def surface( if optim.x0 is not None: fig.add_trace( go.Scatter( - x=[optim.x0[0]], - y=[optim.x0[1]], + x=[optim.log["x0"][0]], + y=[optim.log["x0"][1]], mode="markers", - marker_symbol="circle", + marker_symbol="x", marker=dict( - color="mediumspringgreen", - line_color="mediumspringgreen", + color="white", + line_color="black", line_width=1, size=14, showscale=False, ), - showlegend=False, + name="Initial values", ) ) + + # Plot optimised value + if optim.log["x_best"] is not None: + fig.add_trace( + go.Scatter( + x=[optim.log["x_best"][-1][0]], + y=[optim.log["x_best"][-1][1]], + mode="markers", + marker_symbol="cross", + marker=dict( + color="black", + line_color="white", + line_width=1, + size=14, + showscale=False, + ), + name="Final values", + ) + ) + names = optim.cost.parameters.keys() fig.update_layout( title="Voronoi Cost Landscape", title_x=0.5, - title_y=0.9, + title_y=0.905, xaxis_title=names[0], yaxis_title=names[1], width=600, height=600, xaxis=dict(range=xlim, showexponent="last", exponentformat="e"), yaxis=dict(range=ylim, showexponent="last", exponentformat="e"), + legend=dict(orientation="h", yanchor="bottom", y=1, xanchor="right", x=1), ) fig.update_layout(**layout_kwargs) if show: diff --git a/pybop/samplers/mcmc_summary.py b/pybop/samplers/mcmc_summary.py index c67a3b0fd..851d84259 100644 --- a/pybop/samplers/mcmc_summary.py +++ b/pybop/samplers/mcmc_summary.py @@ -65,7 +65,7 @@ def get_summary_statistics(self): for key, func in summary_funs.items() } - def plot_trace(self): + def plot_trace(self, **kwargs): """ Plot trace plots for the posterior samples. """ @@ -83,9 +83,10 @@ def plot_trace(self): xaxis_title="Sample Index", yaxis_title="Value", ) + fig.update_layout(**kwargs) fig.show() - def plot_chains(self): + def plot_chains(self, **kwargs): """ Plot posterior distributions for each chain. """ @@ -117,9 +118,10 @@ def plot_chains(self): xaxis_title="Value", yaxis_title="Density", ) + fig.update_layout(**kwargs) fig.show() - def plot_posterior(self): + def plot_posterior(self, **kwargs): """ Plot the summed posterior distribution across chains. """ @@ -142,7 +144,9 @@ def plot_posterior(self): xaxis_title="Value", yaxis_title="Density", ) + fig.update_layout(**kwargs) fig.show() + return fig def summary_table(self): """ @@ -192,21 +196,45 @@ def _autocorrelate_negative(self, autocorrelation): negative_indices[0] if negative_indices.size > 0 else len(autocorrelation) ) - def effective_sample_size(self): - """ - Computes the effective sample size for each parameter in each chain. + def effective_sample_size(self, mixed_chains=False): """ + Computes the effective sample size (ESS) for each parameter in each chain. + + Parameters + ---------- + mixed_chains : bool, optional + If True, the ESS is computed for all samplers mixed into a single chain. + Defaults to False. + Returns + ------- + list + A list of effective sample sizes for each parameter in each chain, + or for the mixed chain if `mixed_chains` is True. + + Raises + ------ + ValueError + If there are fewer than two samples in the data. + """ if self.all_samples.shape[0] < 2: raise ValueError("At least two samples must be given.") - ess = [] - for _, chain in enumerate(self.chains): + def compute_ess(samples): + """Helper function to compute the ESS for a single set of samples.""" + ess = [] for j in range(self.num_parameters): - rho = self.autocorrelation(chain[:, j]) + rho = self.autocorrelation(samples[:, j]) T = self._autocorrelate_negative(rho) - ess.append(len(chain[:, j]) / (1 + 2 * rho[:T].sum())) + ess.append(len(samples[:, j]) / (1 + 2 * rho[:T].sum())) + return ess + + if mixed_chains: + return compute_ess(self.all_samples) + ess = [] + for chain in self.chains: + ess.extend(compute_ess(chain)) return ess def rhat(self): diff --git a/tests/unit/test_sampling.py b/tests/unit/test_sampling.py index e410242e6..5d4141910 100644 --- a/tests/unit/test_sampling.py +++ b/tests/unit/test_sampling.py @@ -153,18 +153,32 @@ def test_initialisation_and_run( assert samples.shape == (chains, 1, 2) @pytest.mark.unit - def test_effectie_sample_size(self, log_posterior): - sampler = pybop.SliceStepoutMCMC( - log_pdf=log_posterior, - chains=1, - max_iterations=1, - ) - results = sampler.run() - summary = pybop.PosteriorSummary(results) + def test_effective_sample_size(self, log_posterior): + chains = np.asarray([[[0, 0]]]) + summary = pybop.PosteriorSummary(chains) with pytest.raises(ValueError, match="At least two samples must be given."): summary.effective_sample_size() + n_chains = 3 + sampler = pybop.HaarioBardenetACMC( + log_pdf=log_posterior, + chains=n_chains, + max_iterations=3, + ) + chains = sampler.run() + summary = pybop.PosteriorSummary(chains) + + # Non mixed chains + ess = summary.effective_sample_size() + assert len(ess) == log_posterior.n_parameters * n_chains + assert all(e > 0 for e in ess) # ESS should be positive + + # Mixed chains + ess = summary.effective_sample_size(mixed_chains=True) + assert len(ess) == log_posterior.n_parameters + assert all(e > 0 for e in ess) + @pytest.mark.unit def test_single_parameter_sampling(self, model, dataset, MCMC, chains): parameters = pybop.Parameters(