Merge pull request #3 from databio/core_utils

Core utilities

Author: nleroy917

bindings/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "genimtools-py"
-version = "0.0.4"
+version = "0.0.7"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -9,5 +9,5 @@ name = "genimtools"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = "0.19.0"
 genimtools = { path = "../genimtools" }
+pyo3 = "0.20.0"

bindings/README.md

@@ -14,3 +14,5 @@ import genimtools as gt
 
 gt.prune_universe(...)
 ```
+## Developer docs
+Write the develop docs here...

bindings/genimtools/__init__.py

@@ -0,0 +1 @@
+from .genimtools import *

bindings/genimtools/__init__.pyi

@@ -0,0 +1,7 @@
+PAD_CHR: str
+PAD_START: int
+PAD_END: int
+
+UNKNOWN_CHR: str
+UNKNOWN_START: int
+UNKNOWN_END: int

bindings/genimtools/tokenizers/__init__.py

@@ -0,0 +1 @@
+from .genimtools.tokenizers import *

bindings/genimtools/tokenizers/__init__.pyi

@@ -0,0 +1,38 @@
+from typing import List
+
+class Region:
+    chr: str
+    start: int
+    end: int
+
+class TokenizedRegion:
+    chr: str
+    start: int
+    end: int
+    id: int
+
+class TokenizedRegionSet:
+    def regions(self) -> List[TokenizedRegion]:
+        """
+        Get the list of regions in this set.
+        """
+        
+    def ids(self) -> List[int]:
+        """
+        Get the list of IDs in this set.
+        """
+
+class TreeTokenizer:
+    def __init__(self, universe: str) -> None:
+        """
+        Instatiate a new TreeTokenizer. This tokenizer takes advantage of
+        interval trees to compute genomic interval overlaps efficiently.
+        
+        :param universe: The universe of characters to use for tokenization.
+        """
+        pass
+    
+    def tokenize(self, regions: List[Region]) -> TokenizedRegionSet:
+        """
+        Tokenize a list of regions.
+        """

bindings/genimtools/vocab/__init__.py

@@ -0,0 +1 @@
+from .genimtools.vocab import *

bindings/genimtools/vocab/__init__.pyi

@@ -0,0 +1,13 @@
+from typing import Optional
+
+def prune_universe(data: str, universe: str, min_count: Optional[int], output: Optional[str]) -> None:
+    """
+    Given a universe and a path to its corresponding data, "prune" the universe to remove regions that
+    are not present in the data. This is useful for reducing the size of the universe to remove unused
+    regions. It also can be used to remove regions that fall below a certain frequency threshold.
+
+    :param str data: The path to the data file.
+    :param str universe: The path to the universe file.
+    :param int min_count: The minimum number of times a region must appear in the data to be kept.
+    :param str output: The path to the output file.
+    """

bindings/pyproject.toml

@@ -14,8 +14,6 @@ dynamic = ["version"]
 
 [tool.maturin]
 features = ["pyo3/extension-module"]
-python-source = "python" # necessary to distract maturin from the stub files
-module-name = "my_project._lib_name"
 
 [tool.setuptools.dynamic]
 readme = {file = "README.md", content-type = "text/markdown"}

bindings/src/consts.rs

@@ -0,0 +1 @@
+pub use genimtools::common::consts::*;

bindings/src/lib.rs

@@ -1,34 +1,38 @@
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 
-use crate::functions::{prune_universe, print_uniwig};
+mod vocab;
+mod tokenizers;
+mod models;
+mod consts;
 
-mod functions;
+pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 
 #[pymodule]
-fn vocab(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_wrapped(wrap_pyfunction!(prune_universe))?;
-    Ok(())
-}
-
-#[pymodule]
-fn uniwig(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_wrapped(wrap_pyfunction!(print_uniwig))?;
-    Ok(())
-}
+fn genimtools(py: Python, m: &PyModule) -> PyResult<()> {
 
-#[pymodule]
-fn genimtools(py: Python, m: &PyModule) -> PyResult<()> {    
-    let vocab_module = pyo3::wrap_pymodule!(vocab);
-    let uniwig_module = pyo3::wrap_pymodule!(uniwig);
+    let vocab_module = pyo3::wrap_pymodule!(vocab::vocab);
+    let tokenize_module = pyo3::wrap_pymodule!(tokenizers::tokenizers);
 
     m.add_wrapped(vocab_module)?;
-    m.add_wrapped(uniwig_module)?;
+    m.add_wrapped(tokenize_module)?;
 
     let sys = PyModule::import(py, "sys")?;
     let sys_modules: &PyDict = sys.getattr("modules")?.downcast()?;
+
+    // set names of submodules
     sys_modules.set_item("genimtools.vocab", m.getattr("vocab")?)?;
-    sys_modules.set_item("genimtools.uniwig", m.getattr("uniwig")?)?;
+    sys_modules.set_item("genimtools.tokenizers", m.getattr("tokenizers")?)?;
+
+    // add constants
+    m.add("PAD_CHR", consts::PAD_CHR)?;
+    m.add("PAD_START", consts::PAD_START)?;
+    m.add("PAD_END", consts::PAD_END)?;
+    m.add("UNKNOWN_CHR", consts::UNKNOWN_CHR)?;
+    m.add("UNKNOWN_START", consts::UNKNOWN_START)?;
+    m.add("UNKNOWN_END", consts::UNKNOWN_END)?;
+    m.add("__version__", VERSION)?;
+
     Ok(())
 }
 

bindings/src/models/mod.rs

@@ -0,0 +1,7 @@
+mod region;
+mod region_set;
+mod universe;
+
+pub use self::region::{PyRegion, PyTokenizedRegion};
+pub use self::region_set::PyTokenizedRegionSet;
+pub use self::universe::PyUniverse;

bindings/src/models/region.rs

@@ -0,0 +1,112 @@
+use std::hash::Hash;
+
+use pyo3::class::basic::CompareOp;
+use pyo3::exceptions::PyTypeError;
+use pyo3::prelude::*;
+
+use genimtools::common::models::region::Region;
+
+#[pyclass(name = "Region")]
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+pub struct PyRegion {
+    pub chr: String,
+    pub start: u32,
+    pub end: u32,
+}
+
+impl PyRegion {
+    pub fn to_region(&self) -> Region {
+        Region {
+            chr: self.chr.clone(),
+            start: self.start,
+            end: self.end,
+        }
+    }
+}
+
+#[pymethods]
+impl PyRegion {
+    #[new]
+    pub fn new(chr: String, start: u32, end: u32) -> Self {
+        PyRegion { chr, start, end }
+    }
+
+    #[getter]
+    pub fn chr(&self) -> PyResult<&str> {
+        Ok(&self.chr)
+    }
+
+    #[getter]
+    pub fn start(&self) -> PyResult<u32> {
+        Ok(self.start)
+    }
+
+    #[getter]
+    pub fn end(&self) -> PyResult<u32> {
+        Ok(self.end)
+    }
+    pub fn __repr__(&self) -> String {
+        format!("Region({}, {}, {})", self.chr, self.start, self.end)
+    }
+
+    pub fn __richcmp__(&self, other: PyRef<PyRegion>, op: CompareOp) -> PyResult<bool> {
+        match op {
+            CompareOp::Eq => {
+                Ok(self.chr == other.chr && self.start == other.start && self.end == other.end)
+            }
+            CompareOp::Ne => {
+                Ok(self.chr != other.chr || self.start != other.start || self.end != other.end)
+            }
+            _ => Err(PyTypeError::new_err("Unsupported comparison operator")),
+        }
+    }
+}
+
+#[pyclass(name = "TokenizedRegion")]
+#[derive(Clone, Debug)]
+pub struct PyTokenizedRegion {
+    pub region: PyRegion,
+    pub id: u32,
+}
+
+#[pymethods]
+impl PyTokenizedRegion {
+    #[new]
+    pub fn new(region: PyRegion, id: u32) -> Self {
+        PyTokenizedRegion {
+            region,
+            id,
+        }
+    }
+
+    #[getter]
+    pub fn chr(&self) -> PyResult<&str> {
+        Ok(&self.region.chr)
+    }
+
+    #[getter]
+    pub fn start(&self) -> PyResult<u32> {
+        Ok(self.region.start)
+    }
+
+    #[getter]
+    pub fn end(&self) -> PyResult<u32> {
+        Ok(self.region.end)
+    }
+
+    #[getter]
+    pub fn region(&self) -> PyResult<PyRegion> {
+        Ok(self.region.clone())
+    }
+    #[getter]
+    pub fn id(&self) -> PyResult<u32> {
+        Ok(self.id)
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!(
+            "TokenizedRegion({}, {}, {})",
+            self.region.chr, self.region.start, self.region.end
+        )
+    }
+}

bindings/src/models/region_set.rs

@@ -0,0 +1,92 @@
+use pyo3::exceptions::PyIndexError;
+use pyo3::prelude::*;
+
+use genimtools::common::consts::{PAD_CHR, PAD_START, PAD_END};
+
+use crate::models::{PyRegion, PyTokenizedRegion};
+
+#[pyclass(name = "TokenizedRegionSet")]
+#[derive(Clone, Debug)]
+pub struct PyTokenizedRegionSet {
+    pub regions: Vec<PyRegion>,
+    pub ids: Vec<u32>,
+    curr: usize,
+}
+
+#[pymethods]
+impl PyTokenizedRegionSet {
+    #[new]
+    pub fn new(regions: Vec<PyRegion>, ids: Vec<u32>) -> Self {
+        PyTokenizedRegionSet {
+            regions,
+            ids,
+            curr: 0,
+        }
+    }
+
+    #[getter]
+    pub fn regions(&self) -> PyResult<Vec<PyRegion>> {
+        Ok(self.regions.to_owned())
+    }
+
+    #[getter]
+    pub fn ids(&self) -> PyResult<Vec<u32>> {
+        Ok(self.ids.clone())
+    }
+
+    // this is wrong: the padding token might not be in the universe
+    pub fn pad(&mut self, len: usize) {
+        let pad_region = PyRegion {
+            chr: PAD_CHR.to_string(),
+            start: PAD_START as u32,
+            end: PAD_END as u32,
+        };
+        let pad_id = self.ids[0];
+        let pad_region_set = PyTokenizedRegionSet {
+            regions: vec![pad_region; len],
+            ids: vec![pad_id; len],
+            curr: 0,
+        };
+        self.regions.extend(pad_region_set.regions);
+        self.ids.extend(pad_region_set.ids);
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!("TokenizedRegionSet({} regions)", self.regions.len())
+    }
+
+    pub fn __len__(&self) -> usize {
+        self.regions.len()
+    }
+
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    pub fn __next__(&mut self) -> Option<PyTokenizedRegion> {
+        if self.curr < self.regions.len() {
+            let region = self.regions[self.curr].clone();
+            let id = self.ids[self.curr];
+
+            self.curr += 1;
+            Some(PyTokenizedRegion::new(region, id))
+        } else {
+            None
+        }
+    }
+
+    pub fn __getitem__(&self, indx: isize) -> PyResult<PyTokenizedRegion> {
+        let indx = if indx < 0 {
+            self.regions.len() as isize + indx
+        } else {
+            indx
+        };
+        if indx < 0 || indx >= self.regions.len() as isize {
+            Err(PyIndexError::new_err("Index out of bounds"))
+        } else {
+            let region = self.regions[indx as usize].clone();
+            let id = self.ids[indx as usize];
+            Ok(PyTokenizedRegion::new(region, id))
+        }
+    }
+}