mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 92 Wiki Activity

RotationScaleMerge: stream the resident ingest instead of materialising a record per observation

Browse Source 
The ingest-time geometry smoothing (delta_phi from the smoothed lattice, one exact-Bragg angle per
rocking event, partiality from the smoothed mosaicity) used to run over a 32-byte host record per
observation, copied from the source reflections and held beside them. It now reads hkl, frame, d
and zeta straight from the source reflections, frame by frame, and keeps arrays only for what it
rewrites (delta_phi, partiality) and what the rocking-event walks read in raw-hkl order
(image_number, and on the resident path the usability flag): 13 bytes an observation instead of 32.
The full-Obs path runs the same code and copies the two rewritten fields back.

Same arithmetic in the same order; merged output byte-identical.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013nW6FNRP1bBJJ8pfHiByAT
This commit is contained in:
leonarski_fandClaude Opus 5 committed 2026-09-20 18:45:17 +02:00
1 parent 8bc4c6e8a6
commit bb55715d28
2 files changed
+107 -97
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/scale_merge/RotationScaleMerge.cpp
+77 -70
View File
@@ -303,8 +303,8 @@ void RotationScaleMerge::Ingest() {
partials_released = false;
resident_ingest = false;
#ifdef JFJOCH_USE_CUDA
// Probed before the observation array is built: which record the build lays down (the narrow
// resident one or the full Obs) depends on whether the device pipeline will be active.
// Probed before the observations are built: whether the build lays down the full Obs array or
// only the ingest arrays depends on whether the device pipeline will be active.
gpu_ = std::make_unique<RotationScaleMergeGPU>();
gpu_active_ = gpu_->Available();
resident_ingest = gpu_active_ && observation_dump_path.empty();
@@ -500,6 +500,15 @@ void RotationScaleMerge::Ingest() {
n_partials_obs, n_frames, rawrun_start.size());
SmoothMosaicityAndPartiality();
if (!resident_ingest) {
ParallelChunks(static_cast<int>(partials.size()), ThreadsForWork(partials.size(), nthreads),
[&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
partials[i].delta_phi = ingest_delta_phi[i];
partials[i].partiality = ingest_partiality[i];
}
});
}
// The corr every pass starts from (see the header). The resident build filled it already,
// straight from the source reflections (and the smoothing between there and here never
@@ -525,7 +534,7 @@ void RotationScaleMerge::Ingest() {
for (int at = 0; at < n; at += slice) {
const int cnt = std::min(n - at, slice);
ParallelChunks(cnt, nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) buf[i] = partials_ingest[at + i].partiality;
for (int i = lo; i < hi; ++i) buf[i] = ingest_partiality[at + i];
});
gpu_->SetObsField(F::Partiality, at, cnt, buf.get());
}
@@ -577,16 +586,21 @@ void RotationScaleMerge::Ingest() {
gpu_->SetFrameCellOk(frame_cell_ok.data());
logger.Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
}
// The device now holds every per-obs field the resident pipeline reads, so the host record's
// The device now holds every per-obs field the resident pipeline reads, so the ingest arrays'
// one remaining reader is RockingEventFrames - which on this path sees only ingest-time values
// (see the header). Take its answer now and hand the record back. The dump path keeps the full
// Obs array: the CPU combine is what writes the dump.
// (see the header). Take its answer now and hand them back. The dump path keeps the full Obs
// array: the CPU combine is what writes the dump.
if (resident_ingest) {
rocking_event_frames_at_ingest = RockingEventFrames();
std::vector<ObsIngest>().swap(partials_ingest);
partials_released = true;
}
#endif
std::vector<float>().swap(ingest_delta_phi);
std::vector<float>().swap(ingest_partiality);
std::vector<float>().swap(ingest_image_number);
std::vector<uint8_t>().swap(ingest_rock_ok);
std::vector<uint8_t>().swap(ingest_keep);
std::vector<int32_t>().swap(ingest_src_start);
}
// Build `partials` (and the per-frame CSR, the finiteness mask and `perm`) from the source reflections,
@@ -623,14 +637,15 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
// Where each frame's reflections sit in the key numbering, and how many of them survive. Building in
// index order keeps each frame's observations contiguous and in the order they had, so the per-frame
// ranges still describe them and every per-frame sum keeps its sequence of roundings.
std::vector<int32_t> src_start(n_frames, 0);
ingest_src_start.assign(n_frames, 0);
{
int32_t at = 0;
for (int o = 0; o < n_frames; ++o) {
src_start[o] = at;
ingest_src_start[o] = at;
at += static_cast<int32_t>(partials_out[o].reflections.size());
}
}
const std::vector<int32_t> &src_start = ingest_src_start;
frame_count.assign(n_frames, 0);
ParallelFor(n_frames, nthreads, [&](int o) {
int32_t c = 0;
@@ -696,14 +711,16 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
// finiteness (immutable) - it lets ComputeAsuGroups stamp the ASU-group id per obs from a flat
// 1-byte array instead of re-reading the fat Obs struct for every space group.
finite_ok.assign(n_keep, 0);
ingest_delta_phi.resize(n_keep);
ingest_partiality.resize(n_keep);
ingest_image_number.resize(n_keep);
#ifdef JFJOCH_USE_CUDA
if (resident_ingest) {
// The resident path: build the narrow host record, then stage every immutable per-obs field
// to the device straight from the source reflections, in frame-aligned bounded slices. The
// 80-byte Obs is never materialised here - beside the source reflections it was another
// copy of the whole payload - and the slice buffers are bounded whatever the run.
partials_ingest.clear();
partials_ingest.resize(n_keep);
// The resident path: fill the ingest arrays, then stage every immutable per-obs field to the
// device straight from the source reflections, in frame-aligned bounded slices. The 80-byte
// Obs is never materialised here - beside the source reflections it was another copy of the
// whole payload - and the slice buffers are bounded whatever the run.
ingest_rock_ok.resize(n_keep);
corr_ingested.resize(n_keep);
ParallelFor(n_frames, nthreads, [&](int o) {
int32_t src = src_start[o];
@@ -711,25 +728,18 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
for (const auto &r : partials_out[o].reflections) {
if (!keep[src]) { ++src; continue; }
++src;
ObsIngest &obs = partials_ingest[at];
obs.h = static_cast<int16_t>(r.h);
obs.k = static_cast<int16_t>(r.k);
obs.l = static_cast<int16_t>(r.l);
obs.frame = o;
obs.image_number = r.image_number;
obs.d = r.d;
obs.zeta = r.zeta;
obs.delta_phi = r.delta_phi_deg;
obs.partiality = r.partiality;
ingest_delta_phi[at] = r.delta_phi_deg;
ingest_partiality[at] = r.partiality;
ingest_image_number[at] = r.image_number;
const float prescaling_corr = static_cast<float>(
r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]);
corr_ingested[at] = r.image_scale_corr;
finite_ok[at] = (std::isfinite(r.I) && std::isfinite(prescaling_corr)
&& prescaling_corr != 0.0f
&& std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0;
obs.rock_ok = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f
&& std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f)
? 1 : 0;
ingest_rock_ok[at] = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f
&& std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f)
? 1 : 0;
++at;
}
});
@@ -816,6 +826,9 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
obs.group = -1;
finite_ok[at] = (std::isfinite(obs.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f
&& std::isfinite(obs.sigma) && obs.sigma > 0.0f) ? 1 : 0;
ingest_delta_phi[at] = r.delta_phi_deg;
ingest_partiality[at] = r.partiality;
ingest_image_number[at] = r.image_number;
++at;
}
});
@@ -824,6 +837,7 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
#endif
n_partials_obs = n_keep;
ingest_keep = std::move(keep);
if (n_keep < n_obs)
logger.Info("RotationScaleMerge: dropped {} of {} observations ({} of {} distinct hkl) outside the "
@@ -891,7 +905,7 @@ namespace {
// window would over-smooth a fast-slipping crystal and under-smooth a sparse one. The frame's own
// value is excluded when scoring, so a window is only preferred if it PREDICTS that frame better,
// which is exactly the noise-versus-drift trade.
template <class T> void RotationScaleMerge::SmoothGeometry(std::vector<T> &obs) {
void RotationScaleMerge::SmoothGeometry() {
const auto gon = x.GetGoniometer();
if (!gon || n_frames < 5)
return;
@@ -998,37 +1012,31 @@ template <class T> void RotationScaleMerge::SmoothGeometry(std::vector<T> &obs)
}
std::atomic<int64_t> changed{0};
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
ParallelFor(n_frames, nthreads, [&](int f) {
if (!ok[f])
return;
const std::array<Coord, 3> &rc = frame_recip[f];
int64_t local = 0;
for (int i = lo; i < hi; ++i) {
auto &o = obs[i];
if (o.frame < 0 || o.frame >= n_frames || !ok[o.frame])
continue;
const std::array<Coord, 3> &r = frame_recip[o.frame];
const Coord p0 = r[0] * static_cast<float>(o.h) + r[1] * static_cast<float>(o.k)
+ r[2] * static_cast<float>(o.l);
const float phi_near = o.delta_phi * static_cast<float>(PI) / 180.0f;
ForEachKeptIn(f, [&](int i, const Reflection &r) {
const Coord p0 = rc[0] * static_cast<float>(r.h) + rc[1] * static_cast<float>(r.k)
+ rc[2] * static_cast<float>(r.l);
const float phi_near = ingest_delta_phi[i] * static_cast<float>(PI) / 180.0f;
const auto phi = SolvePhi(p0, S0, m1, m2, m3, m2_S0, m3_S0, phi_near);
if (!phi || !std::isfinite(*phi))
continue;
o.delta_phi = *phi * 180.0f / static_cast<float>(PI);
return;
ingest_delta_phi[i] = *phi * 180.0f / static_cast<float>(PI);
++local;
}
});
changed += local;
});
logger.Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
"recomputed delta_phi for {} of {} partials", best_half, changed.load(),
obs.size());
n_partials_obs);
}
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
if (resident_ingest) SmoothMosaicityAndPartialityOn(partials_ingest);
else SmoothMosaicityAndPartialityOn(partials);
}
template <class T> void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::vector<T> &obs) {
SmoothGeometry(obs);
SmoothGeometry();
// One rocking event, one exact-Bragg angle. Every partial's delta_phi is solved from its OWN
// frame's lattice - by the predictor, and again by SmoothGeometry above - so the per-frame
@@ -1054,25 +1062,25 @@ template <class T> void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::
int i = lo;
while (i < hi) {
int j = i + 1;
while (j < hi && obs[perm[j]].image_number
- obs[perm[j - 1]].image_number <= max_frame_gap)
while (j < hi && ingest_image_number[perm[j]]
- ingest_image_number[perm[j - 1]] <= max_frame_gap)
++j;
if (j - i >= 2) {
const float f0 = obs[perm[i]].image_number;
const float f0 = ingest_image_number[perm[i]];
double sum = 0.0;
int n = 0;
for (int m = i; m < j; ++m) {
const auto &o = obs[perm[m]];
if (!std::isfinite(o.delta_phi)) continue;
sum += o.delta_phi + increment_deg * (o.image_number - f0);
const int p = perm[m];
if (!std::isfinite(ingest_delta_phi[p])) continue;
sum += ingest_delta_phi[p] + increment_deg * (ingest_image_number[p] - f0);
++n;
}
if (n >= 2) {
const double c = sum / n;
for (int m = i; m < j; ++m) {
auto &o = obs[perm[m]];
o.delta_phi = static_cast<float>(c - increment_deg
* (o.image_number - f0));
const int p = perm[m];
ingest_delta_phi[p] = static_cast<float>(c - increment_deg
* (ingest_image_number[p] - f0));
}
++l_ev;
l_pa += j - i;
@@ -1145,20 +1153,19 @@ template <class T> void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::
// partiality recomputed here matches the one integration was based on. sin(theta_B) = lambda/(2d).
const double bandwidth_sigma = x.GetBandwidthFWHM().value_or(0.0f) / 2.3548;
const double half_wavelength_A = x.GetWavelength_A() / 2.0;
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
auto &o = obs[i];
float mos = mos_smooth[o.frame];
if (bandwidth_sigma > 0.0 && o.d > 0.0f) {
const double sin_theta = half_wavelength_A / o.d;
ParallelFor(n_frames, nthreads, [&](int f) {
ForEachKeptIn(f, [&](int i, const Reflection &r) {
float mos = mos_smooth[f];
if (bandwidth_sigma > 0.0 && r.d > 0.0f) {
const double sin_theta = half_wavelength_A / r.d;
const double sigma_bw = (180.0 / PI) * bandwidth_sigma * sin_theta
/ std::sqrt(1.0 - sin_theta * sin_theta);
mos = static_cast<float>(std::sqrt(mos * mos + sigma_bw * sigma_bw));
}
if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(o.zeta) && o.zeta > 0.0f
&& std::isfinite(o.delta_phi))
o.partiality = RotationPartiality(o.delta_phi, o.zeta, mos, wedge);
}
if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(r.zeta) && r.zeta > 0.0f
&& std::isfinite(ingest_delta_phi[i]))
ingest_partiality[i] = RotationPartiality(ingest_delta_phi[i], r.zeta, mos, wedge);
});
});
logger.Info("Recomputed partiality from frame-order-smoothed mosaicity");
}
@@ -2139,10 +2146,10 @@ namespace {
int RotationScaleMerge::RockingEventFrames() const {
if (partials_released)
return rocking_event_frames_at_ingest;
if (resident_ingest) // Ingest's own call, before the narrow record is handed back
if (resident_ingest) // Ingest's own call, before the ingest arrays are handed back
return RockingEventFramesOver(
[&](int i) { return partials_ingest[i].rock_ok != 0; },
[&](int i) { return partials_ingest[i].image_number; });
[&](int i) { return ingest_rock_ok[i] != 0; },
[&](int i) { return ingest_image_number[i]; });
return RockingEventFramesOver(
[&](int i) {
const Obs &o = partials[i];
image_analysis/scale_merge/RotationScaleMerge.h
+30 -27
View File
@@ -140,26 +140,6 @@ private:
// re-run over the samples that survive the automatic resolution cutoff. See MergeAndStats.
struct Sample { double s2, I2, dev2; float d; };
// The record the ingest keeps on the host when the device pipeline is resident (a GPU and no
// observation dump): every immutable per-obs field is uploaded straight from the source
// reflections inside BuildInRangeObservations, and the host keeps only what its own remaining
// stages touch - the geometry smoothing (delta_phi, then the recomputed partiality), plus the
// rocking-event walk's inputs. Thirty-two bytes against the Obs's eighty, and on that path the fat
// record is never materialised at all: on a fine-sliced long axis it alone is many gigabytes,
// held beside the source reflections it was copied from.
struct ObsIngest {
// Sixteen bits are plenty for a Miller index - |h| <= a / d_min, in the hundreds even on the
// longest axis at atomic resolution - and on a long axis this array sets the ingest's memory
// high-water mark, so each index costs two bytes here, not four.
int16_t h, k, l;
// The rocking-event usability test (RockingEventFrames), evaluated on the ingest-time
// values - which is what that walk reads on the resident path, where the host corr is
// never refreshed (see partials_released). Beside the indices, where it fills their padding.
uint8_t rock_ok;
int32_t frame;
float image_number, d, zeta, delta_phi, partiality;
};
// The narrow per-observation record the ingest sort orders: the raw hkl the runs are cut on, the
// frame position that breaks a tie inside one, the observation's own index (which makes the order
// total - see the .cpp), and the resolution the range test reads. Twenty-four bytes against the
@@ -223,12 +203,36 @@ private:
int report_shell_count = 9;
std::optional<ReportResolutionRange> report_range; // the second, reference-range table
// Flat buffers, allocated once by Ingest() and reused across Run() calls. Exactly one of the
// two observation arrays is built: the narrow resident record when the device pipeline is
// active and no observation dump was asked for (resident_ingest), the full Obs otherwise.
// Flat buffers, allocated once by Ingest() and reused across Run() calls. The full Obs array is
// built only when the device pipeline is NOT resident (no GPU, or an observation dump was asked
// for); on the resident path (resident_ingest) the device holds the observations and the host
// keeps only the ingest arrays below, until the end of Ingest.
std::vector<Obs> partials; // all per-frame partials, grouped by frame
std::vector<ObsIngest> partials_ingest; // the resident-path record (same order)
bool resident_ingest = false;
// What the ingest-time geometry smoothing (SmoothMosaicityAndPartiality) needs beyond the source
// reflections, one array per field in `partials` order: the two fields it rewrites, and the two the
// rocking-event walks read in raw-hkl order, where a source reflection is out of reach. Every other
// field it reads - hkl, frame, d, zeta - is read from the source reflection itself (ForEachKeptIn),
// so on the resident path no per-observation record is built at all: on a fine-sliced long axis
// that record was gigabytes, held beside the source reflections it was copied from. Filled by
// BuildInRangeObservations and handed back at the end of Ingest.
std::vector<float> ingest_delta_phi, ingest_partiality, ingest_image_number;
// The rocking-event usability test (RockingEventFrames) on the ingest-time values - which is what
// that walk reads on the resident path, where the host corr is never refreshed (see
// partials_released). Resident path only.
std::vector<uint8_t> ingest_rock_ok;
// Which source reflections the resolution range kept (indexed like the sort keys: frame by frame,
// from ingest_src_start), so that the i-th kept one of a frame can be found again.
std::vector<uint8_t> ingest_keep;
std::vector<int32_t> ingest_src_start;
// Call fn(observation index, source reflection) for each reflection of frame o that the ingest
// kept, in order.
template <class Fn> void ForEachKeptIn(int o, Fn fn) const {
int32_t src = ingest_src_start[o];
int32_t at = frame_start[o];
for (const auto &r : partials_out[o].reflections)
if (ingest_keep[src++]) fn(at++, r);
}
// How many observations Ingest built (== partials.size() until any release below), so the
// sizes survive when the array itself does not.
int n_partials_obs = 0;
@@ -445,9 +449,8 @@ private:
// per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the
// 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average
// that prediction applied. SG-independent, so done once in Ingest.
template <class T> void SmoothGeometry(std::vector<T> &obs);
void SmoothMosaicityAndPartiality(); // dispatch: the resident narrow record or the full Obs
template <class T> void SmoothMosaicityAndPartialityOn(std::vector<T> &obs);
void SmoothGeometry();
void SmoothMosaicityAndPartiality();
void Combine(); // partials -> fulls (CPU)