More Adventures in AstroPy

Or, The Perils of General-Purpose Code

Quick one today: I've had occasion to look in on a brown dwarf spectroscopy codebase and examine it for potential optimizations.

The main problem highlighted up front lay in the interpolation code, which was running into issues with heterogenous spectra ranges, and the maintainer wanted to work on a method to cleanly separate extrapolation in unsupported regions from well-supported interpolation.

This is more of a correctness issue than a performance issue, but I still had a suspicion that we might be able to get some speedups by switching from scipy's interpolation library to interpn. I actually had occasion to speak with James Logan at RustConf a few weeks ago, and told him I was itching to use interpn in a project, so it seems the stars have just aligned.

After establishing the issue with heterogenous extrapolation, I grabbed some sample spectra and ran a some profiles on the problematic analyses to figure out exactly how much runtime was spent in interpolation, so I could estimate how big a speedup we could get out of some reasonable interventions.

And, uhhhh

Out of a 7.5 second sample, virtually all runtime not spent on importing modules was spent on reading files.

These aren't huge files either: the Sonora Diamondback spectra models come in at around 10MB each, and I was only using a few of them. Where is all that coming from.

As it turns out, it's my old nemesis, AstroPy.

I don't envy the task of the AstroPy maintainers: the library needs to be very flexible, easy to use for a wide variety of users, and very general-purpose.

In this case, that generality is murdering the CPU.

Let's take a look at that line, ui.py:510(_guess). I'm not sure why the ascii table code is located in the ui module, but there you have it.

It's attempting to parse the file with a combination of reader-supplied kwargs and default options. There's an option to use a fast reader engine written in C, but it's not clear if our existing scan is using that? It's definitely routing through the core read implementation, which processes individual lines via the DefaultSplitter process_line() implementation, and THAT is spending pretty much all its time on _replace_tab_with_space. Sure enough, the Sonora Diamondback files are tab-separated.

Each call to that averages about 3 microseconds. More than I'd expect for sure, but not that horrible. Only problem is we're calling it 1.58 million times, since each file has a few hundred rows. Not great.

The best solution here would just be to switch to using parquet files, and then we'd be in a land of sunshine and rainbows. However, these files are distributed as ascii, and that's not gonna change, and this project isn't willing to take on the extra complexity of transforming all its inputs into a new format, at least not at present. So we need a better ascii reader.

Luckily, we can do the same thing we did last time, taking a general-purpose library function and replacing it with a narrowly-tailored implementation specific to our use case. We know exactly what these files are supposed to look like, so we can constrain our solution quite effectively.

Specifically, I just made a buffered reader that splits by newlines and splits lines by whitespace/tab, then parses the floating point numbers into separate array columns.

I did get a little fancy and used this as an excuse to test out fearless_simd, which just hit 1.0 earlier this week. It's very nice, though the existing documentation and examples are a bit sparse. But once I got my head around the model, it's quite straightforward!

#[simd]
fn classify_block<S: Simd>(simd: S, block: &[u8]) -> (u64, u64) {
    let n = S::u8s::LEN;
    debug_assert!(n <= 64 && 64usize.is_multiple_of(n));
    let lane_mask = if n == 64 { u64::MAX } else { (1u64 << n) - 1 };

    let space = S::u8s::splat(simd, b' ');
    let tab = S::u8s::splat(simd, b'\t');
    let cr = S::u8s::splat(simd, b'\r');
    let newline = S::u8s::splat(simd, b'\n');

    let (mut word, mut nl) = (0u64, 0u64);
    for (k, chunk) in block.chunks_exact(n).enumerate() {
        let v = S::u8s::from_slice(simd, chunk);
        let is_newline = v.simd_eq(newline);
        let is_whitespace =
            is_newline | v.simd_eq(space) | v.simd_eq(tab) | v.simd_eq(cr);

        let shift = k * n;
        nl |= (is_newline.to_bitmask() & lane_mask) << shift;
        word |= ((!is_whitespace).to_bitmask() & lane_mask) << shift;
    }
    (word, nl)
}

In local testing, this gets us from around 1 second for a 10MB file, to 18 milliseconds for the same (amortized across multiple calls, since there's an up-front cost when importing this from the Python side).

We could go faster by parallelizing, but at this point I don't expect it will be worth it. This project should be able to use its thread budget more effectively elsewhere, and this speedup should be enough to make reading the spectrum files no longer a major concern.