Normalize, limit, broadcasting.

Author: tychedelia

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

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

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