Normalize, limit, broadcasting.

Author: tychedelia

crates/processing_core/src/constants.rs

@@ -78,6 +78,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";

crates/processing_pyo3/src/constants.rs

@@ -51,7 +51,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);

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<u32> {
         _ 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)?;

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};

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<u32>) {
     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<storage, read_write> a:   array<f32>;
@@ -177,7 +196,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     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)
 }
 

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::{

examples/combine_mix_test.rs

@@ -94,6 +94,28 @@ fn sketch() -> error::Result<bool> {
     buffer_destroy(tm)?;
     buffer_destroy(dm)?;
 
+    // one scalar per particle scales each whole vector
+    let vectors: Vec<f32> = (0..12).map(|i| i as f32).collect();
+    let v = buffer_create_with_data(to_bytes(&vectors))?;
+    let scalars = buffer_create_with_data(to_bytes(&[0.0, 1.0, 2.0, -1.0]))?;
+    combine(v, v, scalars, 3, COMBINE_MUL, 1.0, 0.0)?;
+    let got = f32s(&buffer_read(v)?);
+    let want: Vec<f32> = vectors
+        .iter()
+        .enumerate()
+        .map(|(i, x)| x * [0.0, 1.0, 2.0, -1.0][i / 3])
+        .collect();
+    let mismatch = buffer_create_with_data(to_bytes(&[1.0, 2.0]))?;
+    let rejected = combine(v, v, mismatch, 3, COMBINE_MUL, 1.0, 0.0).is_err();
+    if approx(&got, &want) && rejected {
+        println!("  PASS combine broadcasts a scalar per particle: {got:?}");
+    } else {
+        ok = false;
+        println!(
+            "  FAIL combine broadcast: got {got:?}, want {want:?}, mismatch rejected {rejected}"
+        );
+    }
+
     Ok(ok)
 }
 

examples/map_test.rs

@@ -62,6 +62,37 @@ fn sketch() -> error::Result<bool> {
     buffer_destroy(a2)?;
     buffer_destroy(dst)?;
 
+    // lengths 5, 0.5 and 0
+    let vectors = [3.0f32, 4.0, 0.0, 0.3, 0.0, 0.4, 0.0, 0.0, 0.0];
+    let cases = [
+        (
+            "normalize to 2",
+            MAP_NORMALIZE,
+            2.0,
+            0.0,
+            [1.2, 1.6, 0.0, 1.2, 0.0, 1.6, 0.0, 0.0, 0.0],
+        ),
+        (
+            "limit to 1..4",
+            MAP_LIMIT,
+            1.0,
+            4.0,
+            [2.4, 3.2, 0.0, 0.6, 0.0, 0.8, 0.0, 0.0, 0.0],
+        ),
+    ];
+    for (what, op, p0, p1, want) in cases {
+        let v = buffer_create_with_data(to_bytes(&vectors))?;
+        map(v, v, 3, op, p0, p1)?;
+        let got = f32s(&buffer_read(v)?);
+        if approx(&got, &want) {
+            println!("  PASS {what}: {got:?}");
+        } else {
+            ok = false;
+            println!("  FAIL {what}: got {got:?}, want {want:?}");
+        }
+        buffer_destroy(v)?;
+    }
+
     Ok(ok)
 }