From 41213e45761fc1dd795d462ad7bc719533efd09e Mon Sep 17 00:00:00 2001
From: Sam Nystrom <sam@samnystrom.dev>
Date: Fri, 8 Mar 2024 14:52:06 +0000
Subject: Add Unzip node

---
 assembly/index.ts   | 13 +++++++++++--
 src/nodes/Unzip.tsx | 30 ++++++++++++++++++++++++++++++
 src/nodes/index.ts  |  4 ++++
 src/wasm.ts         |  1 +
 4 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100644 src/nodes/Unzip.tsx

diff --git a/assembly/index.ts b/assembly/index.ts
index fd4fdd5..7af2896 100644
--- a/assembly/index.ts
+++ b/assembly/index.ts
@@ -285,6 +285,15 @@ export function intersperse(a: StaticArray<f32>, b: StaticArray<f32>): StaticArr
 	return out;
 }
 
+export function unzip(x: StaticArray<f32>): StaticArray<f32> {
+	const out = new StaticArray<f32>(x.length);
+	for (let i = 0; i < x.length/2; i++) {
+		out[i] = x[i*2];
+		out[i + x.length/2] = x[i*2+1];
+	}
+	return out;
+}
+
 export function dft(x: StaticArray<f32>): StaticArray<f32> {
 	const out = new StaticArray<f32>(x.length);
 	for (let k = 0; k < out.length - out.length % 2; k += 2) {
@@ -292,8 +301,8 @@ export function dft(x: StaticArray<f32>): StaticArray<f32> {
 			const y = -<f32>2.0 * Mathf.PI * <f32>k / <f32>x.length * <f32>n;
 			const u = Mathf.cos(y);
 			const v = Mathf.sin(y);
-			out[k] = x[n] * u - x[n+1] * v;
-			out[k+1] = x[n] * v + x[n+1] * u;
+			out[k] += x[n] * u - x[n+1] * v;
+			out[k+1] += x[n] * v + x[n+1] * u;
 		}
 	}
 	return out;
diff --git a/src/nodes/Unzip.tsx b/src/nodes/Unzip.tsx
new file mode 100644
index 0000000..ff5832d
--- /dev/null
+++ b/src/nodes/Unzip.tsx
@@ -0,0 +1,30 @@
+import { NodeShell, InputArray, OutputNumber, NodeComponentProps, NodeInfo } from '../node.tsx';
+import { unzip } from '../wasm.ts';
+
+export interface UnzipInputs {
+	data: Float32Array,
+}
+
+export interface UnzipOutputs {
+	a: Float32Array,
+	b: Float32Array,
+}
+
+export const Unzip = ({ id, x, y, inputs }: NodeComponentProps<UnzipInputs>) => {
+	return (
+		<NodeShell name="Unzip" id={id} x={x} y={y}>
+			<OutputNumber name="a" label="Even" />
+			<OutputNumber name="b" label="Odd" />
+			<InputArray name="data" label="Data" />
+		</NodeShell>
+	);
+};
+
+export const UnzipNode: NodeInfo<UnzipInputs, UnzipOutputs> = {
+	component: Unzip,
+	func: ({ data }) => {
+		const out = data ? unzip(data) : null;
+		return { a: out?.slice(0, out.length/2), b: out?.slice(out.length/2) };
+	},
+	inputs: { data: null },
+};
\ No newline at end of file
diff --git a/src/nodes/index.ts b/src/nodes/index.ts
index a5678fa..09310bb 100644
--- a/src/nodes/index.ts
+++ b/src/nodes/index.ts
@@ -5,9 +5,12 @@ import { ViewerNode } from './Viewer.tsx';
 import { FourierNode } from './Fourier.tsx';
 import { LinspaceNode } from './Linspace.tsx';
 import { IntersperseNode } from './Intersperse.tsx';
+import { UnzipNode } from './Unzip.tsx';
 import { MathNode } from './Math.tsx';
 import { PlotNode } from './Plot.tsx';
 
+// TODO: ComplexMath
+
 const nodeRegistry: Record<string, NodeInfo<any>> = {
 	'Combine XYZ': CombineXYZNode,
 	'Separate XYZ': SeparateXYZNode,
@@ -15,6 +18,7 @@ const nodeRegistry: Record<string, NodeInfo<any>> = {
 	'Fourier Transform': FourierNode,
 	'Linspace': LinspaceNode,
 	'Intersperse': IntersperseNode,
+	'Unzip': UnzipNode,
 	'Math': MathNode,
 	'Plot': PlotNode,
 }
diff --git a/src/wasm.ts b/src/wasm.ts
index 311a088..a62a8fc 100644
--- a/src/wasm.ts
+++ b/src/wasm.ts
@@ -5,6 +5,7 @@ export const {
 	mathS, mathV, mathSS, mathSV, mathVS, mathVV,
 	linspace,
 	intersperse,
+	unzip,
 	dft,
 	fft,
 } = await instantiate(await WebAssembly.compileStreaming(fetch(url)), {});
\ No newline at end of file
-- 
cgit v1.2.3