1use 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
39pub type GenerationResult<T> = Result<T, GenerationError>;
41
42#[derive(Debug, Error)]
44pub enum GenerationError {
45 #[error("invalid generation channel: {0}")]
47 InvalidChannel(String),
48 #[error("channel validation failed: {source}")]
50 ChannelValidation {
51 #[source]
53 source: LadduPhysicsError,
54 },
55 #[error("invalid generation configuration: {0}")]
57 InvalidConfiguration(String),
58 #[error("initial-state proposal {index} failed: {source}")]
60 InitialState {
61 index: u64,
63 #[source]
65 source: LadduPhysicsError,
66 },
67 #[error("mass proposal for edge `{edge}` at proposal {index} failed: {source}")]
69 MassProposal {
70 index: u64,
72 edge: String,
74 #[source]
76 source: LadduPhysicsError,
77 },
78 #[error("vertex `{vertex}` at proposal {index} failed: {source}")]
80 VertexProposal {
81 index: u64,
83 vertex: String,
85 #[source]
87 source: LadduPhysicsError,
88 },
89 #[error("scalar column `{column}` at proposal {index} failed: {source}")]
91 ScalarProposal {
92 index: u64,
94 column: String,
96 #[source]
98 source: LadduPhysicsError,
99 },
100 #[error("kinematic validation at proposal {index} failed: {source}")]
102 Kinematics {
103 index: u64,
105 #[source]
107 source: LadduPhysicsError,
108 },
109 #[error("model evaluation failed: {0}")]
111 Model(String),
112 #[error("target weight {weight} exceeds envelope {envelope} at proposal {index}")]
114 EnvelopeOverflow {
115 index: u64,
117 weight: f64,
119 envelope: f64,
121 },
122 #[error(
124 "accepted {accepted} events after exhausting {proposals} proposals (requested {requested})"
125 )]
126 Exhausted {
127 requested: usize,
129 accepted: usize,
131 proposals: usize,
133 },
134 #[error(transparent)]
136 Data(#[from] laddu_data::LadduDataError),
137 #[error(transparent)]
139 Runtime(#[from] laddu_runtime::RuntimeError),
140 #[error(transparent)]
142 Physics(#[from] LadduPhysicsError),
143}
144
145#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
147pub struct WeightedConfig {
148 pub events: usize,
150 pub memory: MemoryBudget,
152 pub seed: u64,
154 pub diagnostics: bool,
156}
157
158impl WeightedConfig {
159 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#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
172pub struct UnweightedConfig {
173 pub events: usize,
175 pub max_proposals: Option<usize>,
180 pub memory: MemoryBudget,
182 pub seed: u64,
184 pub diagnostics: bool,
186 pub envelope: EnvelopeMode,
188 pub envelope_overflow: EnvelopeOverflow,
190}
191
192impl UnweightedConfig {
193 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 pub fn with_max_proposals(mut self, max_proposals: usize) -> Self {
209 self.max_proposals = Some(max_proposals);
210 self
211 }
212}
213
214#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
216pub enum EnvelopeOverflow {
217 #[default]
219 Error,
220 Grow {
225 safety_factor: f64,
227 },
228}
229
230#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
232pub enum EnvelopeMode {
233 Strict {
235 max_weight: f64,
237 },
238 Pilot {
243 proposals: usize,
245 safety_factor: f64,
247 },
248 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#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
266pub enum EnvelopeKind {
267 Strict,
269 Pilot,
271 ProvenPhaseSpace,
273}
274
275#[derive(Clone, Copy, Debug, PartialEq)]
277pub struct ProvenEnvelopeReport {
278 pub weight_interval: Interval,
280 pub maximum_weight: f64,
282 pub continuous_dimensions: usize,
284 pub piecewise_regions: usize,
286 pub subdivisions: usize,
288}
289
290impl ProvenEnvelopeReport {
291 pub fn weight_bounds(&self) -> (f64, f64) {
293 self.weight_interval.bounds()
294 }
295}
296
297#[derive(Clone, Debug, Default, Serialize, Deserialize)]
299pub struct GenerationReport {
300 pub requested: usize,
302 pub produced: usize,
304 pub proposals: usize,
306 pub pilot_proposals: usize,
308 pub rejected: usize,
310 pub envelope: Option<f64>,
312 pub envelope_kind: Option<EnvelopeKind>,
314 pub envelope_updates: usize,
316 #[serde(default)]
318 pub proven_weight_interval: Option<(f64, f64)>,
319 #[serde(default)]
321 pub proven_continuous_dimensions: Option<usize>,
322 #[serde(default)]
324 pub proven_piecewise_regions: Option<usize>,
325 #[serde(default)]
327 pub proven_subdivisions: Option<usize>,
328 pub maximum_weight: f64,
330 pub minimum_weight: f64,
332 pub sum_weights: f64,
334 pub sum_squared_weights: f64,
336 pub seed: u64,
338 pub chunk_events: usize,
340 pub estimated_peak_bytes: u64,
342 pub actual_high_water_bytes: Option<u64>,
344}
345
346impl GenerationReport {
347 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#[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 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 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#[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 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 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 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 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 pub fn weighted_output_schema(&self, diagnostics: bool) -> GenerationResult<Arc<Schema>> {
857 self.output_schema(true, diagnostics)
858 }
859
860 pub fn unweighted_output_schema(&self, diagnostics: bool) -> GenerationResult<Arc<Schema>> {
866 self.output_schema(false, diagnostics)
867 }
868
869 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 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 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 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 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 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 let _ = sink.abort();
1188 }
1189
1190 result
1191 }
1192
1193 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 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 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}