Fix parent capacity for dispatch

src/asset.rs

@@ -1100,6 +1100,41 @@ impl AssetRef {
             f(Some(&self), child);
         }
     }
+
+    /// Get an [`AssetRef`] representing a subset of this parent's children.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this asset is not a parent asset or `num_units` is zero or exceeds the total
+    /// capacity of this asset.
+    pub fn make_partial_parent(&self, num_units: u32) -> Self {
+        assert!(
+            self.is_parent(),
+            "Cannot make a partial parent from a non-parent asset"
+        );
+        assert!(
+            num_units > 0,
+            "Cannot make a partial parent with zero units"
+        );
+
+        let (max_num_units, unit_size) = match self.capacity() {
+            AssetCapacity::Discrete(max_num_units, unit_size) => (max_num_units, unit_size),
+            // We know asset capacity type is discrete as this is a parent asset
+            AssetCapacity::Continuous(_) => unreachable!(),
+        };
+        match num_units.cmp(&max_num_units) {
+            // Make a new Asset with fewer units
+            Ordering::Less => Self::from(Asset {
+                capacity: Cell::new(AssetCapacity::Discrete(num_units, unit_size)),
+                ..Asset::clone(self)
+            }),
+            // Same number of units as self
+            Ordering::Equal => self.clone(),
+            Ordering::Greater => {
+                panic!("Cannot make a partial parent with more units than original")
+            }
+        }
+    }
 }
 
 impl From<Rc<Asset>> for AssetRef {
@@ -1429,6 +1464,64 @@ mod tests {
         assert_eq!(count, 1);
     }
 
+    #[fixture]
+    fn parent_asset(asset_divisible: Asset) -> AssetRef {
+        let asset = AssetRef::from(asset_divisible);
+        let mut parent = None;
+
+        asset.into_for_each_child(&mut 0, |maybe_parent, _| {
+            if parent.is_none() {
+                parent = maybe_parent.cloned();
+            }
+        });
+
+        parent.expect("Divisible asset should create a parent")
+    }
+
+    #[rstest]
+    #[case::subset_of_children(2, false)]
+    #[case::all_children(3, true)]
+    fn make_partial_parent(
+        parent_asset: AssetRef,
+        #[case] num_units: u32,
+        #[case] expect_same_asset: bool,
+    ) {
+        let parent = parent_asset;
+        assert!(parent.is_parent());
+
+        let partial_parent = parent.make_partial_parent(num_units);
+
+        assert!(partial_parent.is_parent());
+        assert_eq!(
+            partial_parent.capacity(),
+            AssetCapacity::Discrete(num_units, Capacity(4.0))
+        );
+        assert_eq!(partial_parent.num_children(), Some(num_units));
+        assert_eq!(partial_parent.group_id(), parent.group_id());
+        assert_eq!(partial_parent.agent_id(), parent.agent_id());
+        assert_eq!(Rc::ptr_eq(&partial_parent.0, &parent.0), expect_same_asset);
+        assert_eq!(parent.capacity(), AssetCapacity::Discrete(3, Capacity(4.0)));
+    }
+
+    #[rstest]
+    #[should_panic(expected = "Cannot make a partial parent from a non-parent asset")]
+    fn make_partial_parent_panics_for_non_parent_asset(asset_divisible: Asset) {
+        let asset = AssetRef::from(asset_divisible);
+        asset.make_partial_parent(1);
+    }
+
+    #[rstest]
+    #[should_panic(expected = "Cannot make a partial parent with zero units")]
+    fn make_partial_parent_panics_for_zero_units(parent_asset: AssetRef) {
+        parent_asset.make_partial_parent(0);
+    }
+
+    #[rstest]
+    #[should_panic(expected = "Cannot make a partial parent with more units than original")]
+    fn make_partial_parent_panics_for_too_many_units(parent_asset: AssetRef) {
+        parent_asset.make_partial_parent(4);
+    }
+
     #[rstest]
     fn asset_commission(process: Process) {
         // Test successful commissioning of Future asset

src/simulation/optimisation.rs

@@ -17,9 +17,9 @@ use crate::units::{
 use anyhow::{Result, bail, ensure};
 use highs::{HighsModelStatus, HighsStatus, RowProblem as Problem, Sense};
 use indexmap::{IndexMap, IndexSet};
-use itertools::{Itertools, chain, iproduct};
+use itertools::{chain, iproduct};
 use std::cell::Cell;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::error::Error;
 use std::fmt;
 use std::ops::Range;
@@ -434,15 +434,30 @@ fn filter_input_prices(
 ///
 /// Child assets are converted to their parents and non-divisible assets are returned as is. Each
 /// parent asset is returned only once.
-fn get_parent_or_self(assets: &[AssetRef]) -> impl Iterator<Item = AssetRef> {
-    let mut parents = HashSet::new();
-    assets
-        .iter()
-        .filter_map(move |asset| match asset.parent() {
-            Some(parent) => parents.insert(parent.clone()).then_some(parent),
-            None => Some(asset),
-        })
-        .cloned()
+///
+/// If only a subset of a parent's children are present in `assets`, a new parent asset representing
+/// a portion of the total capacity will be created. This will have the same hash as the original
+/// parent.
+fn get_parent_or_self(assets: &[AssetRef]) -> Vec<AssetRef> {
+    let mut child_counts: IndexMap<&AssetRef, u32> = IndexMap::new();
+    let mut out = Vec::new();
+
+    for asset in assets {
+        if let Some(parent) = asset.parent() {
+            // For child assets, keep count of number of children per parent
+            *child_counts.entry(parent).or_default() += 1;
+        } else {
+            // Non-divisible assets can be returned as is
+            out.push(asset.clone());
+        }
+    }
+
+    for (parent, child_count) in child_counts {
+        // Convert to an object representing the appropriate portion of the parent's capacity
+        out.push(parent.make_partial_parent(child_count));
+    }
+
+    out
 }
 
 /// Provides the interface for running the dispatch optimisation.
@@ -621,7 +636,7 @@ impl<'model, 'run> DispatchRun<'model, 'run> {
         allow_unmet_demand: bool,
         input_prices: Option<&CommodityPrices>,
     ) -> Result<Solution<'model>, ModelError> {
-        let parent_assets = get_parent_or_self(self.existing_assets).collect_vec();
+        let parent_assets = get_parent_or_self(self.existing_assets);
 
         // Set up problem
         let mut problem = Problem::default();