rsnumpy 1.2.1__cp39-cp39-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
rsnumpy/__init__.py ADDED
@@ -0,0 +1,4670 @@
1
+ """rsnumpy - A Rust-powered rsnumpy-compatible array library.
2
+
3
+ rsnumpy provides a rsnumpy-compatible API implemented in Rust for better
4
+ performance and memory efficiency while maintaining full API compatibility.
5
+
6
+ Examples:
7
+ >>> import rsnumpy as np
8
+ >>> a = np.ndarray([1, 2, 3])
9
+ >>> np.sum(a)
10
+ 6
11
+ >>> np.mean(a)
12
+ 2.0
13
+ """
14
+ # ========== 子模块导入和函数挂载 ==========
15
+ # 这些导入放在 ndarray 类和 _ensure/_wrap_result/_scalar 函数定义之后以避免循环导入
16
+ from . import array_methods
17
+ from . import math_functions as _math_functions_module
18
+ from . import statistics as _statistics_module
19
+ from . import array_ops as _array_ops_module
20
+
21
+ from .io import save, load, loadtxt, savetxt, savez, load_npz, genfromtxt
22
+ from .polynomial import Poly, polyval, polyfit, polyder, polyint, polyroots
23
+ from .linalg import linalg_module as _linalg_module
24
+ from .random import random_module as _random_module
25
+ from . import char as _char_module
26
+ from . import matlib as _matlib_module
27
+ # ========== rec / ma 子模块(在顶层完全初始化后导入以避免循环依赖) ==========
28
+ from . import ma
29
+ from . import rec
30
+ from . import lib
31
+ # 子模块
32
+ from . import _core
33
+ from . import fft # noqa: E402 FFT 子模块(numpy 风格 np.fft.*)
34
+ from . import typing # noqa: E402 类型注解子模块(numpy 风格 np.typing.*)
35
+ from . import exceptions
36
+ import datetime as _datetime
37
+ import sys as _sys
38
+ from . import _extra as _extra_module
39
+ # ---------- dtype 提升与浮点判定 ----------
40
+ import rsnumpy.num_core as np
41
+ from rsnumpy.num_core import ndarray_iter as NdArrayIter
42
+ from ._dtypes import (
43
+ DType, dtype, _make_flexible, _make_subarray,
44
+ _build_struct, _scalar_typestr_short,
45
+ )
46
+ _current_module = _sys.modules[__name__]
47
+ _sys.modules['rsnumpy'] = _current_module
48
+ _sys.modules['rsnumpy.__init__'] = _current_module
49
+ _current_module.__name__ = 'rsnumpy'
50
+
51
+ __version__ = "1.2.1"
52
+
53
+ # 捕获内建函数别名:_extra 挂载会向本模块 globals 注入同名的 rsnumpy 函数
54
+ # (all/any/round),会遮蔽内建函数。以下别名保证本文件内部逻辑始终使用内建实现。
55
+ _py_round = round
56
+ _py_all = all
57
+ _py_any = any
58
+ _py_max = max
59
+ _py_min = min
60
+
61
+
62
+ class ArrayFlags:
63
+ """数组内存布局信息,与 rsnumpy 的 np.ndarray.flags 兼容。"""
64
+
65
+ def __init__(self, arr):
66
+ self._arr = arr
67
+
68
+ @property
69
+ def C_CONTIGUOUS(self):
70
+ return True
71
+
72
+ @property
73
+ def F_CONTIGUOUS(self):
74
+ return True
75
+
76
+ @property
77
+ def OWNDATA(self):
78
+ return True
79
+
80
+ @property
81
+ def WRITEABLE(self):
82
+ return True
83
+
84
+ @property
85
+ def ALIGNED(self):
86
+ return True
87
+
88
+ @property
89
+ def WRITEBACKIFCOPY(self):
90
+ return False
91
+
92
+ @property
93
+ def UPDATEIFCOPY(self):
94
+ return False
95
+
96
+ def __repr__(self):
97
+ return self.__str__()
98
+
99
+ def __str__(self):
100
+ lines = [
101
+ f" C_CONTIGUOUS : {self.C_CONTIGUOUS}",
102
+ f" F_CONTIGUOUS : {self.F_CONTIGUOUS}",
103
+ f" OWNDATA : {self.OWNDATA}",
104
+ f" WRITEABLE : {self.WRITEABLE}",
105
+ f" ALIGNED : {self.ALIGNED}",
106
+ f" WRITEBACKIFCOPY : {self.WRITEBACKIFCOPY}",
107
+ f" UPDATEIFCOPY : {self.UPDATEIFCOPY}",
108
+ ]
109
+ return "\n".join(lines)
110
+
111
+
112
+ def _format_float_scalar_str(v):
113
+ """0 维浮点标量的字符串,与 rsnumpy 标量 str 对齐(如 75.0、85850000000.0)。"""
114
+ if v != v:
115
+ return 'nan'
116
+ if v == float('inf'):
117
+ return 'inf'
118
+ if v == float('-inf'):
119
+ return '-inf'
120
+ return str(v)
121
+
122
+
123
+ class ndarray:
124
+ """
125
+ rsnumpy.ndarray - 多维数组对象。
126
+
127
+ 这是一个 Rust 后端支持的数组类型,提供与 rsnumpy 兼容的功能。
128
+
129
+ 参数:
130
+ data (array_like): 数组的初始数据,可以是列表、元组或可迭代对象。
131
+
132
+ 属性:
133
+ shape: 数组维度的元组。
134
+ ndim: 维度数量。
135
+ size: 元素总数。
136
+ dtype: 元素数据类型。
137
+ itemsize: 每个元素的字节大小。
138
+ nbytes: 总字节数。
139
+ T: 数组的转置。
140
+ real: 数组的实部。
141
+ imag: 数组的虚部。
142
+ """
143
+
144
+ def __init__(self, data, _dtype="float64", _fields=None, _raw_data=None):
145
+ if _is_ndarray(data):
146
+ self._array = data._array
147
+ self._dtype = data._dtype
148
+ self._fields = getattr(data, '_fields', None)
149
+ self._raw_data = getattr(data, '_raw_data', None)
150
+ self._complex_data = getattr(data, '_complex_data', None)
151
+ elif hasattr(data, '__class__') and data.__class__.__name__ == 'ndarray':
152
+ self._array = data
153
+ if getattr(data, 'is_complex', False):
154
+ self._dtype = "complex128"
155
+ else:
156
+ self._dtype = _dtype
157
+ self._fields = _fields
158
+ self._raw_data = _raw_data
159
+ else:
160
+ self._raw_data = _raw_data
161
+ if _raw_data is not None:
162
+ # 原始数据存储:用一个虚拟 Rust 数组占位
163
+ self._array = np.zeros((len(data),) if isinstance(data, (list, tuple)) else (1,))
164
+ elif not isinstance(data, (list, tuple)) and hasattr(data, '__array_interface__'):
165
+ # 实现数组接口协议的对象(如 PIL 图像)→ 由 Rust 底层解码缓冲区
166
+ self._array, _dtype = _from_array_interface(data)
167
+ else:
168
+ if isinstance(data, tuple):
169
+ data = list(data)
170
+ if isinstance(data, list) and not _is_rectangular(data):
171
+ # 不规则列表(子列表长度不同),存储为原始数据
172
+ self._raw_data = data
173
+ self._array = np.zeros((len(data),))
174
+ else:
175
+ # 单次遍历同时获取扁平数据与复数/字符串标记
176
+ flat, has_c, has_s = _flatten_check(data)
177
+ if has_s:
178
+ # 字符串数据 → 存储为原始 Python 数据
179
+ self._raw_data = list(data) if isinstance(data, (list, tuple)) else [data]
180
+ self._array = np.zeros((len(self._raw_data),))
181
+ _dtype = "string_"
182
+ # 记录 (kind, width),供 dtype/itemsize 与整数索引取标量使用:
183
+ # Python str 按 rsnumpy 惯例归为 unicode 'U',bytes 归为 'S',宽度取最长元素。
184
+ if flat and _py_all(isinstance(v, (bytes, bytearray)) for v in flat):
185
+ self._str_dtype = ('S', _py_max((len(bytes(v)) for v in flat), default=0))
186
+ else:
187
+ self._str_dtype = ('U', _py_max((len(str(v)) for v in flat), default=0))
188
+ elif has_c:
189
+ # 包含复数 → 交由 Rust 原生复数解析(保留形状与虚部)
190
+ self._array = np.ndarray(data)
191
+ _dtype = "complex128"
192
+ else:
193
+ raw = np.ndarray(data)
194
+ # 实数输入但显式指定复数 dtype → 提升为原生复数(零虚部)
195
+ if _dtype in ('complex128', 'complex64'):
196
+ raw = _maybe_native_complex(raw, _dtype)
197
+ self._array = raw
198
+ self._dtype = _dtype
199
+ self._fields = _fields
200
+
201
+ @staticmethod
202
+ def _wrap(raw_array, _dtype="float64", _fields=None, _raw_data=None):
203
+ """包装原始 Rust ndarray 到 Python 类。"""
204
+ obj = ndarray.__new__(ndarray)
205
+ obj._array = raw_array
206
+ obj._dtype = _dtype
207
+ obj._fields = _fields
208
+ obj._raw_data = _raw_data
209
+ return obj
210
+
211
+ def _wrap_raw(raw_data):
212
+ """包装原始 Python 数据(如字符串列表)到 ndarray。"""
213
+ obj = ndarray.__new__(ndarray)
214
+ obj._array = np.zeros((len(raw_data),))
215
+ obj._dtype = "string_"
216
+ obj._raw_data = list(raw_data)
217
+ return obj
218
+
219
+ def __repr__(self):
220
+ fields = getattr(self, '_fields', None)
221
+ if fields:
222
+ return _format_structured_repr(self)
223
+ if getattr(self, '_dtype', "float64") == "string_":
224
+ raw_data = getattr(self, '_raw_data', None)
225
+ if raw_data is not None:
226
+ val_strs = [repr(x) for x in raw_data]
227
+ return "array([" + " ".join(val_strs) + "])"
228
+ raw = getattr(self, '_raw_data', None)
229
+ if raw is not None:
230
+ return _format_ragged_repr(self)
231
+ cpx = getattr(self, '_complex_data', None)
232
+ if cpx is not None:
233
+ inner = _format_complex_repr_1d(cpx)
234
+ return f"array({inner})"
235
+ if getattr(self._array, 'is_complex', False):
236
+ if self.ndim == 0:
237
+ return _format_complex_scalar(self.tolist())
238
+ inner = _format_complex_nested(self.tolist(), ", ")
239
+ return f"array({inner})"
240
+ dt = getattr(self, '_dtype', "float64")
241
+ if dt == "float64" and self.ndim == 1:
242
+ values = self._array.tolist()
243
+ inner = _format_float_repr_1d(values)
244
+ return f"array({inner})"
245
+ return f"array({self.__str__()})"
246
+
247
+ def __str__(self):
248
+ fields = getattr(self, '_fields', None)
249
+ if fields:
250
+ return _format_structured_str(self)
251
+ if getattr(self, '_dtype', "float64") == "string_":
252
+ raw_data = getattr(self, '_raw_data', None)
253
+ if raw_data is not None:
254
+ sd = getattr(self, '_str_dtype', None)
255
+ if sd is not None and sd[0] == 'V':
256
+ _, width = sd
257
+ parts = []
258
+ for v in raw_data:
259
+ bs = bytes(v) if isinstance(v, (bytes, bytearray)) else str(v).encode('latin-1')
260
+ if width:
261
+ bs = bs[:width].ljust(width, b'\x00')
262
+ parts.append(_void_repr(bs))
263
+ return "[" + " ".join(parts) + "]"
264
+ nested = isinstance(raw_data, list) and len(raw_data) > 0
265
+ if nested and isinstance(raw_data[0], list):
266
+ lines = []
267
+ for row in raw_data:
268
+ inner = " ".join(repr(x) for x in row)
269
+ lines.append("[" + inner + "]")
270
+ return "[" + "\n ".join(lines) + "]"
271
+ val_strs = [repr(x) for x in raw_data]
272
+ return "[" + " ".join(val_strs) + "]"
273
+ return str(raw_data)
274
+ raw = getattr(self, '_raw_data', None)
275
+ if raw is not None:
276
+ return _format_ragged_str(self)
277
+ cpx = getattr(self, '_complex_data', None)
278
+ if cpx is not None:
279
+ inner = _format_complex_repr_1d(cpx)
280
+ return inner
281
+ if getattr(self._array, 'is_complex', False):
282
+ if self.ndim == 0:
283
+ return _format_complex_scalar(self.tolist())
284
+ return _format_complex_nested(self.tolist(), " ")
285
+ if getattr(self, '_dtype', "float64") in ("int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"):
286
+ if getattr(self, '_is_empty', False):
287
+ return np._format_int_str(self._array)
288
+ return np._format_int_val_str(self._array)
289
+ if getattr(self, '_dtype', "float64") == "bool":
290
+ def format_bool_list(data):
291
+ if isinstance(data, list):
292
+ inner = " ".join(format_bool_list(v) for v in data)
293
+ return "[" + inner + "]"
294
+ else:
295
+ return "True" if data else "False"
296
+ raw = self._array.tolist()
297
+ if isinstance(raw, list) and len(raw) > 0 and isinstance(raw[0], list):
298
+ lines = []
299
+ for row in raw:
300
+ inner = " ".join(format_bool_list(v) for v in row)
301
+ lines.append("[" + inner + "]")
302
+ return "[" + "\n ".join(lines) + "]"
303
+ return format_bool_list(raw)
304
+ if self.ndim == 0:
305
+ # 0 维浮点数组的 str 与 rsnumpy 标量一致(保留末尾 .0,如 75.0)。
306
+ return _format_float_scalar_str(float(self._array.tolist()))
307
+ return np._format_float_str(self._array)
308
+
309
+ def __format__(self, fmt):
310
+ if self.ndim == 0:
311
+ val = self.tolist()
312
+ if isinstance(val, complex):
313
+ return format(val, fmt)
314
+ return format(val, fmt)
315
+ return format(str(self), fmt)
316
+
317
+ def __len__(self):
318
+ raw = getattr(self, '_raw_data', None)
319
+ if raw is not None:
320
+ return len(raw)
321
+ if self.ndim == 0:
322
+ return 1
323
+ return len(self._array)
324
+
325
+ def tolist(self):
326
+ """转换为 Python 列表,元素类型与 _dtype 一致。"""
327
+ raw = getattr(self, '_raw_data', None)
328
+ if raw is not None:
329
+ return raw
330
+ cpx = getattr(self, '_complex_data', None)
331
+ if cpx is not None:
332
+ return list(cpx)
333
+ dt = getattr(self, '_dtype', 'float64')
334
+ raw_list = self._array.tolist()
335
+ if dt == 'int64':
336
+ return _convert_nested(raw_list, int)
337
+ if dt == 'bool':
338
+ return _convert_nested(raw_list, bool)
339
+ if dt in ('uint8', 'uint16', 'uint32', 'uint64'):
340
+ return _convert_nested(raw_list, int)
341
+ if dt == 'string_':
342
+ raw = getattr(self, '_raw_data', None)
343
+ if raw is not None:
344
+ return list(raw)
345
+ return raw_list
346
+
347
+ def tobytes(self, order='C'):
348
+ """返回数组原始字节(与 rsnumpy.ndarray.tobytes 兼容)。
349
+
350
+ 字符串数组(S/U/V dtype)按 itemsize 补零对齐;数值数组按 dtype
351
+ 字节布局编码(小端)。
352
+ """
353
+ raw = getattr(self, '_raw_data', None)
354
+ sd = getattr(self, '_str_dtype', None)
355
+ if raw is not None and sd is not None:
356
+ return _string_array_tobytes(_flatten_data(raw), sd)
357
+ return _numeric_tobytes(self, order)
358
+
359
+ def __iter__(self):
360
+ # 一维(及标量)按元素迭代产生 Python 标量;高维按首轴迭代产生子数组,
361
+ # 与 rsnumpy 行为一致(如 a, b = np.random.randn(2, n) 得到两个数组)。
362
+ if self.ndim <= 1:
363
+ return iter(self.tolist())
364
+ return (self[i] for i in range(self.shape[0]))
365
+
366
+ @property
367
+ def __array_interface__(self):
368
+ """NumPy 数组接口协议:让真实 rsnumpy/matplotlib 能直接消费 rsnumpy 数组。
369
+
370
+ 底层数据恒为 f64,按追踪的 _dtype 编码为对应类型字节(在 Rust 层完成)。
371
+ 字符串/复数/不规则数组无法用简单缓冲区表示,抛 AttributeError 交回默认处理。
372
+ """
373
+ if getattr(self, '_raw_data', None) is not None:
374
+ raise AttributeError('__array_interface__')
375
+ if getattr(self, '_complex_data', None) is not None:
376
+ raise AttributeError('__array_interface__')
377
+ if getattr(self._array, 'is_complex', False):
378
+ # 原生复数无法用单一 f64 缓冲表示,交回默认处理。
379
+ raise AttributeError('__array_interface__')
380
+ typestr = _DTYPE_TO_TYPESTR.get(getattr(self, '_dtype', 'float64'))
381
+ if typestr is None:
382
+ raise AttributeError('__array_interface__')
383
+ return np.array_interface(self._array, typestr)
384
+
385
+ def __buffer__(self, flags):
386
+ """PEP 3118 缓冲协议(PEP 688 / Python 3.12+):转发到底层 Rust `_array`,
387
+ 让 memoryview(a) 与下游(如 rsplotlib 的 PyBuffer::get)零拷贝读取底层连续
388
+ f64 内存,免去 __array_interface__ 的 bytes 副本开销。
389
+
390
+ 仅对 float64 暴露缓冲:底层存储恒为 f64,其字节布局与 float64 dtype 完全一致,
391
+ 缓冲 format 'd' 语义正确。int/bool 等 dtype 的底层仍是 f64,若也暴露缓冲,
392
+ rsnumpy 会优先按缓冲(f8)而非 __array_interface__(如 <i8)解读,导致 dtype 失真;
393
+ 故这些 dtype 不暴露缓冲,消费方自动回退到 dtype 精确的 __array_interface__ bytes。
394
+ 绘图坐标以 float64 为主,快路径覆盖热点场景。
395
+
396
+ 注:__buffer__ 由 CPython 在 3.12+ 才识别;更早版本上 memoryview(a) 会失败,
397
+ 消费方自动回退到 __array_interface__ bytes,行为安全。
398
+ """
399
+ if getattr(self, '_dtype', 'float64') != 'float64':
400
+ raise BufferError("only float64 arrays expose a zero-copy buffer")
401
+ if getattr(self, '_raw_data', None) is not None:
402
+ raise BufferError("string/ragged array does not support the buffer protocol")
403
+ if getattr(self, '_complex_data', None) is not None:
404
+ raise BufferError("complex array does not support the buffer protocol")
405
+ return memoryview(self._array)
406
+
407
+ def __release_buffer__(self, view):
408
+ view.release()
409
+
410
+ def __bool__(self):
411
+ if self.size == 1:
412
+ return bool(self.item())
413
+ raise ValueError("The truth value of an array with more than one element is ambiguous.")
414
+
415
+ def __invert__(self):
416
+ """逐元素取反(~ 运算符)。"""
417
+ return bitwise_not(self)
418
+
419
+ def __getitem__(self, key):
420
+ # 结构化数组索引:字段 / 多字段 / 记录标量 / 子结构切片
421
+ if getattr(self, '_dtype_obj', None) is not None:
422
+ handled, res = _struct_getitem(self, key)
423
+ if handled:
424
+ return res
425
+ # 字符串/字节/void(S/U/V)dtype 数组的整数索引:返回标量。
426
+ # S→去尾零 bytes,U→去尾零 str,V→void。
427
+ sd = getattr(self, '_str_dtype', None)
428
+ raw = getattr(self, '_raw_data', None)
429
+ if sd is not None and raw is not None and isinstance(key, int) and not isinstance(key, bool):
430
+ flat = _flatten_data(raw)
431
+ return _str_scalar(flat[key], sd)
432
+ # 复数数组的布尔索引
433
+ cpx = getattr(self, '_complex_data', None)
434
+ if _is_ndarray(key) and cpx is not None:
435
+ mask = key.tolist()
436
+ filtered = [v for v, m in zip(cpx, mask) if m > 0.5]
437
+ return ndarray(filtered)
438
+ # 布尔掩码索引:a[bool_mask]。掩码形状需与 self 的前若干维一致,
439
+ # 返回被掩码选中的元素(按 C 序展开),剩余维度保留。
440
+ if (_is_ndarray(key) and getattr(key, '_dtype', None) == 'bool' and key.ndim >= 1):
441
+ mask_shape = tuple(key.shape)
442
+ self_shape = tuple(self.shape)
443
+ k = len(mask_shape)
444
+ if mask_shape == self_shape[:k]:
445
+ rest = self_shape[k:]
446
+ block = 1
447
+ for s in rest:
448
+ block *= s
449
+ raw = np.masked_select(self._array, key._array, block)
450
+ result = ndarray._wrap(raw, _dtype=self._dtype)
451
+ if rest:
452
+ return result.reshape((len(result) // block,) + rest)
453
+ return result
454
+ # np.newaxis (None) 支持:a[np.newaxis, :] / a[:, np.newaxis]
455
+ # 先用去掉 None 的键做常规索引,再在结果的相应输出位置插入 size-1 维度。
456
+ _key_seq = key if isinstance(key, tuple) else (key,)
457
+ if _py_any(k is None for k in _key_seq):
458
+ newaxis_positions = []
459
+ base_key = []
460
+ out_pos = 0
461
+ for k in _key_seq:
462
+ if k is None:
463
+ newaxis_positions.append(out_pos)
464
+ out_pos += 1
465
+ elif isinstance(k, bool):
466
+ base_key.append(k)
467
+ out_pos += 1
468
+ elif isinstance(k, int):
469
+ base_key.append(k) # 整数索引消耗输入轴但不产生输出维度
470
+ else:
471
+ base_key.append(k)
472
+ out_pos += 1
473
+ base = self[tuple(base_key)]
474
+ new_shape = list(base.shape)
475
+ for p in newaxis_positions:
476
+ new_shape.insert(p, 1)
477
+ return base.reshape(tuple(new_shape))
478
+ # 展开省略号(...),补充完整切片以匹配数组维度
479
+ if isinstance(key, tuple):
480
+ new_key = []
481
+ ellipsis_count = np.isum([1 for k in key if k is Ellipsis])
482
+ if ellipsis_count > 0:
483
+ non_ellipsis = [k for k in key if k is not Ellipsis]
484
+ fill = self.ndim - len(non_ellipsis)
485
+ for k in key:
486
+ if k is Ellipsis:
487
+ for _ in range(fill):
488
+ new_key.append(slice(None, None, None))
489
+ else:
490
+ new_key.append(k)
491
+ key = tuple(new_key)
492
+ # 将 Python ndarray 索引展平为 Python list,确保 Rust 端可正确解析
493
+ if isinstance(key, tuple):
494
+ key = tuple(
495
+ _ndarray_to_index_list(k)
496
+ if hasattr(k, '_array') else k
497
+ for k in key
498
+ )
499
+ else:
500
+ key = (_ndarray_to_index_list(key)
501
+ if hasattr(key, '_array') else key,)
502
+ # 委托给 Rust getitem_multi 处理所有索引逻辑
503
+ result = np.getitem_multi(self._array, key, list(self.shape))
504
+ return _wrap_result(result, self._dtype)
505
+
506
+ def __setitem__(self, key, value):
507
+ # 结构化数组赋值:字段 / 多字段 / 记录 / 切片
508
+ if getattr(self, '_dtype_obj', None) is not None:
509
+ if _struct_setitem(self, key, value):
510
+ return
511
+ # 布尔掩码空选择:numpy 语义下为 no-op,避免 Rust 端空索引 panic
512
+ if _empty_bool_mask_key(key):
513
+ return
514
+ # 展开省略号(...),补充完整切片以匹配数组维度
515
+ if key is Ellipsis:
516
+ key = tuple(slice(None, None, None) for _ in range(self.ndim))
517
+ elif isinstance(key, tuple):
518
+ new_key = []
519
+ ellipsis_count = np.isum([1 for k in key if k is Ellipsis])
520
+ if ellipsis_count > 0:
521
+ non_ellipsis = [k for k in key if k is not Ellipsis]
522
+ fill = self.ndim - len(non_ellipsis)
523
+ for k in key:
524
+ if k is Ellipsis:
525
+ for _ in range(fill):
526
+ new_key.append(slice(None, None, None))
527
+ else:
528
+ new_key.append(k)
529
+ key = tuple(new_key)
530
+ if isinstance(key, tuple):
531
+ # 将 Python ndarray 索引展平为 list(布尔类型保留给 Rust 处理)
532
+ key = tuple(
533
+ _ndarray_to_index_list(k)
534
+ if hasattr(k, '_array') and k.dtype != 'bool' else k
535
+ for k in key
536
+ )
537
+ elif not isinstance(key, str):
538
+ key = (_ndarray_to_index_list(key)
539
+ if hasattr(key, '_array') and key.dtype != 'bool' else key,)
540
+ if isinstance(key, tuple):
541
+ # 计算目标形状:将切片/整数索引应用到 self.shape 得到赋值目标的形状
542
+ target_shape = []
543
+ key_idx = 0
544
+ for dim_size in self.shape:
545
+ if key_idx < len(key):
546
+ k = key[key_idx]
547
+ key_idx += 1
548
+ if isinstance(k, slice):
549
+ start = k.start if k.start is not None else 0
550
+ stop = k.stop if k.stop is not None else dim_size
551
+ step = k.step if k.step is not None else 1
552
+ if hasattr(start, '__index__'):
553
+ start = start.__index__()
554
+ if hasattr(stop, '__index__'):
555
+ stop = stop.__index__()
556
+ target_shape.append(_py_max(0, (stop - start + step - 1) // step))
557
+ elif isinstance(k, int):
558
+ continue
559
+ elif isinstance(k, list):
560
+ target_shape.append(len(k))
561
+ elif _is_ndarray(k):
562
+ if k.dtype == 'bool':
563
+ flat_k = k.flatten()
564
+ true_count = 0
565
+ for x in flat_k:
566
+ if x:
567
+ true_count += 1
568
+ if k.shape == self.shape:
569
+ target_shape = [true_count]
570
+ break
571
+ else:
572
+ target_shape.append(true_count)
573
+ else:
574
+ target_shape.append(dim_size)
575
+ else:
576
+ target_shape.append(dim_size)
577
+ else:
578
+ target_shape.append(dim_size)
579
+ # 广播支持:若 value 是数组且形状可广播到目标形状,则先广播
580
+ if _is_ndarray(value):
581
+ try:
582
+ bcast_shape = _broadcast_shape(tuple(target_shape), value.shape)
583
+ if bcast_shape != value.shape:
584
+ value = broadcast_to(value, bcast_shape)
585
+ except ValueError:
586
+ pass
587
+ val = _cast_setitem_value(_setitem_value(value), self._dtype)
588
+ np.setitem_multi(self._array, key, list(self.shape), val)
589
+ else:
590
+ self._array[key] = value
591
+
592
+ def __add__(self, other):
593
+ dt = _promote_dtype(self._dtype, other)
594
+ if self.size == 0:
595
+ return self.copy()
596
+ if _is_ndarray(other):
597
+ return _wrap_result(self._array + other._array, dt)
598
+ if isinstance(other, (tuple, list)):
599
+ other_arr = ndarray(other)
600
+ return _wrap_result(self._array + other_arr._array, dt)
601
+ return _wrap_result(self._array + other, dt)
602
+
603
+ def __radd__(self, other):
604
+ dt = _promote_dtype(self._dtype, other)
605
+ if self.size == 0:
606
+ return self.copy()
607
+ return _wrap_result(other + self._array, dt)
608
+
609
+ def __sub__(self, other):
610
+ dt = _promote_dtype(self._dtype, other)
611
+ if self.size == 0:
612
+ return self.copy()
613
+ if _is_ndarray(other):
614
+ return _wrap_result(self._array - other._array, dt)
615
+ return _wrap_result(self._array - other, dt)
616
+
617
+ def __rsub__(self, other):
618
+ dt = _promote_dtype(self._dtype, other)
619
+ if self.size == 0:
620
+ return self.copy()
621
+ return _wrap_result(other - self._array, dt)
622
+
623
+ def __mul__(self, other):
624
+ dt = _promote_dtype(self._dtype, other)
625
+ if self.size == 0:
626
+ return self.copy()
627
+ if _is_ndarray(other):
628
+ return _wrap_result(self._array * other._array, dt)
629
+ return _wrap_result(self._array * other, dt)
630
+
631
+ def __rmul__(self, other):
632
+ dt = _promote_dtype(self._dtype, other)
633
+ if self.size == 0:
634
+ return self.copy()
635
+ return _wrap_result(other * self._array, dt)
636
+
637
+ def __truediv__(self, other):
638
+ dt = _truediv_dtype(self._dtype)
639
+ if self.size == 0:
640
+ return self.copy()
641
+ if _is_ndarray(other):
642
+ return _wrap_result(self._array / other._array, dt)
643
+ return _wrap_result(self._array / other, dt)
644
+
645
+ def __rtruediv__(self, other):
646
+ dt = _truediv_dtype(self._dtype)
647
+ if self.size == 0:
648
+ return self.copy()
649
+ return _wrap_result(other / self._array, dt)
650
+
651
+ def __mod__(self, other):
652
+ """逐元素取模(% 运算符)。"""
653
+ return mod(self, other)
654
+
655
+ def __rmod__(self, other):
656
+ """逐元素取模(右操作数,% 运算符)。"""
657
+ return mod(other, self)
658
+
659
+ def __matmul__(self, other):
660
+ if _is_ndarray(other):
661
+ try:
662
+ return _wrap_result(np.linalg.matmul(self._array, other._array), self._dtype)
663
+ except ValueError:
664
+ return _matmul_fallback(self, other)
665
+ return _wrap_result(np.linalg.matmul(self._array, other), self._dtype)
666
+
667
+ def __pow__(self, other):
668
+ if _is_ndarray(other):
669
+ return _wrap_result(np.power(self._array, other._array), self._dtype)
670
+ return _wrap_result(np.power(self._array, np.ndarray([other])), self._dtype)
671
+
672
+ def __rpow__(self, other):
673
+ if _is_ndarray(other):
674
+ return _wrap_result(np.power(other._array, self._array), self._dtype)
675
+ return _wrap_result(np.power(np.ndarray([other]), self._array), self._dtype)
676
+
677
+ def __neg__(self):
678
+ """逐元素取负(- 运算符)。"""
679
+ return _wrap_result(self._array * -1, self._dtype)
680
+
681
+ def __pos__(self):
682
+ """一元正号(+ 运算符),返回副本。"""
683
+ return _wrap_result(self._array.copy(), self._dtype)
684
+
685
+ def __abs__(self):
686
+ """逐元素绝对值(内建 abs() 函数)。"""
687
+ return _wrap_result(np.abs(self._array), self._dtype)
688
+
689
+ def __eq__(self, other):
690
+ if getattr(self, '_dtype_obj', None) is not None and \
691
+ _is_ndarray(other) and getattr(other, '_dtype_obj', None) is not None:
692
+ return _struct_eq(self, other)
693
+ if isinstance(other, str):
694
+ flat_list = self.flatten().tolist()
695
+ flat_result = [x == other for x in flat_list]
696
+ result = ndarray(flat_result)
697
+ result._dtype = "bool"
698
+ return result.reshape(self.shape)
699
+ if _is_ndarray(other):
700
+ return _wrap_result(self._array.__eq__(other._array), "bool")
701
+ return _wrap_result(self._array.__eq__(other), "bool")
702
+
703
+ def __ne__(self, other):
704
+ if _is_ndarray(other):
705
+ return _wrap_result(self._array.__ne__(other._array), "bool")
706
+ return _wrap_result(self._array.__ne__(other), "bool")
707
+
708
+ def __lt__(self, other):
709
+ if _is_ndarray(other):
710
+ return _wrap_result(self._array.__lt__(other._array), "bool")
711
+ return _wrap_result(self._array.__lt__(other), "bool")
712
+
713
+ def __le__(self, other):
714
+ if _is_ndarray(other):
715
+ return _wrap_result(self._array.__le__(other._array), "bool")
716
+ return _wrap_result(self._array.__le__(other), "bool")
717
+
718
+ def __gt__(self, other):
719
+ if _is_ndarray(other):
720
+ return _wrap_result(self._array.__gt__(other._array), "bool")
721
+ return _wrap_result(self._array.__gt__(other), "bool")
722
+
723
+ def __ge__(self, other):
724
+ if _is_ndarray(other):
725
+ return _wrap_result(self._array.__ge__(other._array), "bool")
726
+ return _wrap_result(self._array.__ge__(other), "bool")
727
+
728
+ def __round__(self, ndigits=None):
729
+ if ndigits is None:
730
+ return _wrap_result(np.round(self._array, 0), self._dtype)
731
+ return _wrap_result(np.round(self._array, ndigits), self._dtype)
732
+
733
+ def __and__(self, other):
734
+ """按位与(& 运算符)。"""
735
+ return bitwise_and(self, other)
736
+
737
+ def __or__(self, other):
738
+ """按位或(| 运算符)。"""
739
+ return bitwise_or(self, other)
740
+
741
+ def __xor__(self, other):
742
+ """按位异或(^ 运算符)。"""
743
+ return bitwise_xor(self, other)
744
+
745
+ def __lshift__(self, other):
746
+ """左移(<< 运算符)。"""
747
+ return left_shift(self, other)
748
+
749
+ def __rshift__(self, other):
750
+ """右移(>> 运算符)。"""
751
+ return right_shift(self, other)
752
+
753
+ # ========== 属性 ==========
754
+
755
+ @property
756
+ def shape(self):
757
+ """返回数组维度的元组。"""
758
+ raw = getattr(self, '_raw_data', None)
759
+ arr = getattr(self, '_array', None)
760
+ if raw is not None and arr is not None:
761
+ rust_shape = self._array.shape
762
+ return tuple(int(s) if hasattr(s, '__index__') else s for s in rust_shape)
763
+ if raw is not None:
764
+ # 从嵌套列表递归计算形状
765
+ shape = []
766
+ obj = raw
767
+ while isinstance(obj, list):
768
+ shape.append(len(obj))
769
+ if obj:
770
+ obj = obj[0]
771
+ else:
772
+ break
773
+ return tuple(shape)
774
+ return self._array.shape
775
+
776
+ @shape.setter
777
+ def shape(self, new_shape):
778
+ """设置数组形状,相当于原地 reshape。"""
779
+ raw = getattr(self, '_raw_data', None)
780
+ if raw is not None:
781
+ total = len(raw)
782
+ new_flat = list(raw)
783
+ if isinstance(new_shape, int):
784
+ new_shape = (new_shape,)
785
+ new_size = 1
786
+ for s in new_shape:
787
+ new_size *= s
788
+ if new_size != total:
789
+ raise ValueError(f"cannot reshape array of size {total} into shape {new_shape}")
790
+ # 扁平化 raw_data 并按新形状展平
791
+ if len(new_shape) == 1:
792
+ self._raw_data = new_flat
793
+ else:
794
+ self._raw_data = new_flat
795
+ else:
796
+ total = self._array.size
797
+ if isinstance(new_shape, int):
798
+ new_shape = (new_shape,)
799
+ new_size = 1
800
+ for s in new_shape:
801
+ new_size *= s
802
+ if new_size != total:
803
+ raise ValueError(f"cannot reshape array of size {total} into shape {new_shape}")
804
+ self._array = self._array.reshape(new_shape)
805
+
806
+ @property
807
+ def flags(self):
808
+ """返回数组的内存布局信息,与 rsnumpy 的 ndarray.flags 兼容。"""
809
+ return ArrayFlags(self)
810
+
811
+ @property
812
+ def base(self):
813
+ """若数组是另一数组的视图则返回其基数组,否则返回 None。"""
814
+ return getattr(self, '_base', None)
815
+
816
+ @property
817
+ def ndim(self):
818
+ """返回数组维度数量。"""
819
+ raw = getattr(self, '_raw_data', None)
820
+ if raw is not None:
821
+ return len(self.shape)
822
+ return self._array.ndim
823
+
824
+ @property
825
+ def size(self):
826
+ """返回元素总数。"""
827
+ raw = getattr(self, '_raw_data', None)
828
+ if raw is not None:
829
+ s = self.shape
830
+ n = 1
831
+ for v in s:
832
+ n *= v
833
+ return n
834
+ return self._array.size
835
+
836
+ @property
837
+ def dtype(self):
838
+ """返回元素数据类型。"""
839
+ dobj = getattr(self, '_dtype_obj', None)
840
+ if dobj is not None:
841
+ return dobj
842
+ sd = getattr(self, '_str_dtype', None)
843
+ if sd is not None:
844
+ kind, width = sd
845
+ return _make_flexible(kind, width, None)
846
+ return DType(getattr(self, '_dtype', "float64"))
847
+
848
+ @property
849
+ def itemsize(self):
850
+ """每个元素的字节大小。"""
851
+ try:
852
+ return self.dtype.itemsize
853
+ except Exception:
854
+ return 8
855
+
856
+ @property
857
+ def strides(self):
858
+ """C 连续布局下各维度的字节步长。"""
859
+ itemsize = self.itemsize
860
+ shape = self.shape
861
+ strides = []
862
+ acc = itemsize
863
+ for s in reversed(shape):
864
+ strides.append(acc)
865
+ acc *= s
866
+ return tuple(reversed(strides))
867
+
868
+ @property
869
+ def nbytes(self):
870
+ """总字节数。"""
871
+ return self.size * self.itemsize
872
+
873
+ @property
874
+ def T(self):
875
+ """数组的转置。"""
876
+ return _ndarray_methods().transpose(self)
877
+
878
+ @property
879
+ def flat(self):
880
+ """返回数组元素的扁平迭代器,逐个访问每个元素。"""
881
+ return _NdFlatIter(self)
882
+
883
+ @flat.setter
884
+ def flat(self, values):
885
+ """设置数组元素的扁平值。"""
886
+ size = self.size
887
+ flat_values = list(values)
888
+ indices = np.ndarray(list(range(size)))
889
+ vals = np.ndarray([float(v) for v in flat_values])
890
+ self._array.put(indices, vals)
891
+
892
+ @property
893
+ def real(self):
894
+ """数组的实部。"""
895
+ result = self._array.real
896
+ return _wrap_result(result, "float64")
897
+
898
+ @property
899
+ def imag(self):
900
+ """数组的虚部(实数数组为全零)。"""
901
+ return _wrap_result(self._array.imag, "float64")
902
+
903
+ def conj(self):
904
+ """返回数组的复共轭。"""
905
+ if self._dtype == 'complex128':
906
+ return _wrap_result(self._array.conj(), "complex128")
907
+ return self.copy()
908
+
909
+ # ========== 对象方法 ==========
910
+
911
+ def astype(self, dtype):
912
+ """转换数组为指定类型。"""
913
+ return _ndarray_methods().astype(self, dtype)
914
+
915
+ def reshape(self, *shape):
916
+ """改变数组形状。"""
917
+ return _ndarray_methods().reshape(self, *shape)
918
+
919
+ def resize(self, new_shape):
920
+ """改变数组大小。"""
921
+ _ndarray_methods().resize(self, new_shape)
922
+
923
+ def ravel(self, order='C'):
924
+ """展平数组。"""
925
+ return _ndarray_methods().ravel(self, order)
926
+
927
+ def flatten(self, order='C'):
928
+ """返回数组的一维副本。"""
929
+ return _ndarray_methods().flatten(self, order)
930
+
931
+ def copy(self, order='K'):
932
+ """返回数组的副本。"""
933
+ return _wrap_result(self._array.copy(), self._dtype)
934
+
935
+ def transpose(self, *axes):
936
+ """转置数组。"""
937
+ return _ndarray_methods().transpose(self, *axes)
938
+
939
+ def swapaxes(self, axis1, axis2):
940
+ """交换两个轴。"""
941
+ return _ndarray_methods().swapaxes(self, axis1, axis2)
942
+
943
+ def squeeze(self, axis=None):
944
+ """移除长度为1的轴。"""
945
+ return _ndarray_methods().squeeze(self, axis)
946
+
947
+ def max(self, axis=None):
948
+ """返回最大值。"""
949
+ return _ndarray_methods().max(self, axis)
950
+
951
+ def min(self, axis=None):
952
+ """返回最小值。"""
953
+ return _ndarray_methods().min(self, axis)
954
+
955
+ def mean(self, axis=None):
956
+ """计算平均值。"""
957
+ return _ndarray_methods().mean(self, axis)
958
+
959
+ def std(self, axis=None):
960
+ """计算标准差。"""
961
+ return _ndarray_methods().std(self, axis)
962
+
963
+ def var(self, axis=None):
964
+ """计算方差。"""
965
+ return _ndarray_methods().var(self, axis)
966
+
967
+ def sum(self, axis=None):
968
+ """计算和。"""
969
+ return _ndarray_methods().sum(self, axis)
970
+
971
+ def prod(self, axis=None):
972
+ """计算乘积。"""
973
+ return _ndarray_methods().prod(self, axis)
974
+
975
+ def cumsum(self, axis=None):
976
+ """计算累积和。"""
977
+ return _ndarray_methods().cumsum(self, axis)
978
+
979
+ def cumprod(self, axis=None):
980
+ """计算累积乘积。"""
981
+ return _ndarray_methods().cumprod(self, axis)
982
+
983
+ def argmax(self, axis=None):
984
+ """返回最大值的索引。"""
985
+ return _ndarray_methods().argmax(self, axis)
986
+
987
+ def argmin(self, axis=None):
988
+ """返回最小值的索引。"""
989
+ return _ndarray_methods().argmin(self, axis)
990
+
991
+ def argsort(self, axis=-1):
992
+ """返回排序后的索引。"""
993
+ return _ndarray_methods().argsort(self, axis)
994
+
995
+ def sort(self, axis=-1):
996
+ """原地排序。"""
997
+ _ndarray_methods().sort(self, axis)
998
+
999
+ def diagonal(self, offset=0, axis1=0, axis2=1):
1000
+ """返回对角线元素。"""
1001
+ return _ndarray_methods().diagonal(self, offset, axis1, axis2)
1002
+
1003
+ def trace(self, offset=0, axis1=0, axis2=1):
1004
+ """计算迹。"""
1005
+ return _ndarray_methods().trace(self, offset, axis1, axis2)
1006
+
1007
+ def dot(self, b):
1008
+ """矩阵/向量点积。"""
1009
+ return linalg.dot(self, b)
1010
+
1011
+ def fill(self, value):
1012
+ """用值填充数组。"""
1013
+ _ndarray_methods().fill(self, value)
1014
+
1015
+ def item(self, *args):
1016
+ """获取单个元素。"""
1017
+ return _ndarray_methods().item(self, *args)
1018
+
1019
+ def __index__(self):
1020
+ """支持将 0 维数组用作整数索引或切片参数。"""
1021
+ if self.ndim == 0:
1022
+ val = self.item()
1023
+ if isinstance(val, float):
1024
+ return int(val)
1025
+ return val
1026
+ raise TypeError(f"only 0-dimensional arrays can be converted to integers, got {self.ndim}D")
1027
+
1028
+ def __float__(self):
1029
+ """支持将 0 维数组转换为 Python float。"""
1030
+ if self.ndim == 0:
1031
+ val = self.item()
1032
+ return float(val)
1033
+ raise TypeError(f"only 0-dimensional arrays can be converted to scalars, got {self.ndim}D")
1034
+
1035
+ def __int__(self):
1036
+ """支持将 0 维数组转换为 Python int。"""
1037
+ if self.ndim == 0:
1038
+ val = self.item()
1039
+ return int(val)
1040
+ raise TypeError(f"only 0-dimensional arrays can be converted to scalars, got {self.ndim}D")
1041
+
1042
+ def __complex__(self):
1043
+ """支持将 0 维数组转换为 Python complex。"""
1044
+ if self.ndim == 0:
1045
+ val = self.item()
1046
+ return complex(val)
1047
+ raise TypeError(f"only 0-dimensional arrays can be converted to scalars, got {self.ndim}D")
1048
+
1049
+ def take(self, indices, axis=None):
1050
+ """根据索引取元素。"""
1051
+ return _ndarray_methods().take(self, indices, axis)
1052
+
1053
+ def put(self, indices, values):
1054
+ """设置指定位置的元素。"""
1055
+ _ndarray_methods().put(self, indices, values)
1056
+
1057
+ def repeat(self, repeats, axis=None):
1058
+ """重复元素。"""
1059
+ return _ndarray_methods().repeat(self, repeats, axis)
1060
+
1061
+ def nonzero(self):
1062
+ """返回非零元素索引。"""
1063
+ return _ndarray_methods().nonzero(self)
1064
+
1065
+ # ========== 数学方法 ==========
1066
+
1067
+ def sin(self):
1068
+ """计算正弦。"""
1069
+ return _math_funcs().sin(self)
1070
+
1071
+ def cos(self):
1072
+ """计算余弦。"""
1073
+ return _math_funcs().cos(self)
1074
+
1075
+ def tan(self):
1076
+ """计算正切。"""
1077
+ return _math_funcs().tan(self)
1078
+
1079
+ def sqrt(self):
1080
+ """计算平方根。"""
1081
+ return _math_funcs().sqrt(self)
1082
+
1083
+ def abs(self):
1084
+ """计算绝对值。"""
1085
+ return _math_funcs().abs(self)
1086
+
1087
+ def exp(self):
1088
+ """计算指数。"""
1089
+ return _math_funcs().exp(self)
1090
+
1091
+ def log(self):
1092
+ """计算自然对数。"""
1093
+ return _math_funcs().log(self)
1094
+
1095
+ def log10(self):
1096
+ """计算以10为底的对数。"""
1097
+ return _math_funcs().log10(self)
1098
+
1099
+ def log2(self):
1100
+ """计算以2为底的对数。"""
1101
+ return _math_funcs().log2(self)
1102
+
1103
+ def clip(self, a_min, a_max):
1104
+ """限制值范围。"""
1105
+ return _math_funcs().clip(self, a_min, a_max)
1106
+
1107
+ def round(self, decimals=0):
1108
+ """四舍五入到指定小数位。"""
1109
+ if decimals == 0:
1110
+ return _wrap_result(self._array.__round__(None), self._dtype)
1111
+ return _wrap_result(self._array.__round__(decimals), self._dtype)
1112
+
1113
+ def floor(self):
1114
+ """向下取整。"""
1115
+ return _math_funcs().floor(self)
1116
+
1117
+ def ceil(self):
1118
+ """向上取整。"""
1119
+ return _math_funcs().ceil(self)
1120
+
1121
+ def all(self):
1122
+ """测试是否所有元素为真。"""
1123
+ return np.all(self._array)
1124
+
1125
+ def any(self):
1126
+ """测试是否有任何元素为真。"""
1127
+ return np.any(self._array)
1128
+
1129
+ def byteswap(self, inplace=False):
1130
+ """交换数组中每个元素的字节顺序(大小端转换)。"""
1131
+ data = self._array.flatten().tolist()
1132
+ dtype = getattr(self, '_dtype', 'float64')
1133
+
1134
+ swapped = []
1135
+ for v in data:
1136
+ if isinstance(v, float):
1137
+ v = int(v)
1138
+ if dtype in ('int16', 'uint16'):
1139
+ swapped_val = ((v & 0xFF) << 8) | ((v >> 8) & 0xFF)
1140
+ elif dtype in ('int32', 'uint32'):
1141
+ swapped_val = ((v & 0xFF) << 24) | ((v >> 8 & 0xFF) << 16) | ((v >> 16 & 0xFF) << 8) | ((v >> 24) & 0xFF)
1142
+ elif dtype in ('int64', 'uint64'):
1143
+ swapped_val = 0
1144
+ for i in range(8):
1145
+ swapped_val |= ((v >> (i * 8)) & 0xFF) << ((7 - i) * 8)
1146
+ else:
1147
+ swapped_val = v
1148
+ swapped.append(swapped_val)
1149
+
1150
+ if inplace:
1151
+ arr = np.ndarray(swapped)
1152
+ arr = arr.reshape(self._array.shape)
1153
+ self._array = arr
1154
+ return self
1155
+ else:
1156
+ return ndarray(swapped, _dtype=dtype)
1157
+
1158
+ def view(self, dtype=None, type=None):
1159
+ """创建一个新的数组对象,共享相同的数据但拥有不同的视图。"""
1160
+ _ = type
1161
+ if dtype is not None:
1162
+ return _view_dtype(self, dtype)
1163
+ return ndarray(self._array, _dtype=getattr(self, '_dtype', 'float64'))
1164
+
1165
+ def __reduce__(self):
1166
+ """支持 pickle 序列化。"""
1167
+ return (
1168
+ self.__class__,
1169
+ (
1170
+ self.tolist(),
1171
+ getattr(self, '_dtype', 'float64'),
1172
+ getattr(self, '_fields', None),
1173
+ getattr(self, '_raw_data', None),
1174
+ ),
1175
+ )
1176
+
1177
+
1178
+ class recarray(ndarray):
1179
+ """结构化数组,支持以属性方式访问字段(numpy.recarray 兼容子集)。"""
1180
+
1181
+ def __getattr__(self, name):
1182
+ dt = self.__dict__.get('_dtype_obj')
1183
+ if dt is not None and dt._names and name in dt._names:
1184
+ return self[name]
1185
+ raise AttributeError(
1186
+ "'recarray' object has no attribute %r" % name)
1187
+
1188
+ def __getitem__(self, key):
1189
+ res = ndarray.__getitem__(self, key)
1190
+ if (isinstance(res, ndarray) and res.__class__ is ndarray and getattr(res, '_dtype_obj', None) is not None):
1191
+ res.__class__ = recarray
1192
+ return res
1193
+
1194
+
1195
+ def _as_recarray(arr):
1196
+ """将结构化 ndarray 就地提升为 recarray。"""
1197
+ if isinstance(arr, ndarray) and getattr(arr, '_dtype_obj', None) is not None:
1198
+ arr.__class__ = recarray
1199
+ return arr
1200
+
1201
+
1202
+ # ---------- 内部数组判定与结果包装 ----------
1203
+
1204
+ def _ensure(x):
1205
+ """将列表/元组转换为 ndarray。"""
1206
+ if isinstance(x, (list, tuple)):
1207
+ return np.ndarray(x)
1208
+ elif _is_ndarray(x):
1209
+ return x._array
1210
+ elif hasattr(x, '__class__') and x.__class__.__name__ == 'ndarray':
1211
+ return x
1212
+ return x
1213
+
1214
+
1215
+ def _is_ndarray(obj):
1216
+ """检查对象是否为 rsnumpy ndarray(用 hasattr 避免类身份不一致问题)。"""
1217
+ return hasattr(obj, '_array')
1218
+
1219
+
1220
+ def _wrap_result(result, dtype="float64"):
1221
+ """将原始 ndarray 结果包装到 ndarray 类中。"""
1222
+ if hasattr(result, '__class__') and result.__class__.__name__ == 'ndarray':
1223
+ if getattr(result, 'is_complex', False):
1224
+ dtype = "complex128"
1225
+ inner_array = getattr(result, '_array', result)
1226
+ return ndarray._wrap(inner_array, _dtype=dtype)
1227
+ if isinstance(result, (list, tuple)):
1228
+ return ndarray(result, _dtype=dtype)
1229
+ if isinstance(result, float) and dtype == "int64":
1230
+ return int(result)
1231
+ if hasattr(result, 'tolist'):
1232
+ return ndarray._wrap(result, _dtype=dtype)
1233
+ return result
1234
+
1235
+
1236
+ def _scalar(x):
1237
+ """转换为标量。"""
1238
+ if hasattr(x, 'tolist'):
1239
+ return x.tolist()
1240
+ return x
1241
+
1242
+
1243
+ _FLOAT_DTYPES = ("float16", "float32", "float64")
1244
+
1245
+
1246
+ def _is_float_dtype(dt):
1247
+ return dt in _FLOAT_DTYPES
1248
+
1249
+
1250
+ def _operand_is_float(other):
1251
+ """判断算术运算的另一操作数是否为浮点(bool 视为整数,与 rsnumpy 一致)。"""
1252
+ if _is_ndarray(other):
1253
+ return _is_float_dtype(getattr(other, '_dtype', 'float64'))
1254
+ if isinstance(other, bool):
1255
+ return False
1256
+ return isinstance(other, float)
1257
+
1258
+
1259
+ def _operand_is_complex(other):
1260
+ """判断算术运算的另一操作数是否为复数。"""
1261
+ if _is_ndarray(other):
1262
+ return getattr(other, '_dtype', 'float64') == 'complex128'
1263
+ return isinstance(other, complex)
1264
+
1265
+
1266
+ def _promote_dtype(self_dtype, other):
1267
+ """按 rsnumpy 规则推导加/减/乘结果 dtype:整数遇到浮点操作数提升为 float64,
1268
+ 实数遇到复数操作数提升为 complex128。"""
1269
+ if self_dtype == 'complex128':
1270
+ return 'complex128'
1271
+ if _operand_is_complex(other):
1272
+ return 'complex128'
1273
+ if _is_float_dtype(self_dtype):
1274
+ return self_dtype
1275
+ if _operand_is_float(other):
1276
+ return 'float64'
1277
+ return self_dtype
1278
+
1279
+
1280
+ def _matmul_fallback(a, b):
1281
+ """Fallback for batch matrix multiplication when Rust matmul doesn't support the shapes."""
1282
+ a_shape = a.shape
1283
+ b_shape = b.shape
1284
+
1285
+ if len(a_shape) == 3 and len(b_shape) == 3:
1286
+ batch_size = a_shape[0]
1287
+ result = empty((batch_size, a_shape[1], b_shape[2]), dtype=a._dtype)
1288
+ for i in range(batch_size):
1289
+ result[i] = a[i] @ b[i]
1290
+ return result
1291
+
1292
+ if len(a_shape) == 3 and len(b_shape) == 2:
1293
+ batch_size = a_shape[0]
1294
+ a_rows, a_cols = a_shape[1], a_shape[2]
1295
+ b_rows, b_cols = b_shape[0], b_shape[1]
1296
+
1297
+ if a_cols == b_rows:
1298
+ result = empty((batch_size, a_rows, b_cols), dtype=a._dtype)
1299
+ for i in range(batch_size):
1300
+ result[i] = a[i] @ b
1301
+ return result
1302
+
1303
+ if batch_size == b_rows and a_cols == b_cols:
1304
+ result = empty((batch_size, a_rows, 1), dtype=a._dtype)
1305
+ for i in range(batch_size):
1306
+ ai = a[i]
1307
+ bi = b[i:i+1].T
1308
+ result[i] = ai @ bi
1309
+ return result
1310
+
1311
+ if b_rows == 1 and a_cols == b_cols:
1312
+ result = empty((batch_size, a_rows, 1), dtype=a._dtype)
1313
+ for i in range(batch_size):
1314
+ for r in range(a_rows):
1315
+ s = 0.0
1316
+ for c in range(a_cols):
1317
+ s += a[i, r, c] * b[0, c]
1318
+ result[i, r, 0] = s
1319
+ return result
1320
+
1321
+ if len(a_shape) == 2 and len(b_shape) == 3:
1322
+ batch_size = b_shape[0]
1323
+ result = empty((batch_size, a_shape[0], b_shape[2]), dtype=a._dtype)
1324
+ for i in range(batch_size):
1325
+ result[i] = a @ b[i]
1326
+ return result
1327
+
1328
+ if len(a_shape) == 2 and len(b_shape) == 2:
1329
+ return np.linalg.matmul(a._array, b._array)
1330
+
1331
+ raise ValueError(f"Unsupported shapes for matmul: {a_shape} and {b_shape}")
1332
+
1333
+
1334
+ def _truediv_dtype(self_dtype):
1335
+ """真除法结果恒为浮点:整数提升为 float64,浮点保持自身精度。"""
1336
+ return self_dtype if _is_float_dtype(self_dtype) else 'float64'
1337
+
1338
+
1339
+ # ---------- 嵌套数据与索引转换 ----------
1340
+
1341
+ def _convert_nested(data, converter):
1342
+ """递归转换嵌套列表中的每个元素。"""
1343
+ if isinstance(data, list):
1344
+ return [_convert_nested(x, converter) for x in data]
1345
+ return converter(data)
1346
+
1347
+
1348
+ def _ndarray_to_index_list(k):
1349
+ """将 ndarray 索引转为 Rust 可解析的 Python list,保留多维嵌套结构。"""
1350
+ raw = k._array.tolist() # preserve shape: [[1],[5],[7],[2]] or [1,0,1,0,1]
1351
+ dtype_name = getattr(k, '_dtype', 'float64')
1352
+
1353
+ def _convert(v):
1354
+ if isinstance(v, (list, tuple)):
1355
+ return [_convert(x) for x in v]
1356
+ if dtype_name == 'bool':
1357
+ return bool(v)
1358
+ return int(v)
1359
+
1360
+ return _convert(raw)
1361
+
1362
+
1363
+ def _empty_bool_mask_key(key):
1364
+ """检测布尔掩码索引是否选中零个元素;numpy 语义下对空选择赋值为 no-op。"""
1365
+ items = key if isinstance(key, tuple) else (key,)
1366
+
1367
+ def _has_true(v):
1368
+ if isinstance(v, (list, tuple)):
1369
+ return _py_any(_has_true(x) for x in v)
1370
+ return bool(v)
1371
+
1372
+ for k in items:
1373
+ if hasattr(k, '_array') and getattr(k, '_dtype', None) == 'bool':
1374
+ if not _has_true(k._array.tolist()):
1375
+ return True
1376
+ return False
1377
+
1378
+
1379
+ def _nested_zeros(shape):
1380
+ """生成给定形状的嵌套零列表(标量形状返回 0)。"""
1381
+ if not shape:
1382
+ return 0
1383
+ return [_nested_zeros(shape[1:]) for _ in range(shape[0])]
1384
+
1385
+
1386
+ # ---------- 结构化字段规格与视图 ----------
1387
+
1388
+ def _normalize_field(item):
1389
+ """归一化结构化字段规格:子数组字段保留形状为 (name, type, subshape)。"""
1390
+ if len(item) >= 3 and item[2]:
1391
+ sub = item[2] if isinstance(item[2], (list, tuple)) else (item[2],)
1392
+ return (item[0], item[1], tuple(sub))
1393
+ return (item[0], item[1])
1394
+
1395
+
1396
+ def _field_subshape(fspec):
1397
+ """返回字段的子数组形状(标量字段为空元组)。"""
1398
+ if len(fspec) >= 3 and fspec[2]:
1399
+ return tuple(fspec[2])
1400
+ return ()
1401
+
1402
+
1403
+ def _field_dtype_name(ftype):
1404
+ """将字段类型规格解析为本库 dtype 名称(字符串字段归一为 string_)。"""
1405
+ name = _resolve_type_name(ftype)
1406
+ if isinstance(name, str) and name and name[0] in ('S', 'a', 'U'):
1407
+ return 'string_'
1408
+ return name
1409
+
1410
+
1411
+ def _structured_field_kind(code):
1412
+ """判断结构化字段类型:'int' / 'float' / 'str'。"""
1413
+ if not isinstance(code, str):
1414
+ code = _resolve_type_name(code) if code else 'f8'
1415
+ if not isinstance(code, str):
1416
+ return 'float'
1417
+ c = code.lstrip('<>=|')
1418
+ first = c[:1]
1419
+ if first in ('S', 'U', 'a'):
1420
+ return 'str'
1421
+ low = c.lower()
1422
+ if low.startswith('bytes') or low.startswith('str'):
1423
+ return 'str'
1424
+ is_int = _py_any((
1425
+ low.startswith('int'),
1426
+ low.startswith('uint'),
1427
+ low.startswith('bool'),
1428
+ first == 'i',
1429
+ first == 'u' and c[1:2].isdigit(),
1430
+ ))
1431
+ if is_int:
1432
+ return 'int'
1433
+ return 'float'
1434
+
1435
+
1436
+ def _extract_field(data, fi, depth):
1437
+ """从嵌套结构化 raw_data 中递归抽取第 fi 个字段。depth 为外层维度数。"""
1438
+ if depth <= 0:
1439
+ return data[fi]
1440
+ return [_extract_field(row, fi, depth - 1) for row in data]
1441
+
1442
+
1443
+ def _select_fields(data, idxs, depth):
1444
+ """从嵌套结构化 raw_data 中递归保留 idxs 指定的若干字段。"""
1445
+ if depth <= 0:
1446
+ return tuple(data[i] for i in idxs)
1447
+ return [_select_fields(row, idxs, depth - 1) for row in data]
1448
+
1449
+
1450
+ def _structured_field_view(arr, name):
1451
+ """结构化数组的单字段访问:返回与 arr 同形(子数组字段追加子形状)的字段数组。"""
1452
+ fields = arr._fields
1453
+ names = [f[0] for f in fields]
1454
+ fi = names.index(name)
1455
+ fspec = fields[fi]
1456
+ dt = _field_dtype_name(fspec[1])
1457
+ raw = getattr(arr, '_raw_data', None)
1458
+ if raw is None:
1459
+ # 单字段数组:数据已展平进底层 _array
1460
+ if dt == 'string_':
1461
+ return ndarray._wrap(arr._array, _dtype=dt, _raw_data=getattr(arr, '_raw_data', None))
1462
+ return ndarray._wrap(arr._array, _dtype=dt)
1463
+ extracted = _extract_field(raw, fi, arr.ndim)
1464
+ if dt == 'string_':
1465
+ return ndarray(extracted)
1466
+ return ndarray(extracted, _dtype=dt)
1467
+
1468
+
1469
+ def _structured_multifield_view(arr, keys):
1470
+ """结构化数组的多字段访问 x[['f1','f2']]:返回仅含所选字段的新结构化数组。"""
1471
+ fields = arr._fields
1472
+ names = [f[0] for f in fields]
1473
+ idxs = [names.index(k) for k in keys]
1474
+ new_fields = [fields[i] for i in idxs]
1475
+ raw = getattr(arr, '_raw_data', None)
1476
+ if raw is None:
1477
+ return ndarray._wrap(arr._array, _dtype='void', _fields=new_fields)
1478
+ new_raw = _select_fields(raw, idxs, arr.ndim)
1479
+ return ndarray._wrap(arr._array, _dtype='void', _fields=new_fields, _raw_data=new_raw)
1480
+
1481
+
1482
+ # ---------- 结构化数组格式化 ----------
1483
+
1484
+ def _fmt_structured_str_val(v):
1485
+ if isinstance(v, bytes):
1486
+ return "b'" + v.decode("utf-8", errors="replace") + "'"
1487
+ return "b'" + str(v) + "'"
1488
+
1489
+
1490
+ def _trim_positional_float(v):
1491
+ """numpy dragon4 positional (unique, trim='.') 的最短浮点表示。"""
1492
+ s = repr(float(v))
1493
+ if 'e' in s or 'E' in s:
1494
+ return s
1495
+ if '.' not in s:
1496
+ return s + '.'
1497
+ intp, frac = s.split('.', 1)
1498
+ return intp + '.' + frac.rstrip('0')
1499
+
1500
+
1501
+ def _format_float_field_column(col):
1502
+ """按 rsnumpy 浮点列规则将一列浮点值格式化为等宽对齐的字符串。"""
1503
+ n = len(col)
1504
+ finite = [None] * n
1505
+ special = [None] * n
1506
+ for i, v in enumerate(col):
1507
+ fv = float(v)
1508
+ if fv != fv:
1509
+ special[i] = 'nan'
1510
+ elif fv == float('inf'):
1511
+ special[i] = 'inf'
1512
+ elif fv == float('-inf'):
1513
+ special[i] = '-inf'
1514
+ else:
1515
+ finite[i] = _trim_positional_float(fv)
1516
+ out = [None] * n
1517
+ pad_left = 0
1518
+ pad_right = 0
1519
+ for s in finite:
1520
+ if s is not None and 'e' not in s:
1521
+ intp, frac = s.split('.', 1)
1522
+ pad_left = _py_max(pad_left, len(intp))
1523
+ pad_right = _py_max(pad_right, len(frac))
1524
+ for i, s in enumerate(finite):
1525
+ if s is None or 'e' in s:
1526
+ continue
1527
+ intp, frac = s.split('.', 1)
1528
+ out[i] = intp.rjust(pad_left) + '.' + frac.ljust(pad_right)
1529
+ width = 0
1530
+ for i in range(n):
1531
+ s = out[i] if out[i] is not None else (finite[i] if finite[i] is not None else special[i])
1532
+ if s is not None:
1533
+ width = _py_max(width, len(s))
1534
+ for i in range(n):
1535
+ if out[i] is None:
1536
+ out[i] = (finite[i] if finite[i] is not None else special[i])
1537
+ out[i] = out[i].rjust(width)
1538
+ return out
1539
+
1540
+
1541
+ def _format_int_field_column(col):
1542
+ strs = [str(v) if isinstance(v, str) else str(int(_py_round(float(v)))) for v in col]
1543
+ width = _py_max((len(s) for s in strs), default=0)
1544
+ return [s.rjust(width) for s in strs]
1545
+
1546
+
1547
+ def _format_structured_field_column(col, code):
1548
+ kind = _structured_field_kind(code)
1549
+ if kind == 'str':
1550
+ return [_fmt_structured_str_val(v) for v in col]
1551
+ if kind == 'int':
1552
+ return _format_int_field_column(col)
1553
+ return _format_float_field_column(col)
1554
+
1555
+
1556
+ def _format_structured_str(arr):
1557
+ """__str__ 用于结构化数组:基于富 DType 递归格式化,支持嵌套/子数组。"""
1558
+ dt = getattr(arr, '_dtype_obj', None)
1559
+ if dt is None:
1560
+ return _format_structured_str_legacy(arr)
1561
+
1562
+ def build(o):
1563
+ if isinstance(o, tuple):
1564
+ return _fmt_record(o, dt)
1565
+ return "[" + " ".join(build(x) for x in o) + "]"
1566
+
1567
+ raw = arr._raw_data
1568
+ if arr.ndim == 0:
1569
+ return _fmt_record(raw, dt)
1570
+ return build(raw)
1571
+
1572
+
1573
+ def _format_structured_str_legacy(arr):
1574
+ """旧版结构化数组格式化(基于 _fields codes),保留作回退。"""
1575
+ fields = getattr(arr, '_fields', None)
1576
+ data = arr.tolist()
1577
+ ndim = arr.ndim
1578
+ nfields = len(fields)
1579
+ codes = [f[1] for f in fields]
1580
+
1581
+ leaves = []
1582
+ if ndim <= 1:
1583
+ leaves = list(data)
1584
+ else:
1585
+ def _collect(d, depth):
1586
+ if depth == ndim - 1:
1587
+ leaves.extend(d)
1588
+ else:
1589
+ for sub in d:
1590
+ _collect(sub, depth + 1)
1591
+ _collect(data, 0)
1592
+ norm = [t if isinstance(t, (list, tuple)) else (t,) for t in leaves]
1593
+
1594
+ col_strs = []
1595
+ for j in range(nfields):
1596
+ col_strs.append(_format_structured_field_column([r[j] for r in norm], codes[j]))
1597
+
1598
+ tuple_strs = []
1599
+ for i in range(len(norm)):
1600
+ parts = [col_strs[j][i] for j in range(nfields)]
1601
+ if nfields == 1:
1602
+ tuple_strs.append("(" + parts[0] + ",)")
1603
+ else:
1604
+ tuple_strs.append("(" + ", ".join(parts) + ")")
1605
+
1606
+ if ndim <= 1:
1607
+ return "[" + " ".join(tuple_strs) + "]"
1608
+
1609
+ shape = arr.shape
1610
+ pos = [0]
1611
+
1612
+ def _build(dim):
1613
+ if dim == ndim - 1:
1614
+ chunk = tuple_strs[pos[0]:pos[0] + shape[dim]]
1615
+ pos[0] += shape[dim]
1616
+ return "[" + " ".join(chunk) + "]"
1617
+ return "[" + "\n ".join(_build(dim + 1) for _ in range(shape[dim])) + "]"
1618
+
1619
+ return _build(0)
1620
+
1621
+
1622
+ def _format_structured_repr(arr):
1623
+ """__repr__ 用于结构化数组。"""
1624
+ inner = _format_structured_str(arr)
1625
+ return f"array({inner})"
1626
+
1627
+
1628
+ # ---------- 不规则/嵌套数组格式化 ----------
1629
+
1630
+ def _format_ragged_str(arr):
1631
+ """__str__ 用于不规则数组(包含列表/元组元素)。"""
1632
+ data = arr.tolist()
1633
+ return _format_nested_iterable(data)
1634
+
1635
+
1636
+ def _format_ragged_repr(arr):
1637
+ """__repr__ 用于不规则数组。"""
1638
+ inner = _format_ragged_str(arr)
1639
+ return f"array({inner})"
1640
+
1641
+
1642
+ def _format_nested_iterable(data):
1643
+ """递归格式化嵌套可迭代对象为数组字符串。"""
1644
+ if isinstance(data, (list, tuple)):
1645
+ if not data:
1646
+ return "[]"
1647
+ items = list(data)
1648
+ has_nested = _py_any(isinstance(x, (list, tuple)) for x in items)
1649
+ if not has_nested:
1650
+ return "[" + " ".join(str(x) for x in items) + "]"
1651
+ parts = [_format_nested_iterable(x) for x in items]
1652
+ return "[" + "\n ".join(parts) + "]"
1653
+ return str(data)
1654
+
1655
+
1656
+ # ---------- 浮点/复数标量格式化 ----------
1657
+
1658
+ class _float64:
1659
+ """float64 标量,显示为 np.float64(value)。"""
1660
+ def __init__(self, value):
1661
+ self._value = float(value)
1662
+
1663
+ def __repr__(self):
1664
+ return f"np.float64({self._value})"
1665
+
1666
+ def __str__(self):
1667
+ return repr(self)
1668
+
1669
+ def __float__(self):
1670
+ return self._value
1671
+
1672
+
1673
+ def _format_float_repr_1d(values):
1674
+ """格式化 1D float64 数组的 repr,带逗号和对齐(匹配 rsnumpy)。"""
1675
+ fmt = [format_float_scalar(v) for v in values]
1676
+ max_w = np.imax([len(f) for f in fmt]) if fmt else 0
1677
+ parts = [f.rjust(max_w) for f in fmt]
1678
+ return "[" + ", ".join(parts) + "]"
1679
+
1680
+
1681
+ def _format_complex_repr_1d(values):
1682
+ """格式化 1D complex128 数组的 repr。"""
1683
+ parts = [_format_complex_scalar(v) for v in values]
1684
+ return "[" + ", ".join(parts) + "]"
1685
+
1686
+
1687
+ def _format_complex_nested(data, sep):
1688
+ """递归格式化原生复数数组(nested complex)为 rsnumpy 风格字符串。"""
1689
+ if isinstance(data, list):
1690
+ inner = sep.join(_format_complex_nested(x, sep) for x in data)
1691
+ return "[" + inner + "]"
1692
+ return _format_complex_scalar(data)
1693
+
1694
+
1695
+ def _format_complex_scalar(val):
1696
+ """格式化单个复数为字符串(如 1.+0.j, 2.+6.j)。"""
1697
+ real = val.real
1698
+ imag = val.imag
1699
+ if imag >= 0:
1700
+ return f"({real}{imag:+}j)"
1701
+ return f"({real}{imag}j)"
1702
+
1703
+
1704
+ def format_float_scalar(val):
1705
+ """Python 版的浮点数格式化(四舍五入,保留有效小数)。"""
1706
+ if val != val:
1707
+ return "nan"
1708
+ if val == float("inf"):
1709
+ return "inf"
1710
+ if val == float("-inf"):
1711
+ return "-inf"
1712
+ if val == 0:
1713
+ return "-0." if str(val)[0] == '-' else "0."
1714
+ val_rounded = _py_round(val, 8)
1715
+ if val_rounded == int(val_rounded) and abs(val_rounded) < 1e16:
1716
+ v = int(val_rounded)
1717
+ if float(v) == val_rounded:
1718
+ return f"{v}."
1719
+ if abs(val_rounded) >= 1e10 or (abs(val_rounded) < 1e-10 and val_rounded != 0):
1720
+ return f"{val_rounded:.10e}"
1721
+ s = f"{val_rounded:.8f}"
1722
+ if '.' in s:
1723
+ s = s.rstrip('0').rstrip('.')
1724
+ return s
1725
+
1726
+
1727
+ # ---------- 子模块延迟访问 ----------
1728
+
1729
+ def _ndarray_methods():
1730
+ """延迟导入 array_methods 模块。"""
1731
+ return array_methods.NdArrayMethods
1732
+
1733
+
1734
+ def _math_funcs():
1735
+ """延迟导入 math_functions 模块。"""
1736
+ from . import math_functions as mf
1737
+ return mf
1738
+
1739
+
1740
+ def _stat_funcs():
1741
+ """延迟导入 statistics 模块。"""
1742
+ from . import statistics as st
1743
+ return st
1744
+
1745
+
1746
+ def _array_ops():
1747
+ """延迟导入 array_ops 模块。"""
1748
+ from . import array_ops as ao
1749
+ return ao
1750
+
1751
+
1752
+ # ========== 类型代码映射 ==========
1753
+
1754
+ _RS_DTYPE_CODES = {
1755
+ # 有符号整数
1756
+ 'i1': 'int8', 'int8': 'int8',
1757
+ 'i2': 'int16', 'int16': 'int16',
1758
+ 'i4': 'int32', 'int32': 'int32',
1759
+ 'i8': 'int64', 'int64': 'int64',
1760
+ # 无符号整数
1761
+ 'u1': 'uint8', 'uint8': 'uint8',
1762
+ 'u2': 'uint16', 'uint16': 'uint16',
1763
+ 'u4': 'uint32', 'uint32': 'uint32',
1764
+ 'u8': 'uint64', 'uint64': 'uint64',
1765
+ # 浮点
1766
+ 'f2': 'float16', 'float16': 'float16',
1767
+ 'f4': 'float32', 'float32': 'float32',
1768
+ 'f8': 'float64', 'float64': 'float64', 'float': 'float64',
1769
+ # 复数
1770
+ 'c8': 'complex64', 'c64': 'complex64', 'complex64': 'complex64',
1771
+ 'c16': 'complex128', 'c128': 'complex128', 'complex128': 'complex128', 'complex': 'complex128',
1772
+ # 其他
1773
+ 'bool': 'bool', 'b1': 'bool',
1774
+ 'object': 'object', 'O': 'object',
1775
+ }
1776
+
1777
+ _PY_TYPE_TO_RS = {
1778
+ int: 'int64',
1779
+ float: 'float64',
1780
+ complex: 'complex128',
1781
+ bool: 'bool',
1782
+ bytes: 'bytes',
1783
+ str: 'str',
1784
+ }
1785
+
1786
+ # 运行时收集已知的 rsnumpy 类型名称
1787
+ _numpy_type_names = {}
1788
+
1789
+
1790
+ def _init_numpy_types():
1791
+ """尝试导入 rsnumpy 并注册其标量类型名称。"""
1792
+ try:
1793
+ import rsnnumpy as np
1794
+ except ImportError:
1795
+ return
1796
+ for name in dir(np):
1797
+ if name.startswith('int') or name.startswith('uint') or name.startswith('float') \
1798
+ or name.startswith('complex') or name == 'bool_' or name == 'bool':
1799
+ _numpy_type_names[name] = name
1800
+
1801
+
1802
+ def _resolve_type_name(tp):
1803
+ """将各种 dtype 输入解析为类型名字符串。"""
1804
+ if isinstance(tp, str):
1805
+ # 直接映射
1806
+ lower = tp.lower()
1807
+ if lower in _RS_DTYPE_CODES:
1808
+ return _RS_DTYPE_CODES[lower]
1809
+ return tp # 未知字符串,原样返回
1810
+ if isinstance(tp, DType):
1811
+ return tp.name
1812
+ if type(tp).__name__ == 'dtype' and hasattr(tp, 'name'):
1813
+ return _RS_DTYPE_CODES.get(tp.name, tp.name)
1814
+ # Python 内置类型
1815
+ if tp in _PY_TYPE_TO_RS:
1816
+ return _PY_TYPE_TO_RS[tp]
1817
+ # 检查类名
1818
+ name = getattr(tp, '__name__', None) or getattr(tp, 'name', str(tp))
1819
+ # 先查已知的 rsnumpy 类型名
1820
+ if name in _numpy_type_names:
1821
+ return _numpy_type_names[name]
1822
+ # 再查代码表
1823
+ if name in _RS_DTYPE_CODES:
1824
+ return _RS_DTYPE_CODES[name]
1825
+ # int_ / float_ 等 rsnumpy 特殊名称
1826
+ name_clean = name.rstrip('_')
1827
+ if name_clean in _RS_DTYPE_CODES:
1828
+ return _RS_DTYPE_CODES[name_clean]
1829
+ return name
1830
+
1831
+
1832
+ # 类型名 → 类型码 逆向映射(用于结构化 dtype)
1833
+ _RS_NAME_TO_CODE = {}
1834
+ for _code, _name in _RS_DTYPE_CODES.items():
1835
+ # 只保留短码映射,忽略 'int32': 'int32' 这样的自映射
1836
+ if len(_code) <= 4 and _code not in ('float', 'complex', 'bool') and _code != _name:
1837
+ _RS_NAME_TO_CODE[_name] = _code
1838
+
1839
+
1840
+ # 立即初始化 rsnumpy 类型名称
1841
+ _init_numpy_types()
1842
+
1843
+
1844
+ def _legacy_fields_from_dtype(dt):
1845
+ """从富 DType 结构化类型构建旧版 _fields 列表:(name, code[, subshape])。
1846
+
1847
+ 仅用于与依赖旧格式的模块(io/statistics/array_ops)保持兼容;
1848
+ 结构化格式化与索引优先使用 _dtype_obj。嵌套结构字段存 DType 本身。
1849
+ """
1850
+ if dt is None or dt._names is None:
1851
+ return None
1852
+ out = []
1853
+ for n in dt._names:
1854
+ fdt = dt._fields[n][0]
1855
+ if fdt._subdtype is not None:
1856
+ base, shape = fdt._subdtype
1857
+ out.append((n, _scalar_typestr_short(base), tuple(shape)))
1858
+ elif fdt._names is not None:
1859
+ out.append((n, fdt))
1860
+ else:
1861
+ out.append((n, _scalar_typestr_short(fdt)))
1862
+ return out
1863
+
1864
+
1865
+ def _struct_dt_of(arr):
1866
+ """返回数组的结构化富 DType:优先 _dtype_obj,否则从旧版 _fields 重建。"""
1867
+ dt = getattr(arr, '_dtype_obj', None)
1868
+ if dt is not None:
1869
+ return dt
1870
+ fields = getattr(arr, '_fields', None)
1871
+ if not fields:
1872
+ return None
1873
+ names, formats = [], []
1874
+ for f in fields:
1875
+ names.append(f[0])
1876
+ typ = f[1]
1877
+ fdt = typ if isinstance(typ, DType) else dtype(typ)
1878
+ if len(f) > 2 and f[2]:
1879
+ sub = f[2] if isinstance(f[2], (list, tuple)) else (f[2],)
1880
+ fdt = _make_subarray(fdt, tuple(sub))
1881
+ formats.append(fdt)
1882
+ return _build_struct(names, formats, None, None, None, False)
1883
+
1884
+
1885
+ # ========== 结构化数组统一数据模型 ==========
1886
+ # 结构化数组统一存储:_dtype_obj(富 DType,语义真源)、_fields(旧版列表,供
1887
+ # 既有格式化/视图代码)、_raw_data(Python 记录嵌套列表)、_array(占位取形状)。
1888
+
1889
+ def _struct_shape(raw):
1890
+ """从结构化 raw_data 的嵌套列表推断形状(记录为 tuple,视作叶子)。"""
1891
+ shape = []
1892
+ obj = raw
1893
+ while isinstance(obj, list):
1894
+ shape.append(len(obj))
1895
+ obj = obj[0] if obj else None
1896
+ return tuple(shape)
1897
+
1898
+
1899
+ def _zero_scalar_for(fdt):
1900
+ k = fdt._kind
1901
+ if k in ('i', 'u', 'b'):
1902
+ return 0
1903
+ if k == 'f':
1904
+ return 0.0
1905
+ if k == 'c':
1906
+ return 0j
1907
+ if k == 'S':
1908
+ return b''
1909
+ if k == 'U':
1910
+ return ''
1911
+ if k == 'V':
1912
+ return b'\x00' * fdt._itemsize
1913
+ return 0
1914
+
1915
+
1916
+ def _one_scalar_for(fdt):
1917
+ k = fdt._kind
1918
+ if k in ('i', 'u', 'b'):
1919
+ return 1
1920
+ if k == 'f':
1921
+ return 1.0
1922
+ if k == 'c':
1923
+ return 1 + 0j
1924
+ if k == 'S':
1925
+ return b'1'
1926
+ if k == 'U':
1927
+ return '1'
1928
+ return _zero_scalar_for(fdt)
1929
+
1930
+
1931
+ def _coerce_scalar_by_dt(v, fdt):
1932
+ k = fdt._kind
1933
+ if k in ('i', 'u'):
1934
+ return int(v)
1935
+ if k == 'b':
1936
+ return bool(v)
1937
+ if k == 'f':
1938
+ return float(v)
1939
+ if k == 'c':
1940
+ return complex(v)
1941
+ if k == 'S':
1942
+ if isinstance(v, (bytes, bytearray)):
1943
+ return bytes(v)
1944
+ return str(v).encode('latin-1')
1945
+ if k == 'U':
1946
+ if isinstance(v, (bytes, bytearray)):
1947
+ return bytes(v).decode('latin-1')
1948
+ return str(v)
1949
+ return v
1950
+
1951
+
1952
+ def _coerce_field_value(val, fdt):
1953
+ if fdt._names is not None:
1954
+ seq = tuple(val) if isinstance(val, (list, tuple)) else (val,)
1955
+ return tuple(_coerce_field_value(seq[i], fdt._fields[n][0])
1956
+ for i, n in enumerate(fdt._names))
1957
+ if fdt._subdtype is not None:
1958
+ base, shape = fdt._subdtype
1959
+ return _coerce_subarray(val, base, shape)
1960
+ return _coerce_scalar_by_dt(val, fdt)
1961
+
1962
+
1963
+ def _coerce_subarray(val, base, shape):
1964
+ if not shape:
1965
+ return _coerce_field_value(val, base)
1966
+ if not isinstance(val, (list, tuple)):
1967
+ return [_coerce_subarray(val, base, shape[1:]) for _ in range(shape[0])]
1968
+ return [_coerce_subarray(val[i], base, shape[1:]) for i in range(shape[0])]
1969
+
1970
+
1971
+ def _coerce_record(rec, dt):
1972
+ seq = tuple(rec) if isinstance(rec, (list, tuple)) else (rec,)
1973
+ return tuple(_coerce_field_value(seq[i], dt._fields[n][0])
1974
+ for i, n in enumerate(dt._names))
1975
+
1976
+
1977
+ def _coerce_struct_data(data, dt):
1978
+ if isinstance(data, tuple):
1979
+ return _coerce_record(data, dt)
1980
+ if isinstance(data, list):
1981
+ return [_coerce_struct_data(x, dt) for x in data]
1982
+ # 标量广播到单字段记录
1983
+ return _coerce_record((data,), dt)
1984
+
1985
+
1986
+ def _default_field(fdt, filler):
1987
+ if fdt._names is not None:
1988
+ return tuple(_default_field(fdt._fields[n][0], filler) for n in fdt._names)
1989
+ if fdt._subdtype is not None:
1990
+ base, shape = fdt._subdtype
1991
+ return _default_subarray(base, shape, filler)
1992
+ return filler(fdt)
1993
+
1994
+
1995
+ def _default_subarray(base, shape, filler):
1996
+ if not shape:
1997
+ return _default_field(base, filler)
1998
+ return [_default_subarray(base, shape[1:], filler) for _ in range(shape[0])]
1999
+
2000
+
2001
+ def _default_record(dt, filler):
2002
+ return tuple(_default_field(dt._fields[n][0], filler) for n in dt._names)
2003
+
2004
+
2005
+ def _make_struct_filled(shape, dt, filler):
2006
+ rec = _default_record(dt, filler)
2007
+ if isinstance(shape, int):
2008
+ shape = (shape,)
2009
+ total = 1
2010
+ for s in shape:
2011
+ total *= s
2012
+ flat = [rec for _ in range(total)]
2013
+ if len(shape) <= 1:
2014
+ return flat
2015
+ result = flat
2016
+ for dim in reversed(shape[1:]):
2017
+ result = [result[i:i + dim] for i in range(0, len(result), dim)]
2018
+ return result
2019
+
2020
+
2021
+ def _wrap_structured(dt, raw_data):
2022
+ """用富 DType 与已归一的 raw_data 构造结构化 ndarray。"""
2023
+ raw_data = list(raw_data)
2024
+ shape = _struct_shape(raw_data)
2025
+ obj = ndarray.__new__(ndarray)
2026
+ obj._array = np.zeros(shape if shape else (0,))
2027
+ obj._dtype = 'void'
2028
+ obj._dtype_obj = dt
2029
+ obj._fields = _legacy_fields_from_dtype(dt)
2030
+ obj._raw_data = raw_data
2031
+ return obj
2032
+
2033
+
2034
+ def _make_structured_array(data, dt):
2035
+ """从任意输入数据 + 富结构化 DType 构造 ndarray。"""
2036
+ raw = _coerce_struct_data(list(data) if isinstance(data, (list, tuple)) else data, dt)
2037
+ if not isinstance(raw, list):
2038
+ raw = [raw]
2039
+ return _wrap_structured(dt, raw)
2040
+
2041
+
2042
+ def _as_struct_dtype(dt):
2043
+ """若 dt 可解析为结构化 DType 则返回富 DType,否则 None。"""
2044
+ if dt is None:
2045
+ return None
2046
+ try:
2047
+ d = dt if isinstance(dt, DType) else dtype(dt)
2048
+ except Exception:
2049
+ return None
2050
+ if isinstance(d, DType) and d._names is not None:
2051
+ return d
2052
+ return None
2053
+
2054
+
2055
+ # ========== 结构化数组操作(索引/赋值/视图/比较/格式化) ==========
2056
+
2057
+ def _scalar_name_of(fdt):
2058
+ """标量 DType → 内部 dtype 名称(int32/float64/...)。"""
2059
+ if fdt._typename is not None:
2060
+ return fdt._typename
2061
+ return fdt.name
2062
+
2063
+
2064
+ def _flatten_records(raw):
2065
+ """将结构化 raw_data(嵌套列表,记录为 tuple 叶子)展平为记录列表(行主序)。"""
2066
+ out = []
2067
+
2068
+ def rec(o):
2069
+ if isinstance(o, tuple):
2070
+ out.append(o)
2071
+ elif isinstance(o, list):
2072
+ for x in o:
2073
+ rec(x)
2074
+ else:
2075
+ out.append(o)
2076
+
2077
+ rec(raw)
2078
+ return out
2079
+
2080
+
2081
+ def _reshape_flat(flat, shape):
2082
+ """将扁平列表按 shape 重塑为嵌套列表(叶子视为原子)。"""
2083
+ flat = list(flat)
2084
+ if not shape or len(shape) <= 1:
2085
+ return flat
2086
+ result = flat
2087
+ for dim in reversed(shape[1:]):
2088
+ result = [result[i:i + dim] for i in range(0, len(result), dim)]
2089
+ return result
2090
+
2091
+
2092
+ def _flat_to_multi(flat, shape):
2093
+ idx = []
2094
+ for s in reversed(shape):
2095
+ idx.append(flat % s)
2096
+ flat //= s
2097
+ return tuple(reversed(idx))
2098
+
2099
+
2100
+ def _get_flat_record(arr, flat):
2101
+ idx = _flat_to_multi(flat, arr.shape)
2102
+ o = arr._raw_data
2103
+ for i in idx:
2104
+ o = o[i]
2105
+ return o
2106
+
2107
+
2108
+ def _set_flat_record(arr, flat, new_rec):
2109
+ shape = arr.shape
2110
+ if len(shape) <= 1:
2111
+ arr._raw_data[flat] = new_rec
2112
+ return
2113
+ idx = _flat_to_multi(flat, shape)
2114
+ o = arr._raw_data
2115
+ for i in idx[:-1]:
2116
+ o = o[i]
2117
+ o[idx[-1]] = new_rec
2118
+
2119
+
2120
+ def _build_plain_array(nested, base_dt, shape):
2121
+ """从嵌套列表 + 标量基础 DType 构造非结构化 ndarray。"""
2122
+ k = base_dt._kind
2123
+ if k in ('S', 'U', 'V'):
2124
+ flat = _flatten_data(nested) if isinstance(nested, list) else [nested]
2125
+ arr = ndarray._wrap(np.zeros((len(flat),) if flat else (0,)),
2126
+ _dtype='string_', _raw_data=list(flat))
2127
+ width = base_dt._itemsize // 4 if k == 'U' else base_dt._itemsize
2128
+ arr._str_dtype = (k, width)
2129
+ return arr
2130
+ name = _scalar_name_of(base_dt)
2131
+ if k == 'c':
2132
+ flat = _flatten_data(nested) if isinstance(nested, list) else [nested]
2133
+ cflat = [complex(v) for v in flat]
2134
+ if not cflat:
2135
+ return ndarray._wrap(np.zeros(tuple(shape) if shape else (0,)),
2136
+ _dtype='complex128')
2137
+ nested_c = _reshape_flat(cflat, shape) if shape else cflat
2138
+ return ndarray._wrap(np.ndarray(nested_c), _dtype='complex128')
2139
+ return ndarray(nested, _dtype=name)
2140
+
2141
+
2142
+ def _field_values_to_array(values, fdt, outer):
2143
+ """把每条记录的某字段值列表按字段 DType 组装成 ndarray。"""
2144
+ if fdt._names is not None:
2145
+ return _wrap_structured(fdt, _reshape_flat(values, outer))
2146
+ if fdt._subdtype is not None:
2147
+ base, sub = fdt._subdtype
2148
+ nested = _reshape_flat(values, outer)
2149
+ return _build_plain_array(nested, base, tuple(outer) + tuple(sub))
2150
+ nested = _reshape_flat(values, outer)
2151
+ return _build_plain_array(nested, fdt, tuple(outer))
2152
+
2153
+
2154
+ def _get_field(arr, name):
2155
+ """结构化数组的单字段访问 a['f']。"""
2156
+ dt = _struct_dt_of(arr)
2157
+ fi = dt._names.index(name)
2158
+ fdt = dt._fields[name][0]
2159
+ recs = _flatten_records(arr._raw_data)
2160
+ values = [r[fi] for r in recs]
2161
+ return _field_values_to_array(values, fdt, arr.shape)
2162
+
2163
+
2164
+ def _get_multifield(arr, keys):
2165
+ """多字段访问 a[['f1','f2']]:返回保留原偏移/itemsize 的结构化视图。"""
2166
+ dt = _struct_dt_of(arr)
2167
+ newdt = dt[list(keys)]
2168
+ idxs = [dt._names.index(k) for k in keys]
2169
+ recs = _flatten_records(arr._raw_data)
2170
+ new_recs = [tuple(r[i] for i in idxs) for r in recs]
2171
+ return _wrap_structured(newdt, _reshape_flat(new_recs, arr.shape))
2172
+
2173
+
2174
+ def _field_leaf_to_scalar(v, fdt):
2175
+ """将记录中单个字段值转换为对外标量/数组。"""
2176
+ if fdt._names is not None:
2177
+ return _make_record(v, fdt, None, None)
2178
+ if fdt._subdtype is not None:
2179
+ base, sub = fdt._subdtype
2180
+ return _build_plain_array(v, base, tuple(sub))
2181
+ k = fdt._kind
2182
+ if k == 'S':
2183
+ return _str_scalar(v, ('S', fdt._itemsize))
2184
+ if k == 'U':
2185
+ return _str_scalar(v, ('U', fdt._itemsize // 4))
2186
+ if k == 'V':
2187
+ return _str_scalar(v, ('V', fdt._itemsize))
2188
+ return v
2189
+
2190
+
2191
+ def _make_record(values, dt, parent, flat_index):
2192
+ return _Record(list(values), dt, parent, flat_index)
2193
+
2194
+
2195
+ def _get_record(arr, key):
2196
+ dt = _struct_dt_of(arr)
2197
+ recs = _flatten_records(arr._raw_data)
2198
+ n = len(recs)
2199
+ k = key + n if key < 0 else key
2200
+ return _make_record(recs[k], dt, arr, k)
2201
+
2202
+
2203
+ def _struct_getitem(arr, key):
2204
+ """结构化数组索引分发,返回 (handled, result)。"""
2205
+ dt = _struct_dt_of(arr)
2206
+ if dt is None:
2207
+ return (False, None)
2208
+ if isinstance(key, str):
2209
+ if dt._names and key in dt._names:
2210
+ return (True, _get_field(arr, key))
2211
+ return (False, None)
2212
+ if isinstance(key, list) and key and _py_all(isinstance(k, str) for k in key):
2213
+ return (True, _get_multifield(arr, key))
2214
+ if isinstance(key, int) and not isinstance(key, bool):
2215
+ if arr.ndim <= 1:
2216
+ return (True, _get_record(arr, key))
2217
+ sub = arr._raw_data[key]
2218
+ return (True, _wrap_structured(dt, sub if isinstance(sub, list) else [sub]))
2219
+ if isinstance(key, slice):
2220
+ return (True, _wrap_structured(dt, arr._raw_data[key]))
2221
+ return (False, None)
2222
+
2223
+
2224
+ def _broadcast_field_values(value, n):
2225
+ if _is_ndarray(value):
2226
+ value = value.tolist()
2227
+ if isinstance(value, list) and len(value) == n:
2228
+ return value
2229
+ if isinstance(value, tuple) and len(value) == n:
2230
+ return list(value)
2231
+ return [value] * n
2232
+
2233
+
2234
+ def _set_field(arr, name, value):
2235
+ dt = _struct_dt_of(arr)
2236
+ fi = dt._names.index(name)
2237
+ fdt = dt._fields[name][0]
2238
+ recs = _flatten_records(arr._raw_data)
2239
+ vals = _broadcast_field_values(value, len(recs))
2240
+ new = []
2241
+ for r, v in zip(recs, vals):
2242
+ lst = list(r)
2243
+ lst[fi] = _coerce_field_value(v, fdt)
2244
+ new.append(tuple(lst))
2245
+ arr._raw_data = _reshape_flat(new, arr.shape)
2246
+
2247
+
2248
+ def _set_multifield(arr, keys, value):
2249
+ dt = _struct_dt_of(arr)
2250
+ idxs = [dt._names.index(k) for k in keys]
2251
+ recs = _flatten_records(arr._raw_data)
2252
+ if _is_ndarray(value):
2253
+ value = value.tolist()
2254
+ new = []
2255
+ for r in recs:
2256
+ lst = list(r)
2257
+ for j, k in enumerate(keys):
2258
+ fv = value[j] if isinstance(value, (list, tuple)) else value
2259
+ lst[idxs[j]] = _coerce_field_value(fv, dt._fields[k][0])
2260
+ new.append(tuple(lst))
2261
+ arr._raw_data = _reshape_flat(new, arr.shape)
2262
+
2263
+
2264
+ def _set_record(arr, key, value):
2265
+ dt = _struct_dt_of(arr)
2266
+ recs = _flatten_records(arr._raw_data)
2267
+ n = len(recs)
2268
+ k = key + n if key < 0 else key
2269
+ if _is_ndarray(value):
2270
+ value = value.tolist()
2271
+ seq = value if isinstance(value, (list, tuple)) else (value,)
2272
+ _set_flat_record(arr, k, _coerce_record(seq, dt))
2273
+
2274
+
2275
+ def _set_slice(arr, key, value):
2276
+ dt = _struct_dt_of(arr)
2277
+ recs = _flatten_records(arr._raw_data)
2278
+ n = len(recs)
2279
+ indices = list(range(*key.indices(n))) if isinstance(key, slice) else list(range(n))
2280
+ names = dt._names
2281
+ if _is_ndarray(value) and getattr(value, '_dtype_obj', None) is not None:
2282
+ vrecs = _flatten_records(value._raw_data)
2283
+ for pos, i in enumerate(indices):
2284
+ src = vrecs[pos % len(vrecs)]
2285
+ recs[i] = tuple(_coerce_field_value(src[j], dt._fields[names[j]][0])
2286
+ for j in range(len(names)))
2287
+ elif _is_ndarray(value):
2288
+ vals = value.tolist()
2289
+ for pos, i in enumerate(indices):
2290
+ sv = vals[pos % len(vals)]
2291
+ recs[i] = tuple(_coerce_field_value(sv, dt._fields[nm][0]) for nm in names)
2292
+ elif isinstance(value, (list, tuple)) and len(value) == len(names):
2293
+ rec = _coerce_record(value, dt)
2294
+ for i in indices:
2295
+ recs[i] = rec
2296
+ else:
2297
+ for i in indices:
2298
+ recs[i] = tuple(_coerce_field_value(value, dt._fields[nm][0]) for nm in names)
2299
+ arr._raw_data = _reshape_flat(recs, arr.shape)
2300
+
2301
+
2302
+ def _struct_setitem(arr, key, value):
2303
+ dt = _struct_dt_of(arr)
2304
+ if dt is None:
2305
+ return False
2306
+ if isinstance(key, str):
2307
+ if dt._names and key in dt._names:
2308
+ _set_field(arr, key, value)
2309
+ return True
2310
+ return False
2311
+ if isinstance(key, list) and key and _py_all(isinstance(k, str) for k in key):
2312
+ _set_multifield(arr, key, value)
2313
+ return True
2314
+ if isinstance(key, int) and not isinstance(key, bool):
2315
+ _set_record(arr, key, value)
2316
+ return True
2317
+ if isinstance(key, slice):
2318
+ _set_slice(arr, key, value)
2319
+ return True
2320
+ return False
2321
+
2322
+
2323
+ # ---------- 字节打包/解包(供 view 重解释使用) ----------
2324
+
2325
+ def _pack_scalar(buf, off, val, fdt):
2326
+ import struct
2327
+ k = fdt._kind
2328
+ sz = fdt._itemsize
2329
+ if k in ('i', 'u', 'b'):
2330
+ iv = int(val)
2331
+ if k == 'i' and iv < 0:
2332
+ iv &= (1 << (sz * 8)) - 1
2333
+ buf[off:off + sz] = (iv & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
2334
+ elif k == 'f':
2335
+ fmt = {2: 'e', 4: 'f', 8: 'd'}[sz]
2336
+ buf[off:off + sz] = struct.pack('<' + fmt, float(val))
2337
+ elif k == 'c':
2338
+ half = sz // 2
2339
+ fmt = {4: 'f', 8: 'd'}[half]
2340
+ c = complex(val)
2341
+ buf[off:off + half] = struct.pack('<' + fmt, c.real)
2342
+ buf[off + half:off + sz] = struct.pack('<' + fmt, c.imag)
2343
+ elif k in ('S', 'V'):
2344
+ bs = bytes(val) if isinstance(val, (bytes, bytearray)) else (
2345
+ val.tobytes() if isinstance(val, void) else str(val).encode('latin-1'))
2346
+ buf[off:off + sz] = bs[:sz].ljust(sz, b'\x00')
2347
+ elif k == 'U':
2348
+ s = val if isinstance(val, str) else (
2349
+ bytes(val).decode('latin-1') if isinstance(val, (bytes, bytearray)) else str(val))
2350
+ buf[off:off + sz] = s.encode('utf-32-le')[:sz].ljust(sz, b'\x00')
2351
+ else:
2352
+ buf[off:off + sz] = b'\x00' * sz
2353
+
2354
+
2355
+ def _pack_field(buf, off, val, fdt):
2356
+ if fdt._names is not None:
2357
+ for i, n in enumerate(fdt._names):
2358
+ _pack_field(buf, off + fdt._fields[n][1], val[i], fdt._fields[n][0])
2359
+ elif fdt._subdtype is not None:
2360
+ base, shape = fdt._subdtype
2361
+ flat = _flatten_data(val)
2362
+ esize = base._itemsize
2363
+ for j, e in enumerate(flat):
2364
+ _pack_field(buf, off + j * esize, e, base)
2365
+ else:
2366
+ _pack_scalar(buf, off, val, fdt)
2367
+
2368
+
2369
+ def _pack_record(rec, dt):
2370
+ buf = bytearray(dt._itemsize)
2371
+ for i, n in enumerate(dt._names):
2372
+ _pack_field(buf, dt._fields[n][1], rec[i], dt._fields[n][0])
2373
+ return bytes(buf)
2374
+
2375
+
2376
+ def _unpack_scalar(buf, off, fdt):
2377
+ import struct
2378
+ k = fdt._kind
2379
+ sz = fdt._itemsize
2380
+ seg = bytes(buf[off:off + sz])
2381
+ if k in ('i', 'u', 'b'):
2382
+ v = int.from_bytes(seg, 'little', signed=(k == 'i'))
2383
+ return bool(v) if k == 'b' else v
2384
+ if k == 'f':
2385
+ return struct.unpack('<' + {2: 'e', 4: 'f', 8: 'd'}[sz], seg)[0]
2386
+ if k == 'c':
2387
+ half = sz // 2
2388
+ fmt = {4: 'f', 8: 'd'}[half]
2389
+ return complex(struct.unpack('<' + fmt, seg[:half])[0],
2390
+ struct.unpack('<' + fmt, seg[half:])[0])
2391
+ if k == 'S':
2392
+ return seg.rstrip(b'\x00')
2393
+ if k == 'V':
2394
+ return void(seg)
2395
+ if k == 'U':
2396
+ return seg.decode('utf-32-le').rstrip('\x00')
2397
+ return 0
2398
+
2399
+
2400
+ def _unpack_field(buf, off, fdt):
2401
+ if fdt._names is not None:
2402
+ return tuple(_unpack_field(buf, off + fdt._fields[n][1], fdt._fields[n][0])
2403
+ for n in fdt._names)
2404
+ if fdt._subdtype is not None:
2405
+ base, shape = fdt._subdtype
2406
+ total = 1
2407
+ for s in shape:
2408
+ total *= s
2409
+ esize = base._itemsize
2410
+ flat = [_unpack_field(buf, off + j * esize, base) for j in range(total)]
2411
+ return _reshape_flat(flat, shape)
2412
+ return _unpack_scalar(buf, off, fdt)
2413
+
2414
+
2415
+ def _unpack_record(buf, off, dt):
2416
+ return tuple(_unpack_field(buf, off + dt._fields[n][1], dt._fields[n][0])
2417
+ for n in dt._names)
2418
+
2419
+
2420
+ def _struct_to_bytes(arr):
2421
+ dt = _struct_dt_of(arr)
2422
+ recs = _flatten_records(arr._raw_data)
2423
+ return b''.join(_pack_record(r, dt) for r in recs), len(recs)
2424
+
2425
+
2426
+ def _view_dtype(arr, target):
2427
+ """按目标 dtype 重解释数组字节。"""
2428
+ tdt = target if isinstance(target, DType) else dtype(target)
2429
+ src_dt = getattr(arr, '_dtype_obj', None)
2430
+ if src_dt is not None:
2431
+ raw, _ = _struct_to_bytes(arr)
2432
+ ssize = src_dt._itemsize
2433
+ else:
2434
+ raw = arr.tobytes()
2435
+ ssize = dtype(getattr(arr, '_dtype', 'float64'))._itemsize
2436
+ tsize = tdt._itemsize
2437
+ if tsize == 0:
2438
+ raise ValueError("cannot view with zero-width dtype")
2439
+ count = len(raw) // tsize
2440
+ # rsnumpy 语义:view 仅改变最后一个轴的长度,其余轴保持不变。
2441
+ src_shape = tuple(arr.shape)
2442
+ if not src_shape:
2443
+ new_shape = (count,)
2444
+ else:
2445
+ last_bytes = src_shape[-1] * ssize
2446
+ if last_bytes % tsize != 0:
2447
+ raise ValueError(
2448
+ "When changing to a larger dtype, its size must be a divisor "
2449
+ "of the total size in bytes of the last axis of the array.")
2450
+ new_shape = src_shape[:-1] + (last_bytes // tsize,)
2451
+ if tdt._names is not None:
2452
+ records = [_unpack_record(raw, i * tsize, tdt) for i in range(count)]
2453
+ return _wrap_structured(tdt, records)
2454
+ vals = [_unpack_field(raw, i * tsize, tdt) for i in range(count)]
2455
+ return _build_plain_array(vals, tdt, new_shape)
2456
+
2457
+
2458
+ # ---------- 比较 ----------
2459
+
2460
+ def _leaf_eq(x, y):
2461
+ if isinstance(x, void) or isinstance(y, void):
2462
+ return bytes(x.tobytes() if isinstance(x, void) else x) == \
2463
+ bytes(y.tobytes() if isinstance(y, void) else y)
2464
+ return x == y
2465
+
2466
+
2467
+ def _struct_eq(a, b):
2468
+ ra = _flatten_records(a._raw_data)
2469
+ rb = _flatten_records(b._raw_data)
2470
+ out = []
2471
+ for x, y in zip(ra, rb):
2472
+ out.append(_py_all(_leaf_eq(x[i], y[i]) for i in range(_py_min(len(x), len(y)))))
2473
+ res = _reshape_flat(out, a.shape)
2474
+ return ndarray(res, _dtype='bool')
2475
+
2476
+
2477
+ # ---------- 格式化 ----------
2478
+
2479
+ def _fmt_scalar_value(v, kind):
2480
+ if kind == 'b':
2481
+ return 'True' if v else 'False'
2482
+ if kind in ('i', 'u'):
2483
+ return str(int(v))
2484
+ if kind == 'f':
2485
+ return _trim_positional_float(float(v))
2486
+ if kind == 'c':
2487
+ return _format_complex_scalar(complex(v))
2488
+ if kind == 'S':
2489
+ b = bytes(v) if isinstance(v, (bytes, bytearray)) else str(v).encode('latin-1')
2490
+ return repr(b)
2491
+ if kind == 'U':
2492
+ return repr(v if isinstance(v, str) else str(v))
2493
+ if kind == 'V':
2494
+ return _void_repr(v.tobytes() if isinstance(v, void) else bytes(v))
2495
+ return repr(v)
2496
+
2497
+
2498
+ def _fmt_sublist(v, base):
2499
+ if isinstance(v, (list, tuple)):
2500
+ return "[" + ", ".join(_fmt_sublist(x, base) for x in v) + "]"
2501
+ return _fmt_scalar_value(v, base._kind)
2502
+
2503
+
2504
+ def _fmt_field(v, fdt):
2505
+ if fdt._names is not None:
2506
+ return _fmt_record(v, fdt)
2507
+ if fdt._subdtype is not None:
2508
+ base, shape = fdt._subdtype
2509
+ return _fmt_sublist(v, base)
2510
+ return _fmt_scalar_value(v, fdt._kind)
2511
+
2512
+
2513
+ def _fmt_record(rec, dt):
2514
+ parts = [_fmt_field(rec[i], dt._fields[n][0]) for i, n in enumerate(dt._names)]
2515
+ if len(parts) == 1:
2516
+ return "(" + parts[0] + ",)"
2517
+ return "(" + ", ".join(parts) + ")"
2518
+
2519
+
2520
+ class _Record:
2521
+ """结构化数组的记录标量(对应 rsnumpy.void 记录),支持字段名/整数索引读写。"""
2522
+
2523
+ def __init__(self, values, dt, parent=None, flat_index=None):
2524
+ self._values = list(values)
2525
+ self._dtype_obj = dt
2526
+ self._parent = parent
2527
+ self._flat_index = flat_index
2528
+
2529
+ @property
2530
+ def dtype(self):
2531
+ return self._dtype_obj
2532
+
2533
+ def __getitem__(self, key):
2534
+ dt = self._dtype_obj
2535
+ if isinstance(key, str):
2536
+ i = dt._names.index(key)
2537
+ return _field_leaf_to_scalar(self._values[i], dt._fields[key][0])
2538
+ if isinstance(key, int):
2539
+ n = dt._names[key]
2540
+ return _field_leaf_to_scalar(self._values[key], dt._fields[n][0])
2541
+ raise KeyError(key)
2542
+
2543
+ def __setitem__(self, key, value):
2544
+ dt = self._dtype_obj
2545
+ if isinstance(key, str):
2546
+ i = dt._names.index(key)
2547
+ fdt = dt._fields[key][0]
2548
+ elif isinstance(key, int):
2549
+ i = key
2550
+ fdt = dt._fields[dt._names[key]][0]
2551
+ else:
2552
+ raise KeyError(key)
2553
+ self._values[i] = _coerce_field_value(value, fdt)
2554
+ if self._parent is not None and self._flat_index is not None:
2555
+ _set_flat_record(self._parent, self._flat_index, tuple(self._values))
2556
+
2557
+ def __getattr__(self, name):
2558
+ try:
2559
+ dt = self.__dict__['_dtype_obj']
2560
+ except KeyError:
2561
+ raise AttributeError(name)
2562
+ if dt is not None and dt._names and name in dt._names:
2563
+ i = dt._names.index(name)
2564
+ return _field_leaf_to_scalar(self.__dict__['_values'][i], dt._fields[name][0])
2565
+ raise AttributeError(name)
2566
+
2567
+ def item(self):
2568
+ return _record_to_pytuple(tuple(self._values), self._dtype_obj)
2569
+
2570
+ def tolist(self):
2571
+ return self.item()
2572
+
2573
+ def __len__(self):
2574
+ return len(self._dtype_obj._names)
2575
+
2576
+ def __eq__(self, other):
2577
+ if isinstance(other, _Record):
2578
+ return tuple(self._values) == tuple(other._values)
2579
+ return NotImplemented
2580
+
2581
+ def __repr__(self):
2582
+ return _fmt_record(tuple(self._values), self._dtype_obj)
2583
+
2584
+ def __str__(self):
2585
+ return _fmt_record(tuple(self._values), self._dtype_obj)
2586
+
2587
+
2588
+ def _record_to_pytuple(rec, dt):
2589
+ out = []
2590
+ for i, n in enumerate(dt._names):
2591
+ fdt = dt._fields[n][0]
2592
+ v = rec[i]
2593
+ if fdt._names is not None:
2594
+ out.append(_record_to_pytuple(v, fdt))
2595
+ else:
2596
+ out.append(v)
2597
+ return tuple(out)
2598
+
2599
+
2600
+ # ========== 标量类型别名 ==========
2601
+ # 与 rsnumpy 一致的标量类型层次结构,使 issubdtype / isinstance 检查按预期工作。
2602
+
2603
+ class generic:
2604
+ """所有 rsnumpy 标量类型的抽象基类。"""
2605
+
2606
+
2607
+ class number(generic):
2608
+ """数值标量的抽象基类。"""
2609
+
2610
+
2611
+ class integer(number):
2612
+ """整数标量的抽象基类。"""
2613
+
2614
+
2615
+ class signedinteger(integer):
2616
+ """有符号整数标量的抽象基类。"""
2617
+
2618
+
2619
+ class unsignedinteger(integer):
2620
+ """无符号整数标量的抽象基类。"""
2621
+
2622
+
2623
+ class inexact(number):
2624
+ """浮点/复数标量的抽象基类。"""
2625
+
2626
+
2627
+ class floating(inexact):
2628
+ """浮点标量的抽象基类。"""
2629
+
2630
+
2631
+ class complexfloating(inexact):
2632
+ """复数标量的抽象基类。"""
2633
+
2634
+
2635
+ class flexible(generic):
2636
+ """可变长度标量(字符串/void)的抽象基类。"""
2637
+
2638
+
2639
+ class character(flexible):
2640
+ """字符标量的抽象基类。"""
2641
+
2642
+
2643
+ int8 = type('int8', (signedinteger,), {})
2644
+ int16 = type('int16', (signedinteger,), {})
2645
+ int32 = type('int32', (signedinteger,), {})
2646
+ int64 = type('int64', (signedinteger,), {})
2647
+ uint8 = type('uint8', (unsignedinteger,), {})
2648
+ uint16 = type('uint16', (unsignedinteger,), {})
2649
+ uint32 = type('uint32', (unsignedinteger,), {})
2650
+ uint64 = type('uint64', (unsignedinteger,), {})
2651
+ float16 = type('float16', (floating,), {})
2652
+ float32 = type('float32', (floating,), {})
2653
+ float64 = type('float64', (floating,), {})
2654
+ float_ = float64 # rsnumpy 别名
2655
+ complex64 = type('complex64', (complexfloating,), {})
2656
+ complex128 = type('complex128', (complexfloating,), {})
2657
+ string_ = type('string_', (character,), {})
2658
+ unicode_ = type('unicode_', (character,), {})
2659
+ bytes_ = string_ # rsnumpy 中 bytes_ 为字节字符串类型
2660
+ str_ = unicode_ # rsnumpy 中 str_ 为 unicode 字符串类型
2661
+ bool_ = type('bool', (generic,), {})
2662
+ object_ = type('object_', (generic,), {})
2663
+
2664
+
2665
+ def _void_repr(bs):
2666
+ """void 标量的显示:每字节 \\xHH(大写十六进制)。"""
2667
+ return "b'" + ''.join('\\x%02X' % b for b in bs) + "'"
2668
+
2669
+
2670
+ class void(flexible):
2671
+ """void 标量:存储原始字节,支持按字节相等比较(兼容 rsnumpy.void)。"""
2672
+
2673
+ def __init__(self, data):
2674
+ if isinstance(data, (bytes, bytearray)):
2675
+ self._bytes = bytes(data)
2676
+ elif isinstance(data, int):
2677
+ self._bytes = b'\x00' * data
2678
+ else:
2679
+ self._bytes = bytes(data)
2680
+
2681
+ def tobytes(self):
2682
+ return self._bytes
2683
+
2684
+ def __eq__(self, other):
2685
+ if isinstance(other, void):
2686
+ return self._bytes == other._bytes
2687
+ if isinstance(other, (bytes, bytearray)):
2688
+ return self._bytes == bytes(other)
2689
+ return NotImplemented
2690
+
2691
+ def __ne__(self, other):
2692
+ result = self.__eq__(other)
2693
+ if result is NotImplemented:
2694
+ return result
2695
+ return not result
2696
+
2697
+ def __hash__(self):
2698
+ return hash(self._bytes)
2699
+
2700
+ def __repr__(self):
2701
+ return _void_repr(self._bytes)
2702
+
2703
+ def __str__(self):
2704
+ return _void_repr(self._bytes)
2705
+
2706
+
2707
+ # C 语言宽度别名(与 rsnumpy 在 64 位平台上的取值一致)
2708
+ byte = int8
2709
+ short = int16
2710
+ intc = int32
2711
+ int_ = int64
2712
+ long = int64
2713
+ longlong = int64
2714
+ intp = int64
2715
+ ubyte = uint8
2716
+ ushort = uint16
2717
+ uintc = uint32
2718
+ uint = uint64
2719
+ ulong = uint64
2720
+ ulonglong = uint64
2721
+ uintp = uint64
2722
+ half = float16
2723
+ single = float32
2724
+ double = float64
2725
+ longdouble = float64
2726
+ csingle = complex64
2727
+ cdouble = complex128
2728
+ clongdouble = complex128
2729
+
2730
+ # 布尔常量与其他特殊常量
2731
+ True_ = True
2732
+ False_ = False
2733
+ little_endian = (_sys.byteorder == 'little')
2734
+
2735
+ # 标量类型名 → 类型对象 映射(供 issubdtype / result_type 等使用)
2736
+ _SCTYPE_BY_NAME = {
2737
+ 'bool': bool_,
2738
+ 'int8': int8, 'int16': int16, 'int32': int32, 'int64': int64,
2739
+ 'uint8': uint8, 'uint16': uint16, 'uint32': uint32, 'uint64': uint64,
2740
+ 'float16': float16, 'float32': float32, 'float64': float64,
2741
+ 'complex64': complex64, 'complex128': complex128,
2742
+ 'string_': string_, 'unicode_': unicode_, 'bytes': bytes_,
2743
+ 'str': str_, 'object': object_, 'void': void,
2744
+ }
2745
+
2746
+ # rsnumpy 兼容的 sctypeDict / ScalarType / typecodes
2747
+ ScalarType = (int, float, complex, bool, bytes, str,
2748
+ int8, int16, int32, int64, uint8, uint16, uint32, uint64,
2749
+ float16, float32, float64, complex64, complex128)
2750
+
2751
+ sctypeDict = {
2752
+ 'bool': bool_, 'int8': int8, 'int16': int16, 'int32': int32, 'int64': int64,
2753
+ 'uint8': uint8, 'uint16': uint16, 'uint32': uint32, 'uint64': uint64,
2754
+ 'float16': float16, 'float32': float32, 'float64': float64,
2755
+ 'complex64': complex64, 'complex128': complex128,
2756
+ 'i1': int8, 'i2': int16, 'i4': int32, 'i8': int64,
2757
+ 'u1': uint8, 'u2': uint16, 'u4': uint32, 'u8': uint64,
2758
+ 'f2': float16, 'f4': float32, 'f8': float64,
2759
+ 'c8': complex64, 'c16': complex128, '?': bool_,
2760
+ }
2761
+
2762
+ typecodes = {
2763
+ 'Character': 'c',
2764
+ 'Integer': 'bhilqp',
2765
+ 'UnsignedInteger': 'BHILQP',
2766
+ 'Float': 'efdg',
2767
+ 'Complex': 'FDG',
2768
+ 'AllInteger': 'bBhHiIlLqQpP',
2769
+ 'AllFloat': 'efdgFDG',
2770
+ 'Datetime': 'Mm',
2771
+ 'All': '?bhilqpBHILQPefdgFDGSUVOMm',
2772
+ }
2773
+
2774
+
2775
+ def _to_scalar_type(x):
2776
+ """将 dtype 类字符串 / 类型 / 数组解析为标量类型对象。"""
2777
+ if isinstance(x, type) and issubclass(x, generic):
2778
+ return x
2779
+ if hasattr(x, '_dtype'):
2780
+ return _SCTYPE_BY_NAME.get(x._dtype, generic)
2781
+ name = _resolve_type_name(x)
2782
+ return _SCTYPE_BY_NAME.get(name, generic)
2783
+
2784
+
2785
+ def issubdtype(arg1, arg2):
2786
+ """判断第一个类型是否为第二个类型(含抽象类别)的子类型。"""
2787
+ c1 = _to_scalar_type(arg1)
2788
+ c2 = arg2 if (isinstance(arg2, type) and issubclass(arg2, generic)) else _to_scalar_type(arg2)
2789
+ return issubclass(c1, c2)
2790
+
2791
+
2792
+ class finfo:
2793
+ """浮点类型的机器精度信息(兼容 rsnumpy.finfo)。"""
2794
+
2795
+ _DATA = {
2796
+ 'float16': dict(bits=16, eps=0.0009765625, epsneg=0.00048828125,
2797
+ max=65504.0, min=-65504.0, tiny=6.103515625e-05,
2798
+ resolution=0.001, nmant=10, nexp=5, precision=3),
2799
+ 'float32': dict(bits=32, eps=1.1920929e-07, epsneg=5.9604645e-08,
2800
+ max=3.4028235e+38, min=-3.4028235e+38, tiny=1.1754944e-38,
2801
+ resolution=1e-06, nmant=23, nexp=8, precision=6),
2802
+ 'float64': dict(bits=64, eps=2.220446049250313e-16, epsneg=1.1102230246251565e-16,
2803
+ max=1.7976931348623157e+308, min=-1.7976931348623157e+308,
2804
+ tiny=2.2250738585072014e-308, resolution=1e-15,
2805
+ nmant=52, nexp=11, precision=15),
2806
+ }
2807
+
2808
+ def __init__(self, dtype):
2809
+ name = _resolve_type_name(dtype)
2810
+ if name in ('complex64',):
2811
+ name = 'float32'
2812
+ elif name in ('complex128',):
2813
+ name = 'float64'
2814
+ if name not in self._DATA:
2815
+ raise ValueError(f"data type {dtype!r} not inexact")
2816
+ self._name = name
2817
+ d = self._DATA[name]
2818
+ self.bits = d['bits']
2819
+ self.eps = d['eps']
2820
+ self.epsneg = d['epsneg']
2821
+ self.max = d['max']
2822
+ self.min = d['min']
2823
+ self.tiny = d['tiny']
2824
+ self.smallest_normal = d['tiny']
2825
+ self.resolution = d['resolution']
2826
+ self.nmant = d['nmant']
2827
+ self.nexp = d['nexp']
2828
+ self.precision = d['precision']
2829
+
2830
+ @property
2831
+ def dtype(self):
2832
+ return dtype(self._name)
2833
+
2834
+ def __repr__(self):
2835
+ return (f"finfo(resolution={self.resolution}, min={self.min}, "
2836
+ f"max={self.max}, dtype={self._name})")
2837
+
2838
+
2839
+ class iinfo:
2840
+ """整数类型的取值范围信息(兼容 rsnumpy.iinfo)。"""
2841
+
2842
+ _BITS = {
2843
+ 'int8': (8, True), 'int16': (16, True), 'int32': (32, True), 'int64': (64, True),
2844
+ 'uint8': (8, False), 'uint16': (16, False), 'uint32': (32, False), 'uint64': (64, False),
2845
+ }
2846
+
2847
+ def __init__(self, int_type):
2848
+ name = _resolve_type_name(int_type)
2849
+ if name not in self._BITS:
2850
+ raise ValueError(f"Invalid integer data type {int_type!r}.")
2851
+ self._name = name
2852
+ self.bits, signed = self._BITS[name]
2853
+ self.kind = 'i' if signed else 'u'
2854
+ if signed:
2855
+ self.min = -(1 << (self.bits - 1))
2856
+ self.max = (1 << (self.bits - 1)) - 1
2857
+ else:
2858
+ self.min = 0
2859
+ self.max = (1 << self.bits) - 1
2860
+
2861
+ @property
2862
+ def dtype(self):
2863
+ return dtype(self._name)
2864
+
2865
+ def __repr__(self):
2866
+ return f"iinfo(min={self.min}, max={self.max}, dtype={self._name})"
2867
+
2868
+ def __str__(self):
2869
+ dashes = '-' * 63
2870
+ return (
2871
+ f"Machine parameters for {self._name}\n"
2872
+ f"{dashes}\n"
2873
+ f"min = {self.min}\n"
2874
+ f"max = {self.max}\n"
2875
+ f"{dashes}\n"
2876
+ )
2877
+
2878
+
2879
+ class ndindex:
2880
+ """按 C 顺序遍历给定形状的多维索引(兼容 rsnumpy.ndindex)。"""
2881
+
2882
+ def __init__(self, *shape):
2883
+ if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
2884
+ shape = tuple(shape[0])
2885
+ self._shape = tuple(int(s) for s in shape)
2886
+
2887
+ def __iter__(self):
2888
+ import itertools as _it
2889
+ if not self._shape:
2890
+ yield ()
2891
+ return
2892
+ yield from _it.product(*[range(s) for s in self._shape])
2893
+
2894
+
2895
+ class ndenumerate:
2896
+ """同时给出多维索引与元素值的迭代器(兼容 rsnumpy.ndenumerate)。"""
2897
+
2898
+ def __init__(self, arr):
2899
+ self._arr = asarray(arr)
2900
+
2901
+ def __iter__(self):
2902
+ shape = self._arr.shape
2903
+ flat = _flatten_to_list(self._arr.tolist())
2904
+ import itertools as _it
2905
+ coords = _it.product(*[range(s) for s in shape]) if shape else [()]
2906
+ for c, v in zip(coords, flat):
2907
+ yield c, v
2908
+
2909
+
2910
+ def _flatten_to_list(data):
2911
+ out = []
2912
+ stack = [data]
2913
+ while stack:
2914
+ v = stack.pop()
2915
+ if isinstance(v, list):
2916
+ stack.extend(reversed(v))
2917
+ else:
2918
+ out.append(v)
2919
+ return out
2920
+
2921
+
2922
+ class index_exp_class:
2923
+ """构造索引元组:index_exp[1:3, 2] -> (slice(1,3), 2)。"""
2924
+
2925
+ def __getitem__(self, item):
2926
+ if isinstance(item, tuple):
2927
+ return item
2928
+ return (item,)
2929
+
2930
+
2931
+ index_exp = index_exp_class()
2932
+
2933
+
2934
+ # ========== 日期时间类型 ==========
2935
+ # 内部统一以“自 1970-01-01 (UTC) 起的天数”存储,与 matplotlib/rsplotlib 的日期约定
2936
+ # 一致,从而可直接参与绘图并被 ConciseDateFormatter 正确格式化为刻度标签。
2937
+
2938
+ _DT_EPOCH = _datetime.datetime(1970, 1, 1)
2939
+
2940
+ # 时间单位 → 天数换算(Y/M 因日历长度可变,取近似值)
2941
+ _TD_UNIT_DAYS = {
2942
+ 'W': 7.0,
2943
+ 'D': 1.0,
2944
+ 'h': 1.0 / 24.0,
2945
+ 'm': 1.0 / 1440.0,
2946
+ 's': 1.0 / 86400.0,
2947
+ 'ms': 1.0 / 86400.0e3,
2948
+ 'us': 1.0 / 86400.0e6,
2949
+ 'ns': 1.0 / 86400.0e9,
2950
+ 'Y': 365.0,
2951
+ 'M': 30.0,
2952
+ }
2953
+
2954
+
2955
+ def _parse_datetime_string(s, unit=None):
2956
+ """解析 ISO 日期字符串为 (自纪元起天数, 推断出的单位)。"""
2957
+ s = s.strip()
2958
+ date_part, time_part = s, None
2959
+ if 'T' in s:
2960
+ date_part, time_part = s.split('T', 1)
2961
+ elif ':' in s and ' ' in s:
2962
+ date_part, time_part = s.split(' ', 1)
2963
+ ymd = date_part.split('-')
2964
+ year = int(ymd[0])
2965
+ month = int(ymd[1]) if len(ymd) > 1 else 1
2966
+ day = int(ymd[2]) if len(ymd) > 2 else 1
2967
+ resolved = 'D' if len(ymd) >= 3 else ('M' if len(ymd) == 2 else 'Y')
2968
+ hour = minute = second = micro = 0
2969
+ if time_part:
2970
+ tparts = time_part.split(':')
2971
+ hour = int(tparts[0])
2972
+ resolved = 'h'
2973
+ if len(tparts) > 1:
2974
+ minute = int(tparts[1])
2975
+ resolved = 'm'
2976
+ if len(tparts) > 2:
2977
+ sec = tparts[2]
2978
+ if '.' in sec:
2979
+ sec_i, frac = sec.split('.', 1)
2980
+ second = int(sec_i)
2981
+ micro = int(_py_round(float('0.' + frac) * 1e6))
2982
+ resolved = 'us'
2983
+ else:
2984
+ second = int(sec)
2985
+ resolved = 's'
2986
+ dt = _datetime.datetime(year, month, day, hour, minute, second, micro)
2987
+ days = (dt - _DT_EPOCH).total_seconds() / 86400.0
2988
+ return days, (unit or resolved)
2989
+
2990
+
2991
+ def _parse_datetime_to_days(value, unit=None):
2992
+ """将 datetime64 的各种输入统一解析为 (自纪元起天数, 单位)。"""
2993
+ if isinstance(value, datetime64):
2994
+ return value._days, (unit or value._unit)
2995
+ if isinstance(value, _datetime.datetime):
2996
+ return (value - _DT_EPOCH).total_seconds() / 86400.0, (unit or 'us')
2997
+ if isinstance(value, _datetime.date):
2998
+ dt = _datetime.datetime(value.year, value.month, value.day)
2999
+ return (dt - _DT_EPOCH).total_seconds() / 86400.0, (unit or 'D')
3000
+ if isinstance(value, str):
3001
+ return _parse_datetime_string(value, unit)
3002
+ # 数值 + 单位:表示纪元之后的偏移量
3003
+ u = unit or 'us'
3004
+ return float(value) * _TD_UNIT_DAYS.get(u, _TD_UNIT_DAYS['us']), u
3005
+
3006
+
3007
+ class timedelta64:
3008
+ """时间间隔标量,内部以天数存储,兼容 rsnumpy.timedelta64 的常用构造与运算。"""
3009
+
3010
+ def __init__(self, value=0, unit=None):
3011
+ if isinstance(value, timedelta64):
3012
+ self._days = value._days
3013
+ self._unit = unit or value._unit
3014
+ return
3015
+ u = unit or 'us'
3016
+ self._unit = u
3017
+ self._days = float(value) * _TD_UNIT_DAYS.get(u, _TD_UNIT_DAYS['us'])
3018
+
3019
+ @classmethod
3020
+ def _from_days(cls, days, unit='us'):
3021
+ obj = cls.__new__(cls)
3022
+ obj._days = float(days)
3023
+ obj._unit = unit
3024
+ return obj
3025
+
3026
+ def _value_in_unit(self):
3027
+ return self._days / _TD_UNIT_DAYS.get(self._unit, 1.0)
3028
+
3029
+ def __float__(self):
3030
+ return self._days
3031
+
3032
+ def __int__(self):
3033
+ return int(self._value_in_unit())
3034
+
3035
+ def __add__(self, other):
3036
+ if isinstance(other, datetime64):
3037
+ return datetime64._from_days(self._days + other._days, other._unit)
3038
+ if isinstance(other, timedelta64):
3039
+ return timedelta64._from_days(self._days + other._days, self._unit)
3040
+ return timedelta64._from_days(self._days + float(other), self._unit)
3041
+ __radd__ = __add__
3042
+
3043
+ def __sub__(self, other):
3044
+ d = other._days if isinstance(other, timedelta64) else float(other)
3045
+ return timedelta64._from_days(self._days - d, self._unit)
3046
+
3047
+ def __mul__(self, other):
3048
+ return timedelta64._from_days(self._days * float(other), self._unit)
3049
+ __rmul__ = __mul__
3050
+
3051
+ def __truediv__(self, other):
3052
+ if isinstance(other, timedelta64):
3053
+ return self._days / other._days
3054
+ return timedelta64._from_days(self._days / float(other), self._unit)
3055
+
3056
+ def __eq__(self, other):
3057
+ if isinstance(other, timedelta64):
3058
+ return self._days == other._days
3059
+ return NotImplemented
3060
+
3061
+ def __hash__(self):
3062
+ return hash(self._days)
3063
+
3064
+ def __lt__(self, other):
3065
+ return self._days < float(other)
3066
+
3067
+ def __le__(self, other):
3068
+ return self._days <= float(other)
3069
+
3070
+ def __gt__(self, other):
3071
+ return self._days > float(other)
3072
+
3073
+ def __ge__(self, other):
3074
+ return self._days >= float(other)
3075
+
3076
+ def __repr__(self):
3077
+ v = self._value_in_unit()
3078
+ vi = int(_py_round(v))
3079
+ vout = vi if abs(v - vi) < 1e-9 else v
3080
+ return "numpy.timedelta64(%r,'%s')" % (vout, self._unit)
3081
+
3082
+
3083
+ class datetime64:
3084
+ """日期时间标量,内部以“自 1970-01-01 起的天数”存储,兼容 rsnumpy.datetime64 的常用用法。"""
3085
+
3086
+ def __init__(self, value=None, unit=None):
3087
+ if value is None:
3088
+ self._days = 0.0
3089
+ self._unit = unit or 'D'
3090
+ return
3091
+ self._days, self._unit = _parse_datetime_to_days(value, unit)
3092
+
3093
+ @classmethod
3094
+ def _from_days(cls, days, unit='D'):
3095
+ obj = cls.__new__(cls)
3096
+ obj._days = float(days)
3097
+ obj._unit = unit
3098
+ return obj
3099
+
3100
+ def to_datetime(self):
3101
+ return _DT_EPOCH + _datetime.timedelta(days=self._days)
3102
+
3103
+ def __float__(self):
3104
+ return self._days
3105
+
3106
+ def __int__(self):
3107
+ return int(self._days)
3108
+
3109
+ def __add__(self, other):
3110
+ if isinstance(other, timedelta64):
3111
+ return datetime64._from_days(self._days + other._days, self._unit)
3112
+ return datetime64._from_days(self._days + float(other), self._unit)
3113
+ __radd__ = __add__
3114
+
3115
+ def __sub__(self, other):
3116
+ if isinstance(other, datetime64):
3117
+ return timedelta64._from_days(self._days - other._days, self._unit)
3118
+ if isinstance(other, timedelta64):
3119
+ return datetime64._from_days(self._days - other._days, self._unit)
3120
+ return datetime64._from_days(self._days - float(other), self._unit)
3121
+
3122
+ def __eq__(self, other):
3123
+ if isinstance(other, datetime64):
3124
+ return self._days == other._days
3125
+ return NotImplemented
3126
+
3127
+ def __hash__(self):
3128
+ return hash(self._days)
3129
+
3130
+ def __lt__(self, other):
3131
+ return self._days < float(other)
3132
+
3133
+ def __le__(self, other):
3134
+ return self._days <= float(other)
3135
+
3136
+ def __gt__(self, other):
3137
+ return self._days > float(other)
3138
+
3139
+ def __ge__(self, other):
3140
+ return self._days >= float(other)
3141
+
3142
+ def __repr__(self):
3143
+ dt = self.to_datetime()
3144
+ u = self._unit
3145
+ if u == 'Y':
3146
+ s = dt.strftime('%Y')
3147
+ elif u == 'M':
3148
+ s = dt.strftime('%Y-%m')
3149
+ elif u == 'D':
3150
+ s = dt.strftime('%Y-%m-%d')
3151
+ elif u in ('h', 'm'):
3152
+ s = dt.strftime('%Y-%m-%dT%H:%M')
3153
+ elif u == 's':
3154
+ s = dt.strftime('%Y-%m-%dT%H:%M:%S')
3155
+ else:
3156
+ s = dt.isoformat()
3157
+ return "numpy.datetime64('%s')" % s
3158
+
3159
+
3160
+ # ========== 构造/工厂函数 ==========
3161
+
3162
+ # np.build_array 返回的 dtype 编码 → rsnumpy dtype 名称(与 _infer_int_dtype 一致)
3163
+ _BUILD_ARRAY_DTYPES = ('float64', 'int64', 'bool')
3164
+
3165
+
3166
+ def _flatten_check(data):
3167
+ """单次遍历:展平数据并检查是否包含复数/字符串。"""
3168
+ flat = []
3169
+ has_complex = False
3170
+ has_string = False
3171
+ stack = [data]
3172
+ while stack:
3173
+ item = stack.pop()
3174
+ if isinstance(item, (list, tuple)):
3175
+ stack.extend(reversed(item))
3176
+ else:
3177
+ flat.append(item)
3178
+ if isinstance(item, complex):
3179
+ has_complex = True
3180
+ elif isinstance(item, str):
3181
+ has_string = True
3182
+ return flat, has_complex, has_string
3183
+
3184
+
3185
+ def _flatten_data(data):
3186
+ """展平嵌套的可迭代对象为扁平列表。"""
3187
+ flat, _, _ = _flatten_check(data)
3188
+ return flat
3189
+
3190
+
3191
+ def _setitem_value(value):
3192
+ """规范化赋值右值供 Rust setitem_multi 使用:
3193
+ 标量原样返回;ndarray 或嵌套列表展平为 C 序浮点列表(逐元素赋值)。
3194
+ 复数右值原样透传,交由 Rust coerce_value_to_nd 保留虚部。"""
3195
+ if _is_ndarray(value):
3196
+ rust = value._array
3197
+ if getattr(rust, 'is_complex', False):
3198
+ return rust
3199
+ return [float(v) for v in _flatten_data(rust.tolist())]
3200
+ if value.__class__.__name__ == 'ndarray' and hasattr(value, 'tolist'):
3201
+ if getattr(value, 'is_complex', False):
3202
+ return value
3203
+ return [float(v) for v in _flatten_data(value.tolist())]
3204
+ if isinstance(value, complex):
3205
+ return value
3206
+ if isinstance(value, (list, tuple)):
3207
+ flat = _flatten_data(value)
3208
+ if _py_any(isinstance(v, complex) for v in flat):
3209
+ return value
3210
+ return [float(v) for v in flat]
3211
+ return value
3212
+
3213
+
3214
+ _INTEGER_DTYPES = frozenset((
3215
+ "int8", "int16", "int32", "int64",
3216
+ "uint8", "uint16", "uint32", "uint64",
3217
+ ))
3218
+
3219
+
3220
+ def _cast_setitem_value(value, dtype):
3221
+ """按目标数组 dtype 转换赋值右值:整数类型向零截断,bool 类型布尔化,
3222
+ 与 rsnumpy 赋值时的隐式类型转换一致(底层仍存 f64)。"""
3223
+ if dtype in _INTEGER_DTYPES:
3224
+ if isinstance(value, list):
3225
+ return [float(int(v)) for v in value]
3226
+ return float(int(value))
3227
+ if dtype == "bool":
3228
+ if isinstance(value, list):
3229
+ return [1.0 if v else 0.0 for v in value]
3230
+ return 1.0 if value else 0.0
3231
+ return value
3232
+
3233
+
3234
+ # 数组接口协议 typestr 的 (kind, itemsize) → rsnumpy dtype 名称
3235
+ _ARRAY_INTERFACE_DTYPE = {
3236
+ ('b', 1): 'bool',
3237
+ ('u', 1): 'uint8', ('i', 1): 'int8',
3238
+ ('u', 2): 'uint16', ('i', 2): 'int16',
3239
+ ('u', 4): 'uint32', ('i', 4): 'int32',
3240
+ ('u', 8): 'uint64', ('i', 8): 'int64',
3241
+ ('f', 4): 'float32', ('f', 8): 'float64',
3242
+ }
3243
+
3244
+ # rsnumpy dtype 名称 → 数组接口 typestr(供 __array_interface__ 暴露给真实 rsnumpy)
3245
+ _DTYPE_TO_TYPESTR = {
3246
+ 'bool': '|b1',
3247
+ 'uint8': '|u1', 'int8': '|i1',
3248
+ 'uint16': '<u2', 'int16': '<i2',
3249
+ 'uint32': '<u4', 'int32': '<i4',
3250
+ 'uint64': '<u8', 'int64': '<i8',
3251
+ 'float32': '<f4', 'float64': '<f8',
3252
+ }
3253
+
3254
+
3255
+ def _from_array_interface(obj):
3256
+ """将实现数组接口协议的对象(如 PIL 图像)解码为 (Rust 数组, dtype)。
3257
+
3258
+ 缓冲区的字节解码由 Rust 底层 `np.from_buffer_typed` 完成。
3259
+ """
3260
+ ai = obj.__array_interface__
3261
+ shape = tuple(ai['shape'])
3262
+ typestr = ai['typestr']
3263
+ data = ai['data']
3264
+ if isinstance(data, tuple):
3265
+ # (指针, 只读标志) 形式,无法直接读取 → 退回到 tobytes()
3266
+ data = obj.tobytes()
3267
+ elif not isinstance(data, (bytes, bytearray)):
3268
+ data = bytes(data)
3269
+ kind, itemsize = typestr[1], int(typestr[2:])
3270
+ dtype = _ARRAY_INTERFACE_DTYPE.get((kind, itemsize))
3271
+ if dtype is None:
3272
+ raise TypeError(f"Unsupported array interface type: {typestr}")
3273
+ arr = np.from_buffer_typed(bytes(data), typestr, list(shape))
3274
+ return arr, dtype
3275
+
3276
+
3277
+ def _is_rectangular(data):
3278
+ """检查嵌套列表是否为矩形(所有子列表长度相同)。"""
3279
+ if not isinstance(data, (list, tuple)):
3280
+ return True
3281
+ if not data:
3282
+ return True
3283
+ first_len = None
3284
+ for item in data:
3285
+ if isinstance(item, (list, tuple)):
3286
+ if first_len is None:
3287
+ first_len = len(item)
3288
+ elif len(item) != first_len:
3289
+ return False
3290
+ else:
3291
+ if first_len is not None:
3292
+ return False
3293
+ return True
3294
+
3295
+
3296
+ def array(data, dtype=None, copy=True, order='K', subok=False, ndmin=0):
3297
+ """创建数组。"""
3298
+ if dtype is None:
3299
+ # 数值 list/tuple 快速路径:Rust 单次完成展平 + dtype 推断 + 构造。
3300
+ if isinstance(data, (list, tuple)):
3301
+ try:
3302
+ raw, code = np.build_array(data)
3303
+ except (ValueError, TypeError):
3304
+ raw = None
3305
+ if raw is not None:
3306
+ _dtype = _BUILD_ARRAY_DTYPES[code]
3307
+ arr = ndarray._wrap(raw, _dtype=_dtype)
3308
+ if ndmin > arr.ndim:
3309
+ new_shape = (1,) * (ndmin - arr.ndim) + arr.shape
3310
+ arr = ndarray._wrap(arr._array.reshape(new_shape), _dtype=_dtype)
3311
+ return arr
3312
+ # 回退慢路径:复数/字符串/不规则/标量
3313
+ if isinstance(data, tuple):
3314
+ data = list(data)
3315
+ flat, has_c, has_s = _flatten_check(data)
3316
+ if has_c:
3317
+ _dtype = "complex128"
3318
+ else:
3319
+ _dtype = _infer_int_dtype(flat)
3320
+ else:
3321
+ _dtype = _resolve_dtype(dtype)
3322
+ # 字符串/字节/void dtype(S/U/V):存为原始 Python 数据,避免数值化。
3323
+ if dtype is not None:
3324
+ sd = _parse_str_dtype(dtype)
3325
+ if sd is not None:
3326
+ flat = _flatten_data(data) if isinstance(data, (list, tuple)) else [data]
3327
+ arr = ndarray._wrap(np.zeros((len(flat),)), _dtype='string_', _raw_data=list(flat))
3328
+ arr._str_dtype = sd
3329
+ return arr
3330
+ # 结构化 dtype:统一走富 DType 模型
3331
+ _dt_obj = _as_struct_dtype(dtype)
3332
+ if _dt_obj is not None:
3333
+ return _make_structured_array(data, _dt_obj)
3334
+ arr = ndarray(data, _dtype=_dtype)
3335
+ # 显式复数 dtype:即使输入已是实数 ndarray(__init__ 会沿用其 dtype),
3336
+ # 也需提升为原生复数(零虚部)。
3337
+ if _dtype in ('complex128', 'complex64') and not getattr(arr._array, 'is_complex', False):
3338
+ arr = ndarray._wrap(_maybe_native_complex(arr._array, _dtype), _dtype=_dtype)
3339
+ if ndmin > arr.ndim:
3340
+ new_shape = (1,) * (ndmin - arr.ndim) + arr.shape
3341
+ arr = ndarray._wrap(arr._array.reshape(new_shape), _dtype=_dtype)
3342
+ return arr
3343
+
3344
+
3345
+ def asarray(a, dtype=None, order=None):
3346
+ """转换输入为数组。"""
3347
+ if _is_ndarray(a):
3348
+ return a
3349
+ if dtype is not None:
3350
+ if isinstance(dtype, str):
3351
+ dt_str = dtype
3352
+ elif hasattr(dtype, '__name__'):
3353
+ dt_str = dtype.__name__
3354
+ elif hasattr(dtype, 'name'):
3355
+ dt_str = dtype.name
3356
+ else:
3357
+ dt_str = str(dtype)
3358
+ if dt_str in ("complex", "complex128", "complex64", "cfloat", "cdouble"):
3359
+ _dtype = "complex128"
3360
+ else:
3361
+ _dtype = _resolve_dtype(dtype)
3362
+ else:
3363
+ # 数值 list/tuple 快速路径:Rust 单次完成展平 + dtype 推断 + 构造。
3364
+ if isinstance(a, (list, tuple)):
3365
+ try:
3366
+ raw, code = np.build_array(a)
3367
+ except (ValueError, TypeError):
3368
+ raw = None
3369
+ if raw is not None:
3370
+ return ndarray._wrap(raw, _dtype=_BUILD_ARRAY_DTYPES[code])
3371
+ flat, _, _ = _flatten_check(a)
3372
+ _dtype = _infer_int_dtype(flat)
3373
+ else:
3374
+ _dtype = "float64"
3375
+ return ndarray(a, _dtype=_dtype)
3376
+
3377
+
3378
+ def asanyarray(a, dtype=None, order=None):
3379
+ """转换输入为 ndarray。"""
3380
+ return asarray(a, dtype, order)
3381
+
3382
+
3383
+ def matrix(data, dtype=None, copy=True):
3384
+ """从数组或字符串创建矩阵对象。"""
3385
+ if isinstance(data, str):
3386
+ rows = data.split(';')
3387
+ data = []
3388
+ for row in rows:
3389
+ elements = row.split(',')
3390
+ data.append([float(e.strip()) for e in elements])
3391
+ return asarray(data, dtype=dtype)
3392
+
3393
+
3394
+ def asmatrix(data, dtype=None):
3395
+ """将输入转换为矩阵对象。"""
3396
+ return asarray(data, dtype=dtype)
3397
+
3398
+
3399
+ def bmat(obj, ldict=None, gdict=None):
3400
+ """从字符串、嵌套序列或数组构建矩阵。"""
3401
+ if isinstance(obj, str):
3402
+ obj = obj.strip()
3403
+ if obj.startswith('[') and obj.endswith(']'):
3404
+ obj = obj[1:-1].strip()
3405
+ rows = obj.split(';')
3406
+ data = []
3407
+ for row in rows:
3408
+ elements = row.split(',')
3409
+ row_data = []
3410
+ for e in elements:
3411
+ e = e.strip()
3412
+ if e.startswith('[') and e.endswith(']'):
3413
+ inner = e[1:-1].split()
3414
+ row_data.extend([float(x) for x in inner])
3415
+ else:
3416
+ row_data.append(float(e))
3417
+ data.append(row_data)
3418
+ return asarray(data)
3419
+ if isinstance(obj, (list, tuple)):
3420
+ if isinstance(obj[0], (list, tuple)):
3421
+ if isinstance(obj[0][0], (list, tuple)):
3422
+ rows = []
3423
+ for row_block in obj:
3424
+ row_len = sum(len(block[0]) for block in row_block)
3425
+ # current_row = [0.0] * row_len
3426
+ col_start = 0
3427
+ for block in row_block:
3428
+ b_arr = _to_ndarray(block)
3429
+ b_rows, b_cols = b_arr.shape
3430
+ for i in range(b_rows):
3431
+ if i >= len(rows):
3432
+ rows.append([0.0] * row_len)
3433
+ for j in range(b_cols):
3434
+ rows[i][col_start + j] = b_arr[i, j]
3435
+ col_start += b_cols
3436
+ return asarray(rows)
3437
+ else:
3438
+ return asarray(obj)
3439
+ else:
3440
+ return asarray(obj)
3441
+ return _to_ndarray(obj)
3442
+
3443
+
3444
+ def repmat(a, m, n):
3445
+ """将数组或矩阵重复 m 行 n 列。"""
3446
+ a = _to_ndarray(a)
3447
+ if a.ndim == 1:
3448
+ a = a.reshape(1, a.size)
3449
+ return tile(a, (m, n))
3450
+
3451
+
3452
+ def _to_ndarray(obj):
3453
+ """将任意数组类对象转换为 rsnumpy ndarray(不依赖第三方库)。"""
3454
+ if _is_ndarray(obj):
3455
+ return obj
3456
+ if hasattr(obj, 'tolist'):
3457
+ return ndarray(obj.tolist())
3458
+ return ndarray(obj)
3459
+
3460
+
3461
+ def array_equal(a, b):
3462
+ """判断两个数组是否形状和元素完全相同。"""
3463
+ a = _to_ndarray(a)
3464
+ b = _to_ndarray(b)
3465
+ if a.shape != b.shape:
3466
+ return False
3467
+ return bool(np.array_equal(a._array, b._array))
3468
+
3469
+
3470
+ def copy(a, order='K'):
3471
+ """返回数组的副本。"""
3472
+ return ndarray(a).copy(order=order)
3473
+
3474
+
3475
+ def _parse_str_dtype(dtype):
3476
+ """解析字符串类 dtype,返回 (kind, width);非字符串 dtype 返回 None。
3477
+
3478
+ kind: 'S'(字节串,1 字节/元素)、'U'(Unicode,4 字节/元素)、
3479
+ 'V'(原始字节)。width 为元素长度(字符数),未指定为 0。
3480
+ """
3481
+ if dtype is bytes:
3482
+ return ('S', 0)
3483
+ if dtype is str:
3484
+ return ('U', 0)
3485
+ if not isinstance(dtype, str):
3486
+ return None
3487
+ s = dtype.lstrip('<>=|')
3488
+ if not s:
3489
+ return None
3490
+ k = s[0]
3491
+ if k in ('S', 'U', 'V', 'a'):
3492
+ kind = 'S' if k == 'a' else k
3493
+ rest = s[1:]
3494
+ width = int(rest) if rest.isdigit() else 0
3495
+ return (kind, width)
3496
+ return None
3497
+
3498
+
3499
+ def _string_array_tobytes(items, str_dtype):
3500
+ """将字符串/字节元素列表按 (kind, width) 编码为定宽补零字节串。"""
3501
+ kind, width = str_dtype
3502
+ out = bytearray()
3503
+ for v in items:
3504
+ if kind == 'U':
3505
+ if isinstance(v, (bytes, bytearray)):
3506
+ v = bytes(v).decode('latin-1')
3507
+ enc = str(v).encode('utf-32-le')
3508
+ size = width * 4 if width else len(enc)
3509
+ else:
3510
+ if isinstance(v, (bytes, bytearray)):
3511
+ enc = bytes(v)
3512
+ else:
3513
+ enc = str(v).encode('latin-1')
3514
+ size = width if width else len(enc)
3515
+ enc = enc[:size].ljust(size, b'\x00')
3516
+ out += enc
3517
+ return bytes(out)
3518
+
3519
+
3520
+ def _str_scalar(v, str_dtype):
3521
+ """按 (kind, width) 将单个原始元素转换为标量:
3522
+ S→截断到 width 并去尾零的 bytes;U→去尾零的 str;V→定宽补零的 void。
3523
+ """
3524
+ kind, width = str_dtype
3525
+ if kind == 'U':
3526
+ s = bytes(v).decode('latin-1') if isinstance(v, (bytes, bytearray)) else str(v)
3527
+ if width:
3528
+ s = s[:width]
3529
+ return s.rstrip('\x00')
3530
+ bs = bytes(v) if isinstance(v, (bytes, bytearray)) else str(v).encode('latin-1')
3531
+ if kind == 'V':
3532
+ if width:
3533
+ bs = bs[:width].ljust(width, b'\x00')
3534
+ return void(bs)
3535
+ # kind == 'S'
3536
+ if width:
3537
+ bs = bs[:width]
3538
+ return bs.rstrip(b'\x00')
3539
+
3540
+
3541
+ _TOBYTES_STRUCT = {
3542
+ 'float64': 'd', 'float32': 'f',
3543
+ 'int64': 'q', 'int32': 'i', 'int16': 'h', 'int8': 'b',
3544
+ 'uint64': 'Q', 'uint32': 'I', 'uint16': 'H', 'uint8': 'B',
3545
+ 'bool': '?',
3546
+ }
3547
+
3548
+
3549
+ def _numeric_tobytes(arr, order='C'):
3550
+ """将数值数组按 dtype 小端布局编码为字节串。"""
3551
+ import struct
3552
+ dt = getattr(arr, '_dtype', 'float64')
3553
+ # 快路径:float64 数组的字节布局与底层存储一致,直接从缓冲协议 memcpy,
3554
+ # 避免逐元素 tolist()→float()→struct.pack 的 Python 开销(百万点级差异达百毫秒)。
3555
+ if (dt == 'float64' and getattr(arr, '_raw_data', None) is None and getattr(arr, '_complex_data', None) is None):
3556
+ return memoryview(arr._array).tobytes(order)
3557
+ fmt = _TOBYTES_STRUCT.get(dt, 'd')
3558
+ flat = _flatten_data(arr.tolist())
3559
+ if fmt == '?':
3560
+ return bytes(1 if v else 0 for v in flat)
3561
+ if fmt in ('d', 'f'):
3562
+ vals = [float(v) for v in flat]
3563
+ else:
3564
+ vals = [int(v) for v in flat]
3565
+ return struct.pack('<%d%s' % (len(vals), fmt), *vals)
3566
+
3567
+
3568
+ def _resolve_dtype(dtype):
3569
+ """解析 dtype 字符串/Types/DType 为内部表示。"""
3570
+ if dtype is None:
3571
+ return "float64"
3572
+ if isinstance(dtype, DType):
3573
+ return dtype.name
3574
+ if isinstance(dtype, (list, tuple)):
3575
+ return "void"
3576
+ dt_str = dtype if isinstance(dtype, str) else dtype.__name__
3577
+ if dt_str in ("complex", "complex128", "complex64", "cfloat", "cdouble"):
3578
+ return "complex128"
3579
+ if dt_str in ("int", "int_", "intp", "intc"):
3580
+ return "int64"
3581
+ if dt_str in ("int8", "int16", "int32", "int64"):
3582
+ return dt_str
3583
+ if dt_str in ("uint", "uint8", "uint16", "uint32", "uint64"):
3584
+ return dt_str
3585
+ if dt_str in ("bool", "bool_"):
3586
+ return "bool"
3587
+ return "float64"
3588
+
3589
+
3590
+ def _infer_int_dtype(args):
3591
+ """推断参数中的类型:全 bool→bool,全 int→int64,否则 float64。"""
3592
+ has_bool = False
3593
+ has_int = False
3594
+ for a in args:
3595
+ if isinstance(a, bool):
3596
+ has_bool = True
3597
+ elif isinstance(a, int):
3598
+ has_int = True
3599
+ else:
3600
+ return "float64"
3601
+ if has_bool and not has_int:
3602
+ return "bool"
3603
+ return "int64"
3604
+
3605
+
3606
+ def _make_structured_zeros(shape, fields):
3607
+ """为结构化 dtype 生成零填充的原始数据(子数组字段填充嵌套零)。"""
3608
+ elem = tuple(_nested_zeros(_field_subshape(f)) for f in fields)
3609
+ if isinstance(shape, int):
3610
+ total = shape
3611
+ else:
3612
+ total = 1
3613
+ for s in shape:
3614
+ total *= s
3615
+ flat = [elem] * total
3616
+ if isinstance(shape, int):
3617
+ shape = (shape,)
3618
+ if len(shape) <= 1:
3619
+ return flat
3620
+ # 重塑为多维
3621
+ result = flat
3622
+ for dim in reversed(shape[1:]):
3623
+ result = [result[i:i+dim] for i in range(0, len(result), dim)]
3624
+ return result
3625
+
3626
+
3627
+ def _maybe_native_complex(raw, _dtype):
3628
+ """dtype 为复数时,把实数 Rust 数组提升为原生复数(零虚部)。"""
3629
+ if _dtype in ('complex128', 'complex64'):
3630
+ return raw + 0j
3631
+ return raw
3632
+
3633
+
3634
+ def zeros(shape, dtype=None, order='C'):
3635
+ """返回指定形状的零数组。"""
3636
+ dt_obj = _as_struct_dtype(dtype)
3637
+ if dt_obj is not None:
3638
+ return _wrap_structured(dt_obj, _make_struct_filled(shape, dt_obj, _zero_scalar_for))
3639
+ _dtype = _resolve_dtype(dtype)
3640
+ return ndarray(_maybe_native_complex(np.zeros(shape), _dtype), _dtype=_dtype)
3641
+
3642
+
3643
+ def ones(shape, dtype=None, order='C'):
3644
+ """返回指定形状的1数组。"""
3645
+ dt_obj = _as_struct_dtype(dtype)
3646
+ if dt_obj is not None:
3647
+ return _wrap_structured(dt_obj, _make_struct_filled(shape, dt_obj, _one_scalar_for))
3648
+ _dtype = _resolve_dtype(dtype)
3649
+ return ndarray(_maybe_native_complex(np.ones(shape), _dtype), _dtype=_dtype)
3650
+
3651
+
3652
+ def empty(shape, dtype=None, order='C'):
3653
+ """返回指定形状的空数组。"""
3654
+ dt_obj = _as_struct_dtype(dtype)
3655
+ if dt_obj is not None:
3656
+ return _wrap_structured(dt_obj, _make_struct_filled(shape, dt_obj, _zero_scalar_for))
3657
+ _dtype = _resolve_dtype(dtype)
3658
+ arr = ndarray(_maybe_native_complex(np.empty(shape), _dtype), _dtype=_dtype)
3659
+ arr._is_empty = True
3660
+ return arr
3661
+
3662
+
3663
+ def full(shape, fill_value, dtype=None, order='C'):
3664
+ """返回指定形状的填充数组。"""
3665
+ _dtype = _resolve_dtype(dtype)
3666
+ if isinstance(fill_value, complex) or _dtype in ('complex128', 'complex64'):
3667
+ c = complex(fill_value)
3668
+ raw = np.full(shape, c.real) + complex(0.0, c.imag)
3669
+ return ndarray(raw, _dtype='complex128' if _dtype not in ('complex128', 'complex64') else _dtype)
3670
+ return ndarray(np.full(shape, fill_value), _dtype=_dtype)
3671
+
3672
+
3673
+ def zeros_like(a, dtype=None, order='K', subok=True, shape=None):
3674
+ """返回与输入形状相同的零数组。"""
3675
+ arr = ndarray(a)
3676
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
3677
+ return ndarray(np.zeros_like(arr._array), _dtype=_dtype)
3678
+
3679
+
3680
+ def ones_like(a, dtype=None, order='K', subok=True, shape=None):
3681
+ """返回与输入形状相同的1数组。"""
3682
+ arr = ndarray(a)
3683
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
3684
+ return ndarray(np.ones_like(arr._array), _dtype=_dtype)
3685
+
3686
+
3687
+ def empty_like(a, dtype=None, order='K', subok=True, shape=None):
3688
+ """返回与输入形状相同的空数组。"""
3689
+ arr = ndarray(a)
3690
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
3691
+ return ndarray(np.empty_like(arr._array), _dtype=_dtype)
3692
+
3693
+
3694
+ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
3695
+ """返回与输入形状相同的填充数组。"""
3696
+ arr = ndarray(a)
3697
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
3698
+ return ndarray(np.full_like(arr._array, float(fill_value)), _dtype=_dtype)
3699
+
3700
+
3701
+ def eye(N, M=None, k=0, dtype=None, order='C'):
3702
+ """返回对角线为1的二维数组。"""
3703
+ _dtype = _resolve_dtype(dtype)
3704
+ return ndarray(np.eye(N, M, k), _dtype=_dtype)
3705
+
3706
+
3707
+ def identity(n, dtype=None):
3708
+ """返回单位矩阵。"""
3709
+ return eye(n, dtype=dtype)
3710
+
3711
+
3712
+ def arange(start=0, stop=None, step=1, dtype=None):
3713
+ """返回给定间隔内均匀间隔的值。"""
3714
+ if stop is None:
3715
+ stop = start
3716
+ start = 0
3717
+ if isinstance(start, datetime64) or isinstance(stop, datetime64) \
3718
+ or isinstance(step, timedelta64):
3719
+ return _datetime_arange(start, stop, step)
3720
+ _dtype = _resolve_dtype(dtype) if dtype is not None else _infer_int_dtype((start, stop, step))
3721
+ return ndarray(np.arange(start, stop, step), _dtype=_dtype)
3722
+
3723
+
3724
+ def _datetime_arange(start, stop, step):
3725
+ """datetime64/timedelta64 版本的 arange:以天数(自纪元起)生成浮点数组。
3726
+
3727
+ 产出的数组值遵循 matplotlib/rsplotlib 的日期数值约定,可直接绘图并被
3728
+ ConciseDateFormatter 正确格式化。
3729
+ """
3730
+ if isinstance(step, timedelta64):
3731
+ unit = step._unit
3732
+ elif isinstance(start, datetime64):
3733
+ unit = start._unit
3734
+ elif isinstance(stop, datetime64):
3735
+ unit = stop._unit
3736
+ else:
3737
+ unit = 'D'
3738
+ raw = np.arange(float(start), float(stop), float(step))
3739
+ return ndarray(raw, _dtype='datetime64[%s]' % unit)
3740
+
3741
+
3742
+ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
3743
+ """返回指定间隔内均匀间隔的数字。"""
3744
+ start_val = _scalar(_ensure(start))
3745
+ stop_val = _scalar(_ensure(stop))
3746
+ _dtype = _resolve_dtype(dtype)
3747
+ result = ndarray(np.linspace(start_val, stop_val, num, endpoint=endpoint), _dtype=_dtype)
3748
+ if retstep:
3749
+ step = (stop_val - start_val) / (num - 1) if endpoint else (stop_val - start_val) / num
3750
+ return result, _float64(step)
3751
+ return result
3752
+
3753
+
3754
+ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
3755
+ """返回对数刻度上均匀间隔的数字。"""
3756
+ _dtype = _resolve_dtype(dtype) if dtype else "float64"
3757
+ return ndarray(np.logspace(float(start), float(stop), num, float(base)), _dtype=_dtype)
3758
+
3759
+
3760
+ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
3761
+ """返回几何级数上均匀间隔的数字。"""
3762
+ if start <= 0 or stop <= 0:
3763
+ raise ValueError("geomspace requires positive start and stop values")
3764
+ _dtype = _resolve_dtype(dtype) if dtype else "float64"
3765
+ return ndarray(np.geomspace(float(start), float(stop), num), _dtype=_dtype)
3766
+
3767
+
3768
+ # ========== 索引函数 ==========
3769
+
3770
+ def ix_(*args):
3771
+ """从多个序列构造开放网格,用于笛卡尔积花式索引。
3772
+
3773
+ 返回 ndarray 元组,每个数组沿对应维度广播。
3774
+
3775
+ 示例:
3776
+ >>> x = np.arange(32).reshape((8, 4))
3777
+ >>> x[np.ix_([1,5,7,2], [0,3,1,2])]
3778
+ """
3779
+ # Rust 层的 ix_rs 只解析 Python list/tuple,故先把 ndarray 等参数
3780
+ # 归一化为一维下标列表;布尔序列按 rsnumpy 语义取 True 位置的下标。
3781
+ prepared = []
3782
+ for a in args:
3783
+ dtype = getattr(a, '_dtype', None)
3784
+ seq = a.tolist() if hasattr(a, 'tolist') else list(a)
3785
+ if dtype == 'bool' or (len(seq) > 0 and _py_all(isinstance(v, bool) for v in seq)):
3786
+ seq = [i for i, v in enumerate(seq) if v]
3787
+ else:
3788
+ seq = [int(v) for v in seq]
3789
+ prepared.append(seq)
3790
+ raw = np.ix_rs(tuple(prepared))
3791
+ return tuple(ndarray._wrap(r) for r in raw)
3792
+
3793
+
3794
+ def where(condition, x=None, y=None):
3795
+ """根据条件返回元素或索引。"""
3796
+ if x is not None and y is not None:
3797
+ dtype = getattr(x, '_dtype', 'float64')
3798
+ fields = getattr(x, '_fields', None)
3799
+ raw_data = getattr(x, '_raw_data', None)
3800
+ result = np.where(_ensure(condition), _ensure(x), _ensure(y))
3801
+ return ndarray._wrap(result, _dtype=dtype, _fields=fields, _raw_data=raw_data)
3802
+ return nonzero(condition)
3803
+
3804
+
3805
+ def take(a, indices, axis=None, out=None, mode='raise'):
3806
+ """沿指定轴取元素。"""
3807
+ arr = ndarray(a)
3808
+ return arr.take(indices, axis)
3809
+
3810
+
3811
+ def put(a, indices, values, mode='raise'):
3812
+ """设置指定位置的元素。"""
3813
+ arr = ndarray(a)
3814
+ arr.put(indices, values)
3815
+
3816
+
3817
+ def select(condlist, choicelist, default=0):
3818
+ """根据条件列表选择值(Rust 单次遍历实现)。"""
3819
+ cond_arrays = [_ensure(c) for c in condlist]
3820
+ choice_arrays = [_ensure(c) for c in choicelist]
3821
+ result = np.select_rs(cond_arrays, choice_arrays, float(default))
3822
+ return ndarray(result)
3823
+
3824
+
3825
+ def nonzero(a):
3826
+ """返回非零元素的索引,返回元组形式的数组(兼容NumPy)。"""
3827
+ raw = np.nonzero_arrs(_ensure(a))
3828
+ return tuple(ndarray._wrap(arr, _dtype="int64") for arr in raw)
3829
+
3830
+
3831
+ def argwhere(a):
3832
+ """查找非零元素的索引。"""
3833
+ return ndarray(np.argwhere(_ensure(a)))
3834
+
3835
+
3836
+ def flatnonzero(a):
3837
+ """返回扁平化数组中非零元素的索引。"""
3838
+ return ndarray(np.flatnonzero(_ensure(a)))
3839
+
3840
+
3841
+ # ========== 缺失构造函数 ==========
3842
+
3843
+ def fromfunction(function, shape, *, dtype=None, **kwargs):
3844
+ """根据函数和形状创建数组。"""
3845
+ indices = np._arange_arrays(shape)
3846
+ grid = np.meshgrid(*indices, indexing='ij')
3847
+ return function(*grid, **kwargs)
3848
+
3849
+
3850
+ def frombuffer(buffer, dtype=None, count=-1, offset=0, *, like=None):
3851
+ """从缓冲区创建一维数组(使用 Rust 层实现)。"""
3852
+ if isinstance(buffer, bytes):
3853
+ if offset > 0:
3854
+ buffer = buffer[offset:]
3855
+ if isinstance(dtype, str) and len(dtype) > 1 and dtype[0] == 'S' and dtype[1:].isdigit():
3856
+ # Byte string type (e.g. S1, S2, ...)
3857
+ elem_size = int(dtype[1:])
3858
+ total = len(buffer) // elem_size
3859
+ if count >= 0:
3860
+ total = min(total, count)
3861
+ raw = [buffer[i * elem_size:(i + 1) * elem_size] for i in range(total)]
3862
+ return ndarray(np.zeros((len(raw),)), _dtype=dtype, _raw_data=raw)
3863
+ return ndarray(np.bytes_to_floats(buffer, count))
3864
+ return ndarray(list(buffer))
3865
+
3866
+
3867
+ def fromiter(iterable, dtype=None, count=-1, *, like=None):
3868
+ """从可迭代对象创建一维数组。"""
3869
+ _dtype = _resolve_dtype(dtype)
3870
+ if count >= 0:
3871
+ data = [next(iterable) for _ in range(count)]
3872
+ else:
3873
+ data = list(iterable)
3874
+ return ndarray(data, _dtype=_dtype)
3875
+
3876
+
3877
+ class _RClass:
3878
+ """行连接辅助类,模拟 np.r_。"""
3879
+ def __getitem__(self, item):
3880
+ if not isinstance(item, tuple):
3881
+ item = (item,)
3882
+ arrays = []
3883
+ for it in item:
3884
+ arrays.append(ndarray(it))
3885
+ if len(arrays) == 0:
3886
+ return ndarray([])
3887
+ return concatenate(arrays)
3888
+
3889
+
3890
+ r_ = _RClass()
3891
+
3892
+
3893
+ class _CClass:
3894
+ """列连接辅助类,模拟 np.c_。"""
3895
+ def __getitem__(self, item):
3896
+ if not isinstance(item, tuple):
3897
+ item = (item,)
3898
+ arrays = [ndarray(it) for it in item]
3899
+ if _py_all(a.ndim <= 1 for a in arrays):
3900
+ return column_stack(arrays)
3901
+ return concatenate(arrays, axis=1)
3902
+
3903
+
3904
+ c_ = _CClass()
3905
+
3906
+
3907
+ class _SClass:
3908
+ """切片辅助类,模拟 np.s_。"""
3909
+ def __getitem__(self, item):
3910
+ return item
3911
+
3912
+
3913
+ s_ = _SClass()
3914
+
3915
+
3916
+ def mgrid(*ranges):
3917
+ """密集网格索引,返回密集网格数组。"""
3918
+ arrays = []
3919
+ for r in ranges:
3920
+ if isinstance(r, slice):
3921
+ arrays.append(ndarray(arange(r.start or 0, r.stop, r.step or 1)))
3922
+ elif isinstance(r, (int, float)):
3923
+ arrays.append(ndarray([float(r)]))
3924
+ else:
3925
+ arrays.append(ndarray(r))
3926
+ grids = np.meshgrid(*arrays, indexing='ij')
3927
+ if isinstance(grids, tuple) and len(grids) > 1:
3928
+ return stack(grids, axis=0)
3929
+ return grids
3930
+
3931
+
3932
+ def ogrid(*ranges):
3933
+ """开放网格索引,返回开放网格数组。"""
3934
+ arrays = []
3935
+ for r in ranges:
3936
+ if isinstance(r, slice):
3937
+ arrays.append(ndarray(arange(r.start or 0, r.stop, r.step or 1)))
3938
+ elif isinstance(r, (int, float)):
3939
+ arrays.append(ndarray([float(r)]))
3940
+ else:
3941
+ arrays.append(ndarray(r))
3942
+ return np.meshgrid(*arrays, indexing='ij')
3943
+
3944
+
3945
+ def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
3946
+ """从坐标向量返回坐标矩阵。
3947
+
3948
+ 默认使用 'xy' 索引(与 rsnumpy 一致),返回一维输入数组两两组合的网格。
3949
+ """
3950
+ arrays = []
3951
+ for x in xi:
3952
+ if hasattr(x, 'tolist'):
3953
+ arrays.append([int(v) for v in x.tolist()])
3954
+ elif isinstance(x, (list, tuple)):
3955
+ arrays.append([int(v) for v in x])
3956
+ else:
3957
+ arrays.append([int(x)])
3958
+
3959
+ n = len(arrays)
3960
+ if n == 0:
3961
+ return []
3962
+
3963
+ shapes = [len(a) for a in arrays]
3964
+
3965
+ if sparse:
3966
+ result = []
3967
+ for i in range(n):
3968
+ if i == 0:
3969
+ grid = ndarray([[arrays[i][j]] for j in range(shapes[i])])
3970
+ else:
3971
+ grid = ndarray([arrays[i]])
3972
+ result.append(grid)
3973
+ return result
3974
+
3975
+ result = []
3976
+
3977
+ for i in range(n):
3978
+ if indexing == 'xy' and n == 2:
3979
+ if i == 0:
3980
+ grid = ndarray([[arrays[0][k] for k in range(shapes[0])] for _ in range(shapes[1])])
3981
+ else:
3982
+ grid = ndarray([[arrays[1][j] for _ in range(shapes[0])] for j in range(shapes[1])])
3983
+ elif indexing == 'ij':
3984
+ if n == 2:
3985
+ if i == 0:
3986
+ grid = ndarray([[arrays[0][j] for _ in range(shapes[1])] for j in range(shapes[0])])
3987
+ else:
3988
+ grid = ndarray([[arrays[1][k] for k in range(shapes[1])] for _ in range(shapes[0])])
3989
+ elif n == 3:
3990
+ grid = ndarray([[[arrays[i][j] if i == 0 else (arrays[i][k] if i == 1 else arrays[i][l_])
3991
+ for l_ in range(shapes[2])]
3992
+ for k in range(shapes[1])]
3993
+ for j in range(shapes[0])])
3994
+ else:
3995
+ grid = ndarray([[arrays[i][j] for _ in range(shapes[1])] for j in range(shapes[0])])
3996
+ else:
3997
+ grid = ndarray([[arrays[i][j] for _ in range(shapes[1])] for j in range(shapes[0])])
3998
+ result.append(grid)
3999
+
4000
+ return result
4001
+
4002
+
4003
+ # ========== FFT ==========
4004
+ # FFT 相关函数均由子模块 `rsnumpy.fft` 提供(与 rsnumpy 一致,
4005
+ # 通过 `np.fft.fft(...)` 或 `from rsnumpy.fft import fft` 使用)。
4006
+ # 子模块的绑定见文件末尾的“子模块”一节。
4007
+
4008
+
4009
+ # ========== 常量 ==========
4010
+ pi = 3.141592653589793
4011
+ e = 2.718281828459045
4012
+ euler_gamma = 0.5772156649015328606
4013
+ inf = float('inf')
4014
+ nan = float('nan')
4015
+ newaxis = None
4016
+
4017
+
4018
+ # 字符串操作模块
4019
+ char = _char_module
4020
+
4021
+
4022
+ # ========== 判断函数 ==========
4023
+
4024
+ class _NdFlatIter:
4025
+ """ndarray.flat 的扁平迭代器,逐元素访问数组。"""
4026
+ def __init__(self, arr):
4027
+ self._arr = arr
4028
+ self._flat = arr._array.flatten()
4029
+ self._size = self._flat.size
4030
+ self._idx = 0
4031
+ self._is_int = getattr(arr, '_dtype', 'float64') == 'int64'
4032
+
4033
+ def __iter__(self):
4034
+ return self
4035
+
4036
+ def __next__(self):
4037
+ if self._idx >= self._size:
4038
+ raise StopIteration
4039
+ val = self._flat[self._idx]
4040
+ self._idx += 1
4041
+ return int(val) if self._is_int else float(val)
4042
+
4043
+
4044
+ class _NdIterElement:
4045
+ """nditer 的可写元素包装,支持 op_flags=['readwrite'] 时修改原数组。"""
4046
+ __slots__ = ('_arr', '_idx', '_val')
4047
+
4048
+ def __init__(self, arr, idx, val):
4049
+ self._arr = arr
4050
+ self._idx = idx
4051
+ self._val = val
4052
+
4053
+ def __repr__(self):
4054
+ return repr(self._val)
4055
+
4056
+ def __float__(self):
4057
+ return float(self._val)
4058
+
4059
+ def __int__(self):
4060
+ return int(self._val)
4061
+
4062
+ def __bool__(self):
4063
+ return bool(self._val)
4064
+
4065
+ def __eq__(self, other):
4066
+ return self._val == other
4067
+
4068
+ def __ne__(self, other):
4069
+ return self._val != other
4070
+
4071
+ def __lt__(self, other):
4072
+ return self._val < other
4073
+
4074
+ def __le__(self, other):
4075
+ return self._val <= other
4076
+
4077
+ def __gt__(self, other):
4078
+ return self._val > other
4079
+
4080
+ def __ge__(self, other):
4081
+ return self._val >= other
4082
+
4083
+ def __add__(self, other):
4084
+ return self._val + other
4085
+
4086
+ def __radd__(self, other):
4087
+ return other + self._val
4088
+
4089
+ def __sub__(self, other):
4090
+ return self._val - other
4091
+
4092
+ def __rsub__(self, other):
4093
+ return other - self._val
4094
+
4095
+ def __mul__(self, other):
4096
+ return self._val * other
4097
+
4098
+ def __rmul__(self, other):
4099
+ return other * self._val
4100
+
4101
+ def __truediv__(self, other):
4102
+ return self._val / other
4103
+
4104
+ def __rtruediv__(self, other):
4105
+ return other / self._val
4106
+
4107
+ def __neg__(self):
4108
+ return -self._val
4109
+
4110
+ def __abs__(self):
4111
+ return abs(self._val)
4112
+
4113
+ def __getitem__(self, key):
4114
+ return self._val
4115
+
4116
+ def __setitem__(self, key, value):
4117
+ self._val = float(value)
4118
+ self._arr.put([self._idx], [self._val])
4119
+
4120
+
4121
+ def nditer(a, order='C', op_flags=None, flags=None):
4122
+ """创建数组元素的迭代器,逐个访问数组元素。
4123
+
4124
+ 参数:
4125
+ a: 输入数组或数组列表(广播迭代)
4126
+ order: 遍历顺序,'C'(行序优先)或 'F'(列序优先)
4127
+ op_flags: 操作标志列表,如 ['readwrite'] 表示可读写
4128
+ flags: 迭代器标志,如 ['external_loop'] 表示返回一维数组块而非标量
4129
+
4130
+ >>> a = np.array([[0, 1, 2], [3, 4, 5]])
4131
+ >>> for x in np.nditer(a):
4132
+ ... print(x, end=", ")
4133
+ 0, 1, 2, 3, 4, 5,
4134
+ """
4135
+ # 多数组广播迭代
4136
+ if isinstance(a, (list, tuple)):
4137
+ arrays = [ndarray(x) for x in a]
4138
+ # 计算广播形状:每个维度取最大值
4139
+ all_shapes = [x.shape for x in arrays]
4140
+ max_ndim = np.imax([len(s) for s in all_shapes])
4141
+ padded = []
4142
+ for s in all_shapes:
4143
+ pad = [1] * (max_ndim - len(s)) + list(s)
4144
+ padded.append(pad)
4145
+ bcast_shape = tuple(np.imax([p[i] for p in padded]) for i in range(max_ndim))
4146
+ # 广播每个数组到统一形状
4147
+ bcast_arrays = []
4148
+ for arr in arrays:
4149
+ if arr.shape == bcast_shape:
4150
+ bcast_arrays.append(arr)
4151
+ else:
4152
+ bcast_arrays.append(ndarray(np.broadcast_to(arr._array, bcast_shape)))
4153
+ # 展平并迭代
4154
+ flats = [x._array.flatten().tolist() for x in bcast_arrays]
4155
+ is_ints = [getattr(x, '_dtype', 'float64') == 'int64' for x in bcast_arrays]
4156
+ for i in range(len(flats[0])):
4157
+ yield tuple(int(f[i]) if is_ints[j] else f[i] for j, f in enumerate(flats))
4158
+ return
4159
+
4160
+ arr = ndarray(a)
4161
+ ext_loop = flags and 'external_loop' in flags
4162
+
4163
+ # external_loop: 按 chunk 返回一维数组
4164
+ if ext_loop:
4165
+ if arr.ndim <= 1:
4166
+ yield arr
4167
+ return
4168
+ if order == 'F':
4169
+ # F-order: 每次返回一列(最后一维)
4170
+ ncols = arr.shape[-1]
4171
+ for col in range(ncols):
4172
+ key = tuple([slice(None)] * (arr.ndim - 1) + [col])
4173
+ col_arr = arr[key]
4174
+ yield ndarray(col_arr)
4175
+ else:
4176
+ # C-order: 每次返回一行(第一维)
4177
+ for i in range(arr.shape[0]):
4178
+ yield ndarray(arr[i])
4179
+ return
4180
+
4181
+ # 原有逐元素逻辑
4182
+ flat_arr = arr.copy().flatten()
4183
+ size = flat_arr.size
4184
+ is_int = getattr(arr, '_dtype', 'float64') == 'int64'
4185
+ readwrite = op_flags and 'readwrite' in op_flags
4186
+
4187
+ def _to_val(v, ci):
4188
+ if readwrite:
4189
+ return _NdIterElement(arr, ci, int(v) if is_int else v)
4190
+ return int(v) if is_int else v
4191
+
4192
+ if order == 'F' and arr.ndim > 1:
4193
+ # Fortran-order: 按列序优先(首维最快)生成多维索引,转为 C-flat 索引
4194
+ shape = arr.shape
4195
+ ndim = len(shape)
4196
+ # C-order strides: stride[i] = product(shape[i+1:])
4197
+ c_strides = [1] * ndim
4198
+ for i in range(ndim - 2, -1, -1):
4199
+ c_strides[i] = c_strides[i + 1] * shape[i + 1]
4200
+
4201
+ # F-order 等同于最后一维最慢,首维最快
4202
+ # 递归遍历 dim = ndim-1, ndim-2, ..., 0
4203
+ def _iter_f_order(dim, coords):
4204
+ if dim < 0:
4205
+ # coords 是反向顺序 [last_dim, ..., first_dim],需反转
4206
+ idx = np.isum([c * s for c, s in zip(reversed(coords), c_strides)])
4207
+ if 0 <= idx < size:
4208
+ val = flat_arr[idx].item()
4209
+ yield _to_val(val, idx)
4210
+ else:
4211
+ for v in range(shape[dim]):
4212
+ yield from _iter_f_order(dim - 1, coords + [v])
4213
+
4214
+ yield from _iter_f_order(ndim - 1, [])
4215
+ else:
4216
+ # C-order: 行序优先(默认)
4217
+ for i in range(size):
4218
+ val = flat_arr[i].item()
4219
+ yield _to_val(val, i)
4220
+
4221
+ def isnan(x):
4222
+ """逐元素检测是否为 NaN。"""
4223
+ arr = ndarray(x)
4224
+ return ndarray._wrap(np.isnan(arr._array), _dtype="bool")
4225
+
4226
+
4227
+ def binary_repr(num, width=None):
4228
+ """返回整数的二进制表示字符串。"""
4229
+ return np.binary_repr(num, width)
4230
+
4231
+
4232
+ def isinf(x):
4233
+ """逐元素检测是否为无穷大。"""
4234
+ arr = ndarray(x)
4235
+ return ndarray._wrap(np.isinf(arr._array), _dtype="bool")
4236
+
4237
+
4238
+ def isfinite(x):
4239
+ """逐元素检测是否为有限值。"""
4240
+ arr = ndarray(x)
4241
+ return ndarray._wrap(np.isfinite(arr._array), _dtype="bool")
4242
+
4243
+
4244
+ def iscomplex(x):
4245
+ """逐元素检测元素的虚部是否非零。"""
4246
+ arr = ndarray(x)
4247
+ cpx = getattr(arr, '_complex_data', None)
4248
+ if cpx is not None:
4249
+ return ndarray(np.iscomplex_cpx(cpx))
4250
+ # 不是复数数组 → 全零
4251
+ shape = arr.shape
4252
+ return ndarray(np.zeros(shape))
4253
+
4254
+
4255
+ def _broadcast_shape(*shapes):
4256
+ """计算多个数组的广播形状。"""
4257
+ max_ndim = np.imax([len(s) for s in shapes])
4258
+ result = []
4259
+ for i in range(max_ndim):
4260
+ dims = set()
4261
+ for s in shapes:
4262
+ if i < max_ndim - len(s):
4263
+ dims.add(1)
4264
+ else:
4265
+ dims.add(s[i - (max_ndim - len(s))])
4266
+ non_one = [d for d in dims if d != 1]
4267
+ if len(non_one) > 1:
4268
+ raise ValueError("operands could not be broadcast together")
4269
+ result.append(non_one[0] if non_one else 1)
4270
+ return tuple(result)
4271
+
4272
+
4273
+ def broadcast(*args):
4274
+ """产生模仿广播的对象,将多个数组广播到相同形状。"""
4275
+ if not args:
4276
+ raise ValueError("broadcast requires at least one argument")
4277
+ return _Broadcast(*args)
4278
+
4279
+
4280
+ class _Broadcast:
4281
+ """广播对象,将多个数组广播到相同形状。"""
4282
+
4283
+ def __init__(self, *args):
4284
+ self._arrays = [ndarray(a) for a in args]
4285
+ # 计算广播形状
4286
+ shapes = [a.shape for a in self._arrays]
4287
+ self._shape = _broadcast_shape(*shapes)
4288
+ self._size = 1
4289
+ for s in self._shape:
4290
+ self._size *= s
4291
+
4292
+ @property
4293
+ def shape(self):
4294
+ """广播后的形状。"""
4295
+ return self._shape
4296
+
4297
+ @property
4298
+ def iters(self):
4299
+ """返回各数组的迭代器元组。"""
4300
+ return tuple(_BroadcastIter(a, self._shape) for a in self._arrays)
4301
+
4302
+ def __iter__(self):
4303
+ """迭代广播后的值元组。"""
4304
+ iters = self.iters
4305
+ for _ in range(self._size):
4306
+ yield tuple(next(it) for it in iters)
4307
+
4308
+
4309
+ class _BroadcastIter:
4310
+ """广播迭代器,按广播形状遍历单个数组。"""
4311
+
4312
+ def __init__(self, arr, shape):
4313
+ self._arr = arr
4314
+ self._flat = arr._array.flatten().tolist()
4315
+ self._size = len(self._flat)
4316
+ self._arr_shape = arr.shape
4317
+ self._broadcast_shape = shape
4318
+ # Left-pad array shape for broadcasting
4319
+ ndim_pad = len(shape) - len(arr.shape)
4320
+ self._padded = (1,) * ndim_pad + arr.shape
4321
+ self._indices = [0] * len(shape)
4322
+ self._done = False
4323
+
4324
+ def __iter__(self):
4325
+ return self
4326
+
4327
+ def __next__(self):
4328
+ if self._done:
4329
+ raise StopIteration
4330
+ # Map broadcast index to flat array index
4331
+ arr_idx = 0
4332
+ stride = 1
4333
+ for d in range(len(self._arr_shape) - 1, -1, -1):
4334
+ dim = self._indices[len(self._broadcast_shape) - len(self._arr_shape) + d]
4335
+ dim = dim % self._arr_shape[d]
4336
+ arr_idx += dim * stride
4337
+ stride *= self._arr_shape[d]
4338
+ val = self._flat[arr_idx]
4339
+
4340
+ # Advance to next C-order position
4341
+ for i in range(len(self._indices) - 1, -1, -1):
4342
+ self._indices[i] += 1
4343
+ if self._indices[i] < self._broadcast_shape[i]:
4344
+ break
4345
+ self._indices[i] = 0
4346
+ if i == 0:
4347
+ self._done = True
4348
+
4349
+ dtype = getattr(self._arr, '_dtype', 'float64')
4350
+ if dtype == 'int64':
4351
+ return int(val)
4352
+ return float(val)
4353
+
4354
+
4355
+ # 数组操作函数
4356
+ reshape = _array_ops_module.reshape
4357
+ ravel = _array_ops_module.ravel
4358
+ moveaxis = _array_ops_module.moveaxis
4359
+ rollaxis = _array_ops_module.rollaxis
4360
+ broadcast_to = _array_ops_module.broadcast_to
4361
+ transpose = _array_ops_module.transpose
4362
+ swapaxes = _array_ops_module.swapaxes
4363
+ expand_dims = _array_ops_module.expand_dims
4364
+ squeeze = _array_ops_module.squeeze
4365
+ concatenate = _array_ops_module.concatenate
4366
+ stack = _array_ops_module.stack
4367
+ vstack = _array_ops_module.vstack
4368
+ hstack = _array_ops_module.hstack
4369
+ dstack = _array_ops_module.dstack
4370
+ column_stack = _array_ops_module.column_stack
4371
+ split = _array_ops_module.split
4372
+ hsplit = _array_ops_module.hsplit
4373
+ vsplit = _array_ops_module.vsplit
4374
+ dsplit = _array_ops_module.dsplit
4375
+ tile = _array_ops_module.tile
4376
+ repeat = _array_ops_module.repeat
4377
+ flip = _array_ops_module.flip
4378
+ fliplr = _array_ops_module.fliplr
4379
+ flipud = _array_ops_module.flipud
4380
+ roll = _array_ops_module.roll
4381
+ rot90 = _array_ops_module.rot90
4382
+ append = _array_ops_module.append
4383
+ insert = _array_ops_module.insert
4384
+ delete = _array_ops_module.delete
4385
+ unique = _array_ops_module.unique
4386
+ resize = _array_ops_module.resize
4387
+
4388
+ # 位运算
4389
+ def _to_raw(a):
4390
+ """将输入转为 Rust ndarray,返回 (ndarray, dtype) 元组。"""
4391
+ if hasattr(a, '_array'):
4392
+ return a._array, getattr(a, '_dtype', 'float64')
4393
+ raw = np.ndarray(a)
4394
+ if isinstance(a, int) and not isinstance(a, bool):
4395
+ return raw, 'int64'
4396
+ return raw, 'float64'
4397
+
4398
+ def bitwise_and(x1, x2):
4399
+ """按位与"""
4400
+ r1, dt = _to_raw(x1)
4401
+ r2, _ = _to_raw(x2)
4402
+ return _wrap_result(np.bitwise_and(r1, r2), dt)
4403
+
4404
+ def bitwise_or(x1, x2):
4405
+ """按位或"""
4406
+ r1, dt = _to_raw(x1)
4407
+ r2, _ = _to_raw(x2)
4408
+ return _wrap_result(np.bitwise_or(r1, r2), dt)
4409
+
4410
+ def bitwise_xor(x1, x2):
4411
+ """按位异或"""
4412
+ r1, dt = _to_raw(x1)
4413
+ r2, _ = _to_raw(x2)
4414
+ return _wrap_result(np.bitwise_xor(r1, r2), dt)
4415
+
4416
+ def bitwise_not(x):
4417
+ """按位取反(等效于 ~ 运算符,对布尔数组使用逻辑取反)"""
4418
+ r, dt = _to_raw(x)
4419
+ if dt == 'bool':
4420
+ r = np.invert(r)
4421
+ elif dt == 'uint8':
4422
+ raw = np.bitwise_not(r)
4423
+ r = np.bitwise_and(raw, np.ndarray([255]))
4424
+ dt = 'uint8'
4425
+ else:
4426
+ r = np.bitwise_not(r)
4427
+ return _wrap_result(r, dt)
4428
+
4429
+
4430
+ # invert 与 bitwise_not 行为完全一致,直接作为别名,避免重复实现。
4431
+ invert = bitwise_not
4432
+
4433
+ def left_shift(x1, x2):
4434
+ """左移"""
4435
+ r1, dt = _to_raw(x1)
4436
+ r2, _ = _to_raw(x2)
4437
+ return _wrap_result(np.left_shift(r1, r2), dt)
4438
+
4439
+ def right_shift(x1, x2):
4440
+ """右移"""
4441
+ r1, dt = _to_raw(x1)
4442
+ r2, _ = _to_raw(x2)
4443
+ return _wrap_result(np.right_shift(r1, r2), dt)
4444
+
4445
+
4446
+ # 数学函数
4447
+ sin = _math_functions_module.sin
4448
+ cos = _math_functions_module.cos
4449
+ tan = _math_functions_module.tan
4450
+ arcsin = _math_functions_module.arcsin
4451
+ arccos = _math_functions_module.arccos
4452
+ arctan = _math_functions_module.arctan
4453
+ arctan2 = _math_functions_module.arctan2
4454
+ deg2rad = _math_functions_module.deg2rad
4455
+ rad2deg = _math_functions_module.rad2deg
4456
+ degrees = rad2deg
4457
+ sinh = _math_functions_module.sinh
4458
+ cosh = _math_functions_module.cosh
4459
+ tanh = _math_functions_module.tanh
4460
+ asinh = _math_functions_module.asinh
4461
+ acosh = _math_functions_module.acosh
4462
+ atanh = _math_functions_module.atanh
4463
+ exp = _math_functions_module.exp
4464
+ expm1 = _math_functions_module.expm1
4465
+ log = _math_functions_module.log
4466
+ log10 = _math_functions_module.log10
4467
+ log2 = _math_functions_module.log2
4468
+ log1p = _math_functions_module.log1p
4469
+ around = _math_functions_module.around
4470
+ floor = _math_functions_module.floor
4471
+ ceil = _math_functions_module.ceil
4472
+ trunc = _math_functions_module.trunc
4473
+ fix = _math_functions_module.fix
4474
+ sqrt = _math_functions_module.sqrt
4475
+ square = _math_functions_module.square
4476
+ cbrt = _math_functions_module.cbrt
4477
+ abs = _math_functions_module.abs
4478
+ sign = _math_functions_module.sign
4479
+ reciprocal = _math_functions_module.reciprocal
4480
+ _clip_core = _math_functions_module.clip
4481
+
4482
+ def clip(a, a_min, a_max, out=None):
4483
+ """限制值范围。"""
4484
+ a_arr = asarray(a)
4485
+ if a_min is None:
4486
+ a_min = a_arr.min()
4487
+ if a_max is None:
4488
+ a_max = a_arr.max()
4489
+ result = _clip_core(a, a_min, a_max)
4490
+ if out is not None:
4491
+ for i in range(len(result)):
4492
+ out[i] = result[i]
4493
+ return out
4494
+ return result
4495
+
4496
+
4497
+ sinc = _math_functions_module.sinc
4498
+ heaviside = _math_functions_module.heaviside
4499
+ add = _math_functions_module.add
4500
+ subtract = _math_functions_module.subtract
4501
+ multiply = _math_functions_module.multiply
4502
+ divide = _math_functions_module.divide
4503
+ power = _math_functions_module.power
4504
+ mod = _math_functions_module.mod
4505
+ remainder = _math_functions_module.remainder
4506
+ greater = _math_functions_module.greater
4507
+ less = _math_functions_module.less
4508
+ equal = _math_functions_module.equal
4509
+ not_equal = _math_functions_module.not_equal
4510
+ greater_equal = _math_functions_module.greater_equal
4511
+ less_equal = _math_functions_module.less_equal
4512
+ logical_and = _math_functions_module.logical_and
4513
+ logical_or = _math_functions_module.logical_or
4514
+ logical_xor = _math_functions_module.logical_xor
4515
+ isclose = _math_functions_module.isclose
4516
+ allclose = _math_functions_module.allclose
4517
+
4518
+ # 统计函数
4519
+ sum = _statistics_module.sum
4520
+ mean = _statistics_module.mean
4521
+ std = _statistics_module.std
4522
+ var = _statistics_module.var
4523
+ min = _statistics_module.min
4524
+ max = _statistics_module.max
4525
+ amin = _statistics_module.amin
4526
+ amax = _statistics_module.amax
4527
+ ptp = _statistics_module.ptp
4528
+ median = _statistics_module.median
4529
+ average = _statistics_module.average
4530
+ percentile = _statistics_module.percentile
4531
+ quantile = _statistics_module.quantile
4532
+ nanmedian = _statistics_module.nanmedian
4533
+ nanpercentile = _statistics_module.nanpercentile
4534
+ argmax = _statistics_module.argmax
4535
+ argmin = _statistics_module.argmin
4536
+ argsort = _statistics_module.argsort
4537
+ sort = _statistics_module.sort
4538
+ msort = _statistics_module.msort
4539
+ sort_complex = _statistics_module.sort_complex
4540
+ lexsort = _statistics_module.lexsort
4541
+ partition = _statistics_module.partition
4542
+ argpartition = _statistics_module.argpartition
4543
+ searchsorted = _statistics_module.searchsorted
4544
+ extract = _statistics_module.extract
4545
+ cov = _statistics_module.cov
4546
+ corrcoef = _statistics_module.corrcoef
4547
+ histogram = _statistics_module.histogram
4548
+ histogram2d = _statistics_module.histogram2d
4549
+ histogramdd = _statistics_module.histogramdd
4550
+ digitize = _statistics_module.digitize
4551
+
4552
+
4553
+ def cumsum(a, axis=None):
4554
+ """计算数组元素的累积和。"""
4555
+ return asarray(a).cumsum(axis)
4556
+
4557
+
4558
+ def cumprod(a, axis=None):
4559
+ """计算数组元素的累积乘积。"""
4560
+ return asarray(a).cumprod(axis)
4561
+
4562
+
4563
+ linalg = _linalg_module()
4564
+ random = _random_module()
4565
+ matlib = _matlib_module
4566
+
4567
+ dot = linalg.dot
4568
+ vdot = linalg.vdot
4569
+ inner = linalg.inner
4570
+ matmul = linalg.matmul
4571
+
4572
+
4573
+ # ========== 导出列表 ==========
4574
+ __all__ = [
4575
+ 'ndarray', 'NdArrayIter',
4576
+ 'array', 'asarray', 'asanyarray', 'copy',
4577
+ 'matrix', 'asmatrix', 'bmat', 'repmat',
4578
+ 'dot', 'vdot', 'inner', 'matmul',
4579
+ 'zeros', 'zeros_like', 'ones', 'ones_like', 'full', 'full_like',
4580
+ 'empty', 'empty_like', 'eye', 'identity',
4581
+ 'arange', 'linspace', 'logspace', 'geomspace',
4582
+ 'fromfunction', 'frombuffer', 'r_', 'c_', 's_', 'mgrid', 'ogrid', 'meshgrid',
4583
+ 'reshape', 'ravel', 'moveaxis', 'rollaxis', 'broadcast_to',
4584
+ 'transpose', 'swapaxes', 'expand_dims', 'squeeze',
4585
+ 'concatenate', 'stack', 'vstack', 'hstack', 'dstack', 'column_stack',
4586
+ 'split', 'hsplit', 'vsplit', 'dsplit',
4587
+ 'tile', 'repeat',
4588
+ 'flip', 'fliplr', 'flipud', 'roll', 'rot90',
4589
+ 'append', 'insert', 'delete', 'unique', 'resize',
4590
+ 'where', 'take', 'put', 'select', 'nonzero', 'argwhere', 'flatnonzero',
4591
+ 'binary_repr',
4592
+ 'sin', 'cos', 'tan', 'arcsin', 'arccos', 'arctan', 'arctan2',
4593
+ 'deg2rad', 'rad2deg', 'degrees',
4594
+ 'sinh', 'cosh', 'tanh', 'asinh', 'acosh', 'atanh',
4595
+ 'bitwise_and', 'bitwise_or', 'bitwise_xor', 'bitwise_not',
4596
+ 'invert', 'left_shift', 'right_shift',
4597
+ 'string_', 'unicode_', 'char',
4598
+ 'datetime64', 'timedelta64',
4599
+ 'exp', 'expm1', 'log', 'log10', 'log2', 'log1p',
4600
+ 'around', 'floor', 'ceil', 'trunc', 'fix',
4601
+ 'sqrt', 'square', 'cbrt', 'abs', 'sign', 'reciprocal', 'clip', 'sinc', 'heaviside',
4602
+ 'add', 'subtract', 'multiply', 'divide', 'power', 'mod', 'remainder',
4603
+ 'greater', 'less', 'equal', 'logical_and', 'logical_or', 'isclose', 'allclose',
4604
+ 'sum', 'cumsum', 'cumprod', 'mean', 'std', 'var', 'min', 'max', 'amin', 'amax', 'ptp',
4605
+ 'median', 'average', 'percentile', 'quantile', 'nanmedian', 'nanpercentile',
4606
+ 'argmax', 'argmin', 'argsort', 'sort', 'searchsorted', 'extract',
4607
+ 'cov', 'corrcoef',
4608
+ 'histogram', 'histogram2d', 'histogramdd', 'digitize',
4609
+ 'fft', 'typing',
4610
+ 'pi', 'e', 'euler_gamma', 'inf', 'nan', 'newaxis',
4611
+ 'nditer', '_core', 'exceptions',
4612
+ 'isnan', 'isinf', 'isfinite',
4613
+ 'save', 'load', 'loadtxt', 'savetxt', 'savez', 'genfromtxt',
4614
+ 'Poly', 'polyval', 'polyfit', 'polyder', 'polyint', 'polyroots',
4615
+ 'linalg', 'random', 'matlib', 'load_npz', 'lib',
4616
+ # 标量类型层次
4617
+ 'generic', 'number', 'integer', 'signedinteger', 'unsignedinteger',
4618
+ 'inexact', 'floating', 'complexfloating', 'flexible', 'character',
4619
+ 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64',
4620
+ 'float16', 'float32', 'float64', 'float_', 'complex64', 'complex128',
4621
+ 'bytes_', 'str_', 'bool_', 'object_', 'void',
4622
+ 'byte', 'short', 'intc', 'int_', 'long', 'longlong', 'intp',
4623
+ 'ubyte', 'ushort', 'uintc', 'uint', 'ulong', 'ulonglong', 'uintp',
4624
+ 'half', 'single', 'double', 'longdouble', 'csingle', 'cdouble', 'clongdouble',
4625
+ 'True_', 'False_', 'little_endian', 'ScalarType', 'sctypeDict', 'typecodes',
4626
+ 'issubdtype', 'finfo', 'iinfo', 'ndindex', 'ndenumerate', 'index_exp',
4627
+ 'dtype', 'DType', 'rec', 'ma', 'recarray'
4628
+ ]
4629
+
4630
+
4631
+ # ========== 补充 API(_extra)挂载 ==========
4632
+ # _extra 中的函数均由现有 Rust 原语组合实现;已存在的原生实现(如 iscomplex)优先保留。
4633
+
4634
+ for _extra_name in _extra_module.__all__:
4635
+ if _extra_name not in globals():
4636
+ globals()[_extra_name] = getattr(_extra_module, _extra_name)
4637
+ __all__.append(_extra_name)
4638
+ del _extra_name
4639
+
4640
+
4641
+ # ========== ufunc 归约方法(reduce/accumulate) ==========
4642
+ def _make_ufunc_reduce(reducer):
4643
+ def reduce(a, axis=0, dtype=None, out=None, keepdims=False):
4644
+ _ = keepdims
4645
+ res = reducer(a, axis)
4646
+ if not hasattr(res, '_array'):
4647
+ res = array(res)
4648
+ if dtype is not None:
4649
+ res = res.astype(dtype)
4650
+ if out is not None:
4651
+ out[:] = res
4652
+ return out
4653
+ return res
4654
+ return reduce
4655
+
4656
+
4657
+ def _make_ufunc_accumulate(accumulator):
4658
+ def accumulate(a, axis=0, dtype=None, out=None):
4659
+ _ = dtype, out
4660
+ return accumulator(a, axis)
4661
+ return accumulate
4662
+
4663
+
4664
+ add.reduce = _make_ufunc_reduce(lambda a, ax: sum(a, ax))
4665
+ multiply.reduce = _make_ufunc_reduce(lambda a, ax: _extra_module.prod(a, ax))
4666
+ add.accumulate = _make_ufunc_accumulate(lambda a, ax: cumsum(a, ax))
4667
+ multiply.accumulate = _make_ufunc_accumulate(lambda a, ax: cumprod(a, ax))
4668
+
4669
+ _sys.modules['rsnumpy'] = _sys.modules[__name__]
4670
+ _sys.modules['rsnumpy.__init__'] = _sys.modules[__name__]