gammalooprs/processes/

mod.rs

use std::{fs, path::Path};

// use bincode::{Decode, Encode};
use bincode_trait_derive::{Decode, Encode};
use color_eyre::Result;
use eyre::{Context, eyre};
use rayon::ThreadPool;
use schemars::JsonSchema;
use symbolica::evaluate::OptimizationSettings;
use tracing::debug;

use crate::{
    GammaLoopContext, GammaLoopContextContainer,
    settings::{GlobalSettings, runtime::LockedRuntimeSettings},
    utils::serde_utils::{IsDefault, is_false, is_true, is_usize, show_defaults_helper},
    uv::export::UVForestExportSettings,
};
use serde::{Deserialize, Serialize};

use crate::model::Model;

mod generation_report;
pub use generation_report::{
    EvaluatorBuildTimings, GeneratedGraphKey, GeneratedGraphReport, GraphGenerationStats,
    NamedGraphGenerationReport, merge_generated_graph_reports,
};
mod generation_progress;
pub use generation_progress::{
    GenerationProcessKind, GenerationProgressMode, GenerationProgressModeGuard,
    GenerationProgressObserver, GenerationProgressObserverGuard, GenerationProgressPhase,
    begin_phase, cut_finished, detailed_progress_enabled, detailed_progress_message,
    enter_detailed_progress_span, enter_progress_context,
};
mod selection;
pub use selection::{
    CycleSignature, GraphGroupSelectionMode, GraphGroupSelectionPlan, GraphGroupSelectionReport,
    GraphGroupSelectionSpec, GraphSelectionSignatureInventory, ParticleSignature,
    RaisedCutSignatureInventory, RaisedPropagatorScope, RaisedPropagatorSignature,
    SelectionPolarity, VertexSignature,
};
pub(crate) use selection::{GraphCutSelectionSubject, GraphSelectionSubject};

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(
    Debug, Clone, Copy, Default, Serialize, Deserialize, Encode, Decode, PartialEq, Eq, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum TensorNetworkContractionOrder {
    #[default]
    SparseAtomAware,
    AtomAware,
    ResultRankOnly,
    EntryAware,
}

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
pub enum ExecutionMode {
    Sequential,
    Parallel,
    SequentialRef,
    SequentialExtract,
}

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
pub enum ContractionMode {
    SmallestDegree,
    MinResultRank,
}
#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Encode, Decode, PartialEq, JsonSchema)]
pub struct EvaluatorSettings {
    /// Perform symbolic tensor and scalar algebra before lowering to an evaluator.
    #[serde(default, skip_serializing_if = "is_false")]
    pub do_algebra: bool,
    /// Re-optimize orientation expressions iteratively while building the evaluator.
    #[serde(
        default = "evaluator_default_iterative_orientation_optimization",
        skip_serializing_if = "is_true"
    )]
    pub iterative_orientation_optimization: bool,
    /// Build one evaluator for the sum of all contributions instead of separate evaluators.
    #[serde(default, skip_serializing_if = "is_false")]
    pub summed: bool,
    /// Represent a summed evaluator through a shared Symbolica function map.
    #[serde(default, skip_serializing_if = "is_false")]
    pub summed_function_map: bool,
    /// Compile the lowered evaluator using the configured generation compilation backend.
    #[serde(default, skip_serializing_if = "is_false")]
    pub compile: bool,
    /// Retain the intermediate Symbolica atom alongside the executable evaluator.
    #[serde(default, skip_serializing_if = "is_false")]
    pub store_atom: bool,
    /// Apply function-map replacements before numerical lowering.
    #[serde(default, skip_serializing_if = "is_false")]
    pub do_fn_map_replacements: bool,
    /// Translate compatible symbolic expressions directly into Spenso networks.
    #[serde(
        default = "evaluator_default_direct_translation",
        skip_serializing_if = "is_true"
    )]
    pub direct_translation: bool,
    /// Number of Horner-scheme optimization passes applied to scalar expressions.
    #[serde(
        default = "evaluator_default_horner_iterations",
        skip_serializing_if = "is_usize::<1>"
    )]
    pub horner_iterations: usize,
    /// Worker threads used while constructing one evaluator.
    #[serde(
        default = "evaluator_default_n_cores",
        skip_serializing_if = "is_usize::<1>"
    )]
    pub n_cores: usize,
    /// Optional number of common-pair-elimination optimization passes.
    #[serde(default, skip_serializing_if = "IsDefault::is_default")]
    pub cpe_iterations: Option<usize>,
    /// Optimization stage at which evaluator construction aborts for diagnostic isolation.
    #[serde(default, skip_serializing_if = "IsDefault::is_default")]
    pub abort_level: usize,

    /// Maximum variable count for attempting a Horner-scheme optimization.
    #[serde(
        default = "evaluator_default_max_horner_scheme_variables",
        skip_serializing_if = "is_usize::<500>"
    )]
    pub max_horner_scheme_variables: usize,

    /// Maximum cached candidate pairs during common-pair elimination.
    #[serde(
        default = "evaluator_default_max_common_pair_cache_entries",
        skip_serializing_if = "is_usize::<1000000>"
    )]
    pub max_common_pair_cache_entries: usize,

    /// Maximum expression-tree distance considered during common-pair elimination.
    #[serde(
        default = "evaluator_default_max_common_pair_distance",
        skip_serializing_if = "is_usize::<1000>"
    )]
    pub max_common_pair_distance: usize,
    /// Heuristic used to order tensor-network contractions.
    #[serde(default, skip_serializing_if = "IsDefault::is_default")]
    pub tensor_network_contraction_order: TensorNetworkContractionOrder,
    /// Emit detailed evaluator-construction diagnostics.
    #[serde(default, skip_serializing_if = "is_false")]
    pub verbose: bool,

    /// Spenso execution and contraction modes used by the lowered tensor network.
    #[serde(
        default = "evaluator_default_spenso_execution_mode",
        skip_serializing_if = "is_default_spenso_execution_mode"
    )]
    pub spenso_execution_mode: (ExecutionMode, ContractionMode),
}

const fn evaluator_default_iterative_orientation_optimization() -> bool {
    true
}

const fn evaluator_default_direct_translation() -> bool {
    true
}

const fn evaluator_default_horner_iterations() -> usize {
    1
}

const fn evaluator_default_n_cores() -> usize {
    1
}

const fn evaluator_default_max_horner_scheme_variables() -> usize {
    500
}

const fn evaluator_default_max_common_pair_cache_entries() -> usize {
    1_000_000
}

const fn evaluator_default_max_common_pair_distance() -> usize {
    1000
}

const fn evaluator_default_spenso_execution_mode() -> (ExecutionMode, ContractionMode) {
    (ExecutionMode::Sequential, ContractionMode::MinResultRank)
}

fn is_default_spenso_execution_mode(mode: &(ExecutionMode, ContractionMode)) -> bool {
    show_defaults_helper(mode == &evaluator_default_spenso_execution_mode())
}

impl Default for EvaluatorSettings {
    fn default() -> Self {
        Self {
            iterative_orientation_optimization:
                evaluator_default_iterative_orientation_optimization(),
            summed: false,
            do_algebra: false,
            summed_function_map: false,
            direct_translation: evaluator_default_direct_translation(),
            compile: false,
            do_fn_map_replacements: false,
            store_atom: false,
            horner_iterations: evaluator_default_horner_iterations(),
            n_cores: evaluator_default_n_cores(),
            cpe_iterations: None,
            abort_level: 0,
            max_horner_scheme_variables: evaluator_default_max_horner_scheme_variables(),
            max_common_pair_cache_entries: evaluator_default_max_common_pair_cache_entries(),
            max_common_pair_distance: evaluator_default_max_common_pair_distance(),
            tensor_network_contraction_order: TensorNetworkContractionOrder::default(),
            verbose: false,
            spenso_execution_mode: evaluator_default_spenso_execution_mode(),
        }
    }
}

impl EvaluatorSettings {
    pub fn optimization_settings(&self) -> OptimizationSettings {
        OptimizationSettings::new()
            .horner_iterations(self.horner_iterations)
            .cores(self.n_cores)
            .cpe_iterations(self.cpe_iterations)
            .abort_check(Some(Box::new(
                crate::is_interrupt_requested as fn() -> bool,
            )))
            .abort_level(self.abort_level)
            .max_horner_scheme_variables(self.max_horner_scheme_variables)
            .max_common_pair_cache_entries(self.max_common_pair_cache_entries)
            .max_common_pair_distance(self.max_common_pair_distance)
            .verbose(self.verbose)
            .direct_translation(self.direct_translation)
    }
}

#[derive(Clone, Encode, Decode)]
#[trait_decode(trait = GammaLoopContext)]
pub struct ProcessList {
    // pub amplitude_process: Vec<Process<Amplitude>>
    // pub cross_section_process: Vec<Process<CrossSection>>
    pub processes: Vec<Process>,
}

