use crate::composition::{Composition, Element};
use crate::error::{ProviderError, require_finite};
use crate::precursor::{InMemoryPrecursorCatalog, PrecursorCandidate, PrecursorId};
use crate::provider::{CandidateGenerator, PrecursorCatalog};
use crate::target::PlanningConstraints;
use std::collections::{BTreeMap, BTreeSet};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize))]
pub struct GeneratorId(pub &'static str);
impl std::fmt::Display for GeneratorId {
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))]
pub struct GeneratedCandidate {
pub candidate: PrecursorCandidate,
pub generator: GeneratorId,
pub rank: usize,
}
pub struct CatalogExactGenerator {
catalog: InMemoryPrecursorCatalog,
}
impl CatalogExactGenerator {
pub fn new(catalog: InMemoryPrecursorCatalog) -> Self {
Self { catalog }
}
}
impl CandidateGenerator for CatalogExactGenerator {
fn id(&self) -> GeneratorId {
GeneratorId("catalog-exact")
}
fn generate(
&self,
target: &Composition,
constraints: &PlanningConstraints,
) -> std::result::Result<Vec<GeneratedCandidate>, ProviderError> {
let candidates = self.catalog.candidates_for(target, constraints)?;
Ok(candidates
.into_iter()
.enumerate()
.map(|(rank, candidate)| GeneratedCandidate {
candidate,
generator: self.id(),
rank,
})
.collect())
}
}
pub struct FrequencyPriorGenerator {
entries: Vec<(PrecursorCandidate, u64)>,
}
impl FrequencyPriorGenerator {
pub fn new(mut entries: Vec<(PrecursorCandidate, u64)>) -> Self {
entries.sort_by(|(a, a_freq), (b, b_freq)| {
b_freq.cmp(a_freq).then_with(|| a.id.0.cmp(&b.id.0))
});
Self { entries }
}
}
impl CandidateGenerator for FrequencyPriorGenerator {
fn id(&self) -> GeneratorId {
GeneratorId("frequency-prior")
}
fn generate(
&self,
target: &Composition,
_constraints: &PlanningConstraints,
) -> std::result::Result<Vec<GeneratedCandidate>, ProviderError> {
let target_elements: BTreeSet<Element> = target.elements().collect();
Ok(self
.entries
.iter()
.filter(|(candidate, _frequency)| {
candidate
.composition
.elements()
.any(|e| target_elements.contains(&e))
})
.enumerate()
.map(|(rank, (candidate, _frequency))| GeneratedCandidate {
candidate: candidate.clone(),
generator: self.id(),
rank,
})
.collect())
}
}
pub struct ThermodynamicStabilityGenerator {
entries: Vec<(PrecursorCandidate, f64)>,
}
impl ThermodynamicStabilityGenerator {
pub fn new(mut entries: Vec<(PrecursorCandidate, f64)>) -> crate::error::Result<Self> {
for (_candidate, formation_energy) in &entries {
require_finite("formation_enthalpy_ev_per_atom", *formation_energy)?;
}
entries.sort_by(|(a, a_energy), (b, b_energy)| {
a_energy
.total_cmp(b_energy)
.then_with(|| a.id.0.cmp(&b.id.0))
});
Ok(Self { entries })
}
}
impl CandidateGenerator for ThermodynamicStabilityGenerator {
fn id(&self) -> GeneratorId {
GeneratorId("thermodynamic-stability")
}
fn generate(
&self,
target: &Composition,
_constraints: &PlanningConstraints,
) -> std::result::Result<Vec<GeneratedCandidate>, ProviderError> {
let target_elements: BTreeSet<Element> = target.elements().collect();
Ok(self
.entries
.iter()
.filter(|(candidate, _formation_energy)| {
candidate
.composition
.elements()
.any(|e| target_elements.contains(&e))
})
.enumerate()
.map(
|(rank, (candidate, _formation_energy))| GeneratedCandidate {
candidate: candidate.clone(),
generator: self.id(),
rank,
},
)
.collect())
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct EnsembleOutput {
pub candidates: Vec<PrecursorCandidate>,
pub provenance: BTreeMap<PrecursorId, Vec<GeneratedCandidate>>,
pub generator_errors: Vec<(GeneratorId, ProviderError)>,
}
pub struct CandidateGeneratorEnsemble {
generators: Vec<Box<dyn CandidateGenerator>>,
}
impl CandidateGeneratorEnsemble {
pub fn new(generators: Vec<Box<dyn CandidateGenerator>>) -> Self {
Self { generators }
}
pub fn generate_with_provenance(
&self,
target: &Composition,
constraints: &PlanningConstraints,
) -> EnsembleOutput {
let mut best: BTreeMap<PrecursorId, (usize, PrecursorCandidate)> = BTreeMap::new();
let mut provenance: BTreeMap<PrecursorId, Vec<GeneratedCandidate>> = BTreeMap::new();
let mut generator_errors = Vec::new();
for generator in &self.generators {
match generator.generate(target, constraints) {
Ok(generated) => {
for gc in generated {
let id = gc.candidate.id.clone();
best.entry(id.clone())
.and_modify(|(rank, _payload)| {
if gc.rank < *rank {
*rank = gc.rank;
}
})
.or_insert_with(|| (gc.rank, gc.candidate.clone()));
provenance.entry(id).or_default().push(gc);
}
}
Err(err) => generator_errors.push((generator.id(), err)),
}
}
let mut fused: Vec<(usize, PrecursorCandidate)> = best.into_values().collect();
fused.sort_by(|(rank_a, candidate_a), (rank_b, candidate_b)| {
rank_a
.cmp(rank_b)
.then_with(|| candidate_a.id.0.cmp(&candidate_b.id.0))
});
EnsembleOutput {
candidates: fused
.into_iter()
.map(|(_rank, candidate)| candidate)
.collect(),
provenance,
generator_errors,
}
}
}
impl PrecursorCatalog for CandidateGeneratorEnsemble {
fn candidates_for(
&self,
target: &Composition,
constraints: &PlanningConstraints,
) -> std::result::Result<Vec<PrecursorCandidate>, ProviderError> {
Ok(self
.generate_with_provenance(target, constraints)
.candidates)
}
}
#[cfg(test)]
mod tests {
use super::*;
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 no_constraints() -> PlanningConstraints {
PlanningConstraints::default()
}
fn barium_titanate_target() -> Composition {
composition(&[("Ba", 1.0), ("Ti", 1.0), ("O", 3.0)])
}
struct AlwaysFailsGenerator;
impl CandidateGenerator for AlwaysFailsGenerator {
fn id(&self) -> GeneratorId {
GeneratorId("always-fails")
}
fn generate(
&self,
_target: &Composition,
_constraints: &PlanningConstraints,
) -> std::result::Result<Vec<GeneratedCandidate>, ProviderError> {
Err(ProviderError::Unavailable("test failure".to_string()))
}
}
#[test]
fn catalog_exact_generator_delegates_and_stamps_rank_by_output_position() {
let catalog = InMemoryPrecursorCatalog::new(vec![
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
candidate("NaCl", &[("Na", 1.0), ("Cl", 1.0)]),
]);
let generator = CatalogExactGenerator::new(catalog);
let generated = generator
.generate(&barium_titanate_target(), &no_constraints())
.unwrap();
let ids: Vec<&str> = generated
.iter()
.map(|gc| gc.candidate.id.0.as_str())
.collect();
assert_eq!(ids, vec!["BaCO3", "TiO2"]);
assert!(
generated
.iter()
.all(|gc| gc.generator == GeneratorId("catalog-exact"))
);
assert_eq!(generated[0].rank, 0);
assert_eq!(generated[1].rank, 1);
}
#[test]
fn frequency_prior_generator_filters_by_element_overlap_and_preserves_frequency_order() {
let generator = FrequencyPriorGenerator::new(vec![
(candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]), 5),
(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
50,
),
(candidate("NaCl", &[("Na", 1.0), ("Cl", 1.0)]), 1000),
]);
let generated = generator
.generate(&barium_titanate_target(), &no_constraints())
.unwrap();
let ids: Vec<&str> = generated
.iter()
.map(|gc| gc.candidate.id.0.as_str())
.collect();
assert_eq!(
ids,
vec!["BaCO3", "TiO2"],
"higher frequency (50) must rank first"
);
assert!(
generated
.iter()
.all(|gc| gc.generator == GeneratorId("frequency-prior"))
);
assert_eq!(generated[0].rank, 0);
assert_eq!(generated[1].rank, 1);
}
#[test]
fn thermodynamic_stability_generator_rejects_a_non_finite_formation_energy() {
assert!(
ThermodynamicStabilityGenerator::new(vec![(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
f64::NAN,
)])
.is_err()
);
assert!(
ThermodynamicStabilityGenerator::new(vec![(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
f64::INFINITY,
)])
.is_err()
);
assert!(
ThermodynamicStabilityGenerator::new(vec![(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
-3.5,
)])
.is_ok()
);
}
#[test]
fn thermodynamic_stability_generator_filters_by_element_overlap_and_ranks_most_stable_first() {
let generator = ThermodynamicStabilityGenerator::new(vec![
(candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]), -3.0),
(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
-3.5,
),
(candidate("NaCl", &[("Na", 1.0), ("Cl", 1.0)]), -10.0),
])
.unwrap();
let generated = generator
.generate(&barium_titanate_target(), &no_constraints())
.unwrap();
let ids: Vec<&str> = generated
.iter()
.map(|gc| gc.candidate.id.0.as_str())
.collect();
assert_eq!(
ids,
vec!["BaCO3", "TiO2"],
"more negative formation energy (-3.5) must rank first"
);
assert!(
generated
.iter()
.all(|gc| gc.generator == GeneratorId("thermodynamic-stability"))
);
assert_eq!(generated[0].rank, 0);
assert_eq!(generated[1].rank, 1);
}
#[test]
fn ensemble_min_rank_fuses_candidates_proposed_by_either_generator() {
let catalog_exact = CatalogExactGenerator::new(InMemoryPrecursorCatalog::new(vec![
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
]));
let frequency_prior = FrequencyPriorGenerator::new(vec![
(candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]), 100),
(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
1,
),
]);
let ensemble = CandidateGeneratorEnsemble::new(vec![
Box::new(catalog_exact),
Box::new(frequency_prior),
]);
let output =
ensemble.generate_with_provenance(&barium_titanate_target(), &no_constraints());
let ids: Vec<&str> = output.candidates.iter().map(|c| c.id.0.as_str()).collect();
assert_eq!(ids, vec!["BaCO3", "TiO2"]);
assert!(output.generator_errors.is_empty());
assert_eq!(
output.provenance[&PrecursorId("BaCO3".to_string())].len(),
2
);
assert_eq!(output.provenance[&PrecursorId("TiO2".to_string())].len(), 2);
}
#[test]
fn ensemble_fuses_a_third_generator_including_a_candidate_only_it_proposed() {
let catalog_exact = CatalogExactGenerator::new(InMemoryPrecursorCatalog::new(vec![
candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]),
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
]));
let frequency_prior = FrequencyPriorGenerator::new(vec![
(candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]), 100),
(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
1,
),
]);
let thermodynamic_stability = ThermodynamicStabilityGenerator::new(vec![
(candidate("TiO2", &[("Ti", 1.0), ("O", 2.0)]), -3.0),
(
candidate("BaCO3", &[("Ba", 1.0), ("C", 1.0), ("O", 3.0)]),
-3.5,
),
(candidate("BaO", &[("Ba", 1.0), ("O", 1.0)]), -2.0),
])
.unwrap();
let ensemble = CandidateGeneratorEnsemble::new(vec![
Box::new(catalog_exact),
Box::new(frequency_prior),
Box::new(thermodynamic_stability),
]);
let output =
ensemble.generate_with_provenance(&barium_titanate_target(), &no_constraints());
let ids: std::collections::BTreeSet<&str> =
output.candidates.iter().map(|c| c.id.0.as_str()).collect();
assert_eq!(
ids,
std::collections::BTreeSet::from(["BaCO3", "TiO2", "BaO"]),
"the union of all three generators' candidates, including the one only \
thermodynamic-stability proposed"
);
assert!(output.generator_errors.is_empty());
assert_eq!(
output.provenance[&PrecursorId("BaCO3".to_string())].len(),
3,
"all three generators proposed BaCO3"
);
assert_eq!(
output.provenance[&PrecursorId("TiO2".to_string())].len(),
3,
"all three generators proposed TiO2"
);
assert_eq!(
output.provenance[&PrecursorId("BaO".to_string())].len(),
1,
"only thermodynamic-stability proposed BaO"
);
}
#[test]
fn ensemble_duplicate_id_conflict_keeps_first_generators_payload_but_records_every_proposer() {
let first = CatalogExactGenerator::new(InMemoryPrecursorCatalog::new(vec![candidate(
"BaCO3",
&[("Ba", 1.0), ("C", 1.0), ("O", 3.0)],
)]));
let second = FrequencyPriorGenerator::new(vec![(
candidate("BaCO3", &[("Ba", 2.0), ("C", 1.0), ("O", 3.0)]),
10,
)]);
let ensemble = CandidateGeneratorEnsemble::new(vec![Box::new(first), Box::new(second)]);
let output =
ensemble.generate_with_provenance(&barium_titanate_target(), &no_constraints());
assert_eq!(output.candidates.len(), 1);
assert_eq!(
output.candidates[0].composition,
composition(&[("Ba", 1.0), ("C", 1.0), ("O", 3.0)])
);
assert_eq!(
output.provenance[&PrecursorId("BaCO3".to_string())].len(),
2
);
}
#[test]
fn ensemble_records_a_failed_generators_error_and_still_returns_the_others_candidates() {
let catalog_exact =
CatalogExactGenerator::new(InMemoryPrecursorCatalog::new(vec![candidate(
"BaCO3",
&[("Ba", 1.0), ("C", 1.0), ("O", 3.0)],
)]));
let ensemble = CandidateGeneratorEnsemble::new(vec![
Box::new(catalog_exact),
Box::new(AlwaysFailsGenerator),
]);
let output =
ensemble.generate_with_provenance(&barium_titanate_target(), &no_constraints());
assert_eq!(output.candidates.len(), 1);
assert_eq!(output.candidates[0].id, PrecursorId("BaCO3".to_string()));
assert_eq!(output.generator_errors.len(), 1);
assert_eq!(output.generator_errors[0].0, GeneratorId("always-fails"));
}
#[test]
fn ensemble_as_precursor_catalog_returns_the_same_candidates_as_generate_with_provenance() {
let catalog_exact =
CatalogExactGenerator::new(InMemoryPrecursorCatalog::new(vec![candidate(
"BaCO3",
&[("Ba", 1.0), ("C", 1.0), ("O", 3.0)],
)]));
let ensemble = CandidateGeneratorEnsemble::new(vec![Box::new(catalog_exact)]);
let via_trait = PrecursorCatalog::candidates_for(
&ensemble,
&barium_titanate_target(),
&no_constraints(),
)
.unwrap();
let via_inherent =
ensemble.generate_with_provenance(&barium_titanate_target(), &no_constraints());
assert_eq!(via_trait, via_inherent.candidates);
}
}