Make it easy to fix and free energy offset

Author: henrikjacobsenfys

src/easydynamics/analysis/analysis.py

@@ -422,6 +422,38 @@ def plot_parameters(
         )
         return fig
 
+    def fix_energy_offset(self, Q_index: int | None = None) -> None:
+        """Fix the energy offset parameter(s) for a specific Q index, or
+        for all Q indices if Q_index is None.
+
+        Args:
+            Q_index (int | None, default=None): Index of the Q value to
+                fix the energy offset for. If None, fixes the energy
+                offset for all Q values. Default is None.
+        """
+        if Q_index is not None:
+            Q_index = self._verify_Q_index(Q_index)
+            self.analysis_list[Q_index].fix_energy_offset()
+        else:
+            for analysis in self.analysis_list:
+                analysis.fix_energy_offset()
+
+    def free_energy_offset(self, Q_index: int | None = None) -> None:
+        """Free the energy offset parameter(s) for a specific Q index,
+        or for all Q indices if Q_index is None.
+
+        Args:
+            Q_index (int | None, default=None): Index of the Q value to
+                free the energy offset for. If None, frees the energy
+                offset for all Q values. Default is None.
+        """
+        if Q_index is not None:
+            Q_index = self._verify_Q_index(Q_index)
+            self.analysis_list[Q_index].free_energy_offset()
+        else:
+            for analysis in self.analysis_list:
+                analysis.free_energy_offset()
+
     #############
     # Private methods - updating models when things change
     #############
@@ -430,7 +462,7 @@ def _on_experiment_changed(self) -> None:
         """Update the Q values in the sample and instrument models when
         the experiment changes.
 
-        Also update all the Analysi1d objects with the new experiment.
+        Also update all the Analysis1d objects with the new experiment.
         """
         if self._call_updaters:
             super()._on_experiment_changed()
@@ -441,7 +473,8 @@ def _on_sample_model_changed(self) -> None:
         """Update the Q values in the sample model when the sample model
         changes.
 
-        Also update all the Analysi1d objects with the new sample model.
+        Also update all the Analysis1d objects with the new sample
+        model.
         """
         if self._call_updaters:
             super()._on_sample_model_changed()
@@ -452,7 +485,7 @@ def _on_instrument_model_changed(self) -> None:
         """Update the Q values in the instrument model when the
         instrument model changes.
 
-        Also update all the Analysi1d objects with the new instrument
+        Also update all the Analysis1d objects with the new instrument
         model.
         """
         if self._call_updaters:

src/easydynamics/analysis/analysis1d.py

@@ -305,6 +305,14 @@ def plot_data_and_model(
         )
         return fig
 
+    def fix_energy_offset(self) -> None:
+        """Fix the energy offset parameter for the current Q index."""
+        self.instrument_model.fix_energy_offset(Q_index=self._require_Q_index())
+
+    def free_energy_offset(self) -> None:
+        """Free the energy offset parameter for the current Q index."""
+        self.instrument_model.free_energy_offset(Q_index=self._require_Q_index())
+
     #############
     # Private methods: small utilities
     #############
