vegas_lattice/
lattice.rs

1//! Lattice data structure
2
3use super::util::Axis;
4use crate::alloy::Alloy;
5use crate::error::{Result, VegasLatticeError};
6use crate::mask::Mask;
7use crate::site::Site;
8use crate::vertex::Vertex;
9use rand::Rng;
10use serde::{Deserialize, Serialize};
11use std::iter::repeat;
12use std::str::FromStr;
13
14/// A lattice is a collection of sites and vertices.
15///
16/// For now it only supports rectangular lattices. This is Orthorombic, Tetragonal and Cubic
17/// Bravais lattices. We assume the lattice vectors are aligned with the cartesian axes. While you
18/// can choose the lattice parameters _a_, _b_, and _c_ to be different.
19#[derive(Clone, Debug, Serialize, Deserialize)]
20pub struct Lattice {
21    size: (f64, f64, f64),
22    sites: Vec<Site>,
23    vertices: Vec<Vertex>,
24}
25
26impl Lattice {
27    /// Create a new lattice with the given size
28    pub fn try_new(size: (f64, f64, f64)) -> Result<Self> {
29        if size.0 < 0.0 || size.1 < 0.0 || size.2 < 0.0 {
30            return Err(VegasLatticeError::NegativeSize);
31        }
32        Ok(Lattice {
33            size,
34            sites: Vec::new(),
35            vertices: Vec::new(),
36        })
37    }
38
39    /// Create a simple cubic lattice with the given size _a_
40    pub fn sc(a: f64) -> Self {
41        let sites = vec![Site::new("A")];
42        let vertices = vec![
43            Vertex::new(0, 0, (1, 0, 0)),
44            Vertex::new(0, 0, (0, 1, 0)),
45            Vertex::new(0, 0, (0, 0, 1)),
46        ];
47        Lattice {
48            size: (a, a, a),
49            sites,
50            vertices,
51        }
52    }
53
54    /// Create a body centered cubic lattice with the given size _a_
55    pub fn bcc(a: f64) -> Self {
56        let sites = vec![
57            Site::new("A"),
58            Site::new("B").with_position((0.5 * a, 0.5 * a, 0.5 * a)),
59        ];
60        let vertices = vec![
61            Vertex::new(0, 1, (0, 0, 0)),
62            Vertex::new(0, 1, (0, -1, 0)),
63            Vertex::new(0, 1, (-1, 0, 0)),
64            Vertex::new(0, 1, (-1, -1, 0)),
65            Vertex::new(0, 1, (0, 0, -1)),
66            Vertex::new(0, 1, (0, -1, -1)),
67            Vertex::new(0, 1, (-1, 0, -1)),
68            Vertex::new(0, 1, (-1, -1, -1)),
69        ];
70        Lattice {
71            size: (a, a, a),
72            sites,
73            vertices,
74        }
75    }
76
77    /// Create a face centered cubic lattice with lattice parameter _a_
78    pub fn fcc(a: f64) -> Self {
79        let sites = vec![
80            Site::new("A"),
81            Site::new("B").with_position((0.5 * a, 0.5 * a, 0.0)),
82            Site::new("C").with_position((0.5 * a, 0.0, 0.5 * a)),
83            Site::new("D").with_position((0.0, 0.5 * a, 0.5 * a)),
84        ];
85        let vertices = vec![
86            // xy plane
87            Vertex::new(0, 1, (0, 0, 0)),
88            Vertex::new(0, 1, (-1, 0, 0)),
89            Vertex::new(0, 1, (-1, -1, 0)),
90            Vertex::new(0, 1, (0, -1, 0)),
91            // xz plane
92            Vertex::new(0, 2, (0, 0, 0)),
93            Vertex::new(0, 2, (-1, 0, 0)),
94            Vertex::new(0, 2, (-1, 0, -1)),
95            Vertex::new(0, 2, (0, 0, -1)),
96            // yz plane
97            Vertex::new(0, 3, (0, 0, 0)),
98            Vertex::new(0, 3, (0, -1, 0)),
99            Vertex::new(0, 3, (0, -1, -1)),
100            Vertex::new(0, 3, (0, 0, -1)),
101        ];
102        Lattice {
103            size: (a, a, a),
104            sites,
105            vertices,
106        }
107    }
108
109    /// Get the size of the lattice
110    pub fn size(&self) -> (f64, f64, f64) {
111        self.size
112    }
113
114    /// Get the sites of the lattice
115    pub fn sites(&self) -> &[Site] {
116        &self.sites
117    }
118
119    /// Get the vertices of the lattice
120    pub fn vertices(&self) -> &[Vertex] {
121        &self.vertices
122    }
123
124    /// Changes the size of the lattice
125    pub fn try_with_size(mut self, size: (f64, f64, f64)) -> Result<Self> {
126        self.size = size;
127        self.validate()
128    }
129
130    /// Changes the sites of the lattice
131    pub fn try_with_sites(mut self, sites: Vec<Site>) -> Result<Self> {
132        self.sites = sites;
133        self.validate()
134    }
135
136    /// Changes the vertices of the lattice
137    pub fn try_with_vertices(mut self, vertices: Vec<Vertex>) -> Result<Self> {
138        self.vertices = vertices;
139        self.validate()
140    }
141
142    fn are_vertices_consistent(&self) -> bool {
143        self.vertices
144            .iter()
145            .map(|vertex| vertex.source())
146            .chain(self.vertices.iter().map(|vertex| vertex.target()))
147            .all(|id| id < self.sites.len())
148    }
149
150    /// Validates the lattice
151    fn validate(self) -> Result<Self> {
152        if !self.are_vertices_consistent() {
153            return Err(VegasLatticeError::InconsistentVertices);
154        }
155        if self.size.0 < 0.0 || self.size.1 < 0.0 || self.size.2 < 0.0 {
156            return Err(VegasLatticeError::NegativeSize);
157        }
158        Ok(self)
159    }
160
161    /// Drops all the vertices that are periodic along the given axis
162    fn drop_along(mut self, axis: Axis) -> Self {
163        self.vertices.retain(|v| {
164            let delta = v.delta();
165            match axis {
166                Axis::X => delta.0 == 0,
167                Axis::Y => delta.1 == 0,
168                Axis::Z => delta.2 == 0,
169            }
170        });
171        self
172    }
173
174    /// Drop periodic boundary conditions along the x axis
175    pub fn drop_x(self) -> Self {
176        self.drop_along(Axis::X)
177    }
178
179    /// Drop periodic boundary conditions along the y axis
180    pub fn drop_y(self) -> Self {
181        self.drop_along(Axis::Y)
182    }
183
184    /// Drop periodic boundary conditions along the z axis
185    pub fn drop_z(self) -> Self {
186        self.drop_along(Axis::Z)
187    }
188
189    /// Drop periodic boundary conditions along all axes
190    pub fn drop_all(self) -> Self {
191        self.drop_x().drop_y().drop_z()
192    }
193
194    #[inline]
195    fn size_along(&self, axis: Axis) -> f64 {
196        match axis {
197            Axis::X => self.size.0,
198            Axis::Y => self.size.1,
199            Axis::Z => self.size.2,
200        }
201    }
202
203    /// Expands the lattice along the given axis
204    fn expand_along(mut self, axis: Axis, amount: usize) -> Self {
205        let size = self.size_along(axis);
206        let n_sites = self.sites.len();
207        let n_vertices = self.vertices.len();
208
209        self.sites = (0..amount)
210            .flat_map(|i| repeat(i).take(n_sites))
211            .zip(self.sites().iter().cycle())
212            .map(|(index, site)| match axis {
213                Axis::X => site.clone().move_x((index as f64) * size),
214                Axis::Y => site.clone().move_y((index as f64) * size),
215                Axis::Z => site.clone().move_z((index as f64) * size),
216            })
217            .collect();
218
219        self.vertices = (0..amount)
220            .flat_map(|i| repeat(i).take(n_vertices))
221            .zip(self.vertices.iter().cycle())
222            .map(|(index, vertex)| match axis {
223                Axis::X => vertex.clone().move_x(index, n_sites, amount),
224                Axis::Y => vertex.clone().move_y(index, n_sites, amount),
225                Axis::Z => vertex.clone().move_z(index, n_sites, amount),
226            })
227            .collect();
228
229        match axis {
230            Axis::X => self.size.0 *= amount as f64,
231            Axis::Y => self.size.1 *= amount as f64,
232            Axis::Z => self.size.2 *= amount as f64,
233        }
234
235        self
236    }
237
238    /// Expand lattice along the x axis
239    pub fn expand_x(self, amount: usize) -> Self {
240        self.expand_along(Axis::X, amount)
241    }
242
243    /// Expand lattice along the y axis
244    pub fn expand_y(self, amount: usize) -> Self {
245        self.expand_along(Axis::Y, amount)
246    }
247
248    /// Expand lattice along the z axis
249    pub fn expand_z(self, amount: usize) -> Self {
250        self.expand_along(Axis::Z, amount)
251    }
252
253    /// Expand lattice by the same ammount along all axes
254    pub fn expand_all(self, amount: usize) -> Self {
255        self.expand_x(amount).expand_y(amount).expand_z(amount)
256    }
257
258    /// Expand lattice by the given ammount along all axes
259    pub fn expand(self, x: usize, y: usize, z: usize) -> Self {
260        self.expand_x(x).expand_y(y).expand_z(z)
261    }
262
263    /// Removes sites from the lattice according to the given mask
264    ///
265    /// TODO: This only removes points in the xy plane, and it should be generalized
266    pub fn apply_mask<R: Rng + ?Sized>(mut self, mask: Mask, rng: &mut R) -> Self {
267        let site_mask: Vec<_> = self
268            .sites
269            .iter()
270            .map(|s| {
271                let (x, y, _) = s.position();
272                mask.keep(x, y, rng)
273            })
274            .collect();
275        let mut counter = 0;
276        let new_indices: Vec<_> = (0..self.sites.len())
277            .map(|i| {
278                if site_mask[i] {
279                    counter += 1;
280                    counter - 1
281                } else {
282                    i
283                }
284            })
285            .collect();
286        self.sites = self
287            .sites
288            .into_iter()
289            .enumerate()
290            .filter(|&(i, ref _s)| site_mask[i])
291            .map(|(_i, s)| s)
292            .collect();
293        self.vertices = self
294            .vertices
295            .into_iter()
296            .filter(|v| site_mask[v.source()] && site_mask[v.target()])
297            .map(|v| v.reindex(&new_indices))
298            .collect();
299        self
300    }
301
302    /// Replaces the sites labeled as `source` with sites in the `target` alloy
303    pub fn alloy_sites<R: Rng + ?Sized>(
304        mut self,
305        source: &str,
306        target: Alloy,
307        rng: &mut R,
308    ) -> Self {
309        self.sites = self
310            .sites
311            .into_iter()
312            .map(|site| {
313                if site.kind() != source {
314                    site
315                } else {
316                    site.with_kind(target.pick(rng))
317                }
318            })
319            .collect();
320        self
321    }
322}
323
324impl FromStr for Lattice {
325    type Err = VegasLatticeError;
326    fn from_str(source: &str) -> Result<Lattice> {
327        let lattice: Lattice = serde_json::from_str(source)?;
328        lattice.validate()
329    }
330}
331
332#[cfg(test)]
333mod test {
334    use crate::{Lattice, Site, Vertex};
335
336    #[test]
337    fn drop_example() {
338        let lattice = Lattice::sc(1.0);
339        let lattice = lattice.drop_x();
340        assert!(lattice.vertices().len() == 2);
341    }
342
343    #[test]
344    fn drop_example_actually_dropping() {
345        let lattice = Lattice::sc(1.0);
346        let lattice = lattice.drop_all();
347        assert!(lattice.vertices().is_empty());
348    }
349
350    #[test]
351    fn single_lattice_expansion_1d() {
352        let lattice = Lattice::sc(1.0);
353        let output = lattice.expand_x(2);
354        assert_eq!(output.sites.len(), 2);
355        assert!((output.sites[1].position().0 - 1.0).abs() < 1e-10);
356    }
357
358    #[test]
359    fn double_lattice_expansion_1d() {
360        let lattice = Lattice::sc(1.0);
361        let lattice = lattice.expand_x(2);
362        let output = lattice.expand_x(2);
363        assert_eq!(output.sites.len(), 4);
364        assert!((output.sites[1].position().0 - 1.0).abs() < 1e-10);
365        assert!((output.sites[2].position().0 - 2.0).abs() < 1e-10);
366        assert!((output.sites[3].position().0 - 3.0).abs() < 1e-10);
367    }
368
369    #[test]
370    fn single_lattice_expansion_1d_vertices() {
371        let lattice = Lattice::sc(1.0)
372            .try_with_vertices(vec![Vertex::new(0, 0, (1, 0, 0))])
373            .unwrap();
374        let output = lattice.expand_x(2);
375        assert_eq!(output.vertices.len(), 2);
376        assert_eq!(output.vertices[0].source(), 0);
377        assert_eq!(output.vertices[0].target(), 1);
378        assert_eq!(output.vertices[0].delta().0, 0);
379        assert_eq!(output.vertices[1].source(), 1);
380        assert_eq!(output.vertices[1].target(), 0);
381        assert_eq!(output.vertices[1].delta().0, 1);
382    }
383
384    #[test]
385    fn single_lattice_expansion_1d_negative_vertices() {
386        let lattice = Lattice::sc(1.0)
387            .try_with_vertices(vec![Vertex::new(0, 0, (-1, 0, 0))])
388            .unwrap();
389        let output = lattice.expand_x(2);
390        assert_eq!(output.vertices.len(), 2);
391        assert_eq!(output.vertices[0].source(), 0);
392        assert_eq!(output.vertices[0].target(), 1);
393        assert_eq!(output.vertices[0].delta().0, -1);
394        assert_eq!(output.vertices[1].source(), 1);
395        assert_eq!(output.vertices[1].target(), 0);
396        assert_eq!(output.vertices[1].delta().0, 0);
397    }
398
399    #[test]
400    fn test_sc_lattice() {
401        let lattice = Lattice::sc(1.0);
402        assert_eq!(lattice.size(), (1.0, 1.0, 1.0));
403        assert_eq!(lattice.sites().len(), 1);
404        assert_eq!(lattice.vertices().len(), 3)
405    }
406
407    #[test]
408    fn test_bcc_lattice() {
409        let lattice = Lattice::bcc(1.0);
410        assert_eq!(lattice.size(), (1.0, 1.0, 1.0));
411        assert_eq!(lattice.sites().len(), 2);
412        assert_eq!(lattice.vertices().len(), 8)
413    }
414
415    #[test]
416    fn test_fcc_lattice() {
417        let lattice = Lattice::fcc(1.0);
418        assert_eq!(lattice.size(), (1.0, 1.0, 1.0));
419        assert_eq!(lattice.sites().len(), 4);
420        assert_eq!(lattice.vertices().len(), 12)
421    }
422
423    #[test]
424    fn test_wit_size() {
425        let lattice = Lattice::sc(1.0).try_with_size((2.0, 2.0, 2.0)).unwrap();
426        assert_eq!(lattice.size(), (2.0, 2.0, 2.0));
427    }
428
429    #[test]
430    fn test_with_sites() {
431        let lattice = Lattice::sc(1.0)
432            .try_with_sites(vec![Site::new("Fe"), Site::new("Ni")])
433            .unwrap();
434        assert_eq!(lattice.sites().len(), 2);
435    }
436
437    #[test]
438    fn test_with_vertices() {
439        let lattice = Lattice::sc(1.0)
440            .try_with_vertices(vec![
441                Vertex::new(0, 0, (1, 0, 0)),
442                Vertex::new(0, 0, (0, 1, 0)),
443            ])
444            .unwrap();
445        assert_eq!(lattice.vertices().len(), 2);
446    }
447
448    #[test]
449    fn test_lattice_with_inconsistent_vertices() {
450        // The vertex target is not in the list of sites
451        let result = Lattice::sc(1.0).try_with_vertices(vec![Vertex::new(0, 1, (1, 0, 0))]);
452        assert!(result.is_err());
453    }
454
455    #[test]
456    fn test_lattice_with_negative_size() {
457        let result = Lattice::try_new((-1.0, 1.0, 1.0));
458        assert!(result.is_err());
459    }
460
461    #[test]
462    fn test_lattice_can_be_read_from_string() {
463        let lattice = r#"{
464            "size": [1.0, 1.0, 1.0],
465            "sites": [
466                {
467                    "kind": "Fe",
468                    "position": [0.0, 0.0, 0.0]
469                }
470            ],
471            "vertices": [
472                {
473                    "source": 0,
474                    "target": 0,
475                    "delta": [0, 0, 1]
476                }
477            ]
478        }"#;
479        let lattice: Lattice = lattice.parse().unwrap();
480        assert_eq!(lattice.size(), (1.0, 1.0, 1.0));
481        assert_eq!(lattice.sites().len(), 1);
482        assert_eq!(lattice.vertices().len(), 1);
483    }
484}
vegas_lattice/lattice.rs

vegas_lattice/
lattice.rs
