stronger testing for (single-pass) TV denoising

.github/workflows/mypy.yml

@@ -24,4 +24,5 @@ jobs:
       run: |
         pip install mypy
         pip install .
+        mkdir .mypy_cache
         mypy --install-types --non-interactive meteor/ test/

meteor/tv.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Literal, Sequence, overload
 
 import numpy as np
@@ -26,6 +26,7 @@ class TvDenoiseResult:
     optimal_lambda: float
     optimal_negentropy: float
     map_sampling_used_for_tv: float
+    negetropies_for_lambdas_scanned: list[float] = field(default_factory=list)
 
 
 def _tv_denoise_array(*, map_as_array: np.ndarray, weight: float) -> np.ndarray:
@@ -42,10 +43,10 @@ def _tv_denoise_array(*, map_as_array: np.ndarray, weight: float) -> np.ndarray:
 @overload
 def tv_denoise_difference_map(
     difference_map_coefficients: rs.DataSet,
-    full_output: Literal[False],
+    full_output: Literal[False] = False,
     difference_map_amplitude_column: str = "DF",
     difference_map_phase_column: str = "PHIC",
-    lambda_values_to_scan: Sequence[float] | None = None,
+    lambda_values_to_scan: Sequence[float] | np.ndarray | None = None,
 ) -> rs.DataSet: ...
 
 
@@ -55,7 +56,7 @@ def tv_denoise_difference_map(
     full_output: Literal[True],
     difference_map_amplitude_column: str = "DF",
     difference_map_phase_column: str = "PHIC",
-    lambda_values_to_scan: Sequence[float] | None = None,
+    lambda_values_to_scan: Sequence[float] | np.ndarray | None = None,
 ) -> tuple[rs.DataSet, TvDenoiseResult]: ...
 
 
@@ -140,18 +141,20 @@ def negentropy_objective(tv_lambda: float):
         return -1.0 * negentropy(denoised_map.flatten())
 
     optimal_lambda: float
+    negetropies_for_lambdas_scanned = []
 
     # scan a specific set of lambda values and find the best one
     if lambda_values_to_scan is not None:
         # use no denoising as the default to beat
-        optimal_lambda = 0.0  # initialization
-        highest_negentropy = negentropy(difference_map_as_array.flatten())
+        optimal_lambda = lambda_values_to_scan[0]  # initialization
+        optimal_negentropy = negentropy(difference_map_as_array.flatten())
 
         for tv_lambda in lambda_values_to_scan:
             trial_negentropy = -1.0 * negentropy_objective(tv_lambda)
-            if trial_negentropy > highest_negentropy:
+            negetropies_for_lambdas_scanned.append(trial_negentropy)
+            if trial_negentropy > optimal_negentropy:
                 optimal_lambda = tv_lambda
-                highest_negentropy = trial_negentropy
+                optimal_negentropy = trial_negentropy
 
     # use golden ratio optimization to pick an optimal lambda
     else:
@@ -160,7 +163,7 @@ def negentropy_objective(tv_lambda: float):
         )
         assert optimizer_result.success, "Golden minimization failed to find optimal TV lambda"
         optimal_lambda = optimizer_result.x
-        highest_negentropy = negentropy_objective(optimal_lambda)
+        optimal_negentropy = -1.0 * negentropy_objective(optimal_lambda)
 
     # denoise using the optimized parameters and convert to an rs.DataSet
     final_map = _tv_denoise_array(map_as_array=difference_map_as_array, weight=optimal_lambda)
@@ -180,8 +183,9 @@ def negentropy_objective(tv_lambda: float):
     if full_output:
         tv_result = TvDenoiseResult(
             optimal_lambda=optimal_lambda,
-            optimal_negentropy=highest_negentropy,
+            optimal_negentropy=optimal_negentropy,
             map_sampling_used_for_tv=TV_MAP_SAMPLING,
+            negetropies_for_lambdas_scanned=negetropies_for_lambdas_scanned,
         )
         return final_map_coefficients, tv_result
     else:

