feat: Implement isCloseTo function (#1012)

Author: aleksanderkatan

packages/typegpu/src/data/vectorOps.ts

@@ -903,4 +903,25 @@ export const VectorOps = {
     vec4i: unary4i((value) => value - Math.floor(value)),
     vec4u: unary4u((value) => value - Math.floor(value)),
   } as Record<VecKind, <T extends vBase>(v: T) => T>,
+
+  isCloseToZero: {
+    vec2f: (v: wgsl.v2f, n: number) => Math.abs(v.x) <= n && Math.abs(v.y) <= n,
+    vec2h: (v: wgsl.v2h, n: number) => Math.abs(v.x) <= n && Math.abs(v.y) <= n,
+
+    vec3f: (v: wgsl.v3f, n: number) =>
+      Math.abs(v.x) <= n && Math.abs(v.y) <= n && Math.abs(v.z) <= n,
+    vec3h: (v: wgsl.v3h, n: number) =>
+      Math.abs(v.x) <= n && Math.abs(v.y) <= n && Math.abs(v.z) <= n,
+
+    vec4f: (v: wgsl.v4f, n: number) =>
+      Math.abs(v.x) <= n &&
+      Math.abs(v.y) <= n &&
+      Math.abs(v.z) <= n &&
+      Math.abs(v.w) <= n,
+    vec4h: (v: wgsl.v4h, n: number) =>
+      Math.abs(v.x) <= n &&
+      Math.abs(v.y) <= n &&
+      Math.abs(v.z) <= n &&
+      Math.abs(v.w) <= n,
+  } as Record<VecKind, <T extends vBase>(v: T, n: number) => boolean>,
 };

packages/typegpu/src/std/index.ts

@@ -23,6 +23,7 @@ export {
   mix,
   pow,
   reflect,
+  isCloseTo,
 } from './numeric.js';
 
 export {

packages/typegpu/src/std/numeric.ts

@@ -1,4 +1,4 @@
-import { f32 } from '../data/numeric';
+import { bool, f32 } from '../data/numeric';
 import { VectorOps } from '../data/vectorOps';
 import type {
   AnyMatInstance,
@@ -15,6 +15,19 @@ import type {
   vBaseForMat,
 } from '../data/wgslTypes';
 import { createDualImpl } from '../shared/generators';
+import type { Resource } from '../types';
+
+function isNumeric(element: Resource) {
+  const type = element.dataType.type;
+  return (
+    type === 'abstractInt' ||
+    type === 'abstractFloat' ||
+    type === 'f32' ||
+    type === 'f16' ||
+    type === 'i32' ||
+    type === 'u32'
+  );
+}
 
 type vBase = { kind: VecKind };
 
@@ -82,18 +95,17 @@ export const mul: MulOverload = createDualImpl(
   },
   // GPU implementation
   (s, v) => {
-    const returnType =
-      typeof s === 'number'
-        ? // Scalar * Vector/Matrix
-          (v.dataType as AnyWgslData)
-        : !s.dataType.type.startsWith('mat')
-          ? // Vector * Matrix
-            (s.dataType as AnyWgslData)
-          : !v.dataType.type.startsWith('mat')
-            ? // Matrix * Vector
-              (v.dataType as AnyWgslData)
-            : // Vector * Vector or Matrix * Matrix
-              (s.dataType as AnyWgslData);
+    const returnType = isNumeric(s)
+      ? // Scalar * Vector/Matrix
+        (v.dataType as AnyWgslData)
+      : !s.dataType.type.startsWith('mat')
+        ? // Vector * Matrix
+          (s.dataType as AnyWgslData)
+        : !v.dataType.type.startsWith('mat')
+          ? // Matrix * Vector
+            (v.dataType as AnyWgslData)
+          : // Vector * Vector or Matrix * Matrix
+            (s.dataType as AnyWgslData);
     return { value: `(${s.value} * ${v.value})`, dataType: returnType };
   },
 );
@@ -394,3 +406,51 @@ export const reflect = createDualImpl(
     };
   },
 );
+
+/**
+ * Checks whether the given elements differ by at most 0.01.
+ * Component-wise if arguments are vectors.
+ * @example
+ * isCloseTo(0, 0.1) // returns false
+ * isCloseTo(vec3f(0, 0, 0), vec3f(0.002, -0.009, 0)) // returns true
+ *
+ * @param {number} precision argument that specifies the maximum allowed difference, 0.01 by default.
+ */
+
+export const isCloseTo = createDualImpl(
+  // CPU implementation
+  <T extends v2f | v3f | v4f | v2h | v3h | v4h | number>(
+    e1: T,
+    e2: T,
+    precision = 0.01,
+  ) => {
+    if (typeof e1 === 'number' && typeof e2 === 'number') {
+      return Math.abs(e1 - e2) < precision;
+    }
+    if (typeof e1 !== 'number' && typeof e2 !== 'number') {
+      return VectorOps.isCloseToZero[e1.kind](sub(e1, e2), precision);
+    }
+    return false;
+  },
+  // GPU implementation
+  (e1, e2, precision = { value: 0.01, dataType: f32 }) => {
+    if (isNumeric(e1) && isNumeric(e2)) {
+      return {
+        value: `abs(f32(${e1.value})-f32(${e2.value})) <= ${precision.value}`,
+        dataType: bool,
+      };
+    }
+    if (!isNumeric(e1) && !isNumeric(e2)) {
+      return {
+        // https://www.w3.org/TR/WGSL/#vector-multi-component:~:text=Binary%20arithmetic%20expressions%20with%20mixed%20scalar%20and%20vector%20operands
+        // (a-a)+prec creates a vector of a.length elements, all equal to prec
+        value: `all(abs(${e1.value}-${e2.value}) <= (${e1.value} - ${e1.value})+${precision.value})`,
+        dataType: bool,
+      };
+    }
+    return {
+      value: 'false',
+      dataType: bool,
+    };
+  },
+);

packages/typegpu/tests/std/atan2.test.ts

@@ -0,0 +1,21 @@
+import { describe, expect, it } from 'vitest';
+import { vec4f } from '../../src/data';
+import { atan2, isCloseTo } from '../../src/std';
+
+describe('atan2', () => {
+  it('computes atan2 of two values', () => {
+    expect(atan2(0, 1)).toBeCloseTo(0);
+    expect(atan2(1, 0)).toBeCloseTo(Math.PI / 2);
+    expect(atan2(0, -1)).toBeCloseTo(Math.PI);
+    expect(atan2(-1, 0)).toBeCloseTo(-Math.PI / 2);
+  });
+
+  it('computes atan2 for two vectors', () => {
+    expect(
+      isCloseTo(
+        atan2(vec4f(0, 1, 0, -1), vec4f(1, 0, -1, 0)),
+        vec4f(0, Math.PI / 2, Math.PI, -Math.PI / 2),
+      ),
+    ).toBeTruthy();
+  });
+});

packages/typegpu/tests/std/isCloseTo.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from 'vitest';
+import { vec2f, vec2h, vec3f, vec3h, vec4f, vec4h } from '../../src/data';
+import { isCloseTo } from '../../src/std/numeric';
+
+describe('isCloseTo', () => {
+  it('returns true for close f32 containers', () => {
+    expect(isCloseTo(vec2f(0, 0), vec2f(0.0012, -0.009))).toBeTruthy();
+    expect(
+      isCloseTo(vec3f(1.05, -2.18, 1.22), vec3f(1.05, -2.17777, 1.229)),
+    ).toBeTruthy();
+    expect(
+      isCloseTo(
+        vec4f(3.8, -4.87, -2.42, -1.97),
+        vec4f(3.794, -4.861, -2.412, -1.971),
+      ),
+    ).toBeTruthy();
+  });
+
+  it('returns true for close f16 containers', () => {
+    expect(isCloseTo(vec2h(0, 0), vec2h(0.0012, -0.009))).toBeTruthy();
+    expect(
+      isCloseTo(vec3h(1.05, -2.18, 1.22), vec3h(1.05, -2.17777, 1.229)),
+    ).toBeTruthy();
+    expect(
+      isCloseTo(
+        vec4h(3.8, -4.87, -2.42, -1.97),
+        vec4h(3.794, -4.861, -2.412, -1.971),
+      ),
+    ).toBeTruthy();
+  });
+
+  it('returns true for close numbers', () => {
+    expect(isCloseTo(0, 0.009)).toBeTruthy();
+    expect(isCloseTo(0, 0.0009)).toBeTruthy();
+  });
+
+  it('returns false for distant f32 containers', () => {
+    expect(isCloseTo(vec2f(0, 0), vec2f(0, 1))).toBeFalsy();
+    expect(isCloseTo(vec3f(100, 100, 100), vec3f(101, 100, 100))).toBeFalsy();
+    expect(
+      isCloseTo(vec4f(1, 2, 3, 4), vec4f(1.02, 2.02, 3.02, 4.02)),
+    ).toBeFalsy();
+  });
+
+  it('returns false for distant f16 containers', () => {
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 1))).toBeFalsy();
+    expect(isCloseTo(vec3h(100, 100, 100), vec3h(101, 100, 100))).toBeFalsy();
+    expect(
+      isCloseTo(vec4h(1, 2, 3, 4), vec4h(1.02, 2.02, 3.02, 4.02)),
+    ).toBeFalsy();
+  });
+
+  it('returns false for distant numbers', () => {
+    expect(isCloseTo(0, 0.9)).toBeFalsy();
+    expect(isCloseTo(0, 0.09)).toBeFalsy();
+  });
+
+  it('applies precision correctly', () => {
+    // default precision of 0.01
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.009))).toBeTruthy();
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.011))).toBeFalsy();
+
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.09), 0.1)).toBeTruthy();
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.11), 0.1)).toBeFalsy();
+
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.0009), 0.001)).toBeTruthy();
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 0.0011), 0.001)).toBeFalsy();
+
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 9), 10)).toBeTruthy();
+    expect(isCloseTo(vec2h(0, 0), vec2h(0, 11), 10)).toBeFalsy();
+  });
+});