/// Controls which graph details and algebraic transformations are included in DOT exports.
#[derive(Serialize, Deserialize, Debug, Clone)]
#[cfg_attr(feature = "python_stubgen", pyo3_stub_gen::derive::gen_stub_pyclass)]
#[cfg_attr(
    feature = "python_api",
    pyo3::pyclass(from_py_object, get_all, set_all)
)]
pub struct DotExportSettings {
    // pub root_folder: PathBuf,
    /// Write all diagrams of an integrand to one file during filesystem export.
    pub combine_diagrams: bool,
    /// Request UV counterterm graph output; currently retained for compatibility.
    pub with_uv: bool,
    /// Add the complete symbolic numerator as the graph-level `full_num` field.
    pub output_full_numerator: bool,
    /// Split cross-section graphs by distinct initial-state assignments.
    pub split_xs_by_initial_states: bool,
    /// Simplify gamma-matrix algebra in the full numerator before formatting it.
    pub do_gamma_algebra: bool,
    /// Simplify color algebra in the full numerator before formatting it.
    pub do_color_algebra: bool,
    #[serde(default, skip_serializing_if = "is_false")]
    /// Include vertex and edge fields whose values GammaLoop generated automatically.
    pub include_autogenerated_fields: bool,
}

#[derive(Serialize, Deserialize, Debug, Clone)]
#[cfg_attr(
    feature = "python_api",
    pyo3::pyclass(from_py_object, get_all, set_all)
)]
pub struct StandaloneExportSettings {
    #[serde(default)]
    pub mode: StandaloneExportMode,
    #[serde(default)]
    pub format: StandaloneDataFormat,
    #[serde(default)]
    pub precision: StandaloneNumericTarget,
}

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(
    Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default, Encode, Decode, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum StandaloneExportMode {
    #[default]
    Rust,
    Python,
}

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(
    Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default, Encode, Decode, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum StandaloneDataFormat {
    #[default]
    Binary,
    Json,
}

#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
#[derive(
    Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Default, Encode, Decode, JsonSchema,
)]
#[serde(rename_all = "snake_case")]
pub enum StandaloneNumericTarget {
    #[default]
    Double,
    Quad,
    Arb,
}

#[cfg(feature = "python_api")]
#[cfg_attr(feature = "python_stubgen", pyo3_stub_gen::derive::gen_stub_pymethods)]
#[cfg_attr(feature = "python_api", pyo3::pymethods)]
impl DotExportSettings {
    /// Construct DOT export settings with GammaLoop's defaults.
    #[new]
    fn new() -> Self {
        DotExportSettings::default()
    }
}

impl Default for DotExportSettings {
    fn default() -> Self {
        Self {
            // root_folder: PathBuf::from("."),
            split_xs_by_initial_states: false,
            output_full_numerator: false,
            do_gamma_algebra: false,
            do_color_algebra: true,
            combine_diagrams: false,
            with_uv: false,
            include_autogenerated_fields: false,
        }
    }
}

impl Default for StandaloneExportSettings {
    fn default() -> Self {
        Self {
            format: StandaloneDataFormat::Binary,
            mode: StandaloneExportMode::Rust,
            precision: StandaloneNumericTarget::Double,
        }
    }
}

impl Default for ProcessList {
    fn default() -> Self {
        Self::new()
    }
}

// Python Api
impl ProcessList {
    /// Generates a new empty process list
    pub fn new() -> Self {
        ProcessList { processes: vec![] }
    }

    pub fn warm_up_all_processes(&mut self, model: &Model) -> Result<()> {
        for process in &mut self.processes.iter_mut() {
            process.warm_up(model)?;
        }

        Ok(())
    }

    pub fn load(path: impl AsRef<Path>, context: GammaLoopContextContainer) -> Result<Self> {
        let mut process_list = Self::new();

        let path = path.as_ref().join("processes");
        let amplitudes_path = path.join("amplitudes");
        if amplitudes_path.exists() {
            debug!("Looking for amplitudes in {}", amplitudes_path.display());

            for entry in fs::read_dir(amplitudes_path).context("Error reading dir")? {
                let Ok(entry) = entry else {
                    debug!("Skipping invalid entry");

                    continue;
                };
                let path = entry.path();
                process_list.processes.push(
                    Process::load_amplitude(path, context).context("Error loading amplitude")?,
                );
            }
        }

        let cross_sections_path = path.join("cross_sections");
        if cross_sections_path.exists() {
            debug!(
                "Looking for cross sections in {}",
                cross_sections_path.display()
            );

            for entry in fs::read_dir(cross_sections_path)? {
                let entry = entry?;
                let path = entry.path();
                process_list
                    .processes
                    .push(Process::load_cross_section(path, context)?);
            }
        }
        process_list
            .processes
            .sort_by_key(|p| p.definition.process_id);

        Ok(process_list)
    }

