array_write_read_ndarray/
array_write_read_ndarray.rs

1#![allow(missing_docs)]
2
3use ndarray::{Array2, ArrayD, array};
4use rayon::iter::{IntoParallelIterator, ParallelIterator};
5use zarrs::storage::ReadableWritableListableStorage;
6use zarrs::storage::storage_adapter::usage_log::UsageLogStorageAdapter;
7
8fn array_write_read() -> Result<(), Box<dyn std::error::Error>> {
9    use std::sync::Arc;
10
11    use zarrs::array::{ArraySubset, ZARR_NAN_F32, data_type};
12    use zarrs::node::Node;
13    use zarrs::storage::store;
14
15    // Create a store
16    // let path = tempfile::TempDir::new()?;
17    // let mut store: ReadableWritableListableStorage =
18    //     Arc::new(zarrs::filesystem::FilesystemStore::new(path.path())?);
19    // let mut store: ReadableWritableListableStorage = Arc::new(
20    //     zarrs::filesystem::FilesystemStore::new("zarrs/tests/data/array_write_read.zarr")?,
21    // );
22    let mut store: ReadableWritableListableStorage = Arc::new(store::MemoryStore::new());
23    if let Some(arg1) = std::env::args().collect::<Vec<_>>().get(1)
24        && arg1 == "--usage-log"
25    {
26        let log_writer = Arc::new(std::sync::Mutex::new(
27            // std::io::BufWriter::new(
28            std::io::stdout(),
29            //    )
30        ));
31        store = Arc::new(UsageLogStorageAdapter::new(store, log_writer, || {
32            chrono::Utc::now().format("[%T%.3f] ").to_string()
33        }));
34    }
35
36    // Create the root group
37    zarrs::group::GroupBuilder::new()
38        .build(store.clone(), "/")?
39        .store_metadata()?;
40
41    // Create a group with attributes
42    let group_path = "/group";
43    let mut group = zarrs::group::GroupBuilder::new().build(store.clone(), group_path)?;
44    group
45        .attributes_mut()
46        .insert("foo".into(), serde_json::Value::String("bar".into()));
47    group.store_metadata()?;
48
49    println!(
50        "The group metadata is:\n{}\n",
51        group.metadata().to_string_pretty()
52    );
53
54    // Create an array
55    let array_path = "/group/array";
56    let array = zarrs::array::ArrayBuilder::new(
57        vec![8, 8], // array shape
58        vec![4, 4], // regular chunk shape
59        data_type::float32(),
60        ZARR_NAN_F32,
61    )
62    // .bytes_to_bytes_codecs(vec![]) // uncompressed
63    .dimension_names(["y", "x"].into())
64    // .storage_transformers(vec![].into())
65    .build(store.clone(), array_path)?;
66
67    // Write array metadata to store
68    array.store_metadata()?;
69
70    println!(
71        "The array metadata is:\n{}\n",
72        array.metadata().to_string_pretty()
73    );
74
75    // Write some chunks
76    (0..2).into_par_iter().try_for_each(|i| {
77        let chunk_indices: Vec<u64> = vec![0, i];
78        let chunk_subset = array.chunk_grid().subset(&chunk_indices)?.ok_or_else(|| {
79            zarrs::array::ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec())
80        })?;
81        array.store_chunk(
82            &chunk_indices,
83            ArrayD::<f32>::from_shape_vec(
84                chunk_subset.shape_usize(),
85                vec![i as f32 * 0.1; chunk_subset.num_elements() as usize],
86            )
87            .unwrap(),
88        )
89    })?;
90
91    let subset_all = array.subset_all();
92    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
93    println!("store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n");
94
95    // Store multiple chunks
96    let ndarray_chunks: Array2<f32> = array![
97        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
98        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
99        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
100        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
101    ];
102    array.store_chunks(&[1..2, 0..2], ndarray_chunks)?;
103    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
104    println!("store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n");
105
106    // Write a subset spanning multiple chunks, including updating chunks already written
107    let ndarray_subset: Array2<f32> =
108        array![[-3.3, -3.4, -3.5,], [-4.3, -4.4, -4.5,], [-5.3, -5.4, -5.5],];
109    array.store_array_subset(&[3..6, 3..6], ndarray_subset)?;
110    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
111    println!("store_array_subset [3..6, 3..6]:\n{data_all:+4.1}\n");
112
113    // Store array subset
114    let ndarray_subset: Array2<f32> = array![
115        [-0.6],
116        [-1.6],
117        [-2.6],
118        [-3.6],
119        [-4.6],
120        [-5.6],
121        [-6.6],
122        [-7.6],
123    ];
124    array.store_array_subset(&[0..8, 6..7], ndarray_subset)?;
125    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
126    println!("store_array_subset [0..8, 6..7]:\n{data_all:+4.1}\n");
127
128    // Store chunk subset
129    let ndarray_chunk_subset: Array2<f32> = array![[-7.4, -7.5, -7.6, -7.7],];
130    array.store_chunk_subset(
131        // chunk indices
132        &[1, 1],
133        // subset within chunk
134        &[3..4, 0..4],
135        ndarray_chunk_subset,
136    )?;
137    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
138    println!("store_chunk_subset [3..4, 0..4] of chunk [1, 1]:\n{data_all:+4.1}\n");
139
140    // Erase a chunk
141    array.erase_chunk(&[0, 0])?;
142    let data_all: ArrayD<f32> = array.retrieve_array_subset(&subset_all)?;
143    println!("erase_chunk [0, 0]:\n{data_all:+4.1}\n");
144
145    // Read a chunk
146    let chunk_indices = vec![0, 1];
147    let data_chunk: ArrayD<f32> = array.retrieve_chunk(&chunk_indices)?;
148    println!("retrieve_chunk [0, 1]:\n{data_chunk:+4.1}\n");
149
150    // Read chunks
151    let chunks = ArraySubset::new_with_ranges(&[0..2, 1..2]);
152    let data_chunks: ArrayD<f32> = array.retrieve_chunks(&chunks)?;
153    println!("retrieve_chunks [0..2, 1..2]:\n{data_chunks:+4.1}\n");
154
155    // Retrieve an array subset
156    let subset = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
157    let data_subset: ArrayD<f32> = array.retrieve_array_subset(&subset)?;
158    println!("retrieve_array_subset [2..6, 3..5]:\n{data_subset:+4.1}\n");
159
160    // Show the hierarchy
161    let node = Node::open(&store, "/").unwrap();
162    let tree = node.hierarchy_tree();
163    println!("hierarchy_tree:\n{}", tree);
164
165    Ok(())
166}
167
168fn main() {
169    if let Err(err) = array_write_read() {
170        println!("{:?}", err);
171    }
172}
array_write_read_ndarray/array_write_read_ndarray.rs

array_write_read_ndarray/
array_write_read_ndarray.rs
