Merge pull request #63 from axelhzf/complex-dtype-support

Implement loading and dumping of complex64 and complex128 dtypes

Author: j6k4m8

README.md

@@ -103,6 +103,8 @@ console.log(tensor.get(10, 15));
 -   `float32`
 -   `float64`
 -   `float16` (converted to float32 by default)
+-   `complex64` (as `Float32Array` with interleaved real/imag)
+-   `complex128` (as `Float64Array` with interleaved real/imag)
 
 ### Float16 Control
 
@@ -114,6 +116,56 @@ const n1 = new npyjs();
 const n2 = new npyjs({ convertFloat16: false });
 ```
 
+### Complex Numbers
+
+Complex arrays are returned as typed arrays with interleaved real and imaginary parts: `[real0, imag0, real1, imag1, ...]`
+
+```ts
+import { load } from "npyjs";
+
+const { data, shape } = await load("complex-array.npy");
+// For a shape of [3], data will have 6 elements: [re0, im0, re1, im1, re2, im2]
+
+// Access the first complex number
+const real0 = data[0];
+const imag0 = data[1];
+```
+
+---
+
+## Writing .npy Files
+
+Use the `dump` function to create `.npy` files:
+
+```ts
+import { dump } from "npyjs";
+import { writeFileSync } from "fs";
+
+// Dump a typed array
+const arr = new Float32Array([1.0, 2.0, 3.0, 4.0]);
+const bytes = dump(arr, [2, 2]); // 2x2 shape
+writeFileSync("output.npy", Buffer.from(bytes));
+
+// Dump a plain array (dtype is inferred)
+const plain = [1, 2, 3, 4];
+const bytes2 = dump(plain, [4]);
+```
+
+### Dumping Complex Arrays
+
+Since complex types cannot be inferred from plain number arrays, use the `dtype` option:
+
+```ts
+import { dump } from "npyjs";
+
+// Complex array: 1+2j, 3-4j as interleaved [real, imag, ...]
+const complexData = [1, 2, 3, -4];
+const bytes = dump(complexData, [2], { dtype: "c8" });  // complex64
+
+// Or use c16 for complex128
+const bytes128 = dump(complexData, [2], { dtype: "c16" });
+```
+
 ---
 
 ## Development

docs/index.md

@@ -97,6 +97,8 @@ console.log(tensor.get(10, 15));
 -   `float32`
 -   `float64`
 -   `float16` (converted to float32 by default)
+-   `complex64` (as `Float32Array` with interleaved real/imag)
+-   `complex128` (as `Float64Array` with interleaved real/imag)
 
 ### Float16 Control
 
@@ -108,6 +110,56 @@ const n1 = new npyjs();
 const n2 = new npyjs({ convertFloat16: false });
 ```
 
+### Complex Numbers
+
+Complex arrays are returned as typed arrays with interleaved real and imaginary parts: `[real0, imag0, real1, imag1, ...]`
+
+```ts
+import { load } from "npyjs";
+
+const { data, shape } = await load("complex-array.npy");
+// For a shape of [3], data will have 6 elements: [re0, im0, re1, im1, re2, im2]
+
+// Access the first complex number
+const real0 = data[0];
+const imag0 = data[1];
+```
+
+---
+
+## Writing .npy Files
+
+Use the `dump` function to create `.npy` files:
+
+```ts
+import { dump } from "npyjs";
+import { writeFileSync } from "fs";
+
+// Dump a typed array
+const arr = new Float32Array([1.0, 2.0, 3.0, 4.0]);
+const bytes = dump(arr, [2, 2]); // 2x2 shape
+writeFileSync("output.npy", Buffer.from(bytes));
+
+// Dump a plain array (dtype is inferred)
+const plain = [1, 2, 3, 4];
+const bytes2 = dump(plain, [4]);
+```
+
+### Dumping Complex Arrays
+
+Since complex types cannot be inferred from plain number arrays, use the `dtype` option:
+
+```ts
+import { dump } from "npyjs";
+
+// Complex array: 1+2j, 3-4j as interleaved [real, imag, ...]
+const complexData = [1, 2, 3, -4];
+const bytes = dump(complexData, [2], { dtype: "c8" });  // complex64
+
+// Or use c16 for complex128
+const bytes128 = dump(complexData, [2], { dtype: "c16" });
+```
+
 ---
 
 ## Development

src/index.ts

@@ -1,6 +1,6 @@
 export type DType =
     | "i1" | "u1" | "i2" | "u2" | "i4" | "u4" | "i8" | "u8"
-    | "f2" | "f4" | "f8" | "b1" | `U${number}`; // e.g., U10 for strings of length 10
+    | "f2" | "f4" | "f8" | "b1" | "c8" | "c16" | `U${number}`; // e.g., U10 for strings of length 10
 
 export type TypedArray =
     | Int8Array
@@ -27,6 +27,14 @@ export interface Options {
     convertFloat16?: boolean;
 }
 
