Add ability to load epistatic effects

Author: sirno

Cargo.toml

@@ -27,7 +27,7 @@ evalexpr = "8.2.0"
 indicatif = "0.17.5"
 itertools = "0.10.5"
 log = "0.4.19"
-npyz = "0.7.4"
+npyz = { version = "0.8.3", features = ["derive"] }
 parking_lot = "0.12.1"
 phf = { version = "0.11.2", features = ["macros"] }
 rand = { version = "0.8.5", features = ["alloc"] }

src/core/fitness/epistasis.rs

@@ -4,31 +4,51 @@ use crate::core::haplotype::Symbol;
 use crate::errors::VirolutionError;
 use crate::references::HaplotypeRef;
 
-use super::init::FitnessModel;
+use super::init::{FitnessDistribution, FitnessModel};
 
 type EpistasisTableKey = (usize, Symbol);
 
+#[derive(npyz::Deserialize)]
+pub struct EpiEntry {
+    pos1: u64,
+    base1: u8,
+    pos2: u64,
+    base2: u8,
+    value: f64,
+}
+
 #[derive(Clone, Debug)]
 pub struct EpistasisTable {
     table: HashMap<EpistasisTableKey, HashMap<EpistasisTableKey, f64>>,
 }
 
 impl EpistasisTable {
     pub fn from_model(model: &FitnessModel) -> Result<Self, VirolutionError> {
-        todo!()
+        let table = match &model.distribution {
+            FitnessDistribution::Epistatic(_, epi_params) => {
+                let entries = epi_params.load_table();
+                Self::from_vec(entries)
+            }
+            _ => {
+                return Err(VirolutionError::ImplementationError(
+                    "Model is not epistatic".to_string(),
+                ));
+            }
+        };
+        Ok(table)
     }
 
-    pub fn from_vec(entries: Vec<(usize, Symbol, usize, Symbol, f64)>) -> Self {
+    pub fn from_vec(entries: Vec<EpiEntry>) -> Self {
         let mut table: HashMap<EpistasisTableKey, HashMap<EpistasisTableKey, f64>> = HashMap::new();
-        for (pos1, base1, pos2, base2, value) in entries.iter() {
+        for entry in entries.iter() {
             table
-                .entry((*pos1, *base1))
+                .entry((entry.pos1 as usize, Some(entry.base1)))
                 .or_default()
-                .insert((*pos2, *base2), *value);
+                .insert((entry.pos2 as usize, Some(entry.base2)), entry.value);
             table
-                .entry((*pos2, *base2))
+                .entry((entry.pos2 as usize, Some(entry.base2)))
                 .or_default()
-                .insert((*pos1, *base1), *value);
+                .insert((entry.pos1 as usize, Some(entry.base1)), entry.value);
         }
         Self { table }
     }

src/core/fitness/init.rs

@@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::core::haplotype::Symbol;
 
+use super::epistasis::EpiEntry;
 use super::utility::UtilityFunction;
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
@@ -167,3 +168,14 @@ impl FileParameters {
             .collect()
     }
 }
+
+impl EpiFileParameters {
+    pub fn load_table(&self) -> Vec<EpiEntry> {
+        let reader = NpyFile::new(std::fs::File::open(&self.path).unwrap()).unwrap();
+        reader
+            .data::<EpiEntry>()
+            .unwrap()
+            .map(|entry| entry.unwrap())
+            .collect()
+    }
+}

src/core/fitness/table.rs

@@ -29,13 +29,13 @@ impl FitnessTable {
     ) -> Result<Self, VirolutionError> {
         let n_sites = sequence.len();
         let table = match fitness_model.distribution {
+            FitnessDistribution::Neutral => vec![1.; n_sites * n_symbols],
             FitnessDistribution::Exponential(ref params) => {
                 params.create_table(n_symbols, sequence)
             }
             FitnessDistribution::Lognormal(ref params) => params.create_table(n_symbols, sequence),
             FitnessDistribution::File(ref params) => params.load_table(),
-            FitnessDistribution::Neutral => vec![1.; n_sites * n_symbols],
-            FitnessDistribution::Epistatic(_, _) => todo!(),
+            FitnessDistribution::Epistatic(ref params, _) => params.load_table(),
         };
 
         if table.len() != n_sites * n_symbols {