Skip to main content

laddu_core/variables/
selection.rs

1use serde::{Deserialize, Serialize};
2
3use crate::{data::DatasetMetadata, LadduError, LadduResult};
4
5fn names_to_string(names: &[String]) -> String {
6    names.join(", ")
7}
8
9/// A reusable selection that may span one or more four-momentum names.
10///
11/// Instances are constructed from metadata-facing identifiers and later bound to
12/// column indices so that variable evaluators can resolve aliases or grouped
13/// particles efficiently.
14#[derive(Clone, Debug, Serialize, Deserialize)]
15pub struct P4Selection {
16    names: Vec<String>,
17    #[serde(skip, default)]
18    indices: Vec<usize>,
19}
20
21impl P4Selection {
22    pub(crate) fn new_many<I, S>(names: I) -> Self
23    where
24        I: IntoIterator<Item = S>,
25        S: Into<String>,
26    {
27        Self {
28            names: names.into_iter().map(Into::into).collect(),
29            indices: Vec::new(),
30        }
31    }
32
33    pub(crate) fn with_indices<I, S>(names: I, indices: Vec<usize>) -> Self
34    where
35        I: IntoIterator<Item = S>,
36        S: Into<String>,
37    {
38        Self {
39            names: names.into_iter().map(Into::into).collect(),
40            indices,
41        }
42    }
43
44    /// Returns the metadata names contributing to this selection.
45    pub fn names(&self) -> &[String] {
46        &self.names
47    }
48
49    pub(crate) fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
50        let mut resolved = Vec::with_capacity(self.names.len());
51        for name in &self.names {
52            metadata.append_indices_for_name(name, &mut resolved)?;
53        }
54        self.indices = resolved;
55        Ok(())
56    }
57
58    /// The resolved column indices backing this selection.
59    pub fn indices(&self) -> &[usize] {
60        &self.indices
61    }
62}
63
64/// Helper trait to convert common particle specifications into [`P4Selection`] instances.
65pub trait IntoP4Selection {
66    /// Convert the input into a [`P4Selection`].
67    fn into_selection(self) -> P4Selection;
68}
69
70impl IntoP4Selection for P4Selection {
71    fn into_selection(self) -> P4Selection {
72        self
73    }
74}
75
76impl IntoP4Selection for &P4Selection {
77    fn into_selection(self) -> P4Selection {
78        self.clone()
79    }
80}
81
82impl IntoP4Selection for String {
83    fn into_selection(self) -> P4Selection {
84        P4Selection::new_many(vec![self])
85    }
86}
87
88impl IntoP4Selection for &String {
89    fn into_selection(self) -> P4Selection {
90        P4Selection::new_many(vec![self.clone()])
91    }
92}
93
94impl IntoP4Selection for &str {
95    fn into_selection(self) -> P4Selection {
96        P4Selection::new_many(vec![self.to_string()])
97    }
98}
99
100impl<S> IntoP4Selection for Vec<S>
101where
102    S: Into<String>,
103{
104    fn into_selection(self) -> P4Selection {
105        P4Selection::new_many(self.into_iter().map(Into::into).collect::<Vec<_>>())
106    }
107}
108
109impl<S> IntoP4Selection for &[S]
110where
111    S: Clone + Into<String>,
112{
113    fn into_selection(self) -> P4Selection {
114        P4Selection::new_many(self.iter().cloned().map(Into::into).collect::<Vec<_>>())
115    }
116}
117
118impl<S, const N: usize> IntoP4Selection for [S; N]
119where
120    S: Into<String>,
121{
122    fn into_selection(self) -> P4Selection {
123        P4Selection::new_many(self.into_iter().map(Into::into).collect::<Vec<_>>())
124    }
125}
126
127impl<S, const N: usize> IntoP4Selection for &[S; N]
128where
129    S: Clone + Into<String>,
130{
131    fn into_selection(self) -> P4Selection {
132        P4Selection::new_many(self.iter().cloned().map(Into::into).collect::<Vec<_>>())
133    }
134}
135
136#[derive(Clone, Debug, Serialize, Deserialize)]
137pub(crate) struct AuxSelection {
138    name: String,
139    #[serde(skip, default)]
140    index: Option<usize>,
141}
142
143impl AuxSelection {
144    pub(crate) fn new<S: Into<String>>(name: S) -> Self {
145        Self {
146            name: name.into(),
147            index: None,
148        }
149    }
150
151    pub(crate) fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
152        let idx = metadata
153            .aux_index(&self.name)
154            .ok_or_else(|| LadduError::UnknownName {
155                category: "aux",
156                name: self.name.clone(),
157            })?;
158        self.index = Some(idx);
159        Ok(())
160    }
161
162    pub(crate) fn index(&self) -> usize {
163        self.index.expect("AuxSelection must be bound before use")
164    }
165
166    pub(crate) fn name(&self) -> &str {
167        &self.name
168    }
169}
170
171pub(crate) fn format_names(names: &[String]) -> String {
172    names_to_string(names)
173}