# src/StreamingPipeline.jl # ============================================================================ # The Streaming Pipeline Executor # # This module provides `process_dataset!`, the Sprint 2 master function that # streams spectral data through an in-place kernel chain and accumulates # results into a SparseMatrixCSC without ever holding more than 1 spectrum # per thread in RAM. # # Architecture: # 1. _iterate_spectra_fast → Mmap zero-copy views # 2. copyto!(writable_buf, view) → makes mutable copy for kernels # 3. Kernel chain: smooth! → baseline! → peaks → bin # 4. Thread-local (I, J, V) sparse accumulators # 5. Final sparse(I, J, V, num_bins, num_spectra) assembly # # This works alongside the existing execute_full_preprocessing in # PreprocessingPipeline.jl — it does NOT replace the app.jl integration. # ============================================================================ using SparseArrays using Printf # ============================================================================= # Configuration Structs # ============================================================================= """ StreamingStep Represents a single step in the streaming pipeline. """ struct StreamingStep name::Symbol params::Dict{Symbol, Any} end """ PipelineConfig Holds the complete configuration for a streaming pipeline execution. # Fields - `steps::Vector{StreamingStep}` — ordered sequence of processing steps - `reference_peaks::Vector{Float64}` — fixed m/z values for calibration (Category B) - `num_bins::Int` — number of bins for the output feature matrix - `min_peaks_per_bin::Int` — minimum peak count to keep a bin - `frequency_threshold::Float64` — minimum fraction of spectra a bin must appear in (0.0-1.0) # Example ```julia config = PipelineConfig( steps = [ StreamingStep(:smoothing, Dict(:method => :savitzky_golay, :window => 9, :order => 2)), StreamingStep(:baseline_correction, Dict(:method => :snip, :iterations => 100)), StreamingStep(:normalization, Dict(:method => :tic)), StreamingStep(:peak_picking, Dict(:method => :profile, :snr_threshold => 3.0)), ], num_bins = 2000 ) ``` """ struct PipelineConfig steps::Vector{StreamingStep} reference_peaks::Vector{Float64} num_bins::Int min_peaks_per_bin::Int frequency_threshold::Float64 end # Convenience constructor with defaults function PipelineConfig(; steps::Vector{StreamingStep}=StreamingStep[], reference_peaks::Vector{Float64}=Float64[], num_bins::Int=2000, min_peaks_per_bin::Int=3, frequency_threshold::Float64=0.0) return PipelineConfig(steps, reference_peaks, num_bins, min_peaks_per_bin, frequency_threshold) end # ============================================================================= # Sparse Accumulator (Thread-Local) # ============================================================================= """ SparseAccumulator Thread-local accumulator for sparse matrix construction. Collects (row, col, val) triplets that will be assembled into a SparseMatrixCSC at the end of the pipeline. """ mutable struct SparseAccumulator I::Vector{Int} # Row indices (bin indices) J::Vector{Int} # Column indices (spectrum indices) V::Vector{Float64} # Values (intensities) lck::Base.Threads.SpinLock function SparseAccumulator(capacity_hint::Int=10000) acc = new( Vector{Int}(undef, 0), Vector{Int}(undef, 0), Vector{Float64}(undef, 0), Base.Threads.SpinLock() ) sizehint!(acc.I, capacity_hint) sizehint!(acc.J, capacity_hint) sizehint!(acc.V, capacity_hint) return acc end end """ accumulate!(acc::SparseAccumulator, spectrum_idx::Int, bin_indices::AbstractVector{Int}, intensities::AbstractVector{Float64}) Appends peak data for one spectrum into the sparse accumulator. """ @inline function accumulate!(acc::SparseAccumulator, spectrum_idx::Int, bin_indices::AbstractVector{Int}, intensities::AbstractVector{Float64}) n = length(bin_indices) for k in 1:n @inbounds begin push!(acc.I, bin_indices[k]) push!(acc.J, spectrum_idx) push!(acc.V, intensities[k]) end end end # ============================================================================= # The Pipeline Executor # ============================================================================= """ process_dataset!(data::MSIData, config::PipelineConfig; progress_callback::Union{Function, Nothing}=nothing, masked_indices::Union{AbstractVector{Int}, Nothing}=nothing) The Sprint 2 master streaming function. Processes an entire MSI dataset through a kernel chain without holding more than 1 spectrum per thread in RAM. # Returns - `SparseMatrixCSC{Float64, Int}`: The feature matrix (bins × spectra) - `Vector{Float64}`: The m/z bin centers # Architecture 1. Ensures analytics are computed (for global m/z range) 2. Creates thread-local SparseAccumulators 3. Streams spectra via `_iterate_spectra_fast` 4. Per spectrum: copy view → kernel chain → peak detect → bin → accumulate 5. Merges accumulators → `sparse(I, J, V)` """ function process_dataset!(data::MSIData, config::PipelineConfig; progress_callback::Union{Function, Nothing}=nothing, masked_indices::Union{AbstractVector{Int}, Nothing}=nothing) # --- Step 1: Ensure analytics are computed (provides global m/z range) --- if !is_set(data.analytics_ready) println("Pre-computing analytics for streaming pipeline...") precompute_analytics(data) end # Determine global m/z range for binning global_min_mz = Base.Threads.atomic_add!(data.global_min_mz, 0.0) global_max_mz = Base.Threads.atomic_add!(data.global_max_mz, 0.0) if !isfinite(global_min_mz) || !isfinite(global_max_mz) || global_min_mz >= global_max_mz @warn "Invalid global m/z range: [$global_min_mz, $global_max_mz]. Cannot bin peaks." return spzeros(0, 0), Float64[] end num_bins = config.num_bins bin_edges = range(global_min_mz, stop=global_max_mz, length=num_bins + 1) bin_centers = [(bin_edges[i] + bin_edges[i+1]) / 2 for i in 1:num_bins] inv_bin_width = 1.0 / step(bin_edges) num_spectra = length(data.spectra_metadata) indices_to_process = masked_indices === nothing ? nothing : masked_indices # --- Step 2: Create thread-local accumulators --- n_threads = Base.Threads.nthreads() accumulators = [SparseAccumulator(num_spectra * 10) for _ in 1:n_threads] spectra_processed = Base.Threads.Atomic{Int}(0) # NEW: Create dedicated workspace buffers for each thread. # This completely eliminates the need for acquire/release and prevents deadlocks. workspaces_mz = [Vector{Float64}(undef, 0) for _ in 1:n_threads] workspaces_int = [Vector{Float64}(undef, 0) for _ in 1:n_threads] workspaces_scratch = [Vector{Float64}(undef, 0) for _ in 1:n_threads] # Pre-parse step configuration for fast dispatch in the hot loop has_smoothing = false has_baseline = false has_normalization = false has_transform = false has_peak_picking = false has_calibration = false smooth_params = Dict{Symbol, Any}() baseline_params = Dict{Symbol, Any}() norm_params = Dict{Symbol, Any}() transform_params = Dict{Symbol, Any}() peak_params = Dict{Symbol, Any}() for s in config.steps if s.name === :smoothing has_smoothing = true smooth_params = s.params elseif s.name === :baseline_correction has_baseline = true baseline_params = s.params elseif s.name === :normalization has_normalization = true norm_params = s.params elseif s.name === :stabilization || s.name === :intensity_transformation has_transform = true transform_params = s.params elseif s.name === :peak_picking has_peak_picking = true peak_params = s.params elseif s.name === :calibration has_calibration = true end end reference_masses = config.reference_peaks # --- Step 3: Stream and process --- start_time = time_ns() # Use let block to capture all variables cleanly for the closure let data=data, accumulators=accumulators, spectra_processed=spectra_processed, bin_edges=bin_edges, num_bins=num_bins, inv_bin_width=inv_bin_width, global_min_mz=global_min_mz, workspaces_mz=workspaces_mz, workspaces_int=workspaces_int, workspaces_scratch=workspaces_scratch, has_smoothing=has_smoothing, has_baseline=has_baseline, has_normalization=has_normalization, has_transform=has_transform, has_peak_picking=has_peak_picking, has_calibration=has_calibration, smooth_params=smooth_params, baseline_params=baseline_params, norm_params=norm_params, transform_params=transform_params, peak_params=peak_params, reference_masses=reference_masses _iterate_spectra_fast(data, indices_to_process) do idx, mz_view, int_view thread_id = Base.Threads.threadid() acc = accumulators[thread_id] # --- Grab Thread-Local Workspaces --- # No locking, no blocking, guaranteed to be available mz_buf = workspaces_mz[thread_id] int_buf = workspaces_int[thread_id] scratch_buf = workspaces_scratch[thread_id] resize!(mz_buf, length(mz_view)) resize!(int_buf, length(int_view)) resize!(scratch_buf, length(int_view)) copyto!(mz_buf, mz_view) copyto!(int_buf, int_view) # --- Kernel Chain (in pipeline order) --- # Category B: Fixed-reference calibration if has_calibration && !isempty(reference_masses) calibrate_inplace!(mz_buf, int_buf, reference_masses) end # Category A: Intensity transformation if has_transform transform_inplace!(int_buf, get(transform_params, :method, :sqrt)) end # Category A: Smoothing if has_smoothing smooth_inplace!(int_buf, scratch_buf, data; method=get(smooth_params, :method, :savitzky_golay), window=get(smooth_params, :window, 9), order=get(smooth_params, :order, 2)) end # Category A: Baseline correction if has_baseline baseline_subtract_inplace!(int_buf, scratch_buf, data; method=get(baseline_params, :method, :snip), iterations=get(baseline_params, :iterations, 100), window=get(baseline_params, :window, 20)) end # Category A: Normalization if has_normalization normalize_inplace!(int_buf, get(norm_params, :method, :tic)) end # --- Peak Detection & Binning --- if has_peak_picking detect_peaks_streaming(mz_buf, int_buf, scratch_buf; method=get(peak_params, :method, :profile), snr_threshold=Float64(get(peak_params, :snr_threshold, 3.0)), half_window=Int(get(peak_params, :half_window, 10)), min_peak_prominence=Float64(get(peak_params, :min_peak_prominence, 0.1)), merge_peaks_tolerance=Float64(get(peak_params, :merge_peaks_tolerance, 0.002))) do peak_mz, peak_int # Bin each discovered peak directly bin_idx = trunc(Int, (peak_mz - global_min_mz) * inv_bin_width) + 1 bin_idx = clamp(bin_idx, 1, num_bins) push!(acc.I, bin_idx) push!(acc.J, idx) push!(acc.V, peak_int) end else # No peak picking: bin raw intensity directly @inbounds for i in eachindex(mz_buf) bin_idx = trunc(Int, (mz_buf[i] - global_min_mz) * inv_bin_width) + 1 bin_idx = clamp(bin_idx, 1, num_bins) push!(acc.I, bin_idx) push!(acc.J, idx) push!(acc.V, int_buf[i]) end end Base.Threads.atomic_add!(spectra_processed, 1) end end # --- Step 4: Merge thread-local accumulators --- total_entries = sum(length(acc.I) for acc in accumulators) merged_I = Vector{Int}(undef, total_entries) merged_J = Vector{Int}(undef, total_entries) merged_V = Vector{Float64}(undef, total_entries) offset = 0 for acc in accumulators n = length(acc.I) if n > 0 copyto!(merged_I, offset + 1, acc.I, 1, n) copyto!(merged_J, offset + 1, acc.J, 1, n) copyto!(merged_V, offset + 1, acc.V, 1, n) offset += n end end # --- Step 5: Assemble sparse matrix --- # Use max combiner: when multiple peaks map to the same bin for same spectrum, # keep the maximum intensity feature_matrix = sparse(merged_I, merged_J, merged_V, num_bins, num_spectra, max) # --- Step 6: Apply frequency threshold if configured --- if config.frequency_threshold > 0.0 # Count how many spectra have a non-zero value in each bin bin_presence = vec(sum(feature_matrix .> 0, dims=2)) min_count = ceil(Int, config.frequency_threshold * num_spectra) keep_bins = findall(bin_presence .>= min_count) feature_matrix = feature_matrix[keep_bins, :] bin_centers = bin_centers[keep_bins] end duration = (time_ns() - start_time) / 1e9 n_processed = spectra_processed[] n_nonzeros = nnz(feature_matrix) sparsity = 1.0 - n_nonzeros / (size(feature_matrix, 1) * size(feature_matrix, 2) + 1) @printf "Streaming pipeline complete: %d spectra processed in %.2f seconds.\n" n_processed duration @printf "Feature matrix: %d bins × %d spectra, %d non-zeros (%.1f%% sparse)\n" size(feature_matrix, 1) size(feature_matrix, 2) n_nonzeros sparsity * 100 @printf "RAM: %.1f MB (vs %.1f MB dense)\n" (n_nonzeros * 16) / 1e6 (size(feature_matrix, 1) * size(feature_matrix, 2) * 8) / 1e6 if progress_callback !== nothing progress_callback(1.0) end return feature_matrix, collect(Float64, bin_centers) end """ save_sparse_matrix(matrix::SparseMatrixCSC, output_path::String) Exports a highly optimized SparseMatrixCSC array to disk using the standard Matrix Market Coordinate format (`.mtx`), guaranteeing no bottleneck or OOM crashes for extremely large MS dataset persistence. """ function save_sparse_matrix(matrix::SparseMatrixCSC{Float64, Int}, output_path::String) m, n = size(matrix) nnz_val = nnz(matrix) # Use streaming I/O with a large buffer for ultra-fast persistence open(output_path, "w") do io # Write Matrix Market Header write(io, "%%MatrixMarket matrix coordinate real general\n") write(io, "$m $n $nnz_val\n") # Directly extract CSC properties (O(1) memory, zero allocation) row_indices = rowvals(matrix) values_array = nonzeros(matrix) @inbounds for filter_j in 1:n # nzrange returns the index bounds for non-zero elements in column 'j' for idx in nzrange(matrix, filter_j) i = row_indices[idx] v = values_array[idx] write(io, "$i $filter_j $v\n") end end end end