test/unit/conftest.py

@@ -0,0 +1,88 @@
+import gemmi
+import numpy as np
+import pytest
+import reciprocalspaceship as rs
+
+from meteor.utils import compute_coefficients_from_map
+
+DEFAULT_RESOLUTION = 1.0
+DEFAULT_CELL_SIZE = 10.0
+DEFAULT_CARBON1_POSITION = (5.0, 5.0, 5.0)
+DEFAULT_CARBON2_POSITION = (5.0, 5.2, 5.0)
+
+
+def generate_single_carbon_density(
+    carbon_position: tuple[float, float, float],
+    space_group: gemmi.SpaceGroup,
+    unit_cell: gemmi.UnitCell,
+    d_min: float,
+) -> gemmi.FloatGrid:
+    model = gemmi.Model("single_atom")
+    chain = gemmi.Chain("A")
+
+    residue = gemmi.Residue()
+    residue.name = "X"
+    residue.seqid = gemmi.SeqId("1")
+
+    atom = gemmi.Atom()
+    atom.name = "C"
+    atom.element = gemmi.Element("C")
+    atom.pos = gemmi.Position(*carbon_position)
+
+    residue.add_atom(atom)
+    chain.add_residue(residue)
+    model.add_chain(chain)
+
+    structure = gemmi.Structure()
+    structure.add_model(model)
+    structure.cell = unit_cell
+    structure.spacegroup_hm = space_group.hm
+
+    density_map = gemmi.DensityCalculatorX()
+    density_map.d_min = d_min
+    density_map.grid.setup_from(structure)
+    density_map.put_model_density_on_grid(structure[0])
+
+    return density_map.grid
+
+
+def displaced_single_atom_difference_map_coefficients(
+    *,
+    noise_sigma: float,
+    d_min: float = DEFAULT_RESOLUTION,
+    cell_size: float = DEFAULT_CELL_SIZE,
+    carbon_position1: tuple[float, float, float] = DEFAULT_CARBON1_POSITION,
+    carbon_position2: tuple[float, float, float] = DEFAULT_CARBON2_POSITION,
+) -> rs.DataSet:
+    unit_cell = gemmi.UnitCell(a=cell_size, b=cell_size, c=cell_size, alpha=90, beta=90, gamma=90)
+    space_group = gemmi.find_spacegroup_by_name("P1")
+
+    density1 = generate_single_carbon_density(carbon_position1, space_group, unit_cell, d_min)
+    density2 = generate_single_carbon_density(carbon_position2, space_group, unit_cell, d_min)
+
+    ccp4_map = gemmi.Ccp4Map()
+    grid_values = (
+        np.array(density2) - np.array(density1) + noise_sigma * np.random.randn(*density2.shape)
+    )
+    ccp4_map.grid = gemmi.FloatGrid(grid_values.astype(np.float32), unit_cell, space_group)
+    ccp4_map.update_ccp4_header()
+
+    difference_map_coefficients = compute_coefficients_from_map(
+        ccp4_map=ccp4_map,
+        high_resolution_limit=d_min,
+        amplitude_label="DF",
+        phase_label="PHIC",
+    )
+    assert (difference_map_coefficients.max() > 0.0).any()
+
+    return difference_map_coefficients
+
+
+@pytest.fixture
+def noise_free_map() -> rs.DataSet:
+    return displaced_single_atom_difference_map_coefficients(noise_sigma=0.0)
+
+
+@pytest.fixture
+def noisy_map() -> rs.DataSet:
+    return displaced_single_atom_difference_map_coefficients(noise_sigma=0.03)

test/unit/test_tv.py

@@ -1,98 +1,35 @@
 from typing import Sequence
 
-import gemmi
 import numpy as np
 import pytest
 import reciprocalspaceship as rs
 
 from meteor import tv
-from meteor.utils import MapLabels, compute_coefficients_from_map, compute_map_from_coefficients
+from meteor.utils import MapLabels, compute_map_from_coefficients
 
