<commit_before>use crate::language_client::LanguageClient;
use anyhow::Result;
use lsp_types::Command;
use serde::Deserialize;
use serde_json::Value;

pub mod command {
    pub const TEST: &str = "test";
    pub const GENERATE: &str = "generate";
}

impl LanguageClient {
    pub fn handle_gopls_command(&self, cmd: &Command) -> Result<bool> {
        match cmd.command.as_str() {
            command::TEST => {
                if let Some(args) = &cmd.arguments {
                    if let Some(file) = args.get(0) {
                        let file = String::deserialize(file)?;
                        let path = parse_package_path(file.as_str()).unwrap_or("./...".into());
                        let run = <Option<Vec<String>>>::deserialize(
                            args.get(1).unwrap_or(&Value::Null),
                        )?;

                        let bench = <Option<Vec<String>>>::deserialize(
                            args.get(2).unwrap_or(&Value::Null),
                        )?;

                        let run = run.unwrap_or_default();
                        let bench = bench.unwrap_or_default();

                        if run.len() > 0 {
                            let cmd = format!("term go test -run '{}' {}", run.join("|"), path);
                            self.vim()?.command(cmd)?;
                        } else if bench.len() > 0 {
                            let cmd = format!("term go test -bench '{}' {}", bench.join("|"), path);
                            self.vim()?.command(cmd)?;
                        } else {
                            self.vim()?.echoerr("No tests to run")?;
                        }
                    }
                }
            }
            command::GENERATE => {
                if let Some(arguments) = &cmd.arguments {
                    if let Some(package) = arguments.get(0) {
                        let package = String::deserialize(package)?;
                        let recursive =
                            bool::deserialize(arguments.get(1).unwrap_or(&Value::Bool(false)))?;
                        let cmd = match (package, recursive) {
                            (package, false) => format!("term go generate -x {}", package),
                            (_, true) => "term go generate -x ./...".into(),
                        };
                        self.vim()?.command(cmd)?;
                    }
                }
            }

            _ => return Ok(false),
        }

        Ok(true)
    }
}

fn parse_package_path(path: &str) -> Option<String> {
    let path = if path.starts_with("file://") {
        path.strip_prefix("file://")?
    } else {
        path
    };
    let path = std::path::PathBuf::from(path);
    Some(path.parent()?.to_str()?.to_owned())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_parse_package_path() {
        let folder = parse_package_path("file:///home/dev/someone/project/file.go");
        assert!(folder.is_some());
        assert_eq!("/home/dev/someone/project", folder.unwrap());

        let folder = parse_package_path("/home/dev/someone/project/file.go");
        assert!(folder.is_some());
        assert_eq!("/home/dev/someone/project", folder.unwrap());
    }
}
<commit_msg>Fix gopls codeLens commands to include the `gopls` prefix that it now sends (#1131)<commit_after>use crate::language_client::LanguageClient;
use anyhow::Result;
use lsp_types::Command;
use serde::Deserialize;
use serde_json::Value;

pub mod command {
    // older versions of gopls send code lens commands without the `gopls` prefix
    // so probably a good idea to maintain both for a while.
    pub(super) const TEST: &str = "test";
    pub(super) const GOPLS_TEST: &str = "gopls.test";
    pub(super) const GENERATE: &str = "generate";
    pub(super) const GOPLS_GENERATE: &str = "gopls.generate";
}

impl LanguageClient {
    pub fn handle_gopls_command(&self, cmd: &Command) -> Result<bool> {
        match cmd.command.as_str() {
            command::TEST | command::GOPLS_TEST => {
                if let Some(args) = &cmd.arguments {
                    if let Some(file) = args.get(0) {
                        let file = String::deserialize(file)?;
                        let path = parse_package_path(file.as_str()).unwrap_or("./...".into());
                        let run = <Option<Vec<String>>>::deserialize(
                            args.get(1).unwrap_or(&Value::Null),
                        )?;

                        let bench = <Option<Vec<String>>>::deserialize(
                            args.get(2).unwrap_or(&Value::Null),
                        )?;

                        let run = run.unwrap_or_default();
                        let bench = bench.unwrap_or_default();

                        if run.len() > 0 {
                            let cmd = format!("term go test -run '{}' {}", run.join("|"), path);
                            self.vim()?.command(cmd)?;
                        } else if bench.len() > 0 {
                            let cmd = format!("term go test -bench '{}' {}", bench.join("|"), path);
                            self.vim()?.command(cmd)?;
                        } else {
                            self.vim()?.echoerr("No tests to run")?;
                        }
                    }
                }
            }
            command::GENERATE | command::GOPLS_GENERATE => {
                if let Some(arguments) = &cmd.arguments {
                    if let Some(package) = arguments.get(0) {
                        let package = String::deserialize(package)?;
                        let recursive =
                            bool::deserialize(arguments.get(1).unwrap_or(&Value::Bool(false)))?;
                        let cmd = match (package, recursive) {
                            (package, false) => format!("term go generate -x {}", package),
                            (_, true) => "term go generate -x ./...".into(),
                        };
                        self.vim()?.command(cmd)?;
                    }
                }
            }

            _ => return Ok(false),
        }

        Ok(true)
    }
}

fn parse_package_path(path: &str) -> Option<String> {
    let path = if path.starts_with("file://") {
        path.strip_prefix("file://")?
    } else {
        path
    };
    let path = std::path::PathBuf::from(path);
    Some(path.parent()?.to_str()?.to_owned())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_parse_package_path() {
        let folder = parse_package_path("file:///home/dev/someone/project/file.go");
        assert!(folder.is_some());
        assert_eq!("/home/dev/someone/project", folder.unwrap());

        let folder = parse_package_path("/home/dev/someone/project/file.go");
        assert!(folder.is_some());
        assert_eq!("/home/dev/someone/project", folder.unwrap());
    }
}
<|endoftext|>


<commit_before>// Used internally for evalulating expressions
// Note: for now, negative is falsy and positive (including zero) is truthy

use interpreter::lexer::Token;
use interpreter::lexer;
use std::collections::HashMap;
use std::collections::hashmap::{Occupied, Vacant};
use std::num;

#[deriving(Show)]
enum ExprToken {
	Op(Operator),
	Value(f32),
	Var(uint),
	LParen,
	RParen,
}

#[deriving(Show)]
enum Operator {
	Add,
	Sub,
	Mul,
	Div,
	Exp,
	Mod,
	Neg,
	Less,
	Greater,
	Equ,
	NotEqu,
	ApproxEqu,
	Not,
	And,
	Or,
	Xor,
	GreaterEqual,
	LessEqual,
}

fn precedence(op: Operator) -> int {
	match op {
		And | Or | Xor => 0,
		Equ | NotEqu | ApproxEqu => 1,
		Less | Greater | GreaterEqual | LessEqual => 2,
		Add | Sub => 3,
		Mul | Div | Mod => 4,
		Neg | Not | Exp => 6,
	}
}

fn nargs(op: Operator) -> uint {
	match op {
		Neg | Not => 1,
		_ => 2,
	}
}

#[deriving(PartialEq)]
enum Associativity {
	Left,
	Right,
}

fn associativity(op: Operator) -> Associativity {
	match op {
		Exp => Right,
		_ => Left,
	}
}

#[deriving(Show)]
pub struct Expression<'a> {
	rpn: Vec<ExprToken>, // reverse polish notation
	var_values: Vec<f32>, // values of vars by id
	vars: HashMap<&'a str, uint>, // map of var name -> var id
	value: Option<f32>, // holds the last value of the expression, if it's been evaluated and vars have not been updated
}

