1use crate::types::{CoordinateVector, OffsetVector};
2
3#[derive(Clone, Debug, PartialEq)]
4pub struct Site {
5 pub coordinate: CoordinateVector,
6 pub site_type: i32,
7}
8
9#[derive(Clone, Debug, PartialEq)]
10pub struct Bond {
11 pub source: usize,
12 pub target: usize,
13 pub target_offset: OffsetVector,
14 pub bond_type: i32,
15}
16
17#[derive(Clone, Debug, PartialEq)]
18pub struct Unitcell {
19 dim: usize,
20 sites: Vec<Site>,
21 bonds: Vec<Bond>,
22}
23
24impl Unitcell {
25 pub fn new(dim: usize) -> Self {
26 Self { dim, sites: Vec::new(), bonds: Vec::new() }
27 }
28
29 pub fn simple(dim: usize) -> Self {
30 let mut cell = Self::new(dim);
31 cell.add_site(CoordinateVector::zeros(dim), 0);
32 for axis in 0..dim {
33 let mut offset = OffsetVector::zeros(dim);
34 offset[axis] = 1;
35 cell.add_bond(0, 0, offset, 0);
36 }
37 cell
38 }
39
40 pub fn dimension(&self) -> usize {
41 self.dim
42 }
43
44 pub fn num_sites(&self) -> usize {
45 self.sites.len()
46 }
47
48 pub fn num_bonds(&self) -> usize {
49 self.bonds.len()
50 }
51
52 pub fn site(&self, index: usize) -> &Site {
53 &self.sites[index]
54 }
55
56 pub fn bond(&self, index: usize) -> &Bond {
57 &self.bonds[index]
58 }
59
60 pub fn max_neighbors(&self) -> usize {
61 let mut counts = vec![0usize; self.num_sites()];
62 for bond in &self.bonds {
63 counts[bond.source] += 1;
64 counts[bond.target] += 1;
65 }
66 counts.into_iter().max().unwrap_or(0)
67 }
68
69 pub fn add_site(&mut self, coordinate: CoordinateVector, site_type: i32) -> usize {
70 assert_eq!(coordinate.len(), self.dim, "site coordinate dimension mismatch");
71 for value in coordinate.iter() {
72 if *value < 0.0 || *value >= 1.0 {
73 panic!("site coordinate out of range");
74 }
75 }
76 let index = self.sites.len();
77 self.sites.push(Site { coordinate, site_type });
78 index
79 }
80
81 pub fn add_bond(&mut self, source: usize, target: usize, target_offset: OffsetVector, bond_type: i32) -> usize {
82 if source >= self.num_sites() || target >= self.num_sites() {
83 panic!("site index out of range");
84 }
85 if target_offset.len() != self.dim {
86 panic!("unitcell offset dimension mismatch");
87 }
88 let index = self.bonds.len();
89 self.bonds.push(Bond { source, target, target_offset, bond_type });
90 index
91 }
92}