Skip to main content

laddu_generation/
lib.rs

1//! Channel-aware Monte Carlo event generation.
2
3use std::{
4    collections::{HashMap, HashSet},
5    mem::size_of,
6    sync::Arc,
7};
8
9use laddu_compile::{CompiledModel, ReductionPlan};
10use laddu_data::{
11    BatchLayout,
12    data::{Dataset, EventBatch},
13    io::{EventSink, WritePlan, memory::MemorySink},
14    schema::{Precision as DataPrecision, Schema},
15};
16use laddu_expr::parameters::ParamValues;
17use laddu_memory::{MemoryFitRequest, MemoryFootprint};
18use laddu_physics::{
19    LadduPhysicsError,
20    channel::Channel,
21    generation::{PiecewiseDensity, proven_two_body_decay_weight},
22    vectors::RealVec4,
23};
24use laddu_runtime::{
25    Execution, MemoryBudget, MemoryDecision, MemoryLease, MemoryState, PreparedModel,
26};
27use maryada::{EnclosureOps, GlobalMinimizer, GlobalMinimizerOptions, Interval, IntervalOps};
28use rayon::prelude::*;
29use serde::{Deserialize, Serialize};
30use smallvec::SmallVec;
31use thiserror::Error;
32
33pub use laddu_physics::generation::{
34    AdaptiveTwoBodyDecay, InitialMomentum, InitialMomentumResult, MassProposal, MassProposalResult,
35    NamedMass, NamedMomentum, ProposalResult, ProposalRng, ScalarProposalResult, ScalarSource,
36    TComponent, TDistribution, TwoBodyScattering, VertexProposal,
37};
38
39/// Result type returned by event-generation operations.
40pub type GenerationResult<T> = Result<T, GenerationError>;
41
42/// Error produced while configuring or running event generation.
43#[derive(Debug, Error)]
44pub enum GenerationError {
45    /// The channel topology or output definition is invalid.
46    #[error("invalid generation channel: {0}")]
47    InvalidChannel(String),
48    /// Physics-level validation of the channel failed.
49    #[error("channel validation failed: {source}")]
50    ChannelValidation {
51        /// Underlying channel validation error.
52        #[source]
53        source: LadduPhysicsError,
54    },
55    /// A generation option is inconsistent or outside its valid range.
56    #[error("invalid generation configuration: {0}")]
57    InvalidConfiguration(String),
58    /// Initial-state momentum generation failed for a proposal.
59    #[error("initial-state proposal {index} failed: {source}")]
60    InitialState {
61        /// Global proposal index.
62        index: u64,
63        /// Underlying physics error.
64        #[source]
65        source: LadduPhysicsError,
66    },
67    /// An intermediate mass proposal failed.
68    #[error("mass proposal for edge `{edge}` at proposal {index} failed: {source}")]
69    MassProposal {
70        /// Global proposal index.
71        index: u64,
72        /// Channel edge being sampled.
73        edge: String,
74        /// Underlying physics error.
75        #[source]
76        source: LadduPhysicsError,
77    },
78    /// A vertex kinematics proposal failed.
79    #[error("vertex `{vertex}` at proposal {index} failed: {source}")]
80    VertexProposal {
81        /// Global proposal index.
82        index: u64,
83        /// Channel vertex being sampled.
84        vertex: String,
85        /// Underlying physics error.
86        #[source]
87        source: LadduPhysicsError,
88    },
89    /// A derived scalar proposal failed.
90    #[error("scalar column `{column}` at proposal {index} failed: {source}")]
91    ScalarProposal {
92        /// Global proposal index.
93        index: u64,
94        /// Scalar column being sampled.
95        column: String,
96        /// Underlying physics error.
97        #[source]
98        source: LadduPhysicsError,
99    },
100    /// Final-state kinematic validation failed.
101    #[error("kinematic validation at proposal {index} failed: {source}")]
102    Kinematics {
103        /// Global proposal index.
104        index: u64,
105        /// Underlying physics error.
106        #[source]
107        source: LadduPhysicsError,
108    },
109    /// Target-model evaluation failed.
110    #[error("model evaluation failed: {0}")]
111    Model(String),
112    /// A proposal exceeded the strict rejection-sampling envelope.
113    #[error("target weight {weight} exceeds envelope {envelope} at proposal {index}")]
114    EnvelopeOverflow {
115        /// Global proposal index.
116        index: u64,
117        /// Target weight of the overflowing proposal.
118        weight: f64,
119        /// Active envelope bound.
120        envelope: f64,
121    },
122    /// The proposal limit was reached before enough events were accepted.
123    #[error(
124        "accepted {accepted} events after exhausting {proposals} proposals (requested {requested})"
125    )]
126    Exhausted {
127        /// Requested event count.
128        requested: usize,
129        /// Number of accepted events.
130        accepted: usize,
131        /// Number of production proposals attempted.
132        proposals: usize,
133    },
134    /// Dataset input or output failed.
135    #[error(transparent)]
136    Data(#[from] laddu_data::LadduDataError),
137    /// Model preparation or execution failed.
138    #[error(transparent)]
139    Runtime(#[from] laddu_runtime::RuntimeError),
140    /// A physics-level proposal or validation failed.
141    #[error(transparent)]
142    Physics(#[from] LadduPhysicsError),
143}
144
145/// Configuration for weighted event generation.
146#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
147pub struct WeightedConfig {
148    /// Number of events to generate.
149    pub events: usize,
150    /// Memory available to proposal and output staging.
151    pub memory: MemoryBudget,
152    /// Deterministic random seed.
153    pub seed: u64,
154    /// Whether to include diagnostic weight columns.
155    pub diagnostics: bool,
156}
157
158impl WeightedConfig {
159    /// Creates a configuration for `events` with default batching and seed.
160    pub fn new(events: usize) -> Self {
161        Self {
162            events,
163            memory: MemoryBudget::Auto,
164            seed: 0,
165            diagnostics: false,
166        }
167    }
168}
169
170/// Configuration for rejection-sampled unweighted event generation.
171#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
172pub struct UnweightedConfig {
173    /// Number of accepted events to generate.
174    pub events: usize,
175    /// Optional safeguard limiting production proposals.
176    ///
177    /// `None` allows generation to continue until the requested event count is
178    /// reached. Pilot proposals are not included in this limit.
179    pub max_proposals: Option<usize>,
180    /// Memory available to proposal and output staging.
181    pub memory: MemoryBudget,
182    /// Deterministic random seed.
183    pub seed: u64,
184    /// Whether to include diagnostic weight columns.
185    pub diagnostics: bool,
186    /// Strategy used to establish the rejection-sampling envelope.
187    pub envelope: EnvelopeMode,
188    /// Policy applied when a proposal exceeds the active envelope.
189    pub envelope_overflow: EnvelopeOverflow,
190}
191
192impl UnweightedConfig {
193    /// Create an unweighted-generation configuration without a proposal limit.
194    pub fn new(events: usize) -> Self {
195        Self {
196            events,
197            max_proposals: None,
198            memory: MemoryBudget::Auto,
199            seed: 0,
200            diagnostics: false,
201            envelope: EnvelopeMode::default(),
202            envelope_overflow: EnvelopeOverflow::Error,
203        }
204    }
205
206    /// Stop with [`GenerationError::Exhausted`] after at most `max_proposals`
207    /// production proposals.
208    pub fn with_max_proposals(mut self, max_proposals: usize) -> Self {
209        self.max_proposals = Some(max_proposals);
210        self
211    }
212}
213
214/// Policy used when an unweighting proposal exceeds the current envelope.
215#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
216pub enum EnvelopeOverflow {
217    /// Stop immediately and report [`GenerationError::EnvelopeOverflow`].
218    #[default]
219    Error,
220    /// Grow the envelope and retrospectively thin previously accepted events.
221    ///
222    /// Adaptive generation buffers accepted events until the run completes,
223    /// because events already written to a sink cannot be withdrawn safely.
224    Grow {
225        /// Factor by which the observed overflow expands the envelope.
226        safety_factor: f64,
227    },
228}
229
230/// Strategy used to establish a rejection-sampling envelope.
231#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
232pub enum EnvelopeMode {
233    /// Use a caller-supplied fixed maximum weight.
234    Strict {
235        /// Fixed upper bound for target weights.
236        max_weight: f64,
237    },
238    /// Estimate an envelope from pilot proposals.
239    ///
240    /// Density-aware built-in proposals may use one deterministic pilot pass
241    /// for importance adaptation and a second pass for the final envelope.
242    Pilot {
243        /// Number of pilot proposals used to estimate the maximum.
244        proposals: usize,
245        /// Multiplier applied to the maximum pilot weight.
246        safety_factor: f64,
247    },
248    /// Prove a fixed envelope for the phase-space proposal weight using
249    /// outward-rounded interval arithmetic.
250    ///
251    /// This mode is valid only for unit-model generation.
252    ProvenPhaseSpace,
253}
254
255impl Default for EnvelopeMode {
256    fn default() -> Self {
257        Self::Pilot {
258            proposals: 10_000,
259            safety_factor: 2.0,
260        }
261    }
262}
263
264/// Source from which the final rejection envelope was obtained.
265#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
266pub enum EnvelopeKind {
267    /// A caller-supplied strict bound.
268    Strict,
269    /// A bound estimated from pilot proposals.
270    Pilot,
271    /// A maryada interval enclosure of the unit-model phase-space weight.
272    ProvenPhaseSpace,
273}
274
275/// Diagnostics produced while proving a unit-model phase-space envelope.
276#[derive(Clone, Copy, Debug, PartialEq)]
277pub struct ProvenEnvelopeReport {
278    /// Full nonnegative enclosure of every phase-space proposal weight.
279    pub weight_interval: Interval,
280    /// Upper endpoint used as the rejection-sampling maximum.
281    pub maximum_weight: f64,
282    /// Number of continuous proposal coordinates represented by the domain.
283    pub continuous_dimensions: usize,
284    /// Number of analytical piecewise regions represented by the enclosure.
285    pub piecewise_regions: usize,
286    /// Number of branch-and-bound subdivisions used to tighten the enclosure.
287    pub subdivisions: usize,
288}
289
290impl ProvenEnvelopeReport {
291    /// Return the outward-rounded lower and upper weight endpoints.
292    pub fn weight_bounds(&self) -> (f64, f64) {
293        self.weight_interval.bounds()
294    }
295}
296
297/// Diagnostics and aggregate statistics from an event-generation run.
298#[derive(Clone, Debug, Default, Serialize, Deserialize)]
299pub struct GenerationReport {
300    /// Requested event count.
301    pub requested: usize,
302    /// Produced event count.
303    pub produced: usize,
304    /// Number of production proposals.
305    pub proposals: usize,
306    /// Number of pilot proposals.
307    pub pilot_proposals: usize,
308    /// Number of rejected production proposals.
309    pub rejected: usize,
310    /// Final rejection envelope, if unweighting was used.
311    pub envelope: Option<f64>,
312    /// Method used to establish the envelope.
313    pub envelope_kind: Option<EnvelopeKind>,
314    /// Number of adaptive envelope expansions.
315    pub envelope_updates: usize,
316    /// Proven interval endpoints when the phase-space envelope was established analytically.
317    #[serde(default)]
318    pub proven_weight_interval: Option<(f64, f64)>,
319    /// Continuous proposal-coordinate count for a proven phase-space envelope.
320    #[serde(default)]
321    pub proven_continuous_dimensions: Option<usize>,
322    /// Analytical piecewise-region count for a proven phase-space envelope.
323    #[serde(default)]
324    pub proven_piecewise_regions: Option<usize>,
325    /// Adaptive subdivision count for a proven phase-space envelope.
326    #[serde(default)]
327    pub proven_subdivisions: Option<usize>,
328    /// Maximum target weight encountered.
329    pub maximum_weight: f64,
330    /// Minimum target weight encountered.
331    pub minimum_weight: f64,
332    /// Sum of generated target weights.
333    pub sum_weights: f64,
334    /// Sum of squared generated target weights.
335    pub sum_squared_weights: f64,
336    /// Random seed used for the run.
337    pub seed: u64,
338    /// Memory-derived internal event chunk size.
339    pub chunk_events: usize,
340    /// Estimated peak tracked bytes for one generation chunk.
341    pub estimated_peak_bytes: u64,
342    /// Actual tracked high-water bytes when generation used an execution pool.
343    pub actual_high_water_bytes: Option<u64>,
344}
345
346impl GenerationReport {
347    /// Returns the fraction of production proposals that were accepted.
348    pub fn acceptance_rate(&self) -> f64 {
349        if self.proposals == 0 {
350            0.0
351        } else {
352            self.produced as f64 / self.proposals as f64
353        }
354    }
355}
356
357/// A prepared model and parameters used to weight generated proposals.
358#[derive(Clone, Debug)]
359pub struct ModelEvaluator {
360    prepared: PreparedModel,
361    params: ParamValues,
362    required_scalars: HashSet<String>,
363    execution: Execution,
364}
365
366impl ModelEvaluator {
367    /// Prepares a compiled model for generation-time batch evaluation.
368    ///
369    /// # Errors
370    ///
371    /// Returns [`GenerationError`] when model preparation for the selected
372    /// execution backend fails.
373    pub fn prepare(
374        model: &CompiledModel,
375        params: ParamValues,
376        execution: &Execution,
377    ) -> GenerationResult<Self> {
378        let prepared = PreparedModel::prepare(model, execution)?;
379        let required_scalars = prepared.required_event_scalars().iter().cloned().collect();
380        Ok(Self {
381            prepared,
382            params,
383            required_scalars,
384            execution: execution.clone(),
385        })
386    }
387
388    /// Evaluate the positive-real model value for every event in a batch.
389    ///
390    /// This is useful for projecting a fitted model over weighted Monte Carlo
391    /// without regenerating events.
392    ///
393    /// # Errors
394    ///
395    /// Returns [`GenerationError`] when batch evaluation fails or a model value
396    /// is non-finite or not strictly positive.
397    pub fn evaluate_batch(&self, batch: &EventBatch) -> GenerationResult<Vec<f64>> {
398        let reduction = ReductionPlan::weighted_positive_real();
399        self.prepared
400            .evaluate_batch(&self.params, batch)?
401            .into_iter()
402            .map(|value| {
403                if !value.re.is_finite() {
404                    return Err(GenerationError::Model(format!(
405                        "positive-real model produced nonfinite value {}",
406                        value.re
407                    )));
408                }
409                reduction
410                    .apply(value)
411                    .map(|out| out.value())
412                    .map_err(|err| GenerationError::Model(err.to_string()))
413            })
414            .collect()
415    }
416}
417
418/// Generates kinematically valid events for a physics channel.
419#[derive(Debug)]
420pub struct ChannelGenerator {
421    edges: Vec<EdgePlan>,
422    vertices: Vec<VertexPlan>,
423    edge_names: Vec<String>,
424    output_indices: Vec<usize>,
425    output_names: Vec<String>,
426    root_indices: Vec<usize>,
427    scalar_sources: Vec<(String, ScalarSource)>,
428}
429
430#[derive(Clone, Debug)]
431struct EdgePlan {
432    name: String,
433    initial: Option<InitialMomentum>,
434    mass: EdgeMassPlan,
435}
436
437#[derive(Clone, Debug)]
438enum EdgeMassPlan {
439    Fixed(f64),
440    Proposed(MassProposal),
441}
442
443#[derive(Clone, Debug)]
444struct AdaptiveMassProposal {
445    base: MassProposal,
446    density: PiecewiseDensity,
447    defensive_fraction: f64,
448}
449
450impl AdaptiveMassProposal {
451    fn truncated_total(&self, minimum: f64, maximum: f64) -> f64 {
452        self.density.truncated_total(minimum, maximum)
453    }
454
455    fn adaptive_density(&self, minimum: f64, maximum: f64, mass: f64) -> f64 {
456        self.density.density(minimum, maximum, mass)
457    }
458
459    fn sample_adaptive(&self, minimum: f64, maximum: f64, rng: &mut ProposalRng) -> Option<f64> {
460        self.density.sample(minimum, maximum, rng)
461    }
462}
463
464impl AdaptiveMassProposal {
465    fn propose(
466        &self,
467        minimum: f64,
468        maximum: f64,
469        rng: &mut ProposalRng,
470    ) -> laddu_physics::LadduPhysicsResult<MassProposalResult> {
471        let adaptive_available = self.truncated_total(minimum, maximum) > 0.0;
472        let use_base = !adaptive_available || rng.uniform() < self.defensive_fraction;
473        let mass = if use_base {
474            self.base.propose(minimum, maximum, rng)?.mass
475        } else {
476            self.sample_adaptive(minimum, maximum, rng)
477                .expect("positive adaptive mass support must be sampleable")
478        };
479        let Some(base_density) = self.base.density(minimum, maximum, mass)? else {
480            return Err(LadduPhysicsError::invalid_relation(
481                "adaptive mass proposal lost access to its base density",
482            ));
483        };
484        let density = if adaptive_available {
485            self.defensive_fraction * base_density
486                + (1.0 - self.defensive_fraction) * self.adaptive_density(minimum, maximum, mass)
487        } else {
488            base_density
489        };
490        if !density.is_finite() || density <= 0.0 {
491            return Err(LadduPhysicsError::invalid_value(
492                "adaptive mass-proposal density",
493                "finite and positive",
494                density,
495            ));
496        }
497        Ok(MassProposalResult {
498            mass,
499            weight: density.recip(),
500        })
501    }
502
503    fn density(
504        &self,
505        minimum: f64,
506        maximum: f64,
507        mass: f64,
508    ) -> laddu_physics::LadduPhysicsResult<Option<f64>> {
509        let Some(base_density) = self.base.density(minimum, maximum, mass)? else {
510            return Ok(None);
511        };
512        let adaptive_available = self.truncated_total(minimum, maximum) > 0.0;
513        Ok(Some(if adaptive_available {
514            self.defensive_fraction * base_density
515                + (1.0 - self.defensive_fraction) * self.adaptive_density(minimum, maximum, mass)
516        } else {
517            base_density
518        }))
519    }
520}
521
522#[derive(Clone, Debug)]
523struct VertexPlan {
524    name: String,
525    incoming: Vec<usize>,
526    outgoing: Vec<usize>,
527    proposal: VertexProposal,
528    adaptive_decay: bool,
529}
530
531#[derive(Clone, Debug)]
532struct ProposalAdaptations {
533    masses: Vec<Option<AdaptiveMassProposal>>,
534    vertices: Vec<Option<AdaptiveTwoBodyDecay>>,
535}
536
537#[derive(Clone, Debug)]
538struct GeneratedEvent {
539    p4s: Vec<RealVec4>,
540    scalars: Vec<f64>,
541    proposal_weight: f64,
542    model_weight: f64,
543    target_weight: f64,
544    index: u64,
545}
546
547#[derive(Clone, Copy, Debug)]
548struct IntervalInitialMomentum {
549    energy: Interval,
550    momentum: [Interval; 3],
551    mass: f64,
552    weight: Interval,
553    continuous_dimensions: usize,
554    piecewise_regions: usize,
555}
556
557#[derive(Clone, Copy, Debug)]
558struct IntervalSample {
559    value: Interval,
560    weight: Interval,
561    continuous_dimensions: usize,
562    piecewise_regions: usize,
563}
564
565#[derive(Clone, Copy, Debug)]
566struct GenerationMemoryUse {
567    event_limit: usize,
568    resident_events: usize,
569    retained_output_events: usize,
570}
571
572impl ChannelGenerator {
573    fn generation_memory(
574        &self,
575        schema: &Schema,
576        budget: MemoryBudget,
577        model: Option<&ModelEvaluator>,
578        usage: GenerationMemoryUse,
579        label: &str,
580    ) -> GenerationResult<(MemoryDecision, MemoryLease)> {
581        let generated_layout = BatchLayout::new(
582            schema.n_p4s(),
583            schema.n_scalars(),
584            schema.has_weight(),
585            false,
586        );
587        let generated = generated_layout
588            .footprint(DataPrecision::F64)
589            .and_then(|footprint| {
590                footprint.checked_add(MemoryFootprint::per_event(
591                    u64::try_from(size_of::<GeneratedEvent>())
592                        .map_err(|_| laddu_memory::FootprintOverflow::Conversion)?,
593                ))
594            })
595            .and_then(|footprint| {
596                footprint.checked_add(MemoryFootprint::per_event(
597                    u64::try_from(4 * size_of::<f64>())
598                        .map_err(|_| laddu_memory::FootprintOverflow::Conversion)?,
599                ))
600            })
601            .map_err(|error| {
602                GenerationError::Runtime(laddu_runtime::RuntimeError::Data(format!(
603                    "generation working-set overflow: {error}"
604                )))
605            })?;
606        let output = generated_layout
607            .schema_footprint(DataPrecision::F64)
608            .map_err(|error| {
609                GenerationError::Runtime(laddu_runtime::RuntimeError::Data(format!(
610                    "output working-set overflow: {error}"
611                )))
612            })?;
613        let bytes_per_event = generated.checked_add(output).map_err(|error| {
614            GenerationError::Runtime(laddu_runtime::RuntimeError::Data(format!(
615                "generation working-set overflow: {error}"
616            )))
617        })?;
618        let fixed_bytes = generated
619            .checked_peak_bytes(usage.resident_events)
620            .and_then(|resident| {
621                output
622                    .checked_peak_bytes(usage.retained_output_events)
623                    .and_then(|retained| {
624                        resident
625                            .checked_add(retained)
626                            .ok_or(laddu_memory::FootprintOverflow::Addition)
627                    })
628            })
629            .map_err(|error| {
630                GenerationError::Runtime(laddu_runtime::RuntimeError::Data(format!(
631                    "generation working-set overflow: {error}"
632                )))
633            })?;
634        let state = model
635            .map(|model| model.execution.memory_state().clone())
636            .unwrap_or_else(MemoryState::current);
637        state.refresh();
638        let operation_cap = budget
639            .resolve(&state.host())
640            .map_err(laddu_runtime::RuntimeError::from)?;
641        let owned_pool;
642        let pool = if let Some(model) = model {
643            model.execution.host_memory()
644        } else {
645            owned_pool = state
646                .pool("host", budget)
647                .map_err(laddu_runtime::RuntimeError::from)?;
648            &owned_pool
649        };
650        let available = pool.remaining().min(operation_cap);
651        let decision = MemoryFitRequest {
652            label: label.into(),
653            footprint: MemoryFootprint::new(fixed_bytes, bytes_per_event.bytes_per_event),
654            available_bytes: available,
655            event_limit: usage.event_limit,
656            strategy: "memory-derived generation".into(),
657        }
658        .evaluate()
659        .map_err(laddu_runtime::RuntimeError::from)?;
660        let lease = pool
661            .reserve(decision.estimated_peak_bytes)
662            .map_err(laddu_runtime::RuntimeError::from)?;
663        if let Some(model) = model {
664            model.execution.record_memory_decision(decision.clone());
665        }
666        Ok((decision, lease))
667    }
668
669    /// Validates a channel and constructs its topological generation plan.
670    ///
671    /// # Errors
672    ///
673    /// Returns [`GenerationError`] when channel topology, edge metadata,
674    /// particle masses, output names, or vertex proposals are invalid.
675    pub fn new(channel: Channel) -> GenerationResult<Self> {
676        channel
677            .validate()
678            .map_err(|source| GenerationError::ChannelValidation { source })?;
679        let edge_names = channel
680            .edges()
681            .map(|edge| edge.name().to_owned())
682            .collect::<Vec<_>>();
683        let output_names = channel
684            .edges()
685            .filter(|edge| edge.is_output())
686            .map(|edge| edge.name().to_owned())
687            .collect::<Vec<_>>();
688        if output_names.is_empty() {
689            return Err(GenerationError::InvalidChannel(
690                "at least one edge must be marked as output".into(),
691            ));
692        }
693        const DIAGNOSTICS: [&str; 3] = [
694            "__laddu_proposal_weight",
695            "__laddu_model_weight",
696            "__laddu_target_weight",
697        ];
698        if output_names
699            .iter()
700            .any(|name| DIAGNOSTICS.contains(&name.as_str()))
701        {
702            return Err(GenerationError::InvalidChannel(
703                "an output edge uses a reserved generation-diagnostic name".into(),
704            ));
705        }
706        let root_edges = channel
707            .initial_edges()
708            .map(|edge| edge.name().to_owned())
709            .collect::<HashSet<_>>();
710        let plan = topological_plan(&channel, &root_edges)?;
711        let edge_indices = edge_names
712            .iter()
713            .enumerate()
714            .map(|(index, name)| (name.as_str(), index))
715            .collect::<HashMap<_, _>>();
716        let output_indices = output_names
717            .iter()
718            .map(|name| edge_indices[name.as_str()])
719            .collect::<Vec<_>>();
720        let root_indices = channel
721            .initial_edges()
722            .map(|edge| edge_indices[edge.name()])
723            .collect::<Vec<_>>();
724        let edges = channel
725            .edges()
726            .map(|edge| {
727                let mass = if let Some(proposal) = edge.mass_proposal() {
728                    EdgeMassPlan::Proposed(*proposal)
729                } else {
730                    let properties = edge.properties().ok_or_else(|| {
731                        GenerationError::InvalidChannel(format!(
732                            "edge `{}` has neither particle properties nor a mass proposal",
733                            edge.name()
734                        ))
735                    })?;
736                    EdgeMassPlan::Fixed(
737                        properties
738                            .mass()
739                            .map_err(|source| GenerationError::ChannelValidation { source })?,
740                    )
741                };
742                Ok(EdgePlan {
743                    name: edge.name().to_owned(),
744                    initial: edge.initial_momentum().cloned(),
745                    mass,
746                })
747            })
748            .collect::<GenerationResult<Vec<_>>>()?;
749        let channel_vertices = channel.vertices().collect::<Vec<_>>();
750        let vertices = plan
751            .into_iter()
752            .map(|vertex_index| {
753                let vertex = channel_vertices[vertex_index];
754                let incoming = vertex
755                    .incoming()
756                    .iter()
757                    .map(|name| edge_indices[name.as_str()])
758                    .collect::<Vec<_>>();
759                let outgoing = vertex
760                    .outgoing()
761                    .iter()
762                    .map(|name| edge_indices[name.as_str()])
763                    .collect::<Vec<_>>();
764                let (proposal, adaptive_decay) = match vertex.generation() {
765                    Some(proposal) => (proposal.clone(), false),
766                    None if incoming.len() == 1 && outgoing.len() == 2 => {
767                        (VertexProposal::TwoBodyDecay, true)
768                    }
769                    None => {
770                        return Err(GenerationError::InvalidChannel(format!(
771                            "vertex `{}` has no generation proposal",
772                            vertex.name()
773                        )));
774                    }
775                };
776                Ok(VertexPlan {
777                    name: vertex.name().to_owned(),
778                    incoming,
779                    outgoing,
780                    proposal,
781                    adaptive_decay,
782                })
783            })
784            .collect::<GenerationResult<Vec<_>>>()?;
785        Ok(Self {
786            edges,
787            vertices,
788            edge_names,
789            output_indices,
790            output_names,
791            root_indices,
792            scalar_sources: Vec::new(),
793        })
794    }
795
796    /// Adds a generated scalar column and returns the updated generator.
797    ///
798    /// # Errors
799    ///
800    /// Returns [`GenerationError`] when `name` is empty, reserved, duplicated,
801    /// or conflicts with an edge name.
802    pub fn with_scalar(
803        mut self,
804        name: impl Into<String>,
805        source: ScalarSource,
806    ) -> GenerationResult<Self> {
807        self.add_scalar(name, source)?;
808        Ok(self)
809    }
810
811    /// Adds a generated scalar column.
812    ///
813    /// # Errors
814    ///
815    /// Returns [`GenerationError`] when `name` is empty, reserved, duplicated,
816    /// or conflicts with an edge name.
817    pub fn add_scalar(
818        &mut self,
819        name: impl Into<String>,
820        source: ScalarSource,
821    ) -> GenerationResult<&mut Self> {
822        let name = name.into();
823        const DIAGNOSTICS: [&str; 3] = [
824            "__laddu_proposal_weight",
825            "__laddu_model_weight",
826            "__laddu_target_weight",
827        ];
828        if name.is_empty()
829            || DIAGNOSTICS.contains(&name.as_str())
830            || self.edge_names.contains(&name)
831            || self
832                .scalar_sources
833                .iter()
834                .any(|(existing, _)| existing == &name)
835        {
836            return Err(GenerationError::InvalidConfiguration(format!(
837                "scalar column name `{name}` is empty, reserved, or duplicated"
838            )));
839        }
840        self.scalar_sources.push((name, source));
841        Ok(self)
842    }
843
844    /// Returns the configured named scalar sources in output-column order.
845    pub fn scalar_sources(&self) -> impl ExactSizeIterator<Item = (&str, &ScalarSource)> {
846        self.scalar_sources
847            .iter()
848            .map(|(name, source)| (name.as_str(), source))
849    }
850
851    /// Returns the schema produced by weighted generation.
852    ///
853    /// # Errors
854    ///
855    /// Returns [`GenerationError`] if the output names do not form a valid schema.
856    pub fn weighted_output_schema(&self, diagnostics: bool) -> GenerationResult<Arc<Schema>> {
857        self.output_schema(true, diagnostics)
858    }
859
860    /// Returns the schema produced by unweighted generation.
861    ///
862    /// # Errors
863    ///
864    /// Returns [`GenerationError`] if the output names do not form a valid schema.
865    pub fn unweighted_output_schema(&self, diagnostics: bool) -> GenerationResult<Arc<Schema>> {
866        self.output_schema(false, diagnostics)
867    }
868
869    /// Prove an upper envelope for the model-less phase-space proposal weight.
870    ///
871    /// The returned interval is outward-rounded by maryada. A finite upper
872    /// endpoint is required before it can be used for rejection sampling.
873    ///
874    /// # Errors
875    ///
876    /// Returns [`GenerationError`] when a proposal domain is invalid, a
877    /// scattering vertex does not consume initial edges, or interval
878    /// propagation cannot establish a finite positive upper endpoint.
879    pub fn phase_space_envelope(&self) -> GenerationResult<ProvenEnvelopeReport> {
880        let mut initials = HashMap::new();
881        let mut weight = Interval::ONE;
882        let mut continuous_dimensions = 0_usize;
883        let mut piecewise_regions = 1_usize;
884
885        for &edge_index in &self.root_indices {
886            let edge = &self.edges[edge_index];
887            let mass = match edge.mass {
888                EdgeMassPlan::Fixed(mass) => mass,
889                EdgeMassPlan::Proposed(_) => {
890                    return Err(GenerationError::InvalidConfiguration(format!(
891                        "initial edge `{}` cannot use a generated mass",
892                        edge.name
893                    )));
894                }
895            };
896            let source = edge.initial.as_ref().ok_or_else(|| {
897                GenerationError::InvalidConfiguration(format!(
898                    "initial edge `{}` has no momentum source",
899                    edge.name
900                ))
901            })?;
902            let initial = interval_initial_momentum(source, mass)?;
903            weight *= initial.weight;
904            continuous_dimensions =
905                continuous_dimensions.saturating_add(initial.continuous_dimensions);
906            piecewise_regions = piecewise_regions.saturating_mul(initial.piecewise_regions);
907            initials.insert(edge_index, initial);
908        }
909
910        let root_s = interval_invariant_mass(
911            self.root_indices
912                .iter()
913                .map(|edge| initials[edge])
914                .collect::<Vec<_>>()
915                .as_slice(),
916        );
917        if root_s.is_empty() || !root_s.sup().is_finite() || root_s.sup() <= 0.0 {
918            return Err(GenerationError::InvalidConfiguration(format!(
919                "initial-state invariant-mass enclosure {root_s} is not finite and positive"
920            )));
921        }
922
923        // Validate the mass supports once against the root domain. The
924        // branch-and-bound evaluator below narrows these supports for each
925        // root-invariant-mass and generated-mass subdomain.
926        let mut mass_dimensions = vec![None; self.edges.len()];
927        let mut domain = vec![root_s];
928        for edge in &self.edges {
929            match edge.mass {
930                EdgeMassPlan::Fixed(mass)
931                | EdgeMassPlan::Proposed(MassProposal::Fixed { mass }) => {
932                    if mass > root_s.sup() {
933                        return Err(GenerationError::InvalidConfiguration(format!(
934                            "fixed generated mass {mass} for edge `{}` exceeds the initial invariant-mass enclosure {root_s}",
935                            edge.name
936                        )));
937                    }
938                }
939                EdgeMassPlan::Proposed(MassProposal::Uniform { low, high }) => {
940                    let support_low = low.max(0.0);
941                    let support_high = high.min(root_s.sup());
942                    if !support_low.is_finite()
943                        || !support_high.is_finite()
944                        || support_high <= support_low
945                    {
946                        return Err(GenerationError::InvalidConfiguration(format!(
947                            "uniform generated mass for edge `{}` has no finite support inside [0, {}]",
948                            edge.name,
949                            root_s.sup()
950                        )));
951                    }
952                    continuous_dimensions = continuous_dimensions.saturating_add(1);
953                }
954            }
955        }
956
957        // Re-index generated-mass coordinates by edge position. Keeping the
958        // mapping separate from the physical mass intervals prevents the
959        // latent variables from being mistaken for independent masses.
960        for (edge_index, edge) in self.edges.iter().enumerate() {
961            if matches!(
962                edge.mass,
963                EdgeMassPlan::Proposed(MassProposal::Uniform { .. })
964            ) {
965                mass_dimensions[edge_index] = Some(domain.len());
966                domain.push(Interval::new(0.0, 1.0));
967            }
968        }
969
970        let mut transfer_dimensions = vec![None; self.vertices.len()];
971        for vertex in &self.vertices {
972            let (dimensions, regions) = vertex.proposal.proven_domain_metadata();
973            continuous_dimensions = continuous_dimensions.saturating_add(dimensions);
974            piecewise_regions = piecewise_regions.saturating_mul(regions);
975        }
976        for (vertex_index, vertex) in self.vertices.iter().enumerate() {
977            if matches!(vertex.proposal, VertexProposal::TwoBodyScattering { .. }) {
978                transfer_dimensions[vertex_index] = Some(domain.len());
979                // The transfer coordinate is normalized to the configured
980                // physical support. Angular coordinates remain analytically
981                // enclosed by the physics interval formulas.
982                domain.push(Interval::new(0.0, 1.0));
983            }
984        }
985        for (_, source) in &self.scalar_sources {
986            let sample = interval_scalar_sample(source)?;
987            weight *= sample.weight;
988            continuous_dimensions =
989                continuous_dimensions.saturating_add(sample.continuous_dimensions);
990            piecewise_regions = piecewise_regions.saturating_mul(sample.piecewise_regions);
991        }
992
993        let static_weight = weight;
994        // Search the derived root invariant mass together with normalized
995        // generated-mass and transfer coordinates. Dependent physical masses
996        // and transfer values are reconstructed inside each box so the search
997        // does not treat them as independent kinematic quantities.
998        let evaluate = |domain: &Vec<Interval>| -> GenerationResult<Interval> {
999            let root_s = domain.first().copied().ok_or_else(|| {
1000                GenerationError::InvalidConfiguration(
1001                    "branch-and-bound domain has no root invariant-mass coordinate".into(),
1002                )
1003            })?;
1004            let mut weight = static_weight;
1005            let mut masses = Vec::with_capacity(self.edges.len());
1006            for (edge_index, edge) in self.edges.iter().enumerate() {
1007                match edge.mass {
1008                    EdgeMassPlan::Fixed(mass)
1009                    | EdgeMassPlan::Proposed(MassProposal::Fixed { mass }) => {
1010                        masses.push(Interval::from(mass));
1011                    }
1012                    EdgeMassPlan::Proposed(MassProposal::Uniform { low, high }) => {
1013                        let support_low = low.max(0.0);
1014                        let support_high = high.min(root_s.sup());
1015                        let minimum_width = (high.min(root_s.inf()) - support_low).max(0.0);
1016                        let maximum_width = support_high - support_low;
1017                        if maximum_width <= 0.0 {
1018                            return Ok(Interval::EMPTY);
1019                        }
1020                        let support_width = Interval::new(minimum_width, maximum_width);
1021                        weight *= support_width;
1022                        let fraction = domain[mass_dimensions[edge_index].ok_or_else(|| {
1023                            GenerationError::InvalidConfiguration(format!(
1024                                "missing branch coordinate for generated mass edge `{}`",
1025                                edge.name
1026                            ))
1027                        })?];
1028                        masses.push(Interval::from(support_low) + fraction * support_width);
1029                    }
1030                }
1031            }
1032
1033            for (vertex_index, vertex) in self.vertices.iter().enumerate() {
1034                let vertex_weight = match &vertex.proposal {
1035                    VertexProposal::TwoBodyDecay => {
1036                        if vertex.incoming.len() != 1 || vertex.outgoing.len() != 2 {
1037                            return Err(GenerationError::InvalidConfiguration(format!(
1038                                "vertex `{}` is not a one-to-two decay",
1039                                vertex.name
1040                            )));
1041                        }
1042                        proven_two_body_decay_weight(
1043                            masses[vertex.incoming[0]],
1044                            masses[vertex.outgoing[0]],
1045                            masses[vertex.outgoing[1]],
1046                        )
1047                    }
1048                    VertexProposal::TwoBodyScattering { proposal } => {
1049                        if vertex.incoming.len() != 2
1050                            || vertex.outgoing.len() != 2
1051                            || vertex
1052                                .incoming
1053                                .iter()
1054                                .any(|edge| !self.root_indices.contains(edge))
1055                        {
1056                            return Err(GenerationError::InvalidConfiguration(format!(
1057                                "proven two-body scattering at vertex `{}` currently requires two initial incoming edges",
1058                                vertex.name
1059                            )));
1060                        }
1061                        proposal.proven_weight_bound_for_transfer(
1062                            root_s,
1063                            [
1064                                (
1065                                    self.edges[vertex.incoming[0]].name.as_str(),
1066                                    masses[vertex.incoming[0]],
1067                                ),
1068                                (
1069                                    self.edges[vertex.incoming[1]].name.as_str(),
1070                                    masses[vertex.incoming[1]],
1071                                ),
1072                            ],
1073                            [
1074                                (
1075                                    self.edges[vertex.outgoing[0]].name.as_str(),
1076                                    masses[vertex.outgoing[0]],
1077                                ),
1078                                (
1079                                    self.edges[vertex.outgoing[1]].name.as_str(),
1080                                    masses[vertex.outgoing[1]],
1081                                ),
1082                            ],
1083                            domain[transfer_dimensions[vertex_index].ok_or_else(|| {
1084                                GenerationError::InvalidConfiguration(format!(
1085                                    "missing branch coordinate for transfer vertex `{}`",
1086                                    vertex.name
1087                                ))
1088                            })?],
1089                        )?
1090                    }
1091                };
1092                weight *= vertex_weight;
1093            }
1094            Ok(weight)
1095        };
1096
1097        let coarse_weight = evaluate(&domain)?;
1098        let scale = coarse_weight.sup().abs().max(1.0);
1099        let domain_tolerance = 1.0 / 1_024.0;
1100        let options = GlobalMinimizerOptions {
1101            value_tolerance: Some(scale * 1.0e-8),
1102            domain_tolerance: Some(domain_tolerance),
1103            gap_tolerance: Some(scale * 1.0e-8),
1104            max_steps: Some(1_024),
1105        };
1106        let mut minimizer = GlobalMinimizer::with_options(
1107            domain,
1108            move |domain: &Vec<Interval>| evaluate(domain).map(|value| -value),
1109            options,
1110        );
1111        let result = minimizer.solve().map_err(|error| {
1112            GenerationError::InvalidConfiguration(format!(
1113                "branch-and-bound phase-space envelope failed: {error}"
1114            ))
1115        })?;
1116        let upper = -result.minimum.inf();
1117        debug_assert!(upper <= coarse_weight.sup() * (1.0 + 1.0e-12));
1118        if result.minimum.is_empty() || !upper.is_finite() || upper <= 0.0 {
1119            return Err(GenerationError::InvalidConfiguration(format!(
1120                "phase-space proposal-weight enclosure has no finite positive upper endpoint: {}",
1121                result.minimum
1122            )));
1123        }
1124        let weight_interval = Interval::new(0.0, upper);
1125        Ok(ProvenEnvelopeReport {
1126            weight_interval,
1127            maximum_weight: upper,
1128            continuous_dimensions,
1129            piecewise_regions,
1130            subdivisions: result.branched,
1131        })
1132    }
1133
1134    /// Generates weighted events and writes them to `sink`.
1135    ///
1136    /// # Errors
1137    ///
1138    /// Returns [`GenerationError`] when configuration, proposal generation,
1139    /// model evaluation, batch construction, or sink I/O fails.
1140    pub fn generate_weighted_to(
1141        &self,
1142        config: WeightedConfig,
1143        model: Option<&ModelEvaluator>,
1144        sink: &mut dyn EventSink,
1145    ) -> GenerationResult<GenerationReport> {
1146        validate_common(config.events)?;
1147        let schema = self.output_schema(true, config.diagnostics)?;
1148        let (decision, _memory) = self.generation_memory(
1149            &schema,
1150            config.memory,
1151            model,
1152            GenerationMemoryUse {
1153                event_limit: config.events,
1154                resident_events: 0,
1155                retained_output_events: if sink.retains_batches() {
1156                    config.events
1157                } else {
1158                    0
1159                },
1160            },
1161            "weighted generation",
1162        )?;
1163        sink.begin(Arc::clone(&schema), WritePlan::default())?;
1164        let result = (|| -> GenerationResult<GenerationReport> {
1165            let mut report = report(config.events, config.seed, &decision);
1166            let work_batch = decision.chunk_events.max(1);
1167            for start in (0..config.events).step_by(work_batch) {
1168                let count = work_batch.min(config.events - start);
1169                let mut events = self.propose_range(start as u64, count, config.seed, 0)?;
1170                self.apply_model(&mut events, model)?;
1171                update_report(&mut report, &events);
1172                report.proposals += events.len();
1173                report.produced += events.len();
1174                for chunk in events.chunks(decision.chunk_events.max(1)) {
1175                    let batch =
1176                        self.output_batch(chunk, Arc::clone(&schema), true, config.diagnostics)?;
1177                    sink.write_batch(&batch)?;
1178                }
1179            }
1180            sink.finish()?;
1181            Ok(report)
1182        })();
1183
1184        if result.is_err() {
1185            // The generation error is authoritative even if backend cleanup
1186            // also fails. Aborted files are intentionally left in place.
1187            let _ = sink.abort();
1188        }
1189
1190        result
1191    }
1192
1193    /// Generates rejection-sampled unweighted events and writes them to `sink`.
1194    ///
1195    /// # Errors
1196    ///
1197    /// Returns [`GenerationError`] when configuration, envelope estimation,
1198    /// proposal generation, model evaluation, rejection sampling, or sink I/O
1199    /// fails.
1200    pub fn generate_unweighted_to(
1201        &self,
1202        config: UnweightedConfig,
1203        model: Option<&ModelEvaluator>,
1204        sink: &mut dyn EventSink,
1205    ) -> GenerationResult<GenerationReport> {
1206        validate_common(config.events)?;
1207        if matches!(config.envelope, EnvelopeMode::ProvenPhaseSpace) && model.is_some() {
1208            return Err(GenerationError::InvalidConfiguration(
1209                "the proven phase-space envelope is valid only when no model is supplied".into(),
1210            ));
1211        }
1212        if matches!(config.envelope, EnvelopeMode::ProvenPhaseSpace)
1213            && !matches!(config.envelope_overflow, EnvelopeOverflow::Error)
1214        {
1215            return Err(GenerationError::InvalidConfiguration(
1216                "the proven phase-space envelope requires envelope_overflow=Error".into(),
1217            ));
1218        }
1219        if config
1220            .max_proposals
1221            .is_some_and(|max_proposals| max_proposals < config.events)
1222        {
1223            return Err(GenerationError::InvalidConfiguration(
1224                "max_proposals must be at least the requested event count".into(),
1225            ));
1226        }
1227        let mut adaptations = None;
1228        let schema = self.output_schema(false, config.diagnostics)?;
1229        let pilot_limit = match config.envelope {
1230            EnvelopeMode::Pilot { proposals, .. } => proposals,
1231            EnvelopeMode::Strict { .. } | EnvelopeMode::ProvenPhaseSpace => 0,
1232        };
1233        let (decision, _memory) = self.generation_memory(
1234            &schema,
1235            config.memory,
1236            model,
1237            GenerationMemoryUse {
1238                event_limit: config.events.max(pilot_limit),
1239                resident_events: if matches!(
1240                    config.envelope_overflow,
1241                    EnvelopeOverflow::Grow { .. }
1242                ) {
1243                    config.events
1244                } else {
1245                    0
1246                },
1247                retained_output_events: if sink.retains_batches() {
1248                    config.events
1249                } else {
1250                    0
1251                },
1252            },
1253            "unweighted generation",
1254        )?;
1255        if pilot_limit > decision.chunk_events {
1256            return Err(GenerationError::InvalidConfiguration(format!(
1257                "pilot sample requires {pilot_limit} simultaneously resident proposals, but the \
1258                 memory budget fits {}; increase the budget or reduce pilot_proposals",
1259                decision.chunk_events
1260            )));
1261        }
1262        let mut proven_report = None;
1263        let (mut bound, kind, pilot_count) = match config.envelope {
1264            EnvelopeMode::Strict { max_weight } => {
1265                validate_bound(max_weight)?;
1266                (max_weight, EnvelopeKind::Strict, 0)
1267            }
1268            EnvelopeMode::Pilot {
1269                proposals,
1270                safety_factor,
1271            } => {
1272                if proposals == 0 || !safety_factor.is_finite() || safety_factor <= 1.0 {
1273                    return Err(GenerationError::InvalidConfiguration("pilot proposals must be nonzero and safety_factor must be finite and greater than one".into()));
1274                }
1275                let mut adaptation_pilot = self.propose_range(0, proposals, config.seed, 1)?;
1276                self.apply_model(&mut adaptation_pilot, model)?;
1277                let learned = self.learn_mass_adaptations(&adaptation_pilot)?;
1278                let has_adaptation = learned.masses.iter().any(Option::is_some)
1279                    || learned.vertices.iter().any(Option::is_some);
1280                let mut envelope_pilot = if has_adaptation {
1281                    drop(adaptation_pilot);
1282                    self.propose_range_with_adaptation(
1283                        0,
1284                        proposals,
1285                        config.seed,
1286                        2,
1287                        Some(&learned),
1288                    )?
1289                } else {
1290                    adaptation_pilot
1291                };
1292                if has_adaptation {
1293                    self.apply_model(&mut envelope_pilot, model)?;
1294                    adaptations = Some(learned);
1295                }
1296                let observed = envelope_pilot
1297                    .iter()
1298                    .map(|event| event.target_weight)
1299                    .fold(0.0, f64::max);
1300                (
1301                    observed * safety_factor,
1302                    EnvelopeKind::Pilot,
1303                    proposals * if has_adaptation { 2 } else { 1 },
1304                )
1305            }
1306            EnvelopeMode::ProvenPhaseSpace => {
1307                let report = self.phase_space_envelope()?;
1308                let maximum = report.maximum_weight;
1309                proven_report = Some(report);
1310                (maximum, EnvelopeKind::ProvenPhaseSpace, 0)
1311            }
1312        };
1313        if let EnvelopeOverflow::Grow { safety_factor } = config.envelope_overflow
1314            && (!safety_factor.is_finite() || safety_factor <= 1.0)
1315        {
1316            return Err(GenerationError::InvalidConfiguration(
1317                "envelope growth safety_factor must be finite and greater than one".into(),
1318            ));
1319        }
1320        sink.begin(Arc::clone(&schema), WritePlan::default())?;
1321        let result = (|| -> GenerationResult<GenerationReport> {
1322            let mut report = report(config.events, config.seed, &decision);
1323            report.envelope = Some(bound);
1324            report.envelope_kind = Some(kind);
1325            report.pilot_proposals = pilot_count;
1326            if let Some(proven) = proven_report {
1327                report.proven_weight_interval = Some(proven.weight_interval.bounds());
1328                report.proven_continuous_dimensions = Some(proven.continuous_dimensions);
1329                report.proven_piecewise_regions = Some(proven.piecewise_regions);
1330                report.proven_subdivisions = Some(proven.subdivisions);
1331            }
1332            let mut proposal_index = 0_usize;
1333            let mut buffered = Vec::new();
1334            let work_batch = decision.chunk_events.max(1);
1335            while report.produced < config.events
1336                && config
1337                    .max_proposals
1338                    .is_none_or(|max_proposals| proposal_index < max_proposals)
1339            {
1340                let count = config.max_proposals.map_or(work_batch, |max_proposals| {
1341                    work_batch.min(max_proposals - proposal_index)
1342                });
1343                let mut events = self.propose_range_with_adaptation(
1344                    proposal_index as u64,
1345                    count,
1346                    config.seed,
1347                    0,
1348                    adaptations.as_ref(),
1349                )?;
1350                self.apply_model(&mut events, model)?;
1351                let remaining_before_overflow = config.events - report.produced;
1352                if let Some(last_needed) = events
1353                    .iter()
1354                    .enumerate()
1355                    .filter(|(_, event)| {
1356                        acceptance_uniform(config.seed, event.index) * bound <= event.target_weight
1357                    })
1358                    .nth(remaining_before_overflow - 1)
1359                    .map(|(position, _)| position + 1)
1360                    && !events[..last_needed]
1361                        .iter()
1362                        .any(|event| event.target_weight > bound)
1363                {
1364                    events.truncate(last_needed);
1365                }
1366                update_report(&mut report, &events);
1367                if let Some(overflow) = events
1368                    .iter()
1369                    .filter(|event| event.target_weight > bound)
1370                    .max_by(|a, b| a.target_weight.total_cmp(&b.target_weight))
1371                {
1372                    match config.envelope_overflow {
1373                        EnvelopeOverflow::Error => {
1374                            return Err(GenerationError::EnvelopeOverflow {
1375                                index: overflow.index,
1376                                weight: overflow.target_weight,
1377                                envelope: bound,
1378                            });
1379                        }
1380                        EnvelopeOverflow::Grow { safety_factor } => {
1381                            bound = overflow.target_weight * safety_factor;
1382                            validate_bound(bound)?;
1383                            report.envelope = Some(bound);
1384                            report.envelope_updates += 1;
1385                            buffered.retain(|event: &GeneratedEvent| {
1386                                acceptance_uniform(config.seed, event.index) * bound
1387                                    <= event.target_weight
1388                            });
1389                            report.produced = buffered.len();
1390                        }
1391                    }
1392                }
1393                let remaining = config.events - report.produced;
1394                let proposal_count = events.len();
1395                let accepted = events
1396                    .into_iter()
1397                    .filter(|event| {
1398                        acceptance_uniform(config.seed, event.index) * bound <= event.target_weight
1399                    })
1400                    .take(remaining)
1401                    .collect::<Vec<_>>();
1402                report.proposals += proposal_count;
1403                report.produced += accepted.len();
1404                report.rejected = report.proposals - report.produced;
1405                proposal_index += proposal_count;
1406                match config.envelope_overflow {
1407                    EnvelopeOverflow::Error if !accepted.is_empty() => {
1408                        for chunk in accepted.chunks(decision.chunk_events.max(1)) {
1409                            let batch = self.output_batch(
1410                                chunk,
1411                                Arc::clone(&schema),
1412                                false,
1413                                config.diagnostics,
1414                            )?;
1415                            sink.write_batch(&batch)?;
1416                        }
1417                    }
1418                    EnvelopeOverflow::Grow { .. } => {
1419                        buffered.extend(accepted);
1420                        report.produced = buffered.len();
1421                        report.rejected = report.proposals - report.produced;
1422                    }
1423                    EnvelopeOverflow::Error => {}
1424                }
1425            }
1426            if report.produced != config.events {
1427                return Err(GenerationError::Exhausted {
1428                    requested: config.events,
1429                    accepted: report.produced,
1430                    proposals: report.proposals,
1431                });
1432            }
1433            if matches!(config.envelope_overflow, EnvelopeOverflow::Grow { .. }) {
1434                for events in buffered.chunks(decision.chunk_events.max(1)) {
1435                    let batch =
1436                        self.output_batch(events, Arc::clone(&schema), false, config.diagnostics)?;
1437                    sink.write_batch(&batch)?;
1438                }
1439            }
1440            sink.finish()?;
1441            Ok(report)
1442        })();
1443
1444        if result.is_err() {
1445            let _ = sink.abort();
1446        }
1447
1448        result
1449    }
1450
1451    /// Generates weighted events into an in-memory dataset.
1452    ///
1453    /// # Errors
1454    ///
1455    /// Returns [`GenerationError`] when weighted generation fails or generated
1456    /// batches cannot form a valid dataset.
1457    pub fn generate_weighted_dataset(
1458        &self,
1459        config: WeightedConfig,
1460        model: Option<&ModelEvaluator>,
1461    ) -> GenerationResult<(Dataset, GenerationReport)> {
1462        let mut sink = MemorySink::new();
1463        let report = self.generate_weighted_to(config, model, &mut sink)?;
1464        Ok((Dataset::from_batches(sink.into_batches())?, report))
1465    }
1466
1467    /// Generates unweighted events into an in-memory dataset.
1468    ///
1469    /// # Errors
1470    ///
1471    /// Returns [`GenerationError`] when unweighted generation fails or
1472    /// generated batches cannot form a valid dataset.
1473    pub fn generate_unweighted_dataset(
1474        &self,
1475        config: UnweightedConfig,
1476        model: Option<&ModelEvaluator>,
1477    ) -> GenerationResult<(Dataset, GenerationReport)> {
1478        let mut sink = MemorySink::new();
1479        let report = self.generate_unweighted_to(config, model, &mut sink)?;
1480        Ok((Dataset::from_batches(sink.into_batches())?, report))
1481    }
1482
1483    fn propose_range(
1484        &self,
1485        start: u64,
1486        count: usize,
1487        seed: u64,
1488        stream: u64,
1489    ) -> GenerationResult<Vec<GeneratedEvent>> {
1490        self.propose_range_with_adaptation(start, count, seed, stream, None)
1491    }
1492
1493    fn propose_range_with_adaptation(
1494        &self,
1495        start: u64,
1496        count: usize,
1497        seed: u64,
1498        stream: u64,
1499        adaptations: Option<&ProposalAdaptations>,
1500    ) -> GenerationResult<Vec<GeneratedEvent>> {
1501        (0..count)
1502            .into_par_iter()
1503            .map(|offset| self.propose(start + offset as u64, seed, stream, adaptations))
1504            .collect()
1505    }
1506
1507    fn propose(
1508        &self,
1509        index: u64,
1510        seed: u64,
1511        stream: u64,
1512        adaptations: Option<&ProposalAdaptations>,
1513    ) -> GenerationResult<GeneratedEvent> {
1514        let mut rng = ProposalRng::new(derive_seed(seed, stream, index, 0));
1515        let mut p4s = vec![RealVec4::new(0.0, 0.0, 0.0, 0.0); self.edges.len()];
1516        let mut proposal_weight = 1.0;
1517        for &edge_index in &self.root_indices {
1518            let edge = &self.edges[edge_index];
1519            let source = edge
1520                .initial
1521                .as_ref()
1522                .ok_or_else(|| GenerationError::InitialState {
1523                    index,
1524                    source: LadduPhysicsError::invalid_relation(format!(
1525                        "initial edge `{}` has no momentum source",
1526                        edge.name
1527                    )),
1528                })?;
1529            let sampled = source
1530                .sample_prevalidated(
1531                    match edge.mass {
1532                        EdgeMassPlan::Fixed(mass) => mass,
1533                        EdgeMassPlan::Proposed(_) => {
1534                            return Err(GenerationError::InitialState {
1535                                index,
1536                                source: LadduPhysicsError::invalid_relation(format!(
1537                                    "initial edge `{}` cannot use a generated mass",
1538                                    edge.name
1539                                )),
1540                            });
1541                        }
1542                    },
1543                    &mut rng,
1544                )
1545                .map_err(|source| GenerationError::InitialState { index, source })?;
1546            p4s[edge_index] = sampled.p4;
1547            proposal_weight *= sampled.weight;
1548        }
1549        if !proposal_weight.is_finite() || proposal_weight <= 0.0 {
1550            return Err(GenerationError::InitialState {
1551                index,
1552                source: LadduPhysicsError::invalid_value(
1553                    "initial-state proposal weight",
1554                    "finite and positive",
1555                    proposal_weight,
1556                ),
1557            });
1558        }
1559        let total_initial: RealVec4 = self.root_indices.iter().map(|&edge| p4s[edge]).sum();
1560        let maximum_mass = total_initial
1561            .m()
1562            .map_err(|source| GenerationError::Kinematics { index, source })?;
1563        let mut masses = Vec::with_capacity(self.edges.len());
1564        for (edge_index, edge) in self.edges.iter().enumerate() {
1565            let mass = if let EdgeMassPlan::Proposed(base_proposal) = &edge.mass {
1566                let mut mass_rng =
1567                    ProposalRng::new(derive_seed(seed, stream, index, 1 + edge_index as u64));
1568                let result = if let Some(proposal) =
1569                    adaptations.and_then(|adaptations| adaptations.masses[edge_index].as_ref())
1570                {
1571                    proposal.propose(0.0, maximum_mass, &mut mass_rng)
1572                } else {
1573                    base_proposal.propose(0.0, maximum_mass, &mut mass_rng)
1574                }
1575                .map_err(|source| GenerationError::MassProposal {
1576                    index,
1577                    edge: edge.name.clone(),
1578                    source,
1579                })?;
1580                proposal_weight *= result.weight;
1581                result.mass
1582            } else if let EdgeMassPlan::Fixed(mass) = edge.mass {
1583                mass
1584            } else {
1585                unreachable!()
1586            };
1587            if !mass.is_finite() || mass < 0.0 {
1588                return Err(GenerationError::MassProposal {
1589                    index,
1590                    edge: edge.name.clone(),
1591                    source: LadduPhysicsError::invalid_value(
1592                        "mass",
1593                        "finite and nonnegative",
1594                        mass,
1595                    ),
1596                });
1597            }
1598            masses.push(mass);
1599        }
1600        for &edge in &self.root_indices {
1601            validate_p4(&self.edges[edge].name, p4s[edge], masses[edge], index)?;
1602        }
1603        for (step, vertex) in self.vertices.iter().enumerate() {
1604            let incoming = vertex
1605                .incoming
1606                .iter()
1607                .map(|&edge| NamedMomentum {
1608                    name: &self.edges[edge].name,
1609                    p4: p4s[edge],
1610                })
1611                .collect::<SmallVec<[_; 2]>>();
1612            let outgoing = vertex
1613                .outgoing
1614                .iter()
1615                .map(|&edge| NamedMass {
1616                    name: &self.edges[edge].name,
1617                    mass: masses[edge],
1618                })
1619                .collect::<SmallVec<[_; 2]>>();
1620            let mut vertex_rng =
1621                ProposalRng::new(derive_seed(seed, stream, index, 10_000 + step as u64));
1622            let result = if let Some(proposal) =
1623                adaptations.and_then(|adaptations| adaptations.vertices[step].as_ref())
1624            {
1625                proposal.propose(&incoming, &outgoing, &mut vertex_rng)
1626            } else {
1627                vertex
1628                    .proposal
1629                    .propose(&incoming, &outgoing, &mut vertex_rng)
1630            }
1631            .map_err(|source| GenerationError::VertexProposal {
1632                index,
1633                vertex: vertex.name.clone(),
1634                source,
1635            })?;
1636            if result.outgoing.len() != vertex.outgoing.len()
1637                || !result.weight.is_finite()
1638                || result.weight <= 0.0
1639            {
1640                return Err(GenerationError::VertexProposal {
1641                    index,
1642                    vertex: vertex.name.clone(),
1643                    source: LadduPhysicsError::invalid_relation(format!(
1644                        "proposal returned {} outgoing momenta and weight {}",
1645                        result.outgoing.len(),
1646                        result.weight
1647                    )),
1648                });
1649            }
1650            proposal_weight *= result.weight;
1651            for (&edge, p4) in vertex.outgoing.iter().zip(result.outgoing) {
1652                validate_p4(&self.edges[edge].name, p4, masses[edge], index)?;
1653                p4s[edge] = p4;
1654            }
1655            validate_indexed_conservation(vertex, &p4s, index)?;
1656        }
1657        let mut scalars = Vec::with_capacity(self.scalar_sources.len());
1658        for (scalar_index, (column, source)) in self.scalar_sources.iter().enumerate() {
1659            let mut scalar_rng = ProposalRng::new(derive_seed(
1660                seed,
1661                stream,
1662                index,
1663                20_000 + scalar_index as u64,
1664            ));
1665            let result = source.sample(&mut scalar_rng).map_err(|source| {
1666                GenerationError::ScalarProposal {
1667                    index,
1668                    column: column.clone(),
1669                    source,
1670                }
1671            })?;
1672            proposal_weight *= result.weight;
1673            scalars.push(result.value);
1674        }
1675        if !proposal_weight.is_finite() || proposal_weight <= 0.0 {
1676            return Err(GenerationError::Kinematics {
1677                index,
1678                source: LadduPhysicsError::invalid_value(
1679                    "accumulated proposal weight",
1680                    "finite and positive",
1681                    proposal_weight,
1682                ),
1683            });
1684        }
1685        Ok(GeneratedEvent {
1686            p4s,
1687            scalars,
1688            proposal_weight,
1689            model_weight: 1.0,
1690            target_weight: proposal_weight,
1691            index,
1692        })
1693    }
1694
1695    fn apply_model(
1696        &self,
1697        events: &mut [GeneratedEvent],
1698        model: Option<&ModelEvaluator>,
1699    ) -> GenerationResult<()> {
1700        if let Some(model) = model {
1701            let scalar_names = self
1702                .scalar_sources
1703                .iter()
1704                .map(|(name, _)| name.as_str())
1705                .collect::<Vec<_>>();
1706            let available = scalar_names.iter().copied().collect::<HashSet<_>>();
1707            if let Some(missing) = model
1708                .required_scalars
1709                .iter()
1710                .find(|name| !available.contains(name.as_str()))
1711            {
1712                return Err(GenerationError::Model(format!(
1713                    "model requires scalar column `{missing}`, but the generator has no source for it"
1714                )));
1715            }
1716            let schema = Arc::new(Schema::new(
1717                self.edge_names.iter().map(String::as_str),
1718                scalar_names,
1719                false,
1720            )?);
1721            events
1722                .par_chunks_mut(4_096)
1723                .try_for_each(|events| -> GenerationResult<()> {
1724                    let columns = (0..self.edge_names.len())
1725                        .map(|edge| {
1726                            Arc::<[RealVec4]>::from(
1727                                events
1728                                    .iter()
1729                                    .map(|event| event.p4s[edge])
1730                                    .collect::<Vec<_>>(),
1731                            )
1732                        })
1733                        .collect();
1734                    let scalar_columns = (0..self.scalar_sources.len())
1735                        .map(|column| {
1736                            Arc::<[f64]>::from(
1737                                events
1738                                    .iter()
1739                                    .map(|event| event.scalars[column])
1740                                    .collect::<Vec<_>>(),
1741                            )
1742                        })
1743                        .collect();
1744                    let batch =
1745                        EventBatch::new(Arc::clone(&schema), columns, scalar_columns, None)?;
1746                    let weights = model.evaluate_batch(&batch)?;
1747                    for (event, weight) in events.iter_mut().zip(weights) {
1748                        event.model_weight = weight;
1749                        event.target_weight = event.proposal_weight * weight;
1750                    }
1751                    Ok(())
1752                })?;
1753        }
1754        Ok(())
1755    }
1756
1757    fn learn_mass_adaptations(
1758        &self,
1759        pilot: &[GeneratedEvent],
1760    ) -> GenerationResult<ProposalAdaptations> {
1761        const BINS: usize = 64;
1762        const DEFENSIVE_FRACTION: f64 = 0.2;
1763        const MINIMUM_GAIN: f64 = 1.02;
1764
1765        let mut best: Option<(usize, f64, AdaptiveMassProposal)> = None;
1766        let old_sum: f64 = pilot.iter().map(|event| event.target_weight).sum();
1767        let old_max = pilot
1768            .iter()
1769            .map(|event| event.target_weight)
1770            .fold(0.0, f64::max);
1771        if pilot.is_empty() || old_sum <= 0.0 || old_max <= 0.0 {
1772            return Ok(ProposalAdaptations {
1773                masses: vec![None; self.edges.len()],
1774                vertices: vec![None; self.vertices.len()],
1775            });
1776        }
1777        let old_efficiency = old_sum / (pilot.len() as f64 * old_max);
1778
1779        for (edge_index, edge) in self.edges.iter().enumerate() {
1780            let EdgeMassPlan::Proposed(base) = &edge.mass else {
1781                continue;
1782            };
1783            let masses = pilot
1784                .iter()
1785                .map(|event| {
1786                    event.p4s[edge_index]
1787                        .m()
1788                        .map_err(|source| GenerationError::Kinematics {
1789                            index: event.index,
1790                            source,
1791                        })
1792                })
1793                .collect::<GenerationResult<Vec<_>>>()?;
1794            let low = masses.iter().copied().fold(f64::INFINITY, f64::min);
1795            let high = masses.iter().copied().fold(f64::NEG_INFINITY, f64::max);
1796            if !low.is_finite() || !high.is_finite() || high <= low {
1797                continue;
1798            }
1799            let width = (high - low) / BINS as f64;
1800            let mut counts = vec![0.0; BINS];
1801            for (&mass, event) in masses.iter().zip(pilot) {
1802                let bin = (((mass - low) / width) as usize).min(BINS - 1);
1803                counts[bin] += event.target_weight;
1804            }
1805            if counts.iter().filter(|count| **count > 0.0).count() < 2 {
1806                continue;
1807            }
1808            let candidate = AdaptiveMassProposal {
1809                base: *base,
1810                density: PiecewiseDensity::uniform(low, high, counts.into())
1811                    .map_err(GenerationError::Physics)?,
1812                defensive_fraction: DEFENSIVE_FRACTION,
1813            };
1814            let mut new_sum = 0.0;
1815            let mut new_max: f64 = 0.0;
1816            let mut density_available = true;
1817            for (event, &mass) in pilot.iter().zip(&masses) {
1818                let total_initial: RealVec4 =
1819                    self.root_indices.iter().map(|&edge| event.p4s[edge]).sum();
1820                let maximum = total_initial
1821                    .m()
1822                    .map_err(|source| GenerationError::Kinematics {
1823                        index: event.index,
1824                        source,
1825                    })?;
1826                let Some(base_density) = base.density(0.0, maximum, mass)? else {
1827                    density_available = false;
1828                    break;
1829                };
1830                let Some(new_density) = candidate.density(0.0, maximum, mass)? else {
1831                    density_available = false;
1832                    break;
1833                };
1834                if base_density <= 0.0 || new_density <= 0.0 {
1835                    density_available = false;
1836                    break;
1837                }
1838                let adjusted = event.target_weight * base_density / new_density;
1839                new_sum += adjusted;
1840                new_max = new_max.max(adjusted);
1841            }
1842            if !density_available || new_sum <= 0.0 || new_max <= 0.0 {
1843                continue;
1844            }
1845            let new_efficiency = new_sum / (pilot.len() as f64 * new_max);
1846            let gain = new_efficiency / old_efficiency;
1847            if gain >= MINIMUM_GAIN
1848                && best
1849                    .as_ref()
1850                    .is_none_or(|(_, best_gain, _)| gain > *best_gain)
1851            {
1852                best = Some((edge_index, gain, candidate));
1853            }
1854        }
1855
1856        let mut masses = vec![None; self.edges.len()];
1857        if let Some((edge, _, proposal)) = best {
1858            masses[edge] = Some(proposal);
1859        }
1860        let mut vertices = vec![None; self.vertices.len()];
1861        for (vertex_index, vertex) in self.vertices.iter().enumerate() {
1862            if !vertex.adaptive_decay {
1863                continue;
1864            }
1865            let mut counts = vec![0.0; 32];
1866            let mut costhetas = Vec::with_capacity(pilot.len());
1867            for event in pilot {
1868                let parent = event.p4s[vertex.incoming[0]];
1869                let inverse_beta =
1870                    -parent
1871                        .beta()
1872                        .map_err(|source| GenerationError::Kinematics {
1873                            index: event.index,
1874                            source,
1875                        })?;
1876                let rest = event.p4s[vertex.outgoing[0]].boost(&inverse_beta);
1877                let momentum = rest.vec3();
1878                let magnitude = momentum.mag();
1879                if magnitude <= 0.0 || !magnitude.is_finite() {
1880                    costhetas.push(None);
1881                    continue;
1882                }
1883                let costheta = (momentum.pz() / magnitude).clamp(-1.0, 1.0);
1884                costhetas.push(Some(costheta));
1885                let bin =
1886                    (((costheta + 1.0) * 0.5 * counts.len() as f64) as usize).min(counts.len() - 1);
1887                counts[bin] += event.target_weight;
1888            }
1889            let total: f64 = counts.iter().sum();
1890            let width = 2.0 / counts.len() as f64;
1891            let mut new_sum = 0.0;
1892            let mut new_max: f64 = 0.0;
1893            for (event, costheta) in pilot.iter().zip(&costhetas) {
1894                let Some(costheta) = costheta else {
1895                    continue;
1896                };
1897                let bin = (((costheta + 1.0) / width) as usize).min(counts.len() - 1);
1898                let learned_density = counts[bin] / (total * width);
1899                let density =
1900                    DEFENSIVE_FRACTION * 0.5 + (1.0 - DEFENSIVE_FRACTION) * learned_density;
1901                let adjusted = event.target_weight * 0.5 / density;
1902                new_sum += adjusted;
1903                new_max = new_max.max(adjusted);
1904            }
1905            let new_efficiency = new_sum / (pilot.len() as f64 * new_max);
1906            if counts.iter().filter(|count| **count > 0.0).count() >= 2
1907                && new_efficiency / old_efficiency >= MINIMUM_GAIN
1908            {
1909                vertices[vertex_index] = Some(
1910                    AdaptiveTwoBodyDecay::new(counts.into(), DEFENSIVE_FRACTION)
1911                        .map_err(GenerationError::Physics)?,
1912                );
1913            }
1914        }
1915        Ok(ProposalAdaptations { masses, vertices })
1916    }
1917
1918    fn output_schema(&self, weighted: bool, diagnostics: bool) -> GenerationResult<Arc<Schema>> {
1919        let mut scalars = self
1920            .scalar_sources
1921            .iter()
1922            .map(|(name, _)| name.as_str())
1923            .collect::<Vec<_>>();
1924        if diagnostics {
1925            scalars.extend([
1926                "__laddu_proposal_weight",
1927                "__laddu_model_weight",
1928                "__laddu_target_weight",
1929            ]);
1930        }
1931        Ok(Arc::new(Schema::new(
1932            self.output_names.iter().map(String::as_str),
1933            scalars,
1934            weighted,
1935        )?))
1936    }
1937
1938    fn output_batch(
1939        &self,
1940        events: &[GeneratedEvent],
1941        schema: Arc<Schema>,
1942        weighted: bool,
1943        diagnostics: bool,
1944    ) -> GenerationResult<EventBatch> {
1945        let p4s = self
1946            .output_indices
1947            .iter()
1948            .map(|&edge| {
1949                Arc::<[RealVec4]>::from(
1950                    events
1951                        .iter()
1952                        .map(|event| event.p4s[edge])
1953                        .collect::<Vec<_>>(),
1954                )
1955            })
1956            .collect();
1957        let mut scalars = (0..self.scalar_sources.len())
1958            .map(|column| {
1959                Arc::<[f64]>::from(
1960                    events
1961                        .iter()
1962                        .map(|event| event.scalars[column])
1963                        .collect::<Vec<_>>(),
1964                )
1965            })
1966            .collect::<Vec<_>>();
1967        if diagnostics {
1968            scalars.extend([
1969                Arc::<[f64]>::from(
1970                    events
1971                        .iter()
1972                        .map(|event| event.proposal_weight)
1973                        .collect::<Vec<_>>(),
1974                ),
1975                Arc::<[f64]>::from(
1976                    events
1977                        .iter()
1978                        .map(|event| event.model_weight)
1979                        .collect::<Vec<_>>(),
1980                ),
1981                Arc::<[f64]>::from(
1982                    events
1983                        .iter()
1984                        .map(|event| event.target_weight)
1985                        .collect::<Vec<_>>(),
1986                ),
1987            ]);
1988        }
1989        let weights = weighted.then(|| {
1990            Arc::<[f64]>::from(
1991                events
1992                    .iter()
1993                    .map(|event| event.target_weight)
1994                    .collect::<Vec<_>>(),
1995            )
1996        });
1997        Ok(EventBatch::new(schema, p4s, scalars, weights)?)
1998    }
1999}
2000
2001impl TryFrom<Channel> for ChannelGenerator {
2002    type Error = GenerationError;
2003
2004    fn try_from(channel: Channel) -> Result<Self, Self::Error> {
2005        Self::new(channel)
2006    }
2007}
2008
2009fn topological_plan(channel: &Channel, roots: &HashSet<String>) -> GenerationResult<Vec<usize>> {
2010    let vertices = channel.vertices().collect::<Vec<_>>();
2011    let declared = channel
2012        .edges()
2013        .map(|edge| edge.name())
2014        .collect::<HashSet<_>>();
2015    let mut consumers = HashMap::<&str, usize>::new();
2016    for vertex in &vertices {
2017        for edge in vertex.incoming().iter().chain(vertex.outgoing()) {
2018            if !declared.contains(edge.as_str()) {
2019                return Err(GenerationError::InvalidChannel(format!(
2020                    "vertex `{}` references undeclared edge `{edge}`",
2021                    vertex.name()
2022                )));
2023            }
2024        }
2025        for edge in vertex.incoming() {
2026            *consumers.entry(edge).or_default() += 1;
2027        }
2028    }
2029    if let Some((edge, count)) = consumers.iter().find(|(_, count)| **count > 1) {
2030        return Err(GenerationError::InvalidChannel(format!(
2031            "edge `{edge}` is consumed by {count} vertices"
2032        )));
2033    }
2034    let mut available = roots.clone();
2035    let mut remaining = (0..vertices.len()).collect::<Vec<_>>();
2036    let mut plan = Vec::with_capacity(vertices.len());
2037    while !remaining.is_empty() {
2038        let Some(position) = remaining.iter().position(|index| {
2039            vertices[*index]
2040                .incoming()
2041                .iter()
2042                .all(|edge| available.contains(edge))
2043        }) else {
2044            return Err(GenerationError::InvalidChannel(
2045                "channel is cyclic or has an edge without a producer".into(),
2046            ));
2047        };
2048        let index = remaining.remove(position);
2049        let vertex = vertices[index];
2050        if vertex.incoming().is_empty() || vertex.outgoing().is_empty() {
2051            return Err(GenerationError::InvalidChannel(format!(
2052                "vertex `{}` must have incoming and outgoing edges",
2053                vertex.name()
2054            )));
2055        }
2056        if vertex.generation().is_none()
2057            && !(vertex.incoming().len() == 1 && vertex.outgoing().len() == 2)
2058        {
2059            return Err(GenerationError::InvalidChannel(format!(
2060                "vertex `{}` has no generation proposal and is not a two-body decay",
2061                vertex.name()
2062            )));
2063        }
2064        for edge in vertex.outgoing() {
2065            if !available.insert(edge.clone()) {
2066                return Err(GenerationError::InvalidChannel(format!(
2067                    "edge `{edge}` is produced more than once"
2068                )));
2069            }
2070        }
2071        plan.push(index);
2072    }
2073    Ok(plan)
2074}
2075
2076fn interval_scalar_sample(source: &ScalarSource) -> GenerationResult<IntervalSample> {
2077    match source {
2078        ScalarSource::Constant(value) => {
2079            source.support()?;
2080            Ok(IntervalSample {
2081                value: Interval::from(*value),
2082                weight: Interval::ONE,
2083                continuous_dimensions: 0,
2084                piecewise_regions: 1,
2085            })
2086        }
2087        ScalarSource::Uniform { low, high } => {
2088            source.support()?;
2089            Ok(IntervalSample {
2090                value: Interval::new(*low, *high),
2091                weight: Interval::from(high - low),
2092                continuous_dimensions: 1,
2093                piecewise_regions: 1,
2094            })
2095        }
2096        ScalarSource::Histogram(histogram) => {
2097            let (low, high) = source.support()?;
2098            let total: f64 = histogram.counts().iter().sum();
2099            let mut minimum_inverse = f64::INFINITY;
2100            let mut maximum_inverse = 0.0_f64;
2101            let mut regions = 0_usize;
2102            for (count, edges) in histogram
2103                .counts()
2104                .iter()
2105                .zip(histogram.bin_edges().windows(2))
2106            {
2107                if *count <= 0.0 {
2108                    continue;
2109                }
2110                let width = Interval::from(edges[1]) - Interval::from(edges[0]);
2111                let density = Interval::from(*count) / (width * total);
2112                let inverse_density = density.recip();
2113                minimum_inverse = minimum_inverse.min(inverse_density.inf());
2114                maximum_inverse = maximum_inverse.max(inverse_density.sup());
2115                regions += 1;
2116            }
2117            if !minimum_inverse.is_finite()
2118                || !maximum_inverse.is_finite()
2119                || maximum_inverse <= 0.0
2120            {
2121                return Err(GenerationError::InvalidConfiguration(
2122                    "histogram source has no finite positive-density region".into(),
2123                ));
2124            }
2125            Ok(IntervalSample {
2126                value: Interval::new(low, high),
2127                weight: Interval::new(minimum_inverse, maximum_inverse),
2128                continuous_dimensions: 1,
2129                piecewise_regions: regions,
2130            })
2131        }
2132    }
2133}
2134
2135fn interval_initial_momentum(
2136    source: &InitialMomentum,
2137    mass: f64,
2138) -> GenerationResult<IntervalInitialMomentum> {
2139    let (energy, momentum, weight, continuous_dimensions, piecewise_regions) = match source {
2140        InitialMomentum::P4(p4) => (
2141            Interval::from(p4.e()),
2142            [
2143                Interval::from(p4.px()),
2144                Interval::from(p4.py()),
2145                Interval::from(p4.pz()),
2146            ],
2147            Interval::ONE,
2148            0,
2149            1,
2150        ),
2151        InitialMomentum::Momentum(momentum) => {
2152            let p2 = momentum.px() * momentum.px()
2153                + momentum.py() * momentum.py()
2154                + momentum.pz() * momentum.pz();
2155            (
2156                Interval::from((p2 + mass * mass).sqrt()),
2157                [
2158                    Interval::from(momentum.px()),
2159                    Interval::from(momentum.py()),
2160                    Interval::from(momentum.pz()),
2161                ],
2162                Interval::ONE,
2163                0,
2164                1,
2165            )
2166        }
2167        InitialMomentum::EnergyDirection { energy, direction } => {
2168            let sample = interval_scalar_sample(energy)?;
2169            let direction = direction.unit()?;
2170            let magnitude = (sample.value.sqr() - mass * mass).sqrt();
2171            (
2172                sample.value,
2173                [
2174                    magnitude * direction.px(),
2175                    magnitude * direction.py(),
2176                    magnitude * direction.pz(),
2177                ],
2178                sample.weight,
2179                sample.continuous_dimensions,
2180                sample.piecewise_regions,
2181            )
2182        }
2183    };
2184    Ok(IntervalInitialMomentum {
2185        energy,
2186        momentum,
2187        mass,
2188        weight,
2189        continuous_dimensions,
2190        piecewise_regions,
2191    })
2192}
2193
2194fn interval_invariant_mass(initials: &[IntervalInitialMomentum]) -> Interval {
2195    let mut invariant_squared = Interval::ZERO;
2196    for initial in initials {
2197        invariant_squared += initial.mass * initial.mass;
2198    }
2199    for first in 0..initials.len() {
2200        for second in (first + 1)..initials.len() {
2201            let lhs = initials[first];
2202            let rhs = initials[second];
2203            let spatial_dot = lhs.momentum[0] * rhs.momentum[0]
2204                + lhs.momentum[1] * rhs.momentum[1]
2205                + lhs.momentum[2] * rhs.momentum[2];
2206            invariant_squared += 2.0 * (lhs.energy * rhs.energy - spatial_dot);
2207        }
2208    }
2209    invariant_squared.sqrt()
2210}
2211
2212fn validate_common(events: usize) -> GenerationResult<()> {
2213    if events == 0 {
2214        return Err(GenerationError::InvalidConfiguration(
2215            "events must be nonzero".into(),
2216        ));
2217    }
2218    Ok(())
2219}
2220fn validate_bound(bound: f64) -> GenerationResult<()> {
2221    if !bound.is_finite() || bound <= 0.0 {
2222        return Err(GenerationError::InvalidConfiguration(
2223            "envelope must be finite and positive".into(),
2224        ));
2225    }
2226    Ok(())
2227}
2228fn validate_p4(name: &str, p4: RealVec4, mass: f64, index: u64) -> GenerationResult<()> {
2229    if ![p4.e, p4.px, p4.py, p4.pz].into_iter().all(f64::is_finite) || p4.e <= 0.0 {
2230        return Err(GenerationError::Kinematics {
2231            index,
2232            source: LadduPhysicsError::invalid_value(
2233                format!("four-momentum for edge `{name}`"),
2234                "finite components and positive energy",
2235                p4,
2236            ),
2237        });
2238    }
2239    let tolerance = 1e-9 * (1.0 + mass * mass + p4.e * p4.e);
2240    if (p4.m2() - mass * mass).abs() > tolerance {
2241        return Err(GenerationError::Kinematics {
2242            index,
2243            source: LadduPhysicsError::invalid_relation(format!(
2244                "edge is off shell: p²={} but mass²={}",
2245                p4.m2(),
2246                mass * mass
2247            )),
2248        });
2249    }
2250    Ok(())
2251}
2252fn validate_indexed_conservation(
2253    vertex: &VertexPlan,
2254    p4s: &[RealVec4],
2255    index: u64,
2256) -> GenerationResult<()> {
2257    let incoming: RealVec4 = vertex.incoming.iter().map(|&edge| p4s[edge]).sum();
2258    let outgoing: RealVec4 = vertex.outgoing.iter().map(|&edge| p4s[edge]).sum();
2259    let residual = incoming - outgoing;
2260    let scale = incoming.e.abs().max(1.0);
2261    if residual.e.abs() > 1e-9 * scale
2262        || residual.px.abs() > 1e-9 * scale
2263        || residual.py.abs() > 1e-9 * scale
2264        || residual.pz.abs() > 1e-9 * scale
2265    {
2266        return Err(GenerationError::Kinematics {
2267            index,
2268            source: LadduPhysicsError::invalid_relation(format!(
2269                "vertex `{}` violates four-momentum conservation by {residual:?}",
2270                vertex.name
2271            )),
2272        });
2273    }
2274    Ok(())
2275}
2276fn derive_seed(seed: u64, stream: u64, index: u64, object: u64) -> u64 {
2277    let mut rng = ProposalRng::new(
2278        seed ^ stream.wrapping_mul(0xd6e8_feb8_6659_fd93)
2279            ^ index.wrapping_mul(0xa076_1d64_78bd_642f)
2280            ^ object.wrapping_mul(0xe703_7ed1_a0b4_28db),
2281    );
2282    rng.next_u64()
2283}
2284fn acceptance_uniform(seed: u64, index: u64) -> f64 {
2285    ProposalRng::new(derive_seed(seed, 2, index, 0)).uniform()
2286}
2287fn report(requested: usize, seed: u64, decision: &MemoryDecision) -> GenerationReport {
2288    GenerationReport {
2289        requested,
2290        seed,
2291        chunk_events: decision.chunk_events,
2292        estimated_peak_bytes: decision.estimated_peak_bytes,
2293        actual_high_water_bytes: Some(decision.estimated_peak_bytes),
2294        minimum_weight: f64::INFINITY,
2295        ..GenerationReport::default()
2296    }
2297}
2298fn update_report(report: &mut GenerationReport, events: &[GeneratedEvent]) {
2299    for event in events {
2300        report.maximum_weight = report.maximum_weight.max(event.target_weight);
2301        report.minimum_weight = report.minimum_weight.min(event.target_weight);
2302        report.sum_weights += event.target_weight;
2303        report.sum_squared_weights += event.target_weight * event.target_weight;
2304    }
2305}
2306
2307#[cfg(test)]
2308mod tests {
2309    use super::*;
2310    use laddu_physics::{histogram::Histogram, quantum::ParticleProperties, vectors::RealVec3};
2311
2312    struct FailingSink {
2313        aborted: bool,
2314    }
2315
2316    impl EventSink for FailingSink {
2317        fn begin(
2318            &mut self,
2319            _schema: Arc<Schema>,
2320            _plan: WritePlan,
2321        ) -> laddu_data::LadduDataResult<()> {
2322            Ok(())
2323        }
2324
2325        fn write_batch(&mut self, _batch: &EventBatch) -> laddu_data::LadduDataResult<()> {
2326            Err(laddu_data::LadduDataError::Sink(
2327                "injected write failure".into(),
2328            ))
2329        }
2330
2331        fn finish(&mut self) -> laddu_data::LadduDataResult<()> {
2332            Ok(())
2333        }
2334
2335        fn abort(&mut self) -> laddu_data::LadduDataResult<()> {
2336            self.aborted = true;
2337            Ok(())
2338        }
2339    }
2340
2341    fn decay_generator() -> ChannelGenerator {
2342        let mut channel = Channel::new("decay");
2343        channel
2344            .edge("parent")
2345            .properties(&ParticleProperties::unknown().with_mass(2.0))
2346            .initial_p4(RealVec4::new(2.0, 0.0, 0.0, 0.0));
2347        channel
2348            .edge("a")
2349            .properties(&ParticleProperties::unknown().with_mass(0.2))
2350            .output();
2351        channel
2352            .edge("b")
2353            .properties(&ParticleProperties::unknown().with_mass(0.4))
2354            .output();
2355        channel
2356            .vertex("decay")
2357            .incoming(["parent"])
2358            .outgoing(["a", "b"]);
2359        ChannelGenerator::new(channel).unwrap()
2360    }
2361
2362    fn closure_generator() -> ChannelGenerator {
2363        let photon = ParticleProperties::unknown().with_mass(0.0);
2364        let proton = ParticleProperties::unknown().with_mass(0.938_272_088_16);
2365        let kaon = ParticleProperties::unknown().with_mass(0.497_611);
2366        let mut channel = Channel::new("closure envelope");
2367        channel
2368            .edge("gamma")
2369            .properties(&photon)
2370            .initial_energy_source_direction(
2371                ScalarSource::uniform(8.0, 9.0),
2372                RealVec3::new(0.0, 0.0, 1.0),
2373            );
2374        channel
2375            .edge("target")
2376            .properties(&proton)
2377            .initial_momentum(RealVec3::new(0.0, 0.0, 0.0));
2378        channel
2379            .edge("x")
2380            .mass_proposal(MassProposal::uniform(2.0 * 0.497_611, 2.0))
2381            .generated_only();
2382        channel.edge("recoil").properties(&proton);
2383        channel.edge("ks1").properties(&kaon);
2384        channel.edge("ks2").properties(&kaon);
2385        channel
2386            .vertex("production")
2387            .incoming(["gamma", "target"])
2388            .outgoing(["x", "recoil"])
2389            .generation(VertexProposal::t_exchange(
2390                ("gamma", "x"),
2391                TDistribution::mixture([
2392                    (0.2, TComponent::Uniform),
2393                    (0.8, TComponent::Exponential { slope: 4.0 }),
2394                ]),
2395            ));
2396        channel
2397            .vertex("decay")
2398            .incoming(["x"])
2399            .outgoing(["ks1", "ks2"]);
2400        ChannelGenerator::new(channel).unwrap()
2401    }
2402
2403    #[test]
2404    fn interval_scalar_and_initial_bounds_cover_builtin_samples() {
2405        let histogram = Histogram::new(vec![1.0, 0.0, 3.0], vec![8.0, 8.25, 8.75, 9.0]).unwrap();
2406        let sources = [
2407            ScalarSource::constant(8.5),
2408            ScalarSource::uniform(8.0, 9.0),
2409            ScalarSource::histogram(histogram),
2410        ];
2411        for (index, source) in sources.into_iter().enumerate() {
2412            let scalar_bound = interval_scalar_sample(&source).unwrap();
2413            let initial_source = InitialMomentum::energy_source_direction(
2414                source.clone(),
2415                RealVec3::new(0.0, 0.0, 1.0),
2416            );
2417            let initial_bound = interval_initial_momentum(&initial_source, 0.5).unwrap();
2418            let mut rng = ProposalRng::new(500 + index as u64);
2419            for _ in 0..2_000 {
2420                let scalar = source.sample(&mut rng).unwrap();
2421                assert!(scalar_bound.value.contains(scalar.value));
2422                assert!(scalar_bound.weight.contains(scalar.weight));
2423                let initial = initial_source.sample_prevalidated(0.5, &mut rng).unwrap();
2424                assert!(initial_bound.energy.contains(initial.p4.e()));
2425                assert!(initial_bound.momentum[2].contains(initial.p4.pz()));
2426                assert!(initial_bound.weight.contains(initial.weight));
2427            }
2428        }
2429    }
2430
2431    #[test]
2432    fn proven_decay_envelope_contains_proposals_and_unweights_without_a_model() {
2433        let generator = decay_generator();
2434        let proven = generator.phase_space_envelope().unwrap();
2435        let proposals = generator.propose_range(0, 2_000, 41, 0).unwrap();
2436        assert!(
2437            proposals
2438                .iter()
2439                .all(|event| proven.weight_interval.contains(event.proposal_weight))
2440        );
2441
2442        let mut config = UnweightedConfig::new(64).with_max_proposals(10_000);
2443        config.seed = 41;
2444        config.envelope = EnvelopeMode::ProvenPhaseSpace;
2445        let (_, report) = generator.generate_unweighted_dataset(config, None).unwrap();
2446        assert_eq!(report.envelope_kind, Some(EnvelopeKind::ProvenPhaseSpace));
2447        assert_eq!(report.proven_weight_interval, Some(proven.weight_bounds()));
2448        assert!(report.maximum_weight <= proven.maximum_weight);
2449    }
2450
2451    #[test]
2452    fn proven_unweighting_is_thread_and_memory_budget_independent() {
2453        let generator = decay_generator();
2454        let mut first = UnweightedConfig::new(32).with_max_proposals(1_000);
2455        first.seed = 57;
2456        first.memory = MemoryBudget::Bytes(4_096);
2457        first.envelope = EnvelopeMode::ProvenPhaseSpace;
2458        let mut second = first;
2459        second.memory = MemoryBudget::Bytes(16_384);
2460        let one_thread = rayon::ThreadPoolBuilder::new()
2461            .num_threads(1)
2462            .build()
2463            .unwrap();
2464        let four_threads = rayon::ThreadPoolBuilder::new()
2465            .num_threads(4)
2466            .build()
2467            .unwrap();
2468        let (a, _) = one_thread
2469            .install(|| generator.generate_unweighted_dataset(first, None))
2470            .unwrap();
2471        let (b, _) = four_threads
2472            .install(|| generator.generate_unweighted_dataset(second, None))
2473            .unwrap();
2474        let mut av = Vec::new();
2475        a.for_each_event(|event| av.push(event.p4(0))).unwrap();
2476        let mut bv = Vec::new();
2477        b.for_each_event(|event| bv.push(event.p4(0))).unwrap();
2478        assert_eq!(av, bv);
2479    }
2480
2481    #[test]
2482    fn closure_phase_space_bound_contains_sample_and_reports_efficiency() {
2483        let generator = closure_generator();
2484        let proven = generator.phase_space_envelope().unwrap();
2485        let mut weights = generator
2486            .propose_range(0, 20_000, 73, 0)
2487            .unwrap()
2488            .into_iter()
2489            .map(|event| event.proposal_weight)
2490            .collect::<Vec<_>>();
2491        assert!(
2492            weights
2493                .iter()
2494                .all(|weight| proven.weight_interval.contains(*weight))
2495        );
2496        weights.sort_by(f64::total_cmp);
2497        let quantile = |fraction: f64| {
2498            weights[((weights.len() - 1) as f64 * fraction).round() as usize]
2499                / proven.maximum_weight
2500        };
2501        let mean = weights.iter().sum::<f64>() / weights.len() as f64;
2502        eprintln!(
2503            "closure proven-envelope ratios: max={:.6e}, mean={:.6e}, p50={:.6e}, p90={:.6e}, p99={:.6e}",
2504            weights[weights.len() - 1] / proven.maximum_weight,
2505            mean / proven.maximum_weight,
2506            quantile(0.50),
2507            quantile(0.90),
2508            quantile(0.99),
2509        );
2510        assert!(proven.maximum_weight.is_finite());
2511        assert!(
2512            proven.subdivisions > 0,
2513            "a non-singleton root invariant-mass domain should be subdivided"
2514        );
2515        eprintln!(
2516            "closure branch-and-bound envelope: subdivisions={}, proposal_dimensions={}, regions={}, maximum={:.6e}",
2517            proven.subdivisions,
2518            proven.continuous_dimensions,
2519            proven.piecewise_regions,
2520            proven.maximum_weight,
2521        );
2522        assert!(
2523            weights[weights.len() - 1] / proven.maximum_weight > 0.5,
2524            "latent mass and transfer refinement should materially tighten the closure envelope"
2525        );
2526    }
2527
2528    #[test]
2529    fn weighted_generation_is_memory_budget_independent() {
2530        let generator = decay_generator();
2531        let mut first = WeightedConfig::new(32);
2532        first.seed = 9;
2533        first.memory = MemoryBudget::Bytes(4_096);
2534        let mut second = first;
2535        second.memory = MemoryBudget::Bytes(16_384);
2536        let one_thread = rayon::ThreadPoolBuilder::new()
2537            .num_threads(1)
2538            .build()
2539            .unwrap();
2540        let four_threads = rayon::ThreadPoolBuilder::new()
2541            .num_threads(4)
2542            .build()
2543            .unwrap();
2544        let (a, _) = one_thread
2545            .install(|| generator.generate_weighted_dataset(first, None))
2546            .unwrap();
2547        let (b, _) = four_threads
2548            .install(|| generator.generate_weighted_dataset(second, None))
2549            .unwrap();
2550        let mut av = Vec::new();
2551        a.for_each_event(|event| av.push(event.p4(0))).unwrap();
2552        let mut bv = Vec::new();
2553        b.for_each_event(|event| bv.push(event.p4(0))).unwrap();
2554        assert_eq!(av, bv);
2555    }
2556
2557    #[test]
2558    fn unweighted_generation_is_thread_and_memory_budget_independent() {
2559        let generator = decay_generator();
2560        let model = CompiledModel::from_expr(&laddu_expr::Expr::from(1.0)).unwrap();
2561        let evaluator = ModelEvaluator::prepare(
2562            &model,
2563            model.params().default_values(),
2564            &Execution::default(),
2565        )
2566        .unwrap();
2567        let mut first = UnweightedConfig::new(32).with_max_proposals(20_000);
2568        first.seed = 91;
2569        first.memory = MemoryBudget::Bytes(4_096);
2570        first.envelope = EnvelopeMode::Strict { max_weight: 1.0 };
2571        let mut second = first;
2572        second.memory = MemoryBudget::Bytes(16_384);
2573        let one_thread = rayon::ThreadPoolBuilder::new()
2574            .num_threads(1)
2575            .build()
2576            .unwrap();
2577        let four_threads = rayon::ThreadPoolBuilder::new()
2578            .num_threads(4)
2579            .build()
2580            .unwrap();
2581        let (a, _) = one_thread
2582            .install(|| generator.generate_unweighted_dataset(first, Some(&evaluator)))
2583            .unwrap();
2584        let (b, _) = four_threads
2585            .install(|| generator.generate_unweighted_dataset(second, Some(&evaluator)))
2586            .unwrap();
2587        let mut av = Vec::new();
2588        a.for_each_event(|event| av.push(event.p4(0))).unwrap();
2589        let mut bv = Vec::new();
2590        b.for_each_event(|event| bv.push(event.p4(0))).unwrap();
2591        assert_eq!(av, bv);
2592    }
2593
2594    #[test]
2595    fn adaptive_mass_density_is_normalized_and_matches_returned_weight() {
2596        let proposal = AdaptiveMassProposal {
2597            base: MassProposal::uniform(0.0, 4.0),
2598            density: PiecewiseDensity::uniform(1.0, 3.0, Arc::from([1.0, 3.0, 7.0, 1.0])).unwrap(),
2599            defensive_fraction: 0.2,
2600        };
2601        let steps = 20_000;
2602        let dx = 4.0 / steps as f64;
2603        let integral: f64 = (0..steps)
2604            .map(|step| {
2605                let mass = (step as f64 + 0.5) * dx;
2606                proposal.density(0.0, 4.0, mass).unwrap().unwrap() * dx
2607            })
2608            .sum();
2609        assert!((integral - 1.0).abs() < 1e-10);
2610
2611        let mut rng = ProposalRng::new(123);
2612        for _ in 0..1_000 {
2613            let sampled = proposal.propose(0.0, 4.0, &mut rng).unwrap();
2614            let density = proposal.density(0.0, 4.0, sampled.mass).unwrap().unwrap();
2615            assert!((sampled.weight * density - 1.0).abs() < 1e-12);
2616        }
2617    }
2618
2619    #[test]
2620    fn strict_envelope_rejects_overflow() {
2621        let generator = decay_generator();
2622        let model = CompiledModel::from_expr(&laddu_expr::Expr::from(1.0)).unwrap();
2623        let evaluator = ModelEvaluator::prepare(
2624            &model,
2625            model.params().default_values(),
2626            &Execution::default(),
2627        )
2628        .unwrap();
2629        assert!(matches!(
2630            generator.generate_unweighted_dataset(
2631                UnweightedConfig {
2632                    envelope: EnvelopeMode::Strict { max_weight: 1e-12 },
2633                    ..UnweightedConfig::new(2).with_max_proposals(10)
2634                },
2635                Some(&evaluator),
2636            ),
2637            Err(GenerationError::EnvelopeOverflow { .. })
2638        ));
2639    }
2640
2641    #[test]
2642    fn unweighted_generation_is_unlimited_by_default() {
2643        let generator = decay_generator();
2644        let model = CompiledModel::from_expr(&laddu_expr::Expr::from(1.0)).unwrap();
2645        let evaluator = ModelEvaluator::prepare(
2646            &model,
2647            model.params().default_values(),
2648            &Execution::default(),
2649        )
2650        .unwrap();
2651        let config = UnweightedConfig {
2652            envelope: EnvelopeMode::Strict { max_weight: 1.0 },
2653            ..UnweightedConfig::new(32)
2654        };
2655        assert_eq!(config.max_proposals, None);
2656        let (_, report) = generator
2657            .generate_unweighted_dataset(config, Some(&evaluator))
2658            .unwrap();
2659        assert_eq!(report.produced, 32);
2660        assert!(report.proposals >= 32);
2661    }
2662
2663    #[test]
2664    fn explicit_proposal_limit_can_exhaust_generation() {
2665        let generator = decay_generator();
2666        let model = CompiledModel::from_expr(&laddu_expr::Expr::from(1.0)).unwrap();
2667        let evaluator = ModelEvaluator::prepare(
2668            &model,
2669            model.params().default_values(),
2670            &Execution::default(),
2671        )
2672        .unwrap();
2673        assert!(matches!(
2674            generator.generate_unweighted_dataset(
2675                UnweightedConfig {
2676                    envelope: EnvelopeMode::Strict { max_weight: 1.0 },
2677                    ..UnweightedConfig::new(16).with_max_proposals(16)
2678                },
2679                Some(&evaluator),
2680            ),
2681            Err(GenerationError::Exhausted {
2682                requested: 16,
2683                proposals: 16,
2684                ..
2685            })
2686        ));
2687    }
2688
2689    #[test]
2690    fn adaptive_envelope_grows_and_rethins_buffered_events() {
2691        let generator = decay_generator();
2692        let model = CompiledModel::from_expr(&laddu_expr::Expr::from(1.0)).unwrap();
2693        let evaluator = ModelEvaluator::prepare(
2694            &model,
2695            model.params().default_values(),
2696            &Execution::default(),
2697        )
2698        .unwrap();
2699        let mut config = UnweightedConfig::new(16).with_max_proposals(10_000);
2700        config.envelope = EnvelopeMode::Strict { max_weight: 1e-12 };
2701        config.envelope_overflow = EnvelopeOverflow::Grow { safety_factor: 1.5 };
2702        let (dataset, report) = generator
2703            .generate_unweighted_dataset(config, Some(&evaluator))
2704            .unwrap();
2705
2706        let mut events = 0;
2707        dataset.for_each_event(|_| events += 1).unwrap();
2708        assert_eq!(events, 16);
2709        assert_eq!(report.produced, 16);
2710        assert!(report.envelope_updates >= 1);
2711        assert!(report.envelope.unwrap() >= report.maximum_weight);
2712    }
2713
2714    #[test]
2715    fn channel_validation_preserves_initial_source_error() {
2716        let mut channel = Channel::new("decay");
2717        channel
2718            .edge("parent")
2719            .properties(&ParticleProperties::unknown().with_mass(2.0))
2720            .initial_energy_direction(2.0, RealVec3::default());
2721        channel
2722            .edge("a")
2723            .properties(&ParticleProperties::unknown().with_mass(0.2))
2724            .output();
2725        channel
2726            .edge("b")
2727            .properties(&ParticleProperties::unknown().with_mass(0.4))
2728            .output();
2729        channel
2730            .vertex("decay")
2731            .incoming(["parent"])
2732            .outgoing(["a", "b"])
2733            .generation(VertexProposal::isotropic_decay());
2734        assert!(matches!(
2735            ChannelGenerator::new(channel),
2736            Err(GenerationError::ChannelValidation {
2737                source: LadduPhysicsError::InvalidValue { .. },
2738            })
2739        ));
2740    }
2741
2742    #[test]
2743    fn channel_is_validated_when_consumed_by_the_generator() {
2744        let mut channel = Channel::new("invalid");
2745        channel
2746            .edge("same")
2747            .properties(&ParticleProperties::unknown().with_mass(1.0))
2748            .initial_p4(RealVec4::new(1.0, 0.0, 0.0, 0.0))
2749            .output();
2750        channel
2751            .vertex("duplicate")
2752            .incoming(["same"])
2753            .outgoing(["same"]);
2754
2755        assert!(matches!(
2756            ChannelGenerator::new(channel),
2757            Err(GenerationError::ChannelValidation {
2758                source: LadduPhysicsError::InvalidRelation { .. },
2759            })
2760        ));
2761    }
2762
2763    #[test]
2764    fn momentum_initial_state_requires_particle_mass_metadata() {
2765        let mut channel = Channel::new("decay");
2766        channel
2767            .edge("parent")
2768            .initial_momentum(RealVec3::new(0.0, 0.0, 0.0));
2769        channel
2770            .edge("a")
2771            .properties(&ParticleProperties::unknown().with_mass(0.2))
2772            .output();
2773        channel
2774            .edge("b")
2775            .properties(&ParticleProperties::unknown().with_mass(0.4))
2776            .output();
2777        channel
2778            .vertex("decay")
2779            .incoming(["parent"])
2780            .outgoing(["a", "b"])
2781            .generation(VertexProposal::isotropic_decay());
2782        assert!(matches!(
2783            ChannelGenerator::new(channel),
2784            Err(GenerationError::ChannelValidation {
2785                source: LadduPhysicsError::InvalidRelation { .. },
2786            })
2787        ));
2788    }
2789
2790    #[test]
2791    fn initial_energy_sources_sample_energy_and_apply_the_proposal_correction() {
2792        let mut channel = Channel::new("initial state");
2793        channel
2794            .edge("beam")
2795            .properties(&ParticleProperties::unknown().with_mass(0.5));
2796        let source = InitialMomentum::energy_source_direction(
2797            ScalarSource::uniform(1.0, 3.0),
2798            RealVec3::z(),
2799        );
2800        let properties = channel.particle("beam").unwrap();
2801        source.validate("beam", Some(properties)).unwrap();
2802        let sampled = source
2803            .sample("beam", Some(properties), &mut ProposalRng::new(43))
2804            .unwrap();
2805
2806        assert!((1.0..3.0).contains(&sampled.p4.e()));
2807        assert_eq!(sampled.weight, 2.0);
2808    }
2809
2810    #[test]
2811    fn generated_scalar_columns_are_emitted_and_available_to_models() {
2812        let generator = decay_generator()
2813            .with_scalar("aux", ScalarSource::constant(2.0))
2814            .unwrap();
2815        let sources = generator.scalar_sources().collect::<Vec<_>>();
2816        assert_eq!(sources.len(), 1);
2817        assert_eq!(sources[0].0, "aux");
2818        let weighted_schema = generator.weighted_output_schema(false).unwrap();
2819        let weighted_scalars = weighted_schema
2820            .scalars()
2821            .iter()
2822            .map(AsRef::as_ref)
2823            .collect::<Vec<_>>();
2824        assert_eq!(weighted_scalars, ["aux"]);
2825        let unweighted_schema = generator.unweighted_output_schema(true).unwrap();
2826        let unweighted_scalars = unweighted_schema
2827            .scalars()
2828            .iter()
2829            .map(AsRef::as_ref)
2830            .collect::<Vec<_>>();
2831        assert_eq!(
2832            unweighted_scalars,
2833            [
2834                "aux",
2835                "__laddu_proposal_weight",
2836                "__laddu_model_weight",
2837                "__laddu_target_weight",
2838            ]
2839        );
2840        let model = CompiledModel::from_expr(&laddu_expr::event_scalar("aux")).unwrap();
2841        let evaluator = ModelEvaluator::prepare(
2842            &model,
2843            model.params().default_values(),
2844            &Execution::default(),
2845        )
2846        .unwrap();
2847        let (dataset, _) = generator
2848            .generate_weighted_dataset(WeightedConfig::new(8), Some(&evaluator))
2849            .unwrap();
2850
2851        assert_eq!(dataset.schema().unwrap().scalar_index("aux"), Some(0));
2852        dataset
2853            .for_each_event(|event| {
2854                assert_eq!(event.scalar_named("aux"), Some(2.0));
2855                assert!(event.weight() > 0.0);
2856            })
2857            .unwrap();
2858    }
2859
2860    #[test]
2861    fn weighted_generation_aborts_sink_after_post_begin_failure() {
2862        let generator = decay_generator();
2863        let mut sink = FailingSink { aborted: false };
2864
2865        let result = generator.generate_weighted_to(WeightedConfig::new(2), None, &mut sink);
2866
2867        assert!(matches!(
2868            result,
2869            Err(GenerationError::Data(laddu_data::LadduDataError::Sink(message)))
2870                if message.contains("injected write failure")
2871        ));
2872        assert!(sink.aborted);
2873    }
2874}