array_write_read_ndarray/
array_write_read_ndarray.rs

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

array_write_read_ndarray/
array_write_read_ndarray.rs
