1use color_eyre::Result;
2use eyre::{Context, eyre};
3use linnet::half_edge::subgraph::{SubGraphLike, SubSetOps};
4use spenso::shadowing::symbolica_utils::SpensoPrintSettings;
5use symbolica::atom::{Atom, AtomCore};
6
7use crate::{
8 cff::CutCFFIndex,
9 graph::{
10 FeynmanGraph, Graph,
11 cuts::{CutSet, ResidueSelector},
12 parse::string_utils::dot_statement_value,
13 },
14 integrands::process::ProcessIntegrand,
15 processes::DotExportSettings,
16 settings::global::GenerationSettings,
17 uv::{
18 UVOrchestrator,
19 approx::{CutStructure, OrientationProjection},
20 forest::CutForests,
21 hedge_poset::Wood as HedgePosetWood,
22 wood::CutWoods,
23 },
24};
25
26pub struct UVForestExportSettings {
27 pub computed: bool,
28}
29
30pub struct UVForestExport {
31 pub graph_name: String,
32 pub forest_dot: String,
33 pub node_terms: Vec<UVForestNodeTerm>,
34}
35
36pub struct UVForestNodeTerm {
37 pub forest_index: usize,
38 pub node_index: usize,
39 pub node_key: String,
40 pub term_index: usize,
41 pub residue_index: CutCFFIndex,
42 pub dot: String,
43}
44
45impl UVForestNodeTerm {
46 pub fn file_name(&self) -> String {
47 let key = sanitize_file_component(&self.node_key);
48 let residue = residue_suffix(self.residue_index);
49 format!(
50 "node_{:03}_{}_term_{:03}_{}.dot",
51 self.node_index, key, self.term_index, residue
52 )
53 }
54}
55
56pub(crate) struct UVForestNodeExpression {
57 pub forest_index: usize,
58 pub node_index: usize,
59 pub node_key: String,
60 pub term_index: usize,
61 pub residue_index: CutCFFIndex,
62 pub numerator: Atom,
63}
64
65impl ProcessIntegrand {
66 pub fn export_uv_forest_graph(
67 &self,
68 graph_id: usize,
69 generation_settings: &GenerationSettings,
70 export_settings: &UVForestExportSettings,
71 ) -> Result<UVForestExport> {
72 match self {
73 Self::Amplitude(integrand) => {
74 let term = integrand.data.graph_terms.get(graph_id).ok_or_else(|| {
75 eyre!(
76 "Graph id {} is out of range for amplitude integrand {}",
77 graph_id,
78 integrand.data.name
79 )
80 })?;
81 let cut_structure = CutStructure::empty(&term.graph);
82 let post_factor = (Atom::var(crate::utils::GS.pi) * Atom::num(2))
83 .pow(3 * term.graph.get_loop_number() as i64);
84 export_graph(
85 &term.graph,
86 cut_structure,
87 term.orientations.iter().cloned().collect(),
88 generation_settings,
89 export_settings,
90 |atom| atom / &post_factor,
91 )
92 }
93 Self::CrossSection(integrand) => {
94 let term = integrand.data.graph_terms.get(graph_id).ok_or_else(|| {
95 eyre!(
96 "Graph id {} is out of range for cross-section integrand {}",
97 graph_id,
98 integrand.data.name
99 )
100 })?;
101 let cuts = term
102 .cut_group_data
103 .cut_groups
104 .iter()
105 .map(|cuts| CutSet {
106 residue_selector: ResidueSelector {
107 lu_cut: Some(cuts.related_esurface_group.clone()),
108 left_th_cut: None,
109 right_th_cut: None,
110 },
111 union: cuts
112 .cuts
113 .iter()
114 .map(|cut_id| term.cuts[*cut_id].cut.as_subgraph())
115 .reduce(|cut_1, cut_2| cut_1.union(&cut_2))
116 .unwrap_or_else(|| term.graph.empty_subgraph()),
117 canonicalize_external_shifts: false,
118 })
119 .collect();
120 let cut_structure = CutStructure { cuts };
121 let loop_number = term.graph.cyclotomatic_number(&term.graph.full_filter())
122 - term.graph.initial_state_cut.nedges(&term.graph);
123 let loop_3 = loop_number as i64 * 3;
124 let lu_prefactor = Atom::var(crate::utils::GS.rescale_star).pow(loop_3)
125 * Atom::var(crate::utils::GS.hfunction_lu_cut)
126 / (Atom::num(2) * Atom::var(crate::utils::GS.pi)).pow(loop_3 - 1);
127
128 export_graph(
129 &term.graph,
130 cut_structure,
131 term.orientations.iter().cloned().collect(),
132 generation_settings,
133 export_settings,
134 |atom| atom * &lu_prefactor,
135 )
136 }
137 }
138 }
139}
140
141fn export_graph(
142 graph: &Graph,
143 cut_structure: CutStructure,
144 orientations: Vec<
145 linnet::half_edge::involution::EdgeVec<linnet::half_edge::involution::Orientation>,
146 >,
147 generation_settings: &GenerationSettings,
148 export_settings: &UVForestExportSettings,
149 mut post_process: impl FnMut(Atom) -> Atom,
150) -> Result<UVForestExport> {
151 if generation_settings.uv.orchestrator == UVOrchestrator::Compare {
152 return Err(eyre!(
153 "UV forest export does not support uv.orchestrator = compare"
154 ));
155 }
156
157 let mut forest_dot = String::new();
158 let mut node_terms = Vec::new();
159 for (forest_index, cut) in cut_structure.cuts.into_iter().enumerate() {
160 let single_cut = CutStructure { cuts: vec![cut] };
161 let forest_name = format!(
162 "uv_{}_forest_{forest_index:03}",
163 sanitize_file_component(&graph.name)
164 );
165 let mut graph = graph.clone();
166 match generation_settings.uv.orchestrator {
167 UVOrchestrator::LegacyDagForest => export_legacy_forest(
168 forest_index,
169 &forest_name,
170 &mut graph,
171 single_cut,
172 &orientations,
173 generation_settings,
174 export_settings,
175 &mut post_process,
176 &mut forest_dot,
177 &mut node_terms,
178 )?,
179 UVOrchestrator::HedgePoset => export_hedge_poset_forest(
180 forest_index,
181 &forest_name,
182 &mut graph,
183 single_cut,
184 &orientations,
185 generation_settings,
186 export_settings,
187 &mut post_process,
188 &mut forest_dot,
189 &mut node_terms,
190 )?,
191 UVOrchestrator::Compare => unreachable!("compare is rejected before export"),
192 }
193 }
194
195 Ok(UVForestExport {
196 graph_name: graph.name.clone(),
197 forest_dot,
198 node_terms,
199 })
200}
201
202#[allow(clippy::too_many_arguments)]
203fn export_legacy_forest(
204 forest_index: usize,
205 forest_name: &str,
206 graph: &mut Graph,
207 cut_structure: CutStructure,
208 orientations: &[linnet::half_edge::involution::EdgeVec<
209 linnet::half_edge::involution::Orientation,
210 >],
211 generation_settings: &GenerationSettings,
212 export_settings: &UVForestExportSettings,
213 post_process: &mut impl FnMut(Atom) -> Atom,
214 forest_dot: &mut String,
215 node_terms: &mut Vec<UVForestNodeTerm>,
216) -> Result<()> {
217 let cut_woods = CutWoods::new(cut_structure, graph, &generation_settings.uv);
218 let mut cut_forests = cut_woods.unfold(graph);
219 let Some(forest) = cut_forests.forests.first_mut() else {
220 return Err(eyre!("Legacy UV exporter produced no forest"));
221 };
222 forest_dot.push_str(&forest.dot_serialize_for_export(forest_name));
223 forest_dot.push('\n');
224
225 if !export_settings.computed {
226 return Ok(());
227 }
228
229 compute_legacy_forest(graph, &mut cut_forests, orientations, generation_settings)?;
230 let forest = cut_forests
231 .forests
232 .first()
233 .expect("legacy forest exists after compute");
234 node_terms.extend(
235 forest
236 .export_node_expressions(graph, forest_index, post_process)?
237 .into_iter()
238 .map(|term| node_expression_to_dot(graph, forest_name, term))
239 .collect::<Result<Vec<_>>>()?,
240 );
241 Ok(())
242}
243
244fn compute_legacy_forest(
245 graph: &mut Graph,
246 cut_forests: &mut CutForests,
247 orientations: &[linnet::half_edge::involution::EdgeVec<
248 linnet::half_edge::involution::Orientation,
249 >],
250 generation_settings: &GenerationSettings,
251) -> Result<()> {
252 cut_forests.compute(
253 graph,
254 crate::utils::vakint()?,
255 OrientationProjection::new(orientations, &generation_settings.orientation_pattern),
256 &generation_settings.uv,
257 )
258}
259
260#[allow(clippy::too_many_arguments)]
261fn export_hedge_poset_forest(
262 forest_index: usize,
263 forest_name: &str,
264 graph: &mut Graph,
265 cut_structure: CutStructure,
266 orientations: &[linnet::half_edge::involution::EdgeVec<
267 linnet::half_edge::involution::Orientation,
268 >],
269 generation_settings: &GenerationSettings,
270 export_settings: &UVForestExportSettings,
271 post_process: &mut impl FnMut(Atom) -> Atom,
272 forest_dot: &mut String,
273 node_terms: &mut Vec<UVForestNodeTerm>,
274) -> Result<()> {
275 let wood = HedgePosetWood::new(cut_structure, graph, &generation_settings.uv);
276 let mut forests = wood.unfold();
277 forest_dot.push_str(&name_dot_graph(forests.dot_serialize(), forest_name));
278 forest_dot.push('\n');
279
280 if !export_settings.computed {
281 return Ok(());
282 }
283
284 forests.compute(
285 graph,
286 crate::utils::vakint()?,
287 OrientationProjection::new(orientations, &generation_settings.orientation_pattern),
288 &generation_settings.uv,
289 )?;
290 node_terms.extend(
291 forests
292 .export_node_expressions(forest_index, post_process)?
293 .into_iter()
294 .map(|term| node_expression_to_dot(graph, forest_name, term))
295 .collect::<Result<Vec<_>>>()?,
296 );
297 Ok(())
298}
299
300fn node_expression_to_dot(
301 graph: &Graph,
302 forest_name: &str,
303 term: UVForestNodeExpression,
304) -> Result<UVForestNodeTerm> {
305 let graph_name = format!(
306 "{}_node_{:03}_term_{:03}",
307 forest_name, term.node_index, term.term_index
308 );
309 let mut dot_graph = graph
310 .with_global_numerator_only(graph_name, term.numerator.clone())
311 .to_dot_graph_with_settings(&DotExportSettings {
312 split_xs_by_initial_states: true,
313 output_full_numerator: false,
314 ..DotExportSettings::default()
315 });
316 dot_graph
317 .global_data
318 .statements
319 .insert("forest_name".into(), dot_statement_value(forest_name));
320 dot_graph
321 .global_data
322 .statements
323 .insert("forest_index".into(), term.forest_index.to_string());
324 dot_graph
325 .global_data
326 .statements
327 .insert("forest_node_index".into(), term.node_index.to_string());
328 dot_graph.global_data.statements.insert(
329 "forest_node_key".into(),
330 dot_statement_value(&term.node_key),
331 );
332 dot_graph
333 .global_data
334 .statements
335 .insert("forest_term_index".into(), term.term_index.to_string());
336 dot_graph.global_data.statements.insert(
337 "forest_residue_index".into(),
338 dot_statement_value(&residue_suffix(term.residue_index)),
339 );
340 let full_num = term
341 .numerator
342 .printer(SpensoPrintSettings::typst_options())
343 .to_string();
344 dot_graph
345 .global_data
346 .statements
347 .insert("full_num".into(), dot_statement_value(&full_num));
348
349 let mut dot = Vec::new();
350 dot_graph
351 .write_io(&mut dot)
352 .context("while serializing UV forest node DOT")?;
353 Ok(UVForestNodeTerm {
354 forest_index: term.forest_index,
355 node_index: term.node_index,
356 node_key: term.node_key,
357 term_index: term.term_index,
358 residue_index: term.residue_index,
359 dot: String::from_utf8(dot).context("DOT graph serialization was not UTF-8")?,
360 })
361}
362
363fn name_dot_graph(dot: String, name: &str) -> String {
364 dot.replacen("digraph {", &format!("digraph {name} {{"), 1)
365 .replacen("digraph Poset {", &format!("digraph {name} {{"), 1)
366}
367
368pub(crate) fn sanitize_file_component(value: &str) -> String {
369 let mut sanitized = String::new();
370 for c in value.chars() {
371 let next = if c.is_ascii_alphanumeric() || c == '-' {
372 c
373 } else {
374 '_'
375 };
376 if next == '_' {
377 if !sanitized.ends_with('_') {
378 sanitized.push(next);
379 }
380 } else {
381 sanitized.push(next);
382 }
383 }
384 sanitized.truncate(48);
385 sanitized = sanitized.trim_matches('_').to_string();
386 if sanitized.is_empty() {
387 "key".to_string()
388 } else {
389 sanitized
390 }
391}
392
393fn residue_suffix(index: CutCFFIndex) -> String {
394 let suffix = index.to_string();
395 if suffix.is_empty() {
396 "all_none".to_string()
397 } else {
398 sanitize_file_component(&suffix)
399 }
400}
401
402#[cfg(test)]
403mod tests {
404 use symbolica::{atom::Atom, function};
405
406 use super::{
407 UVForestNodeExpression, UVForestNodeTerm, node_expression_to_dot, sanitize_file_component,
408 };
409 use crate::{
410 cff::CutCFFIndex,
411 dot,
412 graph::{Graph, parse::IntoGraph},
413 initialisation::test_initialise,
414 utils::GS,
415 };
416
417 #[test]
418 fn node_term_file_name_contains_stable_indices_key_and_residue() {
419 let term = UVForestNodeTerm {
420 forest_index: 2,
421 node_index: 4,
422 node_key: "{1,2}(3)".to_string(),
423 term_index: 7,
424 residue_index: CutCFFIndex {
425 left_threshold_order: None,
426 right_threshold_order: None,
427 lu_cut_order: Some(1),
428 },
429 dot: String::new(),
430 };
431
432 assert_eq!(term.file_name(), "node_004_1_2_3_term_007_lu_cut_1.dot");
433 }
434
435 #[test]
436 fn node_term_file_name_uses_all_none_residue_suffix() {
437 let term = UVForestNodeTerm {
438 forest_index: 0,
439 node_index: 0,
440 node_key: "!!!".to_string(),
441 term_index: 0,
442 residue_index: CutCFFIndex::new_all_none(),
443 dot: String::new(),
444 };
445
446 assert_eq!(term.file_name(), "node_000_key_term_000_all_none.dot");
447 }
448
449 #[test]
450 fn file_components_cannot_escape_the_export_directory() {
451 assert_eq!(sanitize_file_component("../result"), "result");
452 assert_eq!(sanitize_file_component("/tmp/result"), "tmp_result");
453 assert_eq!(sanitize_file_component("a/b"), "a_b");
454 assert_eq!(sanitize_file_component(""), "key");
455 }
456
457 #[test]
458 fn computed_node_full_numerator_is_a_spenso_aware_typst_fragment() {
459 test_initialise().unwrap();
460 let graph: Graph = dot!(digraph G {
461 ext [style=invis]
462 node [num=1]
463 ext -> A
464 C -> A
465 A -> D
466 D -> B
467 B -> C
468 C -> D
469 B -> ext
470 })
471 .unwrap();
472 let term = UVForestNodeExpression {
473 forest_index: 0,
474 node_index: 1,
475 node_key: "{1}".to_string(),
476 term_index: 2,
477 residue_index: CutCFFIndex::new_all_none(),
478 numerator: function!(GS.uv_truncate, Atom::num(1)),
479 };
480
481 let exported = node_expression_to_dot(&graph, "uv_typst", term).unwrap();
482
483 assert!(exported.dot.contains(r#"full_num = "op(\"Tr\")(1)";"#));
484 assert!(!exported.dot.contains("gammalooprs::Truncate"));
485 }
486}