laddu_python/

data.rs

use crate::utils::variables::{PyVariable, PyVariableExpression};
use laddu_core::{
    data::{
        io::{
            infer_p4_and_aux_names_from_columns, resolve_columns_case_insensitive,
            resolve_optional_weight_column, resolve_p4_component_columns, P4_COMPONENT_SUFFIXES,
        },
        read_parquet as core_read_parquet,
        read_parquet_chunks_with_options as core_read_parquet_chunks_with_options,
        read_root as core_read_root, write_parquet as core_write_parquet,
        write_root as core_write_root, BinnedDataset, Dataset, DatasetArcIter, DatasetMetadata,
        DatasetWriteOptions, Event, EventData, FloatPrecision, SharedDatasetIterExt,
    },
    utils::variables::IntoP4Selection,
    DatasetReadOptions,
};
use numpy::{PyArray1, PyReadonlyArray1};
use pyo3::{
    exceptions::{PyIndexError, PyKeyError, PyTypeError, PyValueError},
    prelude::*,
    types::{PyDict, PyList},
    IntoPyObjectExt,
};
use std::{path::PathBuf, sync::Arc};

use crate::utils::vectors::PyVec4;

fn parse_aliases(aliases: Option<Bound<'_, PyDict>>) -> PyResult<Vec<(String, Vec<String>)>> {
    let Some(aliases) = aliases else {
        return Ok(Vec::new());
    };

    let mut parsed = Vec::new();
    for (key, value) in aliases.iter() {
        let alias_name = key.extract::<String>()?;
        let selection = if let Ok(single) = value.extract::<String>() {
            vec![single]
        } else {
            let seq = value.extract::<Vec<String>>().map_err(|_| {
                PyTypeError::new_err("Alias values must be a string or a sequence of strings")
            })?;
            if seq.is_empty() {
                return Err(PyValueError::new_err(format!(
                    "Alias '{alias_name}' must reference at least one particle",
                )));
            }
            seq
        };
        parsed.push((alias_name, selection));
    }

    Ok(parsed)
}

fn parse_dataset_path(path: Bound<'_, PyAny>) -> PyResult<String> {
    if let Ok(s) = path.extract::<String>() {
        Ok(s)
    } else if let Ok(pathbuf) = path.extract::<PathBuf>() {
        Ok(pathbuf.to_string_lossy().into_owned())
    } else {
        Err(PyTypeError::new_err("Expected str or Path"))
    }
}

fn parse_precision_arg(value: Option<&str>) -> PyResult<FloatPrecision> {
    match value.map(|v| v.to_ascii_lowercase()) {
        None => Ok(FloatPrecision::F64),
        Some(name) if name == "f64" || name == "float64" || name == "double" => {
            Ok(FloatPrecision::F64)
        }
        Some(name) if name == "f32" || name == "float32" || name == "float" => {
            Ok(FloatPrecision::F32)
        }
        Some(other) => Err(PyValueError::new_err(format!(
            "Unsupported precision '{other}' (expected 'f64' or 'f32')"
        ))),
    }
}

fn extract_numeric_column(value: Bound<'_, PyAny>, name: &str) -> PyResult<Vec<f64>> {
    if let Ok(array) = value.extract::<PyReadonlyArray1<'_, f64>>() {
        return Ok(array.as_slice()?.to_vec());
    }
    if let Ok(array) = value.extract::<PyReadonlyArray1<'_, f32>>() {
        return Ok(array.as_slice()?.iter().map(|v| *v as f64).collect());
    }
    if let Ok(values) = value.extract::<Vec<f64>>() {
        return Ok(values);
    }
    if let Ok(values) = value.extract::<Vec<f32>>() {
        return Ok(values.into_iter().map(|v| v as f64).collect());
    }
    if let Ok(list) = value.cast::<PyList>() {
        let mut converted = Vec::with_capacity(list.len());
        for item in list.iter() {
            converted.push(item.extract::<f64>().map_err(|_| {
                PyTypeError::new_err(format!(
                    "Column '{name}' must be numeric (float32/float64/list of floats)"
                ))
            })?);
        }
        return Ok(converted);
    }
    Err(PyTypeError::new_err(format!(
        "Column '{name}' must be numeric (float32/float64/list of floats)"
    )))
}

/// A single event
///
/// Events are composed of a set of 4-momenta of particles in the overall
/// center-of-momentum frame, optional auxiliary scalars (e.g. polarization magnitude or angle),
/// and a weight.
///
/// Parameters
/// ----------
/// p4s : list of Vec4
///     4-momenta of each particle in the event in the overall center-of-momentum frame
/// aux: list of float
///     Scalar auxiliary data associated with the event
/// weight : float
///     The weight associated with this event
/// p4_names : list of str, optional
///     Human-readable aliases for each four-momentum. Providing names enables name-based
///     lookups when evaluating variables.
/// aux_names : list of str, optional
///     Aliases for auxiliary scalars corresponding to ``aux``.
/// aliases : dict of {str: str or list[str]}, optional
///     Additional particle identifiers that reference one or more entries from ``p4_names``.
///
/// Examples
/// --------
/// >>> from laddu import Event, Vec3  # doctest: +SKIP
/// >>> event = Event(  # doctest: +SKIP
/// ...     [Vec3(0.0, 0.0, 1.0).with_mass(0.0), Vec3(0.0, 0.0, 1.0).with_mass(0.0)],
/// ...     [],
/// ...     1.0,
/// ...     p4_names=['kshort1', 'kshort2'],
/// ...     aliases={'pair': ['kshort1', 'kshort2']},
/// ... )
/// >>> event.p4('pair')  # doctest: +SKIP
/// Vec4(px=0.0, py=0.0, pz=2.0, e=2.0)
/// >>> event.aux['pol_angle']  # doctest: +SKIP
/// 0.3
///
#[pyclass(name = "Event", module = "laddu", from_py_object)]
#[derive(Clone)]
pub struct PyEvent {
    pub event: Event,
    has_metadata: bool,
}

