1use castep_cell_io::cell_document::{
2 sections::species_characters::{SpeciesBlock, SpeciesEntry},
3 CellDocument, CellEntries,
4};
5use castep_periodic_table::{data::ELEMENT_TABLE, element::LookupElement};
6use glob::glob;
7use indicatif::{ParallelProgressIterator, ProgressStyle};
8use rayon::prelude::*;
9use std::{
10 collections::HashSet,
11 ops::ControlFlow,
12 path::{Path, PathBuf},
13};
14
15use crate::{
16 error::SeedingErrors, seed::parse_cell_doc_from_path, CellBuilding, ParamBuilding, SeedFolder,
17};
18
19pub trait RootJobs {
21 fn root_path(&self) -> impl AsRef<Path>;
22 fn get_cell_paths(&self) -> Result<Vec<PathBuf>, SeedingErrors> {
23 let binding = self.root_path().as_ref().join("*.cell");
24 let pattern = binding.to_str().ok_or(SeedingErrors::InvalidPattern)?;
25 let cell_paths = glob(pattern)
26 .map_err(SeedingErrors::MatchingCellFiles)?
27 .map(|paths| paths.map_err(SeedingErrors::GlobError))
28 .collect::<Result<Vec<PathBuf>, SeedingErrors>>()?;
29 if cell_paths.is_empty() {
30 Err(SeedingErrors::NoMatchingResults)
31 } else {
32 Ok(cell_paths)
33 }
34 }
35 fn get_cell_entries(&self) -> Result<Vec<CellDocument>, SeedingErrors> {
36 self.get_cell_paths()?
37 .par_iter()
38 .progress_with_style(ProgressStyle::default_bar())
39 .map(parse_cell_doc_from_path)
40 .collect()
41 }
42 fn fetch_potential_files<P: AsRef<Path>>(
43 &self,
44 potentials_loc: P,
45 ) -> Result<(), SeedingErrors> {
46 let potential_files: HashSet<String> =
47 HashSet::from_iter(self.get_cell_entries()?.iter().flat_map(
48 |cell_doc| -> Vec<String> {
49 let species_pots = cell_doc.other_entries().and_then(|v| {
50 v.iter()
51 .find(|entry| matches!(entry, CellEntries::SpeciesPot(_sp)))
52 .and_then(|entry| {
53 if let CellEntries::SpeciesPot(sp) = entry {
54 Some(
55 sp.items()
56 .iter()
57 .map(|s| s.item())
58 .cloned()
59 .collect::<Vec<String>>(),
60 )
61 } else {
62 None
63 }
64 })
65 });
66 species_pots.unwrap_or_else(|| {
67 cell_doc
68 .get_elements()
69 .iter()
70 .map(|elm| ELEMENT_TABLE.get_by_symbol(*elm).potential().into())
71 .collect::<Vec<String>>()
72 })
73 },
74 ));
75 let copy_potentials = potential_files.iter().try_for_each(|pot_file| {
76 let potential_path = potentials_loc.as_ref().join(pot_file);
77 let dst_path = self.root_path().as_ref().join(pot_file);
78 if !dst_path.exists() {
79 let copy =
80 std::fs::copy(potential_path, dst_path).map_err(SeedingErrors::CopyError);
81 match copy {
82 Ok(_) => ControlFlow::Continue(()),
83 Err(e) => ControlFlow::Break(e),
84 }
85 } else {
86 ControlFlow::Continue(())
87 }
88 });
89 match copy_potentials {
90 ControlFlow::Continue(_) => Ok(()),
91 ControlFlow::Break(e) => Err(e),
92 }
93 }
94 fn generate_seed_folders(&self) -> Result<Vec<impl SeedFolder + Sync>, SeedingErrors>;
96 fn build_all<
98 L: AsRef<Path> + Sync,
99 C: CellBuilding + Sync,
100 P: ParamBuilding + std::marker::Sync,
101 >(
102 &self,
103 cell_builder: &C,
104 param_builder: &P,
105 potentials_loc: L,
106 ) -> Result<(), SeedingErrors> {
107 self.fetch_potential_files(&potentials_loc)?;
108 self.generate_seed_folders()
109 .unwrap()
110 .par_iter()
111 .progress_with_style(ProgressStyle::default_bar())
112 .try_for_each(|seed| seed.actions(cell_builder, param_builder, &potentials_loc))
113 }
114}