use std::num::NonZeroU64;
use egui_wgpu::wgpu;
use crate::core::colormap::Colormap;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum InterpolationMode {
#[default]
Nearest,
Linear,
}
impl InterpolationMode {
fn code(self) -> u32 {
match self {
InterpolationMode::Nearest => 0,
InterpolationMode::Linear => 1,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum AggregationMode {
#[default]
None,
Max,
Mean,
Min,
}
pub fn aggregate_blocks(
data: &[f32],
width: u32,
height: u32,
block_x: u32,
block_y: u32,
mode: AggregationMode,
) -> (Vec<f32>, u32, u32) {
let bx = block_x.max(1);
let by = block_y.max(1);
if mode == AggregationMode::None || (bx == 1 && by == 1) {
return (data.to_vec(), width, height);
}
let out_w = width / bx;
let out_h = height / by;
let mut out = Vec::with_capacity((out_w as usize) * (out_h as usize));
for oy in 0..out_h {
for ox in 0..out_w {
let mut acc: Option<f32> = None;
let mut count: u32 = 0; for j in 0..by {
let row = (oy * by + j) as usize;
let base = row * (width as usize) + (ox * bx) as usize;
for i in 0..bx {
let v = data[base + i as usize];
if v.is_nan() {
continue;
}
count += 1;
acc = Some(match (acc, mode) {
(None, _) => v,
(Some(a), AggregationMode::Max) => a.max(v),
(Some(a), AggregationMode::Min) => a.min(v),
(Some(a), AggregationMode::Mean) => a + v,
(Some(a), AggregationMode::None) => a,
});
}
}
let value = match acc {
None => f32::NAN, Some(a) if mode == AggregationMode::Mean => a / count as f32,
Some(a) => a,
};
out.push(value);
}
}
(out, out_w, out_h)
}
const IDENTITY: [[f32; 4]; 4] = [
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
];
fn srgb_u8_to_linear_rgba(c: [u8; 4]) -> [f32; 4] {
let to_linear = |b: u8| {
let s = b as f32 / 255.0;
if s <= 0.04045 {
s / 12.92
} else {
((s + 0.055) / 1.055).powf(2.4)
}
};
[
to_linear(c[0]),
to_linear(c[1]),
to_linear(c[2]),
c[3] as f32 / 255.0,
]
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ImageParams {
ortho: [[f32; 4]; 4],
rect: [f32; 4],
axis_log: [f32; 2],
alpha: f32,
cmap_min: f32,
cmap_one_over_range: f32,
gamma: f32,
norm: u32,
interp: u32,
has_alpha_map: u32,
_pad: [u32; 3],
nan_color: [f32; 4],
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ImageRgbaParams {
ortho: [[f32; 4]; 4],
rect: [f32; 4],
axis_log: [f32; 2],
alpha: f32,
_pad: f32,
}
#[derive(Clone, Debug)]
pub enum ImagePixels {
Scalar {
data: Vec<f32>,
colormap: Box<Colormap>,
},
Rgba { data: Vec<[u8; 4]> },
}
#[derive(Clone, Debug)]
pub struct ImageData {
pub pixels: ImagePixels,
pub width: u32,
pub height: u32,
pub origin: (f64, f64),
pub scale: (f64, f64),
pub alpha: f32,
pub interpolation: InterpolationMode,
pub alpha_map: Option<Vec<f32>>,
}
impl ImageData {
pub fn new(width: u32, height: u32, data: Vec<f32>, colormap: Colormap) -> Self {
assert_eq!(
data.len(),
(width as usize) * (height as usize),
"data length must equal width * height"
);
Self {
pixels: ImagePixels::Scalar {
data,
colormap: Box::new(colormap),
},
width,
height,
origin: (0.0, 0.0),
scale: (1.0, 1.0),
alpha: 1.0,
interpolation: InterpolationMode::default(),
alpha_map: None,
}
}
pub fn rgba(width: u32, height: u32, data: Vec<[u8; 4]>) -> Self {
assert_eq!(
data.len(),
(width as usize) * (height as usize),
"data length must equal width * height"
);
Self {
pixels: ImagePixels::Rgba { data },
width,
height,
origin: (0.0, 0.0),
scale: (1.0, 1.0),
alpha: 1.0,
interpolation: InterpolationMode::default(),
alpha_map: None,
}
}
pub fn colormap(&self) -> Option<&Colormap> {
match &self.pixels {
ImagePixels::Scalar { colormap, .. } => Some(colormap.as_ref()),
ImagePixels::Rgba { .. } => None,
}
}
pub fn with_interpolation(mut self, interpolation: InterpolationMode) -> Self {
self.interpolation = interpolation;
self
}
pub fn with_alpha_map(mut self, alpha_map: Vec<f32>) -> Self {
assert_eq!(
alpha_map.len(),
(self.width as usize) * (self.height as usize),
"alpha_map length must equal width * height"
);
self.alpha_map = Some(alpha_map);
self
}
}
pub struct ImagePipeline {
pub(crate) pipeline: wgpu::RenderPipeline,
pub(crate) rgba_pipeline: wgpu::RenderPipeline,
bind_group_layout: wgpu::BindGroupLayout,
rgba_bind_group_layout: wgpu::BindGroupLayout,
data_sampler: wgpu::Sampler,
lut_sampler: wgpu::Sampler,
linear_sampler: wgpu::Sampler,
dummy_alpha_view: wgpu::TextureView,
}
impl ImagePipeline {
pub fn new(device: &wgpu::Device, target_format: wgpu::TextureFormat) -> Self {
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("rsplot image"),
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/image.wgsl").into()),
});
let rgba_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("rsplot image rgba"),
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/image_rgba.wgsl").into()),
});
let bind_group_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("rsplot image bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: NonZeroU64::new(
std::mem::size_of::<ImageParams>() as u64
),
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::NonFiltering),
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 3,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 4,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 5,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
],
});
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("rsplot image layout"),
bind_group_layouts: &[Some(&bind_group_layout)],
immediate_size: 0,
});
let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("rsplot image pipeline"),
layout: Some(&pipeline_layout),
vertex: wgpu::VertexState {
module: &shader,
entry_point: Some("vs_main"),
buffers: &[],
compilation_options: wgpu::PipelineCompilationOptions::default(),
},
fragment: Some(wgpu::FragmentState {
module: &shader,
entry_point: Some("fs_main"),
targets: &[Some(wgpu::ColorTargetState {
format: target_format,
blend: Some(wgpu::BlendState::ALPHA_BLENDING),
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: wgpu::PipelineCompilationOptions::default(),
}),
primitive: wgpu::PrimitiveState::default(),
depth_stencil: None,
multisample: wgpu::MultisampleState::default(),
multiview_mask: None,
cache: None,
});
let rgba_bind_group_layout =
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("rsplot image rgba bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: NonZeroU64::new(
std::mem::size_of::<ImageRgbaParams>() as u64,
),
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: true },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
wgpu::BindGroupLayoutEntry {
binding: 2,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
count: None,
},
],
});
let rgba_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("rsplot image rgba layout"),
bind_group_layouts: &[Some(&rgba_bind_group_layout)],
immediate_size: 0,
});
let rgba_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
label: Some("rsplot image rgba pipeline"),
layout: Some(&rgba_pipeline_layout),
vertex: wgpu::VertexState {
module: &rgba_shader,
entry_point: Some("vs_main"),
buffers: &[],
compilation_options: wgpu::PipelineCompilationOptions::default(),
},
fragment: Some(wgpu::FragmentState {
module: &rgba_shader,
entry_point: Some("fs_main"),
targets: &[Some(wgpu::ColorTargetState {
format: target_format,
blend: Some(wgpu::BlendState::ALPHA_BLENDING),
write_mask: wgpu::ColorWrites::ALL,
})],
compilation_options: wgpu::PipelineCompilationOptions::default(),
}),
primitive: wgpu::PrimitiveState::default(),
depth_stencil: None,
multisample: wgpu::MultisampleState::default(),
multiview_mask: None,
cache: None,
});
let data_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("rsplot image data sampler"),
mag_filter: wgpu::FilterMode::Nearest,
min_filter: wgpu::FilterMode::Nearest,
..Default::default()
});
let lut_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("rsplot image lut sampler"),
mag_filter: wgpu::FilterMode::Nearest,
min_filter: wgpu::FilterMode::Nearest,
..Default::default()
});
let linear_sampler = device.create_sampler(&wgpu::SamplerDescriptor {
label: Some("rsplot image linear sampler"),
mag_filter: wgpu::FilterMode::Linear,
min_filter: wgpu::FilterMode::Linear,
..Default::default()
});
let dummy_alpha = device.create_texture(&wgpu::TextureDescriptor {
label: Some("rsplot image dummy alpha"),
size: wgpu::Extent3d {
width: 1,
height: 1,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::R32Float,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let dummy_alpha_view = dummy_alpha.create_view(&wgpu::TextureViewDescriptor::default());
Self {
pipeline,
rgba_pipeline,
bind_group_layout,
rgba_bind_group_layout,
data_sampler,
lut_sampler,
linear_sampler,
dummy_alpha_view,
}
}
}
fn tile_bounds(width: u32, height: u32, max_dim: u32) -> Vec<(u32, u32, u32, u32)> {
let max_dim = max_dim.max(1);
let mut tiles = Vec::new();
let mut y0 = 0;
while y0 < height {
let h = (height - y0).min(max_dim);
let mut x0 = 0;
while x0 < width {
let w = (width - x0).min(max_dim);
tiles.push((x0, y0, w, h));
x0 += w;
}
y0 += h;
}
tiles
}
fn extract_subgrid<T: Copy>(
data: &[T],
full_width: u32,
x0: u32,
y0: u32,
w: u32,
h: u32,
) -> Vec<T> {
let mut out = Vec::with_capacity((w as usize) * (h as usize));
for row in 0..h {
let src = ((y0 + row) as usize) * (full_width as usize) + x0 as usize;
out.extend_from_slice(&data[src..src + w as usize]);
}
out
}
struct ImageTile {
bind_group: wgpu::BindGroup,
params: wgpu::Buffer,
data_texture: wgpu::Texture,
rect: [f32; 4],
x0: u32,
y0: u32,
w: u32,
h: u32,
}
#[derive(Clone, Copy)]
enum GpuImageKind {
Scalar {
cmap_min: f32,
cmap_one_over_range: f32,
gamma: f32,
norm: u32,
interp: u32,
has_alpha_map: u32,
nan_color: [f32; 4],
},
Rgba,
}
pub struct GpuImage {
tiles: Vec<ImageTile>,
width: u32,
height: u32,
kind: GpuImageKind,
alpha: f32,
}
impl GpuImage {
pub fn new(
device: &wgpu::Device,
queue: &wgpu::Queue,
pipeline: &ImagePipeline,
image: &ImageData,
) -> Self {
let (ox, oy) = image.origin;
let (sx, sy) = image.scale;
let max_dim = device.limits().max_texture_dimension_2d;
let rect_of = |x0: u32, y0: u32, w: u32, h: u32| {
[
(ox + sx * x0 as f64) as f32,
(oy + sy * y0 as f64) as f32,
(ox + sx * (x0 + w) as f64) as f32,
(oy + sy * (y0 + h) as f64) as f32,
]
};
let (tiles, kind) = match &image.pixels {
ImagePixels::Scalar { data, colormap } => {
let lut_size = wgpu::Extent3d {
width: 256,
height: 1,
depth_or_array_layers: 1,
};
let lut_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("rsplot image lut"),
size: lut_size,
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8UnormSrgb,
usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &lut_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&colormap.lut),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * 256),
rows_per_image: Some(1),
},
lut_size,
);
let lut_view = lut_texture.create_view(&wgpu::TextureViewDescriptor::default());
let tiles = tile_bounds(image.width, image.height, max_dim)
.into_iter()
.map(|(x0, y0, w, h)| {
let tile_size = wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
};
let data_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("rsplot image tile"),
size: tile_size,
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::R32Float,
usage: wgpu::TextureUsages::TEXTURE_BINDING
| wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let sub = extract_subgrid(data, image.width, x0, y0, w, h);
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &data_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&sub),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * w),
rows_per_image: Some(h),
},
tile_size,
);
let data_view =
data_texture.create_view(&wgpu::TextureViewDescriptor::default());
let alpha_view = image.alpha_map.as_ref().map(|alpha| {
let alpha_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("rsplot image tile alpha"),
size: tile_size,
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::R32Float,
usage: wgpu::TextureUsages::TEXTURE_BINDING
| wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let sub = extract_subgrid(alpha, image.width, x0, y0, w, h);
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &alpha_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&sub),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * w),
rows_per_image: Some(h),
},
tile_size,
);
alpha_texture.create_view(&wgpu::TextureViewDescriptor::default())
});
let alpha_binding =
alpha_view.as_ref().unwrap_or(&pipeline.dummy_alpha_view);
let params = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("rsplot image tile params"),
size: std::mem::size_of::<ImageParams>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("rsplot image tile bg"),
layout: &pipeline.bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: params.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&data_view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: wgpu::BindingResource::Sampler(
&pipeline.data_sampler,
),
},
wgpu::BindGroupEntry {
binding: 3,
resource: wgpu::BindingResource::TextureView(&lut_view),
},
wgpu::BindGroupEntry {
binding: 4,
resource: wgpu::BindingResource::Sampler(&pipeline.lut_sampler),
},
wgpu::BindGroupEntry {
binding: 5,
resource: wgpu::BindingResource::TextureView(alpha_binding),
},
],
});
ImageTile {
bind_group,
params,
data_texture,
rect: rect_of(x0, y0, w, h),
x0,
y0,
w,
h,
}
})
.collect();
let (cmap_min, cmap_one_over_range) = colormap.norm_bounds();
(
tiles,
GpuImageKind::Scalar {
cmap_min,
cmap_one_over_range,
gamma: colormap.gamma,
norm: colormap.normalization.code(),
interp: image.interpolation.code(),
has_alpha_map: u32::from(image.alpha_map.is_some()),
nan_color: srgb_u8_to_linear_rgba(colormap.nan_color),
},
)
}
ImagePixels::Rgba { data } => {
let rgba_sampler = match image.interpolation {
InterpolationMode::Nearest => &pipeline.data_sampler,
InterpolationMode::Linear => &pipeline.linear_sampler,
};
let tiles = tile_bounds(image.width, image.height, max_dim)
.into_iter()
.map(|(x0, y0, w, h)| {
let tile_size = wgpu::Extent3d {
width: w,
height: h,
depth_or_array_layers: 1,
};
let data_texture = device.create_texture(&wgpu::TextureDescriptor {
label: Some("rsplot image rgba tile"),
size: tile_size,
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: wgpu::TextureFormat::Rgba8UnormSrgb,
usage: wgpu::TextureUsages::TEXTURE_BINDING
| wgpu::TextureUsages::COPY_DST,
view_formats: &[],
});
let sub = extract_subgrid(data, image.width, x0, y0, w, h);
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &data_texture,
mip_level: 0,
origin: wgpu::Origin3d::ZERO,
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&sub),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * w),
rows_per_image: Some(h),
},
tile_size,
);
let data_view =
data_texture.create_view(&wgpu::TextureViewDescriptor::default());
let params = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("rsplot image rgba tile params"),
size: std::mem::size_of::<ImageRgbaParams>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("rsplot image rgba tile bg"),
layout: &pipeline.rgba_bind_group_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: params.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&data_view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: wgpu::BindingResource::Sampler(rgba_sampler),
},
],
});
ImageTile {
bind_group,
params,
data_texture,
rect: rect_of(x0, y0, w, h),
x0,
y0,
w,
h,
}
})
.collect();
(tiles, GpuImageKind::Rgba)
}
};
let gpu = Self {
tiles,
width: image.width,
height: image.height,
kind,
alpha: image.alpha,
};
gpu.write_uniforms(queue, IDENTITY, [0.0, 0.0]);
gpu
}
pub(crate) fn update_region(
&self,
queue: &wgpu::Queue,
x0: u32,
y0: u32,
w: u32,
h: u32,
data: &[f32],
) {
assert!(
matches!(self.kind, GpuImageKind::Scalar { .. }),
"update_region is only valid for scalar images"
);
assert_eq!(
data.len(),
(w as usize) * (h as usize),
"data length must equal w * h"
);
assert!(
x0 + w <= self.width && y0 + h <= self.height,
"region out of bounds"
);
for tile in &self.tiles {
let ix0 = x0.max(tile.x0);
let iy0 = y0.max(tile.y0);
let ix1 = (x0 + w).min(tile.x0 + tile.w);
let iy1 = (y0 + h).min(tile.y0 + tile.h);
if ix1 <= ix0 || iy1 <= iy0 {
continue; }
let ow = ix1 - ix0;
let oh = iy1 - iy0;
let sub = extract_subgrid(data, w, ix0 - x0, iy0 - y0, ow, oh);
queue.write_texture(
wgpu::TexelCopyTextureInfo {
texture: &tile.data_texture,
mip_level: 0,
origin: wgpu::Origin3d {
x: ix0 - tile.x0,
y: iy0 - tile.y0,
z: 0,
},
aspect: wgpu::TextureAspect::All,
},
bytemuck::cast_slice(&sub),
wgpu::TexelCopyBufferLayout {
offset: 0,
bytes_per_row: Some(4 * ow),
rows_per_image: Some(oh),
},
wgpu::Extent3d {
width: ow,
height: oh,
depth_or_array_layers: 1,
},
);
}
}
pub(crate) fn write_uniforms(
&self,
queue: &wgpu::Queue,
ortho: [[f32; 4]; 4],
axis_log: [f32; 2],
) {
for tile in &self.tiles {
match self.kind {
GpuImageKind::Scalar {
cmap_min,
cmap_one_over_range,
gamma,
norm,
interp,
has_alpha_map,
nan_color,
} => {
let params = ImageParams {
ortho,
rect: tile.rect,
axis_log,
alpha: self.alpha,
cmap_min,
cmap_one_over_range,
gamma,
norm,
interp,
has_alpha_map,
_pad: [0; 3],
nan_color,
};
queue.write_buffer(&tile.params, 0, bytemuck::bytes_of(¶ms));
}
GpuImageKind::Rgba => {
let params = ImageRgbaParams {
ortho,
rect: tile.rect,
axis_log,
alpha: self.alpha,
_pad: 0.0,
};
queue.write_buffer(&tile.params, 0, bytemuck::bytes_of(¶ms));
}
}
}
}
pub(crate) fn draw(&self, render_pass: &mut wgpu::RenderPass<'_>, pipeline: &ImagePipeline) {
let pipe = match self.kind {
GpuImageKind::Scalar { .. } => &pipeline.pipeline,
GpuImageKind::Rgba => &pipeline.rgba_pipeline,
};
render_pass.set_pipeline(pipe);
for tile in &self.tiles {
render_pass.set_bind_group(0, &tile.bind_group, &[]);
render_pass.draw(0..6, 0..1);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn small_image_is_a_single_tile() {
let tiles = tile_bounds(100, 60, 8192);
assert_eq!(tiles, vec![(0, 0, 100, 60)]);
}
#[test]
fn image_params_is_std140_sized() {
assert_eq!(std::mem::size_of::<ImageParams>(), 144);
assert_eq!(std::mem::align_of::<ImageParams>(), 4);
}
#[test]
fn srgb_u8_to_linear_rgba_matches_transfer_function() {
let black = srgb_u8_to_linear_rgba([0, 0, 0, 255]);
assert_eq!(black, [0.0, 0.0, 0.0, 1.0]);
let white = srgb_u8_to_linear_rgba([255, 255, 255, 255]);
for ch in white {
assert!((ch - 1.0).abs() < 1e-6, "white channel {ch} != 1.0");
}
let nan = srgb_u8_to_linear_rgba([255, 255, 255, 0]);
assert_eq!(nan[3], 0.0, "default nan_color alpha must be transparent");
let mid = srgb_u8_to_linear_rgba([188, 188, 188, 128]);
let expected = ((188.0_f32 / 255.0 + 0.055) / 1.055).powf(2.4);
assert!((mid[0] - expected).abs() < 1e-6, "mid linearization");
assert!((mid[3] - 128.0 / 255.0).abs() < 1e-6, "alpha is linear");
}
#[test]
fn oversize_image_splits_into_grid_covering_every_pixel() {
let w = 5;
let h = 3;
let tiles = tile_bounds(w, h, 2);
assert_eq!(tiles.len(), 6);
let mut area = 0u32;
for (x0, y0, tw, th) in &tiles {
assert!(x0 + tw <= w && y0 + th <= h, "tile out of bounds");
assert!(*tw <= 2 && *th <= 2, "tile exceeds max_dim");
area += tw * th;
}
assert_eq!(area, w * h, "tiles must cover every pixel exactly once");
}
#[test]
fn extract_subgrid_copies_the_right_block() {
#[rustfmt::skip]
let data = vec![
0.0, 1.0, 2.0, 3.0,
4.0, 5.0, 6.0, 7.0,
8.0, 9.0, 10.0, 11.0,
];
let sub = extract_subgrid(&data, 4, 1, 1, 2, 2);
assert_eq!(sub, vec![5.0, 6.0, 9.0, 10.0]);
}
#[test]
fn extract_subgrid_works_on_rgba_texels() {
#[rustfmt::skip]
let data = vec![
[0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3],
[4, 4, 4, 4], [5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7],
];
let sub = extract_subgrid(&data, 4, 1, 0, 2, 2);
assert_eq!(
sub,
vec![[1, 1, 1, 1], [2, 2, 2, 2], [5, 5, 5, 5], [6, 6, 6, 6]]
);
}
#[test]
fn scalar_constructor_carries_a_colormap() {
let img = ImageData::new(2, 2, vec![0.0, 1.0, 2.0, 3.0], Colormap::viridis(0.0, 3.0));
assert!(matches!(img.pixels, ImagePixels::Scalar { .. }));
assert!(img.colormap().is_some());
assert_eq!((img.width, img.height), (2, 2));
}
#[test]
fn rgba_constructor_has_no_colormap() {
let px = vec![
[255, 0, 0, 255],
[0, 255, 0, 255],
[0, 0, 255, 255],
[0, 0, 0, 0],
];
let img = ImageData::rgba(2, 2, px);
assert!(matches!(img.pixels, ImagePixels::Rgba { .. }));
assert!(img.colormap().is_none());
assert_eq!((img.width, img.height), (2, 2));
}
#[test]
#[should_panic(expected = "data length must equal width * height")]
fn rgba_constructor_rejects_length_mismatch() {
ImageData::rgba(2, 2, vec![[0, 0, 0, 0]; 3]);
}
#[test]
fn interpolation_default_matches_silx() {
assert_eq!(InterpolationMode::default(), InterpolationMode::Nearest);
}
#[rustfmt::skip]
fn field_4x4() -> Vec<f32> {
vec![
0.0, 1.0, 2.0, 3.0,
4.0, 5.0, 6.0, 7.0,
8.0, 9.0, 10.0, 11.0,
12.0, 13.0, 14.0, 15.0,
]
}
#[test]
fn aggregate_max_over_4x4_into_2x2() {
let (out, w, h) = aggregate_blocks(&field_4x4(), 4, 4, 2, 2, AggregationMode::Max);
assert_eq!((w, h), (2, 2));
assert_eq!(out, vec![5.0, 7.0, 13.0, 15.0]);
}
#[test]
fn aggregate_min_over_4x4_into_2x2() {
let (out, w, h) = aggregate_blocks(&field_4x4(), 4, 4, 2, 2, AggregationMode::Min);
assert_eq!((w, h), (2, 2));
assert_eq!(out, vec![0.0, 2.0, 8.0, 10.0]);
}
#[test]
fn aggregate_mean_over_4x4_into_2x2() {
let (out, w, h) = aggregate_blocks(&field_4x4(), 4, 4, 2, 2, AggregationMode::Mean);
assert_eq!((w, h), (2, 2));
assert_eq!(out, vec![2.5, 4.5, 10.5, 12.5]);
}
#[test]
fn aggregate_none_returns_data_unchanged() {
let (out, w, h) = aggregate_blocks(&field_4x4(), 4, 4, 2, 2, AggregationMode::None);
assert_eq!((w, h), (4, 4));
assert_eq!(out, field_4x4());
}
#[test]
fn aggregate_block_factor_one_is_a_no_op() {
let (out, w, h) = aggregate_blocks(&field_4x4(), 4, 4, 1, 1, AggregationMode::Max);
assert_eq!((w, h), (4, 4));
assert_eq!(out, field_4x4());
}
#[test]
fn aggregate_drops_remainder_rows_and_cols() {
#[rustfmt::skip]
let data = vec![
0.0, 1.0, 2.0,
3.0, 4.0, 5.0,
6.0, 7.0, 8.0,
];
let (out, w, h) = aggregate_blocks(&data, 3, 3, 2, 2, AggregationMode::Max);
assert_eq!((w, h), (1, 1));
assert_eq!(out, vec![4.0]); }
#[test]
fn aggregate_ignores_nans_and_all_nan_block_is_nan() {
let nan = f32::NAN;
#[rustfmt::skip]
let data = vec![
nan, 1.0, nan, nan,
nan, 5.0, nan, nan,
0.0, 0.0, 0.0, 0.0,
0.0, 0.0, 0.0, 0.0,
];
let (out, _, _) = aggregate_blocks(&data, 4, 4, 2, 2, AggregationMode::Max);
assert_eq!(out[0], 5.0); assert!(out[1].is_nan()); assert_eq!(out[2], 0.0);
assert_eq!(out[3], 0.0);
}
#[test]
fn aggregate_mean_ignores_nans_in_count() {
let nan = f32::NAN;
#[rustfmt::skip]
let data = vec![
nan, 2.0,
4.0, nan,
];
let (out, w, h) = aggregate_blocks(&data, 2, 2, 2, 2, AggregationMode::Mean);
assert_eq!((w, h), (1, 1));
assert_eq!(out, vec![3.0]);
}
#[test]
fn aggregation_default_matches_silx() {
assert_eq!(AggregationMode::default(), AggregationMode::None);
}
#[test]
fn interpolation_codes_match_shader() {
assert_eq!(InterpolationMode::Nearest.code(), 0);
assert_eq!(InterpolationMode::Linear.code(), 1);
}
#[test]
fn new_image_defaults_to_nearest_interpolation() {
let img = ImageData::new(2, 2, vec![0.0, 1.0, 2.0, 3.0], Colormap::viridis(0.0, 3.0));
assert_eq!(img.interpolation, InterpolationMode::Nearest);
let img = img.with_interpolation(InterpolationMode::Linear);
assert_eq!(img.interpolation, InterpolationMode::Linear);
}
}