Skip to main content

laddu_core/data/
metadata.rs

1use indexmap::IndexMap;
2
3use crate::{variables::P4Selection, LadduError, LadduResult};
4
5/// A collection of [`EventData`](crate::data::EventData).
6#[derive(Debug, Clone)]
7pub struct DatasetMetadata {
8    pub(crate) p4_names: Vec<String>,
9    pub(crate) aux_names: Vec<String>,
10    pub(crate) p4_lookup: IndexMap<String, usize>,
11    pub(crate) aux_lookup: IndexMap<String, usize>,
12    pub(crate) p4_selections: IndexMap<String, P4Selection>,
13}
14
15impl DatasetMetadata {
16    /// Construct metadata from explicit particle and auxiliary names.
17    pub fn new<P: Into<String>, A: Into<String>>(
18        p4_names: Vec<P>,
19        aux_names: Vec<A>,
20    ) -> LadduResult<Self> {
21        let mut p4_lookup = IndexMap::with_capacity(p4_names.len());
22        let mut aux_lookup = IndexMap::with_capacity(aux_names.len());
23        let mut p4_selections = IndexMap::with_capacity(p4_names.len());
24        let p4_names: Vec<String> = p4_names
25            .into_iter()
26            .enumerate()
27            .map(|(idx, name)| {
28                let name = name.into();
29                if p4_lookup.contains_key(&name) {
30                    return Err(LadduError::DuplicateName {
31                        category: "p4",
32                        name,
33                    });
34                }
35                p4_lookup.insert(name.clone(), idx);
36                p4_selections.insert(
37                    name.clone(),
38                    P4Selection::with_indices(vec![name.clone()], vec![idx]),
39                );
40                Ok(name)
41            })
42            .collect::<Result<_, _>>()?;
43        let aux_names: Vec<String> = aux_names
44            .into_iter()
45            .enumerate()
46            .map(|(idx, name)| {
47                let name = name.into();
48                if aux_lookup.contains_key(&name) {
49                    return Err(LadduError::DuplicateName {
50                        category: "aux",
51                        name,
52                    });
53                }
54                aux_lookup.insert(name.clone(), idx);
55                Ok(name)
56            })
57            .collect::<Result<_, _>>()?;
58        Ok(Self {
59            p4_names,
60            aux_names,
61            p4_lookup,
62            aux_lookup,
63            p4_selections,
64        })
65    }
66
67    /// Create metadata with no registered names.
68    pub fn empty() -> Self {
69        Self {
70            p4_names: Vec::new(),
71            aux_names: Vec::new(),
72            p4_lookup: IndexMap::new(),
73            aux_lookup: IndexMap::new(),
74            p4_selections: IndexMap::new(),
75        }
76    }
77
78    /// Resolve the index of a four-momentum by name.
79    pub fn p4_index(&self, name: &str) -> Option<usize> {
80        self.p4_lookup.get(name).copied()
81    }
82
83    /// Registered four-momentum names in declaration order.
84    pub fn p4_names(&self) -> &[String] {
85        &self.p4_names
86    }
87
88    /// Resolve the index of an auxiliary scalar by name.
89    pub fn aux_index(&self, name: &str) -> Option<usize> {
90        self.aux_lookup.get(name).copied()
91    }
92
93    /// Registered auxiliary scalar names in declaration order.
94    pub fn aux_names(&self) -> &[String] {
95        &self.aux_names
96    }
97
98    pub(crate) fn ensure_new_p4_name(&self, name: &str) -> LadduResult<()> {
99        if self.p4_selections.contains_key(name) {
100            return Err(LadduError::DuplicateName {
101                category: "p4",
102                name: name.to_string(),
103            });
104        }
105        Ok(())
106    }
107
108    pub(crate) fn ensure_new_aux_name(&self, name: &str) -> LadduResult<()> {
109        if self.aux_lookup.contains_key(name) {
110            return Err(LadduError::DuplicateName {
111                category: "aux",
112                name: name.to_string(),
113            });
114        }
115        Ok(())
116    }
117
118    pub(crate) fn add_p4_name<N>(&mut self, name: N) -> LadduResult<()>
119    where
120        N: Into<String>,
121    {
122        let name = name.into();
123        self.ensure_new_p4_name(&name)?;
124        let index = self.p4_names.len();
125        self.p4_lookup.insert(name.clone(), index);
126        self.p4_selections.insert(
127            name.clone(),
128            P4Selection::with_indices(vec![name.clone()], vec![index]),
129        );
130        self.p4_names.push(name);
131        Ok(())
132    }
133
134    pub(crate) fn add_aux_name<N>(&mut self, name: N) -> LadduResult<()>
135    where
136        N: Into<String>,
137    {
138        let name = name.into();
139        self.ensure_new_aux_name(&name)?;
140        let index = self.aux_names.len();
141        self.aux_lookup.insert(name.clone(), index);
142        self.aux_names.push(name);
143        Ok(())
144    }
145
146    /// Look up a resolved four-momentum selection by name (canonical or alias).
147    pub fn p4_selection(&self, name: &str) -> Option<&P4Selection> {
148        self.p4_selections.get(name)
149    }
150
151    /// Register an alias mapping to one or more existing four-momenta.
152    pub fn add_p4_alias<N>(&mut self, alias: N, mut selection: P4Selection) -> LadduResult<()>
153    where
154        N: Into<String>,
155    {
156        let alias = alias.into();
157        if self.p4_selections.contains_key(&alias) {
158            return Err(LadduError::DuplicateName {
159                category: "alias",
160                name: alias,
161            });
162        }
163        selection.bind(self)?;
164        self.p4_selections.insert(alias, selection);
165        Ok(())
166    }
167
168    /// Register multiple aliases at once.
169    pub fn add_p4_aliases<I, N>(&mut self, entries: I) -> LadduResult<()>
170    where
171        I: IntoIterator<Item = (N, P4Selection)>,
172        N: Into<String>,
173    {
174        for (alias, selection) in entries {
175            self.add_p4_alias(alias, selection)?;
176        }
177        Ok(())
178    }
179
180    pub(crate) fn append_indices_for_name(
181        &self,
182        name: &str,
183        target: &mut Vec<usize>,
184    ) -> LadduResult<()> {
185        if let Some(selection) = self.p4_selections.get(name) {
186            target.extend_from_slice(selection.indices());
187            return Ok(());
188        }
189        Err(LadduError::UnknownName {
190            category: "p4",
191            name: name.to_string(),
192        })
193    }
194}
195
196impl Default for DatasetMetadata {
197    fn default() -> Self {
198        Self::empty()
199    }
200}