PixelRefine: Claude fixed my bugs

2026-06-08 20:06:53 +02:00
parent 32e91f7287
commit 698a98be35
2 changed files with 193 additions and 48 deletions
@@ -319,6 +319,56 @@ PixelRefine::PixelRefine(const DiffractionExperiment &experiment,
        reference_data[hkl_key_generator(ref)] = ref.I;
 }

+void PixelRefine::BuildParameterBlocks(const PixelRefineData &data,
+                                       double beam[2], double &dist_mm,
+                                       double detector_rot[2], double rot_vec[3],
+                                       double latt_vec0[3], double latt_vec1[3], double latt_vec2[3]) const {
+    beam[0] = data.geom.GetBeamX_pxl();
+    beam[1] = data.geom.GetBeamY_pxl();
+    dist_mm = data.geom.GetDetectorDistance_mm();
+    detector_rot[0] = data.geom.GetPoniRot1_rad();
+    detector_rot[1] = data.geom.GetPoniRot2_rad();
+    rot_vec[0] = 1.0; rot_vec[1] = 0.0; rot_vec[2] = 0.0;
+    if (auto axis = data.geom.GetRotation()) {
+        rot_vec[0] = axis->GetAxis().x;
+        rot_vec[1] = axis->GetAxis().y;
+        rot_vec[2] = axis->GetAxis().z;
+    }
+
+    for (int i = 0; i < 3; ++i)
+        latt_vec0[i] = latt_vec1[i] = latt_vec2[i] = 0.0;
+
+    double beta = data.latt.GetUnitCell().beta;
+    switch (data.crystal_system) {
+        case gemmi::CrystalSystem::Orthorhombic:
+            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
+            break;
+        case gemmi::CrystalSystem::Tetragonal:
+            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
+            latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
+            break;
+        case gemmi::CrystalSystem::Cubic:
+            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
+            latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
+            break;
+        case gemmi::CrystalSystem::Hexagonal:
+            LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
+            break;
+        case gemmi::CrystalSystem::Monoclinic:
+            LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
+            latt_vec2[0] = beta;
+            break;
+        default: {
+            LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
+            const auto uc = data.latt.GetUnitCell();
+            latt_vec2[0] = uc.alpha * M_PI / 180.0;
+            latt_vec2[1] = uc.beta * M_PI / 180.0;
+            latt_vec2[2] = uc.gamma * M_PI / 180.0;
+            break;
+        }
+    }
+}
+
 template<class T>
 void PixelRefine::Run(const T *image,
                      const AzimuthalIntegrationProfile &profile,
@@ -340,7 +390,10 @@ void PixelRefine::Run(const T *image,
    const auto azim_std = profile.GetStd();
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const auto &corrections = mapping.Corrections();
-    const int azim_bin_count = mapping.GetAzimuthalBinCount();
+    // pixel_to_bin stores the *full* bin index (azimuthal_sector * q_bins + q_bin),
+    // so the valid range is the total number of bins, i.e. the profile size - NOT
+    // GetAzimuthalBinCount() (which is only the number of azimuthal sectors).
+    const int total_bin_count = static_cast<int>(azim_result.size());

    const double angle_rad = data.angle_deg * M_PI / 180.0;
    const int radius = data.shoebox_radius;
@@ -392,11 +445,15 @@ void PixelRefine::Run(const T *image,
                .PoniRot1_rad(data.geom.GetPoniRot1_rad())
                .PoniRot2_rad(data.geom.GetPoniRot2_rad());

-        prediction.Calc(exp_iter, data.latt, settings_prediction);
+        const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);

        // ---- 2. Collect per-reflection shoebox pixels -------------------------
+        // GetReflections() returns the full pre-sized buffer; only the first
+        // nrefl entries are valid for this image (the rest are stale/zeroed).
        groups.clear();
-        for (const auto &refl : prediction.GetReflections()) {
+        const auto &predicted = prediction.GetReflections();
+        for (int ri = 0; ri < nrefl; ++ri) {
+            const auto &refl = predicted[ri];
            const auto hkl = hkl_key_generator(refl);
            if (!reference_data.contains(hkl))
                continue;
@@ -421,10 +478,10 @@ void PixelRefine::Run(const T *image,
                    const size_t npixel = xpixel * y + x;
                    const int azim_bin = pixel_to_bin[npixel];

-                    // Skip pixels not mapped to an azimuthal bin or carrying a
-                    // sentinel (masked / saturated) value. We assume the pixel
-                    // mask is already applied upstream.
-                    if (azim_bin >= azim_bin_count)
+                    // Skip pixels not mapped to a bin or carrying a sentinel
+                    // (masked / saturated) value. We assume the pixel mask is
+                    // already applied upstream.
+                    if (azim_bin >= total_bin_count)
                        continue;
                    if (image[npixel] == std::numeric_limits<T>::max())
                        continue;
@@ -464,47 +521,8 @@ void PixelRefine::Run(const T *image,
            return;

        // ---- 3. Set up parameter blocks (geometry part mirrors XtalOptimizer) -
-        beam[0] = data.geom.GetBeamX_pxl();
-        beam[1] = data.geom.GetBeamY_pxl();
-        dist_mm = data.geom.GetDetectorDistance_mm();
-        detector_rot[0] = data.geom.GetPoniRot1_rad();
-        detector_rot[1] = data.geom.GetPoniRot2_rad();
-        rot_vec[0] = 1.0; rot_vec[1] = 0.0; rot_vec[2] = 0.0;
-        if (auto axis = data.geom.GetRotation()) {
-            rot_vec[0] = axis->GetAxis().x;
-            rot_vec[1] = axis->GetAxis().y;
-            rot_vec[2] = axis->GetAxis().z;
-        }
-
-        double beta = data.latt.GetUnitCell().beta;
-        switch (data.crystal_system) {
-            case gemmi::CrystalSystem::Orthorhombic:
-                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
-                break;
-            case gemmi::CrystalSystem::Tetragonal:
-                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
-                latt_vec1[0] = (latt_vec1[0] + latt_vec1[1]) / 2.0;
-                break;
-            case gemmi::CrystalSystem::Cubic:
-                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
-                latt_vec1[0] = (latt_vec1[0] + latt_vec1[1] + latt_vec1[2]) / 3.0;
-                break;
-            case gemmi::CrystalSystem::Hexagonal:
-                LatticeToRodriguesAndLengths_Hex(data.latt, latt_vec0, latt_vec1);
-                break;
-            case gemmi::CrystalSystem::Monoclinic:
-                LatticeToRodriguesLengthsBeta_Mono(data.latt, latt_vec0, latt_vec1, beta);
-                latt_vec2[0] = beta;
-                break;
-            default: {
-                LatticeToRodriguesAndLengths_GS(data.latt, latt_vec0, latt_vec1);
-                const auto uc = data.latt.GetUnitCell();
-                latt_vec2[0] = uc.alpha * M_PI / 180.0;
-                latt_vec2[1] = uc.beta * M_PI / 180.0;
-                latt_vec2[2] = uc.gamma * M_PI / 180.0;
-                break;
-            }
-        }
+        BuildParameterBlocks(data, beam, dist_mm, detector_rot, rot_vec,
+                             latt_vec0, latt_vec1, latt_vec2);

        // ---- 4. Build the problem ---------------------------------------------
        ceres::Problem problem;
@@ -709,6 +727,114 @@ void PixelRefine::Run(const T *image,
    }
 }

+std::vector<float> PixelRefine::PredictImage(const AzimuthalIntegrationProfile &profile,
+                                             BraggPrediction &prediction,
+                                             const PixelRefineData &data,
+                                             bool include_background) const {
+    std::vector<float> img(xpixel * ypixel, 0.0f);
+
+    const double lambda = data.geom.GetWavelength_A();
+    const double pixel_size = data.geom.GetPixelSize_mm();
+    const auto azim_result = profile.GetResult();
+    const auto &pixel_to_bin = mapping.GetPixelToBin();
+    const auto &corrections = mapping.Corrections();
+    const int total_bin_count = static_cast<int>(azim_result.size());
+    const double angle_rad = data.angle_deg * M_PI / 180.0;
+    const int radius = data.shoebox_radius;
+    const double bw = data.bandwidth;
+
+    auto recip_area = [&](double x, double y) -> double {
+        const Coord qx = data.geom.DetectorToRecip(x + 0.5, y) - data.geom.DetectorToRecip(x - 0.5, y);
+        const Coord qy = data.geom.DetectorToRecip(x, y + 0.5) - data.geom.DetectorToRecip(x, y - 0.5);
+        return (qx % qy).Length();
+    };
+    auto bandwidth_radial_sq = [&](double d) -> double {
+        if (bw <= 0.0 || d <= 0.0)
+            return 0.0;
+        const double bl = bw * lambda;
+        return bl * bl / (2.0 * d * d * d * d);
+    };
+
+    // The model works in solid-angle/polarization-corrected units (as in Run,
+    // where Iobs = raw * correction). Map back to raw detector units (/ correction)
+    // so the predicted image overlays directly on the original image.
+    auto to_raw = [&](size_t npixel, double corrected) -> float {
+        const double corr = corrections[npixel];
+        return (corr > 0.0) ? static_cast<float>(corrected / corr) : 0.0f;
+    };
+
+    // Background base layer (per-pixel azimuthal mean), full-frame pass.
+    if (include_background) {
+        for (size_t p = 0; p < img.size(); ++p) {
+            const int bin = pixel_to_bin[p];
+            if (bin >= 0 && bin < total_bin_count)
+                img[p] = to_raw(p, azim_result[bin]);
+        }
+    }
+
+    double beam[2], dist_mm, detector_rot[2], rot_vec[3];
+    double latt_vec0[3], latt_vec1[3], latt_vec2[3];
+    BuildParameterBlocks(data, beam, dist_mm, detector_rot, rot_vec, latt_vec0, latt_vec1, latt_vec2);
+
+    DiffractionExperiment exp_iter = experiment;
+    exp_iter.BeamX_pxl(data.geom.GetBeamX_pxl())
+            .BeamY_pxl(data.geom.GetBeamY_pxl())
+            .DetectorDistance_mm(data.geom.GetDetectorDistance_mm())
+            .PoniRot1_rad(data.geom.GetPoniRot1_rad())
+            .PoniRot2_rad(data.geom.GetPoniRot2_rad());
+
+    const BraggPredictionSettings settings_prediction{
+        .high_res_A = experiment.GetBraggIntegrationSettings().GetDMinLimit_A(),
+        .max_hkl = 100,
+        .centering = data.centering
+    };
+    const int nrefl = prediction.Calc(exp_iter, data.latt, settings_prediction);
+    const auto &predicted = prediction.GetReflections();
+
+    for (int ri = 0; ri < nrefl; ++ri) {
+        const auto &refl = predicted[ri];
+        const auto it = reference_data.find(hkl_key_generator(refl));
+        if (it == reference_data.end())
+            continue;
+
+        const double Itrue = it->second;
+        const double R_bw_sq = bandwidth_radial_sq(refl.d);
+
+        const int min_y = std::max<int>(refl.predicted_y - radius, 0);
+        const int max_y = std::min<int>(refl.predicted_y + radius, ypixel - 1);
+        const int min_x = std::max<int>(refl.predicted_x - radius, 0);
+        const int max_x = std::min<int>(refl.predicted_x + radius, xpixel - 1);
+
+        for (int y = min_y; y <= max_y; ++y) {
+            for (int x = min_x; x <= max_x; ++x) {
+                const size_t npixel = xpixel * y + x;
+
+                // Pure Bragg signal: Ibkg = 0 so Model() returns signal only; the
+                // background is already laid down above. Same code path as Run.
+                PixelObs obs{
+                    .x = static_cast<double>(x),
+                    .y = static_cast<double>(y),
+                    .Iobs = 0.0,
+                    .Ibkg = 0.0,
+                    .weight = 1.0,
+                    .A_recip = recip_area(x, y),
+                    .angle_rad = angle_rad
+                };
+                PixelResidual pr(obs, Itrue, lambda, pixel_size,
+                                 refl.h, refl.k, refl.l, R_bw_sq, data.crystal_system);
+
+                double signal = 0.0;
+                if (pr.Model(beam, &dist_mm, detector_rot, rot_vec,
+                             latt_vec0, latt_vec1, latt_vec2,
+                             &data.scale_factor, &data.B_factor, data.R, signal))
+                    img[npixel] += to_raw(npixel, signal);
+            }
+        }
+    }
+
+    return img;
+}
+
 // Explicit instantiations for the supported (uncompressed) image pixel types.
 template void PixelRefine::Run<int8_t>(const int8_t *, const AzimuthalIntegrationProfile &, BraggPrediction &, PixelRefineData &);
 template void PixelRefine::Run<int16_t>(const int16_t *, const AzimuthalIntegrationProfile &, BraggPrediction &, PixelRefineData &);