[WIP] generalize parameter scan metadata

examples/customized_meteor_diffmap.ipynb

@@ -13,7 +13,22 @@
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'pydantic'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[1], line 7\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mreciprocalspaceship\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mrs\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmatplotlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m pyplot \u001b[38;5;28;01mas\u001b[39;00m plt\n\u001b[0;32m----> 7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmeteor\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiffmaps\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m compute_difference_map, max_negentropy_kweighted_difference_map\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmeteor\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrsmap\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Map\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmeteor\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtv\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tv_denoise_difference_map\n",
+      "File \u001b[0;32m~/opt/meteor/meteor/diffmaps.py:13\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msettings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DEFAULT_KPARAMS_TO_SCAN\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m assert_isomorphous, filter_common_indices\n\u001b[0;32m---> 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvalidate\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ScalarMaximizer, map_negentropy\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmetadata\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m KparameterScanMetadata\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_difference_map\u001b[39m(derivative: Map, native: Map, \u001b[38;5;241m*\u001b[39m, check_isomorphous: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Map:\n",
+      "File \u001b[0;32m~/opt/meteor/meteor/validate.py:11\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mscipy\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01moptimize\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m minimize_scalar\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mscipy\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mstats\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m differential_entropy\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmetadata\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m EvaluatedPoint\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrsmap\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Map\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msettings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m MAP_SAMPLING\n",
+      "File \u001b[0;32m~/opt/meteor/meteor/metadata.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpydantic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseModel\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msettings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m K_PARAMETER_NAME, TV_WEIGHT_PARAMETER_NAME\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mEvaluatedPoint\u001b[39;00m(BaseModel):\n",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'pydantic'"
+     ]
+    }
+   ],
    "source": [
     "# ruff: noqa: T201\n",
     "\n",
@@ -138,12 +153,13 @@
     }
    ],
    "source": [
-    "k_weighted_diffmap, kweight_parameter = max_negentropy_kweighted_difference_map(\n",
+    "k_weighted_diffmap, kparameter_metadata = max_negentropy_kweighted_difference_map(\n",
     "    derivative_map, native_map\n",
     ")\n",
     "kewighted_negentropy = map_negentropy(k_weighted_diffmap)\n",
     "\n",
-    "print(f\"optimal k-parameter: {kweight_parameter}, negentropy: {kewighted_negentropy:.5f}\")"
+    "print(f\"optimal k-parameter: {kparameter_metadata.optimal_parameter_value}\")\n",
+    "print(f\"negentropy: {kewighted_negentropy:.5f}\")"
    ]
   },
   {

meteor/diffmaps.py

@@ -7,6 +7,7 @@
 import numpy as np
 import reciprocalspaceship as rs
 
+from .metadata import KparameterScanMetadata
 from .rsmap import Map, assert_is_map
 from .settings import DEFAULT_KPARAMS_TO_SCAN
 from .utils import assert_isomorphous, filter_common_indices
@@ -136,7 +137,7 @@ def max_negentropy_kweighted_difference_map(
     *,
     k_parameter_values_to_scan: np.ndarray | Sequence[float] = DEFAULT_KPARAMS_TO_SCAN,
     check_isomorphous: bool = True,
-) -> rs.DataSet:
+) -> tuple[rs.DataSet, KparameterScanMetadata]:
     """
     Compute k-weighted differences between native and derivative amplitudes and phases.
 
@@ -163,8 +164,8 @@ def max_negentropy_kweighted_difference_map(
     kweighted_dataset: rs.DataSet
         dataset with added columns
 
-    opt_k_parameter: float
-        optimized k-weighting parameter
+    kparameter_metadata: KparameterScanMetadata
+        metadata detailing k-weight scan
     """
     assert_is_map(derivative, require_uncertainties=True)
     assert_is_map(native, require_uncertainties=True)
@@ -191,4 +192,15 @@ def negentropy_objective(k_parameter: float) -> float:
         check_isomorphous=check_isomorphous,
     )
 
-    return kweighted_dataset, opt_k_parameter
+    unweighted_diffmap = compute_difference_map(
+        derivative, native, check_isomorphous=check_isomorphous
+    )
+
+    kparameter_metadata = KparameterScanMetadata(
+        initial_negentropy=map_negentropy(unweighted_diffmap),
+        optimal_parameter_value=float(maximizer.argument_optimum),
+        optimal_negentropy=float(maximizer.objective_maximum),
+        parameter_scan_results=maximizer.parameter_scan_results,
+    )
+
+    return kweighted_dataset, kparameter_metadata

meteor/iterative.py

