extern crate pyo3;
use self::pyo3::types::IntoPyDict;
use self::pyo3::{PyResult, Python};
use std::collections::HashMap;
pub use Grid::*;
use PlotOptions::{Domain, Images, Legends, Path};
pub use Markers::{Point, Line, Circle};
type Vector = Vec<f64>;
#[derive(Debug, Copy, Clone, Hash, PartialOrd, PartialEq, Eq)]
pub enum PlotOptions {
Domain,
Images,
Legends,
Path,
}
#[derive(Debug, Copy, Clone, Hash, PartialOrd, PartialEq, Eq)]
pub enum Markers {
Point,
Line,
Circle,
}
#[derive(Debug, Copy, Clone, Hash, PartialOrd, PartialEq, Eq)]
pub enum Grid {
On,
Off,
}
pub trait Plot {
fn set_domain(&mut self, x: Vec<f64>) -> &mut Self;
fn insert_image(&mut self, y: Vec<f64>) -> &mut Self;
fn set_title(&mut self, title: &str) -> &mut Self;
fn set_xlabel(&mut self, xlabel: &str) -> &mut Self;
fn set_ylabel(&mut self, ylabel: &str) -> &mut Self;
fn set_zlabel(&mut self, zlabel: &str) -> &mut Self;
fn set_legends(&mut self, legends: Vec<&str>) -> &mut Self;
fn set_path(&mut self, path: &str) -> &mut Self;
fn set_fig_size(&mut self, fig_size: (usize, usize)) -> &mut Self;
fn set_dpi(&mut self, dpi: usize) -> &mut Self;
fn grid(&mut self, grid: Grid) -> &mut Self;
fn set_marker(&mut self, styles: Vec<Markers>) -> &mut Self;
fn savefig(&self) -> PyResult<()>;
}
#[derive(Debug)]
pub struct Plot2D {
domain: Vector,
images: Vec<Vector>,
title: String,
xlabel: String,
ylabel: String,
legends: Vec<String>,
markers: Vec<Markers>,
path: String,
fig_size: (usize, usize),
dpi: usize,
grid: Grid,
options: HashMap<PlotOptions, bool>,
}
impl Plot2D {
pub fn new() -> Self {
let mut default_options: HashMap<PlotOptions, bool> = HashMap::new();
default_options.insert(Domain, false);
default_options.insert(Images, false);
default_options.insert(Legends, false);
default_options.insert(Path, false);
Plot2D {
domain: vec![],
images: vec![],
title: "Title".to_string(),
xlabel: "$x$".to_string(),
ylabel: "$y$".to_string(),
legends: vec![],
markers: vec![],
path: "".to_string(),
fig_size: (10, 6),
dpi: 300,
grid: On,
options: default_options,
}
}
}
impl Plot for Plot2D {
fn set_domain(&mut self, x: Vec<f64>) -> &mut Self {
if let Some(x) = self.options.get_mut(&Domain) {
*x = true
}
self.domain = x;
self
}
fn insert_image(&mut self, y: Vec<f64>) -> &mut Self {
if let Some(x) = self.options.get_mut(&Images) {
*x = true
}
self.images.push(y);
self
}
fn set_title(&mut self, title: &str) -> &mut Self {
self.title = title.to_owned();
self
}
fn set_xlabel(&mut self, xlabel: &str) -> &mut Self {
self.xlabel = xlabel.to_owned();
self
}
fn set_ylabel(&mut self, ylabel: &str) -> &mut Self {
self.ylabel = ylabel.to_owned();
self
}
fn set_zlabel(&mut self, _zlabel: &str) -> &mut Self {
unimplemented!()
}
fn set_legends(&mut self, legends: Vec<&str>) -> &mut Self {
if let Some(x) = self.options.get_mut(&Legends) {
*x = true
}
self.legends = legends.into_iter().map(|x| x.to_owned()).collect::<Vec<String>>();
self
}
fn set_path(&mut self, path: &str) -> &mut Self {
if let Some(x) = self.options.get_mut(&Path) {
*x = true
}
self.path = path.to_owned();
self
}
fn set_fig_size(&mut self, fig_size: (usize, usize)) -> &mut Self {
self.fig_size = fig_size;
self
}
fn set_dpi(&mut self, dpi: usize) -> &mut Self {
self.dpi = dpi;
self
}
fn grid(&mut self, grid: Grid) -> &mut Self {
self.grid = grid;
self
}
fn set_marker(&mut self, styles: Vec<Markers>) -> &mut Self {
self.markers = styles;
self
}
fn savefig(&self) -> PyResult<()> {
match self.options.get(&Domain) {
Some(x) => {
assert!(*x, "Domain is not defined");
}
None => {
assert!(false, "Domain is None");
}
}
match self.options.get(&Images) {
Some(x) => {
assert!(*x, "Images are not defined");
}
None => {
assert!(false, "Images are None");
}
}
match self.options.get(&Legends) {
Some(x) => {
assert!(*x, "Legends are not defined");
assert!(
self.images.len() == self.legends.len(),
"Legends are not matched with images"
);
}
None => {
assert!(false, "Legends are None");
}
}
let gil = Python::acquire_gil();
let py = gil.python();
let x = self.domain.clone();
let ys = self.images.clone();
let y_length = ys.len();
let title = self.title.clone();
let fig_size = self.fig_size;
let dpi = self.dpi;
let grid = match self.grid {
On => true,
Off => false,
};
let xlabel = self.xlabel.clone();
let ylabel = self.ylabel.clone();
let legends = self.legends.clone();
let path = self.path.clone();
let globals = vec![("plt", py.import("pylab")?)].into_py_dict(py);
globals.set_item("x", x)?;
globals.set_item("y", ys)?;
globals.set_item("n", y_length)?;
globals.set_item("fs", fig_size)?;
globals.set_item("dp", dpi)?;
globals.set_item("gr", grid)?;
globals.set_item("pa", path)?;
let mut plot_string = format!(
"\
plt.rc(\"text\", usetex=True)\n\
plt.rc(\"font\", family=\"serif\")\n\
plt.figure(figsize=fs, dpi=dp)\n\
plt.title(r\"{}\", fontsize=16)\n\
plt.xlabel(r\"{}\", fontsize=14)\n\
plt.ylabel(r\"{}\", fontsize=14)\n\
if gr:\n\
\tplt.grid()\n",
title, xlabel, ylabel
);
if self.markers.len() == 0 {
for i in 0..y_length {
plot_string.push_str(&format!("plt.plot(x,y[{}],label=r\"{}\")\n", i, legends[i])[..])
}
} else {
for i in 0 .. y_length {
match self.markers[i] {
Line => plot_string.push_str(&format!("plt.plot(x,y[{}],label=r\"{}\")\n", i, legends[i])[..]),
Point => plot_string.push_str(&format!("plt.plot(x,y[{}],\".\",label=r\"{}\")\n", i, legends[i])[..]),
Circle => plot_string.push_str(&format!("plt.plot(x,y[{}],\"o\",label=r\"{}\")\n", i, legends[i])[..]),
}
}
}
plot_string.push_str("plt.legend(fontsize=12)\nplt.savefig(pa)");
py.run(&plot_string[..], Some(&globals), None)?;
Ok(())
}
}