phaseboost in, need to work on tests

meteor/scripts/common.py

@@ -1,13 +1,21 @@
+from __future__ import annotations
+
 import argparse
 import re
 from dataclasses import dataclass
 from enum import StrEnum, auto
 from pathlib import Path
 from typing import Any
 
+import numpy as np
 import reciprocalspaceship as rs
 import structlog
 
+from meteor.diffmaps import (
+    compute_difference_map,
+    compute_kweighted_difference_map,
+    max_negentropy_kweighted_difference_map,
+)
 from meteor.io import find_observed_amplitude_column, find_observed_uncertainty_column
 from meteor.rsmap import Map
 from meteor.scale import scale_maps
@@ -235,3 +243,69 @@ def load_difference_maps(args: argparse.Namespace) -> DiffMapSet:
 
         mapset.scale()
         return mapset
+
+
+def kweight_diffmap_according_to_mode(
+    *, mapset: DiffMapSet, kweight_mode: WeightMode, kweight_parameter: float | None = None
+) -> tuple[Map, float | None]:
+    """
+    Make and k-weight a difference map using a specified `WeightMode`.
+
+    Three modes are possible to pick the k-parameter:
+      * `WeightMode.optimize`, max-negentropy value will and picked, this may take some time
+      * `WeightMode.fixed`, `kweight_parameter` is used
+      * `WeightMode.none`, then no k-weighting is done (note this is NOT equivalent to
+         kweight_parameter=0.0)
+
+    Parameters
+    ----------
+    mapset: DiffMapSet
+        The set of `derivative`, `native`, `computed` maps to use to compute the diffmap.
+
+    kweight_mode: WeightMode
+        How to set the k-parameter: {optimize, fixed, none}. See above. If `fixed`, then
+        `kweight_parameter` is required.
+
+    kweight_parameter: float | None
+        If kweight_mode == WeightMode.fixed, then this must be a float that specifies the
+        k-parameter to use.
+
+    Returns
+    -------
+    diffmap: meteor.rsmap.Map
+        The difference map, k-weighted if requested.
+
+    kweight_parameter: float | None
+        The `kweight_parameter` used. Only really interesting if WeightMode.optimize.
+    """
+    log.info("Computing difference map.")
+
+    if kweight_mode == WeightMode.optimize:
+        diffmap, kweight_parameter = max_negentropy_kweighted_difference_map(
+            mapset.derivative, mapset.native
+        )
+        log.info("  using negentropy optimized", kparameter=kweight_parameter)
+        if kweight_parameter is np.nan:
+            msg = "determined `k-parameter` is NaN, something went wrong..."
+            raise RuntimeError(msg)
+
+    elif kweight_mode == WeightMode.fixed:
+        if not isinstance(kweight_parameter, float):
+            msg = f"`kweight_parameter` is type `{type(kweight_parameter)}`, must be `float`"
+            raise TypeError(msg)
+
+        diffmap = compute_kweighted_difference_map(
+            mapset.derivative, mapset.native, k_parameter=kweight_parameter
+        )
+
+        log.info("  using fixed", kparameter=kweight_parameter)
+
+    elif kweight_mode == WeightMode.none:
+        diffmap = compute_difference_map(mapset.derivative, mapset.native)
+        kweight_parameter = None
+        log.info(" requested no k-weighting")
+
+    else:
+        raise InvalidWeightModeError(kweight_mode)
+
+    return diffmap, kweight_parameter

meteor/scripts/compute_difference_map.py

@@ -5,17 +5,17 @@
 import numpy as np
 import structlog
 
-from meteor.diffmaps import (
-    compute_difference_map,
-    compute_kweighted_difference_map,
-    max_negentropy_kweighted_difference_map,
-)
 from meteor.rsmap import Map
 from meteor.settings import MAP_SAMPLING, TV_WEIGHT_DEFAULT
 from meteor.tv import TvDenoiseResult, tv_denoise_difference_map
 from meteor.validate import negentropy
 
-from .common import DiffmapArgParser, DiffMapSet, InvalidWeightModeError, WeightMode
+from .common import (
+    DiffmapArgParser,
+    InvalidWeightModeError,
+    WeightMode,
+    kweight_diffmap_according_to_mode,
+)
 
 log = structlog.get_logger()
 