+export interface DumpOptions {
+    /**
+     * Specify the dtype for the output. Required for complex types (c8, c16)
+     * since they cannot be inferred from a plain number array.
+     */
+    dtype?: DType;
+}
+
 class StringFromCodePoint extends String {
     constructor(buf: ArrayBufferLike, byteOffset?: number, length?: number) {
         const uint32 = new Uint32Array(buf, byteOffset, length);
@@ -109,6 +117,8 @@ function dtypeToArray(dtype: string, buf: ArrayBufferLike, offset: number, opts:
         case "u8": return new BigUint64Array(buf, offset);
         case "f4": return new Float32Array(buf, offset);
         case "f8": return new Float64Array(buf, offset);
+        case "c8": return new Float32Array(buf, offset); 
+        case "c16": return new Float64Array(buf, offset); 
         case "f2": {
             if (opts.convertFloat16 !== false) {
                 const u16 = new Uint16Array(buf, offset);
@@ -235,6 +245,8 @@ export function arrayToTypedArray(dtype: DType, array: ArrayLike<number | string
         case "u8": return new BigUint64Array(array);
         case "f4": return new Float32Array(array);
         case "f8": return new Float64Array(array);
+        case "c8": return new Float32Array(array);
+        case "c16": return new Float64Array(array);
         default: throw new Error(`Unsupported dtype: ${dtype}`);
     }
 }
@@ -332,11 +344,20 @@ function createPyDescription(dtype : DType, shape: number[]) : string {
     return `{'descr':'${descr}','fortran_order':False,'shape':(${pyShape})}`;
 }
 
-export function dump(array: TypedArray | Array<number | string>, shape: number[] | undefined) : ArrayBuffer{
-    const dtype = array instanceof Array ? inferDtypeFromArray(array) : arrayToDtype(array);
+export function dump(array: TypedArray | Array<number | string>, shape?: number[], options?: DumpOptions) : ArrayBuffer{
+    let dtype: DType;
+    if (options?.dtype) {
+        dtype = options.dtype;
+    } else if (array instanceof Array) {
+        dtype = inferDtypeFromArray(array);
+    } else {
+        dtype = arrayToDtype(array);
+    }
     array = array instanceof Array ? arrayToTypedArray(dtype, array) : array;
-    
-    let pyDesc = createPyDescription(dtype, shape ?? [array.length]);
+
+    // For complex types, shape refers to number of complex elements, not the flat array length
+    const effectiveLength = (dtype === "c8" || dtype === "c16") ? array.length / 2 : array.length;
+    let pyDesc = createPyDescription(dtype, shape ?? [effectiveLength]);
     let headerSize = 10 + pyDesc.length;
     const pad = 8 - ((headerSize + 1) % 8);
     pyDesc = pyDesc + " ".repeat(pad) + "\x0A";
@@ -368,7 +389,7 @@ export default class N {
         return f16toF32(u16);
     }
 
-    dump(array: TypedArray | Array<number | string>, shape: number[]) {
-        return dump(array, shape);
+    dump(array: TypedArray | Array<number | string>, shape?: number[], options?: DumpOptions) {
+        return dump(array, shape, options);
     }
 }

test/data/10-complex128.npy

@@ -61,4 +61,54 @@ describe("npyjs dump", () => {
             expect(error.toString()).toContain("MyArray")
         }
     });
+
+    it("complex64 (c8) with plain array", async () => {
+        const complexData = [1, 2, 3, -4, 5, 6];
+        const bytes = npyjs.dump(complexData, [3], { dtype: "c8" });
+        const result = await npyjs.load(bytes);
+
+        expect(result.dtype).toBe("c8");
+        expect(result.shape).toEqual([3]);
+        expect(result.data).toEqual(new Float32Array(complexData));
+    });
+
+    it("complex128 (c16) with plain array", async () => {
+        const complexData = [1.5, 2.5, -3.5, 4.5];
+        const bytes = npyjs.dump(complexData, [2], { dtype: "c16" });
+        const result = await npyjs.load(bytes);
+
+        expect(result.dtype).toBe("c16");
+        expect(result.shape).toEqual([2]);
+        expect(result.data).toEqual(new Float64Array(complexData));
+    });
+
+    it("complex64 (c8) with Float32Array", async () => {
+        const original = new Float32Array([1, 2, 3, -4, 5.5, 6.5]);
+        const bytes = npyjs.dump(original, [3], { dtype: "c8" });
+        const result = await npyjs.load(bytes);
+
+        expect(result.dtype).toBe("c8");
+        expect(result.shape).toEqual([3]);
+        expect(result.data).toEqual(original);
+    });
+
+    it("complex128 (c16) with Float64Array", async () => {
+        const original = new Float64Array([1.1, 2.2, 3.3, -4.4]);
+        const bytes = npyjs.dump(original, [2], { dtype: "c16" });
+        const result = await npyjs.load(bytes);
+
+        expect(result.dtype).toBe("c16");
+        expect(result.shape).toEqual([2]);
+        expect(result.data).toEqual(original);
+    });
+
+    it("2D complex array", async () => {
+        const complexData = [1, 1, 2, 2, 3, 3, 4, 4];
+        const bytes = npyjs.dump(complexData, [2, 2], { dtype: "c8" });
+        const result = await npyjs.load(bytes);
+
+        expect(result.dtype).toBe("c8");
+        expect(result.shape).toEqual([2, 2]);
+        expect(result.data).toEqual(new Float32Array(complexData));
+    });
 });