From 1173f48dd9643be5ffa72550fba668fa123e8b10 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Wed, 25 Mar 2026 10:37:04 +0100 Subject: [PATCH 1/2] Make it easy to fix and free energy offset --- src/easydynamics/analysis/analysis.py | 39 +++++++++- src/easydynamics/analysis/analysis1d.py | 12 ++- .../sample_model/instrument_model.py | 74 +++++++++++++++++-- .../easydynamics/analysis/test_analysis.py | 37 ++++++++++ .../easydynamics/analysis/test_analysis1d.py | 26 ++++++- .../sample_model/test_instrument_model.py | 69 +++++++++++++++-- 6 files changed, 237 insertions(+), 20 deletions(-) diff --git a/src/easydynamics/analysis/analysis.py b/src/easydynamics/analysis/analysis.py index 136cc5be..0fbf760d 100644 --- a/src/easydynamics/analysis/analysis.py +++ b/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: diff --git a/src/easydynamics/analysis/analysis1d.py b/src/easydynamics/analysis/analysis1d.py index ec6c3c8d..fea6aad6 100644 --- a/src/easydynamics/analysis/analysis1d.py +++ b/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 diff --git a/src/easydynamics/sample_model/instrument_model.py b/src/easydynamics/sample_model/instrument_model.py index c30beeb4..820fd995 100644 --- a/src/easydynamics/sample_model/instrument_model.py +++ b/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.""" diff --git a/tests/unit/easydynamics/analysis/test_analysis.py b/tests/unit/easydynamics/analysis/test_analysis.py index b62dc2b6..4cbe32fb 100644 --- a/tests/unit/easydynamics/analysis/test_analysis.py +++ b/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. diff --git a/tests/unit/easydynamics/analysis/test_analysis1d.py b/tests/unit/easydynamics/analysis/test_analysis1d.py index 40d73326..6669c4fc 100644 --- a/tests/unit/easydynamics/analysis/test_analysis1d.py +++ b/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 diff --git a/tests/unit/easydynamics/sample_model/test_instrument_model.py b/tests/unit/easydynamics/sample_model/test_instrument_model.py index 634dea8a..59596ccd 100644 --- a/tests/unit/easydynamics/sample_model/test_instrument_model.py +++ b/tests/unit/easydynamics/sample_model/test_instrument_model.py @@ -200,24 +200,24 @@ def test_energy_offset_setter_raises(self, instrument_model): ): instrument_model.energy_offset = 'invalid_offset' - def test_get_energy_offset_at_Q(self, instrument_model): + def test_get_energy_offset(self, instrument_model): # WHEN # THEN - offset_at_Q0 = instrument_model.get_energy_offset_at_Q(0) + offset_at_Q0 = instrument_model.get_energy_offset(0) # EXPECT assert offset_at_Q0.value == instrument_model.energy_offset.value - def test_get_energy_offset_at_Q_invalid_index_raises(self, instrument_model): + def test_get_energy_offset_invalid_index_raises(self, instrument_model): # WHEN / THEN / EXPECT with pytest.raises( IndexError, match='Q_index 5 is out of bounds', ): - instrument_model.get_energy_offset_at_Q(5) + instrument_model.get_energy_offset(5) - def test_get_energy_offset_at_Q_no_Q_raises(self, instrument_model): + def test_get_energy_offset_no_Q_raises(self, instrument_model): # WHEN instrument_model.clear_Q(confirm=True) @@ -226,7 +226,7 @@ def test_get_energy_offset_at_Q_no_Q_raises(self, instrument_model): ValueError, match='No Q values are set', ): - instrument_model.get_energy_offset_at_Q(0) + instrument_model.get_energy_offset(0) def test_convert_unit_calls_all_children(self, instrument_model): # WHEN @@ -414,6 +414,63 @@ def test_Q_setter(self, instrument_model_without_Q): np.testing.assert_array_equal(instrument_model_without_Q.background_model.Q, first_new_Q) np.testing.assert_array_equal(instrument_model_without_Q.resolution_model.Q, first_new_Q) + def test_fix_and_free_offset(self, instrument_model): + # WHEN + # EXPECT + for offset in instrument_model._energy_offsets: + assert offset.fixed is False + + # THEN + instrument_model.fix_energy_offset() + + # EXPECT + for offset in instrument_model._energy_offsets: + assert offset.fixed is True + # THEN + instrument_model.free_energy_offset() + + # EXPECT + for offset in instrument_model._energy_offsets: + assert offset.fixed is False + + # THEN + instrument_model.fix_energy_offset(Q_index=1) + + # EXPECT + for i, offset in enumerate(instrument_model._energy_offsets): + if i == 1: + assert offset.fixed is True + else: + assert offset.fixed is False + + def test_fix_or_free_energy_offset_invalid_Q_index_raises(self, instrument_model): + # WHEN / THEN / EXPECT + with pytest.raises( + IndexError, + match='Q_index 5 is out of bounds', + ): + instrument_model.fix_energy_offset(Q_index=5) + + with pytest.raises( + IndexError, + match='Q_index 5 is out of bounds', + ): + instrument_model.free_energy_offset(Q_index=5) + + def test_fix_or_free_energy_offset_nonint_Q_index_raises(self, instrument_model): + # WHEN / THEN / EXPECT + with pytest.raises( + TypeError, + match='Q_index must be an int or None, got str', + ): + instrument_model.fix_energy_offset(Q_index='invalid_index') + + with pytest.raises( + TypeError, + match='Q_index must be an int or None, got str', + ): + instrument_model.free_energy_offset(Q_index='invalid_index') + def test_on_energy_offset_change(self, instrument_model): # WHEN new_offset = 2.0 From 752f1a94c814398640641d4da1346976fedd8878 Mon Sep 17 00:00:00 2001 From: henrikjacobsenfys Date: Wed, 25 Mar 2026 10:47:07 +0100 Subject: [PATCH 2/2] add missing test --- .../easydynamics/sample_model/test_instrument_model.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/unit/easydynamics/sample_model/test_instrument_model.py b/tests/unit/easydynamics/sample_model/test_instrument_model.py index 59596ccd..cec5076a 100644 --- a/tests/unit/easydynamics/sample_model/test_instrument_model.py +++ b/tests/unit/easydynamics/sample_model/test_instrument_model.py @@ -217,6 +217,14 @@ def test_get_energy_offset_invalid_index_raises(self, instrument_model): ): instrument_model.get_energy_offset(5) + def test_get_energy_offset_nonint_index_raises(self, instrument_model): + # WHEN / THEN / EXPECT + with pytest.raises( + TypeError, + match='Q_index must be an int or None, got str', + ): + instrument_model.get_energy_offset('invalid_index') + def test_get_energy_offset_no_Q_raises(self, instrument_model): # WHEN instrument_model.clear_Q(confirm=True)