Skip to content

pyburn: fix quadratic accumulation when reading raw output (10x) - #52

Open
Sohex wants to merge 2 commits into
alphaparrot:masterfrom
Sohex:fix/pyburn-quadratic-read
Open

pyburn: fix quadratic accumulation when reading raw output (10x)#52
Sohex wants to merge 2 commits into
alphaparrot:masterfrom
Sohex:fix/pyburn-quadratic-read

Conversation

@Sohex

@Sohex Sohex commented Aug 19, 2026

Copy link
Copy Markdown

Reading a raw output file grew each variable with np.append(variables[kcode], field) once per record. np.append reallocates and copies the whole accumulated array every call, so reading N records is O(N^2) in copied bytes. Collecting per code and joining once with np.concatenate gives the identical result -- np.append flattens anyway -- in linear time.

readvariablecode() had the same pattern, plus if not variable: on an ndarray, which raises once a second record arrives, so that path completed only for single-record variables.

The derived vertical fields (omega/wap, wa, streamfunction) were a smaller instance: a loop over (time, lat, lon) calling np.append twice and cumulative_trapezoid once per grid cell, ~1.5e6 iterations per T42 orbit. Now integrated along the level axis in one call, same arithmetic per column.

Measured, one T42 orbit, 10 layers, 2.4 GB raw, 120 codes, same input, no profiler:

wall np.append calls
before 304.0 s 3,062,749
after 29.4 s 1

10.3x, and all 123 output variables bitwise identical. A second comparison via a burn7-style namelist with 216 codes (which also exercises wa and stf) gives 363.5 s -> 46.1 s, 281/281 bitwise identical.

Note for re-profiling: cProfile is misleading here. It charges per call event, so the ~6e6 tiny appends in the per-cell loops look dominant while the ~81k large ones in the reader -- where the time actually is -- look cheap. Time without the profiler and count calls separately.

Compiles/imports cleanly. No behaviour change intended or observed.

Sohex and others added 2 commits August 18, 2026 20:21
Reading a raw output file grew each variable one record at a time:

    variables[kcode] = np.append(variables[kcode], field)

np.append allocates a new array and copies the whole accumulated result on
every call, so reading N records costs O(N^2) in copied bytes rather than
O(N). Collecting the records per code and joining once with np.concatenate
gives the identical result -- np.append flattens its arguments anyway, so the
1-D concatenate is the same operation -- in linear time.

readvariablecode() had the same pattern and is fixed the same way. It has a
second problem that this also removes: it accumulated into `variable` and
tested `if not variable:` to decide whether to start or extend, which raises
"truth value of an array with more than one element is ambiguous" as soon as a
second record arrives, so that function could not have worked for any variable
with more than one record.

The derived vertical fields were a smaller instance of the same idea. The
omega/wap, wa and streamfunction branches looped over (time, lat, lon) and
called np.append twice and scipy.integrate.cumulative_trapezoid once per GRID
CELL -- about 1.5 million iterations for a single T42 orbit each. Those are now
integrated along the level axis for the whole array in one call. The arithmetic
per column is unchanged: cumulative_trapezoid takes an `axis` argument and
accepts an `x` of the same shape, so prepending the zero with one concatenate
and integrating on axis 1 is the same computation.

MEASURED on one T42 orbit, 10 layers, 2.4 GB of raw output, 120 requested
codes, same input file both ways and no profiler attached:

    before   304.0 s     3,062,749 np.append calls
    after     29.4 s             1 np.append call

10.3x, and all 123 output variables are BITWISE IDENTICAL. A second comparison
through a burn7-style namelist requesting 216 codes, which additionally
exercises wa and the streamfunction, gives 363.5 s against 46.1 s with all 281
variables bitwise identical.

A note for anyone re-profiling this: cProfile is badly misleading here. It
charges per call event, so the ~6 million tiny np.append calls in the
per-cell loops appear to cost far more than they do, while the ~81k large ones
in the reader -- which is where the time actually goes -- look comparatively
cheap. The two callers differ by three orders of magnitude in call count and in
the opposite direction in per-call cost. Timing without the profiler and
counting calls separately is what distinguishes them.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
With the quadratic accumulation gone, what dominated reading a raw output file
was this, once per record:

    data = struct.unpack(en + datalength*fmt, fbuffer[n:n+datalength*wl])

struct.unpack builds a tuple of Python float OBJECTS. A single T42 grid record
is 8192 values, a T42 orbit is of order 81,000 such records, so that is roughly
6.6e8 Python floats created and destroyed to read one orbit -- and every one of
them is handed straight to np.asarray, which immediately unboxes it again. It
also slices a fresh 32 KB bytes object per record and builds an 8193-character
format string to describe it.

The payload is now decoded with np.frombuffer in a new _decoderecord(), which
creates no Python objects at all. readallvariables() keeps the per-record views
and does the copy once per code at the concatenate, so the file is copied once
rather than once per record.

