Skip to main content

zarrs/array/data_type/
subfloat.rs

1//! Subfloat data types (sub-byte floating point formats).
2
3use super::macros::register_data_type_plugin;
4
5/// Macro to implement `DataTypeTraits` for subfloat types (microfloat only).
6macro_rules! impl_subfloat_data_type {
7    ($marker:ty, $float_type:ty $(, $float8_type:ty)?) => {
8        impl zarrs_data_type::DataTypeTraits for $marker {
9            fn configuration(
10                &self,
11                _version: zarrs_plugin::ZarrVersion,
12            ) -> zarrs_metadata::Configuration {
13                zarrs_metadata::Configuration::default()
14            }
15
16            fn size(&self) -> zarrs_metadata::DataTypeSize {
17                zarrs_metadata::DataTypeSize::Fixed(1)
18            }
19
20            fn fill_value(
21                &self,
22                fill_value_metadata: &zarrs_metadata::FillValueMetadata,
23                _version: zarrs_plugin::ZarrVersion,
24            ) -> Result<zarrs_data_type::FillValue, zarrs_data_type::DataTypeFillValueMetadataError>
25            {
26                if let Some(s) = fill_value_metadata.as_str() {
27                    if let Some(hex) = s.strip_prefix("0x")
28                        && let Ok(byte) = u8::from_str_radix(hex, 16)
29                    {
30                        return Ok(zarrs_data_type::FillValue::from(byte));
31                    }
32                    #[cfg(feature = "microfloat")]
33                    {
34                        match s {
35                            "NaN" => {
36                                return Ok(zarrs_data_type::FillValue::from(
37                                    <$float_type>::NAN.to_bits(),
38                                ));
39                            }
40                            "Infinity" => {
41                                return Ok(zarrs_data_type::FillValue::from(
42                                    <$float_type>::INFINITY.to_bits(),
43                                ));
44                            }
45                            "-Infinity" => {
46                                return Ok(zarrs_data_type::FillValue::from(
47                                    <$float_type>::NEG_INFINITY.to_bits(),
48                                ));
49                            }
50                            _ => {}
51                        }
52                    }
53                }
54                #[cfg(feature = "microfloat")]
55                {
56                    if let Some(f) = fill_value_metadata.as_f64() {
57                        return Ok(zarrs_data_type::FillValue::from(
58                            <$float_type>::from_f64(f).to_bits(),
59                        ));
60                    }
61                }
62                if let Some(int) = fill_value_metadata.as_u64() {
63                    if let Ok(byte) = u8::try_from(int) {
64                        return Ok(zarrs_data_type::FillValue::from(byte));
65                    }
66                }
67                Err(zarrs_data_type::DataTypeFillValueMetadataError)
68            }
69
70            fn metadata_fill_value(
71                &self,
72                fill_value: &zarrs_data_type::FillValue,
73            ) -> Result<zarrs_metadata::FillValueMetadata, zarrs_data_type::DataTypeFillValueError>
74            {
75                let bytes: [u8; 1] = fill_value
76                    .as_ne_bytes()
77                    .try_into()
78                    .map_err(|_| zarrs_data_type::DataTypeFillValueError)?;
79                #[cfg(feature = "microfloat")]
80                {
81                    let value = <$float_type>::from_bits(bytes[0]);
82                    if value.is_nan() {
83                        Ok(zarrs_metadata::FillValueMetadata::from("NaN".to_string()))
84                    } else if value.is_infinite() {
85                        if value.is_sign_negative() {
86                            Ok(zarrs_metadata::FillValueMetadata::from(
87                                "-Infinity".to_string(),
88                            ))
89                        } else {
90                            Ok(zarrs_metadata::FillValueMetadata::from(
91                                "Infinity".to_string(),
92                            ))
93                        }
94                    } else {
95                        Ok(zarrs_metadata::FillValueMetadata::from(value.to_f64()))
96                    }
97                }
98                #[cfg(not(feature = "microfloat"))]
99                {
100                    Ok(zarrs_metadata::FillValueMetadata::from(format!(
101                        "0x{:02x}",
102                        bytes[0]
103                    )))
104                }
105            }
106
107            fn as_any(&self) -> &dyn std::any::Any {
108                self
109            }
110
111            fn compatible_element_types(&self) -> &'static [std::any::TypeId] {
112                #[cfg(all(feature = "float8", feature = "microfloat"))]
113                {
114                    const TYPES: [std::any::TypeId; impl_subfloat_data_type!(@num_types $($float8_type)?)] = [
115                        std::any::TypeId::of::<$float_type>(),
116                        $(std::any::TypeId::of::<$float8_type>(),)?
117                    ];
118                    &TYPES
119                }
120                #[cfg(all(not(feature = "float8"), feature = "microfloat"))]
121                {
122                    const TYPES: [std::any::TypeId; 1] = [std::any::TypeId::of::<$float_type>()];
123                    &TYPES
124                }
125                #[cfg(not(feature = "microfloat"))]
126                {
127                    &[]
128                }
129            }
130        }
131    };
132    (@num_types) => {
133        1
134    };
135    (@num_types $float8_type:ty) => {
136        2
137    };
138}
139
140/// The `float4_e2m1fn` data type.
141#[derive(Debug, Clone, Copy)]
142pub struct Float4E2M1FNDataType;
143register_data_type_plugin!(Float4E2M1FNDataType);
144zarrs_plugin::impl_extension_aliases!(Float4E2M1FNDataType, v3: "float4_e2m1fn");
145
146/// The `float6_e2m3fn` data type.
147#[derive(Debug, Clone, Copy)]
148pub struct Float6E2M3FNDataType;
149register_data_type_plugin!(Float6E2M3FNDataType);
150zarrs_plugin::impl_extension_aliases!(Float6E2M3FNDataType, v3: "float6_e2m3fn");
151
152/// The `float6_e3m2fn` data type.
153#[derive(Debug, Clone, Copy)]
154pub struct Float6E3M2FNDataType;
155register_data_type_plugin!(Float6E3M2FNDataType);
156zarrs_plugin::impl_extension_aliases!(Float6E3M2FNDataType, v3: "float6_e3m2fn");
157
158/// The `float8_e3m4` data type.
159#[derive(Debug, Clone, Copy)]
160pub struct Float8E3M4DataType;
161register_data_type_plugin!(Float8E3M4DataType);
162zarrs_plugin::impl_extension_aliases!(Float8E3M4DataType, v3: "float8_e3m4");
163
164/// The `float8_e4m3` data type.
165#[derive(Debug, Clone, Copy)]
166pub struct Float8E4M3DataType;
167register_data_type_plugin!(Float8E4M3DataType);
168zarrs_plugin::impl_extension_aliases!(Float8E4M3DataType, v3: "float8_e4m3");
169
170/// The `float8_e4m3b11fnuz` data type.
171#[derive(Debug, Clone, Copy)]
172pub struct Float8E4M3B11FNUZDataType;
173register_data_type_plugin!(Float8E4M3B11FNUZDataType);
174zarrs_plugin::impl_extension_aliases!(Float8E4M3B11FNUZDataType, v3: "float8_e4m3b11fnuz");
175
176/// The `float8_e4m3fnuz` data type.
177#[derive(Debug, Clone, Copy)]
178pub struct Float8E4M3FNUZDataType;
179register_data_type_plugin!(Float8E4M3FNUZDataType);
180zarrs_plugin::impl_extension_aliases!(Float8E4M3FNUZDataType, v3: "float8_e4m3fnuz");
181
182/// The `float8_e5m2` data type.
183#[derive(Debug, Clone, Copy)]
184pub struct Float8E5M2DataType;
185register_data_type_plugin!(Float8E5M2DataType);
186zarrs_plugin::impl_extension_aliases!(Float8E5M2DataType, v3: "float8_e5m2");
187
188/// The `float8_e5m2fnuz` data type.
189#[derive(Debug, Clone, Copy)]
190pub struct Float8E5M2FNUZDataType;
191register_data_type_plugin!(Float8E5M2FNUZDataType);
192zarrs_plugin::impl_extension_aliases!(Float8E5M2FNUZDataType, v3: "float8_e5m2fnuz");
193
194/// The `float8_e8m0fnu` data type.
195#[derive(Debug, Clone, Copy)]
196pub struct Float8E8M0FNUDataType;
197register_data_type_plugin!(Float8E8M0FNUDataType);
198zarrs_plugin::impl_extension_aliases!(Float8E8M0FNUDataType, v3: "float8_e8m0fnu");
199
200// DataTypeTraits implementations for subfloats (microfloat only)
201impl_subfloat_data_type!(Float4E2M1FNDataType, microfloat::f4e2m1fn);
202impl_subfloat_data_type!(Float6E2M3FNDataType, microfloat::f6e2m3fn);
203impl_subfloat_data_type!(Float6E3M2FNDataType, microfloat::f6e3m2fn);
204impl_subfloat_data_type!(Float8E3M4DataType, microfloat::f8e3m4);
205impl_subfloat_data_type!(Float8E4M3DataType, microfloat::f8e4m3, float8::F8E4M3);
206impl_subfloat_data_type!(Float8E4M3B11FNUZDataType, microfloat::f8e4m3b11fnuz);
207impl_subfloat_data_type!(Float8E4M3FNUZDataType, microfloat::f8e4m3fnuz);
208impl_subfloat_data_type!(Float8E5M2DataType, microfloat::f8e5m2, float8::F8E5M2);
209impl_subfloat_data_type!(Float8E5M2FNUZDataType, microfloat::f8e5m2fnuz);
210impl_subfloat_data_type!(Float8E8M0FNUDataType, microfloat::f8e8m0fnu);
211
212// PackBits codec implementations for subfloats
213use zarrs_data_type::codec_traits::impl_pack_bits_data_type_traits;
214impl_pack_bits_data_type_traits!(Float4E2M1FNDataType, 4, float, 1);
215impl_pack_bits_data_type_traits!(Float6E2M3FNDataType, 6, float, 1);
216impl_pack_bits_data_type_traits!(Float6E3M2FNDataType, 6, float, 1);
217impl_pack_bits_data_type_traits!(Float8E3M4DataType, 8, float, 1);
218impl_pack_bits_data_type_traits!(Float8E4M3DataType, 8, float, 1);
219impl_pack_bits_data_type_traits!(Float8E4M3B11FNUZDataType, 8, float, 1);
220impl_pack_bits_data_type_traits!(Float8E4M3FNUZDataType, 8, float, 1);
221impl_pack_bits_data_type_traits!(Float8E5M2DataType, 8, float, 1);
222impl_pack_bits_data_type_traits!(Float8E5M2FNUZDataType, 8, float, 1);
223impl_pack_bits_data_type_traits!(Float8E8M0FNUDataType, 8, float, 1);
224
225// Bytes codec implementations for subfloats (passthrough - single byte, no endianness conversion)
226use zarrs_data_type::codec_traits::impl_bytes_data_type_traits;
227impl_bytes_data_type_traits!(Float4E2M1FNDataType, 1);
228impl_bytes_data_type_traits!(Float6E2M3FNDataType, 1);
229impl_bytes_data_type_traits!(Float6E3M2FNDataType, 1);
230impl_bytes_data_type_traits!(Float8E3M4DataType, 1);
231impl_bytes_data_type_traits!(Float8E4M3DataType, 1);
232impl_bytes_data_type_traits!(Float8E4M3B11FNUZDataType, 1);
233impl_bytes_data_type_traits!(Float8E4M3FNUZDataType, 1);
234impl_bytes_data_type_traits!(Float8E5M2DataType, 1);
235impl_bytes_data_type_traits!(Float8E5M2FNUZDataType, 1);
236impl_bytes_data_type_traits!(Float8E8M0FNUDataType, 1);