    pub fn save(&mut self, folder: impl AsRef<Path>, override_existing: bool) -> Result<()> {
        let path = folder.as_ref().join("processes");

        let r = fs::create_dir_all(&path);
        if !override_existing {
            r?;
        }

        for p in self.processes.iter_mut() {
            p.save(&path, override_existing)?;
        }

        Ok(())
    }

    pub fn compile(
        &mut self,
        folder: impl AsRef<Path>,
        override_existing: bool,
        process_id: Option<usize>,
        integrand_name: Option<String>,
        thread_pool: &ThreadPool,
    ) -> Result<Vec<GeneratedGraphReport>> {
        let path = folder.as_ref().join("processes");

        let r = fs::create_dir_all(&path);
        if !override_existing {
            r?;
        }

        let mut reports = Vec::new();
        for p in self.processes.iter_mut() {
            if let Some(id) = process_id
                && p.definition.process_id != id
            {
                continue;
            }
            reports.extend(p.compile(
                &path,
                override_existing,
                integrand_name.clone(),
                thread_pool,
            )?);
        }

        Ok(reports)
    }

    pub fn activate_loaded_integrand_backends(
        &mut self,
        allow_symjit_fallback: bool,
    ) -> Result<()> {
        for process in &mut self.processes {
            process.activate_loaded_integrand_backends(allow_symjit_fallback)?;
        }
        Ok(())
    }

    pub fn get_integrand(
        &self,
        process_id: usize,
        integrand_name: impl AsRef<str>,
    ) -> Result<crate::processes::process::ResolvedIntegrandRef<'_>> {
        let process = &self.processes[process_id];
        process.get_integrand(integrand_name)
    }

    pub fn get_integrand_mut(
        &mut self,
        process_id: usize,
        integrand_name: impl AsRef<str>,
    ) -> Result<&mut crate::integrands::process::ProcessIntegrand> {
        let process = &mut self.processes[process_id];
        process.get_integrand_mut(integrand_name)
    }

    pub fn export_dot(&self, path: impl AsRef<Path>, settings: &DotExportSettings) -> Result<()> {
        let path = path.as_ref().join("processes");
        fs::create_dir_all(&path)?;

        for p in self.processes.iter() {
            p.export_dot(&path, settings)?;
        }
        Ok(())
    }

    pub fn export_uv_forests(
        &self,
        path: impl AsRef<Path>,
        process_id: usize,
        integrand_name: &str,
        graph_ids: &[usize],
        settings: &UVForestExportSettings,
    ) -> Result<()> {
        let path = path.as_ref().join("processes");
        fs::create_dir_all(&path)?;
        let process = self.processes.get(process_id).ok_or_else(|| {
            eyre!(
                "Process id {} is out of range; there are {} processes",
                process_id,
                self.processes.len()
            )
        })?;
        process.export_uv_forests(&path, integrand_name, graph_ids, settings)
    }

    pub fn export_standalone(
        &self,
        path: impl AsRef<Path>,
        settings: &StandaloneExportSettings,
    ) -> Result<()> {
        let path = path.as_ref().join("processes");
        fs::create_dir_all(&path)?;

        for p in self.processes.iter() {
            p.export_standalone(&path, settings)?;
        }
        Ok(())
    }

    pub fn add_process(&mut self, process: Process) {
        self.processes.push(process);
    }

    ///preprocesses the process list according to the settings
    pub fn preprocess(
        &mut self,
        model: &Model,
        settings: &GlobalSettings,
        locked_runtime_settings: &LockedRuntimeSettings,
        thread_pool: &ThreadPool,
    ) -> Result<Vec<GeneratedGraphReport>> {
        let mut reports = Vec::new();
        for process in self.processes.iter_mut() {
            reports.extend(process.preprocess(
                model,
                settings,
                locked_runtime_settings,
                thread_pool,
            )?);
        }
        Ok(reports)
    }

    pub fn generate_integrands(
        &mut self,
        model: &Model,
        global_settings: &GlobalSettings,
        runtime_default: LockedRuntimeSettings,
        thread_pool: &ThreadPool,
    ) -> Result<Vec<GeneratedGraphReport>> {
        let mut reports = Vec::new();
        for process in &mut self.processes {
            reports.extend(process.generate_integrands(
                model,
                global_settings,
                runtime_default,
                thread_pool,
            )?);
        }
        Ok(reports)
    }

    pub fn find_process(&self, process_id: Option<usize>) -> Result<usize> {
        if let Some(id) = process_id {
            if id >= self.processes.len() {
                return Err(color_eyre::eyre::eyre!(
                    "Invalid process id {}. Number of processes: {}",
                    id,
                    self.processes.len()
                ));
            }
            Ok(id)
        } else if self.processes.is_empty() {
            Err(color_eyre::eyre::eyre!("No processes generated yet."))
        } else if self.processes.len() > 1 {
            Err(color_eyre::eyre::eyre!(
                "There are {} processes available. Please specify a process.",
                self.processes.len()
            ))
        } else {
            Ok(0)
        }
    }

    pub fn find_integrand(
        &self,
        process_id: Option<usize>,
        integrand_name: Option<&String>,
    ) -> Result<(usize, String)> {
        let p_id = self.find_process(process_id)?;
        let integrand_name = self.processes[p_id]
            .collection
            .find_integrand(integrand_name.cloned())?;
        Ok((p_id, integrand_name))
    }
}

