use crate::balance;
use crate::composition::{Composition, Element};
use crate::error::{ProviderError, Result};
use crate::provider::PrecursorCatalog;
use crate::rejection::{RejectedCandidate, RejectionCode};
use crate::target::PlanningConstraints;
use std::collections::BTreeSet;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PrecursorId(pub String);
impl std::fmt::Display for PrecursorId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AvailabilityMetadata {
pub source: String,
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PrecursorCandidate {
pub id: PrecursorId,
pub composition: Composition,
pub availability: Option<AvailabilityMetadata>,
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PrecursorSelection {
pub precursor: PrecursorId,
pub formula_units: u64,
}
#[derive(Debug, Clone)]
pub struct InMemoryPrecursorCatalog {
candidates: Vec<PrecursorCandidate>,
}
impl InMemoryPrecursorCatalog {
pub fn new(mut candidates: Vec<PrecursorCandidate>) -> Self {
candidates.sort_by(|a, b| a.id.0.cmp(&b.id.0));
candidates.dedup_by(|a, b| a.id == b.id);
Self { candidates }
}
}
impl PrecursorCatalog for InMemoryPrecursorCatalog {
fn candidates_for(
&self,
target: &Composition,
_constraints: &PlanningConstraints,
) -> std::result::Result<Vec<PrecursorCandidate>, ProviderError> {
let target_elements: BTreeSet<Element> = target.elements().collect();
Ok(self
.candidates
.iter()
.filter(|c| {
c.composition
.elements()
.any(|e| target_elements.contains(&e))
})
.cloned()
.collect())
}
}
#[derive(Debug, Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AcceptedPrecursorSet {
pub precursors: Vec<PrecursorId>,
pub reaction: crate::reaction::BalancedReaction,
}
#[derive(Debug, Clone, PartialEq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PrecursorSearchOutcome {
pub accepted: Vec<AcceptedPrecursorSet>,
pub rejected: Vec<RejectedCandidate>,
}
pub fn search_precursor_sets(
target: &Composition,
candidates: &[PrecursorCandidate],
constraints: &PlanningConstraints,
budget: &crate::config::SearchBudget,
) -> Result<PrecursorSearchOutcome> {
let target_elements: BTreeSet<Element> = target.elements().collect();
let byproducts = balance::curated_byproducts()?;
let byproduct_elements: BTreeSet<Element> =
byproducts.iter().flat_map(Composition::elements).collect();
let byproduct_subsets = power_set(&byproducts);
let (combos, budget_exhausted) = generate_combinations(
candidates.len(),
budget.max_precursors_per_plan,
budget.max_precursor_sets,
);
let mut accepted = Vec::new();
let mut rejected = Vec::new();
for combo in &combos {
let chosen: Vec<&PrecursorCandidate> = combo.iter().map(|&i| &candidates[i]).collect();
let ids: Vec<PrecursorId> = chosen.iter().map(|c| c.id.clone()).collect();
if let Some(bad) = chosen
.iter()
.flat_map(|c| c.composition.elements())
.find(|e| constraints.forbidden_elements.contains(e))
{
rejected.push(RejectedCandidate {
precursors: ids,
reason_codes: vec![RejectionCode::ForbiddenElementPresent],
explanation: format!("precursor set contains forbidden element {bad}"),
});
continue;
}
let combo_elements: BTreeSet<Element> = chosen
.iter()
.flat_map(|c| c.composition.elements())
.collect();
let missing: Vec<Element> = target_elements
.difference(&combo_elements)
.copied()
.collect();
if !missing.is_empty() {
rejected.push(RejectedCandidate {
precursors: ids,
reason_codes: vec![RejectionCode::MissingTargetElement],
explanation: format!(
"precursor set does not cover target element(s): {}",
join_symbols(&missing)
),
});
continue;
}
let unremovable: Vec<Element> = combo_elements
.difference(&target_elements)
.filter(|e| !byproduct_elements.contains(e))
.copied()
.collect();
if !unremovable.is_empty() {
rejected.push(RejectedCandidate {
precursors: ids,
reason_codes: vec![RejectionCode::UnsupportedByproductRequired],
explanation: format!(
"precursor set introduces element(s) with no curated byproduct to remove them: {}",
join_symbols(&unremovable)
),
});
continue;
}
let reactant_compositions: Vec<Composition> =
chosen.iter().map(|c| c.composition.clone()).collect();
let mut found = Vec::new();
for subset in &byproduct_subsets {
let mut products = vec![target.clone()];
products.extend(subset.iter().cloned());
let results = balance::balance(&reactant_compositions, &products)?;
if !results.is_empty() {
found = results;
break;
}
}
if found.is_empty() {
rejected.push(RejectedCandidate {
precursors: ids,
reason_codes: vec![RejectionCode::NoStoichiometricBalance],
explanation:
"no integer balance exists for this precursor set against the target, \
with or without curated byproducts"
.to_string(),
});
continue;
}
for reaction in found {
let matched_ids: Vec<PrecursorId> = reaction
.reactants
.iter()
.map(|species| {
chosen
.iter()
.find(|c| c.composition == species.composition)
.map(|c| c.id.clone())
.expect(
"balance() only returns reactant species drawn from \
the compositions it was given",
)
})
.collect();
let candidate_set = AcceptedPrecursorSet {
precursors: matched_ids,
reaction,
};
if accepted.contains(&candidate_set) {
rejected.push(RejectedCandidate {
precursors: candidate_set.precursors,
reason_codes: vec![RejectionCode::DuplicatePlan],
explanation: "this precursor set and balanced reaction were already \
found via a different combination of candidates (a larger \
combination's extra precursor solved to a zero coefficient, \
collapsing to the same effective reactants)"
.to_string(),
});
continue;
}
accepted.push(candidate_set);
}
}
if budget_exhausted {
rejected.push(RejectedCandidate {
precursors: vec![],
reason_codes: vec![RejectionCode::SearchBudgetExhausted],
explanation: format!(
"stopped after evaluating {} precursor-set combination(s); more were possible",
combos.len()
),
});
}
Ok(PrecursorSearchOutcome { accepted, rejected })
}
fn join_symbols(elements: &[Element]) -> String {
elements
.iter()
.map(Element::symbol)
.collect::<Vec<_>>()
.join(", ")
}
fn power_set<T: Clone>(items: &[T]) -> Vec<Vec<T>> {
let mut subsets = Vec::with_capacity(1 << items.len());
for size in 0..=items.len() {
for combo in index_combinations(items.len(), size) {
subsets.push(combo.iter().map(|&i| items[i].clone()).collect());
}
}
subsets
}
fn index_combinations(n: usize, size: usize) -> Vec<Vec<usize>> {
fn recurse(
start: usize,
n: usize,
size: usize,
current: &mut Vec<usize>,
out: &mut Vec<Vec<usize>>,
) {
if current.len() == size {
out.push(current.clone());
return;
}
for i in start..n {
current.push(i);
recurse(i + 1, n, size, current, out);
current.pop();
}
}
let mut out = Vec::new();
recurse(0, n, size, &mut Vec::new(), &mut out);
out
}
fn generate_combinations(n: usize, max_size: usize, budget: usize) -> (Vec<Vec<usize>>, bool) {
let mut result = Vec::new();
let mut exhausted = false;
'sizes: for size in 1..=max_size.min(n) {
for combo in index_combinations(n, size) {
if result.len() >= budget {
exhausted = true;
break 'sizes;
}
result.push(combo);
}
}
(result, exhausted)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::SearchBudget;
fn element(symbol: &str) -> Element {
Element::new(symbol).unwrap()
}
fn composition(pairs: &[(&str, f64)]) -> Composition {
Composition::new(pairs.iter().map(|&(sym, amt)| (element(sym), amt))).unwrap()
}
fn candidate(id: &str, pairs: &[(&str, f64)]) -> PrecursorCandidate {
PrecursorCandidate {
id: PrecursorId(id.to_string()),
composition: composition(pairs),
availability: None,
}
}
fn barium_titanate_catalog() -> Vec<PrecursorCandidate> {
vec![
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("NaCl", &[("Na", 1.0), ("Cl", 1.0)]),
]
}
fn generous_budget() -> SearchBudget {
SearchBudget {
max_precursor_sets: 10_000,
max_precursors_per_plan: 3,
max_plans_returned: 100,
}
}
#[test]
fn accepts_a_set_that_covers_every_target_element() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = barium_titanate_catalog();
let outcome = search_precursor_sets(
&target,
&catalog,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let ba_ti = outcome.accepted.iter().find(|a| {
let ids: BTreeSet<&str> = a.precursors.iter().map(|p| p.0.as_str()).collect();
ids == BTreeSet::from(["BaCO3", "TiO2"])
});
assert!(
ba_ti.is_some(),
"BaCO3 + TiO2 must be accepted: {:?}",
outcome.accepted
);
}
#[test]
fn rejects_sets_with_unremovable_extra_elements() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = vec![
candidate("SrCO3", &[("Sr", 1.0), ("C", 1.0), ("O", 3.0)]),
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
];
let outcome = search_precursor_sets(
&target,
&catalog,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let bad_combo = outcome.rejected.iter().find(|r| {
let ids: BTreeSet<&str> = r.precursors.iter().map(|p| p.0.as_str()).collect();
ids == BTreeSet::from(["SrCO3", "TiO2", "BaO"])
});
assert_eq!(
bad_combo.map(|r| r.reason_codes.clone()),
Some(vec![RejectionCode::UnsupportedByproductRequired])
);
}
#[test]
fn never_generates_a_combination_larger_than_the_configured_maximum() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = barium_titanate_catalog();
let budget = SearchBudget {
max_precursors_per_plan: 2,
..generous_budget()
};
let outcome =
search_precursor_sets(&target, &catalog, &PlanningConstraints::default(), &budget)
.unwrap();
for a in &outcome.accepted {
assert!(a.precursors.len() <= 2);
}
for r in &outcome.rejected {
assert!(r.precursors.len() <= 2);
}
}
#[test]
fn rejects_combinations_containing_a_forbidden_element() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = barium_titanate_catalog();
let mut constraints = PlanningConstraints::default();
constraints.forbidden_elements.insert(element("C"));
let outcome =
search_precursor_sets(&target, &catalog, &constraints, &generous_budget()).unwrap();
assert!(
outcome
.accepted
.iter()
.all(|a| !a.precursors.iter().any(|p| p.0 == "BaCO3")),
"no accepted set may use BaCO3 once C is forbidden"
);
let forbidden_rejection = outcome
.rejected
.iter()
.find(|r| r.precursors.iter().any(|p| p.0 == "BaCO3"))
.expect("BaCO3-containing combinations must be rejected, not silently dropped");
assert_eq!(
forbidden_rejection.reason_codes,
vec![RejectionCode::ForbiddenElementPresent]
);
}
#[test]
fn duplicate_catalog_entries_do_not_duplicate_results() {
let target = composition(&[("Ba", 1.0), ("O", 1.0)]);
let raw = vec![
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
];
let catalog = InMemoryPrecursorCatalog::new(raw);
let candidates = catalog
.candidates_for(&target, &PlanningConstraints::default())
.unwrap();
assert_eq!(
candidates.len(),
1,
"duplicate PrecursorId entries must collapse to one"
);
let outcome = search_precursor_sets(
&target,
&candidates,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let single_bao_accepts = outcome
.accepted
.iter()
.filter(|a| a.precursors == vec![PrecursorId("BaO".to_string())])
.count();
assert_eq!(single_bao_accepts, 1);
}
#[test]
fn accepted_precursor_ids_stay_aligned_with_reaction_reactants() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = vec![
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
];
let outcome = search_precursor_sets(
&target,
&catalog,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
for accepted in &outcome.accepted {
assert_eq!(
accepted.precursors.len(),
accepted.reaction.reactants.len(),
"precursors and reactants must be the same length: {accepted:?}"
);
for (id, species) in accepted.precursors.iter().zip(&accepted.reaction.reactants) {
let candidate = catalog.iter().find(|c| &c.id == id).unwrap();
assert_eq!(
candidate.composition, species.composition,
"precursor id {id} must match its reactant composition"
);
}
}
assert!(
!outcome.accepted.is_empty(),
"fixture must actually exercise the search, not vacuously pass"
);
}
#[test]
fn a_redundant_larger_combination_is_rejected_as_a_duplicate_not_double_accepted() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = vec![
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]),
];
let outcome = search_precursor_sets(
&target,
&catalog,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let expected_ids = BTreeSet::from([
PrecursorId("BaO".to_string()),
PrecursorId("TiO2".to_string()),
]);
let occurrences = outcome
.accepted
.iter()
.filter(|a| a.precursors.iter().cloned().collect::<BTreeSet<_>>() == expected_ids)
.count();
assert_eq!(
occurrences, 1,
"the {{BaO, TiO2}} precursor set must be accepted exactly once: {:?}",
outcome.accepted
);
assert!(
outcome.rejected.iter().any(|r| {
r.reason_codes == vec![RejectionCode::DuplicatePlan]
&& r.precursors.contains(&PrecursorId("BaO".to_string()))
&& r.precursors.contains(&PrecursorId("TiO2".to_string()))
}),
"the redundant 3-candidate collapse must be explained as DuplicatePlan, not \
silently dropped or silently double-accepted: {:?}",
outcome.rejected
);
}
#[test]
fn result_is_independent_of_catalog_insertion_order() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let mut shuffled = barium_titanate_catalog();
shuffled.reverse();
let a = search_precursor_sets(
&target,
&InMemoryPrecursorCatalog::new(barium_titanate_catalog())
.candidates_for(&target, &PlanningConstraints::default())
.unwrap(),
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let b = search_precursor_sets(
&target,
&InMemoryPrecursorCatalog::new(shuffled)
.candidates_for(&target, &PlanningConstraints::default())
.unwrap(),
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let ids_a: BTreeSet<Vec<String>> = a
.accepted
.iter()
.map(|s| s.precursors.iter().map(|p| p.0.clone()).collect())
.collect();
let ids_b: BTreeSet<Vec<String>> = b
.accepted
.iter()
.map(|s| s.precursors.iter().map(|p| p.0.clone()).collect())
.collect();
assert_eq!(ids_a, ids_b);
}
#[test]
fn budget_exhaustion_is_reported_distinctly_from_no_candidates() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = barium_titanate_catalog();
let tiny_budget = SearchBudget {
max_precursor_sets: 1,
max_precursors_per_plan: 3,
max_plans_returned: 100,
};
let outcome = search_precursor_sets(
&target,
&catalog,
&PlanningConstraints::default(),
&tiny_budget,
)
.unwrap();
let exhaustion = outcome
.rejected
.iter()
.find(|r| r.reason_codes == vec![RejectionCode::SearchBudgetExhausted]);
assert!(
exhaustion.is_some(),
"must report budget exhaustion: {:?}",
outcome.rejected
);
assert!(exhaustion.unwrap().precursors.is_empty());
}
#[test]
fn missing_availability_metadata_does_not_block_acceptance() {
let target = composition(&[("Ba", 1.0), ("O", 1.0)]);
let with_metadata = vec![PrecursorCandidate {
id: PrecursorId("BaO".to_string()),
composition: composition(&[("Ba", 1.0), ("O", 1.0)]),
availability: Some(AvailabilityMetadata {
source: "curated_fixture".to_string(),
}),
}];
let without_metadata = vec![candidate("BaO", &[("Ba", 1.0), ("O", 1.0)])];
let a = search_precursor_sets(
&target,
&with_metadata,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
let b = search_precursor_sets(
&target,
&without_metadata,
&PlanningConstraints::default(),
&generous_budget(),
)
.unwrap();
assert_eq!(a.accepted.len(), 1);
assert_eq!(b.accepted.len(), 1);
assert_eq!(a.accepted[0].reaction, b.accepted[0].reaction);
}
#[test]
fn in_memory_catalog_scopes_to_target_relevant_candidates_and_ignores_insertion_order() {
let target = composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)]);
let catalog = InMemoryPrecursorCatalog::new(barium_titanate_catalog());
let result = catalog
.candidates_for(&target, &PlanningConstraints::default())
.unwrap();
let ids: Vec<&str> = result.iter().map(|c| c.id.0.as_str()).collect();
assert!(
!ids.contains(&"NaCl"),
"NaCl shares no element with Ba-Ti-O and must be scoped out"
);
assert!(ids.contains(&"BaCO3") && ids.contains(&"BaO") && ids.contains(&"TiO2"));
let mut sorted = ids.clone();
sorted.sort();
assert_eq!(ids, sorted);
}
}