diff --git a/rust/loop-reduction/README.md b/rust/loop-reduction/README.md index 73ab8c325..7c03912be 100644 --- a/rust/loop-reduction/README.md +++ b/rust/loop-reduction/README.md @@ -144,6 +144,28 @@ explicit field class. It prepares all selected labelled tree amplitudes in Rust, without Python or process-specific compilation. `ScalarTreeProgram::prepare` retains the scalar entry point by delegating to this path. +`TreeProgram::prepare_with_runtime_parameters` additionally accepts ordered +`RuntimeTreeParameter` declarations. Each parameter names a native model symbol +and a `Real`, `PositiveReal` or `Complex` domain. The existing model scalar +resolver keeps all analytic dependencies on those inputs: for example a runtime +`aS` updates `G=2*sqrt(pi*aS)` and every dependent coupling without regenerating +the process. Frozen exact overrides remain separate. The declared domains are +checked for every lane, including carrier, working precision and row identity. +Runtime parameter batches contain one column per declaration and are supplied +to `evaluate_batch_with_parameters`, `evaluate_coefficients_with_parameters` +or `evaluate_external_components_with_parameters`. Evaluator instructions place +these columns after the momenta and external-state components. No parameter is +converted through binary64. + +Runtime parameters cannot invalidate structural masslessness, zero-width or +vacuum proofs. A previously restricted model that has deleted interactions +cannot recover them by changing a runtime value; use the original model for +varying parameters. Exact graph coupling-order filters currently require each +vertex to have homogeneous coupling-order vectors, because FeynKit labels a +mixed vertex by its componentwise maximum. Mixed vertices are accepted without +that filter and retain their full coherent contribution; requesting the +unsupported order projection fails explicitly. + `MasslessQcdAndSingletScalars` admits massless Dirac fundamental/antifundamental quarks, structurally massless adjoint vectors and singlet scalars. Internal poles must be canonical zero-width vacuum poles. Exotic color, Majorana fields, diff --git a/rust/loop-reduction/src/tree_preparation.rs b/rust/loop-reduction/src/tree_preparation.rs index a50b99be8..2afe9a0cc 100644 --- a/rust/loop-reduction/src/tree_preparation.rs +++ b/rust/loop-reduction/src/tree_preparation.rs @@ -9,10 +9,15 @@ //! numerical instruction path. Unsupported participating fields are rejected, //! never dropped. -use std::{collections::BTreeMap, sync::Arc}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, +}; use feynkit_amplitude::AmplitudeOptions; -use feynkit_generator::{GenerationOptions, NumeratorGrouping, Process}; +use feynkit_generator::{ + FilterScope, GenerationFilter, GenerationOptions, NumeratorGrouping, Process, +}; use feynkit_graph::ExternalState; use feynkit_model::{Model, ParameterCard, ParameterType, Particle}; use idenso::{representations::Bispinor, tensor::SymbolicTensor}; @@ -26,12 +31,12 @@ use spenso::structure::{ slot::{DualSlotTo, IsAbstractSlot, ParseableAind}, }; use symbolica::{ - atom::{Atom, AtomCore, AtomView, Symbol}, + atom::{Atom, AtomCore, AtomView, Symbol, SymbolAttribute, SymbolBuilder}, coefficient::CoefficientView, function, symbol, }; -pub use crate::tree_runtime::{ExternalTreeState, TreeEvaluator}; +pub use crate::tree_runtime::{ExternalTreeState, TreeEvaluator, TreeParameterDomain}; use crate::{ runtime::{Arithmetic, arithmetic_mode}, tree_color::{factor_tree_color, index_label}, @@ -54,6 +59,7 @@ pub struct TreeProgram { norm_state: Vec, scalar_norm_state: Option>, states: Vec, + runtime_parameters: Vec, color_tensors: Vec, } @@ -108,6 +114,14 @@ pub enum TreeFieldClass { MasslessQcdAndSingletScalars, } +/// One model parameter retained as a numerical input. Column order is explicit; +/// analytic dependent parameters/couplings are resolved by the native model. +#[derive(Clone, Debug)] +pub struct RuntimeTreeParameter { + pub symbol: Symbol, + pub domain: TreeParameterDomain, +} + fn error(message: impl std::fmt::Display) -> RusticolError { RusticolError::invalid_argument(format!("native tree preparation: {message}")) } @@ -122,6 +136,43 @@ fn bind_scalars(expression: &Atom, bindings: &BTreeMap) -> Atom { }) } +fn validate_order_filter(model: &Model, options: &GenerationOptions) -> RusticolResult<()> { + if !options.filters(FilterScope::Graph).iter().any( + |filter| matches!(filter, GenerationFilter::CouplingOrders(orders) if !orders.is_empty()), + ) { + return Ok(()); + } + // FeynKit currently labels a vertex by componentwise maximum orders. Exact + // filtering of mixed vertices requires upstream termwise order projection. + for vertex in model.vertex_rules() { + let mut expected = None; + for &id in vertex.couplings.iter().flatten().flatten() { + let orders: BTreeMap<_, _> = model + .coupling_by_id(id) + .map_err(error)? + .orders + .iter() + .filter(|(_, power)| **power != 0) + .map(|(name, &power)| (name.clone(), power)) + .collect(); + if expected + .as_ref() + .is_some_and(|previous| previous != &orders) + { + return Err(RusticolError::unsupported_runtime_capability( + "native-tree-orders", + format!( + "vertex {} mixes coupling-order vectors; order-filtered preparation requires homogeneous vertex orders", + vertex.name + ), + )); + } + expected = Some(orders); + } + } + Ok(()) +} + // FeynKit's amplitude denominator currently represents only p²-m², while the // model permits arbitrary denominator templates. Check the exact specialization // before allowing that representation to replace a participating propagator. @@ -195,6 +246,32 @@ impl TreeProgram { exact_overrides: &BTreeMap, field_class: TreeFieldClass, ) -> RusticolResult { + Self::prepare_with_runtime_parameters( + model, + process, + options, + card, + exact_overrides, + field_class, + &[], + ) + } + + /// Retain ordered model parameters as typed per-row evaluator inputs. + /// Use the unrestricted model for parameters intended to vary: interactions + /// already removed by model restriction cannot be reconstructed here. + /// Domain assumptions are enforced on every runtime parameter batch. + /// Dynamic masses/widths must still satisfy this class's admission proofs. + pub fn prepare_with_runtime_parameters( + model: Arc, + process: &Process, + options: &GenerationOptions, + card: Option<&ParameterCard>, + exact_overrides: &BTreeMap, + field_class: TreeFieldClass, + runtime_parameters: &[RuntimeTreeParameter], + ) -> RusticolResult { + validate_order_filter(&model, options)?; for key in exact_overrides.keys() { if !model .parameters() @@ -206,9 +283,74 @@ impl TreeProgram { ))); } } - let scalar_bindings = model - .scalar_bindings(card, exact_overrides) - .map_err(error)?; + let mut overrides = exact_overrides.clone(); + let mut parameter_inputs = Vec::with_capacity(runtime_parameters.len()); + let mut names = BTreeSet::new(); + for (index, input) in runtime_parameters.iter().enumerate() { + if !names.insert(input.symbol) || overrides.contains_key(&input.symbol) { + return Err(error( + "runtime parameters must be distinct and cannot also have frozen overrides", + )); + } + let parameter = model + .parameters() + .iter() + .find(|parameter| input.symbol == symbol!(&format!("UFO::{}", parameter.name))) + .ok_or_else(|| { + error(format!( + "runtime input {} does not name a model parameter", + input.symbol + )) + })?; + if parameter.parameter_type == ParameterType::Real + && input.domain == TreeParameterDomain::Complex + { + return Err(error(format!( + "real model parameter {} requires a real runtime domain", + parameter.name + ))); + } + // Domains are part of the interned symbol name: Symbolica correctly + // rejects conflicting assumptions on a previously created symbol. + let (domain, attributes) = match input.domain { + TreeParameterDomain::Real => ("real", vec![SymbolAttribute::Real]), + TreeParameterDomain::PositiveReal => ("positive", vec![SymbolAttribute::Positive]), + TreeParameterDomain::Complex => ("complex", vec![]), + }; + let name = format!("rusticol_native_tree::runtime_{domain}_{index}"); + let free = SymbolBuilder::new(symbolica::wrap_symbol!(&name)) + .with_attributes(attributes) + .build() + .map_err(error)?; + let free = Atom::var(free); + overrides.insert(input.symbol, free.clone()); + parameter_inputs.push(free); + } + if exact_overrides + .values() + .chain( + model + .parameters() + .iter() + .filter_map(|parameter| parameter.expression.as_ref()), + ) + .chain( + model + .couplings() + .iter() + .map(|coupling| &coupling.expression), + ) + .any(|expression| { + parameter_inputs + .iter() + .any(|input| expression.contains(input)) + }) + { + return Err(error( + "model scalar definitions use a reserved runtime parameter input symbol", + )); + } + let scalar_bindings = model.scalar_bindings(card, &overrides).map_err(error)?; for (key, binding) in &scalar_bindings { let mut approximate = false; let _ = binding.replace_map(|term, _, _| { @@ -272,6 +414,14 @@ impl TreeProgram { .into_amplitude(AmplitudeOptions::default()) .map_err(|failure| error(format!("amplitude generation: {failure}")))?; let physical_expression = amplitude.expression().map_err(error)?; + if parameter_inputs + .iter() + .any(|input| physical_expression.expression().contains(input)) + { + return Err(error( + "model uses a reserved runtime parameter input symbol", + )); + } for parameter in amplitude.model().parameters() { let variable = symbol!(&format!("UFO::{}", parameter.name)); if parameter.parameter_type == ParameterType::Real @@ -400,6 +550,7 @@ impl TreeProgram { for (_, components) in &external_states { inputs.extend(components.iter().cloned()); } + inputs.extend(parameter_inputs); let factors = factor_tree_color(&source, &color_legs, &external_color)?; let mut color_tensors = Vec::with_capacity(factors.len()); let mut coefficients = Vec::with_capacity(factors.len()); @@ -457,6 +608,7 @@ impl TreeProgram { norm_state: encode_exact_evaluator_state(&norm)?, scalar_norm_state, states, + runtime_parameters: runtime_parameters.to_vec(), color_tensors, }) } @@ -471,9 +623,13 @@ impl TreeProgram { pub fn diagram_count(&self) -> usize { self.diagram_count } + pub fn runtime_parameters(&self) -> &[RuntimeTreeParameter] { + &self.runtime_parameters + } /// Exact raw amplitude-coefficient instructions, before the color norm. /// Inputs are all physical momentum components, followed by four complex - /// external components per entry in [`TreeEvaluator::states`]. + /// external components per entry in [`TreeEvaluator::states`], then the + /// ordered [`Self::runtime_parameters`] columns. pub fn evaluator_state(&self) -> &[u8] { &self.evaluator_state } @@ -507,11 +663,15 @@ impl TreeProgram { arithmetic: Arithmetic, bits: u32, ) -> RusticolResult { - TreeEvaluator::new( + TreeEvaluator::with_parameter_domains( &self.evaluator_state, &self.norm_state, self.states.clone(), self.momentum_legs.len() * 4, + self.runtime_parameters + .iter() + .map(|input| input.domain) + .collect(), arithmetic, bits, ) @@ -815,6 +975,266 @@ mod tests { } } + #[test] + fn runtime_alpha_s_resolves_analytic_couplings_without_repreparation() { + // Cached coupling zeros must not delete analytic interactions: only a + // caller's prior explicit model restriction may do that. + let mut definition = serde_json::to_value(Model::qcd()).unwrap(); + for coupling in definition["couplings"].as_array_mut().unwrap() { + coupling["value"] = serde_json::json!([0., 0.]); + } + let model = Arc::new(Model::from_json(&definition.to_string()).unwrap()); + let program = TreeProgram::prepare_with_runtime_parameters( + model, + &Process::new(["g", "g"], ["g", "g"]), + &GenerationOptions::default().threads(1), + None, + &BTreeMap::new(), + TreeFieldClass::MasslessQcdAndSingletScalars, + &[RuntimeTreeParameter { + symbol: symbol!("UFO::aS"), + domain: TreeParameterDomain::PositiveReal, + }], + ) + .unwrap(); + let bits = 173; + let zero = Float::with_val(bits, 0); + let two = Float::with_val(bits, 2); + let first = two.pow(4).inv(); + let second = first.clone() + two.pow(90).inv(); + let momenta = Batch::typed( + [5, 0, 0, 5, 5, 0, 0, -5, 5, 4, 0, 3, 5, -4, 0, -3] + .into_iter() + .cycle() + .take(32) + .map(|value| Complex::new(Float::with_val(bits, value), zero.clone())) + .collect(), + vec![7, 3], + 16, + bits, + ) + .unwrap(); + let parameters = Batch::typed( + [first.clone(), second.clone()] + .into_iter() + .map(|value| Complex::new(value, zero.clone())) + .collect(), + vec![7, 3], + 1, + bits, + ) + .unwrap(); + let mut evaluator = program + .evaluator_with_precision_bits(Arithmetic::Arbitrary, bits) + .unwrap(); + assert!(evaluator.evaluate_batch(&momenta).is_err()); + let output = evaluator + .evaluate_batch_with_parameters(&momenta, ¶meters) + .unwrap(); + let Data::Arbitrary(values) = output.data() else { + panic!("carrier changed") + }; + let base = Float::with_val(bits, 666792) / Float::with_val(bits, 25); + for (value, alpha) in values.iter().zip([first, second]) { + let g_squared = Float::with_val(bits, 4) * zero.pi() * alpha; + let expected = base.clone() * g_squared.pow(2); + assert!((value.re.clone() - expected).norm() < two.pow(135).inv()); + assert!(value.im.norm() < two.pow(135).inv()); + } + assert!( + values[1].re > values[0].re, + "per-row sub-binary64 alpha_s survives" + ); + let permuted = parameters.take_rows(vec![1, 0], None).unwrap(); + assert!( + evaluator + .evaluate_batch_with_parameters(&momenta, &permuted) + .is_err() + ); + } + + #[test] + fn runtime_complex_coefficients_keep_phases_and_zero_default_vertices() { + let mut definition = scalar_model_definition(4); + definition["parameters"][1]["parameter_type"] = serde_json::json!("complex"); + definition["parameters"][1]["value"] = serde_json::json!([0., 0.]); + definition["couplings"][0]["value"] = serde_json::json!([0., 0.]); + let program = TreeProgram::prepare_with_runtime_parameters( + Arc::new(Model::from_json(&definition.to_string()).unwrap()), + &Process::new(["s", "s"], ["s", "s"]), + &GenerationOptions::default().threads(1), + None, + &BTreeMap::new(), + TreeFieldClass::SingletScalars, + &[RuntimeTreeParameter { + symbol: symbol!("UFO::g_"), + domain: TreeParameterDomain::Complex, + }], + ) + .unwrap(); + assert_eq!(program.diagram_count(), 1); + let momenta = Batch::typed(vec![Complex::new(0., 0.); 32], vec![11, 19], 16, 53).unwrap(); + let parameters = Batch::typed( + vec![Complex::new(1., 2.), Complex::new(3., -4.)], + vec![11, 19], + 1, + 53, + ) + .unwrap(); + let mut evaluator = program.evaluator("double", None).unwrap(); + let amplitudes = evaluator + .evaluate_coefficients_with_parameters(&momenta, ¶meters, &[]) + .unwrap(); + let Data::Double(amplitudes) = amplitudes.data() else { + panic!("carrier changed") + }; + assert_eq!( + amplitudes[0].conj() * amplitudes[1], + Complex::new(-5., -10.) + ); + let norms = evaluator + .evaluate_batch_with_parameters(&momenta, ¶meters) + .unwrap(); + let Data::Double(norms) = norms.data() else { + panic!("carrier changed") + }; + assert_eq!( + norms.as_slice(), + &[Complex::new(5., 0.), Complex::new(25., 0.)] + ); + } + + #[test] + fn runtime_parameter_domains_do_not_erase_model_admission_proofs() { + let model = Arc::new(scalar_model(4)); + for inputs in [ + vec![RuntimeTreeParameter { + symbol: symbol!("UFO::missing"), + domain: TreeParameterDomain::Real, + }], + vec![RuntimeTreeParameter { + symbol: symbol!("UFO::ZERO"), + domain: TreeParameterDomain::Real, + }], + vec![RuntimeTreeParameter { + symbol: symbol!("UFO::g_"), + domain: TreeParameterDomain::Complex, + }], + vec![ + RuntimeTreeParameter { + symbol: symbol!("UFO::g_"), + domain: TreeParameterDomain::Real + }; + 2 + ], + ] { + assert!( + TreeProgram::prepare_with_runtime_parameters( + model.clone(), + &Process::new(["s", "s"], ["s", "s"]), + &GenerationOptions::default().threads(1), + None, + &BTreeMap::new(), + TreeFieldClass::SingletScalars, + &inputs + ) + .is_err() + ); + } + let alias = Atom::var(symbol!("rusticol_native_tree::runtime_real_0"; Real)); + assert!( + TreeProgram::prepare_with_runtime_parameters( + model.clone(), + &Process::new(["s", "s"], ["s", "s"]), + &GenerationOptions::default().threads(1), + None, + &BTreeMap::from([(symbol!("UFO::m"), alias)]), + TreeFieldClass::SingletScalars, + &[RuntimeTreeParameter { + symbol: symbol!("UFO::g_"), + domain: TreeParameterDomain::Real + }] + ) + .is_err() + ); + let qcd = Arc::new(Model::qcd()); + let mass = qcd + .parameter_by_id(qcd.particle("u").unwrap().mass) + .unwrap() + .name + .clone(); + assert!( + TreeProgram::prepare_with_runtime_parameters( + qcd, + &Process::new(["u", "d"], ["u", "d"]), + &GenerationOptions::default().threads(1), + None, + &BTreeMap::new(), + TreeFieldClass::MasslessQcdAndSingletScalars, + &[RuntimeTreeParameter { + symbol: symbol!(&format!("UFO::{mass}")), + domain: TreeParameterDomain::Real + }] + ) + .is_err() + ); + } + + #[test] + fn order_filters_reject_mixed_vertex_orders_without_discarding_terms() { + let mut definition = scalar_model_definition(4); + definition["couplings"] + .as_array_mut() + .unwrap() + .push(serde_json::json!({ + "name":"GC2", "expression":"g_", "orders":[["SELF",2]], "value":[1.,0.] + })); + definition["lorentz_structures"] + .as_array_mut() + .unwrap() + .push(serde_json::json!({ + "name":"SSSS2", "spins":[1,1,1,1], "structure":"1" + })); + definition["vertex_rules"][0]["lorentz_structures"] = serde_json::json!(["SSSS", "SSSS2"]); + definition["vertex_rules"][0]["couplings"] = serde_json::json!([["GC1", "GC2"]]); + let mixed = Arc::new(Model::from_json(&definition.to_string()).unwrap()); + let process = Process::new(["s", "s"], ["s", "s"]); + let options = GenerationOptions::default().threads(1).with_graph_filter( + GenerationFilter::CouplingOrders(BTreeMap::from([("SELF".into(), (1, Some(1)))])), + ); + assert!( + TreeProgram::prepare( + mixed.clone(), + &process, + &options, + None, + &BTreeMap::new(), + TreeFieldClass::SingletScalars + ) + .is_err() + ); + // Without an order projection both terms form a valid coherent sum. + let complete = TreeProgram::prepare( + mixed, + &process, + &GenerationOptions::default().threads(1), + None, + &BTreeMap::new(), + TreeFieldClass::SingletScalars, + ) + .unwrap(); + let input = Batch::typed(vec![Complex::new(0., 0.); 16], vec![4], 16, 53).unwrap(); + let output = complete + .evaluator("double", None) + .unwrap() + .evaluate_batch(&input) + .unwrap(); + let Data::Double(values) = output.data() else { + panic!("carrier changed") + }; + assert_eq!(values[0], Complex::new(4., 0.)); + } + #[test] fn massless_qcd_distinct_quarks_match_pdg_and_boost() { check_massless_qcd_channel(0); diff --git a/rust/loop-reduction/src/tree_runtime.rs b/rust/loop-reduction/src/tree_runtime.rs index fadb0d49a..c430eeb46 100644 --- a/rust/loop-reduction/src/tree_runtime.rs +++ b/rust/loop-reduction/src/tree_runtime.rs @@ -26,16 +26,28 @@ pub struct ExternalTreeState { pub helicities: Vec, } +/// Scalar domains assumed while preparing a runtime model parameter. +/// Runtime batches must obey these assumptions before numerical evaluation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TreeParameterDomain { + Real, + PositiveReal, + Complex, +} + /// Reusable coherent amplitudes and their exact-color one-helicity norm. /// /// Both programs use the shared coefficient evaluator. Amplitude inputs are -/// all momentum columns followed by four components per external state. Its -/// outputs are the color-basis coefficients; the norm program consumes those -/// outputs, including their complex phases, and returns one scalar column. +/// all momentum columns, four components per external state, then the runtime +/// parameter columns in their declared order. Its outputs are the color-basis +/// coefficients; the norm program consumes those outputs, including their +/// complex phases, and returns one scalar column. Evaluation methods without +/// a parameter batch require an empty runtime parameter domain. pub struct TreeEvaluator { amplitude: CoefficientEvaluator, norm: CoefficientEvaluator, states: Vec, + parameter_domains: Vec, momentum_columns: usize, coefficient_columns: usize, arithmetic: Arithmetic, @@ -59,6 +71,31 @@ fn real_finite(values: &[Complex]) -> bool { .all(|value| value.im.is_fully_zero() && value.re.is_finite()) } +fn parameter_values( + values: &[Complex], + domains: &[TreeParameterDomain], +) -> RusticolResult<()> { + for (index, value) in values.iter().enumerate() { + let column = index % domains.len(); + let valid = value.re.is_finite() + && value.im.is_finite() + && match domains[column] { + TreeParameterDomain::Complex => true, + TreeParameterDomain::Real => value.im.is_fully_zero(), + TreeParameterDomain::PositiveReal => { + value.im.is_fully_zero() && value.re > value.re.zero() + } + }; + if !valid { + return Err(error(format!( + "runtime parameter column {column} violates its finite {:?} domain", + domains[column] + ))); + } + } + Ok(()) +} + fn empty_like(batch: &Batch, columns: usize) -> RusticolResult { match batch.data() { Data::Double(_) => Batch::typed::(vec![], vec![], columns, batch.precision_bits()), @@ -124,6 +161,29 @@ impl TreeEvaluator { momentum_columns: usize, arithmetic: Arithmetic, bits: u32, + ) -> RusticolResult { + Self::with_parameter_domains( + coefficient_state, + norm_state, + states, + momentum_columns, + Vec::new(), + arithmetic, + bits, + ) + } + + /// Decode exact instructions with ordered runtime parameter columns. + /// Parameter rows must share the momentum rows, carrier and working precision; + /// each supplied value must be finite and obey its declared scalar domain. + pub fn with_parameter_domains( + coefficient_state: &[u8], + norm_state: &[u8], + states: Vec, + momentum_columns: usize, + parameter_domains: Vec, + arithmetic: Arithmetic, + bits: u32, ) -> RusticolResult { for state in &states { if state @@ -173,6 +233,7 @@ impl TreeEvaluator { .len() .checked_mul(4) .and_then(|count| momentum_columns.checked_add(count)) + .and_then(|count| count.checked_add(parameter_domains.len())) .ok_or_else(|| error("amplitude input dimensions overflow"))?; let coefficient = exact_state(coefficient_state).map_err(error)?; let norm = exact_state(norm_state).map_err(error)?; @@ -196,6 +257,7 @@ impl TreeEvaluator { norm: CoefficientEvaluator::with_precision_bits(norm_state, arithmetic, bits) .map_err(error)?, states, + parameter_domains, momentum_columns, coefficient_columns, arithmetic, @@ -207,6 +269,10 @@ impl TreeEvaluator { &self.states } + pub fn parameter_domains(&self) -> &[TreeParameterDomain] { + &self.parameter_domains + } + pub fn precision_bits(&self) -> u32 { self.amplitude.precision_bits() } @@ -239,6 +305,33 @@ impl TreeEvaluator { Ok(()) } + fn validate_parameters( + &self, + momenta: &Batch, + parameters: Option<&Batch>, + ) -> RusticolResult<()> { + let Some(parameters) = parameters else { + return if self.parameter_domains.is_empty() { + Ok(()) + } else { + Err(error( + "prepared runtime parameters require an explicit parameter batch", + )) + }; + }; + momenta.same_rows(parameters).map_err(error)?; + if parameters.columns() != self.parameter_domains.len() { + return Err(error( + "runtime parameter columns differ from the prepared domain", + )); + } + match parameters.data() { + Data::Double(values) => parameter_values(values, &self.parameter_domains), + Data::DoubleDouble(values) => parameter_values(values, &self.parameter_domains), + Data::Arbitrary(values) => parameter_values(values, &self.parameter_domains), + } + } + /// Evaluate coherent color coefficients for one declared helicity per state. /// /// Momentum components must be real and finite in the selected carrier. @@ -248,8 +341,28 @@ impl TreeEvaluator { &mut self, momenta: &Batch, helicities: &[i32], + ) -> RusticolResult { + self.evaluate_coefficients_inner(momenta, None, helicities) + } + + /// Evaluate chosen helicities with a typed runtime parameter row per momentum row. + pub fn evaluate_coefficients_with_parameters( + &mut self, + momenta: &Batch, + parameters: &Batch, + helicities: &[i32], + ) -> RusticolResult { + self.evaluate_coefficients_inner(momenta, Some(parameters), helicities) + } + + fn evaluate_coefficients_inner( + &mut self, + momenta: &Batch, + parameters: Option<&Batch>, + helicities: &[i32], ) -> RusticolResult { self.validate_momenta(momenta)?; + self.validate_parameters(momenta, parameters)?; if helicities.len() != self.states.len() || self .states @@ -281,7 +394,11 @@ impl TreeEvaluator { .map_err(error) }) .collect::>>()?; - self.evaluate_external_components(momenta, &components.iter().collect::>()) + self.evaluate_external_components_inner( + momenta, + parameters, + &components.iter().collect::>(), + ) } /// Evaluate with caller-supplied complex external vectors or spinors. @@ -294,8 +411,28 @@ impl TreeEvaluator { &mut self, momenta: &Batch, components: &[&Batch], + ) -> RusticolResult { + self.evaluate_external_components_inner(momenta, None, components) + } + + /// Evaluate supplied external components with runtime parameter columns. + pub fn evaluate_external_components_with_parameters( + &mut self, + momenta: &Batch, + parameters: &Batch, + components: &[&Batch], + ) -> RusticolResult { + self.evaluate_external_components_inner(momenta, Some(parameters), components) + } + + fn evaluate_external_components_inner( + &mut self, + momenta: &Batch, + parameters: Option<&Batch>, + components: &[&Batch], ) -> RusticolResult { self.validate_momenta(momenta)?; + self.validate_parameters(momenta, parameters)?; if components.len() != self.states.len() { return Err(error( "external component block count differs from the declared states", @@ -311,9 +448,12 @@ impl TreeEvaluator { if momenta.rows() == 0 { return empty_like(momenta, self.coefficient_columns); } - let mut inputs = Vec::with_capacity(1 + components.len()); + let mut inputs = Vec::with_capacity(2 + components.len()); inputs.push(momenta); inputs.extend_from_slice(components); + if let Some(parameters) = parameters { + inputs.push(parameters); + } let input = Batch::concat_columns(&inputs).map_err(error)?; let coefficients = self.amplitude.evaluate_batch(&input).map_err(error)?; finite(&coefficients)?; @@ -323,7 +463,25 @@ impl TreeEvaluator { /// Sum the exact-color norm over the declared Cartesian helicity domain. /// No averaging, flux, or identical-particle factors are included. pub fn evaluate_batch(&mut self, momenta: &Batch) -> RusticolResult { + self.evaluate_batch_inner(momenta, None) + } + + /// Sum physical helicities with one typed parameter row for each momentum row. + pub fn evaluate_batch_with_parameters( + &mut self, + momenta: &Batch, + parameters: &Batch, + ) -> RusticolResult { + self.evaluate_batch_inner(momenta, Some(parameters)) + } + + fn evaluate_batch_inner( + &mut self, + momenta: &Batch, + parameters: Option<&Batch>, + ) -> RusticolResult { self.validate_momenta(momenta)?; + self.validate_parameters(momenta, parameters)?; if momenta.rows() == 0 { return empty_like(momenta, 1); } @@ -336,7 +494,8 @@ impl TreeEvaluator { .zip(&positions) .map(|(state, &position)| state.helicities[position]) .collect::>(); - let coefficients = self.evaluate_coefficients(momenta, &helicities)?; + let coefficients = + self.evaluate_coefficients_inner(momenta, parameters, &helicities)?; let norm = self.norm.evaluate_batch(&coefficients).map_err(error)?; finite(&norm)?; accumulated = Some(match accumulated { @@ -551,4 +710,267 @@ mod tests { assert_eq!(values.as_slice(), &[Complex::new(5., 0.)]); assert_eq!(output.ids(), &[31]); } + + fn parameter_evaluator(with_state: bool, arithmetic: Arithmetic, bits: u32) -> TreeEvaluator { + let parameter_start = if with_state { 8 } else { 4 }; + let inputs: Vec<_> = (0..parameter_start + 3) + .map(|i| Atom::var(symbol!(&format!("native_tree_runtime_parameters::x{i}")))) + .collect(); + let state = if with_state { + &inputs[4] + Atom::num(2) * &inputs[7] + } else { + Atom::num(1) + }; + let coefficient = state * &inputs[parameter_start] + + &inputs[parameter_start + 1] * &inputs[parameter_start + 2]; + let exact = Atom::evaluator_multiple(&[coefficient], &inputs) + .cores(1) + .build() + .unwrap(); + let a = Atom::var(symbol!("native_tree_runtime_parameters::a")); + let norm = Atom::evaluator_multiple(&[a.clone().conj() * &a], &[a]) + .cores(1) + .build() + .unwrap(); + TreeEvaluator::with_parameter_domains( + &encode_exact_evaluator_state(&exact).unwrap(), + &encode_exact_evaluator_state(&norm).unwrap(), + if with_state { + vec![states().remove(0)] + } else { + vec![] + }, + 4, + vec![ + TreeParameterDomain::Real, + TreeParameterDomain::PositiveReal, + TreeParameterDomain::Complex, + ], + arithmetic, + bits, + ) + .unwrap() + } + + #[test] + fn runtime_parameters_follow_states_and_preserve_row_identity() { + let mut evaluator = parameter_evaluator(true, Arithmetic::Double, 53); + assert_eq!( + evaluator.parameter_domains(), + &[ + TreeParameterDomain::Real, + TreeParameterDomain::PositiveReal, + TreeParameterDomain::Complex, + ] + ); + let p = Batch::typed( + [2., 0., 0., -2., 2., 0., 0., -2.] + .map(|value| Complex::new(value, 0.)) + .to_vec(), + vec![17, 23], + 4, + 53, + ) + .unwrap(); + let parameters = Batch::typed( + vec![ + Complex::new(3., 0.), + Complex::new(2., 0.), + Complex::new(1., 1.), + Complex::new(1., 0.), + Complex::new(3., 0.), + Complex::new(2., -1.), + ], + vec![17, 23], + 3, + 53, + ) + .unwrap(); + let coefficients = evaluator + .evaluate_coefficients_with_parameters(&p, ¶meters, &[-1]) + .unwrap(); + let Data::Double(values) = coefficients.data() else { + panic!("wrong carrier") + }; + assert_eq!( + values.as_slice(), + &[Complex::new(8., 2.), Complex::new(8., -3.)] + ); + let output = evaluator + .evaluate_batch_with_parameters(&p, ¶meters) + .unwrap(); + let Data::Double(values) = output.data() else { + panic!("wrong carrier") + }; + assert_eq!( + values.as_slice(), + &[Complex::new(172., 0.), Complex::new(86., 0.)] + ); + + let components = Batch::typed( + vec![ + Complex::new(2., 1.), + Complex::new(0., 0.), + Complex::new(0., 0.), + Complex::new(0., 0.), + Complex::new(3., -2.), + Complex::new(0., 0.), + Complex::new(0., 0.), + Complex::new(0., 0.), + ], + vec![17, 23], + 4, + 53, + ) + .unwrap(); + let output = evaluator + .evaluate_external_components_with_parameters(&p, ¶meters, &[&components]) + .unwrap(); + let Data::Double(values) = output.data() else { + panic!("wrong carrier") + }; + assert_eq!( + values.as_slice(), + &[Complex::new(8., 5.), Complex::new(9., -5.)] + ); + + let reversed_p = p.take_rows(vec![1, 0], None).unwrap(); + assert!( + evaluator + .evaluate_batch_with_parameters(&reversed_p, ¶meters) + .is_err() + ); + let reversed_parameters = parameters.take_rows(vec![1, 0], None).unwrap(); + let output = evaluator + .evaluate_batch_with_parameters(&reversed_p, &reversed_parameters) + .unwrap(); + let Data::Double(values) = output.data() else { + panic!("wrong carrier") + }; + assert_eq!( + values.as_slice(), + &[Complex::new(86., 0.), Complex::new(172., 0.)] + ); + assert_eq!(output.ids(), &[23, 17]); + } + + #[test] + fn runtime_parameters_preserve_subbinary64_complex_values() { + let bits = 137; + let mut evaluator = parameter_evaluator(false, Arithmetic::Arbitrary, bits); + let zero = Float::with_val(bits, 0); + let two = Float::with_val(bits, 2); + let delta = two.pow(80).inv(); + let p = Batch::typed( + vec![Complex::new(zero.clone(), zero.clone()); 4], + vec![31], + 4, + bits, + ) + .unwrap(); + let parameters = Batch::typed( + vec![ + Complex::new(Float::with_val(bits, 1) + &delta, zero.clone()), + Complex::new(two.clone(), zero.clone()), + Complex::new(Float::with_val(bits, 3), two.clone()), + ], + vec![31], + 3, + bits, + ) + .unwrap(); + let coefficients = evaluator + .evaluate_coefficients_with_parameters(&p, ¶meters, &[]) + .unwrap(); + let Data::Arbitrary(values) = coefficients.data() else { + panic!("wrong carrier") + }; + assert!(values[0].re > Float::with_val(bits, 7)); + let expected_real = Float::with_val(bits, 7) + δ + assert!((values[0].re.clone() - &expected_real).norm() < two.pow(115).inv()); + assert!((values[0].im.clone() - Float::with_val(bits, 4)).norm() < two.pow(115).inv()); + let output = evaluator + .evaluate_batch_with_parameters(&p, ¶meters) + .unwrap(); + let Data::Arbitrary(values) = output.data() else { + panic!("wrong carrier") + }; + let expected = expected_real.clone() * &expected_real + Float::with_val(bits, 16); + assert!(values[0].re > Float::with_val(bits, 65)); + assert!((values[0].re.clone() - expected).norm() < two.pow(115).inv()); + assert!(values[0].im.norm() < two.pow(115).inv()); + assert_eq!(output.precision_bits(), bits); + assert!(values[0].re.get_precision() > 100); + } + + #[test] + fn runtime_parameter_domains_and_empty_batches_are_checked_before_evaluation() { + let mut evaluator = parameter_evaluator(false, Arithmetic::Double, 53); + let p = Batch::typed(vec![Complex::new(0., 0.); 4], vec![31], 4, 53).unwrap(); + let valid = vec![ + Complex::new(1., 0.), + Complex::new(2., 0.), + Complex::new(3., 4.), + ]; + for (column, value) in [ + (0, Complex::new(1., 1.)), + (1, Complex::new(0., 0.)), + (1, Complex::new(-1., 0.)), + (1, Complex::new(1., 1.)), + (0, Complex::new(f64::INFINITY, 0.)), + (1, Complex::new(f64::INFINITY, 0.)), + (2, Complex::new(f64::NAN, 0.)), + (2, Complex::new(1., f64::INFINITY)), + ] { + let mut values = valid.clone(); + values[column] = value; + let parameters = Batch::typed(values, vec![31], 3, 53).unwrap(); + assert!( + evaluator + .evaluate_batch_with_parameters(&p, ¶meters) + .is_err() + ); + } + let valid = Batch::typed(valid, vec![31], 3, 53).unwrap(); + assert!(evaluator.evaluate_batch_with_parameters(&p, &valid).is_ok()); + + let empty = Batch::typed::(vec![], vec![], 4, 53).unwrap(); + assert!(evaluator.evaluate_batch(&empty).is_err()); + assert!(evaluator.evaluate_coefficients(&empty, &[]).is_err()); + assert!(evaluator.evaluate_external_components(&empty, &[]).is_err()); + let parameters = Batch::typed::(vec![], vec![], 3, 53).unwrap(); + assert_eq!( + evaluator + .evaluate_batch_with_parameters(&empty, ¶meters) + .unwrap() + .rows(), + 0 + ); + assert_eq!( + evaluator + .evaluate_coefficients_with_parameters(&empty, ¶meters, &[]) + .unwrap() + .rows(), + 0 + ); + assert_eq!( + evaluator + .evaluate_external_components_with_parameters(&empty, ¶meters, &[]) + .unwrap() + .rows(), + 0 + ); + for invalid in [ + Batch::typed::(vec![], vec![], 3, 53).unwrap(), + Batch::typed::(vec![], vec![], 3, 106).unwrap(), + Batch::typed::(vec![], vec![], 2, 53).unwrap(), + valid, + ] { + assert!( + evaluator + .evaluate_batch_with_parameters(&empty, &invalid) + .is_err() + ); + } + } }