testing in place

meteor/scripts/common.py

@@ -1,19 +1,18 @@
 from __future__ import annotations
 
 import argparse
+import json
 import re
 from dataclasses import dataclass
 from enum import StrEnum, auto
 from pathlib import Path
 from typing import Any
 
 import numpy as np
+import pandas as pd
 import reciprocalspaceship as rs
 import structlog
-import json
-import pandas as pd
 
-from meteor.tv import TvDenoiseResult
 from meteor.diffmaps import (
     compute_difference_map,
     compute_kweighted_difference_map,
@@ -24,6 +23,7 @@
 from meteor.scale import scale_maps
 from meteor.settings import COMPUTED_MAP_RESOLUTION_LIMIT, KWEIGHT_PARAMETER_DEFAULT
 from meteor.sfcalc import structure_file_to_calculated_map
+from meteor.tv import TvDenoiseResult
 
 log = structlog.get_logger()
 
@@ -314,7 +314,9 @@ def kweight_diffmap_according_to_mode(
     return diffmap, kweight_parameter
 
 
-def write_combined_metadata(*, filename: Path, it_tv_metadata: pd.DataFrame, final_tv_metadata: TvDenoiseResult) -> None:
+def write_combined_metadata(
+    *, filename: Path, it_tv_metadata: pd.DataFrame, final_tv_metadata: TvDenoiseResult
+) -> None:
     combined_metadata = {
         "iterative_tv": it_tv_metadata.to_json(),
         "final_tv_pass": final_tv_metadata.json(),
@@ -326,7 +328,6 @@ def write_combined_metadata(*, filename: Path, it_tv_metadata: pd.DataFrame, fin
 def read_combined_metadata(*, filename: Path) -> tuple[pd.DataFrame, TvDenoiseResult]:
     with filename.open("r") as f:
         combined_metadata = json.load(f)
-    print("***", pd.read_json(combined_metadata["iterative_tv"]))
     it_tv_metadata = pd.read_json(combined_metadata["iterative_tv"])
     final_tv_metadata = TvDenoiseResult.from_json(combined_metadata["final_tv_pass"])
     return it_tv_metadata, final_tv_metadata

meteor/scripts/compute_iterative_tv_map.py

@@ -7,8 +7,7 @@
 from meteor.iterative import iterative_tv_phase_retrieval
 from meteor.tv import tv_denoise_difference_map
 
-
-from .common import DiffmapArgParser, kweight_diffmap_according_to_mode
+from .common import DiffmapArgParser, kweight_diffmap_according_to_mode, write_combined_metadata
 
 log = structlog.get_logger()
 
@@ -36,7 +35,14 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 def main(command_line_arguments: list[str] | None = None) -> None:
     parser = IterativeTvArgParser(
         description=(
-            "bla bla"  # TODO
+            "Compute an difference map, where the phases of the derivative structure are estimated "
+            "using the assumption that the resulting map should have a low total variation. Phases "
+            "are estimated using a crystallographic analog of the Gerchberg-Saxton algorithm, with "
+            "TV denoising as the real-space constraint.\n\n K-weighting can optionally be used to "
+            "weight the algorithm input. \n\n In the terminology adopted, this script computes a "
+            "`derivative` minus a `native` map, modifying the derivative phases. Native phases,"
+            "typically from a model of the `native` data, are computed from a CIF/PDB model you "
+            "must provide."
         )
     )
     args = parser.parse_args(command_line_arguments)
@@ -62,6 +68,11 @@ def main(command_line_arguments: list[str] | None = None) -> None:
 
     log.info("Writing metadata.", file=str(args.metadataout))
     final_tv_metadata.k_parameter_used = kparameter_used
+    write_combined_metadata(
+        filename=args.metadataout,
+        it_tv_metadata=it_tv_metadata,
+        final_tv_metadata=final_tv_metadata,
+    )
 
 
 if __name__ == "__main__":

test/functional/test_compute_iterative_tv.py → test/functional/test_compute_iterative_tv_map.py

@@ -6,8 +6,7 @@
 
 from meteor.rsmap import Map
 from meteor.scripts import compute_iterative_tv_map
-from meteor.scripts.common import WeightMode
-from meteor.tv import TvDenoiseResult
+from meteor.scripts.common import WeightMode, read_combined_metadata
 from meteor.utils import filter_common_indices
 
 
@@ -53,19 +52,22 @@ def test_script_produces_consistent_results(
 
     compute_iterative_tv_map.main(cli_args)
 
-    # TODO this simple load metadata won't work, load JSON then make object instead
-    result_metadata = TvDenoiseResult.from_json_file(output_metadata)
+    iterative_tv_metadata, final_tv_metadata = read_combined_metadata(filename=output_metadata)
     result_map = Map.read_mtz_file(output_mtz)
 
-    # 1. make sure negentropy increased
+    # 1. make sure the negentropy increased during iterative TV
+    negentropy_over_iterations = iterative_tv_metadata["negentropy_after_tv"]
+    assert negentropy_over_iterations[-1] > negentropy_over_iterations[0]
+
+    # 2. make sure negentropy increased in the final TV pass
     if kweight_mode == WeightMode.none and tv_weight_mode == WeightMode.none:
         np.testing.assert_allclose(
-            result_metadata.optimal_negentropy, result_metadata.initial_negentropy
+            final_tv_metadata.optimal_negentropy, final_tv_metadata.initial_negentropy
         )
     else:
-        assert result_metadata.optimal_negentropy >= result_metadata.initial_negentropy
+        assert final_tv_metadata.optimal_negentropy >= final_tv_metadata.initial_negentropy
 
-    # 2. make sure computed DF are close to those stored on disk
+    # 3. make sure computed DF are close to those stored on disk
     reference_dataset = rs.read_mtz(str(testing_mtz_file))
     reference_amplitudes = reference_dataset["F_itTV"]
 

test/unit/conftest.py

@@ -16,6 +16,7 @@
 
 NP_RNG = np.random.default_rng()
 
+
 @pytest.fixture
 def tv_denoise_result_source_data() -> dict:
     return {

test/unit/scripts/conftest.py

@@ -17,7 +17,7 @@ def diffmap_set(random_difference_map: Map) -> DiffMapSet:
 
 @pytest.fixture
 def fixed_kparameter() -> float:
-    return 0.75
+    return 0.05
 
 
 @pytest.fixture

test/unit/scripts/test_common.py

@@ -10,15 +10,15 @@
 import reciprocalspaceship as rs
 
 from meteor.rsmap import Map
-from meteor.tv import TvDenoiseResult
 from meteor.scripts.common import (
     DiffmapArgParser,
     DiffMapSet,
     WeightMode,
     kweight_diffmap_according_to_mode,
-    write_combined_metadata,
     read_combined_metadata,
+    write_combined_metadata,
 )
+from meteor.tv import TvDenoiseResult
 
 
 def mocked_read_mtz(dummy_filename: str) -> rs.DataSet:
@@ -158,13 +158,16 @@ def test_kweight_diffmap_according_to_mode(
                 mapset=diffmap_set, kweight_mode=mode, kweight_parameter=None
             )
 
+
 def test_read_write_combined_metadata(tmp_path: Path, tv_denoise_result_source_data: dict) -> None:
     filename = tmp_path / "tmp.json"
 
-    fake_ittv_metadata = pd.DataFrame([1,2,3])
+    fake_ittv_metadata = pd.DataFrame([1, 2, 3])
     fake_tv_metadata = TvDenoiseResult(**tv_denoise_result_source_data)
 
-    write_combined_metadata(filename=filename, it_tv_metadata=fake_ittv_metadata, final_tv_metadata=fake_tv_metadata)
+    write_combined_metadata(
+        filename=filename, it_tv_metadata=fake_ittv_metadata, final_tv_metadata=fake_tv_metadata
+    )
     obtained_ittv_metadata, obtained_tv_metadata = read_combined_metadata(filename=filename)
 
     pd.testing.assert_frame_equal(fake_ittv_metadata, obtained_ittv_metadata)

test/unit/scripts/test_compute_iterative_tv_map.py

@@ -2,20 +2,58 @@
 
 import argparse
 from pathlib import Path
-from typing import Any
+from typing import Any, Sequence
 from unittest import mock
 
+import numpy as np
+import pandas as pd
 import pytest
 
+from meteor.rsmap import Map
 from meteor.scripts import compute_iterative_tv_map
 from meteor.scripts.common import DiffMapSet
 from meteor.scripts.compute_iterative_tv_map import (
     IterativeTvArgParser,
 )
+from meteor.tv import TvDenoiseResult
 
 TV_WEIGHTS_TO_SCAN = [0.01, 0.05]
 
 
+def mock_compute_it_tv(
+    derivative: Map, native: Map, *, tv_weights_to_scan: list[float]
+) -> tuple[Map, pd.DataFrame]:
+    assert isinstance(derivative, Map)
+    assert isinstance(native, Map)
+    fake_map = derivative
+    fake_metadata = pd.DataFrame.from_dict(
+        {
+            "iteration": [0],
+            "tv_weight": [tv_weights_to_scan[0]],
+            "negentropy_after_tv": [0.1],
+            "average_phase_change": [0.0001],
+        }
+    )
+    return fake_map, fake_metadata
+
+
+def mock_tv_denoise_difference_map(
+    diffmap: Map, *, full_output: bool, weights_to_scan: Sequence[float] | np.ndarray | None = None
+) -> tuple[Map, TvDenoiseResult]:
+    fake_metadata = TvDenoiseResult(
+        initial_negentropy=0.001,
+        optimal_tv_weight=0.1,
+        optimal_negentropy=1.0,
+        map_sampling_used_for_tv=3,
+        tv_weights_scanned=[0.1],
+        negentropy_at_weights=[
+            1.0,
+        ],
+        k_parameter_used=None,
+    )
+    return diffmap, fake_metadata
+
+
 @pytest.fixture
 def tv_cli_arguments(base_cli_arguments: list[str]) -> list[str]:
     new_cli_arguments = [
@@ -36,7 +74,7 @@ def test_tv_diffmap_parser(parsed_tv_cli_args: argparse.Namespace) -> None:
 
 
 def test_main(diffmap_set: DiffMapSet, tmp_path: Path, fixed_kparameter: float) -> None:
-    def mock_load_difference_maps(self: Any, args: argparse.Namespace) -> DiffMapSet:
+    def mock_load_maps(self: Any, args: argparse.Namespace) -> DiffMapSet:
         return diffmap_set
 
     output_mtz_path = tmp_path / "out.mtz"
@@ -59,10 +97,19 @@ def mock_load_difference_maps(self: Any, args: argparse.Namespace) -> DiffMapSet
         *[str(weight) for weight in TV_WEIGHTS_TO_SCAN],
     ]
 
-    # TODO is very slow
-
-    fxn_to_mck = "meteor.scripts.compute_iterative_tv_map.IterativeTvArgParser.load_difference_maps"
-    with mock.patch(fxn_to_mck, mock_load_difference_maps):
+    patch1 = mock.patch(
+        "meteor.scripts.compute_iterative_tv_map.IterativeTvArgParser.load_difference_maps",
+        mock_load_maps,
+    )
+    patch2 = mock.patch(
+        "meteor.scripts.compute_iterative_tv_map.iterative_tv_phase_retrieval", mock_compute_it_tv
+    )
+    patch3 = mock.patch(
+        "meteor.scripts.compute_iterative_tv_map.tv_denoise_difference_map",
+        mock_tv_denoise_difference_map,
+    )
+
+    with patch1, patch2, patch3:
         compute_iterative_tv_map.main(cli_arguments)
 
     assert output_mtz_path.exists()

test/unit/test_tv.py

@@ -26,9 +26,6 @@ def rms_between_coefficients(map1: Map, map2: Map) -> float:
     return float(np.linalg.norm(map2_array - map1_array))
 
 
-
-
-
 def test_tv_denoise_result(tv_denoise_result_source_data: dict) -> None:
     tdr_obj = tv.TvDenoiseResult(**tv_denoise_result_source_data)
     assert tv_denoise_result_source_data == asdict(tdr_obj)