Handle nans in data

docs/docs/tutorials/analysis.ipynb

@@ -16,7 +16,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bca91d3c",
+   "id": "f2f898dc",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -49,7 +49,7 @@
     "# Load the vanadium data\n",
     "vanadium_experiment = Experiment('Vanadium')\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
     "    known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n",
     ")\n",
     "\n",
@@ -173,7 +173,8 @@
     "diffusion_experiment = Experiment('Diffusion')\n",
     "\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/diffusion_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/diffusion_data_example.h5',\n",
+    "    known_hash='5fe846b19aacbda8b8b936eb2e5310d025dc56c25b0b353521e7d6b921f229ab',\n",
     ")\n",
     "\n",
     "diffusion_experiment.load_hdf5(filename=file_path)"
@@ -188,6 +189,7 @@
    "source": [
     "# Now we set up the model, similarly to how we set up the model for the\n",
     "# vanadium data.\n",
+    "\n",
     "delta_function = DeltaFunction(display_name='DeltaFunction', area=0.2)\n",
     "lorentzian = Lorentzian(display_name='Lorentzian', area=0.5, width=0.3)\n",
     "component_collection = ComponentCollection(\n",
@@ -342,9 +344,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python (Pixi)",
+   "display_name": "easydynamics_newbase",
    "language": "python",
-   "name": "pixi-kernel-python3"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -356,7 +358,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.13"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,

docs/docs/tutorials/analysis1d.ipynb

@@ -41,11 +41,10 @@
     "vanadium_experiment = Experiment('Vanadium')\n",
     "\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
     "    known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n",
     ")\n",
     "\n",
-    "\n",
     "vanadium_experiment.load_hdf5(filename=file_path)"
    ]
   },

docs/docs/tutorials/experiment.ipynb

@@ -35,7 +35,7 @@
     "vanadium_experiment = Experiment('Vanadium')\n",
     "\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
     "    known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n",
     ")\n",
     "\n",

docs/docs/tutorials/tutorial1_brownian.ipynb

@@ -61,10 +61,11 @@
     "vanadium_experiment = Experiment('Vanadium')\n",
     "\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/vanadium_data_example.h5',\n",
     "    known_hash='16cc1b327c303feeb88fb9dda5390dc4880b62396b1793f98c6fef0b27c7b873',\n",
     ")\n",
     "\n",
+    "\n",
     "vanadium_experiment.load_hdf5(filename=file_path)"
    ]
   },
@@ -315,7 +316,8 @@
     "diffusion_experiment = Experiment('Diffusion')\n",
     "\n",
     "file_path = pooch.retrieve(\n",
-    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/docs/docs/docs/tutorials/data/diffusion_data_example.h5',\n",
+    "    url='https://github.com/easyscience/dynamics-lib/raw/refs/heads/master/docs/docs/tutorials/data/diffusion_data_example.h5',\n",
+    "    known_hash='5fe846b19aacbda8b8b936eb2e5310d025dc56c25b0b353521e7d6b921f229ab',\n",
     ")\n",
     "\n",
     "diffusion_experiment.load_hdf5(filename=file_path)"

src/easydynamics/analysis/analysis.py

@@ -288,6 +288,7 @@ def plot_data_and_model(
             'marker': {'Data': 'o', 'Model': None},
             'color': {'Data': 'black', 'Model': 'red'},
             'markerfacecolor': {'Data': 'none', 'Model': 'none'},
+            'keep': 'energy',
         }
         data_and_model = {
             'Data': self.experiment.binned_data,
@@ -506,7 +507,7 @@ def _fit_all_Q_simultaneously(self) -> FitResults:
         ws = []
 
         for analysis in self.analysis_list:
-            x, y, weight = self._extract_x_y_weights_from_experiment(analysis.Q_index)
+            x, y, weight = self.experiment._extract_x_y_weights_only_finite(analysis.Q_index)
             xs.append(x)
             ys.append(y)
             ws.append(weight)

src/easydynamics/analysis/analysis1d.py

@@ -195,7 +195,9 @@ def fit(self) -> FitResults:
             fit_function=self.as_fit_function(),
         )
 
-        x, y, weights = self._extract_x_y_weights_from_experiment(Q_index=self._require_Q_index())
+        x, y, weights = self.experiment._extract_x_y_weights_only_finite(
+            Q_index=self._require_Q_index()
+        )
         fit_result = fitter.fit(x=x, y=y, weights=weights)
 
         self._fit_result = fit_result

src/easydynamics/analysis/analysis_base.py

