Allow plotting parameters against Q^2

docs/docs/tutorials/tutorial1_brownian.ipynb

@@ -444,7 +444,25 @@
    "id": "eb226c8f",
    "metadata": {},
    "source": [
-    "There is a clear trend: the area is more or less constant, while the width seems to increase with `Q^2`. We therefore try and fit a Brownian translational diffusion model to the data. In this model, the scattering is given by\n",
+    "There is a clear trend: the area is more or less constant, while the width seems to increase with `Q^2`. Let us plot as function of Q^2 to check:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "015815bd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "diffusion_analysis.plot_parameters(names='Lorentzian width', x_axis='Q^2')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "443d3504",
+   "metadata": {},
+   "source": [
+    "We therefore try and fit a Brownian translational diffusion model to the data. In this model, the scattering is given by\n",
     "$$\n",
     "I(Q,E) = S \\frac{\\Gamma(Q)}{\\Gamma(Q)^2 + E^2},\n",
     "$$\n",

src/easydynamics/analysis/analysis.py

@@ -23,6 +23,7 @@
 from easydynamics.sample_model.instrument_model import InstrumentModel
 from easydynamics.settings.convolution_settings import ConvolutionSettings
 from easydynamics.settings.detailed_balance_settings import DetailedBalanceSettings
+from easydynamics.utils.plotting import add_x_axis_coord
 from easydynamics.utils.plotting import slicerplot_with_residuals
 from easydynamics.utils.utils import _in_notebook
 from easydynamics.utils.utils import verify_Q_index
@@ -646,33 +647,38 @@ def parameters_to_dataset(self) -> sc.Dataset:
     def plot_parameters(
         self,
         names: str | list[str] | None = None,
+        x_axis: str = 'Q',
         **kwargs: dict[str, Any],
     ) -> WidgetFigure:
         """
-        Plot fitted parameters as a function of Q.
+        Plot fitted parameters as a function of Q or Q^2.
 
         Parameters
         ----------
         names : str | list[str] | None, default=None
             Name(s) of the parameter(s) to plot. If None, plots all parameters.
+        x_axis : str, default='Q'
+            Quantity to plot the parameters against: 'Q' or 'Q^2'. Plotting against 'Q^2' shows,
+            e.g., a Fickian diffusion width (D*Q^2) as a straight line.
         **kwargs : dict[str, Any]
-            Additional keyword arguments passed to plopp.slicer for customizing the plot (e.g.,
+            Additional keyword arguments passed to plopp.plot for customizing the plot (e.g.,
             title, linestyle, marker, color).
 
         Raises
         ------
         TypeError
             If names is not a string, list of strings, or None.
         ValueError
-            If any of the specified parameter names are not found in the dataset.
+            If any of the specified parameter names are not found in the dataset, or if x_axis is
+            not one of the supported axes.
 
         Returns
         -------
         WidgetFigure
             A Plopp WidgetFigure containing the plot of the parameters.
         """
 
-        ds = self.parameters_to_dataset()
+        ds = add_x_axis_coord(self.parameters_to_dataset(), x_axis)
 
         # None or an empty list both mean 'plot all parameters'.
         if not names:
@@ -690,6 +696,7 @@ def plot_parameters(
 
         data_to_plot = {name: ds[name] for name in names}
         plot_kwargs_defaults = {
+            'coords': [x_axis],
             'linestyle': dict.fromkeys(names, 'none'),
             'marker': dict.fromkeys(names, 'o'),
             'markerfacecolor': dict.fromkeys(names, 'none'),

src/easydynamics/analysis/parameter_analysis.py

@@ -19,6 +19,7 @@
 from easydynamics.analysis.posterior_sampling import PosteriorSampler
 from easydynamics.base_classes.easydynamics_modelbase import EasyDynamicsModelBase
 from easydynamics.utils.fit_target import FitTarget
+from easydynamics.utils.plotting import add_x_axis_coord
 from easydynamics.utils.utils import _in_notebook
 from easydynamics.utils.utils import convert_value_unit
 
@@ -412,14 +413,22 @@ def _chain_parameters(self) -> list[Parameter]:
                 parameters.setdefault(parameter.unique_name, parameter)
         return list(parameters.values())
 
-    def plot(self, names: str | list[str] | None = None, **kwargs: dict[str, Any]) -> WidgetFigure:
+    def plot(
+        self,
+        names: str | list[str] | None = None,
+        x_axis: str = 'Q',
+        **kwargs: dict[str, Any],
+    ) -> WidgetFigure:
         """
-        Plot the parameters and fit results.
+        Plot the parameters and fit results as a function of Q or Q^2.
 
         Parameters
         ----------
         names : str | list[str] | None, default=None
             The names of the parameters to plot. If None, all parameters with bindings are plotted.
+        x_axis : str, default='Q'
+            Quantity to plot the parameters and fit results against: 'Q' or 'Q^2'. Plotting against
+            'Q^2' shows, e.g., a Fickian diffusion width (D*Q^2) as a straight line.
         **kwargs : dict[str, Any]
             Additional keyword arguments to pass to the plotting function.
 
@@ -431,7 +440,8 @@ def plot(self, names: str | list[str] | None = None, **kwargs: dict[str, Any]) -
         Raises
         ------
         ValueError
-            If the units of the specified parameters are not consistent.
+            If the units of the specified parameters are not consistent, or if x_axis is not one of
+            the supported axes.
         RuntimeError
             If plot() is called outside of a Jupyter notebook environment.
         """
@@ -482,6 +492,7 @@ def plot(self, names: str | list[str] | None = None, **kwargs: dict[str, Any]) -
         markers = itertools.cycle(['o', 's', 'D', '^', 'v', '<', '>'])
 
         plot_kwargs = {
+            'coords': [x_axis],
             'title': self.display_name,
             'linestyle': {},
             'marker': {},
@@ -526,7 +537,7 @@ def plot(self, names: str | list[str] | None = None, **kwargs: dict[str, Any]) -
         data_and_model = sc.Dataset(data_arrays)
         data_and_model.update(model_arrays)
 
-        return pp.plot(data_and_model, **plot_kwargs)
+        return pp.plot(add_x_axis_coord(data_and_model, x_axis), **plot_kwargs)
 
     def calculate_model_dataset(self, bindings: list[FitBinding]) -> sc.Dataset:
         """

src/easydynamics/utils/plotting.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: 2025 EasyScience contributors <https://github.com/easyscience>
 # SPDX-License-Identifier: BSD-3-Clause
 
+from collections.abc import Callable
 
 import plopp as pp
 import scipp as sc
@@ -9,6 +10,41 @@
 from plopp.plotting._slicer import _maybe_reduce_dim
 from plopp.widgets import slice_dims
 
+# x-axes that Q-dependent parameters can be plotted against, each mapped to the function
+# computing it from Q.
+Q_AXIS_TRANSFORMS: dict[str, Callable[[sc.Variable], sc.Variable]] = {
+    'Q': lambda Q: Q,
+    'Q^2': lambda Q: Q**2,
+}
+
+
+def add_x_axis_coord(data: sc.Dataset, x_axis: str) -> sc.Dataset:
+    """
+    Add an x_axis coordinate computed from the Q coordinate, for plotting against it.
+
+    Only the coordinate is transformed; the values at each Q are unchanged.
+
+    Parameters
+    ----------
+    data : sc.Dataset
+        Dataset with a 'Q' coordinate.
+    x_axis : str
+        Name of the axis to add, one of the keys of ``Q_AXIS_TRANSFORMS``.
+
+    Returns
+    -------
+    sc.Dataset
+        A copy of data with the x_axis coordinate added.
+
+    Raises
+    ------
+    ValueError
+        If x_axis is not one of the supported axes.
+    """
+    if x_axis not in Q_AXIS_TRANSFORMS:
+        raise ValueError(f"x_axis must be one of {list(Q_AXIS_TRANSFORMS)}, got '{x_axis}'.")
+    return data.assign_coords({x_axis: Q_AXIS_TRANSFORMS[x_axis](data.coords['Q'])})
+
 
 def slicerplot_with_residuals(
     dg: sc.DataGroup,

tests/unit/easydynamics/analysis/test_analysis.py

@@ -655,10 +655,14 @@ def test_plot_parameters(self, analysis):
             'marker': {'amplitude': 'x', 'width': 's'},
         }
 
-        fake_dataset = {
-            'amplitude': object(),
-            'width': object(),
-        }
+        Q = sc.array(dims=['Q'], values=[0.5, 1.0], unit='1/angstrom')
+        fake_dataset = sc.Dataset(
+            data={
+                'amplitude': sc.array(dims=['Q'], values=[1.0, 2.0]),
+                'width': sc.array(dims=['Q'], values=[0.1, 0.2], unit='meV'),
+            },
+            coords={'Q': Q},
+        )
 
         analysis.parameters_to_dataset = MagicMock(return_value=fake_dataset)
 
@@ -674,9 +678,13 @@ def test_plot_parameters(self, analysis):
 
         dataset_passed = args[0]
 
-        assert dataset_passed == fake_dataset
+        assert list(dataset_passed) == list(fake_dataset)
+        for name, data_array in dataset_passed.items():
+            assert sc.identical(data_array, fake_dataset[name])
 
         # Check default kwargs
+        assert kwargs['coords'] == ['Q']
+
         assert 'linestyle' in kwargs
         assert kwargs['linestyle'] == {
             'amplitude': 'none',
@@ -696,6 +704,24 @@ def test_plot_parameters(self, analysis):
         # and that we return the figure
         assert result is fake_fig
 
+    def test_plot_parameters_against_Q_squared(self, analysis):
+        # WHEN
+        Q = sc.array(dims=['Q'], values=[0.5, 1.0], unit='1/angstrom')
+        fake_dataset = sc.Dataset(
+            data={'width': sc.array(dims=['Q'], values=[0.1, 0.2], unit='meV')},
+            coords={'Q': Q},
+        )
+        analysis.parameters_to_dataset = MagicMock(return_value=fake_dataset)
+
+        # THEN
+        with patch('plopp.plot') as mock_plot:
+            analysis.plot_parameters(x_axis='Q^2')
+
+        # EXPECT the parameters are plotted against a Q^2 coordinate
+        args, kwargs = mock_plot.call_args
+        assert kwargs['coords'] == ['Q^2']
+        assert sc.identical(args[0]['width'].coords['Q^2'], Q**2)
+
     def test_fix_and_free_energy_offset(self, analysis):
         # EXPECT
         offsets = analysis.instrument_model.get_energy_offset()

tests/unit/easydynamics/analysis/test_parameter_analysis.py

@@ -611,6 +611,30 @@ def test_plot_evaluates_only_the_bindings_being_plotted(
             parameter_analysis.bindings[0]
         ])
 
+    def test_plot_against_Q_squared(self, parameter_analysis, mock_model_dataset):
+        # WHEN
+        parameter_analysis.calculate_model_dataset = MagicMock(return_value=mock_model_dataset)
+
+        # THEN
+        with (
+            patch(
+                'easydynamics.analysis.parameter_analysis._in_notebook',
+                return_value=True,
+            ),
+            patch('easydynamics.analysis.parameter_analysis.pp.plot') as mock_plot,
+        ):
+            parameter_analysis.plot(names=['parameter1'], x_axis='Q^2')
+
+        # EXPECT data and model are plotted unchanged against a Q^2 coordinate
+        args, kwargs = mock_plot.call_args
+        plotted = args[0]
+        assert kwargs['coords'] == ['Q^2']
+        assert sc.identical(plotted.coords['Q^2'], plotted.coords['Q'] ** 2)
+        assert sc.identical(
+            plotted['parameter1'].data, parameter_analysis.parameters['parameter1'].data
+        )
+        assert sc.identical(plotted['Polynomial'].data, mock_model_dataset['Polynomial'].data)
+
     @pytest.mark.parametrize(
         'set_pars_none, bindings, expected_exception, match',
         [

tests/unit/easydynamics/utils/test_plotting.py

@@ -7,9 +7,50 @@
 import pytest
 import scipp as sc
 
+from easydynamics.utils.plotting import add_x_axis_coord
 from easydynamics.utils.plotting import slicerplot_with_residuals
 
 
+class TestAdd_x_axis_coord:
+    @pytest.fixture
+    def dataset(self):
+        return sc.Dataset(
+            data={'width': sc.array(dims=['Q'], values=[0.1, 0.2], unit='meV')},
+            coords={'Q': sc.array(dims=['Q'], values=[0.5, 2.0], unit='1/angstrom')},
+        )
+
+    @pytest.mark.parametrize(
+        'x_axis, expected_values, expected_unit',
+        [
+            ('Q', [0.5, 2.0], '1/angstrom'),
+            ('Q^2', [0.25, 4.0], '1/angstrom**2'),
+        ],
+        ids=['Q', 'Q_squared'],
+    )
+    def test_add_x_axis_coord(self, dataset, x_axis, expected_values, expected_unit):
+        # THEN
+        result = add_x_axis_coord(dataset, x_axis)
+
+        # EXPECT the coordinate is added and the data is unchanged
+        assert sc.identical(
+            result.coords[x_axis],
+            sc.array(dims=['Q'], values=expected_values, unit=expected_unit),
+        )
+        assert sc.identical(result['width'].data, dataset['width'].data)
+
+    def test_add_x_axis_coord_does_not_modify_input(self, dataset):
+        # THEN
+        add_x_axis_coord(dataset, 'Q^2')
+
+        # EXPECT
+        assert 'Q^2' not in dataset.coords
+
+    def test_add_x_axis_coord_invalid_x_axis_raises(self, dataset):
+        # WHEN THEN EXPECT
+        with pytest.raises(ValueError, match='x_axis must be one of'):
+            add_x_axis_coord(dataset, 'Q^3')
+
+
 class TestSlicerplot_with_residuals:
     @pytest.fixture
     def datagroup(self):