works

rs-station · Aug 21, 2024 · e3cef77 · e3cef77
1 parent 4542b00
commit e3cef77
Show file tree

Hide file tree

Showing 5 changed files with 147 additions and 27 deletions.
diff --git a/meteor/settings.py b/meteor/settings.py
@@ -2,5 +2,3 @@
 TV_STOP_TOLERANCE: float = 0.00000005
 TV_MAX_NUM_ITER: int = 50
 TV_MAP_SAMPLING: int = 3
-TV_AMPLITUDE_LABEL: str = "dFtv"
-TV_PHASE_LABEL: str = "dPHItv"
diff --git a/meteor/tv.py b/meteor/tv.py
@@ -19,8 +19,6 @@
     TV_STOP_TOLERANCE,
     TV_MAP_SAMPLING,
     TV_MAX_NUM_ITER,
-    TV_AMPLITUDE_LABEL,
-    TV_PHASE_LABEL,
 )
 
 
@@ -43,13 +41,7 @@ def tv_denoise_difference_map(
     difference_map_phase_column: str = "PHIC",
     lambda_values_to_scan: Sequence[float] | None = None,
 ) -> tuple[rs.DataSet, float]:
-    """
-
-    lambda_values_to_scan = None --> Golden method
-
-    Returns:
-       rs.Dataset: denoised dataset with new columns `DFtv`, `DPHItv`
-    """
+    """ lambda_values_to_scan = None --> Golden method """
 
     # TODO write decent docstring
 
@@ -63,11 +55,13 @@ def tv_denoise_difference_map(
     difference_map_as_array = np.array(difference_map.grid)
 
     def negentropy_objective(tv_lambda: float):
-        denoised_map = _tv_denoise_array(map_as_array=difference_map_as_array, weight=tv_lambda)
+        denoised_map = _tv_denoise_array(
+            map_as_array=difference_map_as_array, weight=tv_lambda
+        )
         return negentropy(denoised_map.flatten())
 
     optimal_lambda: float
-    if lambda_values_to_scan:
+    if lambda_values_to_scan is not None:
         highest_negentropy = -1e8
         for tv_lambda in lambda_values_to_scan:
             trial_negentropy = negentropy_objective(tv_lambda)
@@ -83,19 +77,21 @@ def negentropy_objective(tv_lambda: float):
         ), "Golden minimization failed to find optimal TV lambda"
         optimal_lambda = optimizer_result.x
 
-    final_map = _tv_denoise_array(map_as_array=difference_map_as_array, weight=optimal_lambda)
+    final_map = _tv_denoise_array(
+        map_as_array=difference_map_as_array, weight=optimal_lambda
+    )
     final_map = numpy_array_to_map(
         final_map,
         spacegroup=difference_map_coefficients.spacegroup,
-        cell=difference_map_coefficients.cell
+        cell=difference_map_coefficients.cell,
     )
 
     _, dmin = resolution_limits(difference_map_coefficients)
     final_map_coefficients = compute_coefficients_from_map(
         ccp4_map=final_map,
         high_resolution_limit=dmin,
-        amplitude_label=TV_AMPLITUDE_LABEL,
-        phase_label=TV_PHASE_LABEL,
+        amplitude_label=difference_map_amplitude_column,
+        phase_label=difference_map_phase_column,
     )
 
     return final_map_coefficients

