Coverage for src/lsqfitgp/_array.py: 93%

1# lsqfitgp/_array.py 

2# 

3# Copyright (c) 2020, 2022, 2023, 2024, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20import math 3 ctx1cba

21import textwrap 3 ctx1cba

22 

23import jax 3 ctx1cba

24import numpy 3 ctx1cba

25from jax import numpy as jnp 3 ctx1cba

26from jax import tree, tree_util 3 ctx1cba

27from numpy.lib import recfunctions 3 ctx1cba

28 

29 

30@tree_util.register_pytree_node_class 3 ctx1cba

31class StructuredArray: 3 ctx1cba

32 """ 3 ctx1cba

33 JAX-friendly imitation of a numpy structured array. 

34 

35 It behaves like a read-only numpy structured array, and you can create 

36 a copy with a modified field with a jax-like syntax. 

37 

38 Parameters 

39 ---------- 

40 array : numpy array, StructuredArray 

41 A structured array. An array qualifies as structured if 

42 ``array.dtype.names is not None``. 

43 

44 Notes 

45 ----- 

46 The StructuredArray is a readonly view on the input array. When you 

47 change the content of a field of the StructuredArray, however, the 

48 reference to the original array for that field is lost. 

49 

50 Examples 

51 -------- 

52 >>> a = numpy.empty(3, dtype=[('f', float), ('g', float)]) 

53 >>> a = StructuredArray(a) 

54 >>> a = a.at['f'].set(numpy.arange(3)) 

55 ... # is equivalent to a['f'] = numpy.arange(3) 

56 """ 

57 

58 @classmethod 3 ctx1cba

59 def _readonlyview_wrapifstructured(cls, x): 3 ctx1cba

60 if x.dtype.names is not None: 3 ctx1cba

61 x = cls(x) 2 ctx1ba

62 if isinstance(x, numpy.ndarray): 3 ctx1cba

63 x = x.view() 3 ctx1cba

64 x.flags.writeable = False 3 ctx1cba

65 # jax arrays and StructuredArrays are already readonly 

66 return x 3 ctx1cba

67 

68 @classmethod 3 ctx1cba

69 def _array(cls, s, t, d, *, check=True): 3 ctx1cba

70 """ 

71 Create a new StructuredArray. 

72 

73 All methods and functions that create a new StructuredArray object 

74 should use this method. 

75 

76 Parameters 

77 ---------- 

78 s : tuple or None 

79 The shape. If None, it is determined automatically from the arrays. 

80 t : dtype or None 

81 The dtype of the array. If None, it is determined automatically 

82 (before the shape). 

83 d : dict str -> array 

84 The _dict of the array. The arrays, if structured, must be already 

85 StructuredArrays. The order of the keys must match the order of the 

86 fields. 

87 check : bool 

88 If True (default), check the passed values are consistent. 

89 

90 Return 

91 ------ 

92 out : StructuredArray 

93 A new StructuredArray object. 

94 """ 

95 if t is None: 3 ctx1cba

96 # infer the data type from the arrays in the dictionary 

97 ndim = min((x.ndim for x in d.values()), default=None) 2 ctx1ba

98 t = numpy.dtype([(name, x.dtype, x.shape[ndim:]) for name, x in d.items()]) 2 ctx1ba

99 

100 # remove offset info since this is actually a columnar format 

101 t = recfunctions.repack_fields(t, align=False, recurse=True) 3 ctx1cba

102 

103 if s is None: 3 ctx1cba

104 # infer the shape from the arrays in the dictionary 

105 assert d, 'can not infer array shape with no fields' 3 ctx1cba

106 f = t.names[0] 3 ctx1cba

107 a = d[f] 3 ctx1cba

108 s = a.shape[: a.ndim - t[0].ndim] 3 ctx1cba

109 

110 if check: 110 ↛ 119line 110 didn't jump to line 119 because the condition on line 110 was always true3 ctx1cba

111 assert len(t) == len(d) 3 ctx1cba

112 assert t.names == tuple(d.keys()) 3 ctx1cba

113 assert all( 3 ctx1cba

114 x.dtype == t[f].base and x.ndim >= t[f].ndim for f, x in d.items() 

115 ) 

116 shapes = [x.shape[: x.ndim - t[f].ndim] for f, x in d.items()] 3 ctx1cba

117 assert all(s == s1 for s1 in shapes) 3 ctx1cba

118 

119 out = super().__new__(cls) 3 ctx1cba

120 out.shape = s 3 ctx1cba

121 out.dtype = t 3 ctx1cba

122 out._dict = d 3 ctx1cba

123 

124 return out 3 ctx1cba

125 

126 def __new__(cls, array): 3 ctx1cba

127 if isinstance(array, cls): 3 ctx1cba

128 return array 3 ctx1cba

129 d = { 3 ctx1cba

130 name: cls._readonlyview_wrapifstructured(array[name]) 

131 for name in array.dtype.names 

132 } 

133 return cls._array(array.shape, array.dtype, d) 3 ctx1cba

134 

135 @classmethod 3 ctx1cba

136 def from_dataframe(cls, df): 3 ctx1cba

137 """ 

138 Make a StructuredArray from a DataFrame. 

139 

140 Data is not copied if not necessary. 

141 """ 

142 d = { 2 ctx1ba

143 col: cls._readonlyview_wrapifstructured(df[col].to_numpy()) 

144 for col in df.columns 

145 } 

146 return cls._array(None, None, d) 2 ctx1ba

147 

148 @classmethod 3 ctx1cba

149 def from_dict(cls, mapping): 3 ctx1cba

150 """Make a StructuredArray from a dictionary of arrays. Data is not copied.""" 

151 d = { 2 ctx1ba

152 name: cls._readonlyview_wrapifstructured(value) 

153 for name, value in mapping.items() 

154 } 

155 return cls._array(None, None, d) 2 ctx1ba

156 

157 @property 3 ctx1cba

158 def size(self): 3 ctx1cba

159 return math.prod(self.shape) 2 ctx1ba

160 

161 @property 3 ctx1cba

162 def ndim(self): 3 ctx1cba

163 return len(self.shape) 2 ctx1ba

164 

165 @property 3 ctx1cba

166 def nbytes(self): 3 ctx1cba

167 return sum(x.nbytes for x in self._dict.values()) 

168 

169 @property 3 ctx1cba

170 def T(self): 3 ctx1cba

171 if self.ndim < 2: 1 ctx1a

172 return self 1 ctx1a

173 return self.swapaxes(self.ndim - 2, self.ndim - 1) 1 ctx1a

174 

175 def swapaxes(self, i, j): 3 ctx1cba

176 shape = jax.eval_shape(lambda: jnp.empty(self.shape).swapaxes(i, j)).shape 1 ctx1a

177 d = {k: v.swapaxes(i, j) for k, v in self._dict.items()} 1 ctx1a

178 return self._array(shape, self.dtype, d) 1 ctx1a

179 

180 def __len__(self): 3 ctx1cba

181 if self.shape: 2 ctx1ba

182 return self.shape[0] 2 ctx1ba

183 else: 

184 msg = 'len() of unsized object' 1 ctx1a

185 raise TypeError(msg) 1 ctx1a

186 

187 def __getitem__(self, key): 3 ctx1cba

188 if isinstance(key, str): 3 ctx1cba

189 return self._dict[key] 3 ctx1cba

190 elif isinstance(key, list) and key and all(isinstance(k, str) for k in key): 3 ctx1cba

191 d = {name: self._dict[name] for name in key} 1 ctx1a

192 return self._array(self.shape, self.dtype[key], d) 1 ctx1a

193 else: 

194 d = { 3 ctx1cba

195 name: x[ 

196 (key if isinstance(key, tuple) else (key,)) 

197 + (slice(None),) * self.dtype[name].ndim 

198 ] 

199 for name, x in self._dict.items() 

200 } 

201 shape = jax.eval_shape(lambda: jnp.empty(self.shape)[key]).shape 3 ctx1cba

202 return self._array(shape, self.dtype, d) 3 ctx1cba

203 

204 @property 3 ctx1cba

205 def at(self): 3 ctx1cba

206 return self._Getter(self) 3 ctx1cba

207 

208 class _Getter: 3 ctx1cba

209 def __init__(self, array): 3 ctx1cba

210 self.array = array 3 ctx1cba

211 

212 def __getitem__(self, key): 3 ctx1cba

213 if key not in self.array.dtype.names: 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true3 ctx1cba

214 raise KeyError(key) 

215 return self.Setter(self.array, key) 3 ctx1cba

216 

217 class Setter: 3 ctx1cba

218 def __init__(self, array, key, parent=None): 3 ctx1cba

219 self.array = array 3 ctx1cba

220 self.key = key 3 ctx1cba

221 self.parent = parent 3 ctx1cba

222 

223 def __getitem__(self, subkey): 3 ctx1cba

224 if subkey not in self.array.dtype[self.key].names: 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true1 ctx1a

225 raise KeyError(subkey) 

226 return self.__class__(self.array[self.key], subkey, self) 1 ctx1a

227 

228 def set(self, val): 3 ctx1cba

229 assert isinstance(val, (numpy.ndarray, jnp.ndarray, StructuredArray)) 3 ctx1cba

230 prev = self.array._dict[self.key] 3 ctx1cba

231 assert prev.dtype == val.dtype 3 ctx1cba

232 assert prev.shape == val.shape 3 ctx1cba

233 d = dict(self.array._dict) 3 ctx1cba

234 d[self.key] = self.array._readonlyview_wrapifstructured(val) 3 ctx1cba

235 out = self.array._array(self.array.shape, self.array.dtype, d) 3 ctx1cba

236 if self.parent: 3 ctx1cba

237 return self.parent.set(out) 1 ctx1a

238 else: 

239 return out 3 ctx1cba

240 

241 def reshape(self, *shape): 3 ctx1cba

242 """ 

243 Reshape the array without changing its contents. 

244 

245 See numpy.ndarray.reshape. 

246 """ 

247 if len(shape) == 1 and hasattr(shape[0], '__len__'): 3 ctx1cba

248 shape = shape[0] 2 ctx1ba

249 shape = tuple(shape) 3 ctx1cba

250 d = { 3 ctx1cba

251 name: x.reshape(shape + self.dtype[name].shape) 

252 for name, x in self._dict.items() 

253 } 

254 shape = numpy.empty(self.shape, []).reshape(shape).shape 3 ctx1cba

255 return self._array(shape, self.dtype, d) 3 ctx1cba

256 

257 def squeeze(self, axis=None): 3 ctx1cba

258 """Remove axes of length 1. See numpy.ndarray.squeeze.""" 

259 if axis is None: 2 ctx1ba

260 axis = tuple(i for i, size in enumerate(self.shape) if size == 1) 1 ctx1a

261 if not hasattr(axis, '__len__'): 2 ctx1ba

262 axis = (axis,) 2 ctx1ba

263 assert all(self.shape[i] == 1 for i in axis) 2 ctx1ba

264 newshape = [size for i, size in enumerate(self.shape) if i not in axis] 2 ctx1ba

265 return self.reshape(newshape) 2 ctx1ba

266 

267 def astype(self, dtype): 3 ctx1cba

268 if dtype != self.dtype: 

269 raise NotImplementedError 

270 return self 

271 

272 def broadcast_to(self, shape, **kw): 3 ctx1cba

273 """ 

274 Return a view of the array broadcasted to another shape. 

275 

276 See numpy.broadcast_to. 

277 """ 

278 # raises if not broadcastable 

279 numpy.broadcast_to(numpy.empty(self.shape, []), shape, **kw) 1 ctx1a

280 d = { 1 ctx1a

281 name: broadcast_to(x, shape + self.dtype[name].shape, **kw) 

282 for name, x in self._dict.items() 

283 } 

284 return self._array(shape, self.dtype, d) 1 ctx1a

285 

286 def tree_flatten(self): 3 ctx1cba

287 """JAX PyTree encoder. See `jax.tree_util.tree_flatten`.""" 

288 children = tuple(self._dict[key] for key in self.dtype.names) 3 ctx1cba

289 aux = dict(shape=self.shape, dtype=self.dtype) 3 ctx1cba

290 return children, aux 3 ctx1cba

291 

292 @classmethod 3 ctx1cba

293 def tree_unflatten(cls, aux, children): 3 ctx1cba

294 """JAX PyTree decoder. See `jax.tree_util.tree_unflatten`.""" 

295 # if there are no fields, keep original shape 

296 if not children: 296 ↛ 297line 296 didn't jump to line 297 because the condition on line 296 was never true3 ctx1cba

297 return cls._array(aux['shape'], aux['dtype'], {}) 

298 

299 # convert children to arrays because tree.flatten unpacks 0d 

300 # arrays 

301 children = list(map(asarray, children)) 3 ctx1cba

302 

303 # if possible, keep original dtype shapes 

304 oldtype = aux['dtype'] 3 ctx1cba

305 compatible_tail_shapes = all( 3 ctx1cba

306 x.shape[max(0, x.ndim - oldtype[i].ndim) :] == oldtype[i].shape 

307 for i, x in enumerate(children) 

308 ) 

309 head_shapes = [ 3 ctx1cba

310 x.shape[: max(0, x.ndim - oldtype[i].ndim)] for i, x in enumerate(children) 

311 ] 

312 compatible_head_shapes = all(head_shapes[0] == s for s in head_shapes) 3 ctx1cba

313 if compatible_tail_shapes and compatible_head_shapes: 3 ctx1cba

314 dtype = numpy.dtype( 3 ctx1cba

315 [ 

316 (oldtype.names[i], x.dtype, oldtype[i].shape) 

317 for i, x in enumerate(children) 

318 ] 

319 ) 

320 else: 

321 dtype = None 1 ctx1a

322 

323 d = dict(zip(oldtype.names, children, strict=True)) 3 ctx1cba

324 

325 return cls._array(None, dtype, d) 3 ctx1cba

326 

327 def __repr__(self): 3 ctx1cba

328 # code from gvar https://github.com/gplepage/gvar 

329 # bufferdict.pyx:BufferDict:__str__ 

330 out = 'StructuredArray({' 1 ctx1a

331 

332 listrepr = [(repr(k), repr(v)) for k, v in self._dict.items()] 1 ctx1a

333 newlinemode = any('\n' in rv for _, rv in listrepr) 1 ctx1a

334 

335 for rk, rv in listrepr: 1 ctx1a

336 if not newlinemode: 1 ctx1a

337 out += f'{rk}: {rv}, ' 1 ctx1a

338 elif '\n' in rv: 1 ctx1a

339 rv = rv.replace('\n', '\n ') 1 ctx1a

340 out += f'\n {rk}:\n {rv},' 1 ctx1a

341 else: 

342 out += f'\n {rk}: {rv},' 1 ctx1a

343 

344 if out.endswith(', '): 1 ctx1a

345 out = out[:-2] 1 ctx1a

346 elif newlinemode: 1 ctx1a

347 out += '\n' 1 ctx1a

348 out += '})' 1 ctx1a

349 

350 return out 1 ctx1a

351 

352 def __array__(self, copy=None, dtype=None): 3 ctx1cba

353 if copy is False: 353 ↛ 354line 353 didn't jump to line 354 because the condition on line 353 was never true1 ctx1a

354 msg = 'StructuredArray has to be copied when converted to a numpy array' 

355 raise ValueError(msg) 

356 if dtype is not None: 356 ↛ 357line 356 didn't jump to line 357 because the condition on line 356 was never true1 ctx1a

357 dtype = numpy.dtype(dtype) 

358 if dtype != self.dtype: 

359 msg = 'StructuredArray can not be converted to a numpy array with a different dtype' 

360 raise ValueError(msg) 

361 array = numpy.empty(self.shape, self.dtype) 1 ctx1a

362 self._copy_into_array(array) 1 ctx1a

363 return array 1 ctx1a

364 

365 def _copy_into_array(self, dest): 3 ctx1cba

366 assert self.dtype == dest.dtype 1 ctx1a

367 assert self.shape == dest.shape 1 ctx1a

368 for name, src in self._dict.items(): 1 ctx1a

369 if isinstance(src, StructuredArray): 1 ctx1a

370 src._copy_into_array(dest[name]) 1 ctx1a

371 else: 

372 dest[name][...] = src 1 ctx1a

373 

374 def __array_function__(self, func, types, args, kwargs): 3 ctx1cba

375 if func not in self._handled_functions: 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true2 ctx1ba

376 return NotImplemented 

377 return self._handled_functions[func](*args, **kwargs) 2 ctx1ba

378 

379 _handled_functions = {} # noqa: RUF012, registry shared by design 3 ctx1cba

380 

381 @classmethod 3 ctx1cba

382 def _implements(cls, np_function): 3 ctx1cba

383 """Register an __array_function__ implementation.""" 

384 

385 def decorator(func): 3 ctx1cba

386 cls._handled_functions[np_function] = func 3 ctx1cba

387 newdoc = f"""\ 3 ctx1cba

388Implementation of `{np_function.__module__}.{np_function.__name__}` for `StructuredArray`. 

389 

390""" 

391 if func.__doc__: 3 ctx1cba

392 newdoc += textwrap.dedent(func.__doc__) + '\n' 3 ctx1cba

393 newdoc += 'Original docstring below:\n\n' 3 ctx1cba

394 newdoc += textwrap.dedent(np_function.__doc__) 3 ctx1cba

395 func.__doc__ = newdoc 3 ctx1cba

396 return func 3 ctx1cba

397 

398 return decorator 3 ctx1cba

399 

400 

401@StructuredArray._implements(numpy.broadcast_to) 3 ctx1cba

402def broadcast_to(x, shape, **kw): 3 ctx1cba

403 """Version of numpy.broadcast_to for StructuredArray and JAX arrays.""" 

404 if isinstance(x, StructuredArray): 1 ctx1a

405 return x.broadcast_to(shape, **kw) 1 ctx1a

406 elif isinstance(x, jnp.ndarray): 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true1 ctx1a

407 return jnp.broadcast_to(x, shape, **kw) 

408 else: 

409 return numpy.broadcast_to(x, shape, **kw) 1 ctx1a

410 

411 

412@StructuredArray._implements(numpy.broadcast_arrays) 3 ctx1cba

413def broadcast_arrays(*arrays, **kw): 3 ctx1cba

414 """Version of numpy.broadcast_arrays for StructuredArray and JAX arrays.""" 

415 shapes = [a.shape for a in arrays] 1 ctx1a

416 shape = numpy.broadcast_shapes(*shapes) 1 ctx1a

417 return [broadcast_to(a, shape, **kw) for a in arrays] 1 ctx1a

418 # numpy.broadcast_arrays returns a list, not a tuple 

419 

420 

421class broadcast: 3 ctx1cba

422 """Version of numpy.broadcast that works with StructuredArray.""" 

423 

424 # not handled by __array_function__ 

425 

426 def __init__(self, *arrays): 3 ctx1cba

427 self.shape = numpy.broadcast_shapes(*(a.shape for a in arrays)) 3 ctx1cba

428 

429 

430def asarray(x, dtype=None): 3 ctx1cba

431 """ 

432 Version of `numpy.asarray` that works with `StructuredArray` and JAX arrays. 

433 

434 If `x` is not an array already, returns a JAX array if possible. 

435 """ 

436 if isinstance(x, (StructuredArray, jnp.ndarray, numpy.ndarray)): 3 ctx1cba

437 return x if dtype is None else x.astype(dtype) 3 ctx1cba

438 if x is None: 3 ctx1cba

439 return numpy.asarray(x, dtype) 1 ctx1a

440 # partial workaround for jax issue #14506, None would be interpreted as 

441 # nan by jax 

442 try: 3 ctx1cba

443 return jnp.asarray(x, dtype) 3 ctx1cba

444 except (TypeError, ValueError): 2 ctx1ca

445 return numpy.asarray(x, dtype) 2 ctx1ca

446 

447 

448def _asarray_jaxifpossible(x): 3 ctx1cba

449 x = asarray(x) 3 ctx1cba

450 if x.dtype.names: 3 ctx1cba

451 return tree.map(_asarray_jaxifpossible, StructuredArray(x)) 3 ctx1cba

452 if isinstance(x, numpy.ndarray): 3 ctx1cba

453 try: 3 ctx1cba

454 return jnp.asarray(x) 3 ctx1cba

455 except (TypeError, ValueError): 3 ctx1cba

456 pass 3 ctx1cba

457 return x 3 ctx1cba

458 

459 

460@StructuredArray._implements(numpy.squeeze) 3 ctx1cba

461def _squeeze(a, axis=None): 3 ctx1cba

462 return a.squeeze(axis) 1 ctx1a

463 

464 

465@StructuredArray._implements(numpy.ix_) 3 ctx1cba

466def _ix(*args): 3 ctx1cba

467 args = tuple(map(asarray, args)) 1 ctx1a

468 assert all(x.ndim == 1 for x in args) 1 ctx1a

469 n = len(args) 1 ctx1a

470 return tuple( 1 ctx1a

471 x.reshape((1,) * i + (-1,) + (1,) * (n - i - 1)) for i, x in enumerate(args) 

472 ) 

473 

474 

475def unstructured_to_structured( 3 ctx1cba

476 arr, dtype=None, names=None, align=False, copy=False, casting='unsafe' 

477): 

478 """ 

479 Convert an unstructured array to a `StructuredArray`. 

480 

481 Like `numpy.lib.recfunctions.unstructured_to_structured`, but outputs a 

482 `StructuredArray`. 

483 """ 

484 arr = asarray(arr) 2 ctx1ba

485 if not arr.ndim: 2 ctx1ba

486 msg = 'arr must have at least one dimension' 1 ctx1a

487 raise ValueError(msg) 1 ctx1a

488 mockup = numpy.empty((0, *arr.shape[-1:]), arr.dtype) 2 ctx1ba

489 dummy = recfunctions.unstructured_to_structured( 2 ctx1ba

490 mockup, dtype=dtype, names=names, align=align, copy=copy, casting=casting 

491 ) 

492 out, length = _unstructured_to_structured_recursive( 2 ctx1ba

493 0, (), arr, dummy.dtype, copy, casting 

494 ) 

495 assert length == arr.shape[-1] 2 ctx1ba

496 return out 2 ctx1ba

497 

498 

499def _unstructured_to_structured_recursive( 3 ctx1cba

500 idx, shape, arr, dtype, copy, casting, *strides 

501): 

502 arrays = {} 2 ctx1ba

503 for i, name in enumerate(dtype.names): 2 ctx1ba

504 base = dtype[i].base 2 ctx1ba

505 subshape = shape + dtype[i].shape 2 ctx1ba

506 size = math.prod(dtype[i].shape) 2 ctx1ba

507 stride = _nd(base) 2 ctx1ba

508 substrides = (*strides, (size, stride)) 2 ctx1ba

509 if base.names is not None: 2 ctx1ba

510 y, newidx = _unstructured_to_structured_recursive( 2 ctx1ba

511 idx, subshape, arr, base, copy, casting, *substrides 

512 ) 

513 shift = newidx - idx 2 ctx1ba

514 assert shift == stride 2 ctx1ba

515 idx += size * stride 2 ctx1ba

516 else: 

517 assert stride == 1 2 ctx1ba

518 if all(size == 1 for size, _ in strides): 2 ctx1ba

519 indices = numpy.s_[idx : idx + size] 2 ctx1ba

520 else: 

521 indices = sum( 1 ctx1a

522 ( 

523 stride * numpy.arange(size)[numpy.s_[:,] + (None,) * i] 

524 for i, (size, stride) in enumerate(reversed(substrides)) 

525 ), 

526 start=idx, 

527 ) 

528 indices = indices.reshape(-1) 1 ctx1a

529 key = numpy.s_[..., indices] 2 ctx1ba

530 x = arr[key] 2 ctx1ba

531 x = x.reshape(arr.shape[:-1] + subshape) 2 ctx1ba

532 if isinstance(x, jnp.ndarray): 2 ctx1ba

533 y = x.astype(base) 2 ctx1ba

534 else: 

535 y = x.astype(base, copy=copy, casting=casting) 1 ctx1a

536 idx += size 2 ctx1ba

537 arrays[name] = y 2 ctx1ba

538 return StructuredArray._array(arr.shape[:-1] + shape, dtype, arrays), idx 2 ctx1ba

539 

540 

541@StructuredArray._implements(recfunctions.structured_to_unstructured) 3 ctx1cba

542def _structured_to_unstructured(arr, dtype=None, casting='unsafe'): 3 ctx1cba

543 mockup = numpy.empty(0, arr.dtype) 2 ctx1ba

544 dummy = recfunctions.structured_to_unstructured( 2 ctx1ba

545 mockup, dtype=dtype, casting=casting 

546 ) 

547 args = (arr.shape + dummy.shape[-1:], dummy.dtype) 2 ctx1ba

548 try: 2 ctx1ba

549 out = jnp.empty(*args) 2 ctx1ba

550 except TypeError: 1 ctx1a

551 out = numpy.empty(*args) 1 ctx1a

552 out, length = _structured_to_unstructured_recursive(0, arr, out) 2 ctx1ba

553 assert length == dummy.shape[-1] 2 ctx1ba

554 return out 2 ctx1ba

555 

556 

557def _nd(dtype): 3 ctx1cba

558 """Count the number of scalars in a dtype.""" 

559 base = dtype.base 3 ctx1cba

560 shape = dtype.shape 3 ctx1cba

561 size = math.prod(shape) 3 ctx1cba

562 if base.names is None: 3 ctx1cba

563 return size 3 ctx1cba

564 else: 

565 return size * sum(_nd(base[name]) for name in base.names) 3 ctx1cba

566 

567 # I use this function in many parts of the package so it should not have an 

568 # underscore, even if I don't export it in the main namespace. And move it 

569 # to utils, it's not specific to StructuredArray. 

570 

571 

572def _structured_to_unstructured_recursive(idx, arr, out, *strides): 3 ctx1cba

573 dtype = arr.dtype 2 ctx1ba

574 for i, name in enumerate(dtype.names): 2 ctx1ba

575 subarr = arr[name] 2 ctx1ba

576 base = dtype[i].base 2 ctx1ba

577 size = math.prod(dtype[i].shape) 2 ctx1ba

578 stride = _nd(base) 2 ctx1ba

579 substrides = (*strides, (size, stride)) 2 ctx1ba

580 if base.names is not None: 2 ctx1ba

581 out, newidx = _structured_to_unstructured_recursive( 2 ctx1ba

582 idx, subarr, out, *substrides 

583 ) 

584 shift = newidx - idx 2 ctx1ba

585 assert shift == stride 2 ctx1ba

586 idx += size * stride 2 ctx1ba

587 else: 

588 assert stride == 1 2 ctx1ba

589 if all(size == 1 for size, _ in strides): 2 ctx1ba

590 indices = numpy.s_[idx : idx + size] 2 ctx1ba

591 srcsize = size 2 ctx1ba

592 else: 

593 indices = sum( 1 ctx1a

594 ( 

595 stride * numpy.arange(size)[numpy.s_[:,] + (None,) * i] 

596 for i, (size, stride) in enumerate(reversed(substrides)) 

597 ), 

598 start=idx, 

599 ) 

600 indices = indices.reshape(-1) 1 ctx1a

601 srcsize = indices.size 1 ctx1a

602 key = numpy.s_[..., indices] 2 ctx1ba

603 src = subarr.reshape((*out.shape[:-1], srcsize)) 2 ctx1ba

604 if hasattr(out, 'at'): 2 ctx1ba

605 out = out.at[key].set(src) 2 ctx1ba

606 else: 

607 out[key] = src 1 ctx1a

608 idx += size 2 ctx1ba

609 return out, idx 2 ctx1ba

610 

611 

612@StructuredArray._implements(numpy.empty_like) 3 ctx1cba

613def _empty_like(prototype, dtype=None, *, shape=None): 3 ctx1cba

614 shape = prototype.shape if shape is None else shape 1 ctx1a

615 dtype = prototype.dtype if dtype is None else dtype 1 ctx1a

616 return _empty(shape, dtype) 1 ctx1a

617 

618 

619@StructuredArray._implements(numpy.empty) 3 ctx1cba

620def _empty(shape, dtype=float): 3 ctx1cba

621 if hasattr(shape, '__len__'): 1 ctx1a

622 shape = tuple(shape) 1 ctx1a

623 else: 

624 shape = (int(shape),) 1 ctx1a

625 dtype = numpy.dtype(dtype) 1 ctx1a

626 arrays = {} 1 ctx1a

627 for i, name in enumerate(dtype.names): 1 ctx1a

628 dtbase = dtype[i].base 1 ctx1a

629 dtshape = shape + dtype[i].shape 1 ctx1a

630 if dtbase.names is not None: 1 ctx1a

631 y = _empty(dtshape, dtbase) 1 ctx1a

632 else: 

633 try: 1 ctx1a

634 y = jnp.empty(dtshape, dtbase) 1 ctx1a

635 except TypeError: 1 ctx1a

636 y = numpy.empty(dtshape, dtbase) 1 ctx1a

637 arrays[name] = y 1 ctx1a

638 return StructuredArray._array(shape, dtype, arrays) 1 ctx1a

639 

640 

641@StructuredArray._implements(numpy.concatenate) 3 ctx1cba

642def _concatenate(arrays, axis=0, dtype=None, casting='same_kind'): 3 ctx1cba