Implement loading and dumping of complex64 and complex128 dtypes

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/dump.test.ts

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

test/generate-test-data.py

@@ -18,19 +18,35 @@
 else:
     records = {}
 
+
+def serialize_tail_values(data, dtype):
+    if dtype in ["complex64", "complex128"]:
+        float_dtype = np.float32 if dtype == "complex64" else np.float64
+        flat_view = data.ravel().view(float_dtype)
+        return flat_view[-5:].tolist()
+    else:
+        return data.ravel()[-5:].tolist()
+
+
 # Generate test data for each combination
 for dimensions in [(10,), (65, 65), (100, 100, 100), (4, 4, 4, 4, 4)]:
-    for dtype in ["int8", "int16", "int64", "float16", "float32", "float64"]:
+    for dtype in ["int8", "int16", "int64", "float16", "float32", "float64", "complex64", "complex128"]:
         name = f"./data/{'x'.join(str(i) for i in dimensions)}-{dtype}"
-        
+
         # Skip if file already exists
         if name in records:
             continue
-
-        data = np.random.randint(0, 255, dimensions).astype(dtype)
+
+        if dtype in ["complex64", "complex128"]:
+            real_part = np.random.randint(-128, 128, dimensions).astype(np.float64)
+            imag_part = np.random.randint(-128, 128, dimensions).astype(np.float64)
+            data = (real_part + 1j * imag_part).astype(dtype)
+        else:
+            data = np.random.randint(0, 255, dimensions).astype(dtype)
+
         # Store the last 5 values consistently for all types
-        records[name] = data.ravel()[-5:].tolist()
-        
+        records[name] = serialize_tail_values(data, dtype)
+
         # Save file using the correct path
         file_path = script_dir / name.lstrip("./")
         file_path.parent.mkdir(parents=True, exist_ok=True)

test/load.test.ts

@@ -39,8 +39,8 @@ function checkTail(expectedDtype: string, actual: Array<number | bigint>, expect
         const isExpectedNaN = typeof expected === "number" && Number.isNaN(expected);
 
         // Detect dtype class
-        const dtype: string = expectedDtype || ""; // e.g., "float32", "f4", "i8", "u8"
-        const isFloat = /^f\d$/.test(dtype) || /float/i.test(dtype);
+        const dtype: string = expectedDtype || ""; // e.g., "float32", "f4", "i8", "u8", "c8", "c16"
+        const isFloat = /^f\d$/.test(dtype) || /float/i.test(dtype) || /^c\d+$/.test(dtype);
         const isI64 = dtype === "i8" || /int64/i.test(dtype);
         const isU64 = dtype === "u8" || /uint64/i.test(dtype);
 

test/records.json

@@ -1,4 +1,18 @@
 {
+    "./data/10-complex128": [
+        -107.0,
+        94.0,
+        102.0,
+        9.0,
+        55.0
+    ],
+    "./data/10-complex64": [
+        85.0,
+        14.0,
+        -104.0,
+        -27.0,
+        -103.0
+    ],
     "./data/10-float16": [
         22.0,
         85.0,
@@ -83,6 +97,20 @@
         89,
         90
     ],
+    "./data/4x4x4x4x4-complex128": [
+        27.0,
+        -54.0,
+        68.0,
+        -76.0,
+        -81.0
+    ],
+    "./data/4x4x4x4x4-complex64": [
+        84.0,
+        12.0,
+        75.0,
+        -15.0,
+        -54.0
+    ],
     "./data/4x4x4x4x4-float16": [
         234.0,
         76.0,
@@ -125,6 +153,20 @@
         119,
         -9
     ],
+    "./data/65x65-complex128": [
+        100.0,
+        -96.0,
+        123.0,
+        13.0,
+        -95.0
+    ],
+    "./data/65x65-complex64": [
+        -3.0,
+        41.0,
+        -101.0,
+        71.0,
+        -24.0
+    ],
     "./data/65x65-float16": [
         163.0,
         254.0,
@@ -167,6 +209,10 @@
         -99,
         -49
     ],
+    "./data/bool": [
+        true,
+        false
+    ],
     "./data/unicode": [
         "123",
         "test",