@@ -1,7 +1,6 @@
 # SPDX-FileCopyrightText: 2026 EasyScience contributors <https://github.com/easyscience>
 # SPDX-License-Identifier: BSD-3-Clause
 
-import numpy as np
 import scipp as sc
 from easyscience.base_classes.model_base import ModelBase as EasyScienceModelBase
 from easyscience.variable import Parameter
@@ -348,29 +347,6 @@ def _verify_Q_index(self, Q_index: int | None) -> int | None:
                 raise IndexError('Q_index must be a valid index for the Q values.')
         return Q_index
 
-    def _extract_x_y_weights_from_experiment(
-        self, Q_index: int
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        """Extract the x, y, and weights arrays from the experiment for
-        the given Q index.
-
-        Args:
-            Q_index (int): The Q index to extract the data for.
-
-        Returns:
-            tuple[np.ndarray, np.ndarray, np.ndarray]: The x, y, and
-                weights arrays extracted from the experiment for the
-                given Q index.
-        """
-        data = self.experiment.data['Q', Q_index]
-        x = data.coords['energy'].values
-        y = data.values
-        e = data.variances**0.5
-        if np.any(e == 0):
-            raise ValueError('Cannot compute weights: some variances are zero.')
-        weights = 1.0 / e
-        return x, y, weights
-
     #############
     # Dunder methods
     #############

src/easydynamics/experiment/experiment.py

@@ -3,6 +3,7 @@
 
 import os
 
+import numpy as np
 import plopp as pp
 import scipp as sc
 from easyscience.base_classes.new_base import NewBase
@@ -380,6 +381,55 @@ def _convert_to_bin_centers(self, data: sc.DataArray) -> sc.DataArray:
                 data = data.assign_coords({dim: sc.midpoints(coord)})
         return data
 
+    def _extract_x_y_e(self, Q_index: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Extract the x, y, and weights arrays from the experiment for
+        the given Q index.
+
+        Args:
+            Q_index (int): The Q index to extract the data for.
+
+        Returns:
+            tuple[np.ndarray, np.ndarray, np.ndarray]: The x, y, and
+                weights arrays extracted from the experiment for the
+                given Q index.
+        """
+        data = self.data['Q', Q_index]
+        x = data.coords['energy'].values
+        y = data.values
+        e = data.variances**0.5
+        return x, y, e
+
+    def _extract_x_y_weights_only_finite(
+        self, Q_index: int
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Extract the x, y, and weights arrays from the experiment for
+        the given Q index, removing any NaN and Inf values.
+
+        Args:
+            Q_index (int): The Q index to extract the data for.
+
+        Returns:
+            tuple[np.ndarray, np.ndarray, np.ndarray]: The x, y, and
+                weights arrays extracted from the experiment for the
+                given Q index, with NaNs and Infs removed.
+
+        Raises:
+            ValueError: If any variances are zero after removing NaNs
+                and Infs, since this would lead to infinite weights.
+        """
+        x, y, e = self._extract_x_y_e(Q_index)
+        mask = np.isfinite(y) & np.isfinite(e) & np.isfinite(x)
+
+        x = x[mask]
+        y = y[mask]
+        e = e[mask]
+
+        if np.any(e == 0):
+            raise ValueError('Cannot compute weights: some variances are zero.')
+        weights = 1.0 / e
+
+        return x, y, weights
+
     ########
     # dunder methods
     ###########

tests/unit/easydynamics/analysis/test_analysis1d.py

@@ -151,7 +151,7 @@ def test_fit_calls_fitter_with_correct_arguments(self, analysis1d):
         fake_y = np.array([10, 20, 30])
         fake_weights = np.array([0.1, 0.2, 0.3])
 
-        analysis1d._extract_x_y_weights_from_experiment = MagicMock(
+        analysis1d.experiment._extract_x_y_weights_only_finite = MagicMock(
             return_value=(fake_x, fake_y, fake_weights)
         )
 
@@ -181,7 +181,7 @@ def test_fit_calls_fitter_with_correct_arguments(self, analysis1d):
             fit_function='fit_func',
         )
 
-        analysis1d._extract_x_y_weights_from_experiment.assert_called_once()
+        analysis1d.experiment._extract_x_y_weights_only_finite.assert_called_once()
 
         fake_fitter_instance.fit.assert_called_once_with(
             x=fake_x,

tests/unit/easydynamics/analysis/test_analysis_base.py

@@ -6,7 +6,6 @@
 
 import numpy as np
 import pytest
-import scipp as sc
 from easyscience.variable import Parameter
 
 from easydynamics.analysis.analysis_base import AnalysisBase
@@ -331,34 +330,6 @@ def test_verify_Q_index_invalid(self, analysis_base):
         with pytest.raises(IndexError, match='Q_index must be a valid index'):
             analysis_base._verify_Q_index(invalid_Q_index)
 
-    def test_extract_x_y_weights_from_experiment(self, analysis_base):
-        # WHEN
-        Q = sc.array(dims=['Q'], values=[1, 2, 3], unit='1/Angstrom')
-        energy = sc.array(dims=['energy'], values=[10.0, 20.0, 30.0], unit='meV')
-        data = sc.array(
-            dims=['Q', 'energy'],
-            values=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
-            variances=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]],
-        )
-
-        data_array = sc.DataArray(data=data, coords={'Q': Q, 'energy': energy})
-
-        experiment = Experiment(data=data_array)
-        analysis_base.experiment = experiment
-
-        Q_index = 0
-
-        # THEN
-        x, y, weights = analysis_base._extract_x_y_weights_from_experiment(Q_index=Q_index)
-
-        # EXPECT
-        assert np.array_equal(x, analysis_base.experiment.energy.values)
-        assert np.array_equal(y, analysis_base.experiment.data.values[Q_index])
-        assert np.array_equal(
-            weights,
-            1 / analysis_base.experiment.data.variances[Q_index] ** 0.5,
-        )
-
     def test_repr(self, analysis_base):
         # WHEN
         repr_str = repr(analysis_base)

tests/unit/easydynamics/experiment/test_experiment.py

@@ -23,6 +23,23 @@ def experiment(self):
         experiment = Experiment(display_name='test_experiment', data=data)
         return experiment
 
+    @pytest.fixture
+    def experiment_with_data(self):
+        "Fixture that provides an Experiment with data for testing methods that require data"
+        Q = sc.array(dims=['Q'], values=[1, 2, 3], unit='1/Angstrom')
+        energy = sc.array(dims=['energy'], values=[10.0, 20.0, 30.0], unit='meV')
+        data = sc.array(
+            dims=['Q', 'energy'],
+            values=[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]],
+            variances=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]],
+        )
+
+        data_array = sc.DataArray(data=data, coords={'Q': Q, 'energy': energy})
+
+        experiment = Experiment(data=data_array)
+
+        return experiment
+
     ##############
     # test init
     ##############