impl<'a> Expression<'a> {
	// converts a token slice from the lexer into an expression that can be evaluated
	pub fn new<'a>(tok: Vec<Token<'a>>) -> Result<Expression, String> {
		let mut vars: HashMap<&'a str, uint> = HashMap::new();
		let mut err: Option<String> = None;
		let out: Vec<ExprToken> = tok.iter().filter_map(|&t| {
			if err.is_some() {
				return None;
			}
			match t {
				// Try to parse constants as f32. subject to maybe change but probably not.
				lexer::Const(v) => {
					if let Some(x) = from_str::<f32>(v) {
						Some(Value(x))
					} else {
						err = Some("internal error: error parsing constant, indicating the lexer is probably broken".to_string());
						None
					}
				},
				lexer::Operator(v) => {
					Some(Op(match v {
						"+" => Add,
						"-" => Sub,
						"*" => Mul,
						"/" => Div,
						"^" => Exp,
						"%" => Mod,
						"==" => Equ,
						"!=" => NotEqu,
						"~=" => ApproxEqu,
						"<" => Less,
						">" => Greater,
						"<=" => LessEqual,
						">=" => GreaterEqual,
						"!" => Not,
						"&&" => And,
						"||" => Or,
						"^^" => Xor,
						x => {
							err = Some(format!("unexpected operator: `{}`", x));
							return None; // report first error only
						}
					}))
				},
				lexer::Paren(v) => {
					match v {
						"(" => Some(LParen),
						")" => Some(RParen),
						x => {
							err = Some(format!("unexpected paren type: `{}`", x));
							return None;
						}
					}
				},
				// An identifier in the context of an expression is always a variable (TODO FALSE!!)
				lexer::Ident(v) => {
					let nv = vars.len();
					match vars.entry(v) {
						Occupied(entry) => {
							Some(Var(*entry.get()))
						},
						Vacant(entry) => {
							entry.set(nv);
							Some(Var(nv))
						},
					}
				},
				// Discard whitesapce
				lexer::Whitespace(_) | lexer::Newline => {
					None
				},
				x => {
					err = Some(format!("unexpected token in expression `{}`", x));
					return None;
				}
			}
		}).collect();

		if let Some(e) = err {
			return Err(e)
		}

		// Handle special case with unary minus
		// If an subtraction operator is preceded by another operator, left paren, or the start of the
		// expression, make it a negation operator.
		let out = out.iter().enumerate().map(|(i, &v)| {
			match v {
				Op(Sub) => {
					if i == 0 || match out[i-1] { Op(_) | LParen => true, _ => false } {
						Op(Neg)
					} else {
						v
					}
				},
				_ => v,
			}
		}).collect();

		let out = try!(shunting_yard(out));

		Ok(Expression {
			rpn: out,
			var_values: Vec::from_elem(vars.len(), 0f32),
			vars: vars,
			value: None,
		})
	}

	// http://en.wikipedia.org/wiki/Reverse_Polish_notation
	pub fn eval(&mut self) -> Result<f32, String> {
		// If the expression is constant, return it's value
		if let Some(v) = self.value {
			return Ok(v);
		}

		let mut stack: Vec<f32> = Vec::new();
		for &t in self.rpn.iter() {
			match t {
				Value(v) => {
					stack.push(v);
				},
				Op(op) => {
					let n = nargs(op);
					if stack.len() < n {
						return Err(format!("invalid expression: not enough args for operator {}", op));
					}
					let args: Vec<f32> = range(0, n).map(|_| stack.pop().unwrap()).collect();
					stack.push(match op {
						Add => {
							args[1] + args[0]
						},
						Sub => {
							args[1] - args[0]
						},
						Mul => {
							args[1] * args[0]
						},
						Div => {
							args[1] / args[0]
						},
						Exp => {
							args[1].powf(args[0])
						},
						Mod => {
							let c = num::abs(args[1]/args[0]).fract()*num::abs(args[0]);
							if args[1] < 0f32 { -c } else { c }
						},
						Neg => {
							-args[0]
						},
						Not => {
							if args[0] > 0f32 { -1f32 } else { 1f32 }
						},
						Less => {
							if args[1] < args[0] { 1f32 } else { -1f32 }
						},
						Greater => {
							if args[1] > args[0] { 1f32 } else { -1f32 }
						},
						LessEqual => {
							if args[1] <= args[0] { 1f32 } else { -1f32 }
						},
						GreaterEqual => {
							if args[1] >= args[0] { 1f32 } else { -1f32 }
						},
						Equ => {
							if args[1] == args[0] { 1f32 } else { -1f32 }
						},
						NotEqu => {
							if args[1] != args[0] { 1f32 } else { -1f32 }
						},
						ApproxEqu => {
							if num::abs(args[1] - args[0]) < 0.0001 { 1f32 } else { -1f32 }
						},
						And => {
							if args[1] > 0f32 && args[0] > 0f32 { 1f32 } else { -1f32 }
						},
						Or => {
							if args[1] > 0f32 || args[0] > 0f32 { 1f32 } else { -1f32 }
						},
						Xor => {
							if (args[1] > 0f32) ^ (args[0] > 0f32) { 1f32 } else { -1f32 }
						},
					});
				},
				Var(id) => {
					stack.push(self.var_values[id]);
				},
				x => return Err(format!("unexpected token in expression: `{}`", x)),
			};
		}
		match stack.len() {
			1 => {
				let val = stack.pop().unwrap();
				self.value = Some(val);
				Ok(val)
			},
			0 => {
				Err("zero values in expression".to_string())
			},
			_ => {
				Err("too many values in expression".to_string())
			},
		}
	}

	// Returns a uint identifier which can be passed to set_var to quickly update the value of a
	// variable.
	pub fn var_id(&self, var: &'a str) -> Option<uint> {
		match self.vars.find(&var) {
			Some(v) => Some(*v),
			None => None,
		}
	}

	// Sets a variable to an id retrieved from var_id. Return value indicates whether it was set
	// sucessfully or not.
	pub fn set_var(&mut self, var_id: uint, value: f32) -> bool {
		if var_id > self.var_values.len() {
			return false;
		}
		self.var_values[var_id] = value;
		self.value = None; // Setting a var clears evaled value
		true
	}

	// Attempts to set a variable by name. If the variable does not exist, nothing is done.
	pub fn try_set_var(&mut self, var: &'a str, value: f32) {
		if let Some(id) = self.var_id(var) {
			self.set_var(id, value);
		}
	}
}

// http://en.wikipedia.org/wiki/Shunting-yard_algorithm
// todo iterators can be used better here
fn shunting_yard(tok: Vec<ExprToken>) -> Result<Vec<ExprToken>, String> {
	let mut out = Vec::new();
	let mut stack: Vec<ExprToken> = Vec::new();
	for &t in tok.iter() {
		match t {
			Value(_) | Var(_) => {
				out.push(t);
			}
			Op(op1) => {
				while stack.len() > 0 {
					let top = stack[stack.len()-1];
					match top {
						Op(op2) => {
							if (associativity(op1) == Left && precedence(op1) <= precedence(op2)) ||
								(precedence(op1) < precedence(op2)) {
								stack.pop();
								out.push(top);
							} else {
								break;
							}
						},
						_ => break,
					}
				}
				stack.push(t);
			},
			LParen => {
				stack.push(t);
			},
			RParen => {
				let mut foundleft = false;
				while stack.len() > 0 {
					let op = stack.pop().unwrap();
					match op {
						Op(_) => {
							out.push(op);
						},
						LParen => {
							foundleft = true;
							break;
						},
						x => return Err(format!("internal error: unexpected value on stack: `{}`", x)),
					}
				}
				if !foundleft {
					return Err("mismatched parens: skewed right".to_string());
				}
			},
		};
	}
	while stack.len() > 0 {
		let top = stack[stack.len()-1];
		match top {
			Op(_) => {
				stack.pop();
				out.push(top);
			},
			LParen | RParen => return Err("mismatched parens: skewed left".to_string()),
			x => return Err(format!("internal error: non operator on stack: `{}`", x)),
		}
	}
	Ok(out)
}
<commit_msg>improve readability with comparison execution<commit_after>// Used internally for evalulating expressions
// Note: for now, negative is falsy and positive (including zero) is truthy

use interpreter::lexer::Token;
use interpreter::lexer;
use std::collections::HashMap;
use std::collections::hashmap::{Occupied, Vacant};
use std::num;

#[deriving(Show)]
enum ExprToken {
	Op(Operator),
	Value(f32),
	Var(uint),
	LParen,
	RParen,
}

#[deriving(Show)]
enum Operator {
	Add,
	Sub,
	Mul,
	Div,
	Exp,
	Mod,
	Neg,
	Less,
	Greater,
	Equ,
	NotEqu,
	ApproxEqu,
	Not,
	And,
	Or,
	Xor,
	GreaterEqual,
	LessEqual,
}

fn precedence(op: Operator) -> int {
	match op {
		And | Or | Xor => 0,
		Equ | NotEqu | ApproxEqu => 1,
		Less | Greater | GreaterEqual | LessEqual => 2,
		Add | Sub => 3,
		Mul | Div | Mod => 4,
		Neg | Not | Exp => 6,
	}
}

fn nargs(op: Operator) -> uint {
	match op {
		Neg | Not => 1,
		_ => 2,
	}
}

#[deriving(PartialEq)]
enum Associativity {
	Left,
	Right,
}

fn associativity(op: Operator) -> Associativity {
	match op {
		Exp => Right,
		_ => Left,
	}
}

#[deriving(Show)]
pub struct Expression<'a> {
	rpn: Vec<ExprToken>, // reverse polish notation
	var_values: Vec<f32>, // values of vars by id
	vars: HashMap<&'a str, uint>, // map of var name -> var id
	value: Option<f32>, // holds the last value of the expression, if it's been evaluated and vars have not been updated
}

