// --- factory functions ---
import { NDArray, DTYPE_MAP } from './ndarray_core.js';
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @param {string} [dtype='float64'] - The desired data type.
* @return {NDArray}
*/
export function array(source, dtype = 'float64') {
const Ctor = DTYPE_MAP[dtype] || Float64Array;
const isTyped = (v) => ArrayBuffer.isView(v) && !(v instanceof DataView);
// 1. Shape Inference & Strict Validation
const getShape = (v) => {
if (isTyped(v)) return [v.length];
if (!Array.isArray(v)) return []; // Scalar
if (v.length === 0) throw new Error("Input array cannot be empty.");
const sub = getShape(v[0]);
const subStr = sub.join(',');
for (let i = 1; i < v.length; i++) {
if (getShape(v[i]).join(',') !== subStr) {
throw new Error(`Jagged array detected at index ${i}.`);
}
}
return [v.length, ...sub];
};
const shape = getShape(source);
const size = shape.reduce((a, b) => a * b, 1);
const data = new Ctor(size);
// 2. Fill Data
// Special handling for scalar (0-dim) to avoid logic complexity in loop
if (shape.length === 0) {
data[0] = source;
} else {
let offset = 0;
const lastDim = shape.length - 1; // The index of the innermost dimension
const fill = (v, dim) => {
if (dim === lastDim) {
// Optimization: We are at the bottom level.
// v is guaranteed to be an Array-like (JS Array or TypedArray)
// data.set is the fastest way to copy a block of memory.
data.set(v, offset);
offset += v.length;
} else {
// We are not at the bottom. v is guaranteed to be an Array.
// Recurse deeper.
for (let i = 0; i < v.length; i++) {
fill(v[i], dim + 1);
}
}
};
fill(source, 0);
}
return new NDArray(data, { shape, dtype });
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function float64(source){
return array(source,"float64");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function float32(source){
return array(source,"float32");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function uint32(source){
return array(source,"uint32");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function int32(source){
return array(source,"int32");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function int16(source){
return array(source,"int16");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function uint16(source){
return array(source,"uint16");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function int8(source){
return array(source,"int8");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function uint8(source){
return array(source,"uint8");
}
/**
* Creates an NDArray from a regular array or TypedArray.
* @param {Array|TypedArray} source - The source data.
* @return {NDArray}
*/
export function uint8c(source){
return array(source,"uint8c");
}
/**
* Creates a new NDArray of the given shape, filled with zeros.
* @param {Array<number>} shape - The shape of the new array.
* @param {string} [dtype='float64'] - The data type of the new array.
* @returns {NDArray}
*/
export function zeros(shape, dtype = 'float64') {
const size = shape.reduce((a, b) => a * b, 1);
const Constructor = DTYPE_MAP[dtype];
return new NDArray(new Constructor(size), { shape, dtype });
}
/**
* Creates a new NDArray of the given shape, filled with ones.
* @param {Array<number>} shape - The shape of the new array.
* @param {string} [dtype='float64'] - The data type of the new array.
* @returns {NDArray}
*/
export function ones(shape, dtype = 'float64') {
const arr = zeros(shape, dtype);
arr.data.fill(1);
return arr;
}
/**
* Creates a new NDArray of the given shape, filled with a specified value.
* @param {Array<number>} shape - The shape of the new array.
* @param {*} value - The value to fill the array with.
* @param {string} [dtype='float64'] - The data type of the new array.
* @returns {NDArray}
*/
export function full(shape, value, dtype = 'float64') {
const arr = zeros(shape, dtype);
arr.data.fill(value);
return arr;
}
/**
* Like Python's arange: [start, stop)
* @param {number} start the starting value of the sequence.
* @param {number | null} stop the end value of the sequence.
* @param {number | null} [step=1] the spacing between values.
* @param {string} [dtype='float64'] the data type of the resulting array.
* @returns {NDArray}
*/
export function arange(start, stop = null, step = 1, dtype = 'float64') {
if (stop === null) {
stop = start;
start = 0;
}
const size = Math.ceil((stop - start) / step);
const Constructor = DTYPE_MAP[dtype];
const data = new Constructor(size);
for (let i = 0; i < size; i++) {
data[i] = start + i * step;
}
return new NDArray(data, { shape: [size], dtype });
}
/**
* Linearly spaced points.
* @param {number} start the starting value of the sequence.
* @param {number | null} stop the end value of the sequence.
* @param {number} [num=50] the number of samples to generate.
* @param {string} [dtype='float64'] the data type of the resulting array.
* @returns {NDArray}
*/
export function linspace(start, stop, num = 50, dtype = 'float64') {
const step = (stop - start) / (num - 1);
const Constructor = DTYPE_MAP[dtype];
const data = new Constructor(num);
for (let i = 0; i < num; i++) {
data[i] = start + i * step;
}
return new NDArray(data, { shape: [num], dtype });
}
/**
* where(condition, x, y)
* Returns a new array with elements chosen from x or y depending on condition.
* Supports NumPy-style broadcasting across all three arguments, including 0-sized dimensions.
*
* @param {NDArray|Array|number} condition - Where True, yield x, otherwise yield y.
* @param {NDArray|Array|number} x - Values from which to choose if condition is True.
* @param {NDArray|Array|number} y - Values from which to choose if condition is False.
* @returns {NDArray} A new contiguous NDArray.
*/
export function where(condition, x, y) {
// 1. Standardize inputs to NDArray
const toNDArray = (val, defaultDtype = 'float64') => {
return (val instanceof NDArray) ? val : array(val, defaultDtype);
};
const cond = toNDArray(condition, 'int8');
const dtypes = [(x instanceof NDArray)?x.dtype:null , (y instanceof NDArray)?y.dtype:null];
let dtype = null;
for(let i of dtypes){
if(dtype === null){
dtype = i;
}else if(dtype !== i && i !== null){
throw new Error(`dtype is different ${dtype} ${i}`);
}
}
if(dtype===null){
dtype="float64";
}
const xArr = toNDArray(x, dtype);
const yArr = toNDArray(y, dtype);
// 2. Determine output shape using strict NumPy broadcasting rules
const outNdim = Math.max(cond.ndim, xArr.ndim, yArr.ndim);
const outShape = new Int32Array(outNdim);
for (let i = 1; i <= outNdim; i++) {
// Dimension sizes from the right (defaulting to 1 for missing dimensions)
const dC = (cond.ndim - i >= 0) ? cond.shape[cond.ndim - i] : 1;
const dX = (xArr.ndim - i >= 0) ? xArr.shape[xArr.ndim - i] : 1;
const dY = (yArr.ndim - i >= 0) ? yArr.shape[yArr.ndim - i] : 1;
// Broadcasting logic:
// A dimension is compatible if all values are equal, or if some are 1.
let maxDim = 1;
const dims = [dC, dX, dY];
for (let j = 0; j < 3; j++) {
const d = dims[j];
if (d !== 1) {
if (maxDim !== 1 && maxDim !== d) {
throw new Error(`Incompatible shapes for where: condition [${cond.shape}], x [${xArr.shape}], y [${yArr.shape}]`);
}
maxDim = d;
}
}
outShape[outNdim - i] = maxDim;
}
// 3. Prepare traversal metadata
// reduce returns 1 for outNdim == 0 (scalars) and 0 if any dim is 0.
const outSize = outShape.reduce((a, b) => a * b, 1);
const sC = new Int32Array(outNdim);
const sX = new Int32Array(outNdim);
const sY = new Int32Array(outNdim);
const wrapC = new Float64Array(outNdim);
const wrapX = new Float64Array(outNdim);
const wrapY = new Float64Array(outNdim);
for (let i = 0; i < outNdim; i++) {
const dim = outShape[i];
const axisFromRight = outNdim - i;
const getStride = (arr, axisRight) => {
const axis = arr.ndim - axisRight;
// Broadcast (stride 0) if dimension is missing (axis < 0) or size is 1
if (axis < 0 || arr.shape[axis] === 1) return 0;
return arr.strides[axis];
};
sC[i] = getStride(cond, axisFromRight);
sX[i] = getStride(xArr, axisFromRight);
sY[i] = getStride(yArr, axisFromRight);
// wrapBack is the distance to jump from the last element back to the first in a dimension
wrapC[i] = (dim - 1) * sC[i];
wrapX[i] = (dim - 1) * sX[i];
wrapY[i] = (dim - 1) * sY[i];
}
// 4. Execution Engine
const Constructor = DTYPE_MAP[xArr.dtype] || Float64Array;
const newData = new Constructor(outSize);
const currentIdx = new Int32Array(outNdim);
let pC = cond.offset;
let pX = xArr.offset;
let pY = yArr.offset;
const dataC = cond.data;
const dataX = xArr.data;
const dataY = yArr.data;
// Loop runs outSize times. If outSize is 0, it doesn't run at all.
// If outSize is 1 (0D scalar), it runs exactly once.
for (let i = 0; i < outSize; i++) {
// Core ternary selection
newData[i] = dataC[pC] ? dataX[pX] : dataY[pY];
// Odometer increment with manual pointer updates
for (let d = outNdim - 1; d >= 0; d--) {
if (++currentIdx[d] < outShape[d]) {
pC += sC[d];
pX += sX[d];
pY += sY[d];
break;
} else {
currentIdx[d] = 0;
pC -= wrapC[d];
pX -= wrapX[d];
pY -= wrapY[d];
}
}
}
return new NDArray(newData, {
shape: outShape,
dtype: xArr.dtype
});
};
/**
* Joins a sequence of arrays along an existing axis.
* @param {Array<NDArray>} arrays The arrays must have the same shape, except in the dimension corresponding to `axis`.
* @param {number} [axis=0] The axis along which the arrays will be joined.
* @returns {NDArray} A new NDArray.
*/
export function concat(arrays, axis = 0) {
if (!arrays || arrays.length === 0) {
throw new Error('Array list cannot be empty.');
}
if (arrays.length === 1) {
return arrays[0].copy(); // .copy() is available in this file.
}
const firstArr = arrays[0];
const { dtype, ndim } = firstArr;
// --- Validation and Metadata ---
if (ndim === 0) {
if (!arrays.every(arr => arr.ndim === 0)) {
throw new Error('All arrays must be scalars to concatenate with a scalar.');
}
const newData = new (DTYPE_MAP[dtype])(arrays.length);
for (let i = 0; i < arrays.length; i++) {
newData[i] = arrays[i].data[arrays[i].offset];
}
return new NDArray(newData, { shape: [arrays.length], dtype });
}
const finalAxis = axis < 0 ? ndim + axis : axis;
if (finalAxis < 0 || finalAxis >= ndim) {
throw new Error(`Axis ${axis} is out of bounds for array of dimension ${ndim}.`);
}
let newShape = [...firstArr.shape];
let concatDimSize = 0;
for (const arr of arrays) {
if (arr.ndim !== ndim) throw new Error('All arrays must have same number of dimensions.');
if (arr.dtype !== dtype) throw new Error('All arrays must have same dtype.');
for (let j = 0; j < ndim; j++) {
if (j !== finalAxis && arr.shape[j] !== newShape[j]) {
throw new Error(`Dimension mismatch on axis ${j}: expected ${newShape[j]} but got ${arr.shape[j]}.`);
}
}
concatDimSize += arr.shape[finalAxis];
}
newShape[finalAxis] = concatDimSize;
// --- End Validation ---
const result = zeros(newShape, dtype);
let destAxisOffset = 0;
for (const arr of arrays) {
if (arr.size === 0) continue;
const dimSize = arr.shape[finalAxis];
const sliceSpec = Array(ndim).fill(null);
sliceSpec[finalAxis] = [destAxisOffset, destAxisOffset + dimSize];
const view = result.slice(...sliceSpec);
view.set(arr);
destAxisOffset += dimSize;
}
return result;
}
/**
* Joins a sequence of arrays along a new axis.
* The `stack` function creates a new dimension, whereas `concat` joins along an existing one.
* All input arrays must have the same shape and dtype.
*
* @param {Array<NDArray>} arrays - The list of arrays to stack.
* @param {number} [axis=0] - The axis in the result array along which the input arrays are stacked.
* @returns {NDArray} A new NDArray.
*/
export function stack(arrays, axis = 0) {
if (!arrays || arrays.length === 0) {
throw new Error('Array list cannot be empty.');
}
const firstArr = arrays[0];
const { shape: oldShape, ndim: oldNdim, dtype } = firstArr;
// The result of stack will have one more dimension than the input arrays
const targetNdim = oldNdim + 1;
// Normalize negative axis (e.g., -1 refers to the last possible axis)
let finalAxis = axis < 0 ? targetNdim + axis : axis;
if (finalAxis < 0 || finalAxis >= targetNdim) {
throw new Error(`Axis ${axis} is out of bounds for stack into ${targetNdim} dimensions.`);
}
// Step 1: Validate shapes and expand dimensions
const expandedArrays = arrays.map((arr, index) => {
// Ensure dimensionality matches
if (arr.ndim !== oldNdim) {
throw new Error(`Array at index ${index} has ${arr.ndim} dimensions, expected ${oldNdim}.`);
}
// Ensure dtypes match
if (arr.dtype !== dtype) {
throw new Error(`Dtype mismatch at index ${index}. Expected ${dtype}.`);
}
// Ensure shapes match exactly
for (let i = 0; i < oldNdim; i++) {
if (arr.shape[i] !== oldShape[i]) {
throw new Error(
`Shape mismatch at index ${index}: expected [${oldShape}] but got [${arr.shape}].`
);
}
}
/**
* Step 2: Reshape the array to add a singleton dimension at the target axis.
* Example: If shape is [3, 4] and axis is 1, new shape becomes [3, 1, 4].
*/
const newShape = [...arr.shape];
newShape.splice(finalAxis, 0, 1);
// return a new view with the expanded shape
return arr.reshape(...newShape);
});
/**
* Step 3: Concatenate the expanded arrays along the newly created axis.
* Since every array now has a size of 1 at 'finalAxis',
* concat will join them to a size of 'arrays.length' at that axis.
*/
return concat(expandedArrays, finalAxis);
};
/**
* Creates a 2D identity matrix.
* @param {number} n - Number of rows.
* @param {number} [m] - Number of columns. Defaults to n if not provided.
* @param {string} [dtype='float64'] - Data type of the array.
* @returns {NDArray} A 2D NDArray with ones on the main diagonal.
*/
export function eye(n, m, dtype = 'float64') {
const rows = n;
const cols = m === undefined ? n : m;
const shape = [rows, cols];
// 1. Initialize an all-zero NDArray.
// This handles data allocation and dtype mapping via the existing zeros method.
const res = zeros(shape, dtype);
const offset = 0;
// 2. Fill the main diagonal with 1.
// For a contiguous 2D array, the element at (i, i) is at index: i * cols + i.
const minDim = Math.min(rows, cols);
for (let i = 0; i < minDim; i++) {
res.data[offset + i * cols + i] = 1;
}
return res;
}