Merge pull request #2 from SingleRust/feature-dev-export

Feature dev export

Author: ianfd

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "anndata-memory"
-version = "1.0.0"
+version = "1.0.1"
 edition = "2021"
 readme = "README.md"
 repository = "https://github.com/SingleRust/Anndata-Memory"

Cargo.toml

@@ -1,6 +1,6 @@
 # AnnData-Memory
 
-![Version](https://img.shields.io/badge/version-1.0.0-blue)
+![Version](https://img.shields.io/badge/version-1.0.1-blue)
 [![License](https://img.shields.io/badge/license-BSD--3--Clause-green)](LICENSE.md)
 
 A high-performance, thread-safe, in-memory implementation of the AnnData data structure for the SingleRust ecosystem.
@@ -31,7 +31,7 @@ Add AnnData-Memory to your `Cargo.toml`:
 
 ```toml
 [dependencies]
-anndata-memory = "1.0.0"
+anndata-memory = "1.0.1"
 ```
 
 ## Usage

README.md

@@ -12,7 +12,6 @@ use anndata::{
 };
 use anyhow::bail;
 
-use nalgebra_sparse::CsrMatrix;
 use ndarray::Array2;
 use polars::{
     frame::DataFrame,
@@ -751,4 +750,11 @@ impl IMElementCollection {
             .map(|element| element.deep_clone())
             .ok_or_else(|| anyhow::anyhow!("Key not found"))
     }
+
+    pub fn keys(&self) -> anyhow::Result<Vec<String>> {
+        let read_guard = self.0.read_inner();
+        let data = read_guard.deref();
+        let keys = data.keys().map(|k| k.clone()).collect();
+        Ok(keys)
+    }
 }

src/ad/helpers.rs

@@ -1,52 +1,154 @@
-use std::ops::Deref;
+use std::ops::{Deref, DerefMut};
 
-use anndata::{AnnData, AnnDataOp, ArrayData, ArrayElemOp, AxisArrays, Backend, ElemCollection};
+use anndata::data::DataFrameIndex;
+use anndata::{
+    AnnData, AnnDataOp, ArrayData, ArrayElemOp, AxisArrays, Backend, ElemCollection,
+    ElemCollectionOp,
+};
+use anndata_hdf5::H5;
 use anyhow::Ok;
 
-use crate::{ad::helpers::{IMElement, IMAxisArrays}, IMAnnData, IMArrayElement, IMElementCollection};
+use crate::{
+    ad::helpers::{IMAxisArrays, IMElement},
+    IMAnnData, IMArrayElement, IMElementCollection,
+};
 
 pub fn convert_to_in_memory<B: Backend>(anndata: AnnData<B>) -> anyhow::Result<IMAnnData> {
-    let obs_df = anndata.read_obs()?; 
+    let obs_df = anndata.read_obs()?;
     let obs_names = anndata.obs_names();
     let var_df = anndata.read_var()?;
     let var_names = anndata.var_names();
     let x = anndata.x().get::<ArrayData>()?.unwrap();
-    let imad = IMAnnData::new_extended(x, obs_names.into_vec(), var_names.into_vec(), obs_df, var_df)?;
+    let imad = IMAnnData::new_extended(
+        x,
+        obs_names.into_vec(),
+        var_names.into_vec(),
+        obs_df,
+        var_df,
+    )?;
     convert_axis_arrays_to_mem(anndata.obsm(), imad.obsm())?;
     convert_axis_arrays_to_mem(anndata.obsp(), imad.obsp())?;
     convert_axis_arrays_to_mem(anndata.varm(), imad.varm())?;
     convert_axis_arrays_to_mem(anndata.varp(), imad.varp())?;
+    convert_axis_arrays_to_mem(anndata.layers(), imad.layers())?;
     convert_uns_to_mem(anndata.uns(), imad.uns())?;
     anndata.close()?;
     Ok(imad)
 }
 
-fn convert_axis_arrays_to_mem<B: Backend>(axis_arr: &AxisArrays<B>, reference_element: IMAxisArrays) -> anyhow::Result<()> {
+fn convert_axis_arrays_to_mem<B: Backend>(
+    axis_arr: &AxisArrays<B>,
+    reference_element: IMAxisArrays,
+) -> anyhow::Result<()> {
     if axis_arr.is_none() {
         return Ok(());
     }
     let x = axis_arr.inner();
     let iax = x.deref();
     let data = iax.deref();
-    for (k,v) in data.iter() {
+    for (k, v) in data.iter() {
         let arr = v.get::<ArrayData>()?.unwrap();
         let im_arr = IMArrayElement::new(arr);
         reference_element.add_array(k.to_string(), im_arr)?;
     }
     Ok(())
 }
 
-fn convert_uns_to_mem<B: Backend>(elem_col: &ElemCollection<B>, reference_element: IMElementCollection) -> anyhow::Result<()> {
+fn convert_uns_to_mem<B: Backend>(
+    elem_col: &ElemCollection<B>,
+    reference_element: IMElementCollection,
+) -> anyhow::Result<()> {
     if elem_col.is_none() {
         return Ok(());
     }
     let x = elem_col.inner();
     let iax = x.deref();
     let data = iax.deref();
-    for (k,v) in data.iter() {
+    for (k, v) in data.iter() {
         let data = v.inner().data();
         let d = IMElement::new(data?);
         reference_element.add_data(k.to_string(), d)?;
     }
     Ok(())
 }
+
+fn convert_axis_arrays_to_backed<B: Backend>(
+    reference: IMAxisArrays,
+    target_axis_arrays: &AxisArrays<B>,
+) -> anyhow::Result<()> {
+    if reference.is_empty() {
+        return Ok(());
+    }
+
+    for key in reference.keys() {
+        let value = reference.get_array(&key)?;
+        let array_data = value.get_data()?;
+        let mut guard = target_axis_arrays.lock();
+        let data = guard.deref_mut();
+        match (data) {
+            None => {}
+            Some(data) => data.add_data(&key, array_data)?,
+        };
+    }
+
+    Ok(())
+}
+
+fn convert_uns_to_backed<B: Backend>(
+    reference: IMElementCollection,
+    target_file: &AnnData<B>,
+) -> anyhow::Result<()> {
+    let uns = target_file.uns();
+
+    let keys = {
+        let read_guard = reference.0.read_inner();
+        read_guard.keys().cloned().collect::<Vec<String>>()
+    };
+
+    for key in keys {
+        let val = reference.get_data(&key)?;
+        {
+            let elem_data = val.get_data()?;
+
+            // Lock the uns structure and update it
+            let mut guard = uns.lock();
+            let data = guard.deref_mut();
+            if let Some(data) = data {
+                data.add_data(&key, elem_data)?;
+            }
+        }
+    }
+    Ok(())
+}
+
+pub fn convert_to_backed<B: Backend>(imad: &IMAnnData, anndata: &AnnData<B>) -> anyhow::Result<()> {
+    let x_data = imad.x().get_data()?;
+    anndata.set_x(x_data)?;
+
+    anndata.set_obs(imad.obs().get_data())?;
+    anndata.set_obs_names(DataFrameIndex::from(imad.obs_names()))?;
+    anndata.set_var(imad.var().get_data())?;
+    anndata.set_var_names(DataFrameIndex::from(imad.var_names()))?;
+
+    convert_axis_arrays_to_backed(imad.obsm(), anndata.obsm())?;
+    convert_axis_arrays_to_backed(imad.obsp(), anndata.obsp())?;
+    convert_axis_arrays_to_backed(imad.varm(), anndata.varm())?;
+    convert_axis_arrays_to_backed(imad.varp(), anndata.varp())?;
+
+    convert_axis_arrays_to_backed(imad.layers(), anndata.layers())?;
+
+    convert_uns_to_backed(imad.uns(), anndata)?;
+
+    Ok(())
+}
+
+pub fn convert_to_new_backed_h5(
+    imad: &IMAnnData,
+    path: impl AsRef<std::path::Path>,
+) -> anyhow::Result<AnnData<H5>> {
+    let anndata = AnnData::<H5>::new(path)?;
+
+    convert_to_backed(imad, &anndata)?;
+
+    Ok(anndata)
+}

src/converter.rs

@@ -10,4 +10,6 @@ pub use ad::helpers::IMElementCollection;
 pub use ad::helpers::IMElement;
 pub use ad::helpers::IMAxisArrays;
 pub use converter::convert_to_in_memory;
+pub use converter::convert_to_backed;
+pub use converter::convert_to_new_backed_h5;
 pub use base::DeepClone;

src/lib.rs

@@ -4,8 +4,8 @@ use anndata::{
 use anndata_hdf5::H5;
 use anndata_memory::*;
 
-#[test]
-//#[allow(dead_code)]
+//#[test]
+#[allow(dead_code)]
 fn test_convert_anndata_to_imanndata_filter() -> anyhow::Result<()> {
     let h5_file = H5::open("/local/bachelor_thesis_ian/single_bench/data/merged_test.h5ad")?;
     let anndata = AnnData::<H5>::open(h5_file)?;

tests/test_conversion.rs

@@ -0,0 +1,125 @@
+use anndata::{AnnData, AnnDataOp, ArrayData, AxisArraysOp, Backend};
+use anndata_hdf5::H5;
+use anndata_memory::*;
+use nalgebra_sparse::{CooMatrix, CsrMatrix};
+use ndarray::Array2;
+use polars::prelude::*;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use tempfile::tempdir;
+
+#[test]
+fn test_round_trip_conversion() -> anyhow::Result<()> {
+    let n_rows = 10;
+    let n_cols = 20;
+    let density = 0.3;
+    let nnz = (n_rows * n_cols) as f64 * density;
+    let mut rng = StdRng::seed_from_u64(42);
+
+    let mut coo_matrix = CooMatrix::new(n_rows, n_cols);
+
+    for _ in 0..(nnz as usize) {
+        let row = rng.gen_range(0..n_rows);
+        let col = rng.gen_range(0..n_cols);
+        let value = rng.gen::<f64>();
+        coo_matrix.push(row, col, value);
+    }
+
+    let csr_matrix: CsrMatrix<f64> = (&coo_matrix).into();
+
+    let im_array = IMArrayElement::new(ArrayData::from(csr_matrix));
+
+    let obs_names: Vec<String> = (0..n_rows).map(|i| format!("cell_{}", i)).collect();
+
+    let cell_types = vec!["T cell", "B cell", "NK cell", "Monocyte", "Dendritic"];
+
+    let obs_cell_types: Vec<&str> = (0..n_rows)
+        .map(|i| cell_types[i % cell_types.len()])
+        .collect();
+
+    let obs_df = DataFrame::new(vec![Column::from(Series::new(
+        PlSmallStr::from("cell_type"),
+        obs_cell_types,
+    ))])?;
+    let obs_element = IMDataFrameElement::new(obs_df, obs_names.clone().into());
+
+    let var_names: Vec<String> = (0..n_cols).map(|i| format!("gene_{}", i)).collect();
+    let gene_types = vec!["protein_coding", "lincRNA", "pseudogene", "miRNA"];
+    let var_gene_types: Vec<&str> = (0..n_cols)
+        .map(|i| gene_types[i % gene_types.len()])
+        .collect();
+
+    let var_df = DataFrame::new(vec![Column::from(Series::new(
+        PlSmallStr::from("gene_type"),
+        var_gene_types,
+    ))])?;
+    let var_element = IMDataFrameElement::new(var_df, var_names.clone().into());
+
+    let mut im_anndata = IMAnnData::new(im_array, obs_element, var_element)?;
+
+    let mut coo_norm = CooMatrix::new(n_rows, n_cols);
+
+    for _ in 0..(nnz as usize) {
+        let row = rng.gen_range(0..n_rows);
+        let col = rng.gen_range(0..n_cols);
+        let value = rng.gen::<f64>() * 0.1; // Smaller values
+        coo_norm.push(row, col, value);
+    }
+
+    let csr_norm: CsrMatrix<f64> = CsrMatrix::from(&coo_norm);
+
+    let layer_im_array = IMArrayElement::new(ArrayData::from(csr_norm));
+    im_anndata.add_layer("normalized".to_string(), layer_im_array)?;
+
+    let n_components = 2;
+    let pca_data: Array2<f64> = Array2::from_shape_fn((n_rows, n_components), |_| rng.gen());
+    let pca_array_data = ArrayData::from(pca_data);
+    let pca_im_array = IMArrayElement::new(pca_array_data);
+    im_anndata
+        .obsm()
+        .add_array("X_pca".to_string(), pca_im_array)?;
+
+    let temp_dir = tempdir()?;
+    let h5_path = temp_dir.path().join("test_anndata.h5ad");
+    let h5_anndata = convert_to_new_backed_h5(&im_anndata, &h5_path)?;
+
+    assert_eq!(h5_anndata.n_obs(), n_rows);
+    assert_eq!(h5_anndata.n_vars(), n_cols);
+
+    let h5_obs_names = h5_anndata.obs_names().into_vec();
+    assert_eq!(h5_obs_names, obs_names);
+
+    let h5_var_names = h5_anndata.var_names().into_vec();
+    assert_eq!(h5_var_names, var_names);
+
+    assert!(h5_anndata
+        .layers()
+        .keys()
+        .contains(&"normalized".to_string()));
+
+    assert!(h5_anndata.obsm().keys().contains(&"X_pca".to_string()));
+
+    h5_anndata.close()?;
+
+    // Now test reading back from H5 to in-memory
+    let h5_file = H5::open(&h5_path)?;
+    let reopened_anndata = AnnData::<H5>::open(h5_file)?;
+    let im_anndata2 = convert_to_in_memory(reopened_anndata)?;
+
+    println!("{:?}", im_anndata2
+        .layers()
+        .keys());
+
+    // Verify round-trip conversion
+    assert_eq!(im_anndata2.n_obs(), n_rows);
+    assert_eq!(im_anndata2.n_vars(), n_cols);
+    assert_eq!(im_anndata2.obs_names(), obs_names);
+    assert_eq!(im_anndata2.var_names(), var_names);
+    assert!(im_anndata2
+        .layers()
+        .keys()
+        .contains(&"normalized".to_string()));
+    assert!(im_anndata2.obsm().keys().contains(&"X_pca".to_string()));
+
+    Ok(())
+}