impl<'a> Expression<'a> {
	// converts a token slice from the lexer into an expression that can be evaluated
	pub fn new<'a>(tok: Vec<Token<'a>>) -> Result<Expression, String> {
		let mut vars: HashMap<&'a str, uint> = HashMap::new();
		let mut err: Option<String> = None;
		let out: Vec<ExprToken> = tok.iter().filter_map(|&t| {
			if err.is_some() {
				return None;
			}
			match t {
				// Try to parse constants as f32. subject to maybe change but probably not.
				lexer::Const(v) => {
					if let Some(x) = from_str::<f32>(v) {
						Some(Value(x))
					} else {
						err = Some("internal error: error parsing constant, indicating the lexer is probably broken".to_string());
						None
					}
				},
				lexer::Operator(v) => {
					Some(Op(match v {
						"+" => Add,
						"-" => Sub,
						"*" => Mul,
						"/" => Div,
						"^" => Exp,
						"%" => Mod,
						"==" => Equ,
						"!=" => NotEqu,
						"~=" => ApproxEqu,
						"<" => Less,
						">" => Greater,
						"<=" => LessEqual,
						">=" => GreaterEqual,
						"!" => Not,
						"&&" => And,
						"||" => Or,
						"^^" => Xor,
						x => {
							err = Some(format!("unexpected operator: `{}`", x));
							return None; // report first error only
						}
					}))
				},
				lexer::Paren(v) => {
					match v {
						"(" => Some(LParen),
						")" => Some(RParen),
						x => {
							err = Some(format!("unexpected paren type: `{}`", x));
							return None;
						}
					}
				},
				// An identifier in the context of an expression is always a variable (TODO FALSE!!)
				lexer::Ident(v) => {
					let nv = vars.len();
					match vars.entry(v) {
						Occupied(entry) => {
							Some(Var(*entry.get()))
						},
						Vacant(entry) => {
							entry.set(nv);
							Some(Var(nv))
						},
					}
				},
				// Discard whitesapce
				lexer::Whitespace(_) | lexer::Newline => {
					None
				},
				x => {
					err = Some(format!("unexpected token in expression `{}`", x));
					return None;
				}
			}
		}).collect();

		if let Some(e) = err {
			return Err(e)
		}

		// Handle special case with unary minus
		// If an subtraction operator is preceded by another operator, left paren, or the start of the
		// expression, make it a negation operator.
		let out = out.iter().enumerate().map(|(i, &v)| {
			match v {
				Op(Sub) => {
					if i == 0 || match out[i-1] { Op(_) | LParen => true, _ => false } {
						Op(Neg)
					} else {
						v
					}
				},
				_ => v,
			}
		}).collect();

		let out = try!(shunting_yard(out));

		Ok(Expression {
			rpn: out,
			var_values: Vec::from_elem(vars.len(), 0f32),
			vars: vars,
			value: None,
		})
	}

	// http://en.wikipedia.org/wiki/Reverse_Polish_notation
	pub fn eval(&mut self) -> Result<f32, String> {
		// If the expression is constant, return it's value
		if let Some(v) = self.value {
			return Ok(v);
		}

		let mut stack: Vec<f32> = Vec::new();
		for &t in self.rpn.iter() {
			match t {
				Value(v) => {
					stack.push(v);
				},
				Op(op) => {
					let n = nargs(op);
					if stack.len() < n {
						return Err(format!("invalid expression: not enough args for operator {}", op));
					}
					let args: Vec<f32> = range(0, n).map(|_| stack.pop().unwrap()).collect();
					stack.push(match op {
						Add => {
							args[1] + args[0]
						},
						Sub => {
							args[1] - args[0]
						},
						Mul => {
							args[1] * args[0]
						},
						Div => {
							args[1] / args[0]
						},
						Exp => {
							args[1].powf(args[0])
						},
						Mod => {
							let c = num::abs(args[1]/args[0]).fract()*num::abs(args[0]);
							if args[1] < 0f32 { -c } else { c }
						},
						Neg => {
							-args[0]
						},
						Not => {
							if args[0] > 0f32
								{ -1f32 } else { 1f32 }
						},
						Less => {
							if args[1] < args[0]
								{ 1f32 } else { -1f32 }
						},
						Greater => {
							if args[1] > args[0]
								{ 1f32 } else { -1f32 }
						},
						LessEqual => {
							if args[1] <= args[0]
								{ 1f32 } else { -1f32 }
						},
						GreaterEqual => {
							if args[1] >= args[0]
								{ 1f32 } else { -1f32 }
						},
						Equ => {
							if args[1] == args[0]
								{ 1f32 } else { -1f32 }
						},
						NotEqu => {
							if args[1] != args[0]
								{ 1f32 } else { -1f32 }
						},
						ApproxEqu => {
							if num::abs(args[1] - args[0]) < 0.0001
								{ 1f32 } else { -1f32 }
						},
						And => {
							if args[1] > 0f32 && args[0] > 0f32
								{ 1f32 } else { -1f32 }
						},
						Or => {
							if args[1] > 0f32 || args[0] > 0f32
								{ 1f32 } else { -1f32 }
						},
						Xor => {
							if (args[1] > 0f32) ^ (args[0] > 0f32)
								{ 1f32 } else { -1f32 }
						},
					});
				},
				Var(id) => {
					stack.push(self.var_values[id]);
				},
				x => return Err(format!("unexpected token in expression: `{}`", x)),
			};
		}
		match stack.len() {
			1 => {
				let val = stack.pop().unwrap();
				self.value = Some(val);
				Ok(val)
			},
			0 => {
				Err("zero values in expression".to_string())
			},
			_ => {
				Err("too many values in expression".to_string())
			},
		}
	}

	// Returns a uint identifier which can be passed to set_var to quickly update the value of a
	// variable.
	pub fn var_id(&self, var: &'a str) -> Option<uint> {
		match self.vars.find(&var) {
			Some(v) => Some(*v),
			None => None,
		}
	}

	// Sets a variable to an id retrieved from var_id. Return value indicates whether it was set
	// sucessfully or not.
	pub fn set_var(&mut self, var_id: uint, value: f32) -> bool {
		if var_id > self.var_values.len() {
			return false;
		}
		self.var_values[var_id] = value;
		self.value = None; // Setting a var clears evaled value
		true
	}

	// Attempts to set a variable by name. If the variable does not exist, nothing is done.
	pub fn try_set_var(&mut self, var: &'a str, value: f32) {
		if let Some(id) = self.var_id(var) {
			self.set_var(id, value);
		}
	}
}

// http://en.wikipedia.org/wiki/Shunting-yard_algorithm
// todo iterators can be used better here
fn shunting_yard(tok: Vec<ExprToken>) -> Result<Vec<ExprToken>, String> {
	let mut out = Vec::new();
	let mut stack: Vec<ExprToken> = Vec::new();
	for &t in tok.iter() {
		match t {
			Value(_) | Var(_) => {
				out.push(t);
			}
			Op(op1) => {
				while stack.len() > 0 {
					let top = stack[stack.len()-1];
					match top {
						Op(op2) => {
							if (associativity(op1) == Left && precedence(op1) <= precedence(op2)) ||
								(precedence(op1) < precedence(op2)) {
								stack.pop();
								out.push(top);
							} else {
								break;
							}
						},
						_ => break,
					}
				}
				stack.push(t);
			},
			LParen => {
				stack.push(t);
			},
			RParen => {
				let mut foundleft = false;
				while stack.len() > 0 {
					let op = stack.pop().unwrap();
					match op {
						Op(_) => {
							out.push(op);
						},
						LParen => {
							foundleft = true;
							break;
						},
						x => return Err(format!("internal error: unexpected value on stack: `{}`", x)),
					}
				}
				if !foundleft {
					return Err("mismatched parens: skewed right".to_string());
				}
			},
		};
	}
	while stack.len() > 0 {
		let top = stack[stack.len()-1];
		match top {
			Op(_) => {
				stack.pop();
				out.push(top);
			},
			LParen | RParen => return Err("mismatched parens: skewed left".to_string()),
			x => return Err(format!("internal error: non operator on stack: `{}`", x)),
		}
	}
	Ok(out)
}
<|endoftext|>


<commit_before>use lazy_static::lazy_static;
use markup5ever::{local_name, LocalName};
use std::collections::HashMap;

use crate::clip_path::NodeClipPath;
use crate::filters::{
    blend::Blend,
    color_matrix::ColorMatrix,
    component_transfer::{ComponentTransfer, FuncA, FuncB, FuncG, FuncR},
    composite::Composite,
    convolve_matrix::ConvolveMatrix,
    displacement_map::DisplacementMap,
    flood::Flood,
    gaussian_blur::GaussianBlur,
    image::Image,
    light::{
        light_source::DistantLight, light_source::PointLight, light_source::SpotLight,
        lighting::DiffuseLighting, lighting::SpecularLighting,
    },
    merge::{Merge, MergeNode},
    morphology::Morphology,
    node::NodeFilter,
    offset::Offset,
    tile::Tile,
    turbulence::Turbulence,
};

use crate::gradient::{NodeLinearGradient, NodeRadialGradient, NodeStop};
use crate::image::NodeImage;
use crate::link::NodeLink;
use crate::marker::NodeMarker;
use crate::mask::NodeMask;
use crate::node::*;
use crate::pattern::NodePattern;
use crate::property_bag::PropertyBag;
use crate::shapes::{
    NodeCircle, NodeEllipse, NodeLine, NodePath, NodePolygon, NodePolyline, NodeRect,
};
use crate::structure::{NodeGroup, NodeNonRendering, NodeSvg, NodeSwitch, NodeSymbol, NodeUse};
use crate::style::NodeStyle;
use crate::text::{NodeTRef, NodeTSpan, NodeText};

macro_rules! n {
    ($name:ident, $node_type:ident, $new_fn:expr) => {
        pub fn $name(element_name: LocalName, id: Option<&str>, class: Option<&str>) -> RsvgNode {
            RsvgNode::new(NodeData::new(
                NodeType::$node_type,
                element_name,
                id,
                class,
                Box::new($new_fn()),
            ))
        }
    };
}

#[cfg_attr(rustfmt, rustfmt_skip)]
mod creators {
    use super::*;

    n!(create_circle,                    Circle,                     NodeCircle::default);
    n!(create_clip_path,                 ClipPath,                   NodeClipPath::default);
    n!(create_blend,                     FeBlend,                    Blend::default);
    n!(create_color_matrix,              FeColorMatrix,              ColorMatrix::default);
    n!(create_component_transfer,        FeComponentTransfer,        ComponentTransfer::default);
    n!(create_component_transfer_func_a, ComponentTransferFunctionA, FuncA::default);
    n!(create_component_transfer_func_b, ComponentTransferFunctionB, FuncB::default);
    n!(create_component_transfer_func_g, ComponentTransferFunctionG, FuncG::default);
    n!(create_component_transfer_func_r, ComponentTransferFunctionR, FuncR::default);
    n!(create_composite,                 FeComposite,                Composite::default);
    n!(create_convolve_matrix,           FeConvolveMatrix,           ConvolveMatrix::default);
    n!(create_defs,                      Defs,                       NodeNonRendering::default);
    n!(create_diffuse_lighting,          FeDiffuseLighting,          DiffuseLighting::default);
    n!(create_distant_light,             FeDistantLight,             DistantLight::default);
    n!(create_displacement_map,          FeDisplacementMap,          DisplacementMap::default);
    n!(create_ellipse,                   Ellipse,                    NodeEllipse::default);
    n!(create_filter,                    Filter,                     NodeFilter::default);
    n!(create_flood,                     FeFlood,                    Flood::default);
    n!(create_gaussian_blur,             FeGaussianBlur,             GaussianBlur::default);
    n!(create_group,                     Group,                      NodeGroup::default);
    n!(create_image,                     Image,                      NodeImage::default);
    n!(create_fe_image,                  FeImage,                    Image::default);
    n!(create_line,                      Line,                       NodeLine::default);
    n!(create_linear_gradient,           LinearGradient,             NodeLinearGradient::default);
    n!(create_link,                      Link,                       NodeLink::default);
    n!(create_marker,                    Marker,                     NodeMarker::default);
    n!(create_mask,                      Mask,                       NodeMask::default);
    n!(create_merge,                     FeMerge,                    Merge::default);
    n!(create_merge_node,                FeMergeNode,                MergeNode::default);
    n!(create_morphology,                FeMorphology,               Morphology::default);
    n!(create_non_rendering,             NonRendering,               NodeNonRendering::default);
    n!(create_offset,                    FeOffset,                   Offset::default);
    n!(create_path,                      Path,                       NodePath::default);
    n!(create_pattern,                   Pattern,                    NodePattern::default);
    n!(create_point_light,               FePointLight,               PointLight::default);
    n!(create_polygon,                   Polygon,                    NodePolygon::default);
    n!(create_polyline,                  Polyline,                   NodePolyline::default);
    n!(create_radial_gradient,           RadialGradient,             NodeRadialGradient::default);
    n!(create_rect,                      Rect,                       NodeRect::default);
    n!(create_specular_lighting,         FeSpecularLighting,         SpecularLighting::default);
    n!(create_spot_light,                FeSpotLight,                SpotLight::default);
    n!(create_stop,                      Stop,                       NodeStop::default);
    n!(create_style,                     Style,                      NodeStyle::default);
    n!(create_svg,                       Svg,                        NodeSvg::default);
    n!(create_switch,                    Switch,                     NodeSwitch::default);
    n!(create_symbol,                    Symbol,                     NodeSymbol::default);
    n!(create_text,                      Text,                       NodeText::default);
    n!(create_tref,                      TRef,                       NodeTRef::default);
    n!(create_tspan,                     TSpan,                      NodeTSpan::default);
    n!(create_tile,                      FeTile,                     Tile::default);
    n!(create_turbulence,                FeTurbulence,               Turbulence::default);
    n!(create_use,                       Use,                        NodeUse::default);

