-- Module:PythonScript/ExtraLibrarys/Numpy
-- 輕量化維基百科 NumPy 核心子集 (支援 ndarray 布林遮罩、Ufuncs、統計聚合、linalg 與 1D/2D FFT)
local NumpyLib = {}
local bit32 = require("bit32")
local function is_inf(x) return x == math.huge or x == -math.huge end
local function is_nan(x) return x ~= x end
function NumpyLib.init()
local Runtime = require("Module:沙盒/a2569875/PythonScript/Runtime")
local PyClass = require("Module:沙盒/a2569875/PythonScript/Class")
local MathCore = require("Module:沙盒/a2569875/PythonScript/Runtime/Cmath")
local Cmath = MathCore.Cmath
local Matrix = MathCore.Matrix
local get_parts = Cmath.get_parts
local None = Runtime.builtins.None
local function raise_err(err_type, msg) Runtime.raise(err_type, msg) end
local function is_none(x) return x == nil or x == None end
-- ==========================================
-- 1. 底層數值與型別工具
-- ==========================================
local function unwrap_real(x)
if type(x) == "number" then return x end
if type(x) == "boolean" then return x and 1 or 0 end
if type(x) == "table" then
if x.__type == "int" then return x.value or tonumber(MathCore.BigInt.to_string(x.sign, x.limbs))
elseif x.__type == "complex" then
if (x.imag and x.imag ~= 0) or (x.k and x.k ~= 0) or (x.jk and x.jk ~= 0) then
raise_err("TypeError", "can't convert complex/quaternion to float")
end
return x.real or 0
elseif x.__type == "instance" and x.__builtin_value ~= nil then
return unwrap_real(x.__builtin_value)
end
end
if type(x) == "string" then return tonumber(x) or raise_err("ValueError", "could not convert string to float") end
raise_err("TypeError", "expected real number")
end
local function is_seq_table(x)
return type(x) == "table" and (x.__type == "list" or x.__type == "tuple" or x.__is_ndarray == true)
end
local function get_elements(x)
if type(x) ~= "table" then return nil end
if x.elements then return x.elements end
if x.__builtin_value and x.__builtin_value.elements then return x.__builtin_value.elements end
return nil
end
-- 遞迴複製並轉化為標準內部 list 結構
local function deep_to_pylist(obj)
if type(obj) == "table" then
if obj.__type == "range" then
local el = {}
if obj.step > 0 then for i = obj.start, obj.stop - 1, obj.step do table.insert(el, i) end
else for i = obj.start, obj.stop + 1, obj.step do table.insert(el, i) end end
return { __type = "list", elements = el, is_matrix = true }
end
local src = get_elements(obj)
if src then
local el = {}
for i = 1, #src do el[i] = deep_to_pylist(src[i]) end
return { __type = "list", elements = el, is_matrix = true }
end
end
return obj
end
-- 計算巢狀 list 的 shape
local function compute_shape(pylist)
local shape = {}
local curr = pylist
while type(curr) == "table" and curr.__type == "list" do
local len = #curr.elements
table.insert(shape, len)
if len == 0 then break end
curr = curr.elements[1]
end
return shape
end
-- 攤平任意維度至 1D Lua 陣列
local function flatten_to_lua(obj, out)
out = out or {}
local el = get_elements(obj)
if el then
for i = 1, #el do flatten_to_lua(el[i], out) end
else
table.insert(out, obj)
end
return out
end
-- 根據 shape 與扁平陣列重建巢狀 pylist
local function build_from_flat(flat, shape, dim_idx, offset)
dim_idx = dim_idx or 1
offset = offset or 1
local n = shape[dim_idx] or 0
if dim_idx == #shape then
local el = {}
for i = 1, n do
el[i] = flat[offset + i - 1]
end
return { __type = "list", elements = el, is_matrix = true }, offset + n
else
local el = {}
local cur_off = offset
for i = 1, n do
el[i], cur_off = build_from_flat(flat, shape, dim_idx + 1, cur_off)
end
return { __type = "list", elements = el, is_matrix = true }, cur_off
end
end
-- 解析 shape 參數 (支援整數 5 或 tuple/list (2, 3))
local function parse_shape_arg(arg)
if type(arg) == "number" or (type(arg) == "table" and arg.__type == "int") then
return { math.floor(unwrap_real(arg)) }
end
local el = get_elements(arg)
if el then
local s = {}
for i = 1, #el do s[i] = math.floor(unwrap_real(el[i])) end
return s
end
raise_err("TypeError", "expected int or sequence of ints for shape")
end
-- ==========================================
-- 2. 定義 ndarray 類別與魔法方法
-- ==========================================
local ndarray_cls
local wrap_ndarray
-- 純量運算分派 (自動支援實數、複數、四元數與布林值)
local function scalar_binary_op(a, b, op)
if op == "==" then
if (type(a) == "table" and a.__type == "complex") or (type(b) == "table" and b.__type == "complex") then
return Cmath.eq(a, b)
end
return a == b
elseif op == "!=" then
if (type(a) == "table" and a.__type == "complex") or (type(b) == "table" and b.__type == "complex") then
return not Cmath.eq(a, b)
end
return a ~= b
elseif op == "<" then return unwrap_real(a) < unwrap_real(b)
elseif op == "<=" then return unwrap_real(a) <= unwrap_real(b)
elseif op == ">" then return unwrap_real(a) > unwrap_real(b)
elseif op == ">=" then return unwrap_real(a) >= unwrap_real(b)
elseif op == "&" then
if type(a) == "boolean" and type(b) == "boolean" then return a and b end
return bit32.band(unwrap_real(a), unwrap_real(b))
elseif op == "|" then
if type(a) == "boolean" and type(b) == "boolean" then return a or b end
return bit32.bor(unwrap_real(a), unwrap_real(b))
elseif op == "^" then
if type(a) == "boolean" and type(b) == "boolean" then return a ~= b end
return bit32.bxor(unwrap_real(a), unwrap_real(b))
end
if (type(a) == "table" and a.__type == "complex") or (type(b) == "table" and b.__type == "complex") then
if op == "+" then return Cmath.add(a, b)
elseif op == "-" then return Cmath.sub(a, b)
elseif op == "*" then return Cmath.mul(a, b)
elseif op == "/" then return Cmath.div(a, b)
elseif op == "**" then return Cmath.pow(a, b) end
end
local na, nb = unwrap_real(a), unwrap_real(b)
if op == "+" then return na + nb
elseif op == "-" then return na - nb
elseif op == "*" then return na * nb
elseif op == "/" then
if nb == 0 then raise_err("ZeroDivisionError", "division by zero") end
return na / nb
elseif op == "//" then
if nb == 0 then raise_err("ZeroDivisionError", "integer division by zero") end
return math.floor(na / nb)
elseif op == "%" then
if nb == 0 then raise_err("ZeroDivisionError", "modulo by zero") end
return na % nb
elseif op == "**" then return na ^ nb
end
raise_err("TypeError", "unsupported operator " .. tostring(op))
end
-- 遞迴廣播二元運算 (支援 純量-陣列、同維陣列、2D與1D列廣播)
local function broadcast_binary_op(A, B, op)
local elA = get_elements(A)
local elB = get_elements(B)
if elA and elB then
local lenA, lenB = #elA, #elB
local subA_is_seq = lenA > 0 and get_elements(elA[1]) ~= nil
local subB_is_seq = lenB > 0 and get_elements(elB[1]) ~= nil
if subA_is_seq and not subB_is_seq then
local res = {}
for i = 1, lenA do res[i] = broadcast_binary_op(elA[i], B, op) end
return { __type = "list", elements = res, is_matrix = true }
elseif not subA_is_seq and subB_is_seq then
local res = {}
for i = 1, lenB do res[i] = broadcast_binary_op(A, elB[i], op) end
return { __type = "list", elements = res, is_matrix = true }
elseif lenA == 1 and lenB > 1 then
local res = {}
for i = 1, lenB do res[i] = broadcast_binary_op(elA[1], elB[i], op) end
return { __type = "list", elements = res, is_matrix = true }
elseif lenB == 1 and lenA > 1 then
local res = {}
for i = 1, lenA do res[i] = broadcast_binary_op(elA[i], elB[1], op) end
return { __type = "list", elements = res, is_matrix = true }
else
if lenA ~= lenB then raise_err("ValueError", "operands could not be broadcast together with different shapes") end
local res = {}
for i = 1, lenA do res[i] = broadcast_binary_op(elA[i], elB[i], op) end
return { __type = "list", elements = res, is_matrix = true }
end
elseif elA then
local res = {}
for i = 1, #elA do res[i] = broadcast_binary_op(elA[i], B, op) end
return { __type = "list", elements = res, is_matrix = true }
elseif elB then
local res = {}
for i = 1, #elB do res[i] = broadcast_binary_op(A, elB[i], op) end
return { __type = "list", elements = res, is_matrix = true }
else
return scalar_binary_op(A, B, op)
end
end
local function apply_unary_op(A, fn)
local el = get_elements(A)
if el then
local res = {}
for i = 1, #el do res[i] = apply_unary_op(el[i], fn) end
return { __type = "list", elements = res, is_matrix = true }
else
return fn(A)
end
end
-- 判斷序列是否為布林遮罩陣列
local function is_boolean_mask(idx)
local el = get_elements(idx)
if not el or #el == 0 then return false end
local first = el[1]
while type(first) == "table" and get_elements(first) do
local sub = get_elements(first)
if #sub == 0 then return false end
first = sub[1]
end
return type(first) == "boolean"
end
-- 解析 Python slice 物件為 0-based (start, stop, step)
local function resolve_py_slice(slice_obj, length)
local d = slice_obj.__dict or {}
local raw_start = not is_none(d.start) and math.floor(unwrap_real(d.start)) or nil
local raw_stop = not is_none(d.stop) and math.floor(unwrap_real(d.stop)) or nil
local raw_step = not is_none(d.step) and math.floor(unwrap_real(d.step)) or 1
if raw_step == 0 then raise_err("ValueError", "slice step cannot be zero") end
local start, stop
if raw_step > 0 then
start = raw_start or 0
if start < 0 then start = math.max(0, length + start) elseif start > length then start = length end
stop = raw_stop or length
if stop < 0 then stop = math.max(0, length + stop) elseif stop > length then stop = length end
else
start = raw_start or (length - 1)
if start < 0 then start = length + start end
if start >= length then start = length - 1 end
if raw_stop == nil then stop = -1
else
stop = raw_stop
if stop < 0 then stop = length + stop end
end
end
return start, stop, raw_step
end
local function format_ndarray_str(obj, indent)
indent = indent or ""
local el = get_elements(obj)
if not el then
if type(obj) == "boolean" then return obj and "True" or "False" end
if type(obj) == "number" then
if math.abs(obj) < 1e-14 then obj = 0 end
if math.floor(obj) == obj and math.abs(obj) < 1e14 then return tostring(obj) end
return string.format("%.6g", obj)
end
return Runtime.builtins.str.__call({ obj })
end
if #el == 0 then return "[]" end
if get_elements(el[1]) then
local rows = {}
for i = 1, #el do
rows[i] = (i == 1 and "" or (indent .. " ")) .. format_ndarray_str(el[i], indent .. " ")
end
return "[" .. table.concat(rows, "\n") .. "]"
else
local items = {}
for i = 1, #el do items[i] = format_ndarray_str(el[i]) end
return "[" .. table.concat(items, " ") .. "]"
end
end
wrap_ndarray = function(data)
if type(data) == "table" and data.__is_ndarray then return data end
local pylist = (type(data) == "table" and data.__type == "list" and data.is_matrix) and data or deep_to_pylist(data)
if type(pylist) ~= "table" or pylist.__type ~= "list" then
return pylist
end
local inst = PyClass.make_instance(ndarray_cls)
inst.__is_ndarray = true
inst.is_matrix = true
inst.__builtin_value = pylist
inst.elements = pylist.elements
return inst
end
-- 同步更新實例內部的 elements 與 __builtin_value
local function sync_ndarray(inst, new_pylist)
inst.__builtin_value = new_pylist
inst.elements = new_pylist.elements
return inst
end
-- ==========================================
-- 3. 統計、變形與 Ufunc 核心實作
-- ==========================================
local np_api = {}
local function reduce_flat_or_axis(arr, axis, reduce_1d_fn)
local pylist = type(arr) == "table" and arr.__builtin_value or deep_to_pylist(arr)
if is_none(axis) then
local flat = flatten_to_lua(pylist)
return reduce_1d_fn(flat)
end
local ax = math.floor(unwrap_real(axis))
local shape = compute_shape(pylist)
if ax < 0 then ax = #shape + ax end
if ax < 0 or ax >= #shape then raise_err("ValueError", "axis out of bounds for array") end
if #shape == 1 then
return reduce_1d_fn(pylist.elements)
elseif #shape == 2 then
local R, C = shape[1], shape[2]
local out = {}
if ax == 0 then
for c = 1, C do
local col = {}
for r = 1, R do col[r] = pylist.elements[r].elements[c] end
out[c] = reduce_1d_fn(col)
end
else
for r = 1, R do
out[r] = reduce_1d_fn(pylist.elements[r].elements)
end
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
end
raise_err("NotImplementedError", "axis reduction only supported for 1D and 2D arrays")
end
np_api.sum = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
local s = 0
for _, v in ipairs(vec) do s = scalar_binary_op(s, v, "+") end
return s
end)
end
np_api.prod = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
local p = 1
for _, v in ipairs(vec) do p = scalar_binary_op(p, v, "*") end
return p
end)
end
np_api.cumsum = function(args, kwargs)
local a = args[1]
local axis = (kwargs and kwargs.axis) or args[2]
local pylist = (type(a) == "table" and a.__builtin_value) or deep_to_pylist(a)
local function cum_1d(vec)
local out = {}
local acc = 0
for i, v in ipairs(vec) do
acc = (i == 1) and v or scalar_binary_op(acc, v, "+")
out[i] = acc
end
return out
end
if is_none(axis) then
local flat = flatten_to_lua(pylist)
return wrap_ndarray({ __type = "list", elements = cum_1d(flat), is_matrix = true })
end
local ax = math.floor(unwrap_real(axis))
local shape = compute_shape(pylist)
if ax < 0 then ax = #shape + ax end
if #shape == 1 then
return wrap_ndarray({ __type = "list", elements = cum_1d(pylist.elements), is_matrix = true })
elseif #shape == 2 then
local R, C = shape[1], shape[2]
if ax == 1 then
local rows = {}
for r = 1, R do
rows[r] = { __type = "list", elements = cum_1d(pylist.elements[r].elements), is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
elseif ax == 0 then
local rows = {}
for r = 1, R do rows[r] = { __type = "list", elements = {}, is_matrix = true } end
for c = 1, C do
local acc = 0
for r = 1, R do
local v = pylist.elements[r].elements[c]
acc = (r == 1) and v or scalar_binary_op(acc, v, "+")
rows[r].elements[c] = acc
end
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
end
raise_err("NotImplementedError", "cumsum axis only supported for 1D and 2D arrays")
end
np_api.cumprod = function(args, kwargs)
local a = args[1]
local flat = flatten_to_lua(a)
local out = {}
local acc = 1
for i, v in ipairs(flat) do
acc = (i == 1) and v or scalar_binary_op(acc, v, "*")
out[i] = acc
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
end
np_api.mean = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
if #vec == 0 then return 0/0 end
local s = 0
for _, v in ipairs(vec) do s = scalar_binary_op(s, v, "+") end
return scalar_binary_op(s, #vec, "/")
end)
end
np_api.var = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
local ddof = (kwargs and kwargs.ddof and unwrap_real(kwargs.ddof)) or (args[3] and unwrap_real(args[3])) or 0
return reduce_flat_or_axis(args[1], axis, function(vec)
local n = #vec
if n - ddof <= 0 then return 0/0 end
local mean = 0
for _, v in ipairs(vec) do mean = scalar_binary_op(mean, v, "+") end
mean = scalar_binary_op(mean, n, "/")
local sq_sum = 0
for _, v in ipairs(vec) do
local diff = scalar_binary_op(v, mean, "-")
local mag = (type(diff) == "table" and diff.__type == "complex") and Cmath.abs(diff) or math.abs(unwrap_real(diff))
sq_sum = sq_sum + mag * mag
end
return sq_sum / (n - ddof)
end)
end
np_api.std = function(args, kwargs)
local v = np_api.var(args, kwargs)
if type(v) == "table" and v.__is_ndarray then
return wrap_ndarray(apply_unary_op(v, math.sqrt))
end
return math.sqrt(v)
end
np_api.median = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
local n = #vec
if n == 0 then return 0/0 end
local copy = {}
for i = 1, n do copy[i] = unwrap_real(vec[i]) end
table.sort(copy)
if n % 2 == 1 then return copy[math.floor(n / 2) + 1]
else return (copy[n / 2] + copy[n / 2 + 1]) * 0.5 end
end)
end
np_api.min = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
if #vec == 0 then raise_err("ValueError", "zero-size array to reduction operation minimum") end
local m = unwrap_real(vec[1])
for i = 2, #vec do
local val = unwrap_real(vec[i])
if val < m then m = val end
end
return m
end)
end
np_api.amin = np_api.min
np_api.max = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
if #vec == 0 then raise_err("ValueError", "zero-size array to reduction operation maximum") end
local m = unwrap_real(vec[1])
for i = 2, #vec do
local val = unwrap_real(vec[i])
if val > m then m = val end
end
return m
end)
end
np_api.amax = np_api.max
np_api.argmin = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
if #vec == 0 then raise_err("ValueError", "attempt to get argmin of an empty sequence") end
local m, idx = unwrap_real(vec[1]), 0
for i = 2, #vec do
local val = unwrap_real(vec[i])
if val < m then m = val; idx = i - 1 end
end
return idx
end)
end
np_api.argmax = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
if #vec == 0 then raise_err("ValueError", "attempt to get argmax of an empty sequence") end
local m, idx = unwrap_real(vec[1]), 0
for i = 2, #vec do
local val = unwrap_real(vec[i])
if val > m then m = val; idx = i - 1 end
end
return idx
end)
end
np_api.all = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
for _, v in ipairs(vec) do
if v == false or v == 0 or v == None then return false end
end
return true
end)
end
np_api.any = function(args, kwargs)
local axis = (kwargs and kwargs.axis) or args[2]
return reduce_flat_or_axis(args[1], axis, function(vec)
for _, v in ipairs(vec) do
if v ~= false and v ~= 0 and v ~= None then return true end
end
return false
end)
end
-- 陣列形狀操作
np_api.reshape = function(args, kwargs)
local arr = args[1]
local new_shape
if #args == 2 and get_elements(args[2]) then
new_shape = parse_shape_arg(args[2])
elseif kwargs and kwargs.newshape then
new_shape = parse_shape_arg(kwargs.newshape)
else
new_shape = {}
for i = 2, #args do table.insert(new_shape, math.floor(unwrap_real(args[i]))) end
end
local flat = flatten_to_lua(arr)
local total = #flat
local neg_idx = nil
local known_prod = 1
for i, d in ipairs(new_shape) do
if d < 0 then
if neg_idx then raise_err("ValueError", "can only specify one unknown dimension") end
neg_idx = i
else
known_prod = known_prod * d
end
end
if neg_idx then
if known_prod == 0 or total % known_prod ~= 0 then
raise_err("ValueError", "cannot reshape array of size " .. total .. " into shape")
end
new_shape[neg_idx] = total / known_prod
known_prod = total
end
if known_prod ~= total then
raise_err("ValueError", "cannot reshape array of size " .. total .. " into shape")
end
local built = build_from_flat(flat, new_shape)
return wrap_ndarray(built)
end
np_api.transpose = function(args)
local a = args[1]
local pylist = (type(a) == "table" and a.__builtin_value) or deep_to_pylist(a)
local shape = compute_shape(pylist)
if #shape <= 1 then return wrap_ndarray(deep_to_pylist(pylist)) end
if #shape == 2 then
local R, C = shape[1], shape[2]
local rows = {}
for c = 1, C do
local new_row = {}
for r = 1, R do new_row[r] = pylist.elements[r].elements[c] end
rows[c] = { __type = "list", elements = new_row, is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
raise_err("NotImplementedError", "transpose currently supports 1D and 2D arrays")
end
np_api.flatten = function(args)
local flat = flatten_to_lua(args[1])
return wrap_ndarray({ __type = "list", elements = flat, is_matrix = true })
end
np_api.ravel = np_api.flatten
-- ==========================================
-- 4. 掛載 ndarray 物件方法與運算子
-- ==========================================
local function make_bin_dunder(op, swapped)
return {
__type = "function",
__call = function(args)
local self, other = args[1], args[2]
local res = swapped and broadcast_binary_op(other, self, op) or broadcast_binary_op(self, other, op)
return wrap_ndarray(res)
end
}
end
-- 處理單一或遞迴布林遮罩取值
local function extract_by_mask(data, mask, out)
out = out or {}
local d_el = get_elements(data)
local m_el = get_elements(mask)
if d_el and m_el then
for i = 1, math.min(#d_el, #m_el) do
if type(m_el[i]) == "boolean" then
if m_el[i] then table.insert(out, deep_to_pylist(d_el[i])) end
else
extract_by_mask(d_el[i], m_el[i], out)
end
end
end
return out
end
-- 處理單一或遞迴布林遮罩原地賦值 (A[x == 5] = 0 或 B[x == 3] = C[y == 4])
local function assign_by_mask(data_el, mask_el, val_is_seq, val_flat, state)
for i = 1, math.min(#data_el, #mask_el) do
local m = mask_el[i]
if type(m) == "boolean" then
if m then
if val_is_seq then
if #val_flat == 0 then raise_err("ValueError", "cannot assign empty sequence to masked array") end
data_el[i] = val_flat[((state.k - 1) % #val_flat) + 1]
state.k = state.k + 1
else
data_el[i] = val_flat
end
end
elseif type(m) == "table" and get_elements(m) and type(data_el[i]) == "table" and get_elements(data_el[i]) then
assign_by_mask(get_elements(data_el[i]), get_elements(m), val_is_seq, val_flat, state)
end
end
end
local ndarray_dict = {
__str__ = { __type = "function", __call = function(args) return format_ndarray_str(args[1]) end },
__repr__ = { __type = "function", __call = function(args) return "array(" .. format_ndarray_str(args[1], " ") .. ")" end },
__len__ = { __type = "function", __call = function(args) return #(args[1].elements or {}) end },
__iter__ = {
__type = "function",
__call = function(args)
local self = args[1]
local els = {}
for i, v in ipairs(self.elements or {}) do
els[i] = (type(v) == "table" and v.__type == "list") and wrap_ndarray(v) or v
end
local list_obj = { __type = "list", elements = els }
return PyClass.get_attr(list_obj, "__iter__").__call({})
end
},
-- 四則與比較運算子 (回傳 ndarray 或布林 ndarray)
__add__ = make_bin_dunder("+", false),
__radd__ = make_bin_dunder("+", true),
__sub__ = make_bin_dunder("-", false),
__rsub__ = make_bin_dunder("-", true),
__mul__ = make_bin_dunder("*", false),
__rmul__ = make_bin_dunder("*", true),
__truediv__ = make_bin_dunder("/", false),
__rtruediv__ = make_bin_dunder("/", true),
__floordiv__ = make_bin_dunder("//", false),
__rfloordiv__= make_bin_dunder("//", true),
__mod__ = make_bin_dunder("%", false),
__rmod__ = make_bin_dunder("%", true),
__pow__ = make_bin_dunder("**", false),
__rpow__ = make_bin_dunder("**", true),
__eq__ = make_bin_dunder("==", false),
__ne__ = make_bin_dunder("!=", false),
__lt__ = make_bin_dunder("<", false),
__le__ = make_bin_dunder("<=", false),
__gt__ = make_bin_dunder(">", false),
__ge__ = make_bin_dunder(">=", false),
__and__ = make_bin_dunder("&", false),
__or__ = make_bin_dunder("|", false),
__xor__ = make_bin_dunder("^", false),
__matmul__ = {
__type = "function",
__call = function(args) return np_api.matmul({ args[1], args[2] }) end
},
__rmatmul__ = {
__type = "function",
__call = function(args) return np_api.matmul({ args[2], args[1] }) end
},
__neg__ = {
__type = "function",
__call = function(args)
return wrap_ndarray(apply_unary_op(args[1], function(x) return scalar_binary_op(0, x, "-") end))
end
},
__pos__ = {
__type = "function",
__call = function(args) return wrap_ndarray(deep_to_pylist(args[1])) end
},
__invert__ = {
__type = "function",
__call = function(args)
return wrap_ndarray(apply_unary_op(args[1], function(x)
if type(x) == "boolean" then return not x end
return -(unwrap_real(x) + 1)
end))
end
},
-- 支援整數、切片、2D Tuple、Fancy Indexing 與布林遮罩的 __getitem__
__getitem__ = {
__type = "function",
__call = function(args)
local self, idx = args[1], args[2]
local el = self.elements
-- 1. 布林遮罩 (A[x == 5])
if is_seq_table(idx) and is_boolean_mask(idx) then
local extracted = extract_by_mask(self, idx)
return wrap_ndarray({ __type = "list", elements = extracted, is_matrix = true })
end
-- 2. 整數陣列 Fancy Indexing (A[[0, 2]])
if type(idx) == "table" and (idx.__type == "list" or idx.__is_ndarray) then
local idx_el = get_elements(idx)
local res = {}
for i = 1, #idx_el do
local k = math.floor(unwrap_real(idx_el[i]))
if k < 0 then k = #el + k end
if k < 0 or k >= #el then raise_err("IndexError", "index out of bounds") end
res[i] = deep_to_pylist(el[k + 1])
end
return wrap_ndarray({ __type = "list", elements = res, is_matrix = true })
end
-- 3. 多維度 Tuple 索引 (A[r, c] 或 A[:, 1])
if type(idx) == "table" and idx.__type == "tuple" then
local r_idx, c_idx = idx.elements[1], idx.elements[2]
local row_part = ndarray_cls.__dict.__getitem__.__call({ self, r_idx })
if c_idx == nil then return row_part end
if type(r_idx) == "table" and r_idx.__class and r_idx.__class.__name == "slice" then
local rows_el = get_elements(row_part)
local out_rows = {}
for i = 1, #rows_el do
local row_nd = wrap_ndarray(rows_el[i])
out_rows[i] = ndarray_cls.__dict.__getitem__.__call({ row_nd, c_idx })
if type(out_rows[i]) == "table" and out_rows[i].__is_ndarray then
out_rows[i] = out_rows[i].__builtin_value
end
end
return wrap_ndarray({ __type = "list", elements = out_rows, is_matrix = true })
else
local row_nd = wrap_ndarray(row_part)
return ndarray_cls.__dict.__getitem__.__call({ row_nd, c_idx })
end
end
-- 4. 切片 Slice (A[1:4])
if type(idx) == "table" and idx.__class and idx.__class.__name == "slice" then
local start, stop, step = resolve_py_slice(idx, #el)
local res = {}
if step > 0 then
for i = start, stop - 1, step do table.insert(res, deep_to_pylist(el[i + 1])) end
else
for i = start, stop + 1, step do table.insert(res, deep_to_pylist(el[i + 1])) end
end
return wrap_ndarray({ __type = "list", elements = res, is_matrix = true })
end
-- 5. 單一整數索引 (A[0])
local k = math.floor(unwrap_real(idx))
if k < 0 then k = #el + k end
if k < 0 or k >= #el then raise_err("IndexError", "index out of bounds for ndarray") end
local item = el[k + 1]
if type(item) == "table" and item.__type == "list" then
return wrap_ndarray(item)
end
return item
end
},
-- 支援布林遮罩賦值 (A[x == 5] = 0, B[x == 3] = C[y == 4])、切片、2D Tuple 與整數賦值的 __setitem__
__setitem__ = {
__type = "function",
__call = function(args)
local self, idx, val = args[1], args[2], args[3]
local el = self.elements
local val_is_seq = is_seq_table(val)
local val_flat = val_is_seq and flatten_to_lua(val) or val
-- 1. 布林遮罩賦值
if is_seq_table(idx) and is_boolean_mask(idx) then
assign_by_mask(el, get_elements(idx), val_is_seq, val_flat, { k = 1 })
return None
end
-- 2. 整數陣列 Fancy Indexing 賦值
if type(idx) == "table" and (idx.__type == "list" or idx.__is_ndarray) then
local idx_el = get_elements(idx)
for i = 1, #idx_el do
local k = math.floor(unwrap_real(idx_el[i]))
if k < 0 then k = #el + k end
if k < 0 or k >= #el then raise_err("IndexError", "index out of bounds") end
el[k + 1] = val_is_seq and val_flat[((i - 1) % #val_flat) + 1] or val_flat
end
return None
end
-- 3. 2D Tuple 賦值 (A[r, c] = val)
if type(idx) == "table" and idx.__type == "tuple" then
local r_idx, c_idx = idx.elements[1], idx.elements[2]
if type(r_idx) == "table" and r_idx.__class and r_idx.__class.__name == "slice" then
local start, stop, step = resolve_py_slice(r_idx, #el)
local count = 1
local val_el = get_elements(val)
local val_is_2d = val_el and #val_el > 0 and get_elements(val_el[1]) ~= nil
local c_is_slice = (type(c_idx) == "table" and c_idx.__class and c_idx.__class.__name == "slice")
local function set_row(r)
local row_nd = wrap_ndarray(el[r + 1])
local sub_val
if val_is_2d then
sub_val = val_el[((count - 1) % #val_el) + 1]
elseif c_is_slice and val_is_seq then
sub_val = val
else
sub_val = val_is_seq and val_flat[((count - 1) % #val_flat) + 1] or val
end
ndarray_cls.__dict.__setitem__.__call({ row_nd, c_idx, sub_val })
count = count + 1
end
if step > 0 then for r = start, stop - 1, step do set_row(r) end
else for r = start, stop + 1, step do set_row(r) end end
else
local r = math.floor(unwrap_real(r_idx))
if r < 0 then r = #el + r end
if r < 0 or r >= #el then raise_err("IndexError", "row index out of bounds") end
local row_nd = wrap_ndarray(el[r + 1])
ndarray_cls.__dict.__setitem__.__call({ row_nd, c_idx, val })
end
return None
end
-- 4. 切片賦值 (A[1:3] = val)
if type(idx) == "table" and idx.__class and idx.__class.__name == "slice" then
local start, stop, step = resolve_py_slice(idx, #el)
local count = 1
local function assign_at(pos)
el[pos + 1] = val_is_seq and val_flat[((count - 1) % #val_flat) + 1] or val_flat
count = count + 1
end
if step > 0 then for i = start, stop - 1, step do assign_at(i) end
else for i = start, stop + 1, step do assign_at(i) end end
return None
end
-- 5. 單一整數賦值
local k = math.floor(unwrap_real(idx))
if k < 0 then k = #el + k end
if k < 0 or k >= #el then raise_err("IndexError", "index out of bounds for ndarray") end
if type(el[k + 1]) == "table" and el[k + 1].__type == "list" and not val_is_seq then
local row_el = el[k + 1].elements
for j = 1, #row_el do row_el[j] = val_flat end
else
el[k + 1] = val_is_seq and deep_to_pylist(val) or val_flat
end
return None
end
},
-- 動態屬性 (.shape, .ndim, .size, .T, .real, .imag)
shape = {
__type = "property",
fget = {
__type = "function",
__call = function(args)
return { __type = "tuple", elements = compute_shape(args[1].__builtin_value) }
end
}
},
ndim = {
__type = "property",
fget = {
__type = "function",
__call = function(args) return #compute_shape(args[1].__builtin_value) end
}
},
size = {
__type = "property",
fget = {
__type = "function",
__call = function(args)
local s = compute_shape(args[1].__builtin_value)
if #s == 0 then return 0 end
local p = 1
for _, d in ipairs(s) do p = p * d end
return p
end
}
},
T = {
__type = "property",
fget = {
__type = "function",
__call = function(args) return np_api.transpose({ args[1] }) end
}
},
real = {
__type = "property",
fget = {
__type = "function",
__call = function(args)
return wrap_ndarray(apply_unary_op(args[1], function(x)
local r = get_parts(x); return r
end))
end
}
},
imag = {
__type = "property",
fget = {
__type = "function",
__call = function(args)
return wrap_ndarray(apply_unary_op(args[1], function(x)
local _, im = get_parts(x); return im
end))
end
}
},
-- 常用實例方法
tolist = {
__type = "function",
__call = function(args)
local function strip_matrix_flag(obj)
if type(obj) == "table" and obj.__type == "list" then
local el = {}
for i = 1, #obj.elements do el[i] = strip_matrix_flag(obj.elements[i]) end
return { __type = "list", elements = el }
end
return obj
end
return strip_matrix_flag(args[1].__builtin_value)
end
},
copy = { __type = "function", __call = function(args) return wrap_ndarray(deep_to_pylist(args[1])) end },
reshape = { __type = "function", __call = function(args, kwargs) return np_api.reshape(args, kwargs) end },
transpose = { __type = "function", __call = function(args) return np_api.transpose(args) end },
flatten = { __type = "function", __call = function(args) return np_api.flatten(args) end },
ravel = { __type = "function", __call = function(args) return np_api.ravel(args) end },
sum = { __type = "function", __call = function(args, kwargs) return np_api.sum(args, kwargs) end },
prod = { __type = "function", __call = function(args, kwargs) return np_api.prod(args, kwargs) end },
cumsum = { __type = "function", __call = function(args, kwargs) return np_api.cumsum(args, kwargs) end },
cumprod = { __type = "function", __call = function(args, kwargs) return np_api.cumprod(args, kwargs) end },
mean = { __type = "function", __call = function(args, kwargs) return np_api.mean(args, kwargs) end },
std = { __type = "function", __call = function(args, kwargs) return np_api.std(args, kwargs) end },
var = { __type = "function", __call = function(args, kwargs) return np_api.var(args, kwargs) end },
min = { __type = "function", __call = function(args, kwargs) return np_api.min(args, kwargs) end },
max = { __type = "function", __call = function(args, kwargs) return np_api.max(args, kwargs) end },
argmin = { __type = "function", __call = function(args, kwargs) return np_api.argmin(args, kwargs) end },
argmax = { __type = "function", __call = function(args, kwargs) return np_api.argmax(args, kwargs) end },
all = { __type = "function", __call = function(args, kwargs) return np_api.all(args, kwargs) end },
any = { __type = "function", __call = function(args, kwargs) return np_api.any(args, kwargs) end },
dot = { __type = "function", __call = function(args) return np_api.dot({ args[1], args[2] }) end }
}
ndarray_cls = PyClass.make_class("ndarray", { Runtime.builtins.object }, ndarray_dict)
np_api.ndarray = ndarray_cls
-- ==========================================
-- 5. 陣列建立與拼接函數 (Array Creation & Manipulation)
-- ==========================================
np_api.array = function(args) return wrap_ndarray(deep_to_pylist(args[1])) end
np_api.asarray = function(args) return (type(args[1]) == "table" and args[1].__is_ndarray) and args[1] or wrap_ndarray(deep_to_pylist(args[1])) end
local function fill_shape(shape, val_fn)
local total = 1
for _, d in ipairs(shape) do total = total * d end
local flat = {}
for i = 1, total do flat[i] = val_fn(i) end
local built = build_from_flat(flat, shape)
return wrap_ndarray(built)
end
np_api.zeros = function(args) return fill_shape(parse_shape_arg(args[1]), function() return 0.0 end) end
np_api.ones = function(args) return fill_shape(parse_shape_arg(args[1]), function() return 1.0 end) end
np_api.empty = np_api.zeros
np_api.full = function(args, kwargs)
local fill_val = (kwargs and kwargs.fill_value) or args[2] or 0
return fill_shape(parse_shape_arg(args[1]), function() return fill_val end)
end
np_api.zeros_like = function(args) return fill_shape(compute_shape(deep_to_pylist(args[1])), function() return 0.0 end) end
np_api.ones_like = function(args) return fill_shape(compute_shape(deep_to_pylist(args[1])), function() return 1.0 end) end
np_api.eye = function(args, kwargs)
local N = math.floor(unwrap_real(args[1]))
local M = (args[2] and not is_none(args[2])) and math.floor(unwrap_real(args[2])) or (kwargs and kwargs.M and math.floor(unwrap_real(kwargs.M))) or N
local k = (args[3] and math.floor(unwrap_real(args[3]))) or (kwargs and kwargs.k and math.floor(unwrap_real(kwargs.k))) or 0
local rows = {}
for r = 1, N do
local row = {}
for c = 1, M do row[c] = (c - r == k) and 1.0 or 0.0 end
rows[r] = { __type = "list", elements = row, is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
np_api.identity = np_api.eye
np_api.diag = function(args, kwargs)
local v = deep_to_pylist(args[1])
local k = (args[2] and math.floor(unwrap_real(args[2]))) or (kwargs and kwargs.k and math.floor(unwrap_real(kwargs.k))) or 0
local shape = compute_shape(v)
if #shape == 1 then
local n = shape[1] + math.abs(k)
local rows = {}
for r = 1, n do
local row = {}
for c = 1, n do row[c] = 0.0 end
rows[r] = { __type = "list", elements = row, is_matrix = true }
end
for i = 1, shape[1] do
local r = (k >= 0) and i or (i - k)
local c = (k >= 0) and (i + k) or i
rows[r].elements[c] = v.elements[i]
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
elseif #shape == 2 then
local R, C = shape[1], shape[2]
local out = {}
for r = 1, R do
local c = r + k
if c >= 1 and c <= C then table.insert(out, v.elements[r].elements[c]) end
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
end
raise_err("ValueError", "Input must be 1- or 2-d")
end
np_api.arange = function(args)
local start_val, stop_val, step
if args[2] == nil then start_val = 0; stop_val = unwrap_real(args[1]); step = 1
else start_val = unwrap_real(args[1]); stop_val = unwrap_real(args[2]); step = args[3] and unwrap_real(args[3]) or 1 end
if step == 0 then raise_err("ValueError", "arange() step argument must not be zero") end
local elements = {}
if step > 0 then
local curr = start_val
while curr < stop_val - 1e-12 * step do table.insert(elements, curr); curr = curr + step end
else
local curr = start_val
while curr > stop_val - 1e-12 * step do table.insert(elements, curr); curr = curr + step end
end
return wrap_ndarray({ __type = "list", elements = elements, is_matrix = true })
end
np_api.linspace = function(args, kwargs)
local start_val = unwrap_real(args[1])
local stop_val = unwrap_real(args[2])
local num = (args[3] and math.floor(unwrap_real(args[3]))) or (kwargs and kwargs.num and math.floor(unwrap_real(kwargs.num))) or 50
local endpoint = true
if kwargs and kwargs.endpoint ~= nil then endpoint = (kwargs.endpoint ~= false and kwargs.endpoint ~= 0) end
if num <= 0 then return wrap_ndarray({ __type = "list", elements = {}, is_matrix = true }) end
if num == 1 then return wrap_ndarray({ __type = "list", elements = { start_val }, is_matrix = true }) end
local div = endpoint and (num - 1) or num
local step = (stop_val - start_val) / div
local elements = {}
for i = 0, num - 1 do table.insert(elements, start_val + i * step) end
if endpoint then elements[#elements] = stop_val end
return wrap_ndarray({ __type = "list", elements = elements, is_matrix = true })
end
np_api.meshgrid = function(args, kwargs)
if #args == 0 then return { __type = "tuple", elements = {} } end
local indexing = (kwargs and kwargs.indexing and tostring(kwargs.indexing)) or "xy"
local x_flat = flatten_to_lua(args[1])
if #args == 1 then
return { __type = "tuple", elements = { wrap_ndarray({ __type = "list", elements = x_flat, is_matrix = true }) } }
end
local y_flat = flatten_to_lua(args[2])
local nx, ny = #x_flat, #y_flat
local rows_X, rows_Y = {}, {}
if indexing == "ij" then
for r = 1, nx do
local rx, ry = {}, {}
for c = 1, ny do
rx[c] = x_flat[r]
ry[c] = y_flat[c]
end
rows_X[r] = { __type = "list", elements = rx, is_matrix = true }
rows_Y[r] = { __type = "list", elements = ry, is_matrix = true }
end
else
for r = 1, ny do
local rx, ry = {}, {}
for c = 1, nx do
rx[c] = x_flat[c]
ry[c] = y_flat[r]
end
rows_X[r] = { __type = "list", elements = rx, is_matrix = true }
rows_Y[r] = { __type = "list", elements = ry, is_matrix = true }
end
end
return {
__type = "tuple",
elements = {
wrap_ndarray({ __type = "list", elements = rows_X, is_matrix = true }),
wrap_ndarray({ __type = "list", elements = rows_Y, is_matrix = true })
}
}
end
np_api.concatenate = function(args, kwargs)
local arrays = get_elements(args[1])
if not arrays or #arrays == 0 then raise_err("ValueError", "need at least one array to concatenate") end
local axis = (kwargs and kwargs.axis and math.floor(unwrap_real(kwargs.axis))) or (args[2] and math.floor(unwrap_real(args[2]))) or 0
local first = deep_to_pylist(arrays[1])
local shape = compute_shape(first)
if axis < 0 then axis = #shape + axis end
if axis == 0 then
local out = {}
for _, a in ipairs(arrays) do
local pa = deep_to_pylist(a)
for _, el in ipairs(pa.elements) do table.insert(out, el) end
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
elseif axis == 1 and #shape == 2 then
local rows = {}
for r = 1, shape[1] do rows[r] = { __type = "list", elements = {}, is_matrix = true } end
for _, a in ipairs(arrays) do
local pa = deep_to_pylist(a)
for r = 1, shape[1] do
for _, v in ipairs(pa.elements[r].elements) do
table.insert(rows[r].elements, v)
end
end
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
raise_err("ValueError", "unsupported axis for concatenate")
end
np_api.vstack = function(args)
local arrays = get_elements(args[1])
local rows = {}
for _, a in ipairs(arrays) do
local pa = deep_to_pylist(a)
local s = compute_shape(pa)
if #s == 1 then
table.insert(rows, pa)
else
for _, r in ipairs(pa.elements) do table.insert(rows, r) end
end
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
np_api.row_stack = np_api.vstack
np_api.hstack = function(args)
local arrays = get_elements(args[1])
local first = deep_to_pylist(arrays[1])
local s = compute_shape(first)
if #s == 1 then return np_api.concatenate({ args[1], 0 })
else return np_api.concatenate({ args[1], 1 }) end
end
np_api.column_stack = function(args)
local arrays = get_elements(args[1])
local first = deep_to_pylist(arrays[1])
local s = compute_shape(first)
if #s == 1 then
local R = s[1]
local rows = {}
for r = 1, R do
local row = {}
for c = 1, #arrays do
local col_el = get_elements(arrays[c])
row[c] = col_el[r]
end
rows[r] = { __type = "list", elements = row, is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
end
return np_api.concatenate({ args[1], 1 })
end
np_api.split = function(args, kwargs)
local ary = deep_to_pylist(args[1])
local sec = args[2]
local axis = (kwargs and kwargs.axis and math.floor(unwrap_real(kwargs.axis))) or (args[3] and math.floor(unwrap_real(args[3]))) or 0
local el = ary.elements
local total = (axis == 0) and #el or #(el[1].elements)
local split_pts = {}
if type(sec) == "number" or (type(sec) == "table" and sec.__type == "int") then
local n_sec = math.floor(unwrap_real(sec))
if total % n_sec ~= 0 then raise_err("ValueError", "array split does not result in an equal division") end
local step = total / n_sec
for i = step, total - step, step do table.insert(split_pts, i) end
else
for _, v in ipairs(get_elements(sec) or {}) do table.insert(split_pts, math.floor(unwrap_real(v))) end
end
table.insert(split_pts, total)
local chunks = {}
local prev = 1
for _, pt in ipairs(split_pts) do
if axis == 0 then
local sub = {}
for i = prev, math.min(pt, #el) do table.insert(sub, deep_to_pylist(el[i])) end
table.insert(chunks, wrap_ndarray({ __type = "list", elements = sub, is_matrix = true }))
else
local sub_rows = {}
for r = 1, #el do
local row_sub = {}
for c = prev, math.min(pt, #el[r].elements) do table.insert(row_sub, el[r].elements[c]) end
sub_rows[r] = { __type = "list", elements = row_sub, is_matrix = true }
end
table.insert(chunks, wrap_ndarray({ __type = "list", elements = sub_rows, is_matrix = true }))
end
prev = pt + 1
end
return { __type = "list", elements = chunks }
end
np_api.array_split = np_api.split
np_api.where = function(args)
local cond, x, y = args[1], args[2], args[3]
if x == nil and y == nil then
local pcond = deep_to_pylist(cond)
local s = compute_shape(pcond)
if #s == 1 then
local idx = {}
for i, v in ipairs(pcond.elements) do
if v ~= false and v ~= 0 and v ~= None then table.insert(idx, i - 1) end
end
return { __type = "tuple", elements = { wrap_ndarray({ __type = "list", elements = idx, is_matrix = true }) } }
elseif #s == 2 then
local r_idx, c_idx = {}, {}
for r = 1, s[1] do
for c = 1, s[2] do
local v = pcond.elements[r].elements[c]
if v ~= false and v ~= 0 and v ~= None then
table.insert(r_idx, r - 1); table.insert(c_idx, c - 1)
end
end
end
return {
__type = "tuple",
elements = {
wrap_ndarray({ __type = "list", elements = r_idx, is_matrix = true }),
wrap_ndarray({ __type = "list", elements = c_idx, is_matrix = true })
}
}
end
end
local function select_rec(c, vx, vy)
local cel = get_elements(c)
if cel then
local xel, yel = get_elements(vx), get_elements(vy)
local res = {}
for i = 1, #cel do
res[i] = select_rec(cel[i], xel and xel[i] or vx, yel and yel[i] or vy)
end
return { __type = "list", elements = res, is_matrix = true }
else
return (c ~= false and c ~= 0 and c ~= None) and vx or vy
end
end
return wrap_ndarray(select_rec(cond, x, y))
end
np_api.clip = function(args)
local a, a_min, a_max = args[1], args[2], args[3]
local lo = not is_none(a_min) and unwrap_real(a_min) or -math.huge
local hi = not is_none(a_max) and unwrap_real(a_max) or math.huge
return wrap_ndarray(apply_unary_op(a, function(v)
local nv = unwrap_real(v)
if nv < lo then return lo elseif nv > hi then return hi else return nv end
end))
end
-- ==========================================
-- 6. 數學與通用元素級函數 (Ufuncs：支援實數、複數與四元數)
-- ==========================================
local cmath_mod = require("Module:沙盒/a2569875/PythonScript/Runtime/Cmath").get_cmath_module()
local cmath_dict = cmath_mod.__dict
local function make_ufunc_from_cmath(func_name, preserve_real)
local cm_fn = cmath_dict[func_name].__call
local function scalar_fn(x)
if preserve_real and type(x) == "number" then
if func_name == "sin" then return math.sin(x)
elseif func_name == "cos" then return math.cos(x)
elseif func_name == "tan" then return math.tan(x)
elseif func_name == "exp" then return math.exp(x)
elseif func_name == "log" and x > 0 then return math.log(x)
elseif func_name == "log10" and x > 0 then return math.log10(x)
elseif func_name == "sqrt" and x >= 0 then return math.sqrt(x)
elseif func_name == "asin" and x >= -1 and x <= 1 then return math.asin(x)
elseif func_name == "acos" and x >= -1 and x <= 1 then return math.acos(x)
elseif func_name == "atan" then return math.atan(x)
end
end
local res = cm_fn({ x })
if type(res) == "table" and res.__type == "complex" and res.imag == 0 and (res.k or 0) == 0 and (res.jk or 0) == 0 and type(x) ~= "table" then
return res.real
end
return res
end
return function(args)
if is_seq_table(args[1]) then
return wrap_ndarray(apply_unary_op(args[1], scalar_fn))
end
return scalar_fn(args[1])
end
end
np_api.pi = math.pi; np_api.e = math.exp(1); np_api.euler_gamma = 0.5772156649015329
np_api.inf = math.huge; np_api.nan = 0/0; np_api.newaxis = None
np_api.sin = make_ufunc_from_cmath("sin", true)
np_api.cos = make_ufunc_from_cmath("cos", true)
np_api.tan = make_ufunc_from_cmath("tan", true)
np_api.arcsin = make_ufunc_from_cmath("asin", true); np_api.asin = np_api.arcsin
np_api.arccos = make_ufunc_from_cmath("acos", true); np_api.acos = np_api.arccos
np_api.arctan = make_ufunc_from_cmath("atan", true); np_api.atan = np_api.arctan
np_api.sinh = make_ufunc_from_cmath("sinh", false)
np_api.cosh = make_ufunc_from_cmath("cosh", false)
np_api.tanh = make_ufunc_from_cmath("tanh", false)
np_api.arcsinh = make_ufunc_from_cmath("asinh", false); np_api.asinh = np_api.arcsinh
np_api.arccosh = make_ufunc_from_cmath("acosh", false); np_api.acosh = np_api.arccosh
np_api.arctanh = make_ufunc_from_cmath("atanh", false); np_api.atanh = np_api.arctanh
np_api.exp = make_ufunc_from_cmath("exp", true)
np_api.log = make_ufunc_from_cmath("log", true)
np_api.log10 = make_ufunc_from_cmath("log10", true)
np_api.sqrt = make_ufunc_from_cmath("sqrt", true)
np_api.gamma = make_ufunc_from_cmath("gamma", false)
np_api.lgamma = make_ufunc_from_cmath("lgamma", false)
np_api.log2 = function(args)
local inv_ln2 = 1 / math.log(2)
local lg = np_api.log(args)
if is_seq_table(lg) then return wrap_ndarray(broadcast_binary_op(lg, inv_ln2, "*")) end
return scalar_binary_op(lg, inv_ln2, "*")
end
np_api.abs = function(args)
local function s_abs(x)
if type(x) == "table" and x.__type == "complex" then return Cmath.abs(x) end
return math.abs(unwrap_real(x))
end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_abs)) end
return s_abs(args[1])
end
np_api.absolute = np_api.abs
np_api.conj = function(args)
local function s_conj(x)
if type(x) == "table" and x.__type == "complex" then return Cmath.conj(x) end
return x
end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_conj)) end
return s_conj(args[1])
end
np_api.conjugate = np_api.conj
np_api.real = function(args)
local function s_re(x) local r = get_parts(x); return r end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_re)) end
return s_re(args[1])
end
np_api.imag = function(args)
local function s_im(x) local _, i = get_parts(x); return i end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_im)) end
return s_im(args[1])
end
np_api.angle = function(args, kwargs)
local deg = kwargs and kwargs.deg
local function s_ang(x)
local a = Cmath.arg(x)
return deg and math.deg(a) or a
end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_ang)) end
return s_ang(args[1])
end
np_api.around = function(args, kwargs)
local dec = (args[2] and math.floor(unwrap_real(args[2]))) or (kwargs and kwargs.decimals and math.floor(unwrap_real(kwargs.decimals))) or 0
local mult = 10 ^ dec
local function s_round(x)
if type(x) == "table" and x.__type == "complex" then
local r, i, k, jk = get_parts(x)
return Cmath.make_complex(math.floor(r * mult + 0.5) / mult, math.floor(i * mult + 0.5) / mult, math.floor(k * mult + 0.5) / mult, math.floor(jk * mult + 0.5) / mult)
end
return math.floor(unwrap_real(x) * mult + 0.5) / mult
end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], s_round)) end
return s_round(args[1])
end
np_api.round = np_api.around
np_api.floor = function(args)
local function fn(x) return math.floor(unwrap_real(x)) end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], fn)) end
return fn(args[1])
end
np_api.ceil = function(args)
local function fn(x) return math.ceil(unwrap_real(x)) end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], fn)) end
return fn(args[1])
end
np_api.sign = function(args)
local function fn(x)
if type(x) == "table" and x.__type == "complex" then return Cmath.sgn(x) end
local n = unwrap_real(x)
return n > 0 and 1 or (n < 0 and -1 or 0)
end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], fn)) end
return fn(args[1])
end
np_api. degrees = function(args)
local function fn(x) return math.deg(unwrap_real(x)) end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], fn)) end
return fn(args[1])
end
np_api.rad2deg = np_api.degrees
np_api.radians = function(args)
local function fn(x) return math.rad(unwrap_real(x)) end
if is_seq_table(args[1]) then return wrap_ndarray(apply_unary_op(args[1], fn)) end
return fn(args[1])
end
np_api.deg2rad = np_api.radians
np_api.isclose = function(args, kwargs)
local rtol = (kwargs and kwargs.rtol) or 1e-05
local atol = (kwargs and kwargs.atol) or 1e-08
local function cmp_close(x, y)
local elx, ely = get_elements(x), get_elements(y)
if elx or ely then
local n = elx and #elx or #ely
local res = {}
for i = 1, n do res[i] = cmp_close(elx and elx[i] or x, ely and ely[i] or y) end
return { __type = "list", elements = res, is_matrix = true }
end
local diff = scalar_binary_op(x, y, "-")
local d_mag = (type(diff) == "table" and diff.__type == "complex") and Cmath.abs(diff) or math.abs(unwrap_real(diff))
local y_mag = (type(y) == "table" and y.__type == "complex") and Cmath.abs(y) or math.abs(unwrap_real(y))
return d_mag <= (atol + rtol * y_mag)
end
local res = cmp_close(args[1], args[2])
return is_seq_table(res) and wrap_ndarray(res) or res
end
np_api.allclose = function(args, kwargs)
local ic = np_api.isclose(args, kwargs)
if is_seq_table(ic) then return np_api.all({ ic }) end
return ic
end
-- ==========================================
-- 7. 線性代數與排序 (Linear Algebra & Sorting)
-- ==========================================
np_api.dot = function(args)
local a, b = args[1], args[2]
local pa = is_seq_table(a) and deep_to_pylist(a) or a
local pb = is_seq_table(b) and deep_to_pylist(b) or b
if not is_seq_table(pa) or not is_seq_table(pb) then
return wrap_ndarray(broadcast_binary_op(pa, pb, "*"))
end
local sa, sb = compute_shape(pa), compute_shape(pb)
if #sa == 1 and #sb == 1 then
if sa[1] ~= sb[1] then raise_err("ValueError", "shapes not aligned for 1D dot") end
local sum = 0
for i = 1, sa[1] do
sum = scalar_binary_op(sum, scalar_binary_op(pa.elements[i], pb.elements[i], "*"), "+")
end
return sum
elseif #sa == 2 and #sb == 1 then
local R, C = sa[1], sa[2]
if C ~= sb[1] then raise_err("ValueError", "shapes not aligned for 2D @ 1D dot") end
local out = {}
for r = 1, R do
local s = 0
for c = 1, C do
s = scalar_binary_op(s, scalar_binary_op(pa.elements[r].elements[c], pb.elements[c], "*"), "+")
end
out[r] = s
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
elseif #sa == 1 and #sb == 2 then
local R, C = sb[1], sb[2]
if sa[1] ~= R then raise_err("ValueError", "shapes not aligned for 1D @ 2D dot") end
local out = {}
for c = 1, C do
local s = 0
for r = 1, R do
s = scalar_binary_op(s, scalar_binary_op(pa.elements[r], pb.elements[r].elements[c], "*"), "+")
end
out[c] = s
end
return wrap_ndarray({ __type = "list", elements = out, is_matrix = true })
else
return wrap_ndarray(Matrix.matmul(pa, pb))
end
end
np_api.matmul = np_api.dot
np_api.cross = function(args)
local res = Cmath.cross(deep_to_pylist(args[1]), deep_to_pylist(args[2]))
return is_seq_table(res) and wrap_ndarray(res) or res
end
np_api.trace = function(args) return Matrix.trace(deep_to_pylist(args[1])) end
np_api.sort = function(args, kwargs)
local pa = deep_to_pylist(args[1])
local s = compute_shape(pa)
if #s == 1 then
local copy = {}
for i = 1, #pa.elements do copy[i] = pa.elements[i] end
table.sort(copy, function(x, y) return unwrap_real(x) < unwrap_real(y) end)
return wrap_ndarray({ __type = "list", elements = copy, is_matrix = true })
elseif #s == 2 then
local axis = (kwargs and kwargs.axis and math.floor(unwrap_real(kwargs.axis))) or (args[2] and math.floor(unwrap_real(args[2]))) or -1
if axis < 0 then axis = 2 + axis end
if axis == 1 then
for r = 1, s[1] do
table.sort(pa.elements[r].elements, function(x, y) return unwrap_real(x) < unwrap_real(y) end)
end
else
for c = 1, s[2] do
local col = {}
for r = 1, s[1] do col[r] = pa.elements[r].elements[c] end
table.sort(col, function(x, y) return unwrap_real(x) < unwrap_real(y) end)
for r = 1, s[1] do pa.elements[r].elements[c] = col[r] end
end
end
return wrap_ndarray(pa)
end
raise_err("NotImplementedError", "sort only supports 1D and 2D arrays")
end
np_api.argsort = function(args)
local flat = flatten_to_lua(args[1])
local idx = {}
for i = 1, #flat do idx[i] = i - 1 end
table.sort(idx, function(a, b) return unwrap_real(flat[a + 1]) < unwrap_real(flat[b + 1]) end)
return wrap_ndarray({ __type = "list", elements = idx, is_matrix = true })
end
np_api.unique = function(args, kwargs)
local flat = flatten_to_lua(args[1])
local return_counts = kwargs and kwargs.return_counts
local copy = {}
for i = 1, #flat do copy[i] = unwrap_real(flat[i]) end
table.sort(copy)
local uniq, counts = {}, {}
for _, v in ipairs(copy) do
if #uniq == 0 or uniq[#uniq] ~= v then
table.insert(uniq, v)
table.insert(counts, 1)
else
counts[#counts] = counts[#counts] + 1
end
end
local u_nd = wrap_ndarray({ __type = "list", elements = uniq, is_matrix = true })
if return_counts then
local c_nd = wrap_ndarray({ __type = "list", elements = counts, is_matrix = true })
return { __type = "tuple", elements = { u_nd, c_nd } }
end
return u_nd
end
-- ==========================================
-- 8. numpy.linalg 子模組
-- ==========================================
local linalg_api = {
inv = function(args) return wrap_ndarray(Matrix.inv(deep_to_pylist(args[1]))) end,
pinv = function(args) return wrap_ndarray(Matrix.pinv(deep_to_pylist(args[1]))) end,
solve = function(args)
local A = deep_to_pylist(args[1])
local b = deep_to_pylist(args[2])
local sb = compute_shape(b)
if #sb == 1 then
local col_b = {}
for i = 1, sb[1] do col_b[i] = { __type = "list", elements = { b.elements[i] }, is_matrix = true } end
local res_2d = Matrix.ldiv(A, { __type = "list", elements = col_b, is_matrix = true })
local flat_x = {}
for i = 1, #res_2d.elements do flat_x[i] = res_2d.elements[i].elements[1] end
return wrap_ndarray({ __type = "list", elements = flat_x, is_matrix = true })
end
return wrap_ndarray(Matrix.ldiv(A, b))
end,
det = function(args) return Matrix.det(deep_to_pylist(args[1])) end,
norm = function(args) return Matrix.frobenius_norm(deep_to_pylist(args[1]), false) end,
matrix_rank = function(args) return Matrix.rank(deep_to_pylist(args[1])) end,
matrix_power = function(args) return wrap_ndarray(Matrix.pow(deep_to_pylist(args[1]), unwrap_real(args[2]))) end,
expm = function(args) return wrap_ndarray(Matrix.expm(deep_to_pylist(args[1]), false)) end,
logm = function(args) return wrap_ndarray(Matrix.logm(deep_to_pylist(args[1]), args[2], false)) end,
sqrtm = function(args) return wrap_ndarray(Matrix.sqrtm(deep_to_pylist(args[1]), false)) end,
svd = function(args)
local tup = Matrix.svd(deep_to_pylist(args[1]))
return {
__type = "tuple",
elements = {
wrap_ndarray(tup.elements[1]),
wrap_ndarray(tup.elements[2]),
wrap_ndarray(tup.elements[3])
}
}
end
}
-- ==========================================
-- 9. numpy.random 子模組 (與 CPython NumPy 位元級一致之 MT19937 + rk_gauss 引擎)
-- ==========================================
local function create_mt19937_state(seed_val)
local st = { mt = {}, mti = 625, has_gauss = false, gauss = 0.0 }
local s = bit32.band(math.floor(tonumber(seed_val) or os.time()), 0xffffffff)
st.mt[0] = s
for i = 1, 623 do
local prev = st.mt[i - 1]
local x = bit32.bxor(prev, bit32.rshift(prev, 30))
local x_hi = bit32.rshift(x, 16)
local x_lo = bit32.band(x, 0xffff)
local prod = bit32.band(1812433253 * x_lo + bit32.lshift(bit32.band(1812433253 * x_hi, 0xffff), 16), 0xffffffff)
st.mt[i] = bit32.band(prod + i, 0xffffffff)
end
st.mti = 624
return st
end
local function rk_random(st)
local mt = st.mt
if st.mti >= 624 then
for kk = 0, 226 do
local y = bit32.bor(bit32.band(mt[kk], 0x80000000), bit32.band(mt[kk + 1], 0x7fffffff))
mt[kk] = bit32.bxor(mt[kk + 397], bit32.rshift(y, 1), (bit32.band(y, 1) ~= 0) and 0x9908b0df or 0)
end
for kk = 227, 622 do
local y = bit32.bor(bit32.band(mt[kk], 0x80000000), bit32.band(mt[kk + 1], 0x7fffffff))
mt[kk] = bit32.bxor(mt[kk - 227], bit32.rshift(y, 1), (bit32.band(y, 1) ~= 0) and 0x9908b0df or 0)
end
local y = bit32.bor(bit32.band(mt[623], 0x80000000), bit32.band(mt[0], 0x7fffffff))
mt[623] = bit32.bxor(mt[396], bit32.rshift(y, 1), (bit32.band(y, 1) ~= 0) and 0x9908b0df or 0)
st.mti = 0
end
local y = mt[st.mti]
st.mti = st.mti + 1
y = bit32.bxor(y, bit32.rshift(y, 11))
y = bit32.bxor(y, bit32.band(bit32.lshift(y, 7), 0x9d2c5680))
y = bit32.bxor(y, bit32.band(bit32.lshift(y, 15), 0xefc60000))
y = bit32.bxor(y, bit32.rshift(y, 18))
return y
end
-- NumPy rk_double: 取 27 bits 與 26 bits 合成 53-bit IEEE 754 雙精度亂數
local function rk_double(st)
local a = bit32.rshift(rk_random(st), 5)
local b = bit32.rshift(rk_random(st), 6)
return (a * 67108864.0 + b) / 9007199254740992.0
end
-- NumPy rk_gauss: Marsaglia Polar Method (快取第一個變量，先回傳第二個變量)
local function rk_gauss(st)
if st.has_gauss then
st.has_gauss = false
return st.gauss
end
local x1, x2, r2
repeat
x1 = 2.0 * rk_double(st) - 1.0
x2 = 2.0 * rk_double(st) - 1.0
r2 = x1 * x1 + x2 * x2
until r2 < 1.0 and r2 ~= 0.0
local f = math.sqrt(-2.0 * math.log(r2) / r2)
st.gauss = f * x1
st.has_gauss = true
return f * x2
end
local default_mt_state = create_mt19937_state(123456789)
local function build_rng_methods(get_st)
return {
seed = function(args)
local st = get_st()
local s_in = (args and args[1] and not is_none(args[1])) and math.floor(unwrap_real(args[1])) or os.time()
local fresh = create_mt19937_state(s_in)
st.mt, st.mti, st.has_gauss, st.gauss = fresh.mt, fresh.mti, false, 0.0
return None
end,
rand = function(args)
local st = get_st()
if #args == 0 then return rk_double(st) end
local shape = (#args == 1 and get_elements(args[1])) and parse_shape_arg(args[1]) or parse_shape_arg({ __type = "list", elements = args })
return fill_shape(shape, function() return rk_double(st) end)
end,
randn = function(args)
local st = get_st()
if #args == 0 then return rk_gauss(st) end
local shape = (#args == 1 and get_elements(args[1])) and parse_shape_arg(args[1]) or parse_shape_arg({ __type = "list", elements = args })
return fill_shape(shape, function() return rk_gauss(st) end)
end,
randint = function(args, kwargs)
local st = get_st()
local low, high, size = args[1], args[2], (kwargs and kwargs.size) or args[3]
if is_none(high) then high = unwrap_real(low); low = 0
else low = unwrap_real(low); high = unwrap_real(high) end
local span = math.max(1, math.floor(high) - math.floor(low))
local function gen_int() return math.floor(low) + math.floor(rk_double(st) * span) end
if is_none(size) then return gen_int() end
return fill_shape(parse_shape_arg(size), gen_int)
end,
uniform = function(args, kwargs)
local st = get_st()
local low = (kwargs and kwargs.low ~= nil and unwrap_real(kwargs.low)) or (args[1] ~= nil and unwrap_real(args[1])) or 0.0
local high = (kwargs and kwargs.high ~= nil and unwrap_real(kwargs.high)) or (args[2] ~= nil and unwrap_real(args[2])) or 1.0
local size = (kwargs and kwargs.size) or args[3]
local function gen_u() return low + (high - low) * rk_double(st) end
if is_none(size) then return gen_u() end
return fill_shape(parse_shape_arg(size), gen_u)
end,
normal = function(args, kwargs)
local st = get_st()
local loc = (kwargs and kwargs.loc ~= nil and unwrap_real(kwargs.loc)) or (args[1] ~= nil and unwrap_real(args[1])) or 0.0
local scale = (kwargs and kwargs.scale ~= nil and unwrap_real(kwargs.scale)) or (args[2] ~= nil and unwrap_real(args[2])) or 1.0
local size = (kwargs and kwargs.size) or args[3]
local function gen_n() return loc + scale * rk_gauss(st) end
if is_none(size) then return gen_n() end
return fill_shape(parse_shape_arg(size), gen_n)
end
}
end
local random_api = build_rng_methods(function() return default_mt_state end)
local RandomState_dict = {}
for k, _ in pairs(random_api) do
RandomState_dict[k] = {
__type = "builtin_method",
__name = k,
__call = function(a, kw)
local self = a[1]
local stripped = {}
for i = 2, #(a or {}) do stripped[i - 1] = a[i] end
local bound_methods = build_rng_methods(function() return self.__dict._mt_state end)
return bound_methods[k](stripped, kw)
end
}
end
local RandomStateClass = PyClass.make_class("RandomState", { Runtime.builtins.object }, RandomState_dict)
random_api.RandomState = function(args)
local s_in = (args and args[1] and not is_none(args[1])) and math.floor(unwrap_real(args[1])) or os.time()
local inst = PyClass.make_instance(RandomStateClass)
inst.__dict._mt_state = create_mt19937_state(s_in)
return inst
end
random_api.default_rng = random_api.RandomState
-- ==========================================
-- 10. numpy.fft 子模組 (1D/2D Cooley-Tukey + Bluestein + 四元數防護)
-- ==========================================
-- 🛡️ 四元數防護網 (Quaternion Guard)
local function check_no_quaternion(obj)
local el = get_elements(obj)
if el then
for i = 1, #el do check_no_quaternion(el[i]) end
elseif type(obj) == "table" and obj.__type == "complex" then
if (obj.k and obj.k ~= 0) or (obj.jk and obj.jk ~= 0) then
raise_err("NotImplementedError", "numpy.fft does not support Quaternions (non-commutative algebra); input must be real or complex")
end
end
end
-- 原地 Radix-2 Cooley-Tukey FFT (要求 N 為 2 的冪次)
local function radix2_fft_inplace(re, im, inverse)
local n = #re
local j = 1
for i = 1, n do
if i < j then
re[i], re[j] = re[j], re[i]
im[i], im[j] = im[j], im[i]
end
local m = math.floor(n / 2)
while m >= 1 and j > m do
j = j - m
m = math.floor(m / 2)
end
j = j + m
end
local len = 2
local sign = inverse and 1 or -1
while len <= n do
local half = math.floor(len / 2)
local ang = sign * 2 * math.pi / len
local wlen_r, wlen_i = math.cos(ang), math.sin(ang)
for i = 1, n, len do
local wr, wi = 1.0, 0.0
for k = 0, half - 1 do
local u_idx = i + k
local v_idx = u_idx + half
local ur, ui = re[u_idx], im[u_idx]
local vr = re[v_idx] * wr - im[v_idx] * wi
local vi = re[v_idx] * wi + im[v_idx] * wr
re[u_idx], im[u_idx] = ur + vr, ui + vi
re[v_idx], im[v_idx] = ur - vr, ui - vi
local next_wr = wr * wlen_r - wi * wlen_i
wi = wr * wlen_i + wi * wlen_r
wr = next_wr
end
end
len = len * 2
end
if inverse then
for i = 1, n do re[i] = re[i] / n; im[i] = im[i] / n end
end
end
-- 支援任意長度 N 的 1D 精確 FFT (2 的冪次走 Radix-2；非 2 的冪次走 Bluestein's Chirp-Z)
local function compute_fft_1d_raw(re_in, im_in, n_target, inverse)
local N = n_target or #re_in
if N <= 0 then return {}, {} end
local re, im = {}, {}
for i = 1, N do
re[i] = re_in[i] or 0.0
im[i] = im_in[i] or 0.0
end
if N == 1 then return re, im end
-- 若 N 為 2 的冪次，直接執行超高速 Radix-2 Cooley-Tukey
if bit32.band(N, N - 1) == 0 then
radix2_fft_inplace(re, im, inverse)
return re, im
end
-- 非 2 的冪次：使用 Bluestein's Chirp-Z 轉換為長度 M >= 2N - 1 的 Radix-2 卷積
local M = 1
while M < 2 * N - 1 do M = M * 2 end
local sign = inverse and 1 or -1
local w_re, w_im = {}, {}
local a_re, a_im = {}, {}
local b_re, b_im = {}, {}
for i = 1, M do a_re[i] = 0.0; a_im[i] = 0.0; b_re[i] = 0.0; b_im[i] = 0.0 end
for k = 0, N - 1 do
local angle = sign * math.pi * ((k * k) % (2 * N)) / N
local cr, ci = math.cos(angle), math.sin(angle)
w_re[k + 1], w_im[k + 1] = cr, ci
local xr, xi = re[k + 1], im[k + 1]
a_re[k + 1] = xr * cr - xi * ci
a_im[k + 1] = xr * ci + xi * cr
-- b 為共軛 Chirp 濾波器
b_re[k + 1], b_im[k + 1] = cr, -ci
if k > 0 then
b_re[M - k + 1], b_im[M - k + 1] = cr, -ci
end
end
radix2_fft_inplace(a_re, a_im, false)
radix2_fft_inplace(b_re, b_im, false)
for i = 1, M do
local ar, ai = a_re[i], a_im[i]
local br, bi = b_re[i], b_im[i]
a_re[i] = ar * br - ai * bi
a_im[i] = ar * bi + ai * br
end
radix2_fft_inplace(a_re, a_im, true)
local out_re, out_im = {}, {}
local scale = inverse and (1.0 / N) or 1.0
for k = 1, N do
local cr, ci = w_re[k], w_im[k]
local ar, ai = a_re[k], a_im[k]
out_re[k] = (ar * cr - ai * ci) * scale
out_im[k] = (ar * ci + ai * cr) * scale
end
return out_re, out_im
end
local function clean_fft_val(r, i)
if math.abs(r) < 1e-13 then r = 0 end
if math.abs(i) < 1e-13 then i = 0 end
return Cmath.make_complex(r, i, 0, 0)
end
local function run_fft_1d_list(vec_elements, n_target, inverse)
local re_in, im_in = {}, {}
for idx, v in ipairs(vec_elements) do
local r, i, k, jk = get_parts(v)
if k ~= 0 or jk ~= 0 then
raise_err("NotImplementedError", "numpy.fft does not support Quaternions (non-commutative algebra); input must be real or complex")
end
re_in[idx], im_in[idx] = r, i
end
local ro, io = compute_fft_1d_raw(re_in, im_in, n_target, inverse)
local out = {}
for idx = 1, #ro do out[idx] = clean_fft_val(ro[idx], io[idx]) end
return { __type = "list", elements = out, is_matrix = true }
end
local function make_fft_1d_fn(inverse)
return function(args, kwargs)
local a = args[1]
check_no_quaternion(a)
local n_target = (args[2] and not is_none(args[2])) and math.floor(unwrap_real(args[2]))
or (kwargs and kwargs.n and not is_none(kwargs.n) and math.floor(unwrap_real(kwargs.n))) or nil
local axis = (args[3] and math.floor(unwrap_real(args[3])))
or (kwargs and kwargs.axis and math.floor(unwrap_real(kwargs.axis))) or -1
local pa = deep_to_pylist(a)
local shape = compute_shape(pa)
if #shape == 1 then
return wrap_ndarray(run_fft_1d_list(pa.elements, n_target, inverse))
elseif #shape == 2 then
if axis < 0 then axis = 2 + axis end
if axis == 1 then
local rows = {}
for r = 1, shape[1] do
rows[r] = run_fft_1d_list(pa.elements[r].elements, n_target, inverse)
end
return wrap_ndarray({ __type = "list", elements = rows, is_matrix = true })
elseif axis == 0 then
local t_pa = np_api.transpose({ pa }).__builtin_value
local cols = {}
for c = 1, shape[2] do
cols[c] = run_fft_1d_list(t_pa.elements[c].elements, n_target, inverse)
end
return np_api.transpose({ { __type = "list", elements = cols, is_matrix = true } })
end
end
raise_err("ValueError", "fft supports 1D and 2D arrays")
end
end
local function make_fft_2d_fn(inverse)
local fn_1d = make_fft_1d_fn(inverse)
return function(args, kwargs)
local a = args[1]
check_no_quaternion(a)
local pa = deep_to_pylist(a)
local shape = compute_shape(pa)
if #shape ~= 2 then raise_err("ValueError", "fft2/ifft2 requires a 2D array") end
local s_arg = (kwargs and kwargs.s) or args[2]
local s_el = not is_none(s_arg) and get_elements(s_arg) or nil
local n_rows = s_el and s_el[1] and math.floor(unwrap_real(s_el[1])) or nil
local n_cols = s_el and s_el[2] and math.floor(unwrap_real(s_el[2])) or nil
-- 先對每一列 (axis=1) 做 1D FFT，再對每一行 (axis=0) 做 1D FFT
local step1 = fn_1d({ pa, n_cols, 1 })
local step2 = fn_1d({ step1, n_rows, 0 })
return step2
end
end
local fft_api = {
fft = make_fft_1d_fn(false),
ifft = make_fft_1d_fn(true),
fft2 = make_fft_2d_fn(false),
ifft2 = make_fft_2d_fn(true),
fftfreq = function(args, kwargs)
local n = math.floor(unwrap_real(args[1]))
local d = (args[2] and unwrap_real(args[2])) or (kwargs and kwargs.d and unwrap_real(kwargs.d)) or 1.0
if n <= 0 then raise_err("ValueError", "Window length n must be a positive integer") end
local val = 1.0 / (n * d)
local N_half = math.floor((n - 1) / 2) + 1
local res = {}
for i = 0, N_half - 1 do table.insert(res, i * val) end
for i = -math.floor(n / 2), -1 do table.insert(res, i * val) end
return wrap_ndarray({ __type = "list", elements = res, is_matrix = true })
end,
rfftfreq = function(args, kwargs)
local n = math.floor(unwrap_real(args[1]))
local d = (args[2] and unwrap_real(args[2])) or (kwargs and kwargs.d and unwrap_real(kwargs.d)) or 1.0
local val = 1.0 / (n * d)
local N_half = math.floor(n / 2) + 1
local res = {}
for i = 0, N_half - 1 do table.insert(res, i * val) end
return wrap_ndarray({ __type = "list", elements = res, is_matrix = true })
end,
fftshift = function(args)
local pa = deep_to_pylist(args[1])
local shape = compute_shape(pa)
local function shift_1d(el)
local n = #el
local mid = math.floor((n + 1) / 2)
local out = {}
for i = mid + 1, n do table.insert(out, el[i]) end
for i = 1, mid do table.insert(out, el[i]) end
return out
end
if #shape == 1 then
return wrap_ndarray({ __type = "list", elements = shift_1d(pa.elements), is_matrix = true })
elseif #shape == 2 then
local shifted_rows = shift_1d(pa.elements)
for r = 1, #shifted_rows do
shifted_rows[r] = { __type = "list", elements = shift_1d(shifted_rows[r].elements), is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = shifted_rows, is_matrix = true })
end
raise_err("NotImplementedError", "fftshift supports 1D and 2D arrays")
end,
ifftshift = function(args)
local pa = deep_to_pylist(args[1])
local shape = compute_shape(pa)
local function ishift_1d(el)
local n = #el
local mid = math.floor(n / 2)
local out = {}
for i = mid + 1, n do table.insert(out, el[i]) end
for i = 1, mid do table.insert(out, el[i]) end
return out
end
if #shape == 1 then
return wrap_ndarray({ __type = "list", elements = ishift_1d(pa.elements), is_matrix = true })
elseif #shape == 2 then
local shifted_rows = ishift_1d(pa.elements)
for r = 1, #shifted_rows do
shifted_rows[r] = { __type = "list", elements = ishift_1d(shifted_rows[r].elements), is_matrix = true }
end
return wrap_ndarray({ __type = "list", elements = shifted_rows, is_matrix = true })
end
raise_err("NotImplementedError", "ifftshift supports 1D and 2D arrays")
end
}
-- ==========================================
-- 11. 組裝並輸出 numpy 模組與子模組
-- ==========================================
local function make_submodule(mod_name, api_tbl)
local sub = PyClass.make_instance(PyClass.make_class("module", {}, {}))
sub.__dict["__name__"] = mod_name
for k, v in pairs(api_tbl) do
if type(v) == "function" then
sub.__dict[k] = { __type = "builtin_method", __call = v }
else
sub.__dict[k] = v
end
end
setmetatable(sub.__dict, {
__index = function(_, key)
if type(key) == "string" and key:sub(1, 2) ~= "__" then
return {
__type = "builtin_method",
__call = function()
raise_err("NotImplementedError", mod_name .. "." .. key .. "() is not implemented in Wiki NumPy subset")
end
}
end
return nil
end
})
return sub
end
local inst = make_submodule("numpy", np_api)
inst.__dict["linalg"] = make_submodule("numpy.linalg", linalg_api)
inst.__dict["random"] = make_submodule("numpy.random", random_api)
inst.__dict["fft"] = make_submodule("numpy.fft", fft_api)
return inst
end
return NumpyLib