#[pymethods]
impl PyEvent {
    #[new]
    #[pyo3(signature = (p4s, aux, weight, *, p4_names=None, aux_names=None, aliases=None))]
    fn new(
        p4s: Vec<PyVec4>,
        aux: Vec<f64>,
        weight: f64,
        p4_names: Option<Vec<String>>,
        aux_names: Option<Vec<String>>,
        aliases: Option<Bound<PyDict>>,
    ) -> PyResult<Self> {
        let event = EventData {
            p4s: p4s.into_iter().map(|arr| arr.0).collect(),
            aux,
            weight,
        };
        let aliases = parse_aliases(aliases)?;

        let missing_p4_names = p4_names
            .as_ref()
            .map(|names| names.is_empty())
            .unwrap_or(true);

        if !aliases.is_empty() && missing_p4_names {
            return Err(PyValueError::new_err(
                "`aliases` requires `p4_names` so selections can be resolved",
            ));
        }

        let metadata_provided = p4_names.is_some() || aux_names.is_some() || !aliases.is_empty();
        let metadata = if metadata_provided {
            let p4_names = p4_names.unwrap_or_default();
            let aux_names = aux_names.unwrap_or_default();
            let mut metadata = DatasetMetadata::new(p4_names, aux_names).map_err(PyErr::from)?;
            if !aliases.is_empty() {
                metadata
                    .add_p4_aliases(
                        aliases.into_iter().map(|(alias_name, selection)| {
                            (alias_name, selection.into_selection())
                        }),
                    )
                    .map_err(PyErr::from)?;
            }
            Arc::new(metadata)
        } else {
            Arc::new(DatasetMetadata::empty())
        };
        let event = Event::new(Arc::new(event), metadata);
        Ok(Self {
            event,
            has_metadata: metadata_provided,
        })
    }
    fn __str__(&self) -> String {
        self.event.data().to_string()
    }
    /// The list of 4-momenta for each particle in the event
    ///
    #[getter]
    fn p4s<'py>(&self, py: Python<'py>) -> PyResult<Py<PyDict>> {
        self.ensure_metadata()?;
        let mapping = PyDict::new(py);
        for (name, vec4) in self.event.p4s() {
            mapping.set_item(name, PyVec4(vec4))?;
        }
        Ok(mapping.into())
    }
    /// The auxiliary scalar values associated with the event
    ///
    #[getter]
    #[pyo3(name = "aux")]
    fn aux_mapping<'py>(&self, py: Python<'py>) -> PyResult<Py<PyDict>> {
        self.ensure_metadata()?;
        let mapping = PyDict::new(py);
        for (name, value) in self.event.aux() {
            mapping.set_item(name, value)?;
        }
        Ok(mapping.into())
    }
    /// The weight of this event relative to others in a Dataset
    ///
    #[getter]
    fn get_weight(&self) -> f64 {
        self.event.weight()
    }
    /// Get the sum of the four-momenta within the event at the given indices
    ///
    /// Parameters
    /// ----------
    /// names : list of str
    ///     The names of the four-momenta to sum
    ///
    /// Returns
    /// -------
    /// Vec4
    ///     The result of summing the given four-momenta
    ///
    fn get_p4_sum(&self, names: Vec<String>) -> PyResult<PyVec4> {
        let indices = self.resolve_p4_indices(&names)?;
        Ok(PyVec4(self.event.data().get_p4_sum(indices)))
    }
    /// Boost all the four-momenta in the event to the rest frame of the given set of
    /// four-momenta by indices.
    ///
    /// Parameters
    /// ----------
    /// names : list of str
    ///     The names of the four-momenta whose rest frame should be used for the boost
    ///
    /// Returns
    /// -------
    /// Event
    ///     The boosted event
    ///
    pub fn boost_to_rest_frame_of(&self, names: Vec<String>) -> PyResult<Self> {
        let indices = self.resolve_p4_indices(&names)?;
        let boosted = self.event.data().boost_to_rest_frame_of(indices);
        Ok(Self {
            event: Event::new(Arc::new(boosted), self.event.metadata_arc()),
            has_metadata: self.has_metadata,
        })
    }
    /// Get the value of a Variable on the given Event
    ///
    /// Parameters
    /// ----------
    /// variable : {laddu.Mass, laddu.CosTheta, laddu.Phi, laddu.PolAngle, laddu.PolMagnitude, laddu.Mandelstam}
    ///
    /// Returns
    /// -------
    /// float
    ///
    /// Notes
    /// -----
    /// Variables that rely on particle names require the event to carry metadata. Provide
    /// ``p4_names``/``aux_names`` when constructing the event or evaluate variables through a
    /// ``laddu.Dataset`` to ensure the metadata is available.
    ///
    /// Examples
    /// --------
    /// >>> from laddu import Event, Vec3  # doctest: +SKIP
    /// >>> from laddu.utils.variables import Mass  # doctest: +SKIP
    /// >>> event = Event(  # doctest: +SKIP
    /// ...     [Vec3(0.0, 0.0, 1.0).with_mass(0.938)],
    /// ...     [],
    /// ...     1.0,
    /// ...     p4_names=['proton'],
    /// ... )
    /// >>> event.evaluate(Mass(['proton']))  # doctest: +SKIP
    /// 0.938
    ///
    fn evaluate(&self, variable: Bound<'_, PyAny>) -> PyResult<f64> {
        let mut variable = variable.extract::<PyVariable>()?;
        let metadata = self.ensure_metadata()?;
        variable.bind_in_place(metadata)?;
        variable.evaluate_event(&self.event)
    }

    /// Retrieve a four-momentum by name.
    fn p4(&self, name: &str) -> PyResult<PyVec4> {
        self.ensure_metadata()?;
        self.event
            .p4(name)
            .map(PyVec4)
            .ok_or_else(|| PyKeyError::new_err(format!("Unknown particle name '{name}'")))
    }
}

impl PyEvent {
    fn ensure_metadata(&self) -> PyResult<&DatasetMetadata> {
        if !self.has_metadata {
            Err(PyValueError::new_err(
                "Event has no associated metadata for name-based operations",
            ))
        } else {
            Ok(self.event.metadata())
        }
    }

    fn resolve_p4_indices(&self, names: &[String]) -> PyResult<Vec<usize>> {
        let metadata = self.ensure_metadata()?;
        let mut resolved = Vec::new();
        for name in names {
            let selection = metadata
                .p4_selection(name)
                .ok_or_else(|| PyKeyError::new_err(format!("Unknown particle name '{name}'")))?;
            resolved.extend_from_slice(selection.indices());
        }
        Ok(resolved)
    }

    pub(crate) fn metadata_opt(&self) -> Option<&DatasetMetadata> {
        self.has_metadata.then(|| self.event.metadata())
    }
}

#[doc(hidden)]
/// A set of Events
///
/// Datasets can be created from lists of Events or by using the constructor helpers
/// such as :func:`laddu.io.read_parquet`, :func:`laddu.io.read_root`, and
/// :func:`laddu.io.read_amptools`
///
/// Datasets can also be indexed directly to access individual Events
///
/// Parameters
/// ----------
/// events : list of Event
/// p4_names : list of str, optional
///     Names assigned to each four-momentum; enables name-based lookups if provided.
/// aux_names : list of str, optional
///     Names for auxiliary scalars stored alongside the events.
/// aliases : dict of {str: str or list[str]}, optional
///     Additional particle identifiers that override aliases stored on the Events.
///
/// Notes
/// -----
/// Explicit metadata provided here takes precedence over metadata embedded in the
/// input Events.
///
/// Examples
/// --------
/// >>> from laddu import Dataset, Event, Vec3  # doctest: +SKIP
/// >>> event = Event(  # doctest: +SKIP
/// ...     [Vec3(0.0, 0.0, 1.0).with_mass(0.0), Vec3(0.0, 0.0, -1.0).with_mass(0.938)],
/// ...     [0.4, 0.3],
/// ...     1.0,
/// ... )
/// >>> dataset = Dataset(  # doctest: +SKIP
/// ...     [event],
/// ...     p4_names=['beam', 'proton'],
/// ...     aux_names=['pol_magnitude', 'pol_angle'],
/// ...     aliases={'target': 'proton'},
/// ... )
/// >>> dataset[0].p4('target')  # doctest: +SKIP
/// Vec4(px=0.0, py=0.0, pz=-1.0, e=1.371073...)
///
#[pyclass(name = "Dataset", module = "laddu", subclass, skip_from_py_object)]
#[derive(Clone)]
pub struct PyDataset(pub Arc<Dataset>);

#[pyclass(
    name = "ParquetChunkIter",
    module = "laddu",
    unsendable,
    skip_from_py_object
)]
pub struct PyParquetChunkIter {
    chunks: Box<dyn Iterator<Item = laddu_core::LadduResult<Arc<Dataset>>> + Send>,
}

#[pymethods]
impl PyParquetChunkIter {
    fn __iter__(slf: PyRef<'_, Self>) -> Py<PyParquetChunkIter> {
        slf.into()
    }

    fn __next__(&mut self) -> PyResult<Option<PyDataset>> {
        match self.chunks.next() {
            Some(Ok(dataset)) => Ok(Some(PyDataset(dataset))),
            Some(Err(err)) => Err(PyErr::from(err)),
            None => Ok(None),
        }
    }
}

#[pyclass(
    name = "DatasetIter",
    module = "laddu",
    unsendable,
    skip_from_py_object
)]
struct PyDatasetIter {
    kind: PyDatasetIterKind,
}

enum PyDatasetIterKind {
    Local { dataset: Arc<Dataset>, index: usize },
    Global(DatasetArcIter),
}

#[pymethods]
impl PyDatasetIter {
    fn __iter__(slf: PyRef<'_, Self>) -> Py<PyDatasetIter> {
        slf.into()
    }

    fn __next__(&mut self) -> Option<PyEvent> {
        let event = match &mut self.kind {
            PyDatasetIterKind::Local { dataset, index } => {
                let event = dataset.events_local().get(*index)?.clone();
                *index += 1;
                event
            }
            PyDatasetIterKind::Global(iterator) => iterator.next()?,
        };
        Some(PyEvent {
            event,
            has_metadata: true,
        })
    }
}

#[pymethods]
impl PyDataset {
    #[new]
    #[pyo3(signature = (events, *, p4_names=None, aux_names=None, aliases=None))]
    fn new(
        events: Vec<PyEvent>,
        p4_names: Option<Vec<String>>,
        aux_names: Option<Vec<String>>,
        aliases: Option<Bound<PyDict>>,
    ) -> PyResult<Self> {
        let inferred_metadata = events
            .iter()
            .find_map(|event| event.has_metadata.then(|| event.event.metadata_arc()));

        let aliases = parse_aliases(aliases)?;
        let use_explicit_metadata =
            p4_names.is_some() || aux_names.is_some() || !aliases.is_empty();

        let metadata =
            if use_explicit_metadata {
                let resolved_p4_names = match (p4_names, inferred_metadata.as_ref()) {
                    (Some(names), _) => names,
                    (None, Some(metadata)) => metadata.p4_names().to_vec(),
                    (None, None) => Vec::new(),
                };
                let resolved_aux_names = match (aux_names, inferred_metadata.as_ref()) {
                    (Some(names), _) => names,
                    (None, Some(metadata)) => metadata.aux_names().to_vec(),
                    (None, None) => Vec::new(),
                };

                if !aliases.is_empty() && resolved_p4_names.is_empty() {
                    return Err(PyValueError::new_err(
                        "`aliases` requires `p4_names` or events with metadata for resolution",
                    ));
                }

                let mut metadata = DatasetMetadata::new(resolved_p4_names, resolved_aux_names)
                    .map_err(PyErr::from)?;
                if !aliases.is_empty() {
                    metadata
                        .add_p4_aliases(aliases.into_iter().map(|(alias_name, selection)| {
                            (alias_name, selection.into_selection())
                        }))
                        .map_err(PyErr::from)?;
                }
                Some(Arc::new(metadata))
            } else {
                inferred_metadata
            };

        let events: Vec<Arc<EventData>> = events
            .into_iter()
            .map(|event| event.event.data_arc())
            .collect();
        let dataset = if let Some(metadata) = metadata {
            Dataset::new_with_metadata(events, metadata)
        } else {
            Dataset::new(events)
        };
        Ok(Self(Arc::new(dataset)))
    }

