castep_seeding/
root.rs

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
19/// A trait to process a folder of many cell files.
20pub 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    /// Implementation of generation of struct that impl `SeedFolder`
95    fn generate_seed_folders(&self) -> Result<Vec<impl SeedFolder + Sync>, SeedingErrors>;
96    /// Execute all actions to build seed files from `.cell` files under the root path.
97    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}