@@ -46,72 +46,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         )
 
 
-def kweight_diffmap_according_to_mode(
-    *, mapset: DiffMapSet, kweight_mode: WeightMode, kweight_parameter: float | None = None
-) -> tuple[Map, float | None]:
-    """
-    Make and k-weight a difference map using a specified `WeightMode`.
-
-    Three modes are possible to pick the k-parameter:
-      * `WeightMode.optimize`, max-negentropy value will and picked, this may take some time
-      * `WeightMode.fixed`, `kweight_parameter` is used
-      * `WeightMode.none`, then no k-weighting is done (note this is NOT equivalent to
-         kweight_parameter=0.0)
-
-    Parameters
-    ----------
-    mapset: DiffMapSet
-        The set of `derivative`, `native`, `computed` maps to use to compute the diffmap.
-
-    kweight_mode: WeightMode
-        How to set the k-parameter: {optimize, fixed, none}. See above. If `fixed`, then
-        `kweight_parameter` is required.
-
-    kweight_parameter: float | None
-        If kweight_mode == WeightMode.fixed, then this must be a float that specifies the
-        k-parameter to use.
-
-    Returns
-    -------
-    diffmap: meteor.rsmap.Map
-        The difference map, k-weighted if requested.
-
-    kweight_parameter: float | None
-        The `kweight_parameter` used. Only really interesting if WeightMode.optimize.
-    """
-    log.info("Computing difference map.")
-
-    if kweight_mode == WeightMode.optimize:
-        diffmap, kweight_parameter = max_negentropy_kweighted_difference_map(
-            mapset.derivative, mapset.native
-        )
-        log.info("  using negentropy optimized", kparameter=kweight_parameter)
-        if kweight_parameter is np.nan:
-            msg = "determined `k-parameter` is NaN, something went wrong..."
-            raise RuntimeError(msg)
-
-    elif kweight_mode == WeightMode.fixed:
-        if not isinstance(kweight_parameter, float):
-            msg = f"`kweight_parameter` is type `{type(kweight_parameter)}`, must be `float`"
-            raise TypeError(msg)
-
-        diffmap = compute_kweighted_difference_map(
-            mapset.derivative, mapset.native, k_parameter=kweight_parameter
-        )
-
-        log.info("  using fixed", kparameter=kweight_parameter)
-
-    elif kweight_mode == WeightMode.none:
-        diffmap = compute_difference_map(mapset.derivative, mapset.native)
-        kweight_parameter = None
-        log.info(" requested no k-weighting")
-
-    else:
-        raise InvalidWeightModeError(kweight_mode)
-
-    return diffmap, kweight_parameter
-
-
 def denoise_diffmap_according_to_mode(
     *,
     diffmap: Map,
@@ -148,9 +82,7 @@ def denoise_diffmap_according_to_mode(
         Information regarding the denoising process.
     """
     if tv_denoise_mode == WeightMode.optimize:
-        log.info(
-            "Searching for max-negentropy TV denoising weight", method="golden-section search"
-        )
+        log.info("Searching for max-negentropy TV denoising weight", method="golden-section search")
         log.info("This may take some time...")
 
         final_map, metadata = tv_denoise_difference_map(diffmap, full_output=True)

meteor/scripts/compute_iterative_tv_map.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import Any
+
+import structlog
+
+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
+
+log = structlog.get_logger()
+
+
+# TODO: test this
+TV_WEIGHTS_TO_SCAN_DEFAULT = [0.01]
+
+
+class IterativeTvArgParser(DiffmapArgParser):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.add_argument(
+            "-x",
+            "--tv-weights-to-scan",
+            nargs="+",
+            type=float,
+            default=TV_WEIGHTS_TO_SCAN_DEFAULT,
+            help=(
+                "Choose what TV weights to evaluate at every iteration. Can be a single float."
+                f"Default: {TV_WEIGHTS_TO_SCAN_DEFAULT}."
+            ),
+        )
+
+
+def main(command_line_arguments: list[str] | None = None) -> None:
+    parser = IterativeTvArgParser(
+        description=(
+            "bla bla"  # TODO
+        )
+    )
+    args = parser.parse_args(command_line_arguments)
+    parser.check_output_filepaths(args)
+    mapset = parser.load_difference_maps(args)
+
+    # First, find improved derivative phases
+    new_derivative_map, it_tv_metadata = iterative_tv_phase_retrieval(
+        mapset.derivative,
+        mapset.native,
+        tv_weights_to_scan=args.tv_weights_to_scan,
+    )
+    mapset.derivative = new_derivative_map
+
+    diffmap, kparameter_used = kweight_diffmap_according_to_mode(
+        kweight_mode=args.kweight_mode, kweight_parameter=args.kweight_parameter, mapset=mapset
+    )
+
+    # TODO: used fixed weight or golden method?
+    final_map, final_tv_metadata = tv_denoise_difference_map(
+        diffmap, full_output=True, weights_to_scan=args.tv_weights_to_scan
+    )
+
+    log.info("Writing output.", file=str(args.mtzout))
+    final_map.write_mtz(args.mtzout)
+
+    # TODO: append it_tv_metadata
+    # log.info("Writing metadata.", file=str(args.metadataout))
+    # metadata.k_parameter_used = kparameter_used
+    #
+
+
+if __name__ == "__main__":
+    main()

test/functional/test_compute_iterative_tv.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+import reciprocalspaceship as rs
+
+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.utils import filter_common_indices
+
+
+@pytest.mark.parametrize("kweight_mode", list(WeightMode))
+@pytest.mark.parametrize("tv_weight_mode", list(WeightMode))
+def test_script_produces_consistent_results(
+    kweight_mode: WeightMode,
+    tv_weight_mode: WeightMode,
+    testing_pdb_file: Path,
+    testing_mtz_file: Path,
+    tmp_path: Path,
+) -> None:
+    # for when WeightMode.fixed; these maximize negentropy in manual testing
+    kweight_parameter = 0.05
+
+    output_mtz = tmp_path / "test-output.mtz"
+    output_metadata = tmp_path / "test-output-metadata.csv"
+
+    cli_args = [
+        str(testing_mtz_file),  # derivative
+        "--derivative-amplitude-column",
+        "F_on",
+        "--derivative-uncertainty-column",
+        "SIGF_on",
+        str(testing_mtz_file),  # native
+        "--native-amplitude-column",
+        "F_off",
+        "--native-uncertainty-column",
+        "SIGF_off",
+        "--structure",
+        str(testing_pdb_file),
+        "-o",
+        str(output_mtz),
+        "-m",
+        str(output_metadata),
+        "--kweight-mode",
+        kweight_mode,
+        "--kweight-parameter",
+        str(kweight_parameter),
+        "-x",
+        "0.01",
+    ]
+
+    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)
+    result_map = Map.read_mtz_file(output_mtz)
+
+    # 1. make sure negentropy increased
+    if kweight_mode == WeightMode.none and tv_weight_mode == WeightMode.none:
+        np.testing.assert_allclose(
+            result_metadata.optimal_negentropy, result_metadata.initial_negentropy
+        )
+    else:
+        assert result_metadata.optimal_negentropy >= result_metadata.initial_negentropy
+
+    # 2. 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"]
+
+    result_amplitudes, reference_amplitudes = filter_common_indices(
+        result_map.amplitudes, reference_amplitudes
+    )
+    rho = np.corrcoef(result_amplitudes.to_numpy(), reference_amplitudes.to_numpy())[0, 1]
+
+    # comparing a correlation coefficienct allows for a global scale factor change, but nothing else
+    if (kweight_mode == WeightMode.none) or (tv_weight_mode == WeightMode.none):  # noqa: PLR1714
+        assert rho > 0.50
+    else:
+        assert rho > 0.98

test/unit/scripts/test_common.py

@@ -10,7 +10,12 @@
 import reciprocalspaceship as rs
 
 from meteor.rsmap import Map
-from meteor.scripts.common import DiffmapArgParser, DiffMapSet, WeightMode
+from meteor.scripts.common import (
+    DiffmapArgParser,
+    DiffMapSet,
+    WeightMode,
+    kweight_diffmap_according_to_mode,
+)
 
 
 def mocked_read_mtz(dummy_filename: str) -> rs.DataSet:
@@ -129,3 +134,23 @@ def return_a_map(*args: Any, **kwargs: Any) -> Map:
         assert isinstance(mapset.native, Map)
         assert isinstance(mapset.derivative, Map)
         assert isinstance(mapset.calculated, Map)
+
+
+@pytest.mark.parametrize("mode", list(WeightMode))
+def test_kweight_diffmap_according_to_mode(
+    mode: WeightMode, diffmap_set: DiffMapSet, fixed_kparameter: float
+) -> None:
+    # ensure the two maps aren't exactly the same to prevent numerical issues
+    diffmap_set.derivative.amplitudes.iloc[0] += 1.0
+
+    diffmap, _ = kweight_diffmap_according_to_mode(
+        mapset=diffmap_set, kweight_mode=mode, kweight_parameter=fixed_kparameter
+    )
+    assert len(diffmap) > 0
+    assert isinstance(diffmap, Map)
+
+    if mode == WeightMode.fixed:
+        with pytest.raises(TypeError):
+            _ = kweight_diffmap_according_to_mode(
+                mapset=diffmap_set, kweight_mode=mode, kweight_parameter=None
+            )