Allow adding and removing coefficients from Polynomial

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,41 +146,35 @@ 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.
 
         Length must match current number of coefficients.
 
         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.