"""Creates the data for the tests."""

data_path = os.path.join(TEST_DIR_PATH, "plotting_data.pickle")

with open(data_path, "rb") as f:

loaded_data = pickle.load(f)

data = PlotEventData()

data.modelType = loaded_data["modelType"]

data.dataPresent = loaded_data["dataPresent"]

data.subRoughs = loaded_data["subRoughs"]

data.resample = loaded_data["resample"]

data.resampledLayers = loaded_data["resampledLayers"]

data.reflectivity = loaded_data["reflectivity"]

data.shiftedData = loaded_data["shiftedData"]

data.sldProfiles = loaded_data["sldProfiles"]

data.contrastNames = ["D2O", "SMW", "H2O"]

"""Creates the fake domains data for the tests."""

domains_data = data()

for sld_list in domains_data.sldProfiles:

sld_list.append(sld_list[0])

"""Creates the fixture for the tests."""

plt.close("all")

figure = plt.subplots(1, 2)[0]

RATplot.plot_ref_sld_helper(fig=figure, data=domains_data() if request.param else data())

plt.close("all")

figure = plt.subplots(1, 2)[0]

dat = data()

confidence_intervals = {

"reflectivity": [

(curve[:, 1] - curve[:, 1] * 0.5, curve[:, 1] + curve[:, 1] * 0.5) for curve in dat.reflectivity

],

"sld": [

[(curve[:, 1] - curve[:, 1] * 0.1, curve[:, 1] + curve[:, 1] * 0.1) for curve in sld]

for sld in dat.sldProfiles

],

}

RATplot.plot_ref_sld_helper(data=dat, fig=figure, confidence_intervals=confidence_intervals)

"""Tests the axis formatting of the figure."""

ref_plot = fig.axes[0]

sld_plot = fig.axes[1]

assert fig.axes[0].get_subplotspec().get_gridspec().get_geometry() == (1, 2)

assert len(fig.axes) == 2

assert ref_plot.get_xlabel() == "$Q_{z} (\u00c5^{-1})$"

assert ref_plot.get_xscale() == "log"

assert ref_plot.get_ylabel() == "Reflectivity"

assert ref_plot.get_yscale() == "log"

assert [label._text for label in ref_plot.get_legend().texts] == ["D2O", "SMW", "H2O"]

assert sld_plot.get_xlabel() == "$Z (\u00c5)$"

assert sld_plot.get_xscale() == "linear"

assert sld_plot.get_ylabel() == "$SLD (\u00c5^{-2})$"

assert sld_plot.get_yscale() == "linear"

labels = [label._text for label in sld_plot.get_legend().texts]

if len(labels) == 3:

assert labels == ["D2O", "SMW", "H2O"]

else:

assert labels == [

"D2O Domain 1",

"D2O Domain 2",

"SMW Domain 1",

"SMW Domain 2",

"H2O Domain 1",

"H2O Domain 2",

"""Tests the color formatting of the figure."""

ref_plot = fig.axes[0]

sld_plot = fig.axes[1]

assert len(ref_plot.get_lines()) == 6

assert len(sld_plot.get_lines()) == 6

for i in range(0, len(ref_plot.get_lines()), 2):

# Tests whether the color of the line and the errorbars match on the ref_plot

assert ref_plot.containers[i // 2][2][0]._original_edgecolor == ref_plot.get_lines()[i].get_color()

# Tests whether the color of the sld and resampled_sld match on the sld_plot

"""Test that shading is correctly added to the figure when confidence intervals are supplied."""

# the shading is of type PolyCollection

for axes in fig.axes:

components = axes.get_children()

assert not any(isinstance(comp, PolyCollection) for comp in components)

for axes in bayes_fig.axes:

components = axes.get_children()

"""Tests the makeSLDProfile function called with

RATplot.plot_ref_sld_helper(data(), plt.subplots(1, 2)[0])

assert mock.call_count == 3

assert mock.call_args_list[0].args[0] == 2.07e-06

assert mock.call_args_list[0].args[1] == 6.28e-06

assert mock.call_args_list[0].args[3] == 0.0

assert mock.call_args_list[0].args[4] == 1

assert mock.call_args_list[1].args[0] == 2.07e-06

assert mock.call_args_list[1].args[1] == 1.83e-06

assert mock.call_args_list[1].args[3] == 0.0

assert mock.call_args_list[1].args[4] == 1

assert mock.call_args_list[2].args[0] == 2.07e-06

assert mock.call_args_list[2].args[1] == -5.87e-07

assert mock.call_args_list[2].args[3] == 0.0

plot_data = data()

with RATplot.LivePlot() as figure:

assert len(figure.axes) == 2

notify(EventTypes.Plot, plot_data)

plt.close(figure)

notify(EventTypes.Plot, plot_data)

assert mock.call_count == 3

assert mock.call_args_list[0].args[0] == 2.07e-06

assert mock.call_args_list[0].args[1] == 6.28e-06

assert mock.call_args_list[0].args[3] == 0.0

assert mock.call_args_list[0].args[4] == 1

assert mock.call_args_list[1].args[0] == 2.07e-06

assert mock.call_args_list[1].args[1] == 1.83e-06

assert mock.call_args_list[1].args[3] == 0.0

assert mock.call_args_list[1].args[4] == 1

assert mock.call_args_list[2].args[0] == 2.07e-06

assert mock.call_args_list[2].args[1] == -5.87e-07

assert mock.call_args_list[2].args[3] == 0.0

RATplot.plot_ref_sld(input_project, reflectivity_calculation_results)

mock.assert_called_once()

data = mock.call_args[0][0]

figure = mock.call_args[0][1]

assert figure.axes[0].get_subplotspec().get_gridspec().get_geometry() == (1, 2)

assert len(figure.axes) == 2

for reflectivity, reflectivity_results in zip(

data.reflectivity, reflectivity_calculation_results.reflectivity, strict=False

):

assert (reflectivity == reflectivity_results).all()

for sldProfile, result_sld_profile in zip(

data.sldProfiles, reflectivity_calculation_results.sldProfiles, strict=False

):

for sld, sld_results in zip(sldProfile, result_sld_profile, strict=False):

assert (sld == sld_results).all()

for resampledLayers, reflectivity_results in zip(

data.resampledLayers, reflectivity_calculation_results.resampledLayers, strict=False

):

for resam, resam_results in zip(resampledLayers, reflectivity_results, strict=False):

assert (resam == resam_results).all()

assert data.modelType == input_project.model

assert data.shiftedData == reflectivity_calculation_results.shiftedData

assert data.dataPresent.size == 0

assert (data.subRoughs == reflectivity_calculation_results.contrastParams.subRoughs).all()

assert data.resample.size == 0

"""Test that if an incorrect number of subplots is corrected in the figure helper."""

fig = plt.subplots(1, 3)[0]

RATplot.plot_ref_sld_helper(data=data(), fig=fig)

assert fig.axes[0].get_subplotspec().get_gridspec().get_geometry() == (1, 2)

"""Test that plot_ref_sld correctly throws errors for bad Bayesian input."""

RATplot.plot_ref_sld(input_project, dream_results)

RATplot.plot_ref_sld(input_project, dream_results, bayes=65)

RATplot.plot_ref_sld(input_project, dream_results, bayes=95)

with pytest.raises(ValueError):

RATplot.plot_ref_sld(input_project, reflectivity_calculation_results, bayes=65)

with pytest.raises(ValueError):

"""Test that the `assert_bayesian` decorator validates correctly."""

@RATplot.assert_bayesian("test")

def test_plot(results):

pass

test_plot(dream_results)

with pytest.raises(

ValueError, match=r"test plots are only available for the results of Bayesian analysis \(NS or DREAM\)"

):

"""Test that the panel plot helper creates a panel with the expected subplots."""

def plot_func(axes, k):

"""Plot k lines on an Axes."""

for i in range(0, k):

axes.plot([i], [i])

nplots = len(indices)

fig = RATplot.panel_plot_helper(plot_func, indices)

# ensure correct number of axes were created

expected_num_axes = ceil(sqrt(nplots)) * round(sqrt(nplots))

assert len(fig.axes) == expected_num_axes

# assert all required axes are visible and have the requested number of lines

for i, index in enumerate(indices):

assert len(fig.axes[i].get_lines()) == index

assert fig.axes[i].get_visible()

# assert remaining axes are not visible

for i in range(nplots, expected_num_axes):

assert fig.axes[i].get_visible() is False

"""Tests the formatting of the histogram plot."""

fig: plt.Figure = RATplot.plot_one_hist(

dream_results, param, estimated_density=est_dens, **hist_settings, return_fig=True

)

ax = fig.axes[0]

components = ax.get_children()

# assert expected number of bins including default

# this ensures default hist_settings are overwritten correctly

# +1 rectangle because the bounds of the plot is a rectangle

expected_bins = hist_settings.get("bins", 25) + 1

assert len([c for c in components if isinstance(c, Rectangle)]) == expected_bins

# assert line is only drawn if estimated density given

assert len(ax.get_lines()) == (0 if est_dens is None else 1)

# assert title is as expected

# also tests string to index conversion

assert ax.get_title(loc="left") == dream_results.fitNames[param] if isinstance(param, int) else param

# assert range is default, unless given

# this tests non-default hist_settings propagates correctly

try:

expected_range = hist_settings["range"]

except KeyError: # if no range given, compute the automatic range

param_index = dream_results.fitNames.index(param) if isinstance(param, str) else param

param_chain = dream_results.chain[:, param_index]

expected_range = (param_chain.min(), param_chain.max())

assert ax.get_xbound() == expected_range

"""Test the formatting of the contour plot."""

fig: plt.Figure = RATplot.plot_contour(dream_results, x_param, y_param, return_fig=True, **hist2d_settings)

ax = fig.axes[0]

components = ax.get_children()

# assert expected number of bins including default

# this ensures default hist2d_settings are overwritten correctly

# +1 as we are counting edges in this case

expected_bins = hist2d_settings.get("bins", 25) + 1

quad_mesh = [c for c in components if isinstance(c, QuadMesh)][0]

assert quad_mesh._coordinates.shape == (expected_bins, expected_bins, 2)

# assert correct axis labels

# this ensures string to index conversion works

assert ax.get_xlabel() == (dream_results.fitNames[x_param] if isinstance(x_param, int) else x_param)

assert ax.get_ylabel() == (dream_results.fitNames[y_param] if isinstance(y_param, int) else y_param)

# assert range is default, unless given

# this tests non-default hist2d_settings propagates correctly

try:

x_expected_range, y_expected_range = hist2d_settings["range"]

except KeyError: # if no range given, compute the automatic range

x_param_index = dream_results.fitNames.index(x_param) if isinstance(x_param, str) else x_param

y_param_index = dream_results.fitNames.index(y_param) if isinstance(y_param, str) else y_param

x_param_chain = dream_results.chain[:, x_param_index]

y_param_chain = dream_results.chain[:, y_param_index]

x_expected_range = (x_param_chain.min(), x_param_chain.max())

y_expected_range = (y_param_chain.min(), y_param_chain.max())

assert ax.get_xbound() == x_expected_range

assert ax.get_ybound() == y_expected_range

"""Test that corner plots are formatted correctly."""

# no use testing hist_settings and hist2d_settings here as they're tested above

fig: plt.Figure = RATplot.plot_corner(dream_results, params, return_fig=True)

axes = fig.axes

if params is None:

params = range(0, len(dream_results.fitNames))

assert len(axes) == len(params) ** 2

# annoyingly, fig.axes doesn't preserve the grid shape from plt.subplots... reconstruct grid

axes = np.array([axes[i : i + len(params)] for i in range(0, len(axes), len(params))])

for i in range(0, len(params)):

for j in range(0, len(params)):

current_axes = axes[i][j]

# ensure upper triangle is invisible

if i < j:

assert current_axes.get_visible() is False

elif i > j:

# check axes are the same along each row and column for contours

assert current_axes.get_ybound() == axes[i][0].get_ybound()

assert current_axes.get_xbound() == axes[-1][j].get_xbound()

elif i == j:

# check title is correct

assert (

current_axes.get_title(loc="left") == dream_results.fitNames[params[i]]

if isinstance(params[i], int)

else params[i]

)

"""Test chain panel name-to-index (panel helper has already been tested)"""

fig = RATplot.plot_hists(dream_results, params, return_fig=True)

plt.close(fig)

if params is None:

params = range(0, len(dream_results.fitNames))

passed_params = mock_panel_helper.call_args.args[1]

assert len(passed_params) == len(params)

for param in passed_params:

"""Test estimated density is correctly standardised."""

_ = RATplot.plot_hists(dream_results, estimated_density=input, return_fig=True)

expected_dict = {

"NONEDICT": {i: None for i in range(0, len(dream_results.fitNames))},

"D2O_DICT": {**{i: None for i in range(0, len(dream_results.fitNames))}, **{16: "kernel"}},

"DEFAULTDICT": {i: "lognor" for i in range(0, len(dream_results.fitNames))},

"DEFAULT_WITH_1CHANGE_DICT": {**{i: "normal" for i in range(0, len(dream_results.fitNames))}, **{1: "kernel"}},

}[expected_dict]

call_args = mock_plot_hist.call_args_list

keys_called = [call[0][1] for call in call_args]

est_density = [call[1]["estimated_density"] for call in call_args]

est_dens_dict = {keys_called[i]: est_density[i] for i in range(0, len(keys_called))}

assert expected_dict == est_dens_dict

"""Ensure a bad estimated density input raises an error."""

# the error message contains the phrase "Parameter {key}" or "Index {key}", so use that

# to ensure we're not getting some random ValueError

with pytest.raises((ValueError, IndexError), match=f"Parameter|Index {(list(input.keys())[0])}"):

"""Test chain panel name-to-index (panel helper has already been tested)"""

# return fig just to avoid plt.show() being called

fig = RATplot.plot_chain(dream_results, params, return_fig=True)

plt.close(fig)

if params is None:

params = range(0, len(dream_results.fitNames))

for param in mock_panel_helper.call_args()[1]:

mock_corner: MagicMock, mock_hists: MagicMock, mock_ref_sld: MagicMock, input_project, dream_results

):

"""Test that the Bayes plot calls the required plotting subroutines."""

RATplot.plot_bayes(input_project, dream_results)

assert mock_ref_sld.call_count == 2

mock_hists.assert_called_once()

"""Ensure that plot_bayes fails if given regular Results."""

with pytest.raises(

ValueError, match=r"Bayes plots are only available for the results of Bayesian analysis \(NS or DREAM\)"

):

plot_data = RATplot._extract_plot_data(data, False, True, 50)

assert len(plot_data["ref"]) == len(data.reflectivity)

assert len(plot_data["sld"]) == len(data.shiftedData)

with pytest.raises(ValueError, match=r"Parameter `shift_value` must be between 0 and 100"):

RATplot._extract_plot_data(data, False, True, -0.1)

with pytest.raises(ValueError, match=r"Parameter `shift_value` must be between 0 and 100"):

"""Test the moving_average function."""

data_to_average = np.arange(0, 20)

mov_avg = RATplot.moving_average(data_to_average)

assert mov_avg == [

1.5,

2.0,

2.5,

3.0,

3.5,

4.5,

5.5,

6.5,

7.5,

8.5,

9.5,

10.5,

11.5,

12.5,

13.5,

14.5,

15.5,

16.0,

16.5,

17.0,

]

mov_avg = RATplot.moving_average(data_to_average, window_size=2)

assert mov_avg == [

0.0,

0.5,

1.5,

2.5,

3.5,

4.5,

5.5,

6.5,

7.5,

8.5,

9.5,

10.5,

11.5,

12.5,

13.5,

14.5,

15.5,

16.5,

17.5,

18.5,

]

with pytest.raises(ValueError):

RATplot.moving_average(data_to_average, window_size=-1)

with pytest.raises(ValueError):