Clean up API (#53)

* Remove axis from the API

* Expand the lattice API a little bit

* Use clap's derive API

* Fix docs

* Extract that as a variable

* Remove unused extern crate.

* Increase the version

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "vegas-lattice"
-version = "0.7.1"
+version = "0.8.0"
 authors = ["Oscar David Arbeláez Echeverri <[email protected]>"]
 description = "CLI and library to work with lattices"
 documentation = "https://docs.rs/vegas-lattice"
@@ -17,5 +17,5 @@ serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 image = "0.24"
 rand = "0.8"
-clap = { version = "4.5.16", features = ["cargo"] }
+clap = { version = "4.5.16", features = ["cargo", "derive"] }
 thiserror = "1.0.63"

src/alloy.rs

@@ -1,5 +1,6 @@
 //! Let's abstract an alloy
 
+use crate::error::{Result, VegasLatticeError};
 use rand::{distributions::WeightedIndex, prelude::Distribution, Rng};
 
 /// An alloy is a collection of kinds of atoms and their ratios
@@ -11,7 +12,7 @@ use rand::{distributions::WeightedIndex, prelude::Distribution, Rng};
 /// ```rust
 /// use vegas_lattice::Alloy;
 ///
-/// let alloy = Alloy::new(vec!["Fe", "Ni"], vec![1, 2]);
+/// let alloy = Alloy::try_new(vec!["Fe", "Ni"], vec![1, 2]).unwrap();
 /// let kind = alloy.pick(&mut rand::thread_rng());
 ///
 /// assert!(kind == "Fe" || kind == "Ni");
@@ -24,17 +25,18 @@ pub struct Alloy {
 
 impl Alloy {
     /// Create a new alloy with a given list of kinds and their ratios
-    pub fn new(kinds: Vec<&str>, ratios: Vec<u32>) -> Self {
-        debug_assert!(kinds.len() == ratios.len());
-        Self {
-            kinds: kinds.into_iter().map(|s| s.to_owned()).collect(),
-            weights: WeightedIndex::new(&ratios).unwrap(),
+    pub fn try_new(kinds: Vec<&str>, ratios: Vec<u32>) -> Result<Self> {
+        if kinds.len() != ratios.len() {
+            return Err(VegasLatticeError::InvalidRatios);
         }
+        let kinds = kinds.into_iter().map(|s| s.to_owned()).collect();
+        let weights = WeightedIndex::new(&ratios)?;
+        Ok(Self { kinds, weights })
     }
 
-    pub fn from_targets(targets: Vec<(&str, u32)>) -> Self {
+    pub fn try_from_targets(targets: Vec<(&str, u32)>) -> Result<Self> {
         let (kinds, ratios): (Vec<_>, Vec<_>) = targets.into_iter().unzip();
-        Self::new(kinds, ratios)
+        Self::try_new(kinds, ratios)
     }
 
     /// Picks a kind of atom from the alloy

src/error.rs

@@ -1,5 +1,6 @@
 //! Error handling for the vegas lattice crate
 
+use rand::distributions::WeightedError;
 use serde_json::Error as SerdeError;
 use std::io::Error as IoError;
 use thiserror::Error;
@@ -17,6 +18,10 @@ pub enum VegasLatticeError {
     InconsistentVertices,
     #[error("negative size")]
     NegativeSize,
+    #[error("inconsistent weights")]
+    InconsistentWeights(#[from] WeightedError),
+    #[error("invalid ratios")]
+    InvalidRatios,
 }
 
 /// Result type for the vegas lattice crate

src/lattice.rs

@@ -1,9 +1,10 @@
 //! Lattice data structure
+
+use super::util::Axis;
 use crate::alloy::Alloy;
 use crate::error::VegasLatticeError;
 use crate::mask::Mask;
 use crate::site::Site;
-use crate::util::Axis;
 use crate::vertex::Vertex;
 use rand::thread_rng;
 use serde::{Deserialize, Serialize};
@@ -117,15 +118,6 @@ impl Lattice {
         &self.vertices
     }
 
-    /// Returns the size of the lattice along the given axis
-    pub fn size_along(&self, axis: Axis) -> f64 {
-        match axis {
-            Axis::X => self.size.0,
-            Axis::Y => self.size.1,
-            Axis::Z => self.size.2,
-        }
-    }
-
     /// Changes the size of the lattice
     pub fn with_size(mut self, size: (f64, f64, f64)) -> Self {
         self.size = size;
@@ -164,13 +156,49 @@ impl Lattice {
     }
 
     /// Drops all the vertices that are periodic along the given axis
-    pub fn drop(mut self, axis: Axis) -> Self {
-        self.vertices.retain(|v| v.delta_along(axis) == 0);
+    fn drop_along(mut self, axis: Axis) -> Self {
+        self.vertices.retain(|v| {
+            let delta = v.delta();
+            match axis {
+                Axis::X => delta.0 == 0,
+                Axis::Y => delta.1 == 0,
+                Axis::Z => delta.2 == 0,
+            }
+        });
         self
     }
 
+    /// Drop periodic boundary conditions along the x axis
+    pub fn drop_x(self) -> Self {
+        self.drop_along(Axis::X)
+    }
+
+    /// Drop periodic boundary conditions along the y axis
+    pub fn drop_y(self) -> Self {
+        self.drop_along(Axis::Y)
+    }
+
+    /// Drop periodic boundary conditions along the z axis
+    pub fn drop_z(self) -> Self {
+        self.drop_along(Axis::Z)
+    }
+
+    /// Drop periodic boundary conditions along all axes
+    pub fn drop_all(self) -> Self {
+        self.drop_x().drop_y().drop_z()
+    }
+
+    #[inline]
+    fn size_along(&self, axis: Axis) -> f64 {
+        match axis {
+            Axis::X => self.size.0,
+            Axis::Y => self.size.1,
+            Axis::Z => self.size.2,
+        }
+    }
+
     /// Expands the lattice along the given axis
-    pub fn expand_along(mut self, axis: Axis, amount: usize) -> Self {
+    fn expand_along(mut self, axis: Axis, amount: usize) -> Self {
         let size = self.size_along(axis);
         let n_sites = self.sites.len();
         let n_vertices = self.vertices.len();
@@ -184,7 +212,11 @@ impl Lattice {
         self.vertices = (0..amount)
             .flat_map(|i| repeat(i).take(n_vertices))
             .zip(self.vertices.iter().cycle())
-            .map(|(index, vertex)| vertex.clone().move_along(axis, index, n_sites, amount))
+            .map(|(index, vertex)| match axis {
+                Axis::X => vertex.clone().move_x(index, n_sites, amount),
+                Axis::Y => vertex.clone().move_y(index, n_sites, amount),
+                Axis::Z => vertex.clone().move_z(index, n_sites, amount),
+            })
             .collect();
 
         match axis {
@@ -196,6 +228,31 @@ impl Lattice {
         self
     }
 
+    /// Expand lattice along the x axis
+    pub fn expand_x(self, amount: usize) -> Self {
+        self.expand_along(Axis::X, amount)
+    }
+
+    /// Expand lattice along the y axis
+    pub fn expand_y(self, amount: usize) -> Self {
+        self.expand_along(Axis::Y, amount)
+    }
+
+    /// Expand lattice along the z axis
+    pub fn expand_z(self, amount: usize) -> Self {
+        self.expand_along(Axis::Z, amount)
+    }
+
+    /// Expand lattice by the same ammount along all axes
+    pub fn expand_all(self, amount: usize) -> Self {
+        self.expand_x(amount).expand_y(amount).expand_z(amount)
+    }
+
+    /// Expand lattice by the given ammount along all axes
+    pub fn expand(self, x: usize, y: usize, z: usize) -> Self {
+        self.expand_x(x).expand_y(y).expand_z(z)
+    }
+
     /// Removes sites from the lattice according to the given mask
     ///
     /// TODO: This only removes points in the xy plane, and it should be generalized
@@ -264,14 +321,14 @@ impl FromStr for Lattice {
 
 #[cfg(test)]
 mod test {
-    use crate::{Axis, Lattice, Site, Vertex};
+    use crate::{Lattice, Site, Vertex};
 
     #[test]
     fn drop_example() {
         let lattice = Lattice::new((1.0, 1.0, 1.0))
             .with_sites(vec![Site::new("Fe")])
             .with_vertices(vec![Vertex::new(0, 0, (0, 0, 1))]);
-        let lattice = lattice.drop(Axis::X);
+        let lattice = lattice.drop_x();
         assert!(lattice.vertices.len() == 1);
     }
 
@@ -280,23 +337,23 @@ mod test {
         let lattice = Lattice::new((1.0, 1.0, 1.0))
             .with_sites(vec![Site::new("Fe")])
             .with_vertices(vec![Vertex::new(0, 0, (0, 0, 1))]);
-        let lattice = lattice.drop(Axis::Z);
+        let lattice = lattice.drop_z();
         assert!(lattice.vertices.is_empty());
     }
 
     #[test]
     fn single_lattice_expansion_1d() {
         let lattice = Lattice::new((1.0, 1.0, 1.0)).with_sites(vec![Site::new("Fe")]);
-        let output = lattice.expand_along(Axis::X, 2);
+        let output = lattice.expand_x(2);
         assert_eq!(output.sites.len(), 2);
         assert!((output.sites[1].position().0 - 1.0).abs() < 1e-10);
     }
 
     #[test]
     fn double_lattice_expansion_1d() {
         let lattice = Lattice::new((1.0, 1.0, 1.0)).with_sites(vec![Site::new("Fe")]);
-        let lattice = lattice.expand_along(Axis::X, 2);
-        let output = lattice.expand_along(Axis::X, 2);
+        let lattice = lattice.expand_x(2);
+        let output = lattice.expand_x(2);
         assert_eq!(output.sites.len(), 4);
         assert!((output.sites[1].position().0 - 1.0).abs() < 1e-10);
         assert!((output.sites[2].position().0 - 2.0).abs() < 1e-10);
@@ -308,29 +365,29 @@ mod test {
         let lattice = Lattice::new((1.0, 1.0, 1.0))
             .with_sites(vec![Site::new("Fe")])
             .with_vertices(vec![Vertex::new(0, 0, (1, 0, 0))]);
-        let output = lattice.expand_along(Axis::X, 2);
+        let output = lattice.expand_x(2);
         assert_eq!(output.vertices.len(), 2);
         assert_eq!(output.vertices[0].source(), 0);
         assert_eq!(output.vertices[0].target(), 1);
-        assert_eq!(output.vertices[0].delta_along(Axis::X), 0);
+        assert_eq!(output.vertices[0].delta().0, 0);
         assert_eq!(output.vertices[1].source(), 1);
         assert_eq!(output.vertices[1].target(), 0);
-        assert_eq!(output.vertices[1].delta_along(Axis::X), 1);
+        assert_eq!(output.vertices[1].delta().0, 1);
     }
 
     #[test]
     fn single_lattice_expansion_1d_negative_vertices() {
         let lattice = Lattice::new((1.0, 1.0, 1.0))
             .with_sites(vec![Site::new("Fe")])
             .with_vertices(vec![Vertex::new(0, 0, (-1, 0, 0))]);
-        let output = lattice.expand_along(Axis::X, 2);
+        let output = lattice.expand_x(2);
         assert_eq!(output.vertices.len(), 2);
         assert_eq!(output.vertices[0].source(), 0);
         assert_eq!(output.vertices[0].target(), 1);
-        assert_eq!(output.vertices[0].delta_along(Axis::X), -1);
+        assert_eq!(output.vertices[0].delta().0, -1);
         assert_eq!(output.vertices[1].source(), 1);
         assert_eq!(output.vertices[1].target(), 0);
-        assert_eq!(output.vertices[1].delta_along(Axis::X), 0);
+        assert_eq!(output.vertices[1].delta().0, 0);
     }
 
     #[test]

src/lib.rs

@@ -42,9 +42,9 @@
 //! lattice and apply the following operation:
 //!
 //! ```rust
-//! use vegas_lattice::{Axis, Lattice};
+//! use vegas_lattice::Lattice;
 //!
-//! let lattice = Lattice::sc(1.0).expand_along(Axis::X, 2);
+//! let lattice = Lattice::sc(1.0).expand_x(2);
 //!
 //! assert_eq!(lattice.size(), (2.0, 1.0, 1.0));
 //! assert_eq!(lattice.sites().len(), 2);
@@ -65,9 +65,9 @@
 //! did before:
 //!
 //! ```rust
-//! use vegas_lattice::{Axis, Lattice};
+//! use vegas_lattice::Lattice;
 //!
-//! let lattice = Lattice::sc(1.0).drop(Axis::X);
+//! let lattice = Lattice::sc(1.0).drop_x();
 //! assert_eq!(lattice.size(), (1.0, 1.0, 1.0));
 //! assert_eq!(lattice.sites().len(), 1);
 //! assert_eq!(lattice.vertices().len(), 2);
@@ -81,16 +81,14 @@
 //! materials by drawing an image of them.
 //!
 //! ```rust
-//! extern crate rand;
-//!
-//! use vegas_lattice::{Axis, Lattice, Mask};
+//! use vegas_lattice::{Lattice, Mask};
 //! use rand::thread_rng;
 //! use std::path::Path;
 //!
 //! let lattice = Lattice::sc(1.0)
-//!     .expand_along(Axis::X, 2)
-//!     .expand_along(Axis::Y, 2)
-//!     .expand_along(Axis::Z, 2)
+//!     .expand_x(2)
+//!     .expand_y(2)
+//!     .expand_z(2)
 //!     .apply_mask(Mask::new(Path::new("docs/pattern.png"), 100.0).unwrap());
 //! assert_eq!(lattice.size(), (2.0, 2.0, 2.0));
 //! assert!(lattice.sites().len() <= 8);
@@ -105,17 +103,12 @@
 //! ```rust
 //! use vegas_lattice::{Alloy, Lattice};
 //!
-//! let lattice = Lattice::sc(1.0).alloy_sites("A", Alloy::new(vec!["B", "C"], vec![50, 50]));
+//! let alloy = Alloy::try_new(vec!["B", "C"], vec![50, 50]).unwrap();
+//! let lattice = Lattice::sc(1.0).alloy_sites("A", alloy);
 //! assert_eq!(lattice.sites().len(), 1);
 //! assert!(lattice.sites()[0].kind() == "B" || lattice.sites()[0].kind() == "C");
 //! ```
 
-extern crate image;
-extern crate rand;
-extern crate serde;
-extern crate serde_json;
-extern crate thiserror;
-
 pub mod error;
 pub mod io;
 
@@ -130,5 +123,5 @@ pub use alloy::Alloy;
 pub use lattice::Lattice;
 pub use mask::Mask;
 pub use site::Site;
-pub use util::{Axis, Tagged};
+pub use util::Tagged;
 pub use vertex::Vertex;