From 04ac0cbf65db991ca4515fec439944f40eb8ce1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?charlotte=20=F0=9F=8C=B8?= Date: Mon, 28 Sep 2026 21:56:39 -0700 Subject: [PATCH] Normalize, limit, broadcasting. --- crates/processing_core/src/constants.rs | 2 + crates/processing_pyo3/src/constants.rs | 2 +- crates/processing_pyo3/src/particles.rs | 15 ++++++-- crates/processing_render/src/lib.rs | 6 +-- .../src/particles/algebra.rs | 37 +++++++++++++++++-- crates/processing_render/src/particles/mod.rs | 6 +-- 6 files changed, 55 insertions(+), 13 deletions(-) diff --git a/crates/processing_core/src/constants.rs b/crates/processing_core/src/constants.rs index d155448..bb61694 100644 --- a/crates/processing_core/src/constants.rs +++ b/crates/processing_core/src/constants.rs @@ -87,6 +87,8 @@ pub const AFFINE: &str = "affine"; pub const ABS: &str = "abs"; pub const NEGATE: &str = "negate"; pub const FLOOR: &str = "floor"; +pub const NORMALIZE: &str = "normalize"; +pub const LIMIT: &str = "limit"; pub const SQRT: &str = "sqrt"; pub const GREATER: &str = "greater"; pub const LESS: &str = "less"; diff --git a/crates/processing_pyo3/src/constants.rs b/crates/processing_pyo3/src/constants.rs index 999a190..bd1a743 100644 --- a/crates/processing_pyo3/src/constants.rs +++ b/crates/processing_pyo3/src/constants.rs @@ -62,7 +62,7 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { add!( m, MAP, COMBINE, MIX, LOOKUP, REDUCE, EXTRACT, PACK, GENERATE ); - add!(m, AFFINE, ABS, NEGATE, FLOOR, SQRT); + add!(m, AFFINE, ABS, NEGATE, FLOOR, SQRT, NORMALIZE, LIMIT); add!(m, GREATER, LESS, GEQ, LEQ, EQ, NEQ); add!(m, SUB, MUL, DIV, POW); add!(m, LENGTH, SUM, SUMSQ, MEAN, MIN, MAX); diff --git a/crates/processing_pyo3/src/particles.rs b/crates/processing_pyo3/src/particles.rs index 4676c1c..ae1497f 100644 --- a/crates/processing_pyo3/src/particles.rs +++ b/crates/processing_pyo3/src/particles.rs @@ -25,7 +25,7 @@ use processing_render::{ use processing_render::{GEN_GAUSSIAN, GEN_SIGNED, GEN_UNIFORM}; use processing_render::{ MAP_ABS, MAP_AFFINE, MAP_CLAMP, MAP_EQ, MAP_FLOOR, MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS, - MAP_NEGATE, MAP_NEQ, MAP_SQRT, MAP_SQUARE, + MAP_LIMIT, MAP_NEGATE, MAP_NEQ, MAP_NORMALIZE, MAP_SQRT, MAP_SQUARE, }; use processing_render::{ REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, REDUCE_SUMSQ, @@ -46,6 +46,8 @@ fn parse_map_op(s: &str) -> PyResult { _ if s.eq_ignore_ascii_case(c::LEQ) => Ok(MAP_LEQ), _ if s.eq_ignore_ascii_case(c::EQ) => Ok(MAP_EQ), _ if s.eq_ignore_ascii_case(c::NEQ) => Ok(MAP_NEQ), + _ if s.eq_ignore_ascii_case(c::NORMALIZE) => Ok(MAP_NORMALIZE), + _ if s.eq_ignore_ascii_case(c::LIMIT) => Ok(MAP_LIMIT), _ => Err(PyValueError::new_err(format!("map: unknown op {s:?}"))), } } @@ -239,6 +241,11 @@ fn map_params(kwargs: Option<&Bound<'_, PyDict>>, op: u32) -> PyResult<(f32, f32 kw_f32(kwargs, "offset", 0.0)?, ), MAP_CLAMP => (kw_f32(kwargs, "lo", 0.0)?, kw_f32(kwargs, "hi", 1.0)?), + MAP_NORMALIZE => (kw_f32(kwargs, "length", 1.0)?, 0.0), + MAP_LIMIT => ( + kw_f32(kwargs, "min_length", 0.0)?, + kw_f32(kwargs, "max_length", f32::MAX)?, + ), MAP_GREATER | MAP_LESS | MAP_GEQ | MAP_LEQ | MAP_EQ | MAP_NEQ => ( kw_f32(kwargs, "threshold", 0.0)?, kw_f32(kwargs, "epsilon", 1.0e-6)?, @@ -251,6 +258,8 @@ fn map_param_keys(op: u32) -> &'static [&'static str] { match op { MAP_AFFINE => &["scale", "offset"], MAP_CLAMP => &["lo", "hi"], + MAP_NORMALIZE => &["length"], + MAP_LIMIT => &["min_length", "max_length"], MAP_GREATER | MAP_LESS | MAP_GEQ | MAP_LEQ | MAP_EQ | MAP_NEQ => &["threshold", "epsilon"], _ => &[], } @@ -567,9 +576,9 @@ impl Particles { reject_unknown_kwargs(kwargs, &["a", "b", "out", "op", "b_scale", "b_offset"])?; let (a, comp) = self.operand(kwargs, "a")?; let (b, b_comp) = self.operand(kwargs, "b")?; - if b_comp != comp { + if b_comp != comp && b_comp != 1 { return Err(PyValueError::new_err(format!( - "apply(combine): `a` has {comp} components but `b` has {b_comp} (must match)" + "apply(combine): the first operand has {comp} components, so the second needs {comp} or 1, not {b_comp}" ))); } let out = self.dest(kwargs, a)?; diff --git a/crates/processing_render/src/lib.rs b/crates/processing_render/src/lib.rs index 31b22b3..10c65cb 100644 --- a/crates/processing_render/src/lib.rs +++ b/crates/processing_render/src/lib.rs @@ -23,9 +23,9 @@ pub mod transform; pub use particles::algebra::{ GEN_GAUSSIAN, GEN_SIGNED, GEN_UNIFORM, MAP_ABS, MAP_AFFINE, MAP_CLAMP, MAP_EQ, MAP_FLOOR, - MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS, MAP_NEGATE, MAP_NEQ, MAP_SQRT, MAP_SQUARE, - REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, REDUCE_SUMSQ, combine, extract, - generate, lookup, map, mix, pack, reduce_components, + MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS, MAP_LIMIT, MAP_NEGATE, MAP_NEQ, MAP_NORMALIZE, + MAP_SQRT, MAP_SQUARE, REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, + REDUCE_SUMSQ, combine, extract, generate, lookup, map, mix, pack, reduce_components, }; pub use particles::compact::compact; pub use particles::grid::{Grid, GridParams, grid_build, grid_create, grid_destroy, grid_get}; diff --git a/crates/processing_render/src/particles/algebra.rs b/crates/processing_render/src/particles/algebra.rs index 313ff98..f2d3829 100644 --- a/crates/processing_render/src/particles/algebra.rs +++ b/crates/processing_render/src/particles/algebra.rs @@ -22,6 +22,8 @@ pub const MAP_GEQ: u32 = 9; pub const MAP_LEQ: u32 = 10; pub const MAP_EQ: u32 = 11; pub const MAP_NEQ: u32 = 12; +pub const MAP_NORMALIZE: u32 = 13; +pub const MAP_LIMIT: u32 = 14; const MAP_SRC: &str = r#" struct Params { @@ -60,10 +62,25 @@ fn main(@builtin(global_invocation_id) gid: vec3) { let i = gid.x; let n = arrayLength(&a) / params.components; if i >= n { return; } + // normalize and limit scale the whole vector by its length + var len2 = 0.0; + for (var c = 0u; c < params.components; c = c + 1u) { + let x = a[i * params.components + c]; + len2 += x * x; + } + let len = sqrt(len2); + var k = 1.0; + if params.op == 13u && len > 0.0 { + k = params.p0 / len; + } else if params.op == 14u && len > 0.0 { + k = clamp(len, params.p0, params.p1) / len; + } for (var c = 0u; c < params.components; c = c + 1u) { let idx = i * params.components + c; - @if(in_place) { a[idx] = apply_op(a[idx]); } - @if(!in_place) { dst[idx] = apply_op(a[idx]); } + var r = a[idx] * k; + if params.op < 13u { r = apply_op(a[idx]); } + @if(in_place) { a[idx] = r; } + @if(!in_place) { dst[idx] = r; } } } "#; @@ -149,6 +166,8 @@ struct Params { op: u32, b_scale: f32, b_offset: f32, + // 1 = one `b` per particle, applied to every component + b_components: u32, } @if(in_place) @group(0) @binding(0) var a: array; @@ -177,7 +196,8 @@ fn main(@builtin(global_invocation_id) gid: vec3) { if i >= n { return; } for (var c = 0u; c < params.components; c = c + 1u) { let idx = i * params.components + c; - let r = combine_op(a[idx], b[idx] * params.b_scale + params.b_offset); + let b_idx = select(idx, i, params.b_components == 1u); + let r = combine_op(a[idx], b[b_idx] * params.b_scale + params.b_offset); @if(in_place) { a[idx] = r; } @if(!in_place) { dst[idx] = r; } } @@ -198,6 +218,16 @@ pub fn combine( check_components("combine", components)?; let v = variants(&COMBINE, COMBINE_SRC)?; let floats = buffer_size(a)? / 4; + let b_floats = buffer_size(b)? / 4; + let b_components = if b_floats == floats { + components + } else if b_floats == floats / components as u64 { + 1 + } else { + return Err(ProcessingError::InvalidArgument(format!( + "combine: `b` needs {components} or 1 components per particle" + ))); + }; let c = if dst == a { ensure_no_alias("combine", a, &[b])?; @@ -216,6 +246,7 @@ pub fn combine( compute_set(c, "op", ShaderValue::UInt(op))?; compute_set(c, "b_scale", ShaderValue::Float(b_scale))?; compute_set(c, "b_offset", ShaderValue::Float(b_offset))?; + compute_set(c, "b_components", ShaderValue::UInt(b_components))?; dispatch_particles(c, floats, components) } diff --git a/crates/processing_render/src/particles/mod.rs b/crates/processing_render/src/particles/mod.rs index 5088726..ddc33a1 100644 --- a/crates/processing_render/src/particles/mod.rs +++ b/crates/processing_render/src/particles/mod.rs @@ -15,9 +15,9 @@ pub mod sort; pub use algebra::{ GEN_GAUSSIAN, GEN_SIGNED, GEN_UNIFORM, MAP_ABS, MAP_AFFINE, MAP_CLAMP, MAP_EQ, MAP_FLOOR, - MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS, MAP_NEGATE, MAP_NEQ, MAP_SQRT, MAP_SQUARE, - REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, REDUCE_SUMSQ, combine, extract, - generate, lookup, map, mix, pack, reduce_components, + MAP_GEQ, MAP_GREATER, MAP_LEQ, MAP_LESS, MAP_LIMIT, MAP_NEGATE, MAP_NEQ, MAP_NORMALIZE, + MAP_SQRT, MAP_SQUARE, REDUCE_LENGTH, REDUCE_MAX, REDUCE_MEAN, REDUCE_MIN, REDUCE_SUM, + REDUCE_SUMSQ, combine, extract, generate, lookup, map, mix, pack, reduce_components, }; pub use compact::compact; pub use emit::{