diff --git a/meteor/scripts/common.py b/meteor/scripts/common.py new file mode 100644 index 0000000..3edd0a2 --- /dev/null +++ b/meteor/scripts/common.py @@ -0,0 +1,132 @@ +import argparse +from dataclasses import dataclass + +import structlog + +from meteor.rsmap import Map +from meteor.scale import scale_maps + +log = structlog.get_logger() + + +# TODO: better name! +@dataclass +class MapSet: + native: Map + derivative: Map + calculated_native: Map + + +class DiffmapArgParser(argparse.ArgumentParser): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + self.add_argument( + "--native_mtz", + nargs=4, + metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), + required=True, + help=("Native MTZ file and associated amplitude, uncertainty labels, and phase label."), + ) + + self.add_argument( + "--derivative_mtz", + nargs=4, + metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), + required=True, + help=( + "Derivative MTZ file and associated amplitude, uncertainty labels, and phase label." + ), + ) + + self.add_argument( + "--calc_native_mtz", + nargs=3, + metavar=("filename", "calc_amplitude_label", "calc_phase_label"), + required=True, + help=( + "Calculated native MTZ file and associated calculated amplitude and phase labels." + ), + ) + + self.add_argument( + "--output", + type=str, + default="meteor_difference_map.mtz", + help="Output file name", + ) + + self.add_argument( + "--use_uncertainties_to_scale", + type=bool, + default=True, + help="Use uncertainties to scale (default: True)", + ) + + k_weight_group = self.add_mutually_exclusive_group() + + k_weight_group.add_argument( + "--k_weight_with_fixed_parameter", + type=float, + default=None, + help="Use k-weighting with a fixed parameter (float between 0 and 1.0)", + ) + + k_weight_group.add_argument( + "--k_weight_with_parameter_optimization", + action="store_true", # This will set the flag to True when specified + help="Use k-weighting with parameter optimization (default: False)", + ) + + +def _log_map_read(map_name: str, args_obj) -> None: + phases = args_obj[2] if len(args_obj) == 3 else None + log.info("Read map", name=map_name, amps=args_obj[1], stds=args_obj[2], phases=phases) + + +def load_and_scale_mapset(args) -> MapSet: + # Create the native map from the native MTZ file + native_map = Map.read_mtz_file( + args.native_mtz[0], + amplitude_column=args.native_mtz[1], + uncertainty_column=args.native_mtz[2], + phase_column=args.native_mtz[3], + ) + _log_map_read("native", args.native_mtz) + + # Create the derivative map from the derivative MTZ file + derivative_map = Map.read_mtz_file( + args.derivative_mtz[0], + amplitude_column=args.derivative_mtz[1], + uncertainty_column=args.derivative_mtz[2], + phase_column=args.derivative_mtz[3], + ) + _log_map_read("derivative", args.derivative_mtz) + + # Create the calculated native map from the calculated native MTZ file + calc_native_map = Map.read_mtz_file( + args.calc_native_mtz[0], + amplitude_column=args.calc_native_mtz[1], + phase_column=args.calc_native_mtz[2], + ) + _log_map_read("calculated native", args.calc_native_mtz) + + # Scale both to common map calculated from native model + native_map_scaled = scale_maps( + reference_map=calc_native_map, + map_to_scale=native_map, + weight_using_uncertainties=False, + ) + log.info("scaling: native --> calculated native") + derivative_map_scaled = scale_maps( + reference_map=calc_native_map, + map_to_scale=derivative_map, + weight_using_uncertainties=False, # FC do not have uncertainties + ) + log.info("scaling: derivative --> calculated native") + + return MapSet( + native=native_map_scaled, + derivative=derivative_map_scaled, + calculated_native=calc_native_map, + ) diff --git a/meteor/scripts/compute_difference_map.py b/meteor/scripts/compute_difference_map.py index 1443430..d276db0 100644 --- a/meteor/scripts/compute_difference_map.py +++ b/meteor/scripts/compute_difference_map.py @@ -1,155 +1,38 @@ -import reciprocalspaceship as rs -from meteor.rsmap import Map, _assert_is_map +import structlog + from meteor.diffmaps import ( compute_difference_map, compute_kweighted_difference_map, max_negentropy_kweighted_difference_map, ) -from meteor.scale import scale_maps -import argparse +from .common import DiffmapArgParser, load_and_scale_mapset + +log = structlog.get_logger() def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser( + parser = DiffmapArgParser( description="Compute a difference map using native, derivative, and calculated MTZ files." ) - - parser.add_argument( - "--native_mtz", - nargs=4, - metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), - required=True, - help=( - "Native MTZ file and associated amplitude, uncertainty labels, and phase label." - ), - ) - - parser.add_argument( - "--derivative_mtz", - nargs=4, - metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), - required=True, - help=( - "Derivative MTZ file and associated amplitude, uncertainty labels, and phase label." - ), - ) - - parser.add_argument( - "--calc_native_mtz", - nargs=3, - metavar=("filename", "calc_amplitude_label", "calc_phase_label"), - required=True, - help=( - "Calculated native MTZ file and associated calculated amplitude and phase labels." - ), - ) - - parser.add_argument( - "--output", - type=str, - default="meteor_difference_map.mtz", - help="Output file name", - ) - - parser.add_argument( - "--use_uncertainties_to_scale", - type=bool, - default=True, - help="Use uncertainties to scale (default: True)", - ) - - k_weight_group = parser.add_mutually_exclusive_group() - - k_weight_group.add_argument( - "--k_weight_with_fixed_parameter", - type=float, - default=None, - help="Use k-weighting with a fixed parameter (float between 0 and 1.0)", - ) - - k_weight_group.add_argument( - "--k_weight_with_parameter_optimization", - action="store_true", # This will set the flag to True when specified - help="Use k-weighting with parameter optimization (default: False)", - ) - return parser.parse_args() -def compute_difference_map_wrapper( - derivative_map, - native_map, - use_k_weighting=False, - k_param=None, - optimize_k=False, -): - """Wrapper for computing either a standard difference map or a k-weighted difference map.""" - if use_k_weighting: - if optimize_k: - result, opt_k = max_negentropy_kweighted_difference_map( - derivative_map, native_map - ) - print(f"Optimal k-parameter determined: {opt_k}") # noqa: T201 - return result - return compute_kweighted_difference_map( - derivative_map, native_map, k_parameter=k_param - ) - return compute_difference_map(derivative_map, native_map) - - def main(): - # Parse command-line arguments args = parse_args() + mapset = load_and_scale_mapset(args) + + if args.k_weight_with_parameter_optimization: + diffmap, opt_k = max_negentropy_kweighted_difference_map(mapset.derivative, mapset.native) + log.info("Optimal k-parameter determined:", value=opt_k) + elif args.k_weight_with_fixed_parameter: + diffmap = compute_kweighted_difference_map( + mapset.derivative, mapset.native, k_parameter=args.k_weight_with_fixed_parameter + ) + else: + diffmap = compute_difference_map(mapset.derivative, mapset.native) - # Create the native map from the native MTZ file - native_map = Map.read_mtz_file( - args.native_mtz[0], - amplitude_column=args.native_mtz[1], - uncertainty_column=args.native_mtz[2], - phase_column=args.native_mtz[3], - ) - - # Create the derivative map from the derivative MTZ file - derivative_map = Map.read_mtz_file( - args.derivative_mtz[0], - amplitude_column=args.derivative_mtz[1], - uncertainty_column=args.derivative_mtz[2], - phase_column=args.derivative_mtz[3], - ) - - # Create the calculated native map from the calculated native MTZ file - calc_native_map = Map.read_mtz_file( - args.calc_native_mtz[0], - amplitude_column=args.calc_native_mtz[1], - phase_column=args.calc_native_mtz[2], - ) - - # Scale both to common map calculated from native model - native_map_scaled = scale_maps( - reference_map=calc_native_map, - map_to_scale=native_map, - weight_using_uncertainties=False, - ) - - derivative_map_scaled = scale_maps( - reference_map=calc_native_map, - map_to_scale=derivative_map, - weight_using_uncertainties=False, # FC do not have uncertainties - ) - - # Calculate diffmap - diffmap = compute_difference_map_wrapper( - derivative_map_scaled, - native_map_scaled, - use_k_weighting=args.k_weight_with_fixed_parameter - or args.k_weight_with_parameter_optimization, - k_param=args.k_weight_with_fixed_parameter, - optimize_k=args.k_weight_with_parameter_optimization, - ) - - print("Writing output file...") # noqa: T201 + print("Writing output", file=args.output) diffmap.write_mtz(args.output) diff --git a/meteor/scripts/compute_tvdenoised_difference_map.py b/meteor/scripts/compute_tvdenoised_difference_map.py index 50d94ba..69fe496 100644 --- a/meteor/scripts/compute_tvdenoised_difference_map.py +++ b/meteor/scripts/compute_tvdenoised_difference_map.py @@ -1,5 +1,6 @@ -import reciprocalspaceship as rs -from meteor.rsmap import Map, _assert_is_map +import numpy as np +import structlog + from meteor.diffmaps import ( compute_difference_map, compute_kweighted_difference_map, @@ -7,162 +8,45 @@ ) from meteor.tv import tv_denoise_difference_map -from meteor.scale import scale_maps -import argparse -import numpy as np +from .common import DiffmapArgParser, load_and_scale_mapset + +log = structlog.get_logger() def parse_args(): """Parse command line arguments.""" - parser = argparse.ArgumentParser( + parser = DiffmapArgParser( description="Compute a difference map using native, derivative, and calculated MTZ files." ) - - parser.add_argument( - "--native_mtz", - nargs=4, - metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), - required=True, - help=( - "Native MTZ file and associated amplitude, uncertainty labels, and phase label." - ), - ) - - parser.add_argument( - "--derivative_mtz", - nargs=4, - metavar=("filename", "amplitude_label", "uncertainty_label", "phase_label"), - required=True, - help=( - "Derivative MTZ file and associated amplitude, uncertainty labels, and phase label." - ), - ) - - parser.add_argument( - "--calc_native_mtz", - nargs=3, - metavar=("filename", "calc_amplitude_label", "calc_phase_label"), - required=True, - help=( - "Calculated native MTZ file and associated calculated amplitude and phase labels." - ), - ) - - parser.add_argument( - "--output", - type=str, - default="meteor_difference_map_TVdenoised.mtz", - help="Output file name", - ) - - parser.add_argument( - "--use_uncertainties_to_scale", - type=bool, - default=True, - help="Use uncertainties to scale (default: True)", - ) - - k_weight_group = parser.add_mutually_exclusive_group() - - k_weight_group.add_argument( - "--k_weight_with_fixed_parameter", - type=float, - default=None, - help="Use k-weighting with a fixed parameter (float between 0 and 1.0)", - ) - - k_weight_group.add_argument( - "--k_weight_with_parameter_optimization", - action="store_true", # This will set the flag to True when specified - help="Use k-weighting with parameter optimization (default: False)", - ) - - # TODO add type of lambda screen input + # TODO: add type of lambda screen input return parser.parse_args() -def compute_difference_map_wrapper( - derivative_map, - native_map, - use_k_weighting=False, - k_param=None, - optimize_k=False, -): - """Wrapper for computing either a standard difference map or a k-weighted difference map.""" - if use_k_weighting: - if optimize_k: - result, opt_k = max_negentropy_kweighted_difference_map( - derivative_map, native_map - ) - print(f"Optimal k-parameter determined: {opt_k}") # noqa: T201 - return result - return compute_kweighted_difference_map( - derivative_map, native_map, k_parameter=k_param - ) - return compute_difference_map(derivative_map, native_map) - - def main(): - # Parse command-line arguments args = parse_args() - - # Create the native map from the native MTZ file - native_map = Map.read_mtz_file( - args.native_mtz[0], - amplitude_column=args.native_mtz[1], - uncertainty_column=args.native_mtz[2], - phase_column=args.native_mtz[3], - ) - - # Create the derivative map from the derivative MTZ file - derivative_map = Map.read_mtz_file( - args.derivative_mtz[0], - amplitude_column=args.derivative_mtz[1], - uncertainty_column=args.derivative_mtz[2], - phase_column=args.derivative_mtz[3], - ) - - # Create the calculated native map from the calculated native MTZ file - calc_native_map = Map.read_mtz_file( - args.calc_native_mtz[0], - amplitude_column=args.calc_native_mtz[1], - phase_column=args.calc_native_mtz[2], - ) - - # Scale both to common map calculated from native model - native_map_scaled = scale_maps( - reference_map=calc_native_map, - map_to_scale=native_map, - weight_using_uncertainties=False, - ) - - derivative_map_scaled = scale_maps( - reference_map=calc_native_map, - map_to_scale=derivative_map, - weight_using_uncertainties=False, # FC do not have uncertainties - ) - - # Calculate diffmap - diffmap = compute_difference_map_wrapper( - derivative_map_scaled, - native_map_scaled, - use_k_weighting=args.k_weight_with_fixed_parameter - or args.k_weight_with_parameter_optimization, - k_param=args.k_weight_with_fixed_parameter, - optimize_k=args.k_weight_with_parameter_optimization, - ) + mapset = load_and_scale_mapset(args) + + # TODO: eliminate this copypasta as well + # question: do we want to enforce k-weighting? if not, I would say maybe we just combine + # the two scripts and add a single --no-tv-denoising flag or something... + if args.k_weight_with_parameter_optimization: + diffmap, opt_k = max_negentropy_kweighted_difference_map(mapset.derivative, mapset.native) + log.info("Optimal k-parameter determined:", value=opt_k) + elif args.k_weight_with_fixed_parameter: + diffmap = compute_kweighted_difference_map( + mapset.derivative, mapset.native, k_parameter=args.k_weight_with_fixed_parameter + ) + else: + diffmap = compute_difference_map(mapset.derivative, mapset.native) # Now TV denoise - + # TODO: hard coded linspace denoised_map, tv_result = tv_denoise_difference_map( diffmap, full_output=True, lambda_values_to_scan=np.linspace(1e-8, 0.05, 10) ) - # denoised_map, tv_result = tv_denoise_difference_map(diffmap, full_output=True) - print( - "Writing output file with optimal lambda weight of ", tv_result.optimal_lambda - ) # noqa: T201 + log.ingo("Writing output file", optimal_lambda=tv_result.optimal_lambda) denoised_map.write_mtz(args.output) diff --git a/pyproject.toml b/pyproject.toml index cb5ebdb..a24e188 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "gemmi", "scikit-image", "reciprocalspaceship", + "structlog", "pytest", "pytest-cov", ]