rsnumpy 1.0.3__cp313-cp313-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,2315 @@
1
+ """rsnumpy - A Rust-powered NumPy-compatible array library.
2
+
3
+ rsnumpy provides a NumPy-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
+ import builtins
16
+ import rsnumpy._core as _core
17
+ from rsnumpy._core import ndarray_iter as NdArrayIter
18
+ # ========== 子模块导入和函数挂载 ==========
19
+ # 这些导入放在 ndarray 类定义之后以避免循环导入
20
+ from . import array_methods
21
+ from . import math_functions as _math_functions_module
22
+ from . import statistics as _statistics_module
23
+ from . import array_ops as _array_ops_module
24
+
25
+ from .io import save, load, loadtxt, savetxt, savez, load_npz
26
+ from .polynomial import Poly, polyval, polyfit, polyder, polyint, polyroots
27
+ from .linalg import linalg_module as _linalg_module
28
+ from .random import random_module as _random_module
29
+ from . import char as _char_module
30
+ from . import matlib as _matlib_module
31
+
32
+ __version__ = "1.0.3"
33
+
34
+ class ArrayFlags:
35
+ """数组内存布局信息,与 NumPy 的 np.ndarray.flags 兼容。"""
36
+
37
+ def __init__(self, arr):
38
+ self._arr = arr
39
+
40
+ @property
41
+ def C_CONTIGUOUS(self):
42
+ return True
43
+
44
+ @property
45
+ def F_CONTIGUOUS(self):
46
+ return True
47
+
48
+ @property
49
+ def OWNDATA(self):
50
+ return True
51
+
52
+ @property
53
+ def WRITEABLE(self):
54
+ return True
55
+
56
+ @property
57
+ def ALIGNED(self):
58
+ return True
59
+
60
+ @property
61
+ def WRITEBACKIFCOPY(self):
62
+ return False
63
+
64
+ @property
65
+ def UPDATEIFCOPY(self):
66
+ return False
67
+
68
+ def __repr__(self):
69
+ return self.__str__()
70
+
71
+ def __str__(self):
72
+ lines = [
73
+ f" C_CONTIGUOUS : {self.C_CONTIGUOUS}",
74
+ f" F_CONTIGUOUS : {self.F_CONTIGUOUS}",
75
+ f" OWNDATA : {self.OWNDATA}",
76
+ f" WRITEABLE : {self.WRITEABLE}",
77
+ f" ALIGNED : {self.ALIGNED}",
78
+ f" WRITEBACKIFCOPY : {self.WRITEBACKIFCOPY}",
79
+ f" UPDATEIFCOPY : {self.UPDATEIFCOPY}",
80
+ ]
81
+ return "\n".join(lines)
82
+
83
+
84
+ class ndarray:
85
+ """
86
+ rsnumpy.ndarray - 多维数组对象。
87
+
88
+ 这是一个 Rust 后端支持的数组类型,提供与 NumPy 兼容的功能。
89
+
90
+ 参数:
91
+ data (array_like): 数组的初始数据,可以是列表、元组或可迭代对象。
92
+
93
+ 属性:
94
+ shape: 数组维度的元组。
95
+ ndim: 维度数量。
96
+ size: 元素总数。
97
+ dtype: 元素数据类型。
98
+ itemsize: 每个元素的字节大小。
99
+ nbytes: 总字节数。
100
+ T: 数组的转置。
101
+ real: 数组的实部。
102
+ imag: 数组的虚部。
103
+ """
104
+
105
+ def __init__(self, data, _dtype="float64", _fields=None, _raw_data=None):
106
+ if _is_ndarray(data):
107
+ self._array = data._array
108
+ self._dtype = data._dtype
109
+ self._fields = getattr(data, '_fields', None)
110
+ self._raw_data = getattr(data, '_raw_data', None)
111
+ self._complex_data = getattr(data, '_complex_data', None)
112
+ elif hasattr(data, '__class__') and data.__class__.__name__ == 'ndarray':
113
+ self._array = data
114
+ self._dtype = _dtype
115
+ self._fields = _fields
116
+ self._raw_data = _raw_data
117
+ else:
118
+ self._raw_data = _raw_data
119
+ if _raw_data is not None:
120
+ # 原始数据存储:用一个虚拟 Rust 数组占位
121
+ self._array = _core.zeros((len(data),) if isinstance(data, (list, tuple)) else (1,))
122
+ else:
123
+ if isinstance(data, tuple):
124
+ data = list(data)
125
+ if isinstance(data, list) and not _is_rectangular(data):
126
+ # 不规则列表(子列表长度不同),存储为原始数据
127
+ self._raw_data = data
128
+ self._array = _core.zeros((len(data),))
129
+ elif _has_complex(data):
130
+ # 包含复数 → 存储为 _complex_data
131
+ flat = _flatten_data(data)
132
+ self._complex_data = [complex(v) for v in flat]
133
+ _dtype = "complex128"
134
+ self._array = _core.zeros((len(flat),))
135
+ else:
136
+ self._array = _core.ndarray(data)
137
+ self._dtype = _dtype
138
+ self._fields = _fields
139
+
140
+ @staticmethod
141
+ def _wrap(raw_array, _dtype="float64", _fields=None, _raw_data=None):
142
+ """包装原始 Rust ndarray 到 Python 类。"""
143
+ obj = ndarray.__new__(ndarray)
144
+ obj._array = raw_array
145
+ obj._dtype = _dtype
146
+ obj._fields = _fields
147
+ obj._raw_data = _raw_data
148
+ return obj
149
+
150
+ def _wrap_raw(raw_data):
151
+ """包装原始 Python 数据(如字符串列表)到 ndarray。"""
152
+ obj = ndarray.__new__(ndarray)
153
+ obj._array = _core.zeros((len(raw_data),))
154
+ obj._dtype = "string_"
155
+ obj._raw_data = list(raw_data)
156
+ return obj
157
+
158
+ def __repr__(self):
159
+ fields = getattr(self, '_fields', None)
160
+ if fields:
161
+ return _format_structured_repr(self)
162
+ if getattr(self, '_dtype', "float64") == "string_":
163
+ raw_data = getattr(self, '_raw_data', None)
164
+ if raw_data is not None:
165
+ val_strs = [repr(x) for x in raw_data]
166
+ return "array([" + " ".join(val_strs) + "])"
167
+ raw = getattr(self, '_raw_data', None)
168
+ if raw is not None:
169
+ return _format_ragged_repr(self)
170
+ cpx = getattr(self, '_complex_data', None)
171
+ if cpx is not None:
172
+ inner = _format_complex_repr_1d(cpx)
173
+ return f"array({inner})"
174
+ dt = getattr(self, '_dtype', "float64")
175
+ if dt == "float64" and self.ndim == 1:
176
+ values = self._array.tolist()
177
+ inner = _format_float_repr_1d(values)
178
+ return f"array({inner})"
179
+ return f"array({self.__str__()})"
180
+
181
+ def __str__(self):
182
+ fields = getattr(self, '_fields', None)
183
+ if fields:
184
+ return _format_structured_str(self)
185
+ if getattr(self, '_dtype', "float64") == "string_":
186
+ raw_data = getattr(self, '_raw_data', None)
187
+ if raw_data is not None:
188
+ val_strs = [repr(x) for x in raw_data]
189
+ return "[" + " ".join(val_strs) + "]"
190
+ return str(raw_data)
191
+ raw = getattr(self, '_raw_data', None)
192
+ if raw is not None:
193
+ return _format_ragged_str(self)
194
+ cpx = getattr(self, '_complex_data', None)
195
+ if cpx is not None:
196
+ inner = _format_complex_repr_1d(cpx)
197
+ return inner
198
+ if getattr(self, '_dtype', "float64") in ("int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"):
199
+ if getattr(self, '_is_empty', False):
200
+ return _core._format_int_str(self._array)
201
+ return _core._format_int_val_str(self._array)
202
+ if getattr(self, '_dtype', "float64") == "bool":
203
+ def format_bool_list(data):
204
+ if isinstance(data, list):
205
+ inner = " ".join(format_bool_list(v) for v in data)
206
+ return "[" + inner + "]"
207
+ else:
208
+ return "True" if data else "False"
209
+ raw = self._array.tolist()
210
+ if isinstance(raw, list) and len(raw) > 0 and isinstance(raw[0], list):
211
+ lines = []
212
+ for row in raw:
213
+ inner = " ".join(format_bool_list(v) for v in row)
214
+ lines.append("[" + inner + "]")
215
+ return "[" + "\n ".join(lines) + "]"
216
+ return format_bool_list(raw)
217
+ return _core._format_float_str(self._array)
218
+
219
+ def __len__(self):
220
+ raw = getattr(self, '_raw_data', None)
221
+ if raw is not None:
222
+ return len(raw)
223
+ if self.ndim == 0:
224
+ raise TypeError("len() of unsized object")
225
+ return len(self._array)
226
+
227
+ def tolist(self):
228
+ """转换为 Python 列表,元素类型与 _dtype 一致。"""
229
+ raw = getattr(self, '_raw_data', None)
230
+ if raw is not None:
231
+ return raw
232
+ cpx = getattr(self, '_complex_data', None)
233
+ if cpx is not None:
234
+ return list(cpx)
235
+ dt = getattr(self, '_dtype', 'float64')
236
+ raw_list = self._array.tolist()
237
+ if dt == 'int64':
238
+ return _convert_nested(raw_list, int)
239
+ if dt == 'bool':
240
+ return _convert_nested(raw_list, bool)
241
+ if dt in ('uint8', 'uint16', 'uint32', 'uint64'):
242
+ return _convert_nested(raw_list, int)
243
+ if dt == 'string_':
244
+ raw = getattr(self, '_raw_data', None)
245
+ if raw is not None:
246
+ return list(raw)
247
+ return raw_list
248
+
249
+ def __iter__(self):
250
+ return iter(self.tolist())
251
+
252
+ def __bool__(self):
253
+ if self.ndim == 0:
254
+ return bool(self.tolist())
255
+ raise ValueError("The truth value of an array with more than one element is ambiguous.")
256
+
257
+ def __invert__(self):
258
+ """逐元素取反(~ 运算符)。"""
259
+ return bitwise_not(self)
260
+
261
+ def __getitem__(self, key):
262
+ # 结构化数组的字段访问:a['age']
263
+ fields = getattr(self, '_fields', None)
264
+ if fields and isinstance(key, str):
265
+ field_names = [f[0] for f in fields]
266
+ if key in field_names:
267
+ return ndarray._wrap(self._array, self._dtype, _fields=None)
268
+ # 复数数组的布尔索引
269
+ cpx = getattr(self, '_complex_data', None)
270
+ if _is_ndarray(key) and cpx is not None:
271
+ mask = key.tolist()
272
+ filtered = [v for v, m in zip(cpx, mask) if m > 0.5]
273
+ return ndarray(filtered)
274
+ # 展开省略号(...),补充完整切片以匹配数组维度
275
+ if isinstance(key, tuple):
276
+ new_key = []
277
+ ellipsis_count = builtins.sum(1 for k in key if k is Ellipsis)
278
+ if ellipsis_count > 0:
279
+ non_ellipsis = [k for k in key if k is not Ellipsis]
280
+ fill = self.ndim - len(non_ellipsis)
281
+ for k in key:
282
+ if k is Ellipsis:
283
+ for _ in range(fill):
284
+ new_key.append(slice(None, None, None))
285
+ else:
286
+ new_key.append(k)
287
+ key = tuple(new_key)
288
+ # 将 Python ndarray 索引展平为 Python list,确保 Rust 端可正确解析
289
+ if isinstance(key, tuple):
290
+ key = tuple(
291
+ _ndarray_to_index_list(k)
292
+ if hasattr(k, '_array') else k
293
+ for k in key
294
+ )
295
+ else:
296
+ key = (_ndarray_to_index_list(key)
297
+ if hasattr(key, '_array') else key,)
298
+ # 委托给 Rust getitem_multi 处理所有索引逻辑
299
+ result = _core.getitem_multi(self._array, key, list(self.shape))
300
+ return _wrap_result(result, self._dtype)
301
+
302
+ def __setitem__(self, key, value):
303
+ if isinstance(key, tuple):
304
+ # 将 Python ndarray 索引展平为 list
305
+ key = tuple(
306
+ _ndarray_to_index_list(k)
307
+ if hasattr(k, '_array') else k
308
+ for k in key
309
+ )
310
+ elif not isinstance(key, str):
311
+ key = (_ndarray_to_index_list(key)
312
+ if hasattr(key, '_array') else key,)
313
+ if isinstance(key, tuple):
314
+ _core.setitem_multi(self._array, key, list(self.shape), value)
315
+ else:
316
+ self._array[key] = value
317
+
318
+ def __add__(self, other):
319
+ if _is_ndarray(other):
320
+ return _wrap_result(self._array + other._array, self._dtype)
321
+ return _wrap_result(self._array + other, self._dtype)
322
+
323
+ def __radd__(self, other):
324
+ return _wrap_result(other + self._array, self._dtype)
325
+
326
+ def __sub__(self, other):
327
+ if _is_ndarray(other):
328
+ return _wrap_result(self._array - other._array, self._dtype)
329
+ return _wrap_result(self._array - other, self._dtype)
330
+
331
+ def __rsub__(self, other):
332
+ return _wrap_result(other - self._array, self._dtype)
333
+
334
+ def __mul__(self, other):
335
+ if _is_ndarray(other):
336
+ return _wrap_result(self._array * other._array, self._dtype)
337
+ return _wrap_result(self._array * other, self._dtype)
338
+
339
+ def __rmul__(self, other):
340
+ return _wrap_result(other * self._array, self._dtype)
341
+
342
+ def __truediv__(self, other):
343
+ if _is_ndarray(other):
344
+ return _wrap_result(self._array / other._array, self._dtype)
345
+ return _wrap_result(self._array / other, self._dtype)
346
+
347
+ def __rtruediv__(self, other):
348
+ return _wrap_result(other / self._array, self._dtype)
349
+
350
+ def __matmul__(self, other):
351
+ if _is_ndarray(other):
352
+ return _wrap_result(_core.linalg.matmul(self._array, other._array), self._dtype)
353
+ return _wrap_result(_core.linalg.matmul(self._array, other), self._dtype)
354
+
355
+ def __pow__(self, other):
356
+ if _is_ndarray(other):
357
+ return _wrap_result(_core.power(self._array, other._array), self._dtype)
358
+ return _wrap_result(_core.power(self._array, _core.ndarray([other])), self._dtype)
359
+
360
+ def __eq__(self, other):
361
+ if _is_ndarray(other):
362
+ return _wrap_result(self._array.__eq__(other._array), "bool")
363
+ return _wrap_result(self._array.__eq__(other), "bool")
364
+
365
+ def __ne__(self, other):
366
+ if _is_ndarray(other):
367
+ return _wrap_result(self._array.__ne__(other._array), "bool")
368
+ return _wrap_result(self._array.__ne__(other), "bool")
369
+
370
+ def __lt__(self, other):
371
+ if _is_ndarray(other):
372
+ return _wrap_result(self._array.__lt__(other._array), "bool")
373
+ return _wrap_result(self._array.__lt__(other), "bool")
374
+
375
+ def __le__(self, other):
376
+ if _is_ndarray(other):
377
+ return _wrap_result(self._array.__le__(other._array), "bool")
378
+ return _wrap_result(self._array.__le__(other), "bool")
379
+
380
+ def __gt__(self, other):
381
+ if _is_ndarray(other):
382
+ return _wrap_result(self._array.__gt__(other._array), "bool")
383
+ return _wrap_result(self._array.__gt__(other), "bool")
384
+
385
+ def __ge__(self, other):
386
+ if _is_ndarray(other):
387
+ return _wrap_result(self._array.__ge__(other._array), "bool")
388
+ return _wrap_result(self._array.__ge__(other), "bool")
389
+
390
+ def __round__(self, ndigits=None):
391
+ if ndigits is None:
392
+ return _wrap_result(_core.round(self._array, 0), self._dtype)
393
+ return _wrap_result(_core.round(self._array, ndigits), self._dtype)
394
+
395
+ def __and__(self, other):
396
+ """按位与(& 运算符)。"""
397
+ return bitwise_and(self, other)
398
+
399
+ def __or__(self, other):
400
+ """按位或(| 运算符)。"""
401
+ return bitwise_or(self, other)
402
+
403
+ def __xor__(self, other):
404
+ """按位异或(^ 运算符)。"""
405
+ return bitwise_xor(self, other)
406
+
407
+ def __lshift__(self, other):
408
+ """左移(<< 运算符)。"""
409
+ return left_shift(self, other)
410
+
411
+ def __rshift__(self, other):
412
+ """右移(>> 运算符)。"""
413
+ return right_shift(self, other)
414
+
415
+ # ========== 属性 ==========
416
+
417
+ @property
418
+ def shape(self):
419
+ """返回数组维度的元组。"""
420
+ raw = getattr(self, '_raw_data', None)
421
+ arr = getattr(self, '_array', None)
422
+ if raw is not None and arr is not None:
423
+ return self._array.shape
424
+ if raw is not None:
425
+ # 从嵌套列表递归计算形状
426
+ shape = []
427
+ obj = raw
428
+ while isinstance(obj, list):
429
+ shape.append(len(obj))
430
+ if obj:
431
+ obj = obj[0]
432
+ else:
433
+ break
434
+ return tuple(shape)
435
+ return self._array.shape
436
+
437
+ @shape.setter
438
+ def shape(self, new_shape):
439
+ """设置数组形状,相当于原地 reshape。"""
440
+ raw = getattr(self, '_raw_data', None)
441
+ if raw is not None:
442
+ total = len(raw)
443
+ new_flat = list(raw)
444
+ if isinstance(new_shape, int):
445
+ new_shape = (new_shape,)
446
+ new_size = 1
447
+ for s in new_shape:
448
+ new_size *= s
449
+ if new_size != total:
450
+ raise ValueError(f"cannot reshape array of size {total} into shape {new_shape}")
451
+ # 扁平化 raw_data 并按新形状展平
452
+ if len(new_shape) == 1:
453
+ self._raw_data = new_flat
454
+ else:
455
+ self._raw_data = new_flat
456
+ else:
457
+ total = self._array.size
458
+ if isinstance(new_shape, int):
459
+ new_shape = (new_shape,)
460
+ new_size = 1
461
+ for s in new_shape:
462
+ new_size *= s
463
+ if new_size != total:
464
+ raise ValueError(f"cannot reshape array of size {total} into shape {new_shape}")
465
+ self._array = self._array.reshape(new_shape)
466
+
467
+ @property
468
+ def flags(self):
469
+ """返回数组的内存布局信息,与 NumPy 的 ndarray.flags 兼容。"""
470
+ return ArrayFlags(self)
471
+
472
+ @property
473
+ def ndim(self):
474
+ """返回数组维度数量。"""
475
+ raw = getattr(self, '_raw_data', None)
476
+ if raw is not None:
477
+ return len(self.shape)
478
+ return self._array.ndim
479
+
480
+ @property
481
+ def size(self):
482
+ """返回元素总数。"""
483
+ raw = getattr(self, '_raw_data', None)
484
+ if raw is not None:
485
+ s = self.shape
486
+ n = 1
487
+ for v in s:
488
+ n *= v
489
+ return n
490
+ return self._array.size
491
+
492
+ @property
493
+ def dtype(self):
494
+ """返回元素数据类型。"""
495
+ return DType(getattr(self, '_dtype', "float64"))
496
+
497
+ @property
498
+ def itemsize(self):
499
+ """每个元素的字节大小。"""
500
+ return 8 # f64
501
+
502
+ @property
503
+ def nbytes(self):
504
+ """总字节数。"""
505
+ return self.size * self.itemsize
506
+
507
+ @property
508
+ def T(self):
509
+ """数组的转置。"""
510
+ return _ndarray_methods().transpose(self)
511
+
512
+ @property
513
+ def flat(self):
514
+ """返回数组元素的扁平迭代器,逐个访问每个元素。"""
515
+ return _NdFlatIter(self)
516
+
517
+ @flat.setter
518
+ def flat(self, values):
519
+ """设置数组元素的扁平值。"""
520
+ size = self.size
521
+ flat_values = list(values)
522
+ indices = _core.ndarray(list(range(size)))
523
+ vals = _core.ndarray([float(v) for v in flat_values])
524
+ self._array.put(indices, vals)
525
+
526
+ @property
527
+ def real(self):
528
+ """数组的实部。"""
529
+ return self.copy()
530
+
531
+ @property
532
+ def imag(self):
533
+ """数组的虚部(全零)。"""
534
+ return ndarray([0.0] * self.size).reshape(self.shape)
535
+
536
+ # ========== 对象方法 ==========
537
+
538
+ def astype(self, dtype):
539
+ """转换数组为指定类型。"""
540
+ return _ndarray_methods().astype(self, dtype)
541
+
542
+ def reshape(self, *shape):
543
+ """改变数组形状。"""
544
+ return _ndarray_methods().reshape(self, *shape)
545
+
546
+ def resize(self, new_shape):
547
+ """改变数组大小。"""
548
+ _ndarray_methods().resize(self, new_shape)
549
+
550
+ def ravel(self, order='C'):
551
+ """展平数组。"""
552
+ return _ndarray_methods().ravel(self, order)
553
+
554
+ def flatten(self, order='C'):
555
+ """返回数组的一维副本。"""
556
+ return _ndarray_methods().flatten(self, order)
557
+
558
+ def copy(self, order='K'):
559
+ """返回数组的副本。"""
560
+ return _wrap_result(self._array.copy(), self._dtype)
561
+
562
+ def transpose(self, *axes):
563
+ """转置数组。"""
564
+ return _ndarray_methods().transpose(self, *axes)
565
+
566
+ def swapaxes(self, axis1, axis2):
567
+ """交换两个轴。"""
568
+ return _ndarray_methods().swapaxes(self, axis1, axis2)
569
+
570
+ def squeeze(self, axis=None):
571
+ """移除长度为1的轴。"""
572
+ return _ndarray_methods().squeeze(self, axis)
573
+
574
+ def max(self, axis=None):
575
+ """返回最大值。"""
576
+ return _ndarray_methods().max(self, axis)
577
+
578
+ def min(self, axis=None):
579
+ """返回最小值。"""
580
+ return _ndarray_methods().min(self, axis)
581
+
582
+ def mean(self, axis=None):
583
+ """计算平均值。"""
584
+ return _ndarray_methods().mean(self, axis)
585
+
586
+ def std(self, axis=None):
587
+ """计算标准差。"""
588
+ return _ndarray_methods().std(self, axis)
589
+
590
+ def var(self, axis=None):
591
+ """计算方差。"""
592
+ return _ndarray_methods().var(self, axis)
593
+
594
+ def sum(self, axis=None):
595
+ """计算和。"""
596
+ return _ndarray_methods().sum(self, axis)
597
+
598
+ def prod(self, axis=None):
599
+ """计算乘积。"""
600
+ return _ndarray_methods().prod(self, axis)
601
+
602
+ def cumsum(self, axis=None):
603
+ """计算累积和。"""
604
+ return _ndarray_methods().cumsum(self, axis)
605
+
606
+ def cumprod(self, axis=None):
607
+ """计算累积乘积。"""
608
+ return _ndarray_methods().cumprod(self, axis)
609
+
610
+ def argmax(self, axis=None):
611
+ """返回最大值的索引。"""
612
+ return _ndarray_methods().argmax(self, axis)
613
+
614
+ def argmin(self, axis=None):
615
+ """返回最小值的索引。"""
616
+ return _ndarray_methods().argmin(self, axis)
617
+
618
+ def argsort(self, axis=-1):
619
+ """返回排序后的索引。"""
620
+ return _ndarray_methods().argsort(self, axis)
621
+
622
+ def sort(self, axis=-1):
623
+ """原地排序。"""
624
+ _ndarray_methods().sort(self, axis)
625
+
626
+ def diagonal(self, offset=0, axis1=0, axis2=1):
627
+ """返回对角线元素。"""
628
+ return _ndarray_methods().diagonal(self, offset, axis1, axis2)
629
+
630
+ def trace(self, offset=0, axis1=0, axis2=1):
631
+ """计算迹。"""
632
+ return _ndarray_methods().trace(self, offset, axis1, axis2)
633
+
634
+ def fill(self, value):
635
+ """用值填充数组。"""
636
+ _ndarray_methods().fill(self, value)
637
+
638
+ def item(self, *args):
639
+ """获取单个元素。"""
640
+ return _ndarray_methods().item(self, *args)
641
+
642
+ def take(self, indices, axis=None):
643
+ """根据索引取元素。"""
644
+ return _ndarray_methods().take(self, indices, axis)
645
+
646
+ def put(self, indices, values):
647
+ """设置指定位置的元素。"""
648
+ _ndarray_methods().put(self, indices, values)
649
+
650
+ def repeat(self, repeats, axis=None):
651
+ """重复元素。"""
652
+ return _ndarray_methods().repeat(self, repeats, axis)
653
+
654
+ def nonzero(self):
655
+ """返回非零元素索引。"""
656
+ return _ndarray_methods().nonzero(self)
657
+
658
+ # ========== 数学方法 ==========
659
+
660
+ def sin(self):
661
+ """计算正弦。"""
662
+ return _math_funcs().sin(self)
663
+
664
+ def cos(self):
665
+ """计算余弦。"""
666
+ return _math_funcs().cos(self)
667
+
668
+ def tan(self):
669
+ """计算正切。"""
670
+ return _math_funcs().tan(self)
671
+
672
+ def sqrt(self):
673
+ """计算平方根。"""
674
+ return _math_funcs().sqrt(self)
675
+
676
+ def abs(self):
677
+ """计算绝对值。"""
678
+ return _math_funcs().abs(self)
679
+
680
+ def exp(self):
681
+ """计算指数。"""
682
+ return _math_funcs().exp(self)
683
+
684
+ def log(self):
685
+ """计算自然对数。"""
686
+ return _math_funcs().log(self)
687
+
688
+ def log10(self):
689
+ """计算以10为底的对数。"""
690
+ return _math_funcs().log10(self)
691
+
692
+ def log2(self):
693
+ """计算以2为底的对数。"""
694
+ return _math_funcs().log2(self)
695
+
696
+ def clip(self, a_min, a_max):
697
+ """限制值范围。"""
698
+ return _math_funcs().clip(self, a_min, a_max)
699
+
700
+ def round(self, decimals=0):
701
+ """四舍五入到指定小数位。"""
702
+ if decimals == 0:
703
+ return _wrap_result(self._array.__round__(None), self._dtype)
704
+ return _wrap_result(self._array.__round__(decimals), self._dtype)
705
+
706
+ def floor(self):
707
+ """向下取整。"""
708
+ return _math_funcs().floor(self)
709
+
710
+ def ceil(self):
711
+ """向上取整。"""
712
+ return _math_funcs().ceil(self)
713
+
714
+ def all(self):
715
+ """测试是否所有元素为真。"""
716
+ return _core.all(self._array)
717
+
718
+ def any(self):
719
+ """测试是否有任何元素为真。"""
720
+ return _core.any(self._array)
721
+
722
+ def byteswap(self, inplace=False):
723
+ """交换数组中每个元素的字节顺序(大小端转换)。"""
724
+ data = self._array.flatten().tolist()
725
+ dtype = getattr(self, '_dtype', 'float64')
726
+
727
+ swapped = []
728
+ for v in data:
729
+ if isinstance(v, float):
730
+ v = int(v)
731
+ if dtype in ('int16', 'uint16'):
732
+ swapped_val = ((v & 0xFF) << 8) | ((v >> 8) & 0xFF)
733
+ elif dtype in ('int32', 'uint32'):
734
+ swapped_val = ((v & 0xFF) << 24) | ((v >> 8 & 0xFF) << 16) | ((v >> 16 & 0xFF) << 8) | ((v >> 24) & 0xFF)
735
+ elif dtype in ('int64', 'uint64'):
736
+ swapped_val = 0
737
+ for i in range(8):
738
+ swapped_val |= ((v >> (i * 8)) & 0xFF) << ((7 - i) * 8)
739
+ else:
740
+ swapped_val = v
741
+ swapped.append(swapped_val)
742
+
743
+ if inplace:
744
+ arr = _core.ndarray(swapped)
745
+ arr = arr.reshape(self._array.shape)
746
+ self._array = arr
747
+ return self
748
+ else:
749
+ return ndarray(swapped, _dtype=dtype)
750
+
751
+ def view(self, dtype=None):
752
+ """创建一个新的数组对象,共享相同的数据但拥有不同的视图。"""
753
+ return ndarray(self._array, _dtype=getattr(self, '_dtype', 'float64'))
754
+
755
+
756
+ def _ensure(x):
757
+ """将列表/元组转换为 ndarray。"""
758
+ if isinstance(x, (list, tuple)):
759
+ return _core.ndarray(x)
760
+ elif _is_ndarray(x):
761
+ return x._array
762
+ elif hasattr(x, '__class__') and x.__class__.__name__ == 'ndarray':
763
+ return x
764
+ return x
765
+
766
+
767
+ def _is_ndarray(obj):
768
+ """检查对象是否为 rsnumpy ndarray(用 hasattr 避免类身份不一致问题)。"""
769
+ return hasattr(obj, '_array')
770
+
771
+
772
+ def _convert_nested(data, converter):
773
+ """递归转换嵌套列表中的每个元素。"""
774
+ if isinstance(data, list):
775
+ return [_convert_nested(x, converter) for x in data]
776
+ return converter(data)
777
+
778
+
779
+ def _ndarray_to_index_list(k):
780
+ """将 ndarray 索引转为 Rust 可解析的 Python list,保留多维嵌套结构。"""
781
+ raw = k._array.tolist() # preserve shape: [[1],[5],[7],[2]] or [1,0,1,0,1]
782
+ dtype_name = getattr(k, '_dtype', 'float64')
783
+
784
+ def _convert(v):
785
+ if isinstance(v, (list, tuple)):
786
+ return [_convert(x) for x in v]
787
+ if dtype_name == 'bool':
788
+ return bool(v)
789
+ return int(v)
790
+
791
+ return _convert(raw)
792
+
793
+
794
+ def _wrap_result(result, dtype="float64"):
795
+ """将原始 ndarray 结果包装到 ndarray 类中。"""
796
+ if hasattr(result, '__class__') and result.__class__.__name__ == 'ndarray':
797
+ return ndarray._wrap(result, _dtype=dtype)
798
+ if isinstance(result, (list, tuple)):
799
+ return ndarray(result, _dtype=dtype)
800
+ if isinstance(result, float) and dtype == "int64":
801
+ return int(result)
802
+ if hasattr(result, 'tolist'):
803
+ return ndarray._wrap(result, _dtype=dtype)
804
+ return result
805
+
806
+
807
+ def _wrap_with_source(result, source):
808
+ """用源数组的 dtype 包装结果。"""
809
+ dtype = getattr(source, '_dtype', "float64")
810
+ return _wrap_result(result, dtype)
811
+
812
+
813
+ def _scalar(x):
814
+ """转换为标量。"""
815
+ if hasattr(x, 'tolist'):
816
+ return x.tolist()
817
+ return x
818
+
819
+
820
+ def _format_structured_val(val, fields=None):
821
+ """格式化结构化数组中的单个元素值。"""
822
+ def fmt(v, field_type=None):
823
+ if isinstance(v, str):
824
+ return "b'" + v + "'"
825
+ if isinstance(v, bytes):
826
+ return "b'" + v.decode("utf-8", errors="replace") + "'"
827
+ if isinstance(v, float):
828
+ if v == int(v) and abs(v) < 1e16:
829
+ return str(int(v)) + "."
830
+ return str(v)
831
+ # int 值且字段类型为浮点 → 显示小数点
832
+ if isinstance(v, int) and field_type and field_type.startswith('f'):
833
+ return str(v) + "."
834
+ return str(v)
835
+ if isinstance(val, (list, tuple)):
836
+ parts = []
837
+ for i, v in enumerate(val):
838
+ ft = fields[i][1] if fields and i < len(fields) else None
839
+ if not isinstance(ft, str):
840
+ ft = _resolve_type_name(ft) if ft else None
841
+ parts.append(fmt(v, ft))
842
+ if len(parts) == 1:
843
+ return "(" + ", ".join(parts) + ",)"
844
+ return "(" + ", ".join(parts) + ")"
845
+ return "(" + fmt(val) + ",)"
846
+
847
+
848
+ def _format_structured_str(arr):
849
+ """__str__ 用于结构化数组。"""
850
+ fields = getattr(arr, '_fields', None)
851
+ flat = arr.tolist()
852
+ if arr.ndim == 1:
853
+ parts = [_format_structured_val(v, fields) for v in flat]
854
+ return "[" + " ".join(parts) + "]"
855
+ # 高维:递归格式化
856
+ return _format_structured_recursive(flat, arr.ndim, fields)
857
+
858
+
859
+ def _format_structured_recursive(data, ndim, fields=None):
860
+ if ndim == 1:
861
+ parts = [_format_structured_val(v, fields) for v in data]
862
+ return "[" + " ".join(parts) + "]"
863
+ parts = []
864
+ for row in data:
865
+ parts.append(_format_structured_recursive(row, ndim - 1, fields))
866
+ return "[" + "\n ".join(parts) + "]"
867
+
868
+
869
+ def _format_structured_repr(arr):
870
+ """__repr__ 用于结构化数组。"""
871
+ inner = _format_structured_str(arr)
872
+ return f"array({inner})"
873
+
874
+
875
+ def _format_ragged_str(arr):
876
+ """__str__ 用于不规则数组(包含列表/元组元素)。"""
877
+ data = arr.tolist()
878
+ return _format_nested_iterable(data)
879
+
880
+
881
+ def _format_ragged_repr(arr):
882
+ """__repr__ 用于不规则数组。"""
883
+ inner = _format_ragged_str(arr)
884
+ return f"array({inner})"
885
+
886
+
887
+ def _format_nested_iterable(data):
888
+ """递归格式化嵌套可迭代对象为数组字符串。"""
889
+ if isinstance(data, (list, tuple)):
890
+ if not data:
891
+ return "[]"
892
+ items = list(data)
893
+ has_nested = any(isinstance(x, (list, tuple)) for x in items)
894
+ if not has_nested:
895
+ return "[" + " ".join(str(x) for x in items) + "]"
896
+ parts = [_format_nested_iterable(x) for x in items]
897
+ return "[" + "\n ".join(parts) + "]"
898
+ return str(data)
899
+
900
+
901
+ class _float64:
902
+ """float64 标量,显示为 np.float64(value)。"""
903
+ def __init__(self, value):
904
+ self._value = float(value)
905
+
906
+ def __repr__(self):
907
+ return f"np.float64({self._value})"
908
+
909
+ def __str__(self):
910
+ return repr(self)
911
+
912
+ def __float__(self):
913
+ return self._value
914
+
915
+
916
+ def _format_float_repr_1d(values):
917
+ """格式化 1D float64 数组的 repr,带逗号和对齐(匹配 NumPy)。"""
918
+ fmt = [format_float_scalar(v) for v in values]
919
+ max_w = builtins.max(len(f) for f in fmt) if fmt else 0
920
+ parts = [f.rjust(max_w) for f in fmt]
921
+ return "[" + ", ".join(parts) + "]"
922
+
923
+
924
+ def _format_complex_repr_1d(values):
925
+ """格式化 1D complex128 数组的 repr。"""
926
+ parts = [_format_complex_scalar(v) for v in values]
927
+ return "[" + ", ".join(parts) + "]"
928
+
929
+
930
+ def _format_complex_scalar(val):
931
+ """格式化单个复数为字符串(如 1.+0.j, 2.+6.j)。"""
932
+ real = val.real
933
+ imag = val.imag
934
+ if real == int(real) and abs(real) < 1e16:
935
+ real_s = f"{int(real)}."
936
+ else:
937
+ real_s = f"{real}"
938
+ if imag == int(imag) and abs(imag) < 1e16:
939
+ imag_s = f"{int(imag)}."
940
+ else:
941
+ imag_s = f"{imag}"
942
+ if imag >= 0:
943
+ return f"{real_s}+{imag_s}j"
944
+ return f"{real_s}{imag_s}j"
945
+
946
+
947
+ def format_float_scalar(val):
948
+ """Python 版的浮点数格式化(四舍五入,保留有效小数)。"""
949
+ if val != val:
950
+ return "nan"
951
+ if val == float("inf"):
952
+ return "inf"
953
+ if val == float("-inf"):
954
+ return "-inf"
955
+ val_rounded = round(val, 10)
956
+ if val_rounded == int(val_rounded) and abs(val_rounded) < 1e16:
957
+ v = int(val_rounded)
958
+ if float(v) == val_rounded:
959
+ return f"{v}."
960
+ if abs(val_rounded) >= 1e10 or (abs(val_rounded) < 1e-10 and val_rounded != 0):
961
+ return f"{val_rounded:.10e}"
962
+ s = f"{val_rounded:.10f}"
963
+ if '.' in s:
964
+ s = s.rstrip('0').rstrip('.')
965
+ return s
966
+
967
+
968
+ def _ndarray_methods():
969
+ """延迟导入 array_methods 模块。"""
970
+ return array_methods.NdArrayMethods
971
+
972
+
973
+ def _math_funcs():
974
+ """延迟导入 math_functions 模块。"""
975
+ from . import math_functions as mf
976
+ return mf
977
+
978
+
979
+ def _stat_funcs():
980
+ """延迟导入 statistics 模块。"""
981
+ from . import statistics as st
982
+ return st
983
+
984
+
985
+ def _array_ops():
986
+ """延迟导入 array_ops 模块。"""
987
+ from . import array_ops as ao
988
+ return ao
989
+
990
+
991
+ # ========== 类型代码映射 ==========
992
+
993
+ _RS_DTYPE_CODES = {
994
+ # 有符号整数
995
+ 'i1': 'int8', 'int8': 'int8',
996
+ 'i2': 'int16', 'int16': 'int16',
997
+ 'i4': 'int32', 'int32': 'int32',
998
+ 'i8': 'int64', 'int64': 'int64',
999
+ # 无符号整数
1000
+ 'u1': 'uint8', 'uint8': 'uint8',
1001
+ 'u2': 'uint16', 'uint16': 'uint16',
1002
+ 'u4': 'uint32', 'uint32': 'uint32',
1003
+ 'u8': 'uint64', 'uint64': 'uint64',
1004
+ # 浮点
1005
+ 'f2': 'float16', 'float16': 'float16',
1006
+ 'f4': 'float32', 'float32': 'float32',
1007
+ 'f8': 'float64', 'float64': 'float64', 'float': 'float64',
1008
+ # 复数
1009
+ 'c8': 'complex64', 'c64': 'complex64', 'complex64': 'complex64',
1010
+ 'c16': 'complex128', 'c128': 'complex128', 'complex128': 'complex128', 'complex': 'complex128',
1011
+ # 其他
1012
+ 'bool': 'bool', 'b1': 'bool',
1013
+ 'object': 'object', 'O': 'object',
1014
+ }
1015
+
1016
+ _PY_TYPE_TO_RS = {
1017
+ int: 'int64',
1018
+ float: 'float64',
1019
+ complex: 'complex128',
1020
+ bool: 'bool',
1021
+ bytes: 'bytes',
1022
+ str: 'str',
1023
+ }
1024
+
1025
+ # 运行时收集已知的 numpy 类型名称
1026
+ _numpy_type_names = {}
1027
+
1028
+
1029
+ def _init_numpy_types():
1030
+ """尝试导入 numpy 并注册其标量类型名称。"""
1031
+ global _numpy_type_names
1032
+ try:
1033
+ import numpy
1034
+ except ImportError:
1035
+ return
1036
+ for name in dir(numpy):
1037
+ if name.startswith('int') or name.startswith('uint') or name.startswith('float') \
1038
+ or name.startswith('complex') or name == 'bool_' or name == 'bool':
1039
+ _numpy_type_names[name] = name
1040
+
1041
+
1042
+ def _resolve_type_name(tp):
1043
+ """将各种 dtype 输入解析为类型名字符串。"""
1044
+ if isinstance(tp, str):
1045
+ # 直接映射
1046
+ lower = tp.lower()
1047
+ if lower in _RS_DTYPE_CODES:
1048
+ return _RS_DTYPE_CODES[lower]
1049
+ return tp # 未知字符串,原样返回
1050
+ if isinstance(tp, DType):
1051
+ return tp.name
1052
+ if type(tp).__name__ == 'dtype' and hasattr(tp, 'name'):
1053
+ return _RS_DTYPE_CODES.get(tp.name, tp.name)
1054
+ # Python 内置类型
1055
+ if tp in _PY_TYPE_TO_RS:
1056
+ return _PY_TYPE_TO_RS[tp]
1057
+ # 检查类名
1058
+ name = getattr(tp, '__name__', None) or getattr(tp, 'name', str(tp))
1059
+ # 先查已知的 numpy 类型名
1060
+ if name in _numpy_type_names:
1061
+ return _numpy_type_names[name]
1062
+ # 再查代码表
1063
+ if name in _RS_DTYPE_CODES:
1064
+ return _RS_DTYPE_CODES[name]
1065
+ # int_ / float_ 等 numpy 特殊名称
1066
+ name_clean = name.rstrip('_')
1067
+ if name_clean in _RS_DTYPE_CODES:
1068
+ return _RS_DTYPE_CODES[name_clean]
1069
+ return name
1070
+
1071
+
1072
+ # 类型名 → 类型码 逆向映射(用于结构化 dtype)
1073
+ _RS_NAME_TO_CODE = {}
1074
+ for _code, _name in _RS_DTYPE_CODES.items():
1075
+ # 只保留短码映射,忽略 'int32': 'int32' 这样的自映射
1076
+ if len(_code) <= 4 and _code not in ('float', 'complex', 'bool') and _code != _name:
1077
+ _RS_NAME_TO_CODE[_name] = _code
1078
+
1079
+
1080
+ class DType:
1081
+ """表示元素数据类型。"""
1082
+
1083
+ def __init__(self, name, fields=None):
1084
+ self._name = name
1085
+ self._fields = fields # [(field_name, type_code), ...] 或 None
1086
+
1087
+ @property
1088
+ def name(self):
1089
+ return self._name
1090
+
1091
+ def _field_to_code(self, tp):
1092
+ """将字段类型转换为短类型码。"""
1093
+ if isinstance(tp, str):
1094
+ return _RS_NAME_TO_CODE.get(_resolve_type_name(tp), tp)
1095
+ resolved = _resolve_type_name(tp)
1096
+ return _RS_NAME_TO_CODE.get(resolved, resolved)
1097
+
1098
+ def __str__(self):
1099
+ if self._fields:
1100
+ parts = [f"('{f[0]}', '{self._field_to_code(f[1])}')" for f in self._fields]
1101
+ return "[" + ", ".join(parts) + "]"
1102
+ return self._name
1103
+
1104
+ def __repr__(self):
1105
+ if self._fields:
1106
+ parts = [f"('{f[0]}', '{self._field_to_code(f[1])}')" for f in self._fields]
1107
+ return "dtype([" + ", ".join(parts) + "])"
1108
+ return f"dtype('{self._name}')"
1109
+
1110
+ def __eq__(self, other):
1111
+ if isinstance(other, DType):
1112
+ if self._fields is not None or other._fields is not None:
1113
+ return self._fields == other._fields
1114
+ return self._name == other._name
1115
+ if isinstance(other, str):
1116
+ return self._name == other
1117
+ # 对外部类型(如 numpy dtype)通过字符串比较
1118
+ try:
1119
+ return self._name == str(other)
1120
+ except Exception:
1121
+ return NotImplemented
1122
+
1123
+
1124
+ # 立即初始化 numpy 类型名称
1125
+ _init_numpy_types()
1126
+
1127
+
1128
+ def dtype(obj):
1129
+ """创建 dtype 对象。
1130
+
1131
+ >>> dt = dtype('i4')
1132
+ >>> print(dt)
1133
+ int32
1134
+
1135
+ >>> dt = dtype([('age', np.int8)])
1136
+ >>> print(dt)
1137
+ [('age', 'i1')]
1138
+ """
1139
+ if isinstance(obj, (list, tuple)):
1140
+ # 结构化 dtype
1141
+ fields = []
1142
+ for item in obj:
1143
+ if isinstance(item, (list, tuple)) and len(item) >= 2:
1144
+ fields.append((item[0], item[1]))
1145
+ if fields:
1146
+ return DType("void", fields=fields)
1147
+ name = _resolve_type_name(obj)
1148
+ return DType(name)
1149
+
1150
+
1151
+ # ========== 标量类型别名 ==========
1152
+
1153
+ int8 = type('int8', (), {})
1154
+ int16 = type('int16', (), {})
1155
+ int32 = type('int32', (), {})
1156
+ int64 = type('int64', (), {})
1157
+ uint8 = type('uint8', (), {})
1158
+ uint16 = type('uint16', (), {})
1159
+ uint32 = type('uint32', (), {})
1160
+ uint64 = type('uint64', (), {})
1161
+ float16 = type('float16', (), {})
1162
+ float32 = type('float32', (), {})
1163
+ float64 = type('float64', (), {})
1164
+ float_ = float64 # NumPy 别名
1165
+ complex64 = type('complex64', (), {})
1166
+ complex128 = type('complex128', (), {})
1167
+ string_ = type('string_', (), {})
1168
+ unicode_ = type('unicode_', (), {})
1169
+
1170
+
1171
+ # ========== 构造/工厂函数 ==========
1172
+
1173
+ def _flatten_check(data):
1174
+ """单次遍历:展平数据并检查是否包含复数。"""
1175
+ flat = []
1176
+ has_complex = False
1177
+ stack = [data]
1178
+ while stack:
1179
+ item = stack.pop()
1180
+ if isinstance(item, (list, tuple)):
1181
+ stack.extend(reversed(item))
1182
+ else:
1183
+ flat.append(item)
1184
+ if isinstance(item, complex):
1185
+ has_complex = True
1186
+ return flat, has_complex
1187
+
1188
+
1189
+ def _has_complex(data):
1190
+ """检查数据中是否包含复数。"""
1191
+ flat, has_c = _flatten_check(data)
1192
+ return has_c
1193
+
1194
+
1195
+ def _flatten_data(data):
1196
+ """展平嵌套的可迭代对象为扁平列表。"""
1197
+ flat, _ = _flatten_check(data)
1198
+ return flat
1199
+
1200
+
1201
+ def _is_rectangular(data):
1202
+ """检查嵌套列表是否为矩形(所有子列表长度相同)。"""
1203
+ if not isinstance(data, (list, tuple)):
1204
+ return True
1205
+ if not data:
1206
+ return True
1207
+ first_len = None
1208
+ for item in data:
1209
+ if isinstance(item, (list, tuple)):
1210
+ if first_len is None:
1211
+ first_len = len(item)
1212
+ elif len(item) != first_len:
1213
+ return False
1214
+ else:
1215
+ if first_len is not None:
1216
+ return False
1217
+ return True
1218
+
1219
+
1220
+ def array(data, dtype=None, copy=True, order='K', subok=False, ndmin=0):
1221
+ """创建数组。"""
1222
+ if dtype is None:
1223
+ # 单次遍历:展平 + 推断 dtype
1224
+ flat, has_c = _flatten_check(data)
1225
+ if has_c:
1226
+ _dtype = "complex128"
1227
+ else:
1228
+ _dtype = _infer_int_dtype(*flat)
1229
+ else:
1230
+ _dtype = _resolve_dtype(dtype)
1231
+ _fields = None
1232
+ _raw_data = None
1233
+ if isinstance(dtype, DType) and dtype._fields:
1234
+ _fields = dtype._fields
1235
+ # 多字段结构化数据 → 存储原始 Python 数据
1236
+ if len(_fields) > 1 or any(isinstance(row, (list, tuple)) and len(row) > 1 for row in data):
1237
+ _raw_data = list(data)
1238
+ data = [tuple(v for v in row) if isinstance(row, (list, tuple)) else row for row in data]
1239
+ else:
1240
+ # 单字段:展平数值
1241
+ flat_data = [row[0] if isinstance(row, (list, tuple)) and len(row) == 1 else row for row in data]
1242
+ data = flat_data
1243
+ arr = ndarray(data, _dtype=_dtype, _fields=_fields, _raw_data=_raw_data)
1244
+ if ndmin > arr.ndim:
1245
+ new_shape = (1,) * (ndmin - arr.ndim) + arr.shape
1246
+ arr = ndarray._wrap(arr._array.reshape(new_shape), _dtype=_dtype, _fields=_fields, _raw_data=_raw_data)
1247
+ return arr
1248
+
1249
+
1250
+ def asarray(a, dtype=None, order=None):
1251
+ """转换输入为数组。"""
1252
+ if _is_ndarray(a):
1253
+ return a
1254
+ if dtype is not None:
1255
+ dt_str = dtype if isinstance(dtype, str) else dtype.__name__
1256
+ if dt_str in ("complex", "complex128", "complex64", "cfloat", "cdouble"):
1257
+ _dtype = "complex128"
1258
+ else:
1259
+ _dtype = _resolve_dtype(dtype)
1260
+ else:
1261
+ # 从数据推断 dtype(单次遍历)
1262
+ if isinstance(a, (list, tuple)):
1263
+ flat, _ = _flatten_check(a)
1264
+ _dtype = _infer_int_dtype(*flat)
1265
+ else:
1266
+ _dtype = "float64"
1267
+ return ndarray(a, _dtype=_dtype)
1268
+
1269
+
1270
+ def asanyarray(a, dtype=None, order=None):
1271
+ """转换输入为 ndarray。"""
1272
+ return asarray(a, dtype, order)
1273
+
1274
+
1275
+ def matrix(data, dtype=None, copy=True):
1276
+ """从数组或字符串创建矩阵对象。"""
1277
+ if isinstance(data, str):
1278
+ rows = data.split(';')
1279
+ data = []
1280
+ for row in rows:
1281
+ elements = row.split(',')
1282
+ data.append([float(e.strip()) for e in elements])
1283
+ return asarray(data, dtype=dtype)
1284
+
1285
+
1286
+ def asmatrix(data, dtype=None):
1287
+ """将输入转换为矩阵对象。"""
1288
+ return asarray(data, dtype=dtype)
1289
+
1290
+
1291
+ def bmat(obj, ldict=None, gdict=None):
1292
+ """从字符串、嵌套序列或数组构建矩阵。"""
1293
+ if isinstance(obj, str):
1294
+ obj = obj.strip()
1295
+ if obj.startswith('[') and obj.endswith(']'):
1296
+ obj = obj[1:-1].strip()
1297
+ rows = obj.split(';')
1298
+ data = []
1299
+ for row in rows:
1300
+ elements = row.split(',')
1301
+ row_data = []
1302
+ for e in elements:
1303
+ e = e.strip()
1304
+ if e.startswith('[') and e.endswith(']'):
1305
+ inner = e[1:-1].split()
1306
+ row_data.extend([float(x) for x in inner])
1307
+ else:
1308
+ row_data.append(float(e))
1309
+ data.append(row_data)
1310
+ return asarray(data)
1311
+ if isinstance(obj, (list, tuple)):
1312
+ if isinstance(obj[0], (list, tuple)):
1313
+ if isinstance(obj[0][0], (list, tuple)):
1314
+ rows = []
1315
+ for row_block in obj:
1316
+ row_len = sum(len(block[0]) for block in row_block)
1317
+ # current_row = [0.0] * row_len
1318
+ col_start = 0
1319
+ for block in row_block:
1320
+ b_arr = _to_ndarray(block)
1321
+ b_rows, b_cols = b_arr.shape
1322
+ for i in range(b_rows):
1323
+ if i >= len(rows):
1324
+ rows.append([0.0] * row_len)
1325
+ for j in range(b_cols):
1326
+ rows[i][col_start + j] = b_arr[i, j]
1327
+ col_start += b_cols
1328
+ return asarray(rows)
1329
+ else:
1330
+ return asarray(obj)
1331
+ else:
1332
+ return asarray(obj)
1333
+ return _to_ndarray(obj)
1334
+
1335
+
1336
+ def repmat(a, m, n):
1337
+ """将数组或矩阵重复 m 行 n 列。"""
1338
+ a = _to_ndarray(a)
1339
+ if a.ndim == 1:
1340
+ a = a.reshape(1, a.size)
1341
+ return tile(a, (m, n))
1342
+
1343
+
1344
+ def _to_ndarray(obj):
1345
+ """将任意数组类对象转换为 rsnumpy ndarray(不依赖第三方库)。"""
1346
+ if _is_ndarray(obj):
1347
+ return obj
1348
+ if hasattr(obj, 'tolist'):
1349
+ return ndarray(obj.tolist())
1350
+ return ndarray(obj)
1351
+
1352
+
1353
+ def array_equal(a, b):
1354
+ """判断两个数组是否形状和元素完全相同。"""
1355
+ a = _to_ndarray(a)
1356
+ b = _to_ndarray(b)
1357
+ if a.shape != b.shape:
1358
+ return False
1359
+ return bool(_core.array_equal(a._array, b._array))
1360
+
1361
+
1362
+ def copy(a, order='K'):
1363
+ """返回数组的副本。"""
1364
+ return ndarray(a).copy(order=order)
1365
+
1366
+
1367
+ def _resolve_dtype(dtype):
1368
+ """解析 dtype 字符串/Types/DType 为内部表示。"""
1369
+ if dtype is None:
1370
+ return "float64"
1371
+ if isinstance(dtype, DType):
1372
+ return dtype.name
1373
+ if isinstance(dtype, (list, tuple)):
1374
+ return "void"
1375
+ dt_str = dtype if isinstance(dtype, str) else dtype.__name__
1376
+ if dt_str in ("complex", "complex128", "complex64", "cfloat", "cdouble"):
1377
+ return "complex128"
1378
+ if dt_str in ("int", "int_", "intp", "intc"):
1379
+ return "int64"
1380
+ if dt_str in ("int8", "int16", "int32", "int64"):
1381
+ return dt_str
1382
+ if dt_str in ("uint", "uint8", "uint16", "uint32", "uint64"):
1383
+ return dt_str
1384
+ if dt_str in ("bool", "bool_"):
1385
+ return "bool"
1386
+ return "float64"
1387
+
1388
+
1389
+ def _infer_int_dtype(*args):
1390
+ """推断参数中的类型:全 bool→bool,全 int→int64,否则 float64。"""
1391
+ has_bool = False
1392
+ has_int = False
1393
+ for a in args:
1394
+ if isinstance(a, bool):
1395
+ has_bool = True
1396
+ elif isinstance(a, int):
1397
+ has_int = True
1398
+ else:
1399
+ return "float64"
1400
+ if has_bool and not has_int:
1401
+ return "bool"
1402
+ return "int64"
1403
+
1404
+
1405
+ def _make_structured_zeros(shape, fields):
1406
+ """为结构化 dtype 生成零填充的原始数据。"""
1407
+ elem = tuple(0 for _ in fields)
1408
+ if isinstance(shape, int):
1409
+ total = shape
1410
+ else:
1411
+ total = 1
1412
+ for s in shape:
1413
+ total *= s
1414
+ flat = [elem] * total
1415
+ if isinstance(shape, int):
1416
+ shape = (shape,)
1417
+ if len(shape) <= 1:
1418
+ return flat
1419
+ # 重塑为多维
1420
+ result = flat
1421
+ for dim in reversed(shape[1:]):
1422
+ result = [result[i:i+dim] for i in range(0, len(result), dim)]
1423
+ return result
1424
+
1425
+
1426
+ def zeros(shape, dtype=None, order='C'):
1427
+ """返回指定形状的零数组。"""
1428
+ _fields = None
1429
+ _raw_data = None
1430
+ if isinstance(dtype, (list, tuple)):
1431
+ _fields = [(item[0], item[1]) for item in dtype]
1432
+ _raw_data = _make_structured_zeros(shape, _fields)
1433
+ _dtype = "void"
1434
+ elif isinstance(dtype, DType) and dtype._fields:
1435
+ _fields = dtype._fields
1436
+ _raw_data = _make_structured_zeros(shape, _fields)
1437
+ _dtype = "void"
1438
+ else:
1439
+ _dtype = _resolve_dtype(dtype)
1440
+ return ndarray(_core.zeros(shape), _dtype=_dtype, _fields=_fields, _raw_data=_raw_data)
1441
+
1442
+
1443
+ def ones(shape, dtype=None, order='C'):
1444
+ """返回指定形状的1数组。"""
1445
+ _dtype = _resolve_dtype(dtype)
1446
+ return ndarray(_core.ones(shape), _dtype=_dtype)
1447
+
1448
+
1449
+ def empty(shape, dtype=None, order='C'):
1450
+ """返回指定形状的空数组。"""
1451
+ _dtype = _resolve_dtype(dtype)
1452
+ arr = ndarray(_core.empty(shape), _dtype=_dtype)
1453
+ arr._is_empty = True
1454
+ return arr
1455
+
1456
+
1457
+ def full(shape, fill_value, dtype=None, order='C'):
1458
+ """返回指定形状的填充数组。"""
1459
+ _dtype = _resolve_dtype(dtype)
1460
+ return ndarray(_core.full(shape, fill_value), _dtype=_dtype)
1461
+
1462
+
1463
+ def zeros_like(a, dtype=None, order='K', subok=True, shape=None):
1464
+ """返回与输入形状相同的零数组。"""
1465
+ arr = ndarray(a)
1466
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
1467
+ return ndarray(_core.zeros_like(arr._array), _dtype=_dtype)
1468
+
1469
+
1470
+ def ones_like(a, dtype=None, order='K', subok=True, shape=None):
1471
+ """返回与输入形状相同的1数组。"""
1472
+ arr = ndarray(a)
1473
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
1474
+ return ndarray(_core.ones_like(arr._array), _dtype=_dtype)
1475
+
1476
+
1477
+ def empty_like(a, dtype=None, order='K', subok=True, shape=None):
1478
+ """返回与输入形状相同的空数组。"""
1479
+ arr = ndarray(a)
1480
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
1481
+ return ndarray(_core.empty_like(arr._array), _dtype=_dtype)
1482
+
1483
+
1484
+ def full_like(a, fill_value, dtype=None, order='K', subok=True, shape=None):
1485
+ """返回与输入形状相同的填充数组。"""
1486
+ arr = ndarray(a)
1487
+ _dtype = dtype if dtype is not None else getattr(arr, '_dtype', 'float64')
1488
+ return ndarray(_core.full_like(arr._array, float(fill_value)), _dtype=_dtype)
1489
+
1490
+
1491
+ def eye(N, M=None, k=0, dtype=None, order='C'):
1492
+ """返回对角线为1的二维数组。"""
1493
+ _dtype = _resolve_dtype(dtype)
1494
+ return ndarray(_core.eye(N, M, k), _dtype=_dtype)
1495
+
1496
+
1497
+ def identity(n, dtype=None):
1498
+ """返回单位矩阵。"""
1499
+ return eye(n, dtype=dtype)
1500
+
1501
+
1502
+ def arange(start=0, stop=None, step=1, dtype=None):
1503
+ """返回给定间隔内均匀间隔的值。"""
1504
+ if stop is None:
1505
+ stop = start
1506
+ start = 0
1507
+ _dtype = _resolve_dtype(dtype) if dtype is not None else _infer_int_dtype(start, stop, step)
1508
+ return ndarray(_core.arange(start, stop, step), _dtype=_dtype)
1509
+
1510
+
1511
+ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
1512
+ """返回指定间隔内均匀间隔的数字。"""
1513
+ start_val = _scalar(_ensure(start))
1514
+ stop_val = _scalar(_ensure(stop))
1515
+ _dtype = _resolve_dtype(dtype)
1516
+ result = ndarray(_core.linspace(start_val, stop_val, num, endpoint=endpoint), _dtype=_dtype)
1517
+ if retstep:
1518
+ step = (stop_val - start_val) / (num - 1) if endpoint else (stop_val - start_val) / num
1519
+ return result, _float64(step)
1520
+ return result
1521
+
1522
+
1523
+ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0):
1524
+ """返回对数刻度上均匀间隔的数字。"""
1525
+ _dtype = _resolve_dtype(dtype) if dtype else "float64"
1526
+ return ndarray(_core.logspace(float(start), float(stop), num, float(base)), _dtype=_dtype)
1527
+
1528
+
1529
+ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
1530
+ """返回几何级数上均匀间隔的数字。"""
1531
+ if start <= 0 or stop <= 0:
1532
+ raise ValueError("geomspace requires positive start and stop values")
1533
+ _dtype = _resolve_dtype(dtype) if dtype else "float64"
1534
+ return ndarray(_core.geomspace(float(start), float(stop), num), _dtype=_dtype)
1535
+
1536
+
1537
+ # ========== 索引函数 ==========
1538
+
1539
+ def ix_(*args):
1540
+ """从多个序列构造开放网格,用于笛卡尔积花式索引。
1541
+
1542
+ 返回 ndarray 元组,每个数组沿对应维度广播。
1543
+
1544
+ 示例:
1545
+ >>> x = np.arange(32).reshape((8, 4))
1546
+ >>> x[np.ix_([1,5,7,2], [0,3,1,2])]
1547
+ """
1548
+ raw = _core.ix_rs(args)
1549
+ return tuple(ndarray._wrap(r) for r in raw)
1550
+
1551
+
1552
+ def where(condition, x=None, y=None):
1553
+ """根据条件返回元素或索引。"""
1554
+ if x is not None and y is not None:
1555
+ dtype = getattr(x, '_dtype', 'float64')
1556
+ fields = getattr(x, '_fields', None)
1557
+ raw_data = getattr(x, '_raw_data', None)
1558
+ result = _core.where(_ensure(condition), _ensure(x), _ensure(y))
1559
+ return ndarray._wrap(result, _dtype=dtype, _fields=fields, _raw_data=raw_data)
1560
+ return nonzero(condition)
1561
+
1562
+
1563
+ def take(a, indices, axis=None, out=None, mode='raise'):
1564
+ """沿指定轴取元素。"""
1565
+ arr = ndarray(a)
1566
+ return arr.take(indices, axis)
1567
+
1568
+
1569
+ def put(a, indices, values, mode='raise'):
1570
+ """设置指定位置的元素。"""
1571
+ arr = ndarray(a)
1572
+ arr.put(indices, values)
1573
+
1574
+
1575
+ def select(condlist, choicelist, default=0):
1576
+ """根据条件列表选择值(Rust 单次遍历实现)。"""
1577
+ cond_arrays = [_ensure(c) for c in condlist]
1578
+ choice_arrays = [_ensure(c) for c in choicelist]
1579
+ result = _core.select_rs(cond_arrays, choice_arrays, float(default))
1580
+ return ndarray(result)
1581
+
1582
+
1583
+ def nonzero(a):
1584
+ """返回非零元素的索引,返回元组形式的数组(兼容NumPy)。"""
1585
+ raw = _core.nonzero_arrs(_ensure(a))
1586
+ return tuple(ndarray._wrap(arr, _dtype="int64") for arr in raw)
1587
+
1588
+
1589
+ def argwhere(a):
1590
+ """查找非零元素的索引。"""
1591
+ return ndarray(_core.argwhere(_ensure(a)))
1592
+
1593
+
1594
+ def flatnonzero(a):
1595
+ """返回扁平化数组中非零元素的索引。"""
1596
+ return ndarray(_core.flatnonzero(_ensure(a)))
1597
+
1598
+
1599
+ # ========== 缺失构造函数 ==========
1600
+
1601
+ def fromfunction(function, shape, *, dtype=None, **kwargs):
1602
+ """根据函数和形状创建数组。"""
1603
+ indices = _core._arange_arrays(shape)
1604
+ grid = _core.meshgrid(*indices, indexing='ij')
1605
+ return function(*grid, **kwargs)
1606
+
1607
+
1608
+ def frombuffer(buffer, dtype=None, count=-1, offset=0, *, like=None):
1609
+ """从缓冲区创建一维数组(使用 Rust 层实现)。"""
1610
+ if isinstance(buffer, bytes):
1611
+ if offset > 0:
1612
+ buffer = buffer[offset:]
1613
+ if isinstance(dtype, str) and len(dtype) > 1 and dtype[0] == 'S' and dtype[1:].isdigit():
1614
+ # Byte string type (e.g. S1, S2, ...)
1615
+ elem_size = int(dtype[1:])
1616
+ total = len(buffer) // elem_size
1617
+ if count >= 0:
1618
+ total = min(total, count)
1619
+ raw = [buffer[i * elem_size:(i + 1) * elem_size] for i in range(total)]
1620
+ return ndarray(_core.zeros((len(raw),)), _dtype=dtype, _raw_data=raw)
1621
+ return ndarray(_core.bytes_to_floats(buffer, count))
1622
+ return ndarray(list(buffer))
1623
+
1624
+
1625
+ def fromiter(iterable, dtype=None, count=-1, *, like=None):
1626
+ """从可迭代对象创建一维数组。"""
1627
+ _dtype = _resolve_dtype(dtype)
1628
+ if count >= 0:
1629
+ data = [next(iterable) for _ in range(count)]
1630
+ else:
1631
+ data = list(iterable)
1632
+ return ndarray(data, _dtype=_dtype)
1633
+
1634
+
1635
+ class _RClass:
1636
+ """行连接辅助类,模拟 np.r_。"""
1637
+ def __getitem__(self, item):
1638
+ if not isinstance(item, tuple):
1639
+ item = (item,)
1640
+ arrays = []
1641
+ for it in item:
1642
+ arrays.append(ndarray(it))
1643
+ return concatenate([a.ravel() for a in arrays])
1644
+
1645
+
1646
+ r_ = _RClass()
1647
+
1648
+
1649
+ class _SClass:
1650
+ """切片辅助类,模拟 np.s_。"""
1651
+ def __getitem__(self, item):
1652
+ return item
1653
+
1654
+
1655
+ s_ = _SClass()
1656
+
1657
+
1658
+ def mgrid(*ranges):
1659
+ """密集网格索引,返回密集网格数组。"""
1660
+ arrays = []
1661
+ for r in ranges:
1662
+ if isinstance(r, slice):
1663
+ arrays.append(ndarray(arange(r.start or 0, r.stop, r.step or 1)))
1664
+ elif isinstance(r, (int, float)):
1665
+ arrays.append(ndarray([float(r)]))
1666
+ else:
1667
+ arrays.append(ndarray(r))
1668
+ grids = _core.meshgrid(*arrays, indexing='ij')
1669
+ if isinstance(grids, tuple) and len(grids) > 1:
1670
+ return stack(grids, axis=0)
1671
+ return grids
1672
+
1673
+
1674
+ def ogrid(*ranges):
1675
+ """开放网格索引,返回开放网格数组。"""
1676
+ arrays = []
1677
+ for r in ranges:
1678
+ if isinstance(r, slice):
1679
+ arrays.append(ndarray(arange(r.start or 0, r.stop, r.step or 1)))
1680
+ elif isinstance(r, (int, float)):
1681
+ arrays.append(ndarray([float(r)]))
1682
+ else:
1683
+ arrays.append(ndarray(r))
1684
+ return _core.meshgrid(*arrays, indexing='ij')
1685
+
1686
+
1687
+ # ========== FFT 函数 ==========
1688
+
1689
+ def fft(a, n=None, axis=-1):
1690
+ """计算一维离散傅里叶变换。"""
1691
+ arr = ndarray(a)
1692
+ return _core.py_fft_ndarray(arr._array)
1693
+
1694
+
1695
+ def ifft(a, n=None, axis=-1):
1696
+ """计算一维逆离散傅里叶变换。"""
1697
+ if _is_ndarray(a):
1698
+ return _core.py_ifft_ndarray(a._array)
1699
+ return _core.py_ifft(a)
1700
+
1701
+
1702
+ def rfft(a, n=None, axis=-1):
1703
+ """计算实输入的一维离散傅里叶变换。"""
1704
+ arr = ndarray(a)
1705
+ return _core.py_rfft_ndarray(arr._array)
1706
+
1707
+
1708
+ def irfft(a, n=None, axis=-1):
1709
+ """计算 rfft 的逆变换。"""
1710
+ if _is_ndarray(a):
1711
+ return _core.py_irfft_ndarray(a._array, n)
1712
+ return _core.py_irfft(a, n)
1713
+
1714
+
1715
+ # ========== 常量 ==========
1716
+ pi = 3.141592653589793
1717
+ e = 2.718281828459045
1718
+ euler_gamma = 0.5772156649015328606
1719
+ inf = float('inf')
1720
+ nan = float('nan')
1721
+ newaxis = None
1722
+
1723
+
1724
+ # 字符串操作模块
1725
+ char = _char_module
1726
+
1727
+
1728
+ # ========== 判断函数 ==========
1729
+
1730
+ class _NdFlatIter:
1731
+ """ndarray.flat 的扁平迭代器,逐元素访问数组。"""
1732
+ def __init__(self, arr):
1733
+ self._arr = arr
1734
+ self._flat = arr._array.flatten()
1735
+ self._size = self._flat.size
1736
+ self._idx = 0
1737
+ self._is_int = getattr(arr, '_dtype', 'float64') == 'int64'
1738
+
1739
+ def __iter__(self):
1740
+ return self
1741
+
1742
+ def __next__(self):
1743
+ if self._idx >= self._size:
1744
+ raise StopIteration
1745
+ val = self._flat[self._idx]
1746
+ self._idx += 1
1747
+ return int(val) if self._is_int else float(val)
1748
+
1749
+
1750
+ class _NdIterElement:
1751
+ """nditer 的可写元素包装,支持 op_flags=['readwrite'] 时修改原数组。"""
1752
+ __slots__ = ('_arr', '_idx', '_val')
1753
+
1754
+ def __init__(self, arr, idx, val):
1755
+ self._arr = arr
1756
+ self._idx = idx
1757
+ self._val = val
1758
+
1759
+ def __repr__(self):
1760
+ return repr(self._val)
1761
+
1762
+ def __float__(self):
1763
+ return float(self._val)
1764
+
1765
+ def __int__(self):
1766
+ return int(self._val)
1767
+
1768
+ def __bool__(self):
1769
+ return bool(self._val)
1770
+
1771
+ def __eq__(self, other):
1772
+ return self._val == other
1773
+
1774
+ def __ne__(self, other):
1775
+ return self._val != other
1776
+
1777
+ def __lt__(self, other):
1778
+ return self._val < other
1779
+
1780
+ def __le__(self, other):
1781
+ return self._val <= other
1782
+
1783
+ def __gt__(self, other):
1784
+ return self._val > other
1785
+
1786
+ def __ge__(self, other):
1787
+ return self._val >= other
1788
+
1789
+ def __add__(self, other):
1790
+ return self._val + other
1791
+
1792
+ def __radd__(self, other):
1793
+ return other + self._val
1794
+
1795
+ def __sub__(self, other):
1796
+ return self._val - other
1797
+
1798
+ def __rsub__(self, other):
1799
+ return other - self._val
1800
+
1801
+ def __mul__(self, other):
1802
+ return self._val * other
1803
+
1804
+ def __rmul__(self, other):
1805
+ return other * self._val
1806
+
1807
+ def __truediv__(self, other):
1808
+ return self._val / other
1809
+
1810
+ def __rtruediv__(self, other):
1811
+ return other / self._val
1812
+
1813
+ def __neg__(self):
1814
+ return -self._val
1815
+
1816
+ def __abs__(self):
1817
+ return abs(self._val)
1818
+
1819
+ def __getitem__(self, key):
1820
+ return self._val
1821
+
1822
+ def __setitem__(self, key, value):
1823
+ self._val = float(value)
1824
+ self._arr.put([self._idx], [self._val])
1825
+
1826
+
1827
+ def nditer(a, order='C', op_flags=None, flags=None):
1828
+ """创建数组元素的迭代器,逐个访问数组元素。
1829
+
1830
+ 参数:
1831
+ a: 输入数组或数组列表(广播迭代)
1832
+ order: 遍历顺序,'C'(行序优先)或 'F'(列序优先)
1833
+ op_flags: 操作标志列表,如 ['readwrite'] 表示可读写
1834
+ flags: 迭代器标志,如 ['external_loop'] 表示返回一维数组块而非标量
1835
+
1836
+ >>> a = np.array([[0, 1, 2], [3, 4, 5]])
1837
+ >>> for x in np.nditer(a):
1838
+ ... print(x, end=", ")
1839
+ 0, 1, 2, 3, 4, 5,
1840
+ """
1841
+ # 多数组广播迭代
1842
+ if isinstance(a, (list, tuple)):
1843
+ arrays = [ndarray(x) for x in a]
1844
+ # 计算广播形状:每个维度取最大值
1845
+ all_shapes = [x.shape for x in arrays]
1846
+ max_ndim = builtins.max(len(s) for s in all_shapes)
1847
+ padded = []
1848
+ for s in all_shapes:
1849
+ pad = [1] * (max_ndim - len(s)) + list(s)
1850
+ padded.append(pad)
1851
+ bcast_shape = tuple(builtins.max(p[i] for p in padded) for i in range(max_ndim))
1852
+ # 广播每个数组到统一形状
1853
+ bcast_arrays = []
1854
+ for arr in arrays:
1855
+ if arr.shape == bcast_shape:
1856
+ bcast_arrays.append(arr)
1857
+ else:
1858
+ bcast_arrays.append(ndarray(_core.broadcast_to(arr._array, bcast_shape)))
1859
+ # 展平并迭代
1860
+ flats = [x._array.flatten().tolist() for x in bcast_arrays]
1861
+ is_ints = [getattr(x, '_dtype', 'float64') == 'int64' for x in bcast_arrays]
1862
+ for i in range(len(flats[0])):
1863
+ yield tuple(int(f[i]) if is_ints[j] else f[i] for j, f in enumerate(flats))
1864
+ return
1865
+
1866
+ arr = ndarray(a)
1867
+ ext_loop = flags and 'external_loop' in flags
1868
+
1869
+ # external_loop: 按 chunk 返回一维数组
1870
+ if ext_loop:
1871
+ if arr.ndim <= 1:
1872
+ yield arr
1873
+ return
1874
+ if order == 'F':
1875
+ # F-order: 每次返回一列(最后一维)
1876
+ ncols = arr.shape[-1]
1877
+ for col in range(ncols):
1878
+ key = tuple([slice(None)] * (arr.ndim - 1) + [col])
1879
+ col_arr = arr[key]
1880
+ yield ndarray(col_arr)
1881
+ else:
1882
+ # C-order: 每次返回一行(第一维)
1883
+ for i in range(arr.shape[0]):
1884
+ yield ndarray(arr[i])
1885
+ return
1886
+
1887
+ # 原有逐元素逻辑
1888
+ flat_arr = arr.copy().flatten()
1889
+ size = flat_arr.size
1890
+ is_int = getattr(arr, '_dtype', 'float64') == 'int64'
1891
+ readwrite = op_flags and 'readwrite' in op_flags
1892
+
1893
+ def _to_val(v, ci):
1894
+ if readwrite:
1895
+ return _NdIterElement(arr, ci, int(v) if is_int else v)
1896
+ return int(v) if is_int else v
1897
+
1898
+ if order == 'F' and arr.ndim > 1:
1899
+ # Fortran-order: 按列序优先(首维最快)生成多维索引,转为 C-flat 索引
1900
+ shape = arr.shape
1901
+ ndim = len(shape)
1902
+ # C-order strides: stride[i] = product(shape[i+1:])
1903
+ c_strides = [1] * ndim
1904
+ for i in range(ndim - 2, -1, -1):
1905
+ c_strides[i] = c_strides[i + 1] * shape[i + 1]
1906
+
1907
+ # F-order 等同于最后一维最慢,首维最快
1908
+ # 递归遍历 dim = ndim-1, ndim-2, ..., 0
1909
+ def _iter_f_order(dim, coords):
1910
+ if dim < 0:
1911
+ # coords 是反向顺序 [last_dim, ..., first_dim],需反转
1912
+ idx = builtins.sum(c * s for c, s in zip(reversed(coords), c_strides))
1913
+ if 0 <= idx < size:
1914
+ val = flat_arr[idx].item()
1915
+ yield _to_val(val, idx)
1916
+ else:
1917
+ for v in range(shape[dim]):
1918
+ yield from _iter_f_order(dim - 1, coords + [v])
1919
+
1920
+ yield from _iter_f_order(ndim - 1, [])
1921
+ else:
1922
+ # C-order: 行序优先(默认)
1923
+ for i in range(size):
1924
+ val = flat_arr[i].item()
1925
+ yield _to_val(val, i)
1926
+
1927
+ def isnan(x):
1928
+ """逐元素检测是否为 NaN。"""
1929
+ arr = ndarray(x)
1930
+ return ndarray(_core.isnan(arr._array))
1931
+
1932
+
1933
+ def binary_repr(num, width=None):
1934
+ """返回整数的二进制表示字符串。"""
1935
+ return _core.binary_repr(num, width)
1936
+
1937
+
1938
+ def isinf(x):
1939
+ """逐元素检测是否为无穷大。"""
1940
+ arr = ndarray(x)
1941
+ return ndarray(_core.isinf(arr._array))
1942
+
1943
+
1944
+ def isfinite(x):
1945
+ """逐元素检测是否为有限值。"""
1946
+ arr = ndarray(x)
1947
+ return ndarray(_core.isfinite(arr._array))
1948
+
1949
+
1950
+ def iscomplex(x):
1951
+ """逐元素检测元素的虚部是否非零。"""
1952
+ arr = ndarray(x)
1953
+ cpx = getattr(arr, '_complex_data', None)
1954
+ if cpx is not None:
1955
+ return ndarray(_core.iscomplex_cpx(cpx))
1956
+ # 不是复数数组 → 全零
1957
+ shape = arr.shape
1958
+ return ndarray(_core.zeros(shape))
1959
+
1960
+
1961
+ def _broadcast_shape(*shapes):
1962
+ """计算多个数组的广播形状。"""
1963
+ max_ndim = builtins.max(len(s) for s in shapes)
1964
+ result = []
1965
+ for i in range(max_ndim):
1966
+ dims = set()
1967
+ for s in shapes:
1968
+ if i < max_ndim - len(s):
1969
+ dims.add(1)
1970
+ else:
1971
+ dims.add(s[i - (max_ndim - len(s))])
1972
+ non_one = [d for d in dims if d != 1]
1973
+ if len(non_one) > 1:
1974
+ raise ValueError("operands could not be broadcast together")
1975
+ result.append(non_one[0] if non_one else 1)
1976
+ return tuple(result)
1977
+
1978
+
1979
+ def broadcast(*args):
1980
+ """产生模仿广播的对象,将多个数组广播到相同形状。"""
1981
+ if not args:
1982
+ raise ValueError("broadcast requires at least one argument")
1983
+ return _Broadcast(*args)
1984
+
1985
+
1986
+ class _Broadcast:
1987
+ """广播对象,将多个数组广播到相同形状。"""
1988
+
1989
+ def __init__(self, *args):
1990
+ self._arrays = [ndarray(a) for a in args]
1991
+ # 计算广播形状
1992
+ shapes = [a.shape for a in self._arrays]
1993
+ self._shape = _broadcast_shape(*shapes)
1994
+ self._size = 1
1995
+ for s in self._shape:
1996
+ self._size *= s
1997
+
1998
+ @property
1999
+ def shape(self):
2000
+ """广播后的形状。"""
2001
+ return self._shape
2002
+
2003
+ @property
2004
+ def iters(self):
2005
+ """返回各数组的迭代器元组。"""
2006
+ return tuple(_BroadcastIter(a, self._shape) for a in self._arrays)
2007
+
2008
+ def __iter__(self):
2009
+ """迭代广播后的值元组。"""
2010
+ iters = self.iters
2011
+ for _ in range(self._size):
2012
+ yield tuple(next(it) for it in iters)
2013
+
2014
+
2015
+ class _BroadcastIter:
2016
+ """广播迭代器,按广播形状遍历单个数组。"""
2017
+
2018
+ def __init__(self, arr, shape):
2019
+ self._arr = arr
2020
+ self._flat = arr._array.flatten().tolist()
2021
+ self._size = len(self._flat)
2022
+ self._arr_shape = arr.shape
2023
+ self._broadcast_shape = shape
2024
+ # Left-pad array shape for broadcasting
2025
+ ndim_pad = len(shape) - len(arr.shape)
2026
+ self._padded = (1,) * ndim_pad + arr.shape
2027
+ self._indices = [0] * len(shape)
2028
+ self._done = False
2029
+
2030
+ def __iter__(self):
2031
+ return self
2032
+
2033
+ def __next__(self):
2034
+ if self._done:
2035
+ raise StopIteration
2036
+ # Map broadcast index to flat array index
2037
+ arr_idx = 0
2038
+ stride = 1
2039
+ for d in range(len(self._arr_shape) - 1, -1, -1):
2040
+ dim = self._indices[len(self._broadcast_shape) - len(self._arr_shape) + d]
2041
+ dim = dim % self._arr_shape[d]
2042
+ arr_idx += dim * stride
2043
+ stride *= self._arr_shape[d]
2044
+ val = self._flat[arr_idx]
2045
+
2046
+ # Advance to next C-order position
2047
+ for i in range(len(self._indices) - 1, -1, -1):
2048
+ self._indices[i] += 1
2049
+ if self._indices[i] < self._broadcast_shape[i]:
2050
+ break
2051
+ self._indices[i] = 0
2052
+ if i == 0:
2053
+ self._done = True
2054
+
2055
+ dtype = getattr(self._arr, '_dtype', 'float64')
2056
+ if dtype == 'int64':
2057
+ return int(val)
2058
+ return float(val)
2059
+
2060
+ def _ravel_c(self, *indices):
2061
+ """将多维索引转为扁平索引(C-order)。"""
2062
+ idx = 0
2063
+ stride = 1
2064
+ for i in range(len(self._arr_shape) - 1, -1, -1):
2065
+ idx += indices[i] * stride
2066
+ stride *= self._arr_shape[i]
2067
+ return idx
2068
+
2069
+
2070
+ # 数组操作函数
2071
+ reshape = _array_ops_module.reshape
2072
+ ravel = _array_ops_module.ravel
2073
+ moveaxis = _array_ops_module.moveaxis
2074
+ rollaxis = _array_ops_module.rollaxis
2075
+ broadcast_to = _array_ops_module.broadcast_to
2076
+ transpose = _array_ops_module.transpose
2077
+ swapaxes = _array_ops_module.swapaxes
2078
+ expand_dims = _array_ops_module.expand_dims
2079
+ squeeze = _array_ops_module.squeeze
2080
+ concatenate = _array_ops_module.concatenate
2081
+ stack = _array_ops_module.stack
2082
+ vstack = _array_ops_module.vstack
2083
+ hstack = _array_ops_module.hstack
2084
+ dstack = _array_ops_module.dstack
2085
+ column_stack = _array_ops_module.column_stack
2086
+ split = _array_ops_module.split
2087
+ hsplit = _array_ops_module.hsplit
2088
+ vsplit = _array_ops_module.vsplit
2089
+ dsplit = _array_ops_module.dsplit
2090
+ tile = _array_ops_module.tile
2091
+ repeat = _array_ops_module.repeat
2092
+ flip = _array_ops_module.flip
2093
+ fliplr = _array_ops_module.fliplr
2094
+ flipud = _array_ops_module.flipud
2095
+ roll = _array_ops_module.roll
2096
+ rot90 = _array_ops_module.rot90
2097
+ append = _array_ops_module.append
2098
+ insert = _array_ops_module.insert
2099
+ delete = _array_ops_module.delete
2100
+ unique = _array_ops_module.unique
2101
+ resize = _array_ops_module.resize
2102
+
2103
+ # 位运算
2104
+ def _to_raw(a):
2105
+ """将输入转为 Rust ndarray,返回 (ndarray, dtype) 元组。"""
2106
+ if hasattr(a, '_array'):
2107
+ return a._array, getattr(a, '_dtype', 'float64')
2108
+ raw = _core.ndarray(a)
2109
+ if isinstance(a, int) and not isinstance(a, bool):
2110
+ return raw, 'int64'
2111
+ return raw, 'float64'
2112
+
2113
+ def bitwise_and(x1, x2):
2114
+ """按位与"""
2115
+ r1, dt = _to_raw(x1)
2116
+ r2, _ = _to_raw(x2)
2117
+ return _wrap_result(_core.bitwise_and(r1, r2), dt)
2118
+
2119
+ def bitwise_or(x1, x2):
2120
+ """按位或"""
2121
+ r1, dt = _to_raw(x1)
2122
+ r2, _ = _to_raw(x2)
2123
+ return _wrap_result(_core.bitwise_or(r1, r2), dt)
2124
+
2125
+ def bitwise_xor(x1, x2):
2126
+ """按位异或"""
2127
+ r1, dt = _to_raw(x1)
2128
+ r2, _ = _to_raw(x2)
2129
+ return _wrap_result(_core.bitwise_xor(r1, r2), dt)
2130
+
2131
+ def bitwise_not(x):
2132
+ """按位取反(对布尔数组使用逻辑取反)"""
2133
+ r, dt = _to_raw(x)
2134
+ if dt == 'bool':
2135
+ r = _core.invert(r)
2136
+ elif dt == 'uint8':
2137
+ raw = _core.bitwise_not(r)
2138
+ r = _core.bitwise_and(raw, _core.ndarray([255]))
2139
+ dt = 'uint8'
2140
+ else:
2141
+ r = _core.bitwise_not(r)
2142
+ return _wrap_result(r, dt)
2143
+
2144
+
2145
+ def invert(x):
2146
+ """按位取反(等效于 ~ 运算符,对布尔数组使用逻辑取反)"""
2147
+ r, dt = _to_raw(x)
2148
+ if dt == 'bool':
2149
+ r = _core.invert(r)
2150
+ elif dt == 'uint8':
2151
+ raw = _core.bitwise_not(r)
2152
+ r = _core.bitwise_and(raw, _core.ndarray([255]))
2153
+ dt = 'uint8'
2154
+ else:
2155
+ r = _core.bitwise_not(r)
2156
+ return _wrap_result(r, dt)
2157
+
2158
+ def left_shift(x1, x2):
2159
+ """左移"""
2160
+ r1, dt = _to_raw(x1)
2161
+ r2, _ = _to_raw(x2)
2162
+ return _wrap_result(_core.left_shift(r1, r2), dt)
2163
+
2164
+ def right_shift(x1, x2):
2165
+ """右移"""
2166
+ r1, dt = _to_raw(x1)
2167
+ r2, _ = _to_raw(x2)
2168
+ return _wrap_result(_core.right_shift(r1, r2), dt)
2169
+
2170
+
2171
+ # 数学函数
2172
+ sin = _math_functions_module.sin
2173
+ cos = _math_functions_module.cos
2174
+ tan = _math_functions_module.tan
2175
+ arcsin = _math_functions_module.arcsin
2176
+ arccos = _math_functions_module.arccos
2177
+ arctan = _math_functions_module.arctan
2178
+ arctan2 = _math_functions_module.arctan2
2179
+ deg2rad = _math_functions_module.deg2rad
2180
+ rad2deg = _math_functions_module.rad2deg
2181
+ degrees = rad2deg
2182
+ sinh = _math_functions_module.sinh
2183
+ cosh = _math_functions_module.cosh
2184
+ tanh = _math_functions_module.tanh
2185
+ asinh = _math_functions_module.asinh
2186
+ acosh = _math_functions_module.acosh
2187
+ atanh = _math_functions_module.atanh
2188
+ exp = _math_functions_module.exp
2189
+ expm1 = _math_functions_module.expm1
2190
+ log = _math_functions_module.log
2191
+ log10 = _math_functions_module.log10
2192
+ log2 = _math_functions_module.log2
2193
+ log1p = _math_functions_module.log1p
2194
+ around = _math_functions_module.around
2195
+ floor = _math_functions_module.floor
2196
+ ceil = _math_functions_module.ceil
2197
+ trunc = _math_functions_module.trunc
2198
+ fix = _math_functions_module.fix
2199
+ sqrt = _math_functions_module.sqrt
2200
+ square = _math_functions_module.square
2201
+ cbrt = _math_functions_module.cbrt
2202
+ abs = _math_functions_module.abs
2203
+ sign = _math_functions_module.sign
2204
+ reciprocal = _math_functions_module.reciprocal
2205
+ clip = _math_functions_module.clip
2206
+ sinc = _math_functions_module.sinc
2207
+ heaviside = _math_functions_module.heaviside
2208
+ add = _math_functions_module.add
2209
+ subtract = _math_functions_module.subtract
2210
+ multiply = _math_functions_module.multiply
2211
+ divide = _math_functions_module.divide
2212
+ power = _math_functions_module.power
2213
+ mod = _math_functions_module.mod
2214
+ remainder = _math_functions_module.remainder
2215
+ greater = _math_functions_module.greater
2216
+ less = _math_functions_module.less
2217
+ equal = _math_functions_module.equal
2218
+ not_equal = _math_functions_module.not_equal
2219
+ greater_equal = _math_functions_module.greater_equal
2220
+ less_equal = _math_functions_module.less_equal
2221
+ logical_and = _math_functions_module.logical_and
2222
+ logical_or = _math_functions_module.logical_or
2223
+ logical_xor = _math_functions_module.logical_xor
2224
+ isclose = _math_functions_module.isclose
2225
+ allclose = _math_functions_module.allclose
2226
+
2227
+ # 统计函数
2228
+ sum = _statistics_module.sum
2229
+ mean = _statistics_module.mean
2230
+ std = _statistics_module.std
2231
+ var = _statistics_module.var
2232
+ min = _statistics_module.min
2233
+ max = _statistics_module.max
2234
+ amin = _statistics_module.amin
2235
+ amax = _statistics_module.amax
2236
+ ptp = _statistics_module.ptp
2237
+ median = _statistics_module.median
2238
+ average = _statistics_module.average
2239
+ percentile = _statistics_module.percentile
2240
+ quantile = _statistics_module.quantile
2241
+ nanmedian = _statistics_module.nanmedian
2242
+ nanpercentile = _statistics_module.nanpercentile
2243
+ argmax = _statistics_module.argmax
2244
+ argmin = _statistics_module.argmin
2245
+ argsort = _statistics_module.argsort
2246
+ sort = _statistics_module.sort
2247
+ msort = _statistics_module.msort
2248
+ sort_complex = _statistics_module.sort_complex
2249
+ lexsort = _statistics_module.lexsort
2250
+ partition = _statistics_module.partition
2251
+ argpartition = _statistics_module.argpartition
2252
+ searchsorted = _statistics_module.searchsorted
2253
+ extract = _statistics_module.extract
2254
+ cov = _statistics_module.cov
2255
+ corrcoef = _statistics_module.corrcoef
2256
+ histogram = _statistics_module.histogram
2257
+ histogram2d = _statistics_module.histogram2d
2258
+ histogramdd = _statistics_module.histogramdd
2259
+ digitize = _statistics_module.digitize
2260
+
2261
+ # 子模块
2262
+
2263
+ linalg = _linalg_module()
2264
+ random = _random_module()
2265
+ matlib = _matlib_module
2266
+
2267
+ dot = linalg.dot
2268
+ vdot = linalg.vdot
2269
+ inner = linalg.inner
2270
+ matmul = linalg.matmul
2271
+
2272
+
2273
+ # ========== 导出列表 ==========
2274
+ __all__ = [
2275
+ 'ndarray', 'NdArrayIter',
2276
+ 'array', 'asarray', 'asanyarray', 'copy',
2277
+ 'matrix', 'asmatrix', 'bmat', 'repmat',
2278
+ 'dot', 'vdot', 'inner', 'matmul',
2279
+ 'zeros', 'zeros_like', 'ones', 'ones_like', 'full', 'full_like',
2280
+ 'empty', 'empty_like', 'eye', 'identity',
2281
+ 'arange', 'linspace', 'logspace', 'geomspace',
2282
+ 'fromfunction', 'frombuffer', 'r_', 's_', 'mgrid', 'ogrid',
2283
+ 'reshape', 'ravel', 'moveaxis', 'rollaxis', 'broadcast_to',
2284
+ 'transpose', 'swapaxes', 'expand_dims', 'squeeze',
2285
+ 'concatenate', 'stack', 'vstack', 'hstack', 'dstack', 'column_stack',
2286
+ 'split', 'hsplit', 'vsplit', 'dsplit',
2287
+ 'tile', 'repeat',
2288
+ 'flip', 'fliplr', 'flipud', 'roll', 'rot90',
2289
+ 'append', 'insert', 'delete', 'unique', 'resize',
2290
+ 'where', 'take', 'put', 'select', 'nonzero', 'argwhere', 'flatnonzero',
2291
+ 'binary_repr',
2292
+ 'sin', 'cos', 'tan', 'arcsin', 'arccos', 'arctan', 'arctan2',
2293
+ 'deg2rad', 'rad2deg', 'degrees',
2294
+ 'sinh', 'cosh', 'tanh', 'asinh', 'acosh', 'atanh',
2295
+ 'bitwise_and', 'bitwise_or', 'bitwise_xor', 'bitwise_not',
2296
+ 'invert', 'left_shift', 'right_shift',
2297
+ 'string_', 'unicode_', 'char',
2298
+ 'exp', 'expm1', 'log', 'log10', 'log2', 'log1p',
2299
+ 'around', 'floor', 'ceil', 'trunc', 'fix',
2300
+ 'sqrt', 'square', 'cbrt', 'abs', 'sign', 'reciprocal', 'clip', 'sinc', 'heaviside',
2301
+ 'add', 'subtract', 'multiply', 'divide', 'power', 'mod', 'remainder',
2302
+ 'greater', 'less', 'equal', 'logical_and', 'logical_or', 'isclose', 'allclose',
2303
+ 'sum', 'mean', 'std', 'var', 'min', 'max', 'amin', 'amax', 'ptp',
2304
+ 'median', 'average', 'percentile', 'quantile', 'nanmedian', 'nanpercentile',
2305
+ 'argmax', 'argmin', 'argsort', 'sort', 'searchsorted', 'extract',
2306
+ 'cov', 'corrcoef',
2307
+ 'histogram', 'histogram2d', 'histogramdd', 'digitize',
2308
+ 'fft', 'ifft', 'rfft', 'irfft',
2309
+ 'pi', 'e', 'euler_gamma', 'inf', 'nan', 'newaxis',
2310
+ 'nditer',
2311
+ 'isnan', 'isinf', 'isfinite',
2312
+ 'save', 'load', 'loadtxt', 'savetxt', 'savez',
2313
+ 'Poly', 'polyval', 'polyfit', 'polyder', 'polyint', 'polyroots',
2314
+ 'linalg', 'random', 'matlib', 'load_npz'
2315
+ ]