Files
Jungfraujoch/tools/jfjoch_scale.cpp
T

373 lines
14 KiB
C++

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <iostream>
#include <vector>
#include <string>
#include <unistd.h>
#include <future>
#include <mutex>
#include <atomic>
#include <chrono>
#include <fstream>
#include <sstream>
#include <getopt.h>
#include "../reader/JFJochHDF5Reader.h"
#include "../common/Logger.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/time_utc.h"
#include "../common/print_license.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../image_analysis/indexing/IndexerFactory.h"
#include "../writer/FileWriter.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../receiver/JFJochReceiverPlots.h"
#include "../compression/JFJochCompressor.h"
#include "../image_analysis/LoadFCalcFromMtz.h"
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/UpdateReflectionResolution.h"
void print_usage() {
std::cout << "Usage ./jfjoch_scale {<options>} <input.h5>" << std::endl;
std::cout << "Options:" << std::endl;
std::cout << " -o, --output-prefix <txt> Output file prefix (default: output)" << std::endl;
std::cout << " -N, --threads <num> Number of threads (default: 1)" << std::endl;
std::cout << " -s, --start-image <num> Start image number (default: 0)" << std::endl;
std::cout << " -e, --end-image <num> End image number (default: all)" << std::endl;
std::cout << " -v, --verbose Verbose output" << std::endl;
std::cout << "" << std::endl;
std::cout << " Scaling and merging" << std::endl;
std::cout << " -S, --space-group <num> Space group number" << std::endl;
std::cout <<
" -D, --scaling-high-resolution <num> High resolution limit for scaling/merging (default: 0.0; no limit)"
<< std::endl;
std::cout << " -P, --partiality <txt> Partiality refinement fixed|rot|unity (default: fixed)" <<
std::endl;
std::cout << " -A, --anomalous Anomalous mode (don't merge Friedel pairs)" << std::endl;
std::cout << " -B, --refine-bfactor Refine per image B-factor" << std::endl;
std::cout << " -w, --wedge[=num] Refine image wedge during scaling with starting wedge value"
<< std::endl;
std::cout << " --min-partiality <num> Minimum partiality to accept reflection (default: 0.02)" <<
std::endl;
std::cout << " --min-image-cc <num> Per-image CC limit in percent (default: no limit)" << std::endl;
std::cout << " --scaling-iterations <num> Number of scaling iterations with no reference data (default: 3)"
<< std::endl;
std::cout << " -z, --reference-mtz <file> Reference MTZ file" << std::endl;
}
enum {
OPT_MIN_PARTIALITY = 1000,
OPT_MIN_IMAGE_CC,
OPT_SCALING_ITERATIONS
};
static option long_options[] = {
{"verbose", no_argument, nullptr, 'v'},
{"output-prefix", required_argument, nullptr, 'o'},
{"threads", required_argument, nullptr, 'N'},
{"start-image", required_argument, nullptr, 's'},
{"end-image", required_argument, nullptr, 'e'},
{"reference-mtz", required_argument, nullptr, 'z'},
{"space-group", required_argument, nullptr, 'S'},
{"partiality", required_argument, nullptr, 'P'},
{"anomalous", no_argument, nullptr, 'A'},
{"refine-bfactor", no_argument, nullptr, 'B'},
{"wedge", optional_argument, nullptr, 'w'},
{"scaling-high-resolution", required_argument, nullptr, 'D'},
{"min-partiality", required_argument, nullptr, OPT_MIN_PARTIALITY},
{"min-image-cc", required_argument, nullptr, OPT_MIN_IMAGE_CC},
{"scaling-iterations", required_argument, nullptr, OPT_SCALING_ITERATIONS},
{nullptr, 0, nullptr, 0}
};
int main(int argc, char **argv) {
RegisterHDF5Filter();
print_license("jfjoch_scale");
Logger logger("jfjoch_scale");
std::string input_file;
std::string output_prefix = "output";
int nthreads = 0;
int start_image = 0;
int end_image = -1; // -1 indicates process until end
bool verbose = false;
bool anomalous_mode = false;
std::optional<int64_t> space_group_number;
bool refine_bfactor = false;
bool refine_wedge = false;
std::optional<double> wedge_for_scaling;
std::string ref_mtz;
double min_partiality = 0.02;
double min_image_cc = 0.0;
int64_t scaling_iter = 3;
PartialityModel partiality_model = PartialityModel::Fixed;
std::optional<float> d_min_scale_merge;
if (argc == 1) {
print_usage();
exit(EXIT_FAILURE);
}
int opt;
int option_index = 0;
const char *short_opts = "vo:N:s:e:z:S:P:ABw::D:p:q:i:";
while ((opt = getopt_long(argc, argv, short_opts, long_options, &option_index)) != -1) {
switch (opt) {
case 'o':
output_prefix = optarg;
break;
case 'v':
verbose = true;
break;
case 'N':
nthreads = atoi(optarg);
break;
case 's':
start_image = atoi(optarg);
break;
case 'e':
end_image = atoi(optarg);
break;
case 'z':
ref_mtz = optarg;
break;
case 'S':
space_group_number = atoi(optarg);
break;
case 'P':
if (strcmp(optarg, "unity") == 0)
partiality_model = PartialityModel::Unity;
else if (strcmp(optarg, "fixed") == 0)
partiality_model = PartialityModel::Fixed;
else if (strcmp(optarg, "rot") == 0)
partiality_model = PartialityModel::Rotation;
else {
logger.Error("Invalid partiality mode: {}", optarg);
print_usage();
exit(EXIT_FAILURE);
}
break;
case 'A':
anomalous_mode = true;
break;
case 'B':
refine_bfactor = true;
break;
case 'w':
refine_wedge = true;
if (optarg)
wedge_for_scaling = std::stod(optarg);
break;
case 'D':
d_min_scale_merge = atof(optarg);
logger.Info("High resolution limit for scaling/merging set to {:.2f} A", d_min_scale_merge.value());
break;
case OPT_MIN_PARTIALITY:
min_partiality = std::stod(optarg);
break;
case OPT_MIN_IMAGE_CC:
min_image_cc = std::stod(optarg);
break;
case OPT_SCALING_ITERATIONS:
scaling_iter = atoi(optarg);
if (scaling_iter <= 0) {
logger.Error("Invalid scaling iteration count: {}", scaling_iter);
exit(EXIT_FAILURE);
}
break;
}
}
if (optind != argc - 1) {
logger.Error("Input file not specified");
print_usage();
exit(EXIT_FAILURE);
}
input_file = argv[optind];
logger.Verbose(verbose);
// Validate space group number early
const gemmi::SpaceGroup *space_group = nullptr;
if (space_group_number.has_value()) {
space_group = gemmi::find_spacegroup_by_number(space_group_number.value());
if (!space_group) {
logger.Error("Unknown space group number {}", space_group_number.value());
exit(EXIT_FAILURE);
}
logger.Info("Using space group {} (number {})", space_group->hm, space_group_number.value());
}
// 1. Read Input File
JFJochHDF5Reader reader;
try {
reader.ReadFile(input_file);
} catch (const std::exception &e) {
logger.Error("Error reading input file: {}", e.what());
exit(EXIT_FAILURE);
}
const auto dataset = reader.GetDataset();
if (!dataset) {
logger.Error("No experiment dataset found in the input file");
exit(EXIT_FAILURE);
}
logger.Info("Loaded dataset from {}", input_file);
std::vector<MergedReflection> reference_data;
if (!ref_mtz.empty()) {
reference_data = LoadFCalcFromMtz(ref_mtz);
logger.Info("Loaded {} reflections from {} MTZ file", reference_data.size(), ref_mtz);
}
uint64_t total_images_in_file = reader.GetNumberOfImages();
if (end_image < 0 || end_image > total_images_in_file)
end_image = total_images_in_file;
int images_to_process = end_image - start_image;
if (images_to_process <= 0) {
logger.Warning("No images to process (Start: {}, End: {}, Total: {})", start_image, end_image,
total_images_in_file);
return 0;
}
logger.Info("Starting analysis of {} images (range {}-{}) using {} threads",
images_to_process, start_image, end_image, nthreads);
// 2. Setup Experiment & Components
DiffractionExperiment experiment(dataset->experiment);
experiment.BitDepthImage(32).Compression(CompressionAlgorithm::BSHUF_LZ4);
experiment.FilePrefix(output_prefix);
experiment.Mode(DetectorMode::Standard); // Ensure full image analysis
experiment.PixelSigned(true);
experiment.OverwriteExistingFiles(true);
experiment.PolarizationFactor(0.99);
experiment.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);
experiment.SpaceGroupNumber(space_group_number);
experiment.ImagesPerTrigger(images_to_process);
experiment.NumTriggers(1);
ScalingSettings scaling_settings;
scaling_settings.SetPartialityModel(partiality_model);
if (d_min_scale_merge)
scaling_settings.HighResolutionLimit_A(d_min_scale_merge.value());
scaling_settings.MergeFriedel(!anomalous_mode);
scaling_settings.RefineB(refine_bfactor);
scaling_settings.RefineRotationWedge(refine_wedge);
if (wedge_for_scaling.has_value())
scaling_settings.RotationWedgeForScaling(wedge_for_scaling);
scaling_settings.MinPartiality(min_partiality);
scaling_settings.MinCCForImage(min_image_cc);
experiment.ImportScalingSettings(scaling_settings);
if (!experiment.GetUnitCell()) {
logger.Error("Experiment unit cell not found, cannot update reflection resolution");
exit(EXIT_FAILURE);
}
logger.Info("Loading reflections");
auto reflections = reader.ReadReflections(start_image, end_image);
auto refl_stats = UpdateReflectionResolution(experiment.GetUnitCell().value(), reflections);
logger.Info("Read {} reflections from {} images", refl_stats.n_reflections, refl_stats.n_images);
std::vector<float> mosaicity(end_image - start_image + 1);
for (int i = 0; i < end_image - start_image + 1; i++) {
mosaicity[i] = dataset->mosaicity_deg[start_image + i];
}
ScalingResult scale_result(0);
auto scale_start = std::chrono::steady_clock::now();
for (int i = 0; i < scaling_iter; i++) {
auto iter_start = std::chrono::steady_clock::now();
if (reference_data.empty()) {
auto merged = MergeAll(experiment, reflections);
ScaleOnTheFly scaling(experiment, merged);
scale_result = scaling.Scale(reflections, mosaicity, nthreads);
} else {
ScaleOnTheFly scaling(experiment, reference_data);
scale_result = scaling.Scale(reflections, mosaicity, nthreads);
}
scale_result.SaveToFile(output_prefix + "_iter" + std::to_string(i) + "_scale.dat");
auto iter_end = std::chrono::steady_clock::now();
double iter_time = std::chrono::duration<double>(iter_end - iter_start).count();
logger.Info("Scaling iteration {} took {:.3f} seconds", i, iter_time);
}
auto scale_end = std::chrono::steady_clock::now();
double scale_time = std::chrono::duration<double>(scale_end - scale_start).count();
logger.Info("Scaling completed in {:.2f} s", scale_time);
auto merge_start = std::chrono::steady_clock::now();
std::vector<uint8_t> merge_mask(images_to_process, 1);
auto rejected = CalcMergeMaskCC(experiment, scale_result.image_cc, merge_mask);
if (rejected > 0)
logger.Info("Rejected {} images due to low CC with reference", rejected);
auto merge_result = MergeAll(experiment, reflections, merge_mask);
auto merge_stats = MergeStats(experiment, merge_result, reflections, experiment.GetUnitCell().value(), merge_mask,
reference_data);
auto merge_end = std::chrono::steady_clock::now();
double merge_time = std::chrono::duration<double>(merge_end - merge_start).count();
logger.Info("Merge completed in {:.2f} s ({} unique reflections)", merge_time,
merge_result.size());
const bool fixed_space_group = space_group || experiment.GetGemmiSpaceGroup().has_value();
if (!fixed_space_group) {
logger.Info("Searching for space group from P1-merged reflections ...");
SearchSpaceGroupOptions sg_opts;
sg_opts.crystal_system.reset();
sg_opts.centering = '\0';
sg_opts.merge_friedel = experiment.GetScalingSettings().GetMergeFriedel();
sg_opts.d_min_limit_A = experiment.GetScalingSettings().GetHighResolutionLimit_A().value_or(0.0);
sg_opts.min_i_over_sigma = 0.0;
sg_opts.min_operator_cc = 0.80;
sg_opts.min_pairs_per_operator = 20;
sg_opts.min_total_compared = 100;
sg_opts.test_systematic_absences = true;
const auto sg_search = SearchSpaceGroup(merge_result, sg_opts);
logger.Info("");
{
std::istringstream iss(SearchSpaceGroupResultToText(sg_search));
std::string line;
while (std::getline(iss, line)) {
if (!line.empty())
logger.Info("{}", line);
}
}
logger.Info("");
}
// Print resolution-shell statistics table
std::cout << merge_stats;
if (!output_prefix.empty())
WriteReflections(merge_result, *experiment.GetUnitCell(), experiment, output_prefix);
}