@@ -428,7 +436,7 @@ def _evaluate_components(
         if energy is None:
             energy = self._masked_energy
 
-        energy_offset = self.instrument_model.get_energy_offset_at_Q(Q_index)
+        energy_offset = self.instrument_model.get_energy_offset(Q_index)
         energy_with_offset = self._calculate_energy_with_offset(
             energy=energy,
             energy_offset=energy_offset,
@@ -601,7 +609,7 @@ def _create_convolver(
             resolution_components=resolution_components,
             energy=energy,
             temperature=self.temperature,
-            energy_offset=self.instrument_model.get_energy_offset_at_Q(Q_index),
+            energy_offset=self.instrument_model.get_energy_offset(Q_index),
         )
         return convolver
 

src/easydynamics/sample_model/instrument_model.py

@@ -351,32 +351,96 @@ def free_resolution_parameters(self) -> None:
         """Free all parameters in the resolution model."""
         self.resolution_model.free_all_parameters()
 
-    def get_energy_offset_at_Q(self, Q_index: int) -> Parameter:
+    def get_energy_offset(
+        self,
+        Q_index: int | None = None,
+    ) -> Parameter | list[Parameter]:
         """Get the energy offset Parameter at a specific Q index.
 
         Args:
-            Q_index (int): The index of the Q value to get the energy
-                offset for.
+            Q_index (int | None, default=None): The index of the Q value to get the energy
+                offset for. If None, get the energy offset for all Q values.
 
         Returns:
-            Parameter: The energy offset Parameter at the specified Q
-                index.
+            Parameter | list[Parameter]: The energy offset Parameter at the specified Q
+                index, or a list of Parameters if Q_index is None.
 
         Raises:
             ValueError: If no Q values are set in the InstrumentModel.
             IndexError: If Q_index is out of bounds.
+            TypeError: If Q_index is not an int or None.
         """
         if self._Q is None:
             raise ValueError('No Q values are set in the InstrumentModel.')
 
+        if Q_index is None:
+            return self._energy_offsets
+
+        if not isinstance(Q_index, int):
+            raise TypeError(f'Q_index must be an int or None, got {type(Q_index).__name__}')
+
         if Q_index < 0 or Q_index >= len(self._Q):
             raise IndexError(f'Q_index {Q_index} is out of bounds for Q of length {len(self._Q)}')
 
         return self._energy_offsets[Q_index]
 
+    def fix_energy_offset(self, Q_index: int | None = None) -> None:
+        """Fix energy offset parameters. If Q_index is specified, only
+        fix the energy offset for that Q value. If Q_index is None, fix
+        energy offsets for all Q values.
+
+        Args:
+            Q_index (int | None, default=None): The index of the Q value
+                to fix the energy offset for. If None, fix energy
+                offsets for all Q values.
+        """
+        self._fix_or_free_energy_offset(Q_index, fixed=True)
+
+    def free_energy_offset(self, Q_index: int | None = None) -> None:
+        """Free energy offset parameters. If Q_index is specified, only
+        free the energy offset for that Q value. If Q_index is None,
+        free energy offsets for all Q values.
+
+        Args:
+            Q_index (int | None, default=None): The index of the Q value
+                to free the energy offset for. If None, free energy
+                offsets for all Q values.
+        """
+        self._fix_or_free_energy_offset(Q_index, fixed=False)
+
     # --------------------------------------------------------------
     # Private methods
     # --------------------------------------------------------------
+    def _fix_or_free_energy_offset(self, Q_index: int | None = None, fixed: bool = True) -> None:
+        """Fix or free energy offset parameters. If Q_index is
+        specified, only fix or free the energy offset for that Q value.
+        If Q_index is None, fix or free energy offsets for all Q values.
+
+        Args:
+            Q_index (int | None, default=None): The index of the Q value
+                to fix or free the energy offset for. If None, fix or
+                free energy offsets for all Q values.
+            fixed (bool, default=True): Whether to fix (True) or free
+                (False) the energy offset.
+
+        Raises:
+            TypeError: If Q_index is not an int or None.
+            IndexError: If Q_index is out of bounds for the Q values in
+                the InstrumentModel.
+        """
+
+        if Q_index is None:
+            for offset in self._energy_offsets:
+                offset.fixed = fixed
+        else:
+            if not isinstance(Q_index, int):
+                raise TypeError(f'Q_index must be an int or None, got {type(Q_index).__name__}')
+
+            if Q_index < 0 or Q_index >= len(self._Q):
+                raise IndexError(
+                    f'Q_index {Q_index} is out of bounds for Q of length {len(self._Q)}'
+                )
+            self._energy_offsets[Q_index].fixed = fixed
 
     def _generate_energy_offsets(self) -> None:
         """Generate energy offset Parameters for each Q value."""

tests/unit/easydynamics/analysis/test_analysis.py

@@ -476,6 +476,43 @@ def test_plot_parameters(self, analysis):
         # and that we return the figure
         assert result is fake_fig
 
+    def test_fix_and_free_energy_offset(self, analysis):
+        # EXPECT
+        offsets = analysis.instrument_model.get_energy_offset()
+        for offset in offsets:
+            assert offset.fixed is False
+
+        # THEN
+        analysis.fix_energy_offset()
+
+        # EXPECT
+        for offset in offsets:
+            assert offset.fixed is True
+
+        # THEN
+        analysis.free_energy_offset()
+
+        # EXPECT
+        for offset in offsets:
+            assert offset.fixed is False
+
+        # THEN
+        analysis.fix_energy_offset(Q_index=1)
+
+        # EXPECT
+        for i, offset in enumerate(offsets):
+            if i == 1:
+                assert offset.fixed is True
+            else:
+                assert offset.fixed is False
+
+        # THEN
+        analysis.free_energy_offset(Q_index=1)
+
+        # EXPECT
+        for offset in offsets:
+            assert offset.fixed is False
+
     def test_on_experiment_changed_similar_Q(self, analysis):
         # WHEN
         # Create a new experiment.

tests/unit/easydynamics/analysis/test_analysis1d.py

@@ -274,6 +274,24 @@ def test_plot_calls_plopp_with_correct_arguments(self, analysis1d):
 
         assert result is fake_fig
 
+    def test_fix_and_free_offset(self, analysis1d):
+        # WHEN
+
+        # EXPECT
+        assert analysis1d.instrument_model.get_energy_offset(Q_index=0).fixed is False
+
+        # THEN
+        analysis1d.fix_energy_offset()
+
+        # EXPECT
+        assert analysis1d.instrument_model.get_energy_offset(Q_index=0).fixed is True
+
+        # THEN
+        analysis1d.free_energy_offset()
+
+        # EXPECT
+        assert analysis1d.instrument_model.get_energy_offset(Q_index=0).fixed is False
+
     #############
     # Private methods: small utilities
     #############
@@ -334,7 +352,7 @@ def test_verify_energy_raises(self, analysis1d):
     def test_calculate_energy_with_offset(self, analysis1d):
         # WHEN
         energy = analysis1d.experiment.energy
-        energy_offset = analysis1d.instrument_model.get_energy_offset_at_Q(analysis1d.Q_index)
+        energy_offset = analysis1d.instrument_model.get_energy_offset(Q_index=analysis1d.Q_index)
         energy_offset.value = 1.0  # override with a simple value for testing
 
         # THEN
@@ -347,7 +365,7 @@ def test_calculate_energy_with_offset(self, analysis1d):
     def test_calculate_energy_with_offset_different_units(self, analysis1d):
         # WHEN
         energy = analysis1d.experiment.energy
-        energy_offset = analysis1d.instrument_model.get_energy_offset_at_Q(analysis1d.Q_index)
+        energy_offset = analysis1d.instrument_model.get_energy_offset(Q_index=analysis1d.Q_index)
         energy_offset.value = 1.0  # override with a simple value for testing
         energy_offset.convert_unit('eV')
 
@@ -452,7 +470,7 @@ def test_evaluate_with_resolution(self, analysis1d):
                 )
             )
 
-            energy_offset = analysis1d.instrument_model.get_energy_offset_at_Q(analysis1d.Q_index)
+            energy_offset = analysis1d.instrument_model.get_energy_offset(analysis1d.Q_index)
 
             # Extract call arguments
             _, kwargs = MockConvolution.call_args
@@ -576,7 +594,7 @@ def test_create_convolver(self, analysis1d):
             return_value=resolution_components
         )
 
-        analysis1d.instrument_model.get_energy_offset_at_Q = MagicMock(return_value=123.0)
+        analysis1d.instrument_model.get_energy_offset = MagicMock(return_value=123.0)
 
         with patch('easydynamics.analysis.analysis1d.Convolution') as MockConvolution:
             # THEN