    // hack to make multiImage sort-of work
    n!(create_multi_image,               Switch,                     NodeSwitch::default);
    n!(create_sub_image,                 Group,                      NodeGroup::default);
    n!(create_sub_image_ref,             Image,                      NodeImage::default);
}

use creators::*;

type NodeCreateFn = fn(element_name: LocalName, id: Option<&str>, class: Option<&str>) -> RsvgNode;

macro_rules! c {
    ($hashset:expr, $str_name:expr, $supports_class:expr, $fn_name:ident) => {
        $hashset.insert($str_name, ($supports_class, $fn_name as NodeCreateFn));
    };
}

lazy_static! {
    // Lines in comments are elements that we don't support.
    #[cfg_attr(rustfmt, rustfmt_skip)]
    static ref NODE_CREATORS: HashMap<&'static str, (bool, NodeCreateFn)> = {
        let mut h = HashMap::new();
        // name, supports_class, create_fn
        c!(h, "a",                   true,  create_link);
        /* c!(h, "altGlyph",         true,  ); */
        /* c!(h, "altGlyphDef",      false, ); */
        /* c!(h, "altGlyphItem",     false, ); */
        /* c!(h, "animate",          false, ); */
        /* c!(h, "animateColor",     false, ); */
        /* c!(h, "animateMotion",    false, ); */
        /* c!(h, "animateTransform", false, ); */
        c!(h, "circle",              true,  create_circle);
        c!(h, "clipPath",            true,  create_clip_path);
        /* c!(h, "color-profile",    false, ); */
        /* c!(h, "cursor",           false, ); */
        c!(h, "defs",                true,  create_defs);
        /* c!(h, "desc",             true,  ); */
        c!(h, "ellipse",             true,  create_ellipse);
        c!(h, "feBlend",             true,  create_blend);
        c!(h, "feColorMatrix",       true,  create_color_matrix);
        c!(h, "feComponentTransfer", true,  create_component_transfer);
        c!(h, "feComposite",         true,  create_composite);
        c!(h, "feConvolveMatrix",    true,  create_convolve_matrix);
        c!(h, "feDiffuseLighting",   true,  create_diffuse_lighting);
        c!(h, "feDisplacementMap",   true,  create_displacement_map);
        c!(h, "feDistantLight",      false, create_distant_light);
        c!(h, "feFuncA",             false, create_component_transfer_func_a);
        c!(h, "feFuncB",             false, create_component_transfer_func_b);
        c!(h, "feFuncG",             false, create_component_transfer_func_g);
        c!(h, "feFuncR",             false, create_component_transfer_func_r);
        c!(h, "feFlood",             true,  create_flood);
        c!(h, "feGaussianBlur",      true,  create_gaussian_blur);
        c!(h, "feImage",             true,  create_fe_image);
        c!(h, "feMerge",             true,  create_merge);
        c!(h, "feMergeNode",         false, create_merge_node);
        c!(h, "feMorphology",        true,  create_morphology);
        c!(h, "feOffset",            true,  create_offset);
        c!(h, "fePointLight",        false, create_point_light);
        c!(h, "feSpecularLighting",  true,  create_specular_lighting);
        c!(h, "feSpotLight",         false, create_spot_light);
        c!(h, "feTile",              true,  create_tile);
        c!(h, "feTurbulence",        true,  create_turbulence);
        c!(h, "filter",              true,  create_filter);
        /* c!(h, "font",             true,  ); */
        /* c!(h, "font-face",        false, ); */
        /* c!(h, "font-face-format", false, ); */
        /* c!(h, "font-face-name",   false, ); */
        /* c!(h, "font-face-src",    false, ); */
        /* c!(h, "font-face-uri",    false, ); */
        /* c!(h, "foreignObject",    true,  ); */
        c!(h, "g",                   true,  create_group);
        /* c!(h, "glyph",            true,  ); */
        /* c!(h, "glyphRef",         true,  ); */
        /* c!(h, "hkern",            false, ); */
        c!(h, "image",               true,  create_image);
        c!(h, "line",                true,  create_line);
        c!(h, "linearGradient",      true,  create_linear_gradient);
        c!(h, "marker",              true,  create_marker);
        c!(h, "mask",                true,  create_mask);
        /* c!(h, "metadata",         false, ); */
        /* c!(h, "missing-glyph",    true,  ); */
        /* c!(h, "mpath",            false, ); */
        c!(h, "multiImage",          false, create_multi_image);
        c!(h, "path",                true,  create_path);
        c!(h, "pattern",             true,  create_pattern);
        c!(h, "polygon",             true,  create_polygon);
        c!(h, "polyline",            true,  create_polyline);
        c!(h, "radialGradient",      true,  create_radial_gradient);
        c!(h, "rect",                true,  create_rect);
        /* c!(h, "script",           false, ); */
        /* c!(h, "set",              false, ); */
        c!(h, "stop",                true,  create_stop);
        c!(h, "style",               false, create_style);
        c!(h, "subImage",            false, create_sub_image);
        c!(h, "subImageRef",         false, create_sub_image_ref);
        c!(h, "svg",                 true,  create_svg);
        c!(h, "switch",              true,  create_switch);
        c!(h, "symbol",              true,  create_symbol);
        c!(h, "text",                true,  create_text);
        /* c!(h, "textPath",         true,  ); */
        /* c!(h, "title",            true,  ); */
        c!(h, "tref",                true,  create_tref);
        c!(h, "tspan",               true,  create_tspan);
        c!(h, "use",                 true,  create_use);
        /* c!(h, "view",             false, ); */
        /* c!(h, "vkern",            false, ); */
        h
    };
}

pub fn create_node_and_register_id(
    name: &str,
    pbag: &PropertyBag,
    ids: &mut HashMap<String, RsvgNode>,
) -> RsvgNode {
    let mut id = None;
    let mut class = None;

    for (attr, value) in pbag.iter() {
        match attr {
            local_name!("id") => id = Some(value),
            local_name!("class") => class = Some(value),
            _ => (),
        }
    }

    let &(supports_class, create_fn) = match NODE_CREATORS.get(name) {
        Some(c) => c,

        // Whenever we encounter a node we don't understand, represent it as a
        // non-rendering node.  This is like a group, but it doesn't do any rendering of
        // children.  The effect is that we will ignore all children of unknown elements.
        None => &(true, create_non_rendering as NodeCreateFn),
    };

    let element_name = LocalName::from(name);

    if !supports_class {
        class = None;
    };

    let node = create_fn(element_name, id, class);

    if let Some(id) = id {
        // This is so we don't overwrite an existing id
        ids.entry(id.to_string()).or_insert_with(|| node.clone());
    }

    node
}
<commit_msg>node_create: move the call to default() in the macro<commit_after>use lazy_static::lazy_static;
use markup5ever::{local_name, LocalName};
use std::collections::HashMap;

use crate::clip_path::NodeClipPath;
use crate::filters::{
    blend::Blend,
    color_matrix::ColorMatrix,
    component_transfer::{ComponentTransfer, FuncA, FuncB, FuncG, FuncR},
    composite::Composite,
    convolve_matrix::ConvolveMatrix,
    displacement_map::DisplacementMap,
    flood::Flood,
    gaussian_blur::GaussianBlur,
    image::Image,
    light::{
        light_source::DistantLight, light_source::PointLight, light_source::SpotLight,
        lighting::DiffuseLighting, lighting::SpecularLighting,
    },
    merge::{Merge, MergeNode},
    morphology::Morphology,
    node::NodeFilter,
    offset::Offset,
    tile::Tile,
    turbulence::Turbulence,
};

use crate::gradient::{NodeLinearGradient, NodeRadialGradient, NodeStop};
use crate::image::NodeImage;
use crate::link::NodeLink;
use crate::marker::NodeMarker;
use crate::mask::NodeMask;
use crate::node::*;
use crate::pattern::NodePattern;
use crate::property_bag::PropertyBag;
use crate::shapes::{
    NodeCircle, NodeEllipse, NodeLine, NodePath, NodePolygon, NodePolyline, NodeRect,
};
use crate::structure::{NodeGroup, NodeNonRendering, NodeSvg, NodeSwitch, NodeSymbol, NodeUse};
use crate::style::NodeStyle;
use crate::text::{NodeTRef, NodeTSpan, NodeText};

macro_rules! n {
    ($name:ident, $node_type:ident, $node_trait:ty) => {
        pub fn $name(element_name: LocalName, id: Option<&str>, class: Option<&str>) -> RsvgNode {
            RsvgNode::new(NodeData::new(
                NodeType::$node_type,
                element_name,
                id,
                class,
                Box::new(<$node_trait>::default()),
            ))
        }
    };
}

#[cfg_attr(rustfmt, rustfmt_skip)]
mod creators {
    use super::*;

    n!(create_circle,                    Circle,                     NodeCircle);
    n!(create_clip_path,                 ClipPath,                   NodeClipPath);
    n!(create_blend,                     FeBlend,                    Blend);
    n!(create_color_matrix,              FeColorMatrix,              ColorMatrix);
    n!(create_component_transfer,        FeComponentTransfer,        ComponentTransfer);
    n!(create_component_transfer_func_a, ComponentTransferFunctionA, FuncA);
    n!(create_component_transfer_func_b, ComponentTransferFunctionB, FuncB);
    n!(create_component_transfer_func_g, ComponentTransferFunctionG, FuncG);
    n!(create_component_transfer_func_r, ComponentTransferFunctionR, FuncR);
    n!(create_composite,                 FeComposite,                Composite);
    n!(create_convolve_matrix,           FeConvolveMatrix,           ConvolveMatrix);
    n!(create_defs,                      Defs,                       NodeNonRendering);
    n!(create_diffuse_lighting,          FeDiffuseLighting,          DiffuseLighting);
    n!(create_distant_light,             FeDistantLight,             DistantLight);
    n!(create_displacement_map,          FeDisplacementMap,          DisplacementMap);
    n!(create_ellipse,                   Ellipse,                    NodeEllipse);
    n!(create_filter,                    Filter,                     NodeFilter);
    n!(create_flood,                     FeFlood,                    Flood);
    n!(create_gaussian_blur,             FeGaussianBlur,             GaussianBlur);
    n!(create_group,                     Group,                      NodeGroup);
    n!(create_image,                     Image,                      NodeImage);
    n!(create_fe_image,                  FeImage,                    Image);
    n!(create_line,                      Line,                       NodeLine);
    n!(create_linear_gradient,           LinearGradient,             NodeLinearGradient);
    n!(create_link,                      Link,                       NodeLink);
    n!(create_marker,                    Marker,                     NodeMarker);
    n!(create_mask,                      Mask,                       NodeMask);
    n!(create_merge,                     FeMerge,                    Merge);
    n!(create_merge_node,                FeMergeNode,                MergeNode);
    n!(create_morphology,                FeMorphology,               Morphology);
    n!(create_non_rendering,             NonRendering,               NodeNonRendering);
    n!(create_offset,                    FeOffset,                   Offset);
    n!(create_path,                      Path,                       NodePath);
    n!(create_pattern,                   Pattern,                    NodePattern);
    n!(create_point_light,               FePointLight,               PointLight);
    n!(create_polygon,                   Polygon,                    NodePolygon);
    n!(create_polyline,                  Polyline,                   NodePolyline);
    n!(create_radial_gradient,           RadialGradient,             NodeRadialGradient);
    n!(create_rect,                      Rect,                       NodeRect);
    n!(create_specular_lighting,         FeSpecularLighting,         SpecularLighting);
    n!(create_spot_light,                FeSpotLight,                SpotLight);
    n!(create_stop,                      Stop,                       NodeStop);
    n!(create_style,                     Style,                      NodeStyle);
    n!(create_svg,                       Svg,                        NodeSvg);
    n!(create_switch,                    Switch,                     NodeSwitch);
    n!(create_symbol,                    Symbol,                     NodeSymbol);
    n!(create_text,                      Text,                       NodeText);
    n!(create_tref,                      TRef,                       NodeTRef);
    n!(create_tspan,                     TSpan,                      NodeTSpan);
    n!(create_tile,                      FeTile,                     Tile);
    n!(create_turbulence,                FeTurbulence,               Turbulence);
    n!(create_use,                       Use,                        NodeUse);

    // hack to make multiImage sort-of work
    n!(create_multi_image,               Switch,                     NodeSwitch);
    n!(create_sub_image,                 Group,                      NodeGroup);
    n!(create_sub_image_ref,             Image,                      NodeImage);
}

use creators::*;

type NodeCreateFn = fn(element_name: LocalName, id: Option<&str>, class: Option<&str>) -> RsvgNode;

macro_rules! c {
    ($hashset:expr, $str_name:expr, $supports_class:expr, $fn_name:ident) => {
        $hashset.insert($str_name, ($supports_class, $fn_name as NodeCreateFn));
    };
}

lazy_static! {
    // Lines in comments are elements that we don't support.
    #[cfg_attr(rustfmt, rustfmt_skip)]
    static ref NODE_CREATORS: HashMap<&'static str, (bool, NodeCreateFn)> = {
        let mut h = HashMap::new();
        // name, supports_class, create_fn
        c!(h, "a",                   true,  create_link);
        /* c!(h, "altGlyph",         true,  ); */
        /* c!(h, "altGlyphDef",      false, ); */
        /* c!(h, "altGlyphItem",     false, ); */
        /* c!(h, "animate",          false, ); */
        /* c!(h, "animateColor",     false, ); */
        /* c!(h, "animateMotion",    false, ); */
        /* c!(h, "animateTransform", false, ); */
        c!(h, "circle",              true,  create_circle);
        c!(h, "clipPath",            true,  create_clip_path);
        /* c!(h, "color-profile",    false, ); */
        /* c!(h, "cursor",           false, ); */
        c!(h, "defs",                true,  create_defs);
        /* c!(h, "desc",             true,  ); */
        c!(h, "ellipse",             true,  create_ellipse);
        c!(h, "feBlend",             true,  create_blend);
        c!(h, "feColorMatrix",       true,  create_color_matrix);
        c!(h, "feComponentTransfer", true,  create_component_transfer);
        c!(h, "feComposite",         true,  create_composite);
        c!(h, "feConvolveMatrix",    true,  create_convolve_matrix);
        c!(h, "feDiffuseLighting",   true,  create_diffuse_lighting);
        c!(h, "feDisplacementMap",   true,  create_displacement_map);
        c!(h, "feDistantLight",      false, create_distant_light);
        c!(h, "feFuncA",             false, create_component_transfer_func_a);
        c!(h, "feFuncB",             false, create_component_transfer_func_b);
        c!(h, "feFuncG",             false, create_component_transfer_func_g);
        c!(h, "feFuncR",             false, create_component_transfer_func_r);
        c!(h, "feFlood",             true,  create_flood);
        c!(h, "feGaussianBlur",      true,  create_gaussian_blur);
        c!(h, "feImage",             true,  create_fe_image);
        c!(h, "feMerge",             true,  create_merge);
        c!(h, "feMergeNode",         false, create_merge_node);
        c!(h, "feMorphology",        true,  create_morphology);
        c!(h, "feOffset",            true,  create_offset);
        c!(h, "fePointLight",        false, create_point_light);
        c!(h, "feSpecularLighting",  true,  create_specular_lighting);
        c!(h, "feSpotLight",         false, create_spot_light);
        c!(h, "feTile",              true,  create_tile);
        c!(h, "feTurbulence",        true,  create_turbulence);
        c!(h, "filter",              true,  create_filter);
        /* c!(h, "font",             true,  ); */
        /* c!(h, "font-face",        false, ); */
        /* c!(h, "font-face-format", false, ); */
        /* c!(h, "font-face-name",   false, ); */
        /* c!(h, "font-face-src",    false, ); */
        /* c!(h, "font-face-uri",    false, ); */
        /* c!(h, "foreignObject",    true,  ); */
        c!(h, "g",                   true,  create_group);
        /* c!(h, "glyph",            true,  ); */
        /* c!(h, "glyphRef",         true,  ); */
        /* c!(h, "hkern",            false, ); */
        c!(h, "image",               true,  create_image);
        c!(h, "line",                true,  create_line);
        c!(h, "linearGradient",      true,  create_linear_gradient);
        c!(h, "marker",              true,  create_marker);
        c!(h, "mask",                true,  create_mask);
        /* c!(h, "metadata",         false, ); */
        /* c!(h, "missing-glyph",    true,  ); */
        /* c!(h, "mpath",            false, ); */
        c!(h, "multiImage",          false, create_multi_image);
        c!(h, "path",                true,  create_path);
        c!(h, "pattern",             true,  create_pattern);
        c!(h, "polygon",             true,  create_polygon);
        c!(h, "polyline",            true,  create_polyline);
        c!(h, "radialGradient",      true,  create_radial_gradient);
        c!(h, "rect",                true,  create_rect);
        /* c!(h, "script",           false, ); */
        /* c!(h, "set",              false, ); */
        c!(h, "stop",                true,  create_stop);
        c!(h, "style",               false, create_style);
        c!(h, "subImage",            false, create_sub_image);
        c!(h, "subImageRef",         false, create_sub_image_ref);
        c!(h, "svg",                 true,  create_svg);
        c!(h, "switch",              true,  create_switch);
        c!(h, "symbol",              true,  create_symbol);
        c!(h, "text",                true,  create_text);
        /* c!(h, "textPath",         true,  ); */
        /* c!(h, "title",            true,  ); */
        c!(h, "tref",                true,  create_tref);
        c!(h, "tspan",               true,  create_tspan);
        c!(h, "use",                 true,  create_use);
        /* c!(h, "view",             false, ); */
        /* c!(h, "vkern",            false, ); */
        h
    };
}

pub fn create_node_and_register_id(
    name: &str,
    pbag: &PropertyBag,
    ids: &mut HashMap<String, RsvgNode>,
) -> RsvgNode {
    let mut id = None;
    let mut class = None;

    for (attr, value) in pbag.iter() {
        match attr {
            local_name!("id") => id = Some(value),
            local_name!("class") => class = Some(value),
            _ => (),
        }
    }

    let &(supports_class, create_fn) = match NODE_CREATORS.get(name) {
        Some(c) => c,

        // Whenever we encounter a node we don't understand, represent it as a
        // non-rendering node.  This is like a group, but it doesn't do any rendering of
        // children.  The effect is that we will ignore all children of unknown elements.
        None => &(true, create_non_rendering as NodeCreateFn),
    };

    let element_name = LocalName::from(name);

    if !supports_class {
        class = None;
    };

    let node = create_fn(element_name, id, class);

    if let Some(id) = id {
        // This is so we don't overwrite an existing id
        ids.entry(id.to_string()).or_insert_with(|| node.clone());
    }

    node
}
<|endoftext|>


<commit_before>//! Gaussian Mixture Models
//!
//! Provides implementation of GMMs using the EM algorithm.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::learning::gmm::{CovOption, GaussianMixtureModel};
//! use rusty_machine::learning::UnSupModel;
//!
//! let inputs = Matrix::new(4, 2, vec![1.0, 2.0, 3.0, 3.0, 2.0, 4.0, 5.0, 2.5]);
//! let test_inputs = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 2.9, 2.4, 2.5]);
//!
//! // Create gmm with k(=2) classes.
//! let mut model = GaussianMixtureModel::new(2);
//! model.set_max_iters(10);
//! model.cov_option = CovOption::Diagonal;
//!
//! // Where inputs is a Matrix with features in columns.
//! model.train(&inputs);
//!
//! // Where pred_data is a Matrix with features in columns.
//! let a = model.predict(&test_inputs);
//! println!("{:?}", a.data());
//! ```

