Skip to main content

neopdf/
metadata.rs

1use std::collections::HashMap;
2
3use pyo3::prelude::*;
4
5use neopdf::metadata::{InterpolatorType, MetaData, SetType};
6
7/// The type of the set.
8#[pyclass(eq, eq_int, from_py_object, name = "SetType")]
9#[derive(Clone, PartialEq, Eq)]
10pub enum PySetType {
11    /// Parton Distribution Function.
12    SpaceLike,
13    /// Fragmentation Function.
14    TimeLike,
15}
16
17impl From<&SetType> for PySetType {
18    fn from(set_type: &SetType) -> Self {
19        match set_type {
20            SetType::SpaceLike => Self::SpaceLike,
21            SetType::TimeLike => Self::TimeLike,
22        }
23    }
24}
25
26impl From<&PySetType> for SetType {
27    fn from(set_type: &PySetType) -> Self {
28        match set_type {
29            PySetType::SpaceLike => Self::SpaceLike,
30            PySetType::TimeLike => Self::TimeLike,
31        }
32    }
33}
34
35/// The interpolation method used for the grid.
36#[pyclass(eq, eq_int, from_py_object, name = "InterpolatorType")]
37#[derive(Clone, PartialEq, Eq)]
38pub enum PyInterpolatorType {
39    /// Bilinear interpolation strategy.
40    Bilinear,
41    /// Bilinear logarithmic interpolation strategy.
42    LogBilinear,
43    /// Bicubic logarithmic interpolation strategy.
44    LogBicubic,
45    /// Tricubic logarithmic interpolation strategy.
46    LogTricubic,
47    /// Linear interpolation for N-dimensional data.
48    NDLinear,
49    /// Chebyshev logarithmic interpolation strategy.
50    LogChebyshev,
51    /// Four-dimensional cubic logarithmic interpolation strategy.
52    LogFourCubic,
53    /// Five-dimensional cubic logarithmic interpolation strategy.
54    LogFiveCubic,
55}
56
57impl From<&InterpolatorType> for PyInterpolatorType {
58    fn from(basis: &InterpolatorType) -> Self {
59        match basis {
60            InterpolatorType::Bilinear => Self::Bilinear,
61            InterpolatorType::LogBilinear => Self::LogBilinear,
62            InterpolatorType::LogBicubic => Self::LogBicubic,
63            InterpolatorType::LogTricubic => Self::LogTricubic,
64            InterpolatorType::InterpNDLinear => Self::NDLinear,
65            InterpolatorType::LogChebyshev => Self::LogChebyshev,
66            InterpolatorType::LogFourCubic => Self::LogFourCubic,
67            InterpolatorType::LogFiveCubic => Self::LogFiveCubic,
68        }
69    }
70}
71
72impl From<&PyInterpolatorType> for InterpolatorType {
73    fn from(basis: &PyInterpolatorType) -> Self {
74        match basis {
75            PyInterpolatorType::Bilinear => Self::Bilinear,
76            PyInterpolatorType::LogBilinear => Self::LogBilinear,
77            PyInterpolatorType::LogBicubic => Self::LogBicubic,
78            PyInterpolatorType::LogTricubic => Self::LogTricubic,
79            PyInterpolatorType::NDLinear => Self::InterpNDLinear,
80            PyInterpolatorType::LogChebyshev => Self::LogChebyshev,
81            PyInterpolatorType::LogFourCubic => Self::LogFourCubic,
82            PyInterpolatorType::LogFiveCubic => Self::LogFiveCubic,
83        }
84    }
85}
86
87/// Physical Parameters of the PDF set.
88#[pyclass(from_py_object, name = "PhysicsParameters")]
89#[derive(Debug, Clone)]
90pub struct PyPhysicsParameters {
91    pub(crate) flavor_scheme: String,
92    pub(crate) order_qcd: u32,
93    pub(crate) alphas_order_qcd: u32,
94    pub(crate) m_w: f64,
95    pub(crate) m_z: f64,
96    pub(crate) m_up: f64,
97    pub(crate) m_down: f64,
98    pub(crate) m_strange: f64,
99    pub(crate) m_charm: f64,
100    pub(crate) m_bottom: f64,
101    pub(crate) m_top: f64,
102    pub(crate) alphas_type: String,
103    pub(crate) number_flavors: u32,
104}
105
106#[pymethods]
107impl PyPhysicsParameters {
108    /// Constructor for `PyPhysicsParameters`.
109    #[new]
110    #[must_use]
111    #[allow(clippy::too_many_arguments)]
112    #[pyo3(signature = (
113        flavor_scheme = "None".to_string(),
114        order_qcd = 0,
115        alphas_order_qcd = 0,
116        m_w = 0.0,
117        m_z = 0.0,
118        m_up = 0.0,
119        m_down = 0.0,
120        m_strange = 0.0,
121        m_charm = 0.0,
122        m_bottom = 0.0,
123        m_top = 0.0,
124        alphas_type = "None".to_string(),
125        number_flavors = 0,
126    ))]
127    pub const fn new(
128        flavor_scheme: String,
129        order_qcd: u32,
130        alphas_order_qcd: u32,
131        m_w: f64,
132        m_z: f64,
133        m_up: f64,
134        m_down: f64,
135        m_strange: f64,
136        m_charm: f64,
137        m_bottom: f64,
138        m_top: f64,
139        alphas_type: String,
140        number_flavors: u32,
141    ) -> Self {
142        Self {
143            flavor_scheme,
144            order_qcd,
145            alphas_order_qcd,
146            m_w,
147            m_z,
148            m_up,
149            m_down,
150            m_strange,
151            m_charm,
152            m_bottom,
153            m_top,
154            alphas_type,
155            number_flavors,
156        }
157    }
158
159    /// Convert to Python dictionary.
160    ///
161    /// # Errors
162    ///
163    /// Raises an error if the values are not Python compatible.
164    pub fn to_dict(&self, py: Python) -> PyResult<Py<PyAny>> {
165        let dict = pyo3::types::PyDict::new(py);
166        dict.set_item("flavor_scheme", &self.flavor_scheme)?;
167        dict.set_item("order_qcd", self.order_qcd)?;
168        dict.set_item("alphas_order_qcd", self.alphas_order_qcd)?;
169        dict.set_item("m_w", self.m_w)?;
170        dict.set_item("m_z", self.m_z)?;
171        dict.set_item("m_up", self.m_up)?;
172        dict.set_item("m_down", self.m_down)?;
173        dict.set_item("m_strange", self.m_strange)?;
174        dict.set_item("m_charm", self.m_charm)?;
175        dict.set_item("m_bottom", self.m_bottom)?;
176        dict.set_item("m_top", self.m_top)?;
177
178        Ok(dict.into())
179    }
180}
181
182impl Default for PyPhysicsParameters {
183    fn default() -> Self {
184        Self {
185            flavor_scheme: String::new(),
186            order_qcd: 0,
187            alphas_order_qcd: 0,
188            m_w: 0.0,
189            m_z: 0.0,
190            m_up: 0.0,
191            m_down: 0.0,
192            m_strange: 0.0,
193            m_charm: 0.0,
194            m_bottom: 0.0,
195            m_top: 0.0,
196            alphas_type: String::new(),
197            number_flavors: 0,
198        }
199    }
200}
201
202fn format_list<T: ToString>(v: &[T]) -> String {
203    format!(
204        "[{}]",
205        v.iter()
206            .map(std::string::ToString::to_string)
207            .collect::<Vec<_>>()
208            .join(", ")
209    )
210}
211
212fn insert_mandatory_fields(map: &mut HashMap<String, String>, meta: &MetaData) {
213    let set_type = match &meta.set_type {
214        SetType::SpaceLike => "PDF",
215        SetType::TimeLike => "FragFn",
216    };
217    let interp_type = match &meta.interpolator_type {
218        InterpolatorType::Bilinear => "Bilinear",
219        InterpolatorType::LogBilinear => "LogBilinear",
220        InterpolatorType::LogBicubic => "LogBicubic",
221        InterpolatorType::LogTricubic => "LogTricubic",
222        InterpolatorType::InterpNDLinear => "NDLinear",
223        InterpolatorType::LogChebyshev => "LogChebyshev",
224        InterpolatorType::LogFourCubic => "LogFourCubic",
225        InterpolatorType::LogFiveCubic => "LogFiveCubic",
226    };
227    map.insert("SetDesc".to_string(), meta.set_desc.clone());
228    map.insert("SetIndex".to_string(), meta.set_index.to_string());
229    map.insert("NumMembers".to_string(), meta.num_members.to_string());
230    map.insert("XMin".to_string(), meta.x_min.to_string());
231    map.insert("XMax".to_string(), meta.x_max.to_string());
232    map.insert("QMin".to_string(), meta.q_min.to_string());
233    map.insert("QMax".to_string(), meta.q_max.to_string());
234    map.insert("Flavors".to_string(), format_list(&meta.flavors));
235    map.insert("Format".to_string(), meta.format.clone());
236    map.insert("SetType".to_string(), set_type.to_string());
237    map.insert("InterpolatorType".to_string(), interp_type.to_string());
238    map.insert("Polarized".to_string(), meta.polarised.to_string());
239    map.insert("Particle".to_string(), meta.hadron_pid.to_string());
240    map.insert("OrderQCD".to_string(), meta.order_qcd.to_string());
241}
242
243fn insert_optional_fields(map: &mut HashMap<String, String>, meta: &MetaData) {
244    for (key, val) in [
245        ("ErrorType", meta.error_type.as_str()),
246        ("GitVersion", meta.git_version.as_str()),
247        ("CodeVersion", meta.code_version.as_str()),
248        ("FlavorScheme", meta.flavor_scheme.as_str()),
249        ("AlphaS_Type", meta.alphas_type.as_str()),
250    ] {
251        if !val.is_empty() {
252            map.insert(key.to_string(), val.to_string());
253        }
254    }
255    for (key, vals) in [
256        ("AlphaS_Qs", &meta.alphas_q_values),
257        ("AlphaS_Vals", &meta.alphas_vals),
258    ] {
259        if !vals.is_empty() {
260            map.insert(key.to_string(), format_list(vals));
261        }
262    }
263    for (key, val) in [
264        ("AlphaS_OrderQCD", meta.alphas_order_qcd),
265        ("NumFlavors", meta.number_flavors),
266    ] {
267        if val != 0 {
268            map.insert(key.to_string(), val.to_string());
269        }
270    }
271}
272
273fn insert_mass_fields(map: &mut HashMap<String, String>, meta: &MetaData) {
274    for (key, val) in [
275        ("MW", meta.m_w),
276        ("MZ", meta.m_z),
277        ("MUp", meta.m_up),
278        ("MDown", meta.m_down),
279        ("MStrange", meta.m_strange),
280        ("MCharm", meta.m_charm),
281        ("MBottom", meta.m_bottom),
282        ("MTop", meta.m_top),
283    ] {
284        if val != 0.0 {
285            map.insert(key.to_string(), val.to_string());
286        }
287    }
288}
289
290fn insert_tmd_fields(map: &mut HashMap<String, String>, meta: &MetaData) {
291    for (key, val) in [
292        ("XiMin", meta.xi_min),
293        ("XiMax", meta.xi_max),
294        ("DeltaMin", meta.delta_min),
295        ("DeltaMax", meta.delta_max),
296    ] {
297        if val != 0.0 {
298            map.insert(key.to_string(), val.to_string());
299        }
300    }
301    if let Some(ecl) = meta.error_conf_level {
302        map.insert("ErrorConfLevel".to_string(), ecl.to_string());
303    }
304}
305
306/// Build a map of LHAPDF-canonical key names to their string values for a `MetaData` object.
307///
308/// Keys that are absent from the original `.info` file (represented by empty strings,
309/// zero values, or `None`) are not included, mirroring LHAPDF's `has_key` semantics.
310pub(crate) fn build_lhapdf_map(meta: &MetaData) -> HashMap<String, String> {
311    let mut map = HashMap::new();
312    insert_mandatory_fields(&mut map, meta);
313    insert_optional_fields(&mut map, meta);
314    insert_mass_fields(&mut map, meta);
315    insert_tmd_fields(&mut map, meta);
316    map
317}
318
319/// Grid metadata.
320#[pyclass(from_py_object, name = "MetaData")]
321#[derive(Debug, Clone)]
322#[repr(transparent)]
323pub struct PyMetaData {
324    pub(crate) meta: MetaData,
325}
326
327#[pymethods]
328impl PyMetaData {
329    /// Constructor for `PyMetaData`.
330    #[new]
331    #[must_use]
332    #[allow(clippy::too_many_arguments)]
333    #[allow(clippy::needless_pass_by_value)]
334    #[pyo3(signature = (
335        set_desc,
336        set_index,
337        num_members,
338        x_min,
339        x_max,
340        q_min,
341        q_max,
342        xsi_min,
343        xsi_max,
344        delta_min,
345        delta_max,
346        flavors,
347        format,
348        alphas_q_values = vec![],
349        alphas_vals = vec![],
350        polarised = false,
351        set_type = PySetType::SpaceLike,
352        interpolator_type = PyInterpolatorType::LogBicubic,
353        error_type = "replicas".to_string(),
354        hadron_pid = 2212,
355        phys_params = PyPhysicsParameters::default(),
356    ))]
357    pub fn new(
358        set_desc: String,
359        set_index: u32,
360        num_members: u32,
361        x_min: f64,
362        x_max: f64,
363        q_min: f64,
364        q_max: f64,
365        xsi_min: f64,
366        xsi_max: f64,
367        delta_min: f64,
368        delta_max: f64,
369        flavors: Vec<i32>,
370        format: String,
371        alphas_q_values: Vec<f64>,
372        alphas_vals: Vec<f64>,
373        polarised: bool,
374        set_type: PySetType,
375        interpolator_type: PyInterpolatorType,
376        error_type: String,
377        hadron_pid: i32,
378        phys_params: PyPhysicsParameters,
379    ) -> Self {
380        Self {
381            meta: MetaData {
382                set_desc,
383                set_index,
384                num_members,
385                x_min,
386                x_max,
387                q_min,
388                q_max,
389                flavors,
390                format,
391                alphas_q_values,
392                alphas_vals,
393                polarised,
394                set_type: SetType::from(&set_type),
395                interpolator_type: InterpolatorType::from(&interpolator_type),
396                error_type,
397                hadron_pid,
398                git_version: String::new(), // placeholder to be overwritten
399                code_version: String::new(), // placeholder to be overwritten
400                flavor_scheme: phys_params.flavor_scheme,
401                order_qcd: phys_params.order_qcd,
402                alphas_order_qcd: phys_params.alphas_order_qcd,
403                m_w: phys_params.m_w,
404                m_z: phys_params.m_z,
405                m_up: phys_params.m_up,
406                m_down: phys_params.m_down,
407                m_strange: phys_params.m_strange,
408                m_charm: phys_params.m_charm,
409                m_bottom: phys_params.m_bottom,
410                m_top: phys_params.m_top,
411                alphas_type: phys_params.alphas_type,
412                number_flavors: phys_params.number_flavors,
413                // New V2 fields with defaults
414                xi_min: xsi_min,
415                xi_max: xsi_max,
416                delta_min,
417                delta_max,
418                error_conf_level: None,
419            },
420        }
421    }
422
423    /// Convert to Python dictionary
424    ///
425    /// # Errors
426    ///
427    /// Raises an erro if the values are not Python compatible.
428    pub fn to_dict(&self, py: Python) -> PyResult<Py<PyAny>> {
429        let dict = pyo3::types::PyDict::new(py);
430
431        let set_type = match &self.meta.set_type {
432            SetType::SpaceLike => "PDF",
433            SetType::TimeLike => "FragFn",
434        };
435
436        let interpolator_type = match &self.meta.interpolator_type {
437            InterpolatorType::Bilinear => "Bilinear",
438            InterpolatorType::LogBilinear => "LogBilinear",
439            InterpolatorType::LogBicubic => "LogBicubic",
440            InterpolatorType::LogTricubic => "LogTricubic",
441            InterpolatorType::InterpNDLinear => "NDLinear",
442            InterpolatorType::LogChebyshev => "LogChebyshev",
443            InterpolatorType::LogFourCubic => "LogFourCubic",
444            InterpolatorType::LogFiveCubic => "LogFiveCubic",
445        };
446
447        dict.set_item("set_desc", &self.meta.set_desc)?;
448        dict.set_item("set_index", self.meta.set_index)?;
449        dict.set_item("num_members", self.meta.num_members)?;
450        dict.set_item("x_min", self.meta.x_min)?;
451        dict.set_item("x_max", self.meta.x_max)?;
452        dict.set_item("q_min", self.meta.q_min)?;
453        dict.set_item("q_max", self.meta.q_max)?;
454        dict.set_item("flavors", &self.meta.flavors)?;
455        dict.set_item("format", &self.meta.format)?;
456        dict.set_item("alphas_q_values", &self.meta.alphas_q_values)?;
457        dict.set_item("alphas_vals", &self.meta.alphas_vals)?;
458        dict.set_item("polarised", self.meta.polarised)?;
459        dict.set_item("set_type", set_type)?;
460        dict.set_item("interpolator_type", interpolator_type)?;
461        dict.set_item("error_type", &self.meta.error_type)?;
462        dict.set_item("hadron_pid", self.meta.hadron_pid)?;
463        dict.set_item("git_version", &self.meta.git_version)?;
464        dict.set_item("code_version", &self.meta.code_version)?;
465        dict.set_item("flavor_scheme", &self.meta.flavor_scheme)?;
466        dict.set_item("order_qcd", self.meta.order_qcd)?;
467        dict.set_item("alphas_order_qcd", self.meta.alphas_order_qcd)?;
468        dict.set_item("m_w", self.meta.m_w)?;
469        dict.set_item("m_z", self.meta.m_z)?;
470        dict.set_item("m_up", self.meta.m_up)?;
471        dict.set_item("m_down", self.meta.m_down)?;
472        dict.set_item("m_strange", self.meta.m_strange)?;
473        dict.set_item("m_charm", self.meta.m_charm)?;
474        dict.set_item("m_bottom", self.meta.m_bottom)?;
475        dict.set_item("m_top", self.meta.m_top)?;
476
477        Ok(dict.into())
478    }
479
480    /// Return `True` if the given LHAPDF-canonical key is present in the metadata.
481    #[must_use]
482    #[pyo3(name = "has_key")]
483    pub fn has_key(&self, key: &str) -> bool {
484        build_lhapdf_map(&self.meta).contains_key(key)
485    }
486
487    /// Return the string value for an LHAPDF-canonical metadata key.
488    ///
489    /// # Errors
490    ///
491    /// Raises `KeyError` if the key is not present.
492    #[pyo3(name = "get_entry")]
493    pub fn get_entry(&self, key: &str) -> PyResult<String> {
494        build_lhapdf_map(&self.meta).remove(key).ok_or_else(|| {
495            pyo3::exceptions::PyKeyError::new_err(format!("Key '{key}' not found in metadata"))
496        })
497    }
498
499    /// Return a sorted list of all available LHAPDF-canonical metadata keys.
500    #[must_use]
501    #[pyo3(name = "keys")]
502    pub fn keys(&self) -> Vec<String> {
503        let mut keys: Vec<String> = build_lhapdf_map(&self.meta).into_keys().collect();
504        keys.sort();
505        keys
506    }
507
508    /// The description of the set.
509    #[must_use]
510    pub const fn set_desc(&self) -> &String {
511        &self.meta.set_desc
512    }
513
514    /// The index of the grid.
515    #[must_use]
516    pub const fn set_index(&self) -> u32 {
517        self.meta.set_index
518    }
519
520    /// The number of sets in the grid.
521    #[must_use]
522    pub const fn number_sets(&self) -> u32 {
523        self.meta.num_members
524    }
525
526    /// The minimum value of `x` in the grid.
527    #[must_use]
528    pub const fn x_min(&self) -> f64 {
529        self.meta.x_min
530    }
531
532    /// The maximum value of `x` in the grid.
533    #[must_use]
534    pub const fn x_max(&self) -> f64 {
535        self.meta.x_max
536    }
537
538    /// The minimum value of `q` in the grid.
539    #[must_use]
540    pub const fn q_min(&self) -> f64 {
541        self.meta.q_min
542    }
543
544    /// The maximum value of `q` in the grid.
545    #[must_use]
546    pub const fn q_max(&self) -> f64 {
547        self.meta.q_max
548    }
549
550    /// The minimum value of `xi` in the grid.
551    #[must_use]
552    pub const fn xi_min(&self) -> f64 {
553        self.meta.xi_min
554    }
555
556    /// The maximum value of `xi` in the grid.
557    #[must_use]
558    pub const fn xi_max(&self) -> f64 {
559        self.meta.xi_max
560    }
561
562    /// The minimum value of `delta` in the grid.
563    #[must_use]
564    pub const fn delta_min(&self) -> f64 {
565        self.meta.delta_min
566    }
567
568    /// The maximum value of `delta` in the grid.
569    #[must_use]
570    pub const fn delta_max(&self) -> f64 {
571        self.meta.delta_max
572    }
573
574    /// The particle IDs of the grid.
575    #[must_use]
576    pub const fn pids(&self) -> &Vec<i32> {
577        &self.meta.flavors
578    }
579
580    /// The format of the grid.
581    #[must_use]
582    pub const fn format(&self) -> &String {
583        &self.meta.format
584    }
585
586    /// The values of `q` for the running of the strong coupling constant.
587    #[must_use]
588    pub const fn alphas_q(&self) -> &Vec<f64> {
589        &self.meta.alphas_q_values
590    }
591
592    /// The values of the running of the strong coupling constant.
593    #[must_use]
594    pub const fn alphas_values(&self) -> &Vec<f64> {
595        &self.meta.alphas_vals
596    }
597
598    /// Whether the grid is polarised.
599    #[must_use]
600    pub const fn is_polarised(&self) -> bool {
601        self.meta.polarised
602    }
603
604    /// The type of the set.
605    #[must_use]
606    pub fn set_type(&self) -> PySetType {
607        PySetType::from(&self.meta.set_type)
608    }
609
610    /// The interpolation method used for the grid.
611    #[must_use]
612    pub fn interpolator_type(&self) -> PyInterpolatorType {
613        PyInterpolatorType::from(&self.meta.interpolator_type)
614    }
615
616    /// The type of error.
617    #[must_use]
618    pub const fn error_type(&self) -> &String {
619        &self.meta.error_type
620    }
621
622    /// The hadron PID.
623    #[must_use]
624    pub const fn hadron_pid(&self) -> i32 {
625        self.meta.hadron_pid
626    }
627}
628
629/// Registers the `metadata` submodule with the parent Python module.
630///
631/// Parameters
632/// ----------
633/// `parent_module` : pyo3.Bound[pyo3.types.PyModule]
634///     The parent Python module to which the `metadata` submodule will be added.
635///
636/// Returns
637/// -------
638/// pyo3.PyResult<()>
639///     `Ok(())` if the registration is successful, or an error if the submodule
640///     cannot be created or added.
641///
642/// # Errors
643///
644/// Raises an error if the (sub)module is not found or cannot be registered.
645pub fn register(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
646    let m = PyModule::new(parent_module.py(), "metadata")?;
647    m.setattr(pyo3::intern!(m.py(), "__doc__"), "Interface for PDF.")?;
648    pyo3::py_run!(
649        parent_module.py(),
650        m,
651        "import sys; sys.modules['neopdf.metadata'] = m"
652    );
653    m.add_class::<PySetType>()?;
654    m.add_class::<PyInterpolatorType>()?;
655    m.add_class::<PyPhysicsParameters>()?;
656    m.add_class::<PyMetaData>()?;
657    parent_module.add_submodule(&m)
658}