Dimitri Lozeve
9a1d251716
Fixes https://github.com/dlozeve/bqn-npy/issues/1. The spec says that •bit._cast should support 32-bit uints [1], but CBQN does not [2]. To get around this, we just use signed integers, which works for numbers smaller than 2^31. Above that, it will wrap around and (probably) fail, but CBQN does not support them anyway (except by converting them to doubles). From Marshall: > All of the provided functions at this time are identical on i32 and > u32, because they work (mod 2^32) and the only difference between > those formats is to add or subtract 2^32 from numbers outside the > range [0,2^31). > If you want to run on u32s stored directly as numbers in BQN (they > won't fit in i32 and will be stored as double!) you can •bit._cast > those into a signed format at the beginning and then back out at the > end. [1] https://mlochbaum.github.io/BQN/spec/system.html#bitwise-operations [2] https://matrix.to/#/!EjsgbQQNuTfHXQoiax:matrix.org/$PL2AYNlLQuYhPftdt7yw7rW9heiVEEeATte2zOt2BYk
50 lines
2 KiB
BQN
50 lines
2 KiB
BQN
⟨SaveNpy,LoadNpy⟩⇐
|
||
|
||
lf←@+10
|
||
magicString←(@+147)∾"NUMPY"
|
||
|
||
# Type conversions
|
||
|
||
# CBQN does not support other widths than 1 for 'u'. However, signed
|
||
# and unsigned integers seem to have the same behavior for numbers
|
||
# below 2⋆31.
|
||
dtypes‿formats←⟨"<f8","<i4","<u4"⟩‿⟨64‿'f',32‿'i',1‿'u'⟩
|
||
DtypeToFormat←{(⊑dtypes⊐<𝕩)⊑formats}
|
||
ArrayToBytes←{⟨DtypeToFormat 𝕨,8‿'c'⟩•bit._cast 𝕩}
|
||
ArrayFromBytes←{⟨8‿'c',DtypeToFormat 𝕨⟩•bit._cast 𝕩}
|
||
|
||
# Saving
|
||
To_i16le←{@+256(|∾⌊∘÷˜)𝕩}
|
||
FormatShape←{'('∾(1↓∾⥊','∾˘•Fmt¨≢𝕩)∾",)"}
|
||
|
||
BuildHeader←{dtype 𝕊 data:
|
||
version←1‿0
|
||
# Hypothesis: BQN arrays are C-contiguous
|
||
headerData←"{'descr':'"∾dtype∾"','fortran_order':False,'shape':"∾(FormatShape data)∾",}"
|
||
padding←' '↑˜64-64|(≠magicString)+2+2+(≠headerData)+1
|
||
! 0=64|(≠magicString)+2+2+(≠padding)+(≠headerData)+1
|
||
headerDataPadded←headerData∾padding∾lf
|
||
magicString∾(@+version)∾(To_i16le ≠headerDataPadded)∾headerDataPadded
|
||
}
|
||
DetectDtype←{(∧´⥊⌊⊸=¨𝕩)⊑⟨"<f8",(∧´⥊0≤𝕩)⊑"<i4"‿"<u4"⟩}
|
||
EncodeNpy←DetectDtype⊸(BuildHeader∾ArrayToBytes⟜⥊)
|
||
SaveNpy←(•wdpath⊸•file.At)⊸•file.Bytes⟜EncodeNpy
|
||
|
||
# Loading
|
||
ParseHeader←{𝕊𝕩:
|
||
"Not a valid NPY file"! magicString≡(≠magicString)↑𝕩
|
||
version←@-˜(0‿1+≠magicString)⊏𝕩
|
||
headerlen←+´1‿256×@-˜(2‿3+≠magicString)⊏𝕩
|
||
header←(¬∊⟜(' '∾lf))⊸/headerlen↑(4+≠magicString)↓𝕩
|
||
dtype←3↑(⌊´∘⊐⟜"<>"↓⊢) (("descr"⍷header)⊐1)↓header
|
||
("Unsupported dtype: "∾dtype) ! ⊑(<dtype)∊dtypes
|
||
shapestr←(⊐⟜')'↑⊢) (1⊸+∘⊐⟜'('↓⊢) (("shape"⍷header)⊐1)↓header
|
||
shape←•BQN¨','((⊢-˜+`׬)∘=⊔⊢)shapestr
|
||
⟨version,dtype,shape,headerlen+2+2+≠magicString⟩
|
||
}
|
||
|
||
LoadNpy←{𝕊𝕩:
|
||
bytes←•file.Bytes 𝕩
|
||
version‿dtype‿shape‿datastart←ParseHeader bytes
|
||
shape⥊dtype ArrayFromBytes datastart↓bytes
|
||
}
|