use linalg::vector::Vector;
use linalg::matrix::Matrix;
use linalg::utils;

use learning::UnSupModel;
use learning::toolkit::rand_utils;

/// Covariance options for GMMs.
///
/// - Full : The full covariance structure.
/// - Regularized : Adds a regularization constant to the covariance diagonal.
/// - Diagonal : Only the diagonal covariance structure.
#[derive(Clone, Copy, Debug)]
pub enum CovOption {
    /// The full covariance structure.
    Full,
    /// Adds a regularization constant to the covariance diagonal.
    Regularized(f64),
    /// Only the diagonal covariance structure.
    Diagonal,
}

/// A Gaussian Mixture Model
#[derive(Debug)]
pub struct GaussianMixtureModel {
    comp_count: usize,
    mix_weights: Vector<f64>,
    model_means: Option<Matrix<f64>>,
    model_covars: Option<Vec<Matrix<f64>>>,
    log_lik: f64,
    max_iters: usize,
    /// The covariance options for the GMM.
    pub cov_option: CovOption,
}

impl UnSupModel<Matrix<f64>, Matrix<f64>> for GaussianMixtureModel {
    /// Train the model using inputs.
    fn train(&mut self, inputs: &Matrix<f64>) {
        // Initialization:
        let k = self.comp_count;

        // TODO: Compute sample covariance of the data.
        // https://en.wikipedia.org/wiki/Sample_mean_and_covariance#Sample_covariance
        let mut cov_vec = Vec::with_capacity(k);
        for _ in 0..k {
            cov_vec.push(Matrix::identity(inputs.cols()));
        }

        self.model_covars = Some(cov_vec);

        let random_rows: Vec<usize> = rand_utils::reservoir_sample(&(0..inputs.rows())
                                                                        .collect::<Vec<usize>>(),
                                                                   k);

        self.model_means = Some(inputs.select_rows(&random_rows));

        for _ in 0..self.max_iters {
            let log_lik_0 = self.log_lik;

            let (weights, log_lik_1) = self.membership_weights(inputs);

            if (log_lik_1 - log_lik_0).abs() < 1e-10 {
                break;
            }

            self.log_lik = log_lik_1;

            self.update_params(inputs, weights);
        }
    }

    /// Predict output from inputs.
    fn predict(&self, inputs: &Matrix<f64>) -> Matrix<f64> {
        if let (&Some(_), &Some(_)) = (&self.model_means, &self.model_covars) {
            self.membership_weights(inputs).0
        } else {
            panic!("Model has not been trained.");
        }

    }
}

impl GaussianMixtureModel {
    /// Constructs a new Gaussian Mixture Model
    ///
    /// Defaults to 100 maximum iterations and
    /// full covariance structure.
    ///
    /// # Examples
    /// ```
    /// use rusty_machine::learning::gmm::GaussianMixtureModel;
    ///
    /// let gmm = GaussianMixtureModel::new(3);
    /// ```
    pub fn new(k: usize) -> GaussianMixtureModel {
        GaussianMixtureModel {
            comp_count: k,
            mix_weights: Vector::ones(k) / (k as f64),
            model_means: None,
            model_covars: None,
            log_lik: 0f64,
            max_iters: 100,
            cov_option: CovOption::Full,
        }
    }

    /// The model means
    ///
    /// Returns a reference to the Option<Matrix<f64>>
    /// which contains the model means. Each row represents
    /// the mean of one of the Gaussians.
    pub fn means(&self) -> &Option<Matrix<f64>> {
        &self.model_means
    }

    /// The model covariances
    ///
    /// Returns a reference to the Option<Vec<Matrix<f64>>>
    /// which contains the model covariances. Each Matrix in
    /// the vector is the covariance of one of the Gaussians.
    pub fn covariances(&self) -> &Option<Vec<Matrix<f64>>> {
        &self.model_covars
    }

    /// Sets the max number of iterations for the EM algorithm.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gmm::GaussianMixtureModel;
    ///
    /// let mut gmm = GaussianMixtureModel::new(2);
    /// gmm.set_max_iters(5);
    /// ```
    pub fn set_max_iters(&mut self, iters: usize) {
        self.max_iters = iters;
    }

    fn membership_weights(&self, inputs: &Matrix<f64>) -> (Matrix<f64>, f64) {
        let n = inputs.rows();

        let mut member_weights_data = Vec::with_capacity(n * self.comp_count);

        // We compute the determinants and inverses now
        let mut cov_sqrt_dets = Vec::with_capacity(self.comp_count);
        let mut cov_invs = Vec::with_capacity(self.comp_count);

        if let Some(ref covars) = self.model_covars {
            for cov in covars {
                // TODO: combine these. We compute det to get the inverse.
                let covar_det = cov.det();
                let covar_inv = cov.inverse();

                cov_sqrt_dets.push(covar_det.sqrt());
                cov_invs.push(covar_inv);
            }
        }

        let mut log_lik = 0f64;

        // Now we compute the membership weights
        if let Some(ref means) = self.model_means {
            for i in 0..n {
                let mut pdfs = Vec::with_capacity(self.comp_count);
                let x_i = inputs.select_rows(&[i]);

                for j in 0..self.comp_count {
                    let mu_j = means.select_rows(&[j]);
                    let diff = &x_i - mu_j;

                    let pdf = (&diff * &cov_invs[j] * diff.transpose() * -0.5).into_vec()[0]
                                  .exp() / cov_sqrt_dets[j];
                    pdfs.push(pdf);
                }

                let weighted_pdf_sum = utils::dot(&pdfs, self.mix_weights.data());

                for (idx, pdf) in pdfs.iter().enumerate() {
                    member_weights_data.push(self.mix_weights[idx] * pdf / (weighted_pdf_sum));
                }

                log_lik += weighted_pdf_sum.ln();
            }
        }

        (Matrix::new(n, self.comp_count, member_weights_data),
         log_lik)
    }

    fn update_params(&mut self, inputs: &Matrix<f64>, membership_weights: Matrix<f64>) {
        let n = membership_weights.rows();
        let d = inputs.cols();

        let sum_weights = membership_weights.sum_rows();

        self.mix_weights = &sum_weights / (n as f64);

        let mut new_means = membership_weights.transpose() * inputs;

        for (idx, mean) in new_means.mut_data().chunks_mut(d).enumerate() {
            for m in mean {
                *m = *m / sum_weights[idx];
            }
        }

        let mut new_covs = Vec::with_capacity(self.comp_count);

        for k in 0..self.comp_count {
            let mut cov_mat = Matrix::zeros(d, d);

            for i in 0..n {
                let diff = inputs.select_rows(&[i]) - new_means.select_rows(&[k]);
                cov_mat = cov_mat + self.compute_cov(diff, membership_weights[[i, k]]);
            }
            new_covs.push(cov_mat / sum_weights[k]);

        }

        self.model_means = Some(new_means);
        self.model_covars = Some(new_covs);
    }

