FEAT: make ParameterFlattener public (#546)

Author: grayson-helmholz

src/tensorwaves/optimizer/minuit.py

@@ -11,9 +11,11 @@
 from tqdm.auto import tqdm
 
 from tensorwaves.interface import Estimator, FitResult, Optimizer, ParameterValue
-
-from ._parameter import ParameterFlattener
-from .callbacks import Callback, _create_log  # pyright: ignore[reportPrivateUsage]
+from tensorwaves.optimizer.callbacks import (
+    Callback,
+    _create_log,  # pyright: ignore[reportPrivateUsage]
+)
+from tensorwaves.optimizer.parameter import ParameterFlattener
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping

src/tensorwaves/optimizer/_parameter.py renamed to src/tensorwaves/optimizer/parameter.py

@@ -1,3 +1,5 @@
+"""Utility to splice (flatten) and merge (unflatten) complex parameters for 'real-only' optimizers."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
@@ -9,6 +11,13 @@
 
 
 class ParameterFlattener:
+    """Utility-class to flatten complex parameters.
+
+    Args:
+        parameters: Original parameter-dictionary (unflattened). Non-complex values will
+            not be affected by any method.
+    """
+
     def __init__(self, parameters: Mapping[str, ParameterValue]) -> None:
         self.__real_imag_to_complex_name: dict[str, str] = {}
         self.__complex_to_real_imag_name: dict[str, tuple[str, str]] = {}
@@ -23,6 +32,17 @@ def __init__(self, parameters: Mapping[str, ParameterValue]) -> None:
     def unflatten(
         self, flattened_parameters: dict[str, float]
     ) -> dict[str, ParameterValue]:
+        """Reverse the flattening operation.
+
+        Takes a parameter-dictionary and merges all real and imaginary values whose
+        respective keys have been registered in the constructor of the
+        `ParameterFlattener` into a complex number. Specifically, while this works also
+        on inputs which have not been generated by :meth:`.flatten` their outputs might
+        be unexpected.
+
+        Args:
+            flattened_parameters: parameter `dict` whose values are to be unflattened.
+        """
         parameters: dict[str, ParameterValue] = {
             k: v
             for k, v in flattened_parameters.items()
@@ -39,6 +59,15 @@ def unflatten(
         return parameters
 
     def flatten(self, parameters: Mapping[str, ParameterValue]) -> dict[str, float]:
+        """Flatten the parameter-values whose keys have been registered in the constructor.
+
+        Splits all complex values whose keys have been registered in the constructor of
+        `ParameterFlattener` into their real and imaginary parts. Their keys are
+        predetermined by the constructor. Other key-value pairs remain unchanged.
+
+        Args:
+            parameters: parameter `dict` whose values are to be flattened.
+        """
         flattened_parameters: dict[str, float] = {}
         for par_name, value in parameters.items():
             if isinstance(value, complex):

src/tensorwaves/optimizer/scipy.py

@@ -11,8 +11,8 @@
 
 from tensorwaves.function._backend import raise_missing_module_error
 from tensorwaves.interface import Estimator, FitResult, Optimizer, ParameterValue
+from tensorwaves.optimizer.parameter import ParameterFlattener
 
-from ._parameter import ParameterFlattener
 from .callbacks import Callback, _create_log  # pyright: ignore[reportPrivateUsage]
 
 if TYPE_CHECKING:

tests/optimizer/test_parameter.py

@@ -1,6 +1,6 @@
 import pytest
 
-from tensorwaves.optimizer._parameter import ParameterFlattener
+from tensorwaves.optimizer.parameter import ParameterFlattener
 
 
 @pytest.fixture(scope="module")