Skip to content

Commit 6fa3355

Browse files
awvwgkawvwgk
committed
Add routines for saving arrays in npy format
1 parent 57c41f9 commit 6fa3355

File tree

3 files changed

+225
-0
lines changed

3 files changed

+225
-0
lines changed

‎src/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ set(fppFiles
77
stdlib_bitsets_64.fypp
88
stdlib_bitsets_large.fypp
99
stdlib_io.fypp
10+
stdlib_io_npy.fypp
1011
stdlib_kinds.fypp
1112
stdlib_linalg.fypp
1213
stdlib_linalg_diag.fypp

‎src/Makefile.manual

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ SRCFYPP = \
44
stdlib_bitsets_large.fypp \
55
stdlib_bitsets.fypp \
66
stdlib_io.fypp \
7+
stdlib_io_npy.fypp \
78
stdlib_kinds.fypp \
89
stdlib_linalg.fypp \
910
stdlib_linalg_diag.fypp \
@@ -87,6 +88,10 @@ stdlib_io.o: \
8788
stdlib_optval.o \
8889
stdlib_kinds.o \
8990
stdlib_ascii.o
91+
stdlib_io_npy.o: \
92+
stdlib_error.o \
93+
stdlib_kinds.o \
94+
stdlib_strings.o
9095
stdlib_linalg.o: \
9196
stdlib_kinds.o
9297
stdlib_linalg_diag.o: \

‎src/stdlib_io_npy.fypp

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
! SPDX-Identifer: MIT
2+
3+
#:include "common.fypp"
4+
#:set RANKS = range(1, MAXRANK + 1)
5+
#:set KINDS_TYPES = REAL_KINDS_TYPES + INT_KINDS_TYPES + CMPLX_KINDS_TYPES
6+
7+
!> Description of the npy format taken from
8+
!> https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html
9+
!>
10+
!>## Format Version 1.0
11+
!>
12+
!> The first 6 bytes are a magic string: exactly \x93NUMPY.
13+
!>
14+
!> The next 1 byte is an unsigned byte:
15+
!> the major version number of the file format, e.g. \x01.
16+
!>
17+
!> The next 1 byte is an unsigned byte:
18+
!> the minor version number of the file format, e.g. \x00.
19+
!> Note: the version of the file format is not tied to the version of the numpy package.
20+
!>
21+
!> The next 2 bytes form a little-endian unsigned short int:
22+
!> the length of the header data HEADER_LEN.
23+
!>
24+
!> The next HEADER_LEN bytes form the header data describing the array’s format.
25+
!> It is an ASCII string which contains a Python literal expression of a dictionary.
26+
!> It is terminated by a newline (\n) and padded with spaces (\x20) to make the total
27+
!> of len(magic string) + 2 + len(length) + HEADER_LEN be evenly divisible by 64 for
28+
!> alignment purposes.
29+
!>
30+
!> The dictionary contains three keys:
31+
!>
32+
!> - “descr”: dtype.descr
33+
!> An object that can be passed as an argument to the numpy.dtype constructor
34+
!> to create the array’s dtype.
35+
!>
36+
!> - “fortran_order”: bool
37+
!> Whether the array data is Fortran-contiguous or not. Since Fortran-contiguous
38+
!> arrays are a common form of non-C-contiguity, we allow them to be written directly
39+
!> to disk for efficiency.
40+
!>
41+
!> - “shape”: tuple of int
42+
!> The shape of the array.
43+
!>
44+
!> For repeatability and readability, the dictionary keys are sorted in alphabetic order.
45+
!> This is for convenience only. A writer SHOULD implement this if possible. A reader MUST
46+
!> NOT depend on this.
47+
!>
48+
!> Following the header comes the array data. If the dtype contains Python objects
49+
!> (i.e. dtype.hasobject is True), then the data is a Python pickle of the array.
50+
!> Otherwise the data is the contiguous (either C- or Fortran-, depending on fortran_order)
51+
!> bytes of the array. Consumers can figure out the number of bytes by multiplying the
52+
!> number of elements given by the shape (noting that shape=() means there is 1 element)
53+
!> by dtype.itemsize.
54+
!>
55+
!>## Format Version 2.0
56+
!>
57+
!> The version 1.0 format only allowed the array header to have a total size of 65535 bytes.
58+
!> This can be exceeded by structured arrays with a large number of columns.
59+
!> The version 2.0 format extends the header size to 4 GiB. numpy.save will automatically
60+
!> save in 2.0 format if the data requires it, else it will always use the more compatible
61+
!> 1.0 format.
62+
!>
63+
!> The description of the fourth element of the header therefore has become:
64+
!> “The next 4 bytes form a little-endian unsigned int: the length of the header data
65+
!> HEADER_LEN.”
66+
!>
67+
!>## Format Version 3.0
68+
!>
69+
!> This version replaces the ASCII string (which in practice was latin1) with a
70+
!> utf8-encoded string, so supports structured types with any unicode field names.
71+
module stdlib_io_npy
72+
use stdlib_error, only : error_stop
73+
use stdlib_kinds, only : int8, int16, int32, int64, sp, dp, xdp, qp
74+
use stdlib_strings, only : to_string
75+
implicit none
76+
private
77+
78+
public :: save_npy
79+
80+
81+
!> Save multidimensional array in npy format
82+
interface save_npy
83+
#:for k1, t1 in KINDS_TYPES
84+
#:for rank in RANKS
85+
module procedure save_npy_${t1[0]}$${k1}$_${rank}$
86+
#:endfor
87+
#:endfor
88+
end interface save_npy
89+
90+
91+
character(len=*), parameter :: nl = char(10)
92+
93+
character(len=*), parameter :: &
94+
type_iint8 = "<i1", type_iint16 = "<i2", type_iint32 = "<i4", type_iint64 = "<i8", &
95+
type_rsp = "<f4", type_rdp = "<f8", type_rxdp = "<f10", type_rqp = "<f16", &
96+
type_csp = "<c8", type_cdp = "<c16", type_cxdp = "<c20", type_cqp = "<c32"
97+
98+
character(len=*), parameter :: &
99+
& magic_number = char(int(z"93")), &
100+
& magic_string = "NUMPY"
101+
102+
contains
103+
104+
105+
!> Generate magic header string for npy format
106+
pure function magic_header(major, minor) result(str)
107+
!> Major version of npy format
108+
integer, intent(in) :: major
109+
!> Minor version of npy format
110+
integer, intent(in) :: minor
111+
!> Magic string for npy format
112+
character(len=8) :: str
113+
114+
str = magic_number // magic_string // char(major) // char(minor)
115+
end function magic_header
116+
117+
118+
!> Generate header for npy format
119+
pure function npy_header(vtype, vshape) result(str)
120+
!> Type of variable
121+
character(len=*), intent(in) :: vtype
122+
!> Shape of variable
123+
integer, intent(in) :: vshape(:)
124+
!> Header string for npy format
125+
character(len=:), allocatable :: str
126+
127+
integer, parameter :: len_v10 = 8 + 2, len_v20 = 8 + 4
128+
129+
str = &
130+
"{'descr': '"//vtype//&
131+
"', 'fortran_order': True, 'shape': "//&
132+
shape_str(vshape)//", }"
133+
134+
if (len(str) + len_v10 >= 65535) then
135+
str = str // &
136+
& repeat(" ", 16 - mod(len(str) + len_v20 + 1, 16)) // nl
137+
str = magic_header(2, 0) // to_bytes_i4(int(len(str))) // str
138+
else
139+
str = str // &
140+
& repeat(" ", 16 - mod(len(str) + len_v10 + 1, 16)) // nl
141+
str = magic_header(1, 0) // to_bytes_i2(int(len(str))) // str
142+
end if
143+
end function npy_header
144+
145+
!> Write integer as byte string in little endian encoding
146+
pure function to_bytes_i4(val) result(str)
147+
!> Integer value to convert to bytes
148+
integer, intent(in) :: val
149+
!> String of bytes
150+
character(len=4), allocatable :: str
151+
152+
str = char(mod(val, 2**8)) // &
153+
& char(mod(val, 2**16) / 2**8) // &
154+
& char(mod(val, 2**32) / 2**16) // &
155+
& char(val / 2**32)
156+
end function to_bytes_i4
157+
158+
159+
!> Write integer as byte string in little endian encoding, 2-byte truncated version
160+
pure function to_bytes_i2(val) result(str)
161+
!> Integer value to convert to bytes
162+
integer, intent(in) :: val
163+
!> String of bytes
164+
character(len=2), allocatable :: str
165+
166+
str = char(mod(val, 2**8)) // &
167+
& char(mod(val, 2**16) / 2**8)
168+
end function to_bytes_i2
169+
170+
171+
!> Print array shape as tuple of int
172+
pure function shape_str(vshape) result(str)
173+
!> Shape of variable
174+
integer, intent(in) :: vshape(:)
175+
!> Shape string for npy format
176+
character(len=:), allocatable :: str
177+
178+
integer :: i
179+
180+
str = "("
181+
do i = 1, size(vshape)
182+
str = str//to_string(vshape(i))//", "
183+
enddo
184+
str = str//")"
185+
end function shape_str
186+
187+
188+
#:for k1, t1 in KINDS_TYPES
189+
#:for rank in RANKS
190+
!> Save ${rank}$-dimensional array in npy format
191+
subroutine save_npy_${t1[0]}$${k1}$_${rank}$(filename, array, iostat)
192+
character(len=*), intent(in) :: filename
193+
${t1}$, intent(in) :: array${ranksuffix(rank)}$
194+
integer, intent(out), optional :: iostat
195+
character(len=*), parameter :: vtype = type_${t1[0]}$${k1}$
196+
197+
integer :: io, stat
198+
199+
open(newunit=io, file=filename, form="unformatted", access="stream", iostat=stat)
200+
if (stat == 0) then
201+
write(io, iostat=stat) npy_header(vtype, shape(array))
202+
end if
203+
if (stat == 0) then
204+
write(io, iostat=stat) array
205+
end if
206+
close(io, iostat=stat)
207+
208+
if (present(iostat)) then
209+
iostat = stat
210+
else if (stat /= 0) then
211+
call error_stop("Failed to write array to file '"//filename//"'")
212+
end if
213+
214+
end subroutine save_npy_${t1[0]}$${k1}$_${rank}$
215+
#:endfor
216+
#:endfor
217+
218+
219+
end module stdlib_io_npy

0 commit comments

Comments
 (0)