    fn __len__(&self) -> usize {
        self.0.n_events()
    }
    /// Iterate over all events in dataset order.
    ///
    /// Notes
    /// -----
    /// This is the default iterator used by ``for event in dataset``.
    /// When MPI is enabled, it preserves global indexing semantics and may fetch
    /// remote events as needed.
    fn __iter__(&self) -> PyDatasetIter {
        self.iter_global()
    }
    /// Get the number of events owned by the current rank.
    #[getter]
    fn n_events_local(&self) -> usize {
        self.0.n_events_local()
    }
    /// Iterate over the events owned by the current rank.
    ///
    /// Notes
    /// -----
    /// When MPI is disabled, this iterates over the full Dataset.
    /// When MPI is enabled, this iterates only over events stored on the current rank.
    /// The yielded order matches the current rank's local storage order.
    fn iter_local(&self) -> PyDatasetIter {
        PyDatasetIter {
            kind: PyDatasetIterKind::Local {
                dataset: self.0.clone(),
                index: 0,
            },
        }
    }
    /// Iterate over all events in the Dataset.
    ///
    /// Notes
    /// -----
    /// This is the default iterator used by ``for event in dataset``.
    /// When MPI is enabled, this preserves global dataset order and performs
    /// explicit cross-rank event fetches as needed.
    fn iter_global(&self) -> PyDatasetIter {
        PyDatasetIter {
            kind: PyDatasetIterKind::Global(self.0.shared_iter_global()),
        }
    }
    fn __add__(&self, other: &Bound<'_, PyAny>) -> PyResult<PyDataset> {
        if let Ok(other_ds) = other.extract::<PyRef<PyDataset>>() {
            Ok(PyDataset(Arc::new(self.0.as_ref() + other_ds.0.as_ref())))
        } else if let Ok(other_int) = other.extract::<usize>() {
            if other_int == 0 {
                Ok(self.clone())
            } else {
                Err(PyTypeError::new_err(
                    "Addition with an integer for this type is only defined for 0",
                ))
            }
        } else {
            Err(PyTypeError::new_err("Unsupported operand type for +"))
        }
    }
    fn __radd__(&self, other: &Bound<'_, PyAny>) -> PyResult<PyDataset> {
        if let Ok(other_ds) = other.extract::<PyRef<PyDataset>>() {
            Ok(PyDataset(Arc::new(other_ds.0.as_ref() + self.0.as_ref())))
        } else if let Ok(other_int) = other.extract::<usize>() {
            if other_int == 0 {
                Ok(self.clone())
            } else {
                Err(PyTypeError::new_err(
                    "Addition with an integer for this type is only defined for 0",
                ))
            }
        } else {
            Err(PyTypeError::new_err("Unsupported operand type for +"))
        }
    }
    /// Get the number of Events in the Dataset
    ///
    /// Notes
    /// -----
    /// When MPI is enabled, this returns the global event count.
    /// It therefore matches ``len(dataset)`` and the valid range for ``dataset[i]``.
    ///
    /// Returns
    /// -------
    /// n_events : int
    ///     The number of Events
    ///
    #[getter]
    fn n_events(&self) -> usize {
        self.0.n_events()
    }
    /// Alias for ``n_events``.
    #[getter]
    fn n_events_global(&self) -> usize {
        self.0.n_events_global()
    }
    /// Particle names used to construct four-momenta when loading from a Parquet file.
    #[getter]
    fn p4_names(&self) -> Vec<String> {
        self.0.p4_names().to_vec()
    }
    /// Auxiliary scalar names associated with this Dataset.
    #[getter]
    fn aux_names(&self) -> Vec<String> {
        self.0.aux_names().to_vec()
    }