+DEFAULT_LAMBDA_VALUES_TO_SCAN = np.logspace(-2, 0, 30)
 
-def _generate_single_carbon_density(
-    carbon_position: tuple[float, float, float],
-    space_group: gemmi.SpaceGroup,
-    unit_cell: gemmi.UnitCell,
-    d_min: float,
-) -> gemmi.FloatGrid:
-    model = gemmi.Model("single_atom")
-    chain = gemmi.Chain("A")
 
-    residue = gemmi.Residue()
-    residue.name = "X"
-    residue.seqid = gemmi.SeqId("1")
-
-    atom = gemmi.Atom()
-    atom.name = "C"
-    atom.element = gemmi.Element("C")
-    atom.pos = gemmi.Position(*carbon_position)
-
-    residue.add_atom(atom)
-    chain.add_residue(residue)
-    model.add_chain(chain)
-
-    structure = gemmi.Structure()
-    structure.add_model(model)
-    structure.cell = unit_cell
-    structure.spacegroup_hm = space_group.hm
-
-    density_map = gemmi.DensityCalculatorX()
-    density_map.d_min = d_min
-    density_map.grid.setup_from(structure)
-    density_map.put_model_density_on_grid(structure[0])
-
-    return density_map.grid
-
-
-def displaced_single_atom_difference_map_coefficients(
-    *,
-    noise_sigma: float,
-) -> rs.DataSet:
-    unit_cell = gemmi.UnitCell(a=10.0, b=10.0, c=10.0, alpha=90, beta=90, gamma=90)
-    space_group = gemmi.find_spacegroup_by_name("P1")
-    d_min = 1.0
-
-    carbon_position1 = (5.0, 5.0, 5.0)
-    carbon_position2 = (5.1, 5.0, 5.0)
-
-    density1 = _generate_single_carbon_density(carbon_position1, space_group, unit_cell, d_min)
-    density2 = _generate_single_carbon_density(carbon_position2, space_group, unit_cell, d_min)
-
-    ccp4_map = gemmi.Ccp4Map()
-    grid_values = (
-        np.array(density2) - np.array(density1) + noise_sigma * np.random.randn(*density2.shape)
-    )
-    ccp4_map.grid = gemmi.FloatGrid(grid_values.astype(np.float32), unit_cell, space_group)
-    ccp4_map.update_ccp4_header()
-
-    difference_map_coefficients = compute_coefficients_from_map(
-        ccp4_map=ccp4_map,
-        high_resolution_limit=1.0,
-        amplitude_label="DF",
-        phase_label="PHIC",
-    )
-    assert (difference_map_coefficients.max() > 0.0).any()
-
-    return difference_map_coefficients
-
-
-def rms_between_coefficients(ds1: rs.DataSet, ds2: rs.DataSet, diffmap_labels: MapLabels) -> float:
+def rms_between_coefficients(
+    ds1: rs.DataSet, ds2: rs.DataSet, diffmap_labels: MapLabels, map_sampling: int = 3
+) -> float:
     map1 = compute_map_from_coefficients(
         map_coefficients=ds1,
         amplitude_label=diffmap_labels.amplitude,
         phase_label=diffmap_labels.phase,
-        map_sampling=3,
+        map_sampling=map_sampling,
     )
     map2 = compute_map_from_coefficients(
         map_coefficients=ds2,
         amplitude_label=diffmap_labels.amplitude,
         phase_label=diffmap_labels.phase,
-        map_sampling=3,
+        map_sampling=map_sampling,
     )
 
-    map1_array = np.array(map2.grid)
-    map2_array = np.array(map1.grid)
+    map1_array = np.array(map1.grid)
+    map2_array = np.array(map2.grid)
 
+    # standardize
     map1_array /= map1_array.std()
     map2_array /= map2_array.std()
 