@@ -3,17 +3,17 @@
 from __future__ import annotations
 
 import numpy as np
-import pandas as pd
 import reciprocalspaceship as rs
 import structlog
 
+from .metadata import TvIterationMetadata, TvScanMetadata
 from .rsmap import Map
 from .settings import (
     DEFAULT_TV_WEIGHTS_TO_SCAN_AT_EACH_ITERATION,
     ITERATIVE_TV_CONVERGENCE_TOLERANCE,
     ITERATIVE_TV_MAX_ITERATIONS,
 )
-from .tv import TvDenoiseResult, tv_denoise_difference_map
+from .tv import tv_denoise_difference_map
 from .utils import (
     CellType,
     SpacegroupType,
@@ -110,7 +110,7 @@ def _tv_denoise_complex_difference_sf(
         *,
         cell: CellType,
         spacegroup: SpacegroupType,
-    ) -> tuple[rs.DataSeries, TvDenoiseResult]:
+    ) -> tuple[rs.DataSeries, TvScanMetadata]:
         """Apply a single iteration of TV denoising to set of complex SFs, return complex SFs"""
         diffmap = Map.from_structurefactor(
             complex_difference_sf,
@@ -134,7 +134,7 @@ def _iteratively_denoise_sf_amplitudes(
         native: rs.DataSeries,
         cell: CellType,
         spacegroup: SpacegroupType,
-    ) -> tuple[rs.DataSeries, pd.DataFrame]:
+    ) -> tuple[rs.DataSeries, list[TvIterationMetadata]]:
         """
         Estimate the derivative phases using the iterative TV algorithm.
 
@@ -160,7 +160,7 @@ def _iteratively_denoise_sf_amplitudes(
         estimated_complex_derivative: rs.DataSeries
             The derivative SFs, with the same amplitudes but phases altered to minimize the TV.
 
-        metadata: pd.DataFrame
+        metadata: list[TvIterationMetadata]
             Information about the algorithm run as a function of iteration. For each step, includes:
             the tv_weight used, the negentropy (after the TV step), and the average phase change in
             degrees.
@@ -170,7 +170,7 @@ def _iteratively_denoise_sf_amplitudes(
 
         converged: bool = False
         num_iterations: int = 0
-        metadata: list[dict[str, float]] = []
+        metadata: list[TvIterationMetadata] = []
 
         # do differences with rs.DataSeries, handles missing indices
         difference: rs.DataSeries = initial_derivative - native
@@ -198,33 +198,33 @@ def _iteratively_denoise_sf_amplitudes(
             num_iterations += 1
 
             metadata.append(
-                {
-                    "iteration": num_iterations,
-                    "tv_weight": tv_metadata.optimal_tv_weight,
-                    "negentropy_after_tv": tv_metadata.optimal_negentropy,
-                    "average_phase_change": phase_change,
-                },
+                TvIterationMetadata(
+                    iteration=num_iterations,
+                    tv_weight=tv_metadata.optimal_parameter_value,
+                    negentropy_after_tv=tv_metadata.optimal_negentropy,
+                    average_phase_change=phase_change,
+                )
             )
             if self.verbose:
                 log.info(
                     f"  iteration {num_iterations:04d}",  # noqa: G004
                     phase_change=round(phase_change, 4),
                     negentropy=round(tv_metadata.optimal_negentropy, 4),
-                    tv_weight=tv_metadata.optimal_tv_weight,
+                    tv_weight=tv_metadata.optimal_parameter_value,
                 )
 
             if num_iterations > self.max_iterations:
                 break
 
-        return derivative, pd.DataFrame(metadata)
+        return derivative, metadata
 
     def __call__(
         self,
         *,
         derivative: Map,
         native: Map,
         check_isomorphous: bool = True,
-    ) -> tuple[Map, pd.DataFrame]:
+    ) -> tuple[Map, list[TvIterationMetadata]]:
         """
         Denoise by estimating new, low-TV phases for the `derivative` dataset.
 
