Skip to main content

gammaloop_api/commands/
inspect.rs

1use std::{fs, path::PathBuf};
2
3use clap::Args;
4use color_eyre::eyre::Context;
5use colored::Colorize;
6use ndarray::Array2;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use spenso::algebra::complex::Complex;
10use tracing::info;
11
12use crate::{
13    commands::evaluate_samples::{evaluate_sample, EvaluateSamples},
14    commands::CliArgumentMetadataExt,
15    completion::CompletionArgExt,
16    state::{ProcessRef, State},
17};
18use color_eyre::Result;
19use eyre::eyre;
20use gammalooprs::utils::F;
21
22#[derive(Debug, Args, Serialize, Deserialize, Clone, JsonSchema, PartialEq, Default)]
23pub struct Inspect {
24    /// Process reference: `#<id>`, `name:<name>`, or `<id>/<name>`
25    #[arg(
26        short = 'p',
27        long = "process",
28        value_name = "PROCESS",
29        completion_process_selector(crate::completion::SelectorKind::Any)
30    )]
31    pub process: Option<ProcessRef>,
32    /// The integrand name to inspect
33    #[arg(
34        short = 'i',
35        long = "integrand-name",
36        value_name = "NAME",
37        completion_integrand_selector(crate::completion::SelectorKind::Any)
38    )]
39    pub integrand_name: Option<String>,
40    /// Integration coordinates, or flattened loop momenta `(px, py, pz) ...` with `--momentum-space`
41    #[arg(
42        short = 'x',
43        long = "point",
44        num_args = 2..,
45        value_name = "POINT",
46        value_delimiter = ',',
47        allow_negative_numbers = true,
48    )]
49    pub point: Vec<f64>,
50
51    /// Force arbitrary-precision (Arb) internal evaluation; output remains f64
52    #[arg(short = 'f', long = "use_arb_prec")]
53    pub use_arb_prec: bool,
54
55    /// Interpret `--point` as spatial loop-momentum triplets; energy components are not accepted
56    #[arg(short = 'm', long)]
57    pub momentum_space: bool,
58
59    /// The discrete dimensions of the sample
60    #[arg(
61        short = 'd',
62        long = "discrete-dim",
63        value_name = "DIMS",
64        num_args = 1..,
65        value_delimiter = ',',
66        conflicts_with_all = ["graph_id", "orientation_id"],
67    )]
68    pub discrete_dim: Vec<usize>,
69
70    /// Select a specific graph in momentum-space inspect
71    #[arg(long = "graph-id", value_name = "GRAPH_ID")]
72    pub graph_id: Option<usize>,
73
74    /// Select a specific orientation of the selected graph in momentum-space inspect
75    #[arg(
76        long = "orientation-id",
77        value_name = "ORIENTATION_ID",
78        cli_requires("graph_id"),
79        conflicts_with = "discrete_dim"
80    )]
81    pub orientation_id: Option<usize>,
82
83    /// Write the inspection result as pretty JSON to this path.
84    #[arg(long = "json-output", value_name = "PATH")]
85    #[serde(default, skip_serializing_if = "Option::is_none")]
86    pub json_output: Option<PathBuf>,
87}
88
89impl Inspect {
90    fn validate_selector_mode(&self) -> Result<()> {
91        if !self.momentum_space && (self.graph_id.is_some() || self.orientation_id.is_some()) {
92            return Err(eyre!(
93                "Graph and orientation selectors are only supported in momentum-space inspect."
94            ));
95        }
96        Ok(())
97    }
98
99    fn resolve_graph_name(
100        &self,
101        state: &mut State,
102        process_id: usize,
103        integrand_name: &str,
104    ) -> Result<Option<String>> {
105        let Some(graph_id) = self.graph_id else {
106            return Ok(None);
107        };
108        let integrand = state
109            .process_list
110            .get_integrand_mut(process_id, integrand_name)?;
111        let graph_name = integrand.graph_name_by_id(graph_id).ok_or_else(|| {
112            eyre!(
113                "Graph id {} is out of range for integrand '{}'; it has {} graphs.",
114                graph_id,
115                integrand_name,
116                integrand.graph_count()
117            )
118        })?;
119        Ok(Some(graph_name.to_string()))
120    }
121
122    fn validate_x_space_point(
123        &self,
124        state: &mut State,
125        process_id: usize,
126        integrand_name: &str,
127    ) -> Result<()> {
128        let integrand = state
129            .process_list
130            .get_integrand_mut(process_id, integrand_name)?;
131        let expected_dimension = integrand.expected_x_space_dimension(&self.discrete_dim)?;
132        if self.point.len() != expected_dimension {
133            return Err(eyre!(
134                "Expected {} x-space coordinates for this integrand selection, got {}.",
135                expected_dimension,
136                self.point.len()
137            ));
138        }
139        Ok(())
140    }
141
142    pub fn run(&self, state: &mut State) -> Result<(Option<f64>, Complex<f64>)> {
143        let (process_id, integrand_name) =
144            state.find_integrand_ref(self.process.as_ref(), self.integrand_name.as_ref())?;
145        self.validate_selector_mode()?;
146        if !self.momentum_space {
147            self.validate_x_space_point(state, process_id, &integrand_name)?;
148        }
149        let graph_name = self.resolve_graph_name(state, process_id, &integrand_name)?;
150        let show_detailed_output = state
151            .process_list
152            .get_integrand_mut(process_id, &integrand_name)?
153            .get_settings()
154            .general
155            .generate_events;
156        let points = Array2::from_shape_vec((1, self.point.len()), self.point.clone())?;
157        let discrete_dims =
158            Array2::from_shape_vec((1, self.discrete_dim.len()), self.discrete_dim.clone())?;
159        let result = evaluate_sample(
160            state,
161            &EvaluateSamples {
162                process_id: Some(process_id),
163                integrand_name: Some(integrand_name.clone()),
164                use_arb_prec: self.use_arb_prec,
165                minimal_output: !show_detailed_output,
166                return_generated_events: None,
167                momentum_space: self.momentum_space,
168                points: points.view(),
169                integrator_weights: None,
170                discrete_dims: Some(discrete_dims.view()),
171                graph_names: graph_name.map(|name| vec![Some(name)]),
172                orientations: self
173                    .orientation_id
174                    .map(|orientation| vec![Some(orientation)]),
175            },
176        )?;
177        self.write_json_output(&result.sample)?;
178        let evaluation = &result.sample.evaluation;
179
180        let raw_result = evaluation.integrand_result;
181        let jacobian = evaluation.parameterization_jacobian.map(|jac| jac.0);
182        let displayed_result = if self.momentum_space {
183            raw_result
184        } else {
185            let jacobian = jacobian.ok_or_else(|| {
186                eyre!("x-space inspect requires a parameterization jacobian, but none was returned")
187            })?;
188            raw_result.map(|entry| entry * F(jacobian))
189        };
190
191        let point_label = if self.momentum_space {
192            "Input point in momentum space"
193        } else {
194            "Input point in unit hypercube xs"
195        };
196        info!(
197            "\n{}:\n\n{}\n\nThe evaluation of integrand '{}' is:\n\n{}\n",
198            point_label,
199            format!(
200                "( {} )",
201                self.point
202                    .iter()
203                    .map(|x| format!("{x:.16}"))
204                    .collect::<Vec<_>>()
205                    .join(", ")
206            )
207            .blue(),
208            integrand_name.green(),
209            format!(
210                "( {:+.16e}, {:+.16e} i)",
211                displayed_result.re, displayed_result.im
212            )
213            .blue(),
214        );
215
216        if let Some(jacobian) = jacobian {
217            info!(
218                "Parameterization jacobian for this point: {:+.16e}",
219                jacobian
220            );
221        }
222        if show_detailed_output {
223            info!("\n{}", result);
224        }
225
226        Ok((jacobian, displayed_result.map(|entry| entry.0)))
227    }
228
229    fn write_json_output<T: Serialize>(&self, value: &T) -> Result<()> {
230        let Some(path) = self.json_output.as_deref() else {
231            return Ok(());
232        };
233        let path = if path.is_absolute() {
234            path.to_path_buf()
235        } else {
236            std::env::current_dir()?.join(path)
237        };
238        if let Some(parent) = path.parent() {
239            fs::create_dir_all(parent)
240                .with_context(|| format!("Could not create {}", parent.display()))?;
241        }
242        fs::write(&path, serde_json::to_vec_pretty(value)?)
243            .with_context(|| format!("Could not write inspect JSON output {}", path.display()))?;
244        info!(
245            "Wrote inspect JSON output to {}",
246            path.display().to_string().green()
247        );
248        Ok(())
249    }
250}