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, FillValue, 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        DataType::Float32,
62        vec![4, 4].try_into()?, // regular chunk shape
63        FillValue::from(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
82            .chunk_grid()
83            .subset(&chunk_indices, array.shape())?
84            .ok_or_else(|| {
85                zarrs::array::ArrayError::InvalidChunkGridIndicesError(chunk_indices.to_vec())
86            })?;
87        array.store_chunk_ndarray(
88            &chunk_indices,
89            ArrayD::<f32>::from_shape_vec(
90                chunk_subset.shape_usize(),
91                vec![i as f32 * 0.1; chunk_subset.num_elements() as usize],
92            )
93            .unwrap(),
94        )
95    })?;
96
97    let subset_all = array.subset_all();
98    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
99    println!("store_chunk [0, 0] and [0, 1]:\n{data_all:+4.1}\n");
100
101    // Store multiple chunks
102    let ndarray_chunks: Array2<f32> = array![
103        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
104        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
105        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
106        [1.0, 1.0, 1.0, 1.0, 1.1, 1.1, 1.1, 1.1,],
107    ];
108    array.store_chunks_ndarray(&ArraySubset::new_with_ranges(&[1..2, 0..2]), ndarray_chunks)?;
109    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
110    println!("store_chunks [1..2, 0..2]:\n{data_all:+4.1}\n");
111
112    // Write a subset spanning multiple chunks, including updating chunks already written
113    let ndarray_subset: Array2<f32> =
114        array![[-3.3, -3.4, -3.5,], [-4.3, -4.4, -4.5,], [-5.3, -5.4, -5.5],];
115    array.store_array_subset_ndarray(
116        ArraySubset::new_with_ranges(&[3..6, 3..6]).start(),
117        ndarray_subset,
118    )?;
119    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
120    println!("store_array_subset [3..6, 3..6]:\n{data_all:+4.1}\n");
121
122    // Store array subset
123    let ndarray_subset: Array2<f32> = array![
124        [-0.6],
125        [-1.6],
126        [-2.6],
127        [-3.6],
128        [-4.6],
129        [-5.6],
130        [-6.6],
131        [-7.6],
132    ];
133    array.store_array_subset_ndarray(
134        ArraySubset::new_with_ranges(&[0..8, 6..7]).start(),
135        ndarray_subset,
136    )?;
137    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
138    println!("store_array_subset [0..8, 6..7]:\n{data_all:+4.1}\n");
139
140    // Store chunk subset
141    let ndarray_chunk_subset: Array2<f32> = array![[-7.4, -7.5, -7.6, -7.7],];
142    array.store_chunk_subset_ndarray(
143        // chunk indices
144        &[1, 1],
145        // subset within chunk
146        ArraySubset::new_with_ranges(&[3..4, 0..4]).start(),
147        ndarray_chunk_subset,
148    )?;
149    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
150    println!("store_chunk_subset [3..4, 0..4] of chunk [1, 1]:\n{data_all:+4.1}\n");
151
152    // Erase a chunk
153    array.erase_chunk(&[0, 0])?;
154    let data_all = array.retrieve_array_subset_ndarray::<f32>(&subset_all)?;
155    println!("erase_chunk [0, 0]:\n{data_all:+4.1}\n");
156
157    // Read a chunk
158    let chunk_indices = vec![0, 1];
159    let data_chunk = array.retrieve_chunk_ndarray::<f32>(&chunk_indices)?;
160    println!("retrieve_chunk [0, 1]:\n{data_chunk:+4.1}\n");
161
162    // Read chunks
163    let chunks = ArraySubset::new_with_ranges(&[0..2, 1..2]);
164    let data_chunks = array.retrieve_chunks_ndarray::<f32>(&chunks)?;
165    println!("retrieve_chunks [0..2, 1..2]:\n{data_chunks:+4.1}\n");
166
167    // Retrieve an array subset
168    let subset = ArraySubset::new_with_ranges(&[2..6, 3..5]); // the center 4x2 region
169    let data_subset = array.retrieve_array_subset_ndarray::<f32>(&subset)?;
170    println!("retrieve_array_subset [2..6, 3..5]:\n{data_subset:+4.1}\n");
171
172    // Show the hierarchy
173    let node = Node::open(&store, "/").unwrap();
174    let tree = node.hierarchy_tree();
175    println!("hierarchy_tree:\n{}", tree);
176
177    Ok(())
178}
179
180fn main() {
181    if let Err(err) = array_write_read() {
182        println!("{:?}", err);
183    }
184}
array_write_read_ndarray/array_write_read_ndarray.rs

array_write_read_ndarray/
array_write_read_ndarray.rs
