Correctly compute limits in `prepare_for_plot` in case of samples with inf value

[zinaStef]

In issue #1037, the sampler produced samples that had inf value and the pair_plot function was failing when computing the x-axis max limit in the prepare_for_plot function.

TODO:
Change prepare_for_plot function to handle inf values.

[janfb]

@Matthijspals are you working on this, or should be un-assign you?

[Matthijspals]

@Matthijspals are you working on this, or should be un-assign you?

This could be a quick fix - but what is expected behaviour here? Do we actually want to ignore nans everywhere in the plotting function? And what if a datapoint is nan for one dimension, but not the others - should it be included in marginal plots for the non nan dimensions?

[janfb]

I would probably be easier to just exclude the entire data point.
I think it makes sense to issue a warning about NaNs in the prepare_for_plot function and to add an flag exclude_invalid_data: bool=True to exclude the values in the downstream plotting functions. Similar to how we do it during inference:

sbi/sbi/utils/sbiutils.py

Lines 360 to 372 in 0b5f931

	if num_nans + num_infs > 0:
	if exclude_invalid_x:
	logging.warning(
	f"Found {num_nans} NaN simulations and {num_infs} Inf simulations. "
	"They will be excluded from training."
	)
	else:
	logging.warning(
	f"Found {num_nans} NaN simulations and {num_infs} Inf simulations. "
	"They are not excluded from training due to `exclude_invalid_x=False`."
	"Training will likely fail, we strongly recommend "
	f"`exclude_invalid_x=True` for {algorithm}."
	)

[janfb]

@Matthijspals has this been fixed in the pairplot PR back then, or not?
If not, would you be up for proving the fix?

[Matthijspals]

@Matthijspals has this been fixed in the pairplot PR back then, or not?
If not, would you be up for proving the fix?

Was not yet fixed - up for doing it later this week!

[manuelgloeckler]

Actually, another improvement would be to cast inputs to numpy explicitly. The plotting does currently fail if you input a JAX array, which is a bit annoying.

[manuelgloeckler]

This should be closed with #1185