@@ -244,7 +244,7 @@ def __call__(
         updated_derivative: Map
             The estimated derivative phases, along with the input amplitudes and input phases.
 
-        metadata: pd.DataFrame
+        metadata: list[TvIterationMetadata]
             Information about the algorithm run as a function of iteration. For each step, includes:
             the tv_weight used, the negentropy (after the TV step), and the average phase change in
             degrees.

meteor/metadata.py

@@ -0,0 +1,43 @@
+from pydantic import BaseModel
+
+from .settings import K_PARAMETER_NAME, TV_WEIGHT_PARAMETER_NAME
+
+
+class EvaluatedPoint(BaseModel):
+    parameter_value: float
+    objective_value: float
+
+
+class MaximizerScanMetadata(BaseModel):
+    parameter_name: str
+    initial_negentropy: float
+    optimal_parameter_value: float
+    optimal_negentropy: float
+    parameter_scan_results: list[EvaluatedPoint]
+
+
+class KparameterScanMetadata(MaximizerScanMetadata):
+    parameter_name: str = K_PARAMETER_NAME
+
+
+class TvScanMetadata(MaximizerScanMetadata):
+    parameter_name: str = TV_WEIGHT_PARAMETER_NAME
+    map_sampling: float
+
+
+class DiffmapMetadata(BaseModel):
+    k_parameter_optimization: KparameterScanMetadata | None
+    tv_weight_optmization: TvScanMetadata
+
+
+class TvIterationMetadata(BaseModel):
+    iteration: int
+    tv_weight: float
+    negentropy_after_tv: float
+    average_phase_change: float
+
+
+class IterativeDiffmapMetadata(BaseModel):
+    kparameter_metadata: KparameterScanMetadata | None
+    iterative_tv_iterations: list[TvIterationMetadata]
+    final_tv_pass: TvScanMetadata

meteor/scripts/common.py

@@ -3,15 +3,12 @@
 from __future__ import annotations
 
 import argparse
-import json
 from dataclasses import dataclass
 from enum import StrEnum, auto
-from io import StringIO
 from pathlib import Path
 from typing import Any
 
 import numpy as np
-import pandas as pd
 import reciprocalspaceship as rs
 import structlog
 
@@ -20,12 +17,15 @@
     compute_kweighted_difference_map,
     max_negentropy_kweighted_difference_map,
 )
+from meteor.metadata import KparameterScanMetadata
 from meteor.mtzio import find_observed_amplitude_column, find_observed_uncertainty_column
 from meteor.rsmap import Map
 from meteor.scale import scale_maps
-from meteor.settings import COMPUTED_MAP_RESOLUTION_LIMIT, KWEIGHT_PARAMETER_DEFAULT
+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
 from meteor.utils import cut_resolution
 
 log = structlog.get_logger()
@@ -283,7 +283,7 @@ def load_difference_maps(args: argparse.Namespace) -> DiffMapSet:
 
 def kweight_diffmap_according_to_mode(
     *, mapset: DiffMapSet, kweight_mode: WeightMode, kweight_parameter: float | None = None
-) -> tuple[Map, float | None]:
+) -> tuple[Map, KparameterScanMetadata | None]:
     """
     Make and k-weight a difference map using a specified `WeightMode`.
 
@@ -311,16 +311,17 @@ def kweight_diffmap_according_to_mode(
     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.
+    kparameter_metadata: KparameterScanMetadata | None
+        The information about the k_weight optimization 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(
+        diffmap, kparameter_metadata = max_negentropy_kweighted_difference_map(
             mapset.derivative, mapset.native
         )
-        log.info("  using negentropy optimized", kparameter=kweight_parameter)
+        log.info("  using negentropy max.", kparameter=kparameter_metadata.optimal_parameter_value)
         if kweight_parameter is np.nan:
             msg = "determined `k-parameter` is NaN, something went wrong..."
             raise RuntimeError(msg)
@@ -333,34 +334,15 @@ def kweight_diffmap_according_to_mode(
         diffmap = compute_kweighted_difference_map(
             mapset.derivative, mapset.native, k_parameter=kweight_parameter
         )
-
+        kparameter_metadata = None
         log.info("  using fixed", kparameter=kweight_parameter)
 
     elif kweight_mode == WeightMode.none:
         diffmap = compute_difference_map(mapset.derivative, mapset.native)
-        kweight_parameter = None
+        kparameter_metadata = None
         log.info(" requested no k-weighting")
 
     else:
         raise InvalidWeightModeError(kweight_mode)
 
-    return diffmap, kweight_parameter
-
-
-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(orient="records", indent=4),
-        "final_tv_pass": final_tv_metadata.json(),
-    }
-    with filename.open("w") as f:
-        json.dump(combined_metadata, f, indent=4)
-
-
-def read_combined_metadata(*, filename: Path) -> tuple[pd.DataFrame, TvDenoiseResult]:
-    with filename.open("r") as f:
-        combined_metadata = json.load(f)
-    it_tv_metadata = pd.read_json(StringIO(combined_metadata["iterative_tv"]))
-    final_tv_metadata = TvDenoiseResult.from_json(combined_metadata["final_tv_pass"])
-    return it_tv_metadata, final_tv_metadata
+    return diffmap, kparameter_metadata