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 #[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 #[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 #[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 #[arg(short = 'f', long = "use_arb_prec")]
53 pub use_arb_prec: bool,
54
55 #[arg(short = 'm', long)]
57 pub momentum_space: bool,
58
59 #[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 #[arg(long = "graph-id", value_name = "GRAPH_ID")]
72 pub graph_id: Option<usize>,
73
74 #[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 #[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}