    fn compute_cov(&self, diff: Matrix<f64>, weight: f64) -> Matrix<f64> {
        match self.cov_option {
            CovOption::Full => (diff.transpose() * diff) * weight,
            CovOption::Regularized(eps) => (diff.transpose() * diff) * weight + eps,
            CovOption::Diagonal => Matrix::from_diag(&diff.elemul(&diff).into_vec()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::GaussianMixtureModel;
    #[test]
    fn test_means_none() {
        let model = GaussianMixtureModel::new(5);

        assert_eq!(model.means(), &None);
    }

    #[test]
    fn test_covars_none() {
        let model = GaussianMixtureModel::new(5);

        assert_eq!(model.covariances(), &None);
    }
}<commit_msg>Adding `get mixture weights` to GMM (#78)<commit_after>//! Gaussian Mixture Models
//!
//! Provides implementation of GMMs using the EM algorithm.
//!
//! # Usage
//!
//! ```
//! use rusty_machine::linalg::matrix::Matrix;
//! use rusty_machine::learning::gmm::{CovOption, GaussianMixtureModel};
//! use rusty_machine::learning::UnSupModel;
//!
//! let inputs = Matrix::new(4, 2, vec![1.0, 2.0, 3.0, 3.0, 2.0, 4.0, 5.0, 2.5]);
//! let test_inputs = Matrix::new(3, 2, vec![1.0, 2.0, 3.0, 2.9, 2.4, 2.5]);
//!
//! // Create gmm with k(=2) classes.
//! let mut model = GaussianMixtureModel::new(2);
//! model.set_max_iters(10);
//! model.cov_option = CovOption::Diagonal;
//!
//! // Where inputs is a Matrix with features in columns.
//! model.train(&inputs);
//!
//! // Where pred_data is a Matrix with features in columns.
//! let a = model.predict(&test_inputs);
//! println!("{:?}", a.data());
//! ```

use linalg::vector::Vector;
use linalg::matrix::Matrix;
use linalg::utils;

use learning::UnSupModel;
use learning::toolkit::rand_utils;

/// Covariance options for GMMs.
///
/// - Full : The full covariance structure.
/// - Regularized : Adds a regularization constant to the covariance diagonal.
/// - Diagonal : Only the diagonal covariance structure.
#[derive(Clone, Copy, Debug)]
pub enum CovOption {
    /// The full covariance structure.
    Full,
    /// Adds a regularization constant to the covariance diagonal.
    Regularized(f64),
    /// Only the diagonal covariance structure.
    Diagonal,
}

/// A Gaussian Mixture Model
#[derive(Debug)]
pub struct GaussianMixtureModel {
    comp_count: usize,
    mix_weights: Vector<f64>,
    model_means: Option<Matrix<f64>>,
    model_covars: Option<Vec<Matrix<f64>>>,
    log_lik: f64,
    max_iters: usize,
    /// The covariance options for the GMM.
    pub cov_option: CovOption,
}

impl UnSupModel<Matrix<f64>, Matrix<f64>> for GaussianMixtureModel {
    /// Train the model using inputs.
    fn train(&mut self, inputs: &Matrix<f64>) {
        // Initialization:
        let k = self.comp_count;

        // TODO: Compute sample covariance of the data.
        // https://en.wikipedia.org/wiki/Sample_mean_and_covariance#Sample_covariance
        let mut cov_vec = Vec::with_capacity(k);
        for _ in 0..k {
            cov_vec.push(Matrix::identity(inputs.cols()));
        }

        self.model_covars = Some(cov_vec);

        let random_rows: Vec<usize> = rand_utils::reservoir_sample(&(0..inputs.rows())
                                                                        .collect::<Vec<usize>>(),
                                                                   k);

        self.model_means = Some(inputs.select_rows(&random_rows));

        for _ in 0..self.max_iters {
            let log_lik_0 = self.log_lik;

            let (weights, log_lik_1) = self.membership_weights(inputs);

            if (log_lik_1 - log_lik_0).abs() < 1e-15 {
                break;
            }

            self.log_lik = log_lik_1;

            self.update_params(inputs, weights);
        }
    }

    /// Predict output from inputs.
    fn predict(&self, inputs: &Matrix<f64>) -> Matrix<f64> {
        if let (&Some(_), &Some(_)) = (&self.model_means, &self.model_covars) {
            self.membership_weights(inputs).0
        } else {
            panic!("Model has not been trained.");
        }

    }
}

impl GaussianMixtureModel {
    /// Constructs a new Gaussian Mixture Model
    ///
    /// Defaults to 100 maximum iterations and
    /// full covariance structure.
    ///
    /// # Examples
    /// ```
    /// use rusty_machine::learning::gmm::GaussianMixtureModel;
    ///
    /// let gmm = GaussianMixtureModel::new(3);
    /// ```
    pub fn new(k: usize) -> GaussianMixtureModel {
        GaussianMixtureModel {
            comp_count: k,
            mix_weights: Vector::ones(k) / (k as f64),
            model_means: None,
            model_covars: None,
            log_lik: 0f64,
            max_iters: 100,
            cov_option: CovOption::Full,
        }
    }

    /// The model means
    ///
    /// Returns a reference to the Option<Matrix<f64>>
    /// which contains the model means. Each row represents
    /// the mean of one of the Gaussians.
    pub fn means(&self) -> &Option<Matrix<f64>> {
        &self.model_means
    }

    /// The model covariances
    ///
    /// Returns a reference to the Option<Vec<Matrix<f64>>>
    /// which contains the model covariances. Each Matrix in
    /// the vector is the covariance of one of the Gaussians.
    pub fn covariances(&self) -> &Option<Vec<Matrix<f64>>> {
        &self.model_covars
    }

    /// The model mixture weights
    ///
    /// Returns a reference to the model mixture weights.
    /// These are the weighted contributions of each underlying
    /// Gaussian to the model distribution.
    pub fn mixture_weights(&self) -> &Vector<f64> {
        &self.mix_weights
    }

    /// Sets the max number of iterations for the EM algorithm.
    ///
    /// # Examples
    ///
    /// ```
    /// use rusty_machine::learning::gmm::GaussianMixtureModel;
    ///
    /// let mut gmm = GaussianMixtureModel::new(2);
    /// gmm.set_max_iters(5);
    /// ```
    pub fn set_max_iters(&mut self, iters: usize) {
        self.max_iters = iters;
    }

    fn membership_weights(&self, inputs: &Matrix<f64>) -> (Matrix<f64>, f64) {
        let n = inputs.rows();

        let mut member_weights_data = Vec::with_capacity(n * self.comp_count);

        // We compute the determinants and inverses now
        let mut cov_sqrt_dets = Vec::with_capacity(self.comp_count);
        let mut cov_invs = Vec::with_capacity(self.comp_count);

        if let Some(ref covars) = self.model_covars {
            for cov in covars {
                // TODO: combine these. We compute det to get the inverse.
                let covar_det = cov.det();
                let covar_inv = cov.inverse();

                cov_sqrt_dets.push(covar_det.sqrt());
                cov_invs.push(covar_inv);
            }
        }

        let mut log_lik = 0f64;

        // Now we compute the membership weights
        if let Some(ref means) = self.model_means {
            for i in 0..n {
                let mut pdfs = Vec::with_capacity(self.comp_count);
                let x_i = inputs.select_rows(&[i]);

                for j in 0..self.comp_count {
                    let mu_j = means.select_rows(&[j]);
                    let diff = &x_i - mu_j;

                    let pdf = (&diff * &cov_invs[j] * diff.transpose() * -0.5).into_vec()[0]
                                  .exp() / cov_sqrt_dets[j];
                    pdfs.push(pdf);
                }

                let weighted_pdf_sum = utils::dot(&pdfs, self.mix_weights.data());

                for (idx, pdf) in pdfs.iter().enumerate() {
                    member_weights_data.push(self.mix_weights[idx] * pdf / (weighted_pdf_sum));
                }

                log_lik += weighted_pdf_sum.ln();
            }
        }

        (Matrix::new(n, self.comp_count, member_weights_data),
         log_lik)
    }

    fn update_params(&mut self, inputs: &Matrix<f64>, membership_weights: Matrix<f64>) {
        let n = membership_weights.rows();
        let d = inputs.cols();

        let sum_weights = membership_weights.sum_rows();

        self.mix_weights = &sum_weights / (n as f64);

        let mut new_means = membership_weights.transpose() * inputs;

        for (idx, mean) in new_means.mut_data().chunks_mut(d).enumerate() {
            for m in mean {
                *m = *m / sum_weights[idx];
            }
        }

        let mut new_covs = Vec::with_capacity(self.comp_count);

        for k in 0..self.comp_count {
            let mut cov_mat = Matrix::zeros(d, d);

            for i in 0..n {
                let diff = inputs.select_rows(&[i]) - new_means.select_rows(&[k]);
                cov_mat = cov_mat + self.compute_cov(diff, membership_weights[[i, k]]);
            }
            new_covs.push(cov_mat / sum_weights[k]);

        }

        self.model_means = Some(new_means);
        self.model_covars = Some(new_covs);
    }

    fn compute_cov(&self, diff: Matrix<f64>, weight: f64) -> Matrix<f64> {
        match self.cov_option {
            CovOption::Full => (diff.transpose() * diff) * weight,
            CovOption::Regularized(eps) => (diff.transpose() * diff) * weight + eps,
            CovOption::Diagonal => Matrix::from_diag(&diff.elemul(&diff).into_vec()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::GaussianMixtureModel;
    #[test]
    fn test_means_none() {
        let model = GaussianMixtureModel::new(5);

        assert_eq!(model.means(), &None);
    }

    #[test]
    fn test_covars_none() {
        let model = GaussianMixtureModel::new(5);

        assert_eq!(model.covariances(), &None);
    }
}<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);
        assert!(*prob_total <= 0.0, "prob_total > 0 ({})", prob_total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        debug!("{:?}", event_probs);

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        *self.prob_dropout
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<commit_msg>Minor.<commit_after>use std::cell::RefCell;
use std::cmp;

use rand;
use rand::Rng;
use itertools::Itertools;
use rgsl::randist::multinomial::multinomial_pdf;

use bio::stats::{Prob, LogProb};

use io::codebook::{Codebook, FeatureID};

use model::readout::{Expressions, Counts, Miscalls, Xi};

pub trait AbstractFeatureModel {
    /// FeatureID of this feature.
    fn feature_id(&self) -> FeatureID;

    /// Neighbors of this feature.
    fn neighbors(&self) -> &[FeatureID];

    /// Minimum hamming distance between any pair of features.
    fn min_dist(&self) -> u8;

    /// Number of events (calls and miscalls) except dropout.
    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32;

    /// Dropout probability.
    fn prob_dropout(&self) -> f64;

    /// Probability of exact miscall to given neighbor index.
    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64;

    /// Probability of mismatching miscall to given neighbor index.
    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64;

    /// Calculate maximum likelihood estimate of miscalls given expression and store in the given
    /// arrays.
    /// Returns the total absolute changes to estimated miscall counts.
    fn mle_miscalls(
        &self,
        expressions: &Expressions,
        miscalls_exact: &mut Miscalls,
        miscalls_mismatch: &mut Miscalls,
        rng: &mut rand::StdRng
    ) {
        let x = expressions[self.feature_id()];

        // We need to ensure that the sum of the estimated miscalls does not exceed x.
        // Hence, we iterate randomly over the neighbors such that we can simply stop once that
        // happens without causing a bias.
        let mut idx = (0..self.neighbors().len()).collect_vec();
        rng.shuffle(&mut idx);

        // by stochastic rounding, we ensure that on average \sum x*p_i = x
        // numeric rounding would not work if many of the probabilities are small enough to yield
        // expectations less than 1
        let mut stochastic_round = |v: f64| v.floor() as u32 + (rng.next_f64() <= v % 1.0) as u32;

        let mut rest = x;
        for i in idx {
            let n = self.neighbors()[i];
            let miscall_exact = self.prob_miscall_exact(i) * x as f64;
            miscalls_exact.set(
                self.feature_id(), n, stochastic_round(miscall_exact), &mut rest
            );

            if self.min_dist() == 4 {
                let miscall_mismatch = self.prob_miscall_mismatch(i) * x as f64;
                let miscall_mismatch = stochastic_round(miscall_mismatch);
                miscalls_mismatch.set(
                    self.feature_id(), n, miscall_mismatch, &mut rest
                );
            }
        }
    }

    /// Calculate maximum likelihood estimate of expression given miscalls.
    /// Returns the absolute amount of change to the expression.
    fn mle_expression(
        &self,
        expressions: &mut Expressions,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> i32 {
        let event_count = self.event_count_non_dropout(miscalls_exact, miscalls_mismatch);

        let x_ = cmp::max(
            (event_count as f64 / (1.0 - self.prob_dropout())).floor() as u32,
            // or at least as many as we have events coming from this feature
            // TODO does this allow x to decrease properly?
            event_count
        );

        let x = expressions.get_mut(self.feature_id()).unwrap();

        let change = *x as i32 - x_ as i32;

        *x = x_;

        change
    }
}

pub struct FeatureModel {
    pub(crate) feature_id: FeatureID,
    event_probs: Vec<f64>,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
    pub(crate) counts: Counts,
    event_counts: RefCell<Vec<u32>>
}

impl FeatureModel {
    pub fn new(
        feature: FeatureID,
        counts: Counts,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        let feature_record = codebook.record(feature);

        let xi_call = xi.prob(feature_record.codeword(), feature_record.codeword());

        let prob_call_exact = xi_call[0];
        let prob_call_mismatch = if codebook.min_dist == 4 {
            xi_call[1]
        } else {
            LogProb::ln_zero()
        };

        let neighbors_4 = codebook.neighbors(feature, 4);
        let neighbors_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2)
        } else {
            vec![]
        };

        let mut neighbors = neighbors_4.clone();
        neighbors.extend(&neighbors_2);

        let xi_miscall_4 = codebook.neighbors(feature, 4).iter().map(|&n| {
            xi.prob(feature_record.codeword(), codebook.record(n).codeword())
        }).collect_vec();
        let xi_miscall_2 = if codebook.min_dist == 2 {
            codebook.neighbors(feature, 2).iter().map(|&n| {
                xi.prob(feature_record.codeword(), codebook.record(n).codeword())
            }).collect_vec()
        } else {
            vec![]
        };

        let prob_miscall_exact = xi_miscall_4.iter().map(|xi| {
            xi[0]
        }).chain(xi_miscall_2.iter().map(|xi| xi[0])).collect_vec();

        let prob_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall_4.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let mut total = vec![prob_call_exact, prob_call_mismatch];
        total.extend(&prob_miscall_exact);
        total.extend(&prob_miscall_mismatch);

        let prob_total = LogProb::ln_sum_exp(&total);
        assert!(*prob_total <= 0.0, "prob_total > 0 ({})", *prob_total);

        let mut event_probs = vec![
            *Prob::from(prob_call_exact),
            *Prob::from(prob_call_mismatch),
            *Prob::from(prob_total.ln_one_minus_exp())
        ];
        event_probs.extend(prob_miscall_exact.iter().map(|&p| *Prob::from(p)));
        event_probs.extend(prob_miscall_mismatch.iter().map(|&p| *Prob::from(p)));

        debug!("{:?}", event_probs);

        FeatureModel {
            feature_id: feature,
            event_probs: event_probs,
            neighbors: neighbors,
            min_dist: codebook.min_dist,
            counts: counts,
            event_counts: RefCell::new(Vec::new())
        }
    }

    /// Calculate likelihood of given expression.
    pub fn likelihood(
        &self,
        x: u32,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> LogProb {
        let count_exact_miscalled = miscalls_exact.total_to(self.feature_id);
        let count_mismatch_miscalled = miscalls_mismatch.total_to(self.feature_id);

        if self.counts.exact < count_exact_miscalled {
            panic!("bug: likelihood cannot be calculated if exact miscalls exceed counts: counts={}, miscalls={}", self.counts.exact, count_exact_miscalled);
        }
        if self.counts.mismatch < count_mismatch_miscalled {
            panic!("bug: likelihood cannot be calculated if mismatch miscalls exceed counts: counts={}, miscalls={}", self.counts.mismatch, count_mismatch_miscalled);
        }
        let total_count = self.counts.exact - count_exact_miscalled +
                          self.counts.mismatch - count_mismatch_miscalled;
        if total_count > x {
            panic!("bug: event count exceeds expression: x={}, event count={}", x, total_count);
        }

        let calls_exact = self.counts.exact - count_exact_miscalled;
        let calls_mismatch = self.counts.mismatch - count_mismatch_miscalled;

        // setup event counts
        {
            let mut event_counts = self.event_counts.borrow_mut();
            event_counts.clear();
            event_counts.push(
                // exact calls
                calls_exact
            );
            event_counts.push(
                // mismatch calls
                calls_mismatch
            );
            event_counts.push(
                // dropouts (fill in later)
                0
            );
            // miscalls
            event_counts.extend(self.neighbors.iter().map(
                |&n| miscalls_exact.get(self.feature_id, n)
            ));
            if self.min_dist == 4 {
                event_counts.extend(
                    self.neighbors.iter().map(
                        |&n| miscalls_mismatch.get(self.feature_id, n)
                    )
                );
            }

            // set dropouts
            let total_counts = event_counts.iter().sum();
            assert!(x >= total_counts, "total event counts exceed expression: x={} {:?}", x, event_counts);
            let dropouts = x - total_counts;
            event_counts[2] = dropouts;
        }

        LogProb::from(Prob(
            multinomial_pdf(&self.event_probs, &self.event_counts.borrow())
        ))
    }

    /// Minimum expression given fixed miscalls.
    pub fn min_expression(&self, miscalls_exact: &Miscalls, miscalls_mismatch: &Miscalls) -> u32 {
        self.counts.exact.saturating_sub(miscalls_exact.total_to(self.feature_id)) +
        self.counts.mismatch.saturating_sub(miscalls_mismatch.total_to(self.feature_id)) +
        miscalls_exact.total_from(self.feature_id) + miscalls_mismatch.total_from(self.feature_id)
    }
}

impl AbstractFeatureModel for FeatureModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let calls_exact = self.counts.exact.saturating_sub(
            miscalls_exact.total_to(self.feature_id)
        );
        let calls_mismatch = self.counts.mismatch.saturating_sub(
            miscalls_mismatch.total_to(self.feature_id)
        );

        calls_exact + calls_mismatch +
        miscalls_exact.total_from(self.feature_id) +
        miscalls_mismatch.total_from(self.feature_id)
    }

    fn prob_dropout(&self) -> f64 {
        self.event_probs[2]
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        self.event_probs[3 + self.neighbors.len() + neighbor_index]
    }
}

pub struct NoiseModel {
    pub(crate) feature_id: FeatureID,
    pub(crate) not_expressed_feature_ids: Vec<FeatureID>,
    pub(crate) not_expressed_counts: Vec<Counts>,
    probs_miscall_exact: Vec<Prob>,
    probs_miscall_mismatch: Vec<Prob>,
    prob_dropout: Prob,
    neighbors: Vec<FeatureID>,
    min_dist: u8,
}

impl NoiseModel {
    pub fn new(
        feature_id: FeatureID,
        not_expressed_feature_ids: Vec<FeatureID>,
        not_expressed_counts: Vec<Counts>,
        codebook: &Codebook,
        xi: &Xi
    ) -> Self {
        assert!(
            codebook.min_dist == 2 || codebook.min_dist == 4,
            "unsupported hamming distance: {}, supported: 2, 4",
            codebook.min_dist
        );

        // miscalls generated from noise
        let xi_miscall = codebook.records().iter().filter_map(|feat_rec| {
            if feat_rec.expressed() {
                Some(xi.prob(codebook.noise().codeword(), feat_rec.codeword()))
            } else { None }
        }).collect_vec();
        let neighbors = codebook.records().iter().filter_map(|rec| {
            if rec.expressed() {
                Some(codebook.get_id(rec.name()))
            } else { None }
        }).collect_vec();

        let probs_miscall_exact = xi_miscall.iter().map(|xi| xi[0]).collect_vec();
        let probs_miscall_mismatch = if codebook.min_dist == 4 {
            xi_miscall.iter().map(|xi| xi[1]).collect_vec()
        } else {
            vec![]
        };

        let prob_total_miscall_exact = LogProb::ln_sum_exp(&probs_miscall_exact);
        let prob_total_miscall_mismatch = LogProb::ln_sum_exp(&probs_miscall_mismatch);

        // probs for calls of not expressed features
        let xi_not_expressed = not_expressed_feature_ids.iter().map(|&f| {
            xi.prob(codebook.noise().codeword(), codebook.record(f).codeword())
        }).collect_vec();

        let prob_not_expressed_exact = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[0]).collect_vec()
        );
        let prob_not_expressed_mismatch = LogProb::ln_sum_exp(
            &xi_not_expressed.iter().map(|xi| xi[1]).collect_vec()
        );

        // dropout
        let prob_dropout = LogProb::ln_sum_exp(&[
            prob_not_expressed_exact,
            prob_not_expressed_mismatch,
            prob_total_miscall_exact,
            prob_total_miscall_mismatch
        ]).ln_one_minus_exp();

        NoiseModel {
            feature_id: feature_id,
            not_expressed_feature_ids: not_expressed_feature_ids,
            not_expressed_counts: not_expressed_counts,
            probs_miscall_exact: probs_miscall_exact.iter().map(|&p| Prob::from(p)).collect_vec(),
            probs_miscall_mismatch: probs_miscall_mismatch.iter().map(|&p| Prob::from(p)).collect_vec(),
            prob_dropout: Prob::from(prob_dropout),
            neighbors: neighbors,
            min_dist: codebook.min_dist
        }
    }
}

impl AbstractFeatureModel for NoiseModel {
    fn feature_id(&self) -> FeatureID {
        self.feature_id
    }

    fn neighbors(&self) -> &[FeatureID] {
        &self.neighbors
    }

    fn min_dist(&self) -> u8 {
        self.min_dist
    }

    fn event_count_non_dropout(
        &self,
        miscalls_exact: &Miscalls,
        miscalls_mismatch: &Miscalls
    ) -> u32 {
        let mut calls_exact = 0;
        let mut calls_mismatch = 0;
        for (&id, counts) in self.not_expressed_feature_ids.iter().zip(&self.not_expressed_counts) {
            calls_exact += counts.exact.saturating_sub(miscalls_exact.total_to(id));
            calls_mismatch += counts.mismatch.saturating_sub(miscalls_mismatch.total_to(id));
        }

        let event_count = calls_exact + calls_mismatch +
                          miscalls_exact.total_from(self.feature_id) +
                          miscalls_mismatch.total_from(self.feature_id);

        event_count
    }

    fn prob_dropout(&self) -> f64 {
        *self.prob_dropout
    }

    fn prob_miscall_exact(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_exact[neighbor_index]
    }

    fn prob_miscall_mismatch(&self, neighbor_index: usize) -> f64 {
        *self.probs_miscall_mismatch[neighbor_index]
    }
}
<|endoftext|>