diff --git a/meteor/utils.py b/meteor/utils.py
@@ -64,11 +64,24 @@ def canonicalize_amplitudes(
         return None
 
 
-def numpy_array_to_map(array: np.ndarray, *, spacegroup: str | int, cell: tuple[float, float, float, float, float, float]) -> gemmi.Ccp4Map:
+def numpy_array_to_map(
+    array: np.ndarray,
+    *,
+    spacegroup: str | int | gemmi.SpaceGroup,
+    cell: tuple[float, float, float, float, float, float] | gemmi.UnitCell,
+) -> gemmi.Ccp4Map:
     ccp4_map = gemmi.Ccp4Map()
-    ccp4_map.grid = gemmi.FloatGrid(array, dtype=array.dtype)
-    ccp4_map.grid.unit_cell.set(*cell)
-    ccp4_map.grid.spacegroup = gemmi.SpaceGroup(spacegroup)
+    ccp4_map.grid = gemmi.FloatGrid(array.astype(np.float32))
+
+    if isinstance(cell, gemmi.UnitCell):
+        ccp4_map.grid.unit_cell = cell
+    else:
+        ccp4_map.grid.unit_cell.set(*cell)
+
+    if not isinstance(spacegroup, gemmi.SpaceGroup):
+        spacegroup = gemmi.SpaceGroup(spacegroup)
+    ccp4_map.grid.spacegroup = spacegroup
+
     return ccp4_map
 
 
@@ -97,7 +110,7 @@ def compute_coefficients_from_map(
     phase_label: str,
 ) -> rs.DataSet:
     # to ensure we include the final shell of reflections, add a small buffer to the resolution
-    
+
     high_resolution_buffer = 1e-8
     gemmi_structure_factors = gemmi.transform_map_to_f_phi(ccp4_map.grid, half_l=False)
     data = gemmi_structure_factors.prepare_asu_data(

diff --git a/test/unit/test_tv.py b/test/unit/test_tv.py
@@ -1,5 +1,106 @@
 import reciprocalspaceship as rs
+from meteor.utils import compute_coefficients_from_map, compute_map_from_coefficients
 from meteor import tv
+import gemmi
+import numpy as np
+from pytest import fixture
+
+# TODO make these universal in the tests
+TEST_AMPLITUDE_LABEL = "DF"
+TEST_PHASE_LABEL = "PHIC"
+
+
+def _generate_single_carbon_density(
+    carbon_position: tuple[float, float, float],
+    space_group: gemmi.SpaceGroup,
+    unit_cell: gemmi.UnitCell,
+    d_min: float,
+) -> gemmi.FloatGrid:
+    # Create a model with a single carbon atom
+    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) -> float:
+    map1 = compute_map_from_coefficients(
+        map_coefficients=ds1,
+        amplitude_label=TEST_AMPLITUDE_LABEL,
+        phase_label=TEST_PHASE_LABEL,
+        map_sampling=3
+    )
+    map2 = compute_map_from_coefficients(
+        map_coefficients=ds2,
+        amplitude_label=TEST_AMPLITUDE_LABEL,
+        phase_label=TEST_PHASE_LABEL,
+        map_sampling=3
+    )
+    difference_map = np.array(map2.grid) - np.array(map1.grid)
+    rms = float(np.linalg.norm(difference_map))
+    return rms
+
+
+@fixture
+def noise_free_map() -> rs.DataSet:
+    return displaced_single_atom_difference_map_coefficients(noise_sigma=0.0)
+
+
+@fixture
+def noisy_map() -> rs.DataSet:
+    return displaced_single_atom_difference_map_coefficients(noise_sigma=3.0)
 
 
 def test_tv_denoise_difference_map_smoke(flat_difference_map: rs.DataSet) -> None:
@@ -8,15 +109,27 @@ def test_tv_denoise_difference_map_smoke(flat_difference_map: rs.DataSet) -> Non
         difference_map_coefficients=flat_difference_map,
         lambda_values_to_scan=[1.0, 2.0],
     )
-
     # test golden optimizer
-    tv.TV_LAMBDA_RANGE = (1.0, 1.01)
+    tv.TV_LAMBDA_RANGE = (1.0, 2.0)
     tv.tv_denoise_difference_map(
         difference_map_coefficients=flat_difference_map,
     )
 
 
-def test_tv_denoise_difference_map_golden(): ...
+def test_tv_denoise_difference_map_golden(noise_free_map: rs.DataSet, noisy_map: rs.DataSet) -> None:
+    rms_before_denoising = rms_between_coefficients(noise_free_map, noisy_map)
+    denoised_map = tv.tv_denoise_difference_map(
+        difference_map_coefficients=noisy_map,
+    )
+    rms_after_denoising = rms_between_coefficients(noise_free_map, denoised_map)
+    assert rms_after_denoising < rms_before_denoising
 
 
-def test_tv_denoise_difference_map_specific_lambdas(): ...
+def test_tv_denoise_difference_map_specific_lambdas(noise_free_map: rs.DataSet, noisy_map: rs.DataSet) -> None:
+    rms_before_denoising = rms_between_coefficients(noise_free_map, noisy_map)
+    denoised_map = tv.tv_denoise_difference_map(
+        difference_map_coefficients=noisy_map,
+        lambda_values_to_scan=np.logspace(-3, 0, 25),
+    )
+    rms_after_denoising = rms_between_coefficients(noise_free_map, denoised_map)
+    assert rms_after_denoising < rms_before_denoising
diff --git a/test/unit/test_utils.py b/test/unit/test_utils.py
@@ -86,7 +86,7 @@ def test_compute_map_from_coefficients(flat_difference_map: rs.DataSet) -> None:
         map_sampling=1,
     )
     assert isinstance(map, gemmi.Ccp4Map)
-    assert map.grid.shape == (6,6,6)
+    assert map.grid.shape == (6, 6, 6)
 
 
 @pytest.mark.parametrize("map_sampling", [1, 2, 2.25, 3, 5])