better pr curve, plotter, store fixes

src/store.rs

@@ -18,12 +18,55 @@ pub enum SaveError {
 }
 
 /// Struct for calculating Recall-Precision per mutation
-struct PRStats {
+pub struct PRStats {
 	true_positives: [u64; 64],
 	false_positives: [u64; 64],
 	false_negatives: [u64; 64],
-	tag: String,
-	name: String
+	pub tag: String,
+	pub name: String
+}
+
+impl PRStats {
+	pub fn precision(&self, threshold: usize) -> f64 {
+		let t = match threshold {
+			0..=64 => threshold,
+			_ => std::cmp::max(0, std::cmp::min(64, threshold))
+		};
+		let tpos = self.true_positives[t] as f64;
+		let fpos = self.false_positives[t] as f64;
+		let p = tpos / (tpos + fpos);
+		p
+	}
+	pub fn recall(&self, threshold: usize) -> f64 {
+		let t = match threshold {
+			0..=64 => threshold,
+			_ => std::cmp::max(0, std::cmp::min(64, threshold))
+		};
+		let tpos = self.true_positives[t] as f64;
+		let fneg = self.false_negatives[t] as f64;
+		let r = tpos / (tpos + fneg);
+		r
+	}
+	pub fn pr(&self, threshold: usize) -> (f64, f64) {
+		let t = match threshold {
+			0..=64 => threshold,
+			_ => std::cmp::max(0, std::cmp::min(64, threshold))
+		};
+		let tpos = self.true_positives[t] as f64;
+		let fpos = self.false_positives[t] as f64;
+		let fneg = self.false_negatives[t] as f64;
+		let p = tpos / (tpos + fpos);
+		let r = tpos / (tpos + fneg);
+		(r, p)
+	}
+	pub fn pr_curve(&self) -> Vec<(f64, f64)> {
+		let mut curve = Vec::new();
+		for i in 0..64 {
+			curve.push(self.pr(i));
+		}
+		curve
+	}
+
 }
 
 /// Uses a hashmap to map descriptors to buckets of files.
@@ -71,7 +114,9 @@ impl DescriptorStore {
 				let mut seen: HashSet<String> = HashSet::new();
 				for bucket in self.map.values() {
 					for element in bucket {
-						seen.insert(element.clone());
+						if !element.starts_with("mut.") {
+							seen.insert(element.clone());
+						}
 					}
 				}
 				self.seen = Some(seen);
@@ -199,15 +244,15 @@ impl DescriptorStore {
 	}
 
 	/// Nearest neighbours
-	pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
-		let mut neighbours = self.bktree.find(from, max_distance);
+	pub fn nn(&self, from: u64, max_distance: usize) -> Vec<(&u64, isize)> {
+		let mut neighbours = self.bktree.find(from, max_distance.try_into().unwrap());
 		neighbours.sort_by(|a, b| a.1.cmp(&b.1));
 		neighbours
 	}
 
 
 	/// Returns a flat vector with all the filenames that are within a given distance
-	pub fn nn_flat_results(&self, from: u64, max_distance: isize) -> Vec<String> {
+	pub fn nn_flat_results(&self, from: u64, max_distance: usize) -> Vec<String> {
 		let nn = self.nn(from, max_distance);
 		let mut results: Vec<String> = Vec::new();
 		for (key, _) in nn {
@@ -218,7 +263,7 @@ impl DescriptorStore {
 		results
 	}
 
-	pub fn print_nn(&self, from: u64, max_distance: isize, show_images: bool) {
+	pub fn print_nn(&self, from: u64, max_distance: usize, show_images: bool) {
 		//println!("{from:064b}");
 		let neighbours = self.nn(from, max_distance);
 		let (term_width, _) = viuer::terminal_size();
@@ -280,7 +325,8 @@ impl DescriptorStore {
 		print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
 	}
 
-	pub fn print_stats(&self) {
+	pub fn get_stats(&self, mut max_threshold: usize) -> Vec<PRStats> {
+		max_threshold = std::cmp::min(64, max_threshold);
 		let mut mutator_stats = Vec::new();
 		for mutator in get_all_mutators() {
 			mutator_stats.push(
@@ -298,7 +344,7 @@ impl DescriptorStore {
 		for image in set.iter() {
 			let phash = self.get(image).unwrap();
 			debug!("Checking {}, with phash: {}", image, phash);
-			for threshold in 0..64 {
+			for threshold in 0..max_threshold {
 				debug!("Threshold: {}", threshold);
 				//Assume miss, therefore a false negative
 				//Undo the miss when there is a true positive
@@ -328,8 +374,13 @@ impl DescriptorStore {
 				}
 			}
 		}
+		mutator_stats
+	}
 
+	pub fn print_stats(&self) {
+		let mut mutator_stats = self.get_stats(64);
 		for mutator in &mut mutator_stats {
+			println!("{}:", mutator.name);
 			for threshold in 0..64 {
 				let tpos = mutator.true_positives[threshold] as f64;
 				let fpos = mutator.false_positives[threshold] as f64;
@@ -337,27 +388,11 @@ impl DescriptorStore {
 				let p = tpos / (tpos + fpos);
 				let r = tpos / (tpos + fneg);
 				let f1 = 2.0*tpos / (2.0*tpos + fpos + fneg);
-				println!("{} / {}:", mutator.name, threshold);
-				println!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
-				println!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
+				debug!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
+				debug!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
+				println!("{}, {}" , p, r)
 			}
-			// println!("\t{} true positives: {}", mutator.name, mutator.true_positives);
-			// println!("\t{} false positives: {}", mutator.name, mutator.false_positives);
-			// println!(
-			// 	"\t{} Precision: {:.3}", mutator.name, 
-			// 	mutator.true_positives as f64 / 
-			// 	(mutator.true_positives + mutator.false_positives) as f64
-			// );
-			// println!(
-			// 	"\t{} Recall: {:.3}", mutator.name, 
-			// 	mutator.true_positives as f64 / 
-			// 	//(mutator.true_positives + mutator.false_negatives) as f64
-			// 	set.len() as f64
-			// );
 		}
-		//println!("Images processed: {}", set.len())
-		// println!("{} true positives", true_positives);
-		// println!("{} false positives", false_positives);
 	}
 }