@@ -101,16 +38,6 @@ def rms_between_coefficients(ds1: rs.DataSet, ds2: rs.DataSet, diffmap_labels: M
     return rms
 
 
-@pytest.fixture
-def noise_free_map() -> rs.DataSet:
-    return displaced_single_atom_difference_map_coefficients(noise_sigma=0.0)
-
-
-@pytest.fixture
-def noisy_map() -> rs.DataSet:
-    return displaced_single_atom_difference_map_coefficients(noise_sigma=0.03)
-
-
 @pytest.mark.parametrize(
     "lambda_values_to_scan",
     [
@@ -139,42 +66,43 @@ def test_tv_denoise_difference_map_smoke(
         assert isinstance(output, rs.DataSet)
 
 
-@pytest.mark.parametrize("lambda_values_to_scan", [None, np.logspace(-3, 2, 100)])
+@pytest.mark.parametrize("lambda_values_to_scan", [None, DEFAULT_LAMBDA_VALUES_TO_SCAN])
 def test_tv_denoise_difference_map(
     lambda_values_to_scan: None | Sequence[float],
     noise_free_map: rs.DataSet,
     noisy_map: rs.DataSet,
     diffmap_labels: MapLabels,
 ) -> None:
-    rms_before_denoising = rms_between_coefficients(noise_free_map, noisy_map, diffmap_labels)
+    def rms_to_noise_free(test_map: rs.DataSet) -> float:
+        return rms_between_coefficients(test_map, noise_free_map, diffmap_labels)
+
+    # Normally, the `tv_denoise_difference_map` function only returns the best result -- since we
+    # know  the ground truth, work around this to test all possible results.
+
+    lowest_rms: float = np.inf
+    best_lambda: float = 0.0
+
+    for trial_lambda in DEFAULT_LAMBDA_VALUES_TO_SCAN:
+        denoised_map, result = tv.tv_denoise_difference_map(
+            difference_map_coefficients=noisy_map,
+            lambda_values_to_scan=[
+                trial_lambda,
+            ],
+            full_output=True,
+        )
+        rms = rms_to_noise_free(denoised_map)
+        if rms < lowest_rms:
+            lowest_rms = rms
+            best_lambda = trial_lambda
+
+    # now run the denoising algorithm and make sure we get a result that's close
+    # to the one that minimizes the RMS error to the ground truth
     denoised_map, result = tv.tv_denoise_difference_map(
         difference_map_coefficients=noisy_map,
         lambda_values_to_scan=lambda_values_to_scan,
         full_output=True,
     )
-    rms_after_denoising = rms_between_coefficients(noise_free_map, denoised_map, diffmap_labels)
-    assert rms_after_denoising < rms_before_denoising
-
-    # print("xyz", result.optimal_lambda, rms_before_denoising, rms_after_denoising)
-
-    # testmap = compute_map_from_coefficients(
-    #     map_coefficients=noise_free_map,
-    #     amplitude_label=diffmap_labels.amplitude,
-    #     phase_label=diffmap_labels.phase,
-    #     map_sampling=1,
-    # )
-    # testmap.write_ccp4_map("original.ccp4")
-    # testmap = compute_map_from_coefficients(
-    #     map_coefficients=noisy_map,
-    #     amplitude_label=diffmap_labels.amplitude,
-    #     phase_label=diffmap_labels.phase,
-    #     map_sampling=1,
-    # )
-    # testmap.write_ccp4_map("noisy.ccp4")
-    # testmap = compute_map_from_coefficients(
-    #     map_coefficients=denoised_map,
-    #     amplitude_label=diffmap_labels.amplitude,
-    #     phase_label=diffmap_labels.phase,
-    #     map_sampling=1,
-    # )
-    # testmap.write_ccp4_map("denoised.ccp4")
+
+    rms_after_denoising = rms_to_noise_free(denoised_map)
+    assert rms_after_denoising < rms_to_noise_free(noisy_map)
+    np.testing.assert_allclose(result.optimal_lambda, best_lambda, rtol=0.2)