use golden optimization for TV weight (#45)

* golden optimizer

* test precision

* @alisiafadini feedback

meteor/diffmaps.py

@@ -174,7 +174,7 @@ def negentropy_objective(k_parameter: float) -> float:
 
     maximizer = ScalarMaximizer(objective=negentropy_objective)
     maximizer.optimize_over_explicit_values(arguments_to_scan=k_parameter_values_to_scan)
-    opt_k_parameter = maximizer.argument_optimum
+    opt_k_parameter = float(maximizer.argument_optimum)
 
     kweighted_dataset = compute_kweighted_difference_map(
         derivative,

meteor/iterative.py

@@ -122,7 +122,7 @@ def _complex_derivative_from_iterative_tv(
         metadata.append(
             {
                 "iteration": num_iterations,
-                "tv_weight": tv_metadata.optimal_weight,
+                "tv_weight": tv_metadata.optimal_tv_weight,
                 "negentropy_after_tv": tv_metadata.optimal_negentropy,
                 "average_phase_change": phase_change,
             },

meteor/scripts/compute_difference_map.py

@@ -20,10 +20,6 @@
 log = structlog.get_logger()
 
 
-# TO OPTMIZE
-TV_WEIGHTS_TO_SCAN = np.linspace(0.0, 0.1, 101)
-
-
 class TvDiffmapArgParser(DiffmapArgParser):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
@@ -153,20 +149,15 @@ def denoise_diffmap_according_to_mode(
     """
     if tv_denoise_mode == WeightMode.optimize:
         log.info(
-            "Searching for max-negentropy TV denoising weight",
-            min=np.min(TV_WEIGHTS_TO_SCAN),
-            max=np.max(TV_WEIGHTS_TO_SCAN),
-            points_to_test=len(TV_WEIGHTS_TO_SCAN),
+            "Searching for max-negentropy TV denoising weight", method="golden ration optimization"
         )
         log.info("This may take some time...")
 
-        final_map, metadata = tv_denoise_difference_map(
-            diffmap, full_output=True, weights_to_scan=TV_WEIGHTS_TO_SCAN
-        )
+        final_map, metadata = tv_denoise_difference_map(diffmap, full_output=True)
 
         log.info(
             "Optimal TV weight found",
-            weight=metadata.optimal_weight,
+            weight=metadata.optimal_tv_weight,
             initial_negentropy=f"{metadata.initial_negentropy:.2e}",
             final_negetropy=f"{metadata.optimal_negentropy:.2e}",
         )
@@ -196,9 +187,9 @@ def denoise_diffmap_according_to_mode(
         metadata = TvDenoiseResult(
             initial_negentropy=map_negetropy,
             optimal_negentropy=map_negetropy,
-            optimal_weight=0.0,
+            optimal_tv_weight=0.0,
             map_sampling_used_for_tv=MAP_SAMPLING,
-            weights_scanned=[0.0],
+            tv_weights_scanned=[0.0],
             negentropy_at_weights=[map_negetropy],
         )
 

meteor/settings.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-TV_WEIGHT_RANGE: tuple[float, float] = (0.0, 1.0)
+BRACKET_FOR_GOLDEN_OPTIMIZATION: tuple[float, float] = (0.0, 0.05)
 TV_STOP_TOLERANCE: float = 0.00000005
 TV_MAX_NUM_ITER: int = 50
 MAP_SAMPLING: int = 3
 
 KWEIGHT_PARAMETER_DEFAULT: float = 0.05
-TV_WEIGHT_DEFAULT: float = 0.01  # TO OPTIMIZE
+TV_WEIGHT_DEFAULT: float = 0.01
 
 COMPUTED_MAP_RESOLUTION_LIMIT: float = 1.0
 GEMMI_HIGH_RESOLUTION_BUFFER: float = 1e-6

meteor/tv.py

@@ -13,7 +13,7 @@
     MAP_SAMPLING,
     TV_MAX_NUM_ITER,
     TV_STOP_TOLERANCE,
-    TV_WEIGHT_RANGE,
+    BRACKET_FOR_GOLDEN_OPTIMIZATION,
 )
 from .validate import ScalarMaximizer, negentropy
 
@@ -26,23 +26,23 @@ class TvDenoiseResult:
     _negentropy_name = "negentropy"
 
     initial_negentropy: float
-    optimal_weight: float
+    optimal_tv_weight: float
     optimal_negentropy: float
     map_sampling_used_for_tv: float
-    weights_scanned: list[float]
+    tv_weights_scanned: list[float]
     negentropy_at_weights: list[float]
     k_parameter_used: float | None = None
 
     def json(self) -> dict:
         json_payload = asdict(self)
-        json_payload.pop("weights_scanned")
+        json_payload.pop("tv_weights_scanned")
         json_payload.pop("negentropy_at_weights")
         json_payload[self._scan_name] = [
             {
-                self._weight_name: float(self.weights_scanned[idx]),
+                self._weight_name: float(self.tv_weights_scanned[idx]),
                 self._negentropy_name: float(self.negentropy_at_weights[idx]),
             }
-            for idx in range(len(self.weights_scanned))
+            for idx in range(len(self.tv_weights_scanned))
         ]
         return json_payload
 
@@ -54,7 +54,7 @@ def to_json_file(self, filename: Path) -> None:
     def from_json(cls, json_payload: dict) -> TvDenoiseResult:
         try:
             data = json_payload.pop(cls._scan_name)
-            json_payload["weights_scanned"] = [float(point[cls._weight_name]) for point in data]
+            json_payload["tv_weights_scanned"] = [float(point[cls._weight_name]) for point in data]
             json_payload["negentropy_at_weights"] = [
                 float(point[cls._negentropy_name]) for point in data
             ]
@@ -157,7 +157,7 @@ def tv_denoise_difference_map(
     -------
     >>> coefficients = Map.read_mtz("./path/to/difference_map.mtz", ...)  # load dataset
     >>> denoised_map, result = tv_denoise_difference_map(coefficients, full_output=True)
-    >>> print(f"Optimal: {result.optimal_weight}, Negentropy: {result.optimal_negentropy}")
+    >>> print(f"Optimal: {result.optimal_tv_weight}, Negentropy: {result.optimal_negentropy}")
     """
     realspace_map = difference_map.to_ccp4_map(map_sampling=MAP_SAMPLING)
     realspace_map_array = np.array(realspace_map.grid)
@@ -170,7 +170,7 @@ def negentropy_objective(tv_weight: float) -> float:
     if weights_to_scan is not None:
         maximizer.optimize_over_explicit_values(arguments_to_scan=weights_to_scan)
     else:
-        maximizer.optimize_with_golden_algorithm(bracket=TV_WEIGHT_RANGE)
+        maximizer.optimize_with_golden_algorithm(bracket=BRACKET_FOR_GOLDEN_OPTIMIZATION)
 
     # denoise using the optimized parameters and convert to an rs.DataSet
     final_realspace_map_as_array = _tv_denoise_array(
@@ -201,10 +201,10 @@ def negentropy_objective(tv_weight: float) -> float:
         initial_negentropy = negentropy(realspace_map_array)
         tv_result = TvDenoiseResult(
             initial_negentropy=initial_negentropy,
-            optimal_weight=maximizer.argument_optimum,
+            optimal_tv_weight=maximizer.argument_optimum,
             optimal_negentropy=maximizer.objective_maximum,
             map_sampling_used_for_tv=MAP_SAMPLING,
-            weights_scanned=maximizer.values_evaluated,
+            tv_weights_scanned=maximizer.values_evaluated,
             negentropy_at_weights=maximizer.objective_at_values,
         )
         return final_map, tv_result

test/functional/test_compute_difference_map.py

@@ -1,5 +1,4 @@
 from pathlib import Path
-from unittest import mock
 
 import numpy as np
 import pytest
@@ -11,10 +10,7 @@
 from meteor.tv import TvDenoiseResult
 from meteor.utils import filter_common_indices
 
-TV_WEIGHTS_TO_SCAN = np.array([0.005, 0.01, 0.025, 0.5])
 
-
-@mock.patch("meteor.scripts.compute_difference_map.TV_WEIGHTS_TO_SCAN", TV_WEIGHTS_TO_SCAN)
 @pytest.mark.parametrize("kweight_mode", list(WeightMode))
 @pytest.mark.parametrize("tv_weight_mode", list(WeightMode))
 def test_script_produces_consistent_results(
@@ -84,14 +80,15 @@ def test_script_produces_consistent_results(
             optimal_tv_no_kweighting = 0.025
             np.testing.assert_allclose(
                 optimal_tv_no_kweighting,
-                result_metadata.optimal_weight,
+                result_metadata.optimal_tv_weight,
                 rtol=0.1,
                 err_msg="tv weight optimium different from expected",
             )
         else:
+            optimal_tv_with_weighting = 0.00867
             np.testing.assert_allclose(
-                tv_weight,
-                result_metadata.optimal_weight,
+                optimal_tv_with_weighting,
+                result_metadata.optimal_tv_weight,
                 rtol=0.1,
                 err_msg="tv weight optimium different from expected",
             )

test/unit/scripts/test_compute_difference_map.py

@@ -5,7 +5,6 @@
 from typing import Any
 from unittest import mock
 
-import numpy as np
 import pytest
 
 from meteor.rsmap import Map
@@ -18,10 +17,6 @@
 )
 from meteor.tv import TvDenoiseResult
 
-# ensure tests complete quickly by monkey-patching a limited number of weights
-compute_difference_map.TV_WEIGHTS_TO_SCAN = np.linspace(0.0, 0.1, 6)
-
-
 TV_WEIGHT = 0.1
 
 
@@ -80,17 +75,18 @@ def test_denoise_diffmap_according_to_mode(mode: WeightMode, random_difference_m
     assert isinstance(metadata, TvDenoiseResult)
 
     if mode == WeightMode.optimize:
-        assert metadata.optimal_weight in [0.04, 0.06]  # random test; alternates
+        # random test; changes
+        assert 0.04 < metadata.optimal_tv_weight < 0.06
 
     elif mode == WeightMode.fixed:
-        assert metadata.optimal_weight == TV_WEIGHT
+        assert metadata.optimal_tv_weight == TV_WEIGHT
         with pytest.raises(TypeError):
             _, _ = denoise_diffmap_according_to_mode(
                 diffmap=random_difference_map, tv_denoise_mode=mode, tv_weight=None
             )
 
     elif mode == WeightMode.none:
-        assert metadata.optimal_weight == 0.0
+        assert metadata.optimal_tv_weight == 0.0
 
 
 def test_main(diffmap_set: DiffMapSet, tmp_path: Path, fixed_kparameter: float) -> None:

test/unit/test_iterative.py

@@ -30,9 +30,9 @@ def simple_tv_function(fourier_array: np.ndarray) -> tuple[np.ndarray, TvDenoise
     denoised = denoise_tv_chambolle(real_space, weight=weight)  # type: ignore[no-untyped-call]
     result = TvDenoiseResult(
         initial_negentropy=0.0,
-        optimal_weight=weight,
+        optimal_tv_weight=weight,
         optimal_negentropy=1.0,
-        weights_scanned=[weight],
+        tv_weights_scanned=[weight],
         negentropy_at_weights=[1.0],
         map_sampling_used_for_tv=0,
     )

test/unit/test_tv.py

@@ -30,10 +30,10 @@ def rms_between_coefficients(map1: Map, map2: Map) -> float:
 def tv_denoise_result_source_data() -> dict:
     return {
         "initial_negentropy": 0.0,
-        "optimal_weight": 1.0,
+        "optimal_tv_weight": 1.0,
         "optimal_negentropy": 5.0,
         "map_sampling_used_for_tv": 5,
-        "weights_scanned": [0.0, 1.0],
+        "tv_weights_scanned": [0.0, 1.0],
         "negentropy_at_weights": [0.0, 5.0],
         "k_parameter_used": 0.0,
     }
@@ -131,4 +131,4 @@ def rms_to_noise_free(test_map: Map) -> float:
     )
 
     assert rms_to_noise_free(denoised_map) < rms_to_noise_free(noisy_map), "error didnt drop"
-    np.testing.assert_allclose(result.optimal_weight, best_weight, rtol=0.5, err_msg="opt weight")
+    np.testing.assert_allclose(result.optimal_tv_weight, best_weight, rtol=0.5, err_msg="opt weight")