pub mod process;
pub use process::*;
pub mod amplitude;
pub use amplitude::*;
pub mod cross_section;
pub use cross_section::*;

#[cfg(test)]
mod tests {
    use std::fs::OpenOptions;

    use linnet::half_edge::{
        involution::EdgeIndex,
        subgraph::{SuBitGraph, SubSetLike},
    };

    use symbolica::state::State;

    use crate::{
        GammaLoopContextContainer, dot,
        graph::{Graph, LoopMomentumBasis, parse::IntoGraph},
        momentum::signature::LoopExtSignature,
        settings::{
            RuntimeSettings,
            global::{
                CompilationMode, CompilationOptimizationLevel, GammaloopCompileOptions,
                GenerationSettings, TropicalSubgraphTableSettings,
            },
            runtime::LockedRuntimeSettings,
        },
        utils::load_generic_model,
    };

    use super::AmplitudeGraph;

    mod failing {
        use super::*;

        #[test]
        fn test_encode_decode_amplitude_graph() {
            // load the model and hack the masses, go through serializable model since arc is not mutable
            let model = load_generic_model("sm");

            let mut graph: Graph = dot!(
                digraph G{
                    e1      [flow=sink]
                    e2      [flow=source]
                    e3      [flow=source]
                    e1 -> n1  [particle=h]
                    e2 -> n4    [particle=h]
                    n1 -> n2    [particle=h]
                    n1 -> n3    [particle=h]
                    n2 -> n3    [particle=t]
                    n3 -> n4    [particle=t]
                    n4 -> n2    [particle=t]
                }
            )
            .unwrap();
            let loop_momentum_basis = LoopMomentumBasis {
                tree: SuBitGraph::empty(0),
                loop_edges: vec![EdgeIndex::from(0), EdgeIndex::from(4)].into(),
                ext_edges: vec![EdgeIndex::from(5), EdgeIndex::from(6)].into(),
                edge_signatures: graph
                    .underlying
                    .new_edgevec(|_, _, _| LoopExtSignature::from((vec![], vec![]))),
            };

            // loop_momentum_basis
            //     .set_edge_signatures(&graph.underlying)
            //     .unwrap();

            graph.loop_momentum_basis = loop_momentum_basis;

            let mut amplitude: AmplitudeGraph = AmplitudeGraph::new(graph.clone());

            amplitude
                .preprocess(
                    &model,
                    &GenerationSettings {
                        compile: GammaloopCompileOptions {
                            compilation_mode: CompilationMode::Cpp,
                            fast_math: false,
                            optimization_level: CompilationOptimizationLevel::O0,
                            unsafe_math: false,
                            compiler: crate::settings::global::default_external_compiler()
                                .to_owned(),
                            custom: Vec::new(),
                        },
                        tropical_subgraph_table: TropicalSubgraphTableSettings {
                            panic_on_fail: false,
                            target_omega: 1.0,
                            ..Default::default()
                        },
                        ..Default::default()
                    },
                    &LockedRuntimeSettings::from(&RuntimeSettings::default()),
                )
                .unwrap();

            let mut temp = OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(true)
                .open("test.bin")
                .unwrap();

            State::export(&mut temp).unwrap();
            drop(temp);

            let mut temp = OpenOptions::new().read(true).open("test.bin").unwrap();
            let state_map = State::import(&mut temp, None).unwrap();

            let context = GammaLoopContextContainer {
                model: &model,
                state_map: &state_map,
            };

            println!("context created");

            let encoded_amplitude =
                bincode::encode_to_vec(&amplitude, bincode::config::standard()).unwrap();

            let _amplitude: AmplitudeGraph = bincode::decode_from_slice_with_context(
                &encoded_amplitude,
                bincode::config::standard(),
                context,
            )
            .expect("amplitude decode failed")
            .0;

            println!("amplitude graph passed");
        }
    }
}