jfjoch_scale: Improve process of reading reflections, particularly with multiple files

reader/JFJochHDF5Reader.cpp

@@ -1150,8 +1150,8 @@ std::vector<HDF5DataSourceMessage> JFJochHDF5Reader::GetHDF5DataSource(
                           "Unsupported HDF5 file layout for source mapping");
 }
 
-std::vector<std::vector<Reflection> > JFJochHDF5Reader::ReadReflections(size_t start_image, std::optional<size_t> end_image) const
-{
+
+void JFJochHDF5Reader::ReadReflections(std::vector<std::vector<Reflection> > &output, size_t start_image, std::optional<size_t> end_image) const {
     std::unique_lock ul(hdf5_mutex);
     if (start_image >= number_of_images)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
@@ -1159,24 +1159,30 @@ std::vector<std::vector<Reflection> > JFJochHDF5Reader::ReadReflections(size_t s
 
     if (end_image.has_value() && end_image.value() < start_image)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
-                              "end_image must be greater than start_image if provided");
+                              "end_image must be greater or equal than start_image if provided");
 
-    int end_image_val = end_image.value_or(number_of_images);
+    int end_image_val = end_image.value_or(number_of_images - 1);
 
-    if (end_image_val > number_of_images)
+    if (end_image_val > number_of_images - 1)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
             "end_image_val must be less than or equal to number_of_images");
+
     if (!master_file)
         throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                               "Cannot read all reflections if file not loaded");
 
-    std::vector<std::vector<Reflection> > ret(end_image_val - start_image);
+    size_t output_start = output.size();
 
-    for (int i = 0; i < end_image_val - start_image; ++i)
+    output.resize(output.size() + end_image_val - start_image + 1);
+
+    for (int i = 0; i < end_image_val - start_image + 1; ++i)
         ReadReflectionsFromGroup(*master_file,
             fmt::format("/entry/reflections/image_{:06d}", start_image + i),
-            ret[i]);
-
-    return ret;
-
+            output[output_start + i]);
+}
+
+std::vector<std::vector<Reflection> > JFJochHDF5Reader::ReadReflections(size_t start_image, std::optional<size_t> end_image) const {
+    std::vector<std::vector<Reflection> > ret;
+    ReadReflections(ret, start_image, end_image);
+    return ret;
 }

reader/JFJochHDF5Reader.h

@@ -53,6 +53,7 @@ public:
     ) const;
 
     std::vector<std::vector<Reflection>> ReadReflections(size_t start_image = 0, std::optional<size_t> end_image = {}) const;
+    void ReadReflections(std::vector<std::vector<Reflection>> &output, size_t start_image = 0, std::optional<size_t> end_image = {}) const;
 
     CompressedImage ReadCalibration(std::vector<uint8_t> &tmp, const std::string &name) const;
 

tools/jfjoch_scale.cpp

@@ -189,15 +189,6 @@ int main(int argc, char **argv) {
         }
     }
 
-    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()) {
@@ -209,22 +200,48 @@ int main(int argc, char **argv) {
         logger.Info("Using space group {} (number {})", space_group->hm, space_group_number.value());
     }
 
-    // 1. Read Input File
+    std::vector<std::vector<Reflection>> reflections;
+    std::vector<float> mosaicity;
+
+    if (optind >= argc) {
+        logger.Error("Input file not specified");
+        print_usage();
+        exit(EXIT_FAILURE);
+    }
+
+    logger.Verbose(verbose);
+    logger.Info("Loading reflections");
+
+    // 1. Read Input files
     JFJochHDF5Reader reader;
-    try {
-        reader.ReadFile(input_file);
-    } catch (const std::exception &e) {
-        logger.Error("Error reading input file: {}", e.what());
-        exit(EXIT_FAILURE);
+    std::shared_ptr<const JFJochReaderDataset> dataset;
+    for (int i = optind; i < argc; i++) {
+        input_file = argv[optind];
+        try {
+            reader.ReadFile(input_file);
+            auto tmp_dataset = reader.GetDataset();
+            if (i == optind)
+                dataset = tmp_dataset;
+
+            uint64_t total_images_in_file = reader.GetNumberOfImages();
+            if (end_image < 0 || end_image >= total_images_in_file)
+                end_image = total_images_in_file - 1;
+
+            size_t mosaicity_size = mosaicity.size();
+            mosaicity.resize(mosaicity_size + end_image - start_image + 1);
+            for (int i = 0; i < end_image - start_image + 1; i++)
+                mosaicity[mosaicity_size + i] = dataset->mosaicity_deg[start_image + i];
+
+            reader.ReadReflections(reflections, start_image, end_image);
+            logger.Info("Loaded dataset from {}", input_file);
+        } catch (const std::exception &e) {
+            logger.Error("Error reading input file: {}", e.what());
+            // Hard exit is only for the first file
+            if (i == optind)
+                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()) {
@@ -232,20 +249,7 @@ int main(int argc, char **argv) {
         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);
+    logger.Info("Starting analysis of {} images", reflections.size());
 
     // 2. Setup Experiment & Components
     DiffractionExperiment experiment(dataset->experiment);
@@ -257,7 +261,6 @@ int main(int argc, char **argv) {
     experiment.PolarizationFactor(0.99);
     experiment.SetFileWriterFormat(FileWriterFormat::NXmxLegacy);
     experiment.SpaceGroupNumber(space_group_number);
-    experiment.ImagesPerTrigger(images_to_process);
     experiment.NumTriggers(1);
 
     ScalingSettings scaling_settings;
@@ -279,18 +282,11 @@ int main(int argc, char **argv) {
         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];
-    }
+    experiment.ImagesPerTrigger(refl_stats.n_images);
 
     ScalingResult scale_result(0);
 
@@ -319,7 +315,7 @@ int main(int argc, char **argv) {
 
     auto merge_start = std::chrono::steady_clock::now();
 
-    std::vector<uint8_t> merge_mask(images_to_process, 1);
+    std::vector<uint8_t> merge_mask(reflections.size(), 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);