diff --git a/pineappl/src/grid.rs b/pineappl/src/grid.rs index c7fcbea7..bcb4b076 100644 --- a/pineappl/src/grid.rs +++ b/pineappl/src/grid.rs @@ -851,12 +851,11 @@ impl Grid { /// TODO pub fn merge(&mut self, mut other: Self) -> Result<(), GridError> { let mut new_orders: Vec = Vec::new(); - let mut new_bins = 0; let mut new_entries: Vec = Vec::new(); - if self.bin_info() != other.bin_info() { + let new_bin_limits = if self.bin_info() != other.bin_info() { let lhs_bins = self.bin_info().bins(); - new_bins = other.bin_info().bins(); + let rhs_bins = other.bin_info().bins(); let lhs_remapper = self.remapper_mut(); let rhs_remapper = other.remapper(); @@ -866,57 +865,38 @@ impl Grid { lhs.merge(rhs).map_err(GridError::MergeBinError)?; let a = u32::try_from(lhs_bins).unwrap_or_else(|_| unreachable!()); - let b = u32::try_from(lhs_bins + new_bins).unwrap_or_else(|_| unreachable!()); + let b = u32::try_from(lhs_bins + rhs_bins).unwrap_or_else(|_| unreachable!()); - self.bin_limits = BinLimits::new((0..=b).map(f64::from).collect()); other.bin_limits = BinLimits::new((a..=b).map(f64::from).collect()); + BinLimits::new((0..=b).map(f64::from).collect()) } else { // Return an error todo!(); } } else if rhs_remapper.is_none() { - self.bin_limits + let mut new_bin_limits = self.bin_limits.clone(); + new_bin_limits .merge(&other.bin_limits) .map_err(GridError::InvalidBinLimits)?; + new_bin_limits } else { // Return an error todo!(); } - } + } else { + self.bin_limits.clone() + }; for ((i, _, k), _) in other .subgrids .indexed_iter_mut() .filter(|((_, _, _), subgrid)| !subgrid.is_empty()) { - let other_order = &other.orders[i]; - let other_entry = &other.lumi[k]; - - if !self - .orders - .iter() - .chain(new_orders.iter()) - .any(|x| x == other_order) - { - new_orders.push(other_order.clone()); - } - - if !self - .lumi - .iter() - .chain(new_entries.iter()) - .any(|y| y == other_entry) - { - new_entries.push(other_entry.clone()); - } + new_orders.push(other.orders[i].clone()); + new_entries.push(other.lumi[k].clone()); } - if !new_orders.is_empty() || !new_entries.is_empty() || (new_bins != 0) { - self.increase_shape(&(new_orders.len(), new_bins, new_entries.len())); - } - - self.orders.append(&mut new_orders); - self.lumi.append(&mut new_entries); + self.extend(&new_orders, &new_bin_limits, &new_entries); let bin_indices: Vec<_> = (0..other.bin_info().bins()) .map(|bin| { @@ -948,6 +928,50 @@ impl Grid { Ok(()) } + /// Extend subgrids array with the given `orders`, luminosity `entries`, and number of `bins`. + /// + /// Returns the inserted orders and entries, i.e. those not already present. + pub fn extend( + &mut self, + orders: &[Order], + bins: &BinLimits, + entries: &[LumiEntry], + ) -> (Vec, Vec) { + /// Extend collection with non-duplicated elements. + /// + /// Notice that this function does not modify the order of elements. + fn unique(current: &[T], update: &[T]) -> Vec { + let mut new = Vec::new(); + + // Since the order has to be preserved, this can't be achieved just by sorting and + // deduplicating. + for el in update.iter() { + if !current + .iter() + .chain(new.iter()) + .any(|x| x == el) + { + new.push((*el).clone()); + } + } + new + } + + let mut new_orders = unique(&self.orders, &orders); + let new_bins = bins.bins() - self.bin_limits.bins(); + let mut new_entries = unique(&self.lumi, &entries); + + if !new_orders.is_empty() || !new_entries.is_empty() || (new_bins != 0) { + self.increase_shape(&(new_orders.len(), new_bins, new_entries.len())); + } + + self.orders.append(&mut new_orders); + self.lumi.append(&mut new_entries); + self.bin_limits = bins.clone(); + + (new_orders, new_entries) + } + fn increase_shape(&mut self, new_dim: &(usize, usize, usize)) { let old_dim = self.subgrids.raw_dim().into_pattern(); let mut new_subgrids = Array3::from_shape_simple_fn(