    /// Get the weighted number of Events in the Dataset
    ///
    /// Notes
    /// -----
    /// When MPI is enabled, this returns the global weighted event count.
    ///
    /// Returns
    /// -------
    /// n_events : float
    ///     The sum of all Event weights
    ///
    #[getter]
    fn n_events_weighted(&self) -> f64 {
        self.0.n_events_weighted()
    }
    /// Alias for ``n_events_weighted``.
    #[getter]
    fn n_events_weighted_global(&self) -> f64 {
        self.0.n_events_weighted_global()
    }
    /// Get the weighted number of local Events in the Dataset
    ///
    /// Notes
    /// -----
    /// When MPI is enabled, this returns the sum of the current rank's Event
    /// weights.
    ///
    /// Returns
    /// -------
    /// n_events : float
    ///     The sum of the current rank's Event weights
    ///
    #[getter]
    fn n_events_weighted_local(&self) -> f64 {
        self.0.n_events_weighted_local()
    }
    /// The weights associated with the Dataset
    ///
    /// Notes
    /// -----
    /// When MPI is enabled, this returns the global weight vector in dataset order.
    ///
    /// Returns
    /// -------
    /// weights : array_like
    ///     A ``numpy`` array of global Event weights
    ///
    #[getter]
    fn weights<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_slice(py, &self.0.weights())
    }
    /// Alias for ``weights``.
    #[getter]
    fn weights_global<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_slice(py, &self.0.weights_global())
    }
    /// The weights associated with the Dataset on the current rank.
    ///
    /// Notes
    /// -----
    /// This is the explicit rank-local counterpart to ``weights``.
    ///
    /// Returns
    /// -------
    /// weights : array_like
    ///     A ``numpy`` array of rank-local Event weights
    ///
    #[getter]
    fn weights_local<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_slice(py, &self.0.weights_local())
    }
    /// The internal list of Events stored in the Dataset
    ///
    /// Notes
    /// -----
    /// When MPI is enabled, this returns the full global event list.
    /// Use ``events_local`` or ``iter_local()`` to access only the current rank's
    /// event ownership.
    /// The returned list matches the order produced by ``for event in dataset``.
    ///
    /// Returns
    /// -------
    /// events : list of Event
    ///     The Events in the Dataset
    ///
    #[getter]
    fn events(&self) -> Vec<PyEvent> {
        self.0
            .shared_iter()
            .map(|rust_event| PyEvent {
                event: rust_event,
                has_metadata: true,
            })
            .collect()
    }
    /// Alias for ``events``.
    #[getter]
    fn events_global(&self) -> Vec<PyEvent> {
        self.events()
    }
    /// The list of Events stored on the current rank.
    ///
    /// Notes
    /// -----
    /// This is the explicit rank-local counterpart to ``events``.
    /// The returned list matches the order produced by ``iter_local()``.
    #[getter]
    fn events_local(&self) -> Vec<PyEvent> {
        self.0
            .events_local()
            .iter()
            .map(|rust_event| PyEvent {
                event: rust_event.clone(),
                has_metadata: true,
            })
            .collect()
    }
    /// Retrieve a four-momentum by particle name for the event at ``index``.
    fn p4_by_name(&self, index: usize, name: &str) -> PyResult<PyVec4> {
        self.0
            .p4_by_name(index, name)
            .map(PyVec4)
            .ok_or_else(|| PyKeyError::new_err(format!("Unknown particle name '{name}'")))
    }
    /// Retrieve an auxiliary scalar by name for the event at ``index``.
    fn aux_by_name(&self, index: usize, name: &str) -> PyResult<f64> {
        self.0
            .aux_by_name(index, name)
            .ok_or_else(|| PyKeyError::new_err(format!("Unknown auxiliary name '{name}'")))
    }
    /// Alias for ``dataset[index]``.
    ///
    /// Notes
    /// -----
    /// This preserves the default global indexing semantics under MPI.
    fn event_global(&self, index: usize) -> PyResult<PyEvent> {
        let event = self
            .0
            .get_event_global(index)
            .ok_or_else(|| PyIndexError::new_err("index out of range"))?;
        Ok(PyEvent {
            event,
            has_metadata: true,
        })
    }
    fn __getitem__<'py>(
        &self,
        py: Python<'py>,
        index: Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyAny>> {
        if let Ok(value) = self.evaluate(py, index.clone()) {
            value.into_bound_py_any(py)
        } else if let Ok(index) = index.extract::<usize>() {
            let event = self
                .0
                .get_event(index)
                .ok_or_else(|| PyIndexError::new_err("index out of range"))?;
            PyEvent {
                event,
                has_metadata: true,
            }
            .into_bound_py_any(py)
        } else {
            Err(PyTypeError::new_err(
                "Unsupported index type (int or Variable)",
            ))
        }
    }
    /// Separates a Dataset into histogram bins by a Variable value
    ///
    /// Parameters
    /// ----------
    /// variable : {laddu.Mass, laddu.CosTheta, laddu.Phi, laddu.PolAngle, laddu.PolMagnitude, laddu.Mandelstam}
    ///     The Variable by which each Event is binned
    /// bins : int
    ///     The number of equally-spaced bins
    /// range : tuple[float, float]
    ///     The minimum and maximum bin edges
    ///
    /// Returns
    /// -------
    /// datasets : BinnedDataset
    ///     A structure containing Datasets binned by the given `variable`
    ///
    /// See Also
    /// --------
    /// laddu.Mass
    /// laddu.CosTheta
    /// laddu.Phi
    /// laddu.PolAngle
    /// laddu.PolMagnitude
    /// laddu.Mandelstam
    ///
    /// Examples
    /// --------
    /// >>> from laddu.utils.variables import Mass  # doctest: +SKIP
    /// >>> binned = dataset.bin_by(Mass(['kshort1']), bins=10, range=(0.9, 1.5))  # doctest: +SKIP
    /// >>> len(binned)  # doctest: +SKIP
    /// 10
    ///
    /// Raises
    /// ------
    /// TypeError
    ///     If the given `variable` is not a valid variable
    ///
    #[pyo3(signature = (variable, bins, range))]
    fn bin_by(
        &self,
        variable: Bound<'_, PyAny>,
        bins: usize,
        range: (f64, f64),
    ) -> PyResult<PyBinnedDataset> {
        let py_variable = variable.extract::<PyVariable>()?;
        let bound_variable = py_variable.bound(self.0.metadata())?;
        Ok(PyBinnedDataset(self.0.bin_by(
            bound_variable,
            bins,
            range,
        )?))
    }
    /// Filter the Dataset by a given VariableExpression, selecting events for which the expression returns ``True``.
    ///
    /// Parameters
    /// ----------
    /// expression : VariableExpression
    ///     The expression with which to filter the Dataset
    ///
    /// Returns
    /// -------
    /// Dataset
    ///     The filtered Dataset
    ///
    /// Examples
    /// --------
    /// >>> from laddu.utils.variables import Mass  # doctest: +SKIP
    /// >>> heavy = dataset.filter(Mass(['kshort1']) > 1.0)  # doctest: +SKIP
    ///
    pub fn filter(&self, expression: &PyVariableExpression) -> PyResult<PyDataset> {
        Ok(PyDataset(
            self.0.filter(&expression.0).map_err(PyErr::from)?,
        ))
    }
    /// Generate a new bootstrapped Dataset by randomly resampling the original with replacement
    ///
    /// The new Dataset is resampled with a random generator seeded by the provided `seed`
    ///
    /// Parameters
    /// ----------
    /// seed : int
    ///     The random seed used in the resampling process
    ///
    /// Returns
    /// -------
    /// Dataset
    ///     A bootstrapped Dataset
    ///
    /// Examples
    /// --------
    /// >>> replica = dataset.bootstrap(2024)  # doctest: +SKIP
    /// >>> len(replica) == len(dataset)  # doctest: +SKIP
    /// True
    ///
    fn bootstrap(&self, seed: usize) -> PyDataset {
        PyDataset(self.0.bootstrap(seed))
    }
    /// Boost all the four-momenta in all events to the rest frame of the given set of
    /// named four-momenta.
    ///
    /// Parameters
    /// ----------
    /// names : list of str
    ///     The names of the four-momenta defining the rest frame
    ///
    /// Returns
    /// -------
    /// Dataset
    ///     The boosted dataset
    ///
    /// Examples
    /// --------
    /// >>> dataset.boost_to_rest_frame_of(['kshort1', 'kshort2'])  # doctest: +SKIP
    ///
    pub fn boost_to_rest_frame_of(&self, names: Vec<String>) -> PyDataset {
        PyDataset(self.0.boost_to_rest_frame_of(&names))
    }
    /// Get the value of a Variable over every event in the Dataset.
    ///
    /// Parameters
    /// ----------
    /// variable : {laddu.Mass, laddu.CosTheta, laddu.Phi, laddu.PolAngle, laddu.PolMagnitude, laddu.Mandelstam}
    ///
    /// Returns
    /// -------
    /// values : array_like
    ///
    /// Examples
    /// --------
    /// >>> from laddu.utils.variables import Mass  # doctest: +SKIP
    /// >>> masses = dataset.evaluate(Mass(['kshort1']))  # doctest: +SKIP
    /// >>> masses.shape  # doctest: +SKIP
    /// (len(dataset),)
    ///
    fn evaluate<'py>(
        &self,
        py: Python<'py>,
        variable: Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
        let variable = variable.extract::<PyVariable>()?;
        let bound_variable = variable.bound(self.0.metadata())?;
        let values = self.0.evaluate(&bound_variable).map_err(PyErr::from)?;
        Ok(PyArray1::from_vec(py, values))
    }
}

/// Read a Dataset from a Parquet file.
///
/// Examples
/// --------
/// >>> import laddu.io as ldio  # doctest: +SKIP
/// >>> dataset = ldio.read_parquet(  # doctest: +SKIP
/// ...     'events.parquet',
/// ...     p4s=['beam', 'proton'],
/// ...     aux=['pol_magnitude', 'pol_angle'],
/// ...     aliases={'target': 'proton'},
/// ... )
/// >>> dataset.p4_names  # doctest: +SKIP
/// ['beam', 'proton']
#[pyfunction]
#[pyo3(signature = (path, *, p4s=None, aux=None, aliases=None))]
pub fn read_parquet(
    path: Bound<PyAny>,
    p4s: Option<Vec<String>>,
    aux: Option<Vec<String>>,
    aliases: Option<Bound<PyDict>>,
) -> PyResult<PyDataset> {
    let path_str = parse_dataset_path(path)?;
    let mut read_options = DatasetReadOptions::default();
    if let Some(p4s) = p4s {
        read_options = read_options.p4_names(p4s);
    }
    if let Some(aux) = aux {
        read_options = read_options.aux_names(aux);
    }
    for (alias_name, selection) in parse_aliases(aliases)?.into_iter() {
        read_options = read_options.alias(alias_name, selection);
    }
    let dataset = core_read_parquet(&path_str, &read_options)?;
    Ok(PyDataset(dataset))
}

/// Read a Dataset from a Parquet file in chunks.
#[pyfunction]
#[pyo3(signature = (path, *, p4s=None, aux=None, aliases=None, chunk_size=None))]
pub fn read_parquet_chunked(
    path: Bound<PyAny>,
    p4s: Option<Vec<String>>,
    aux: Option<Vec<String>>,
    aliases: Option<Bound<PyDict>>,
    chunk_size: Option<usize>,
) -> PyResult<PyParquetChunkIter> {
    let path_str = parse_dataset_path(path)?;
    let mut read_options = DatasetReadOptions::default();
    if let Some(p4s) = p4s {
        read_options = read_options.p4_names(p4s);
    }
    if let Some(aux) = aux {
        read_options = read_options.aux_names(aux);
    }
    if let Some(chunk_size) = chunk_size {
        read_options = read_options.chunk_size(chunk_size);
    }
    for (alias_name, selection) in parse_aliases(aliases)?.into_iter() {
        read_options = read_options.alias(alias_name, selection);
    }

    let chunks = core_read_parquet_chunks_with_options(&path_str, &read_options)?;
    Ok(PyParquetChunkIter {
        chunks: Box::new(chunks),
    })
}

/// Read a Dataset from a ROOT file using the oxyroot backend.
///
/// Examples
/// --------
/// >>> import laddu.io as ldio  # doctest: +SKIP
/// >>> dataset = ldio.read_root(  # doctest: +SKIP
/// ...     'events.root',
/// ...     tree='kin',
/// ...     p4s=['beam', 'proton'],
/// ...     aux=['pol_magnitude', 'pol_angle'],
/// ... )
/// >>> dataset.aux_names  # doctest: +SKIP
/// ['pol_magnitude', 'pol_angle']
#[pyfunction]
#[pyo3(signature = (path, *, tree=None, p4s=None, aux=None, aliases=None))]
pub fn read_root(
    path: Bound<PyAny>,
    tree: Option<String>,
    p4s: Option<Vec<String>>,
    aux: Option<Vec<String>>,
    aliases: Option<Bound<PyDict>>,
) -> PyResult<PyDataset> {
    let path_str = parse_dataset_path(path)?;
    let mut read_options = DatasetReadOptions::default();
    if let Some(p4s) = p4s {
        read_options = read_options.p4_names(p4s);
    }
    if let Some(aux) = aux {
        read_options = read_options.aux_names(aux);
    }
    if let Some(tree) = tree {
        read_options = read_options.tree(tree);
    }
    for (alias_name, selection) in parse_aliases(aliases)?.into_iter() {
        read_options = read_options.alias(alias_name, selection);
    }
    let dataset = core_read_root(&path_str, &read_options)?;
    Ok(PyDataset(dataset))
}

/// Write a Dataset to a Parquet file.
#[pyfunction]
#[pyo3(signature = (dataset, path, *, chunk_size=None, precision="f64"))]
pub fn write_parquet(
    dataset: &PyDataset,
    path: Bound<PyAny>,
    chunk_size: Option<usize>,
    precision: &str,
) -> PyResult<()> {
    let path_str = parse_dataset_path(path)?;
    let mut write_options = DatasetWriteOptions::default();
    if let Some(size) = chunk_size {
        write_options.batch_size = size.max(1);
    }
    write_options.precision = parse_precision_arg(Some(precision))?;
    core_write_parquet(dataset.0.as_ref(), &path_str, &write_options).map_err(PyErr::from)
}

/// Write a Dataset to a ROOT file using the oxyroot backend.
#[pyfunction]
#[pyo3(signature = (dataset, path, *, tree=None, chunk_size=None, precision="f64"))]
pub fn write_root(
    dataset: &PyDataset,
    path: Bound<PyAny>,
    tree: Option<String>,
    chunk_size: Option<usize>,
    precision: &str,
) -> PyResult<()> {
    let path_str = parse_dataset_path(path)?;
    let mut write_options = DatasetWriteOptions::default();
    if let Some(name) = tree {
        write_options.tree = Some(name);
    }
    if let Some(size) = chunk_size {
        write_options.batch_size = size.max(1);
    }
    write_options.precision = parse_precision_arg(Some(precision))?;
    core_write_root(dataset.0.as_ref(), &path_str, &write_options).map_err(PyErr::from)
}

#[doc(hidden)]
/// Build a Dataset from columnar arrays.
///
/// This is the canonical high-throughput ingestion path used by Python reader helpers.
///
/// Examples
/// --------
/// >>> import laddu.io as ldio  # doctest: +SKIP
/// >>> dataset = ldio.from_columns(  # doctest: +SKIP
/// ...     {
/// ...         'beam_px': [0.0],
/// ...         'beam_py': [0.0],
/// ...         'beam_pz': [8.5],
/// ...         'beam_e': [8.5],
/// ...         'proton_px': [0.0],
/// ...         'proton_py': [0.0],
/// ...         'proton_pz': [-0.2],
/// ...         'proton_e': [0.959],
/// ...         'pol_magnitude': [0.4],
/// ...         'pol_angle': [0.3],
/// ...         'weight': [1.0],
/// ...     },
/// ...     p4s=['beam', 'proton'],
/// ...     aux=['pol_magnitude', 'pol_angle'],
/// ...     aliases={'target': 'proton'},
/// ... )
/// >>> dataset[0].p4('target')  # doctest: +SKIP
/// Vec4(px=0.0, py=0.0, pz=-0.2, e=0.959)
#[pyfunction]
#[pyo3(signature = (columns, *, p4s=None, aux=None, aliases=None))]
pub fn from_columns(
    columns: Bound<'_, PyDict>,
    p4s: Option<Vec<String>>,
    aux: Option<Vec<String>>,
    aliases: Option<Bound<'_, PyDict>>,
) -> PyResult<PyDataset> {
    let column_names = columns
        .iter()
        .map(|(key, _)| key.extract::<String>())
        .collect::<PyResult<Vec<_>>>()?;

    let (detected_p4_names, detected_aux_names) =
        infer_p4_and_aux_names_from_columns(&column_names);
    let p4_names = p4s.unwrap_or(detected_p4_names);
    if p4_names.is_empty() {
        let mut partial_components: std::collections::BTreeMap<
            String,
            std::collections::BTreeSet<&str>,
        > = std::collections::BTreeMap::new();
        for column_name in &column_names {
            let lowered = column_name.to_ascii_lowercase();
            for suffix in P4_COMPONENT_SUFFIXES {
                if lowered.ends_with(suffix) && column_name.len() > suffix.len() {
                    let prefix = column_name[..column_name.len() - suffix.len()].to_string();
                    partial_components.entry(prefix).or_default().insert(suffix);
                }
            }
        }
        if let Some((prefix, present)) = partial_components.iter().next() {
            if present.len() < P4_COMPONENT_SUFFIXES.len() {
                let missing = P4_COMPONENT_SUFFIXES
                    .iter()
                    .filter(|suffix| !present.contains(**suffix))
                    .map(|suffix| format!("{prefix}{suffix}"))
                    .collect::<Vec<_>>()
                    .join(", ");
                return Err(PyKeyError::new_err(format!(
                    "Missing components [{missing}] for four-momentum '{prefix}'"
                )));
            }
        }
        return Err(PyValueError::new_err(
            "No four-momentum columns found (expected *_px, *_py, *_pz, *_e)",
        ));
    }

    let aux_names = aux.unwrap_or(detected_aux_names);
    let p4_component_columns =
        resolve_p4_component_columns(&column_names, &p4_names).map_err(PyErr::from)?;
    let resolved_aux_columns =
        resolve_columns_case_insensitive(&column_names, &aux_names).map_err(PyErr::from)?;

    let n_events = {
        let first_name = p4_component_columns
            .first()
            .map(|components| components[0].clone())
            .ok_or_else(|| PyKeyError::new_err("Missing required p4 column"))?;
        let values = extract_numeric_column(
            columns
                .get_item(first_name.as_str())?
                .ok_or_else(|| PyKeyError::new_err("Missing required p4 column"))?,
            &first_name,
        )?;
        values.len()
    };

    let mut p4_columns: Vec<[Vec<f64>; 4]> = Vec::with_capacity(p4_names.len());
    for component_names in &p4_component_columns {
        let px = extract_numeric_column(
            columns
                .get_item(component_names[0].as_str())?
                .ok_or_else(|| PyKeyError::new_err(format!("Missing {}", component_names[0])))?,
            component_names[0].as_str(),
        )?;
        let py = extract_numeric_column(
            columns
                .get_item(component_names[1].as_str())?
                .ok_or_else(|| PyKeyError::new_err(format!("Missing {}", component_names[1])))?,
            component_names[1].as_str(),
        )?;
        let pz = extract_numeric_column(
            columns
                .get_item(component_names[2].as_str())?
                .ok_or_else(|| PyKeyError::new_err(format!("Missing {}", component_names[2])))?,
            component_names[2].as_str(),
        )?;
        let e = extract_numeric_column(
            columns
                .get_item(component_names[3].as_str())?
                .ok_or_else(|| PyKeyError::new_err(format!("Missing {}", component_names[3])))?,
            component_names[3].as_str(),
        )?;
        if px.len() != n_events
            || py.len() != n_events
            || pz.len() != n_events
            || e.len() != n_events
        {
            return Err(PyValueError::new_err(
                "All p4 components must have the same length",
            ));
        }
        p4_columns.push([px, py, pz, e]);
    }

    let mut aux_columns: Vec<Vec<f64>> = Vec::with_capacity(resolved_aux_columns.len());
    for (aux_name, aux_column_name) in aux_names.iter().zip(&resolved_aux_columns) {
        let values = extract_numeric_column(
            columns.get_item(aux_column_name.as_str())?.ok_or_else(|| {
                PyKeyError::new_err(format!("Missing auxiliary column '{aux_name}'"))
            })?,
            aux_name,
        )?;
        if values.len() != n_events {
            return Err(PyValueError::new_err(format!(
                "Auxiliary column '{aux_name}' length does not match p4 columns"
            )));
        }
        aux_columns.push(values);
    }

    let weights = if let Some(weight_column_name) = resolve_optional_weight_column(&column_names) {
        let weight_values = columns
            .get_item(weight_column_name.as_str())?
            .ok_or_else(|| PyKeyError::new_err("Missing weight column"))?;
        let values = extract_numeric_column(weight_values, "weight")?;
        if values.len() != n_events {
            return Err(PyValueError::new_err(
                "Column 'weight' length does not match p4 columns",
            ));
        }
        values
    } else {
        vec![1.0; n_events]
    };

    let parsed_aliases = parse_aliases(aliases)?;
    let mut metadata =
        DatasetMetadata::new(p4_names.clone(), aux_names.clone()).map_err(PyErr::from)?;
    if !parsed_aliases.is_empty() {
        metadata
            .add_p4_aliases(
                parsed_aliases
                    .into_iter()
                    .map(|(alias_name, selection)| (alias_name, selection.into_selection())),
            )
            .map_err(PyErr::from)?;
    }

    let mut events = Vec::with_capacity(n_events);
    for event_idx in 0..n_events {
        let p4s = p4_columns
            .iter()
            .map(|components| {
                laddu_core::utils::vectors::Vec4::new(
                    components[0][event_idx],
                    components[1][event_idx],
                    components[2][event_idx],
                    components[3][event_idx],
                )
            })
            .collect::<Vec<_>>();
        let aux = aux_columns
            .iter()
            .map(|values| values[event_idx])
            .collect::<Vec<_>>();
        events.push(Arc::new(EventData {
            p4s,
            aux,
            weight: weights[event_idx],
        }));
    }

    Ok(PyDataset(Arc::new(Dataset::new_with_metadata(
        events,
        Arc::new(metadata),
    ))))
}

/// A collection of Datasets binned by a Variable
///
/// BinnedDatasets can be indexed directly to access the underlying Datasets by bin
///
/// See Also
/// --------
/// laddu.Dataset.bin_by
///
#[pyclass(name = "BinnedDataset", module = "laddu", skip_from_py_object)]
pub struct PyBinnedDataset(BinnedDataset);

#[pymethods]
impl PyBinnedDataset {
    fn __len__(&self) -> usize {
        self.0.n_bins()
    }
    /// The number of bins in the BinnedDataset
    ///
    #[getter]
    fn n_bins(&self) -> usize {
        self.0.n_bins()
    }
    /// The minimum and maximum values of the binning Variable used to create this BinnedDataset
    ///
    #[getter]
    fn range(&self) -> (f64, f64) {
        self.0.range()
    }
    /// The edges of each bin in the BinnedDataset
    ///
    #[getter]
    fn edges<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<f64>> {
        PyArray1::from_slice(py, &self.0.edges())
    }
    fn __getitem__(&self, index: usize) -> PyResult<PyDataset> {
        self.0
            .get(index)
            .ok_or(PyIndexError::new_err("index out of range"))
            .map(|rust_dataset| PyDataset(rust_dataset.clone()))
    }
}