diff --git a/src/easydynamics/sample_model/components/polynomial.py b/src/easydynamics/sample_model/components/polynomial.py index a36700df..7c7e9ec2 100644 --- a/src/easydynamics/sample_model/components/polynomial.py +++ b/src/easydynamics/sample_model/components/polynomial.py @@ -4,7 +4,7 @@ from __future__ import annotations import warnings -from typing import TYPE_CHECKING +from collections.abc import Sequence import numpy as np import scipp as sc @@ -15,8 +15,7 @@ from easydynamics.sample_model.components.model_component import ModelComponent from easydynamics.utils.utils import Numeric -if TYPE_CHECKING: - from collections.abc import Sequence +_CoefficientsInput = Sequence[Numeric] | dict[int, Numeric] class Polynomial(ModelComponent): @@ -52,7 +51,7 @@ class Polynomial(ModelComponent): def __init__( self, - coefficients: Sequence[Numeric | Parameter] = (0.0,), + coefficients: _CoefficientsInput = (0.0,), unit: str | sc.Unit = 'meV', name: str = 'Polynomial', display_name: str | None = None, @@ -63,8 +62,10 @@ def __init__( Parameters ---------- - coefficients : Sequence[Numeric | Parameter], default=(0.0,) - Coefficients c0, c1, ..., cN. + coefficients : _CoefficientsInput, default=(0.0,) + Coefficients c0, c1, ..., cN as a sequence of numbers, or a sparse dict mapping integer + powers to numbers (e.g. ``{2: 1.5}`` for ``1.5*x^2``). Missing powers in a dict are + filled with fixed-to-zero Parameters. unit : str | sc.Unit, default='meV' Unit of the Polynomial component. name : str, default='Polynomial' @@ -78,10 +79,10 @@ def __init__( Raises ------ TypeError - If coefficients is not a sequence of numbers or Parameters or if any item in - coefficients is not a number or Parameter. + If coefficients is not a sequence of numbers or not a dict with integer keys and + numeric values. ValueError - If coefficients is an empty sequence. + If coefficients is an empty sequence or dict, or if dict keys are negative. """ super().__init__( @@ -91,27 +92,42 @@ def __init__( unique_name=unique_name, ) - if not isinstance(coefficients, (list, tuple, np.ndarray)): - raise TypeError( - 'coefficients must be a sequence (list/tuple/ndarray) \ - of numbers or Parameter objects.' - ) - - if len(coefficients) == 0: - raise ValueError('At least one coefficient must be provided.') - # Internal storage of Parameter objects self._coefficients: list[Parameter] = [] - # Coefficients are treated as dimensionless Parameters - for i, coef in enumerate(coefficients): - if isinstance(coef, Parameter): - param = coef - elif isinstance(coef, Numeric): - param = Parameter(name=f'{name}_c{i}', value=float(coef)) - else: - raise TypeError('Each coefficient must be either a numeric value or a Parameter.') - self._coefficients.append(param) + if isinstance(coefficients, dict): + if len(coefficients) == 0: + raise ValueError('At least one coefficient must be provided.') + for key in coefficients: + if not isinstance(key, int): + raise TypeError('Dict keys must be integers representing polynomial powers.') + if key < 0: + raise ValueError('Dict keys (powers) must be non-negative integers.') + degree = max(coefficients) + for i in range(degree + 1): + coef = coefficients.get(i) + if coef is None: + param = Parameter(name=f'{name}_c{i}', value=0.0, fixed=True) + elif isinstance(coef, Numeric): + param = Parameter(name=f'{name}_c{i}', value=float(coef)) + else: + raise TypeError('Each coefficient value must be a number.') + self._coefficients.append(param) + elif isinstance(coefficients, (list, tuple, np.ndarray)): + if len(coefficients) == 0: + raise ValueError('At least one coefficient must be provided.') + for i, coef in enumerate(coefficients): + if isinstance(coef, Parameter): + param = coef + elif isinstance(coef, Numeric): + param = Parameter(name=f'{name}_c{i}', value=float(coef)) + else: + raise TypeError('Each coefficient value must be a number.') + self._coefficients.append(param) + else: + raise TypeError( + 'coefficients must be a sequence (list/tuple/ndarray) or dict of numbers.' + ) # Helper scipp scalar to track unit conversions # (value initialized to 1 with provided unit) @@ -130,7 +146,7 @@ def coefficients(self) -> list[Parameter]: return list(self._coefficients) @coefficients.setter - def coefficients(self, coeffs: Sequence[Numeric | Parameter]) -> None: + def coefficients(self, coeffs: Sequence[Numeric]) -> None: """ Set the coefficients of the polynomial. @@ -138,33 +154,27 @@ def coefficients(self, coeffs: Sequence[Numeric | Parameter]) -> None: Parameters ---------- - coeffs : Sequence[Numeric | Parameter] - New coefficients as a sequence of numbers or Parameters. + coeffs : Sequence[Numeric] + New coefficient values as a sequence of numbers. Raises ------ TypeError - If coeffs is not a sequence of numbers or Parameters or if any item in coeffs is not a - number or Parameter. + If coeffs is not a sequence of numbers or if any item is not a number. ValueError If the length of coeffs does not match the existing number of coefficients. """ if not isinstance(coeffs, (list, tuple, np.ndarray)): - raise TypeError( - 'coefficients must be a sequence (list/tuple/ndarray) of numbers or Parameter .' - ) + raise TypeError('coefficients must be a sequence (list/tuple/ndarray) of numbers.') if len(coeffs) != len(self._coefficients): raise ValueError( 'Number of coefficients must match the existing number of coefficients.' ) for i, coef in enumerate(coeffs): - if isinstance(coef, Parameter): - # replace parameter - self._coefficients[i] = coef - elif isinstance(coef, Numeric): + if isinstance(coef, Numeric): self._coefficients[i].value = float(coef) else: - raise TypeError('Each coefficient must be either a numeric value or a Parameter.') + raise TypeError('Each coefficient value must be a number.') def coefficient_values(self) -> list[float]: """ @@ -242,6 +252,49 @@ def degree(self, _value: int) -> None: 'and cannot be set directly.' ) + def add_coefficient(self, value: Numeric = 0.0, fixed: bool = False) -> None: + """ + Add a new coefficient at the next highest power, increasing the degree by one. + + Parameters + ---------- + value : Numeric, default=0.0 + The numeric value of the new coefficient. + fixed : bool, default=False + If True, the new coefficient is fixed (not free for fitting). + + Raises + ------ + TypeError + If ``value`` is not a numeric value. + """ + if not isinstance(value, Numeric): + raise TypeError('value must be a numeric value.') + new_power = len(self._coefficients) + param = Parameter(name=f'{self.name}_c{new_power}', value=float(value), fixed=fixed) + self._coefficients.append(param) + + def remove_coefficient(self) -> float: + """ + Remove and return the value of the highest-power coefficient, decreasing the degree by one. + + Returns + ------- + float + The value of the removed coefficient. + + Raises + ------ + ValueError + If there is only one coefficient remaining (at least one must always be present). + """ + if len(self._coefficients) == 1: + raise ValueError( + 'Cannot remove the last coefficient. The Polynomial must have at least one ' + 'coefficient.' + ) + return self._coefficients.pop().value + def get_all_variables(self) -> list[DescriptorBase]: """ Get all variables from the model component. diff --git a/tests/unit/easydynamics/sample_model/components/test_polynomial.py b/tests/unit/easydynamics/sample_model/components/test_polynomial.py index 2ba3f552..1138bf32 100644 --- a/tests/unit/easydynamics/sample_model/components/test_polynomial.py +++ b/tests/unit/easydynamics/sample_model/components/test_polynomial.py @@ -45,7 +45,7 @@ def test_initialization(self, polynomial: Polynomial): ), ( {'coefficients': [1.0, 'invalid', 3.0]}, - 'Each coefficient must be ', + 'Each coefficient', ), ( {'coefficients': [1.0, -2.0, 3.0], 'unit': 123}, @@ -55,7 +55,6 @@ def test_initialization(self, polynomial: Polynomial): {'coefficients': None}, 'coefficients must be ', ), - ({'coefficients': {}}, 'coefficients must be '), ], ) def test_input_type_validation_raises(self, kwargs, expected_message): @@ -94,32 +93,18 @@ def test_degree_setter_raises(self, polynomial: Polynomial): with pytest.raises(AttributeError, match='cannot be set directly'): polynomial.degree = 3 - @pytest.mark.parametrize( - 'values', - [ - [2.0, 0.0, -1.0], # all floats - [ - Parameter('p0', 2.0), - Parameter('p1', 0.0), - Parameter('p2', -1.0), - ], # all Parameters - [2.0, Parameter('p1', 0.0), -1.0], # mixed numbers and Parameters - ], - ) - def test_set_coefficients(self, polynomial: Polynomial, values): - """Test that coefficients can be updated from numeric values - or Parameters.""" + def test_set_coefficients(self, polynomial: Polynomial): + """Test that coefficients can be updated from numeric values.""" # WHEN - polynomial.coefficients = values + polynomial.coefficients = [2.0, 0.0, -1.0] - # THEN EXPECT: Parameter values match the new inputs - for i, val in enumerate(values): - expected = val.value if isinstance(val, Parameter) else val - assert np.isclose(polynomial.coefficients[i].value, expected) + # THEN EXPECT + assert np.isclose(polynomial.coefficients[0].value, 2.0) + assert np.isclose(polynomial.coefficients[1].value, 0.0) + assert np.isclose(polynomial.coefficients[2].value, -1.0) def test_set_coefficients_wrong_length_raises(self, polynomial: Polynomial): - """Ensure that setting coefficients with mismatched length - raises an error.""" + """Ensure that setting coefficients with mismatched length raises an error.""" with pytest.raises(ValueError, match='Number of coefficients'): polynomial.coefficients = [1.0, 2.0] # shorter list @@ -131,8 +116,8 @@ def test_set_coefficients_invalid_type_raises(self, polynomial: Polynomial): @pytest.mark.parametrize( 'invalid_coeffs, expected_message', [ - ([None, 2.0, 3.0], 'Each coefficient must be '), - ([1.0, 2.0, 'invalid'], 'Each coefficient must be '), + ([None, 2.0, 3.0], 'Each coefficient'), + ([1.0, 2.0, 'invalid'], 'Each coefficient'), ('not a list', 'coefficients must be '), ], ) @@ -197,5 +182,207 @@ def test_repr(self, polynomial: Polynomial): repr_str = repr(polynomial) # EXPECT - assert "name='PolynomialName'" in repr_str - assert 'coefficients=' in repr_str + assert 'name = PolynomialName' in repr_str + assert 'coefficients =' in repr_str + + # --- Serialization --- + + def test_to_dict(self, polynomial: Polynomial): + # WHEN + d = polynomial.to_dict() + + # EXPECT: base-class serialisation format + assert d['@class'] == 'Polynomial' + assert d['name'] == 'PolynomialName' + assert d['unit'] == 'meV' + coeffs = d['coefficients'] + assert len(coeffs) == 3 + assert coeffs[0]['@class'] == 'Parameter' + assert coeffs[0]['value'] == pytest.approx(1.0) + assert coeffs[1]['value'] == pytest.approx(-2.0) + assert coeffs[2]['value'] == pytest.approx(3.0) + + def test_from_dict_round_trip(self, polynomial: Polynomial): + # WHEN + d = polynomial.to_dict() + restored = Polynomial.from_dict(d) + + # EXPECT + assert restored.name == polynomial.name + assert restored.display_name == polynomial.display_name + assert restored.unit == polynomial.unit + assert restored.degree == polynomial.degree + np.testing.assert_allclose( + restored.coefficient_values(), polynomial.coefficient_values(), rtol=1e-10 + ) + + def test_from_dict_preserves_fixed(self): + # WHEN: sparse polynomial where c0 and c1 are fixed to zero + p = Polynomial(name='Sparse', coefficients={2: 1.5}) + d = p.to_dict() + restored = Polynomial.from_dict(d) + + # EXPECT + assert restored.coefficients[0].fixed is True + assert np.isclose(restored.coefficients[0].value, 0.0) + assert restored.coefficients[1].fixed is True + assert np.isclose(restored.coefficients[1].value, 0.0) + assert restored.coefficients[2].fixed is False + assert np.isclose(restored.coefficients[2].value, 1.5) + + def test_from_dict_wrong_class_raises(self, polynomial: Polynomial): + # WHEN + d = polynomial.to_dict() + d['@class'] = 'NotAPolynomial' + + # EXPECT + with pytest.raises(ValueError, match='Class name in dictionary does not match'): + Polynomial.from_dict(d) + + def test_from_dict_invalid_dict_raises(self): + with pytest.raises(ValueError, match='must be a dictionary representing'): + Polynomial.from_dict({'not': 'valid'}) + + # --- Sparse dict initialization --- + + def test_sparse_dict_single_term(self): + # WHEN + p = Polynomial(coefficients={2: 1.5}) + + # THEN EXPECT + assert p.degree == 2 + assert np.isclose(p.coefficients[0].value, 0.0) + assert p.coefficients[0].fixed is True + assert np.isclose(p.coefficients[1].value, 0.0) + assert p.coefficients[1].fixed is True + assert np.isclose(p.coefficients[2].value, 1.5) + assert p.coefficients[2].fixed is False + + def test_sparse_dict_evaluate(self): + # WHEN + p = Polynomial(coefficients={2: 1.5}) + x = np.array([0.0, 1.0, 2.0, 3.0]) + + # THEN + result = p.evaluate(x) + + # EXPECT + np.testing.assert_allclose(result, 1.5 * x**2, rtol=1e-5) + + def test_sparse_dict_multiple_powers(self): + # WHEN + p = Polynomial(coefficients={0: 1.0, 2: 3.0}) + + # THEN EXPECT: c1 (power 1) is fixed to zero, c0 and c2 are free + assert np.isclose(p.coefficients[0].value, 1.0) + assert p.coefficients[0].fixed is False + assert np.isclose(p.coefficients[1].value, 0.0) + assert p.coefficients[1].fixed is True + assert np.isclose(p.coefficients[2].value, 3.0) + assert p.coefficients[2].fixed is False + + def test_sparse_dict_empty_raises(self): + with pytest.raises(ValueError, match='At least one coefficient'): + Polynomial(coefficients={}) + + def test_sparse_dict_non_int_key_raises(self): + with pytest.raises(TypeError, match='Dict keys must be integers'): + Polynomial(coefficients={'a': 1.0}) + + def test_sparse_dict_negative_key_raises(self): + with pytest.raises(ValueError, match='non-negative'): + Polynomial(coefficients={-1: 1.0}) + + def test_sparse_dict_invalid_value_raises(self): + with pytest.raises(TypeError, match='Each coefficient'): + Polynomial(coefficients={2: 'bad'}) + + # --- add_coefficient --- + + def test_add_coefficient_increases_degree(self, polynomial: Polynomial): + # WHEN + polynomial.add_coefficient(5.0) + + # THEN EXPECT + assert polynomial.degree == 3 + assert np.isclose(polynomial.coefficients[3].value, 5.0) + + def test_add_coefficient_default_value(self, polynomial: Polynomial): + # WHEN + polynomial.add_coefficient() + + # THEN EXPECT + assert polynomial.degree == 3 + assert np.isclose(polynomial.coefficients[3].value, 0.0) + assert polynomial.coefficients[3].fixed is False + + def test_add_coefficient_fixed(self, polynomial: Polynomial): + # WHEN + polynomial.add_coefficient(0.0, fixed=True) + + # THEN EXPECT + assert polynomial.coefficients[3].fixed is True + + def test_add_coefficient_invalid_type_raises(self, polynomial: Polynomial): + with pytest.raises(TypeError, match='value must be'): + polynomial.add_coefficient('invalid') + + def test_add_coefficient_evaluate(self, polynomial: Polynomial): + # WHEN: polynomial is 1 - 2x + 3x^2; add coefficient 2.0 for x^3 + polynomial.add_coefficient(2.0) + x = np.array([1.0, 2.0]) + + # THEN + result = polynomial.evaluate(x) + + # EXPECT + expected = 1.0 - 2.0 * x + 3.0 * x**2 + 2.0 * x**3 + np.testing.assert_allclose(result, expected, rtol=1e-5) + + # --- remove_coefficient --- + + def test_remove_coefficient_decreases_degree(self, polynomial: Polynomial): + # WHEN + removed = polynomial.remove_coefficient() + + # THEN EXPECT + assert polynomial.degree == 1 + assert np.isclose(removed, 3.0) + + def test_remove_coefficient_returns_float(self, polynomial: Polynomial): + # WHEN + removed = polynomial.remove_coefficient() + + # THEN EXPECT + assert isinstance(removed, float) + + def test_remove_only_coefficient_raises(self): + # WHEN + p = Polynomial(coefficients=[5.0]) + + # THEN EXPECT + with pytest.raises(ValueError, match='Cannot remove the last coefficient'): + p.remove_coefficient() + + def test_remove_coefficient_evaluate(self, polynomial: Polynomial): + # WHEN: polynomial is 1 - 2x + 3x^2; remove x^2 term + polynomial.remove_coefficient() + x = np.array([0.0, 1.0, 2.0]) + + # THEN + result = polynomial.evaluate(x) + + # EXPECT: only 1 - 2x remains + expected = 1.0 - 2.0 * x + np.testing.assert_allclose(result, expected, rtol=1e-5) + + def test_add_then_remove_round_trip(self, polynomial: Polynomial): + # WHEN + polynomial.add_coefficient(4.0) + assert polynomial.degree == 3 + + polynomial.remove_coefficient() + + # THEN EXPECT: back to original degree + assert polynomial.degree == 2 + assert np.isclose(polynomial.coefficients[2].value, 3.0)