@@ -379,6 +396,55 @@ def test_convert_to_bin_centers(self, experiment):
         assert sc.identical(converted_data.coords['energy'], expected_energy)
         assert sc.identical(converted_data.data, binned_data.data)
 
+    def test_extract_x_y_e(self, experiment_with_data):
+        # WHEN
+
+        Q_index = 0
+
+        # THEN
+        x, y, e = experiment_with_data._extract_x_y_e(Q_index=Q_index)
+
+        # EXPECT
+        assert np.array_equal(x, experiment_with_data.energy.values)
+        assert np.array_equal(y, experiment_with_data.data.values[Q_index])
+        assert np.array_equal(
+            e,
+            experiment_with_data.data.variances[Q_index] ** 0.5,
+        )
+
+    def test_extract_x_y_weights_only_finite_zero_variances(self, experiment_with_data):
+        "Test that _extract_x_y_weights_only_finite raises ValueError when variances contain zeros"
+        # WHEN
+        Q_index = 0
+        invalid_data = experiment_with_data._data.copy()
+        invalid_data.data.variances[Q_index] = 0  # Set variances to zero
+        # throw in a nan for good measure
+        invalid_data.data.variances[Q_index][0] = np.nan
+
+        # THEN EXPECT
+        with pytest.raises(ValueError, match='Cannot compute weights: some variances are zero'):
+            Experiment(data=invalid_data)._extract_x_y_weights_only_finite(Q_index=Q_index)
+
+    def test_extract_x_y_weights_only_finite(self, experiment_with_data):
+        "Test that _extract_x_y_weights_only_finite only returns finite values"
+        # WHEN
+        Q_index = 0
+        invalid_data = experiment_with_data._data.copy()
+        invalid_data.data.values[Q_index][0] = np.inf
+        invalid_data.data.variances[Q_index][1] = np.nan
+
+        # THEN
+        x, y, weights = Experiment(data=invalid_data)._extract_x_y_weights_only_finite(
+            Q_index=Q_index
+        )
+
+        # EXPECT
+        assert np.isfinite(x).all()
+        assert np.isfinite(y).all()
+        assert np.isfinite(weights).all()
+        assert weights[0] == 1.0 / (experiment_with_data.data.variances[Q_index][2] ** 0.5)
+        assert len(x) == len(y) == len(weights) == 1  # 2 values should be removed
+
     ##############
     # test dunder methods
     ##############