Word length is still derived exactly as _getknownwordlength() derives it, from
the ratio of the record's length in bytes to its length in words, but without
re-reading the header to do so -- readrecord() unpacked every header twice, once
inside _getknownwordlength() and once for itself.

float64 is deliberate. np.asarray on the old tuple of Python floats produced
float64, so promoting the frombuffer view keeps both the values and the dtype
of every variable unchanged. Returning the file's native float32 instead would
be cheaper again and is tempting, but it would silently change the precision in
which wap, the streamfunction and the surface-pressure gradients are derived
downstream, so it is not done here.

MEASURED on a file with the geometry of one T42 orbit at NLOWIO = 0 -- 77
single-level codes and 36 ten-level codes over 185 output steps, 80,845 records
of 128x64 float32, 2.65 GB:

    readallvariables    before   30.53 s
                        after     0.96 s

32x, and all 116 variables plus all 118 refactored readfile() keys are identical
in VALUE, SHAPE AND DTYPE.

That file is synthetic, and it is worth saying why: raw output is normally
deleted once it has been postprocessed, so there was none left to measure
against, and what the reader's cost depends on -- record count and record size --
is reproducible exactly. It is a valid Fortran sequential unformatted stream and
reads correctly through unmodified pyburn, which is what makes it a bench rather
than a mock.

readvariablecode() and _gettimevar() carry the same reader pattern and are
converted with it. Neither could ever have run: both call

    readrecord(fbuffer, n, en, ml)

without the required `mf` argument, so both raise TypeError on their first line
of work. The previous commit noted the `if not variable:` defect in
readvariablecode() without noticing that this one made it moot. Both now run and
agree with readallvariables() record for record.

No new dependency. A JIT was considered for this and does not fit: struct.unpack
over a bytes object is not supported in numba's nopython mode, and rewriting the
loop so that numpy is handed the bytes directly IS the fix, after which there is
nothing left for a JIT to compile.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@Sohex

Sohex commented Aug 19, 2026

Copy link
Copy Markdown
Author

I've pushed a second commit to this branch (603e563). It's the same defect class in the same reader, and it turns out the first fix only uncovered it.

What was left. Once the np.append accumulation was gone, the remaining cost was almost entirely this line in readrecord:

data = struct.unpack(en + datalength*fmt, fbuffer[n:n+datalength*wl])

struct.unpack builds a tuple of Python float objects. One T42 grid record is 8192 values and a T42 orbit is of order 81,000 such records, so that's roughly 6.6e8 Python floats created and destroyed to read a single orbit -- each handed straight to np.asarray, which unboxes it again. It also slices a fresh 32 KB bytes per record and builds an 8193-character format string to describe it.

The payload is now decoded with np.frombuffer in a new _decoderecord, which creates no Python objects at all, and readallvariables keeps the per-record views and copies once per code at the concatenate rather than once per record. Word length is still derived the way _getknownwordlength derives it, from the ratio of the record's byte count to its word count, but without re-reading the header -- readrecord was unpacking every header twice, once inside _getknownwordlength and once for itself.

Measured, on a file with the geometry of one T42 orbit at NLOWIO = 0 (77 single-level codes and 36 ten-level codes over 185 output steps, 80,845 records of 128x64 float32, 2.65 GB):

readallvariables wall
after the first commit 30.53 s
after this one 0.96 s

32x, and all 116 variables plus all 118 refactored readfile() keys identical in value, shape and dtype. Combined with the first commit, postprocessing an orbit goes from ~304 s to a second or two, against 77-93 s for the model run that produced it.

That file is synthetic and I want to be upfront about why: raw output is normally deleted once it's been postprocessed, so there was nothing left to re-measure against, and what the reader's cost actually depends on -- record count and record size -- is reproducible exactly. It's a valid Fortran sequential unformatted stream and reads correctly through unmodified pyburn, which is what makes it a bench rather than a mock.

float64 is deliberate. np.asarray on the old tuple of Python floats produced float64, so promoting the frombuffer view keeps every variable's dtype as well as its values unchanged. Returning the file's native float32 is cheaper again and tempting, but it would silently change the precision in which wap, the streamfunction and the surface-pressure gradients are derived downstream, so I left it out of this PR. Happy to add it as a separate opt-in if you'd want that.

Two functions that could never have run. readvariablecode and _gettimevar carry the same reader pattern and are converted with it. Both call readrecord(fbuffer, n, en, ml) without the required mf argument, so both raise TypeError on their first line of work. My first commit noted the if not variable: bug in readvariablecode without spotting that this one already made it moot. Both now run and agree with readallvariables record for record.

No new dependency, and nothing outside pyburn.py is touched.

One note in case anyone profiles this further: I looked at whether a JIT would take the rest off, and it doesn't fit. struct.unpack over a bytes object isn't supported in numba's nopython mode, and rewriting the loop so numpy is handed the bytes directly is the fix -- after which there's nothing left to compile. The remaining stages aren't JIT-shaped either: in mode='grid' a grid variable passes through _transformvar untouched, and a spectral one goes to the already-f2py'd pyfft.sp2gp.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant