From c08a3acf95a7bc26e7779e552154387995211520 Mon Sep 17 00:00:00 2001 From: TJ Lane Date: Sun, 27 Oct 2024 15:04:46 +0000 Subject: [PATCH] use golden optimization for TV weight (#45) * golden optimizer * test precision * @alisiafadini feedback --- meteor/diffmaps.py | 2 +- meteor/iterative.py | 2 +- meteor/scripts/compute_difference_map.py | 19 +++++----------- meteor/settings.py | 4 ++-- meteor/tv.py | 22 +++++++++---------- .../functional/test_compute_difference_map.py | 11 ++++------ .../scripts/test_compute_difference_map.py | 12 ++++------ test/unit/test_iterative.py | 4 ++-- test/unit/test_tv.py | 6 ++--- 9 files changed, 33 insertions(+), 49 deletions(-) diff --git a/meteor/diffmaps.py b/meteor/diffmaps.py index db22df5..d905940 100644 --- a/meteor/diffmaps.py +++ b/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, diff --git a/meteor/iterative.py b/meteor/iterative.py index 8a3d034..5e6bab7 100644 --- a/meteor/iterative.py +++ b/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, }, diff --git a/meteor/scripts/compute_difference_map.py b/meteor/scripts/compute_difference_map.py index d8dd379..b564e03 100644 --- a/meteor/scripts/compute_difference_map.py +++ b/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], ) diff --git a/meteor/settings.py b/meteor/settings.py index dd481b6..279b4e6 100644 --- a/meteor/settings.py +++ b/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 diff --git a/meteor/tv.py b/meteor/tv.py index b32366c..1d58670 100644 --- a/meteor/tv.py +++ b/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 diff --git a/test/functional/test_compute_difference_map.py b/test/functional/test_compute_difference_map.py index 325efcd..fc95346 100644 --- a/test/functional/test_compute_difference_map.py +++ b/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", ) diff --git a/test/unit/scripts/test_compute_difference_map.py b/test/unit/scripts/test_compute_difference_map.py index 78bbdb1..70452ad 100644 --- a/test/unit/scripts/test_compute_difference_map.py +++ b/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: diff --git a/test/unit/test_iterative.py b/test/unit/test_iterative.py index 71a8a4c..44b19f6 100644 --- a/test/unit/test_iterative.py +++ b/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, ) diff --git a/test/unit/test_tv.py b/test/unit/test_tv.py index 4ebca8b..4757b44 100644 --- a/test/unit/test_tv.py +++ b/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")