Vectorize string output for Time formats - #19942

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• Do the proposed changes actually accomplish desired goals?
• Do the proposed changes follow the Astropy coding guidelines?
• Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see instructions for rebase and squash.
• Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

Not sure why CI is skipped. There is no skip ci in commit and I cannot find a button to run them.

Thanks @pllim! Not sure why they didn't fire at first either — merging main in to update the branch seems to have kicked the workflows off, and they're running now.

For future reference, we prefer a rebase to avoid merge commit in a PR commit history.

Note to maintainer: Please squash to merge when it is time, if you choose to accept the PR.

Thanks, all.

@taldcroft quick process question — the branch has picked up a couple of "Merge main" commits from keeping it in sync, which I know pllim would rather we avoid. I can clean those out with a rebase, but that's a force-push and it'd scramble your "changes since last review" view while you're still going through it.

Would you rather I rebase to drop the merge commits now, or leave the history as-is and do a single tidy-up rebase right before merge (or just let the squash-merge handle it)? Happy either way — just didn't want to disrupt your review.

@veyron-kairo - sorry for the delay, I'm starting to look at this again. Yes, please rebase. That's the way we do things in astropy and I'm used to dealing with that. For the most part you don't rebase unless there is an actual merge conflict.

Thanks @taldcroft, and no worries about the delay! Good to know — I'll keep that in mind for next time and only rebase when there's an actual conflict. Gone ahead and rebased this one onto main to drop the merge commits, so the history is linear now. 👍

@veyron-kairo - it appears that main was merged into your branch again. Please rebase and do not (ever) merge main into a dev branch. You do not need to be concerned with the This branch is out-of-date with the base branch message.

@veyron-kairo - this is great work, thanks! Just needs a rebase and a couple new tests. I took the liberty of expanding the doc strings since it is so easy now with AI. They look about right but you should double check for accuracy.

Thanks @taldcroft! Rebased onto the latest main with the merge commit removed, and I kept your AI-docs commit in place. I went through the expanded docstrings against the implementation and they all check out. Tests pass locally. Will add the extra test cases you flagged next.

@veyron-kairo - it would be good if you could verify that the new PR is indeed equally fast to what you had.

@taldcroft - with my changes, it would be good to have you do a quick sanity check... A larger question (not for this PR), is whether we should start moving to numpy's new string dtype "T". It uses UTF-8 so uses less space, though the minimum is 16 bytes (everything up to 15 characters get stored directly; 16 and more is stored in a separate "arena"). Still, for Time, I think it will nearly always be better than "U".

Annoyingly, the code failed on oldestdeps, because numpy 2.0 had a bug for zfill. Since this is just for performance and we'll drop 2.0 soon enough, I just let the code bail for numpy < 2.1. I tested that oldestdeps with numpy 2.1 does work.

Thanks @mhvk — and no need to apologise at all, the rewrite is much cleaner and I learned a lot from it (I didn't know about the np.strings functions or the new "T" dtype). And good catch on the +06d truncation.

On speed: I tried to verify as you asked. On numpy 2.4.6 I'm fairly consistently seeing the new version come out a few times slower than the old buffer path for the common uniform-width case — e.g. for a 1M-element iso/isot/fits array, roughly 0.10 s before vs 0.55 s now (both of course still far faster than the previous per-element loop). Profiling points at the per-field int → "T" conversion plus the repeated concatenations.

That said, I don't think it's worth chasing: the old buffer only hit that speed when every year had the same width — it fell back to the per-element path for mixed widths and, as you found, truncated large +06d years. So the simpler, always-correct version looks like the right trade to me, especially since formatting isn't really the critical path. Happy to leave it as is.

Script I used, in case it's handy or doesn't match your numbers:

import time
import numpy as np
import erfa
from astropy.time import Time
from astropy.time.formats import (
    TimeString, _day_of_year, _SUBFMT_FIELD, _INT_SPEC,
)

# Original buffer implementation, for comparison.
def _orig_write_decimal(out, values, width, signed):
    d = np.asarray(values, dtype=np.int64)
    if signed:
        neg = d < 0
        d = np.abs(d)
    for i in range(width - 1, -1, -1):
        out[..., i] = d % 10 + ord("0")
        d //= 10
    if signed:
        out[..., 0] = np.where(neg, ord("-"), ord("+"))

def _orig_field_width(spec, values):
    m = _INT_SPEC.fullmatch(spec)
    if m is None:
        return None, None
    sign, zero, digits = m.group(1, 2, 3)
    if digits:
        return (int(digits), sign == "+") if zero else (None, None)
    if values.size == 0:
        return 1, False
    low = int(values.min())
    if sign or low < 0 or len(str(low)) != len(str(int(values.max()))):
        return None, None
    return len(str(int(values.max()))), False

def _orig_value_fast(self, str_fmt):
    scale = (self.scale.upper().encode("ascii"),)
    iys, ims, ids, ih = erfa.d2dtf(scale, self.precision, self.jd1, self.jd2)
    f = {"year": np.asarray(iys), "mon": np.asarray(ims), "day": np.asarray(ids),
         "hour": np.asarray(ih["h"]), "min": np.asarray(ih["m"]),
         "sec": np.asarray(ih["s"]), "fracsec": np.asarray(ih["f"])}
    if "{yday:" in str_fmt:
        f["yday"] = _day_of_year(f["year"], f["mon"], f["day"])
    pieces, width, pos = [], 0, 0
    for mt in _SUBFMT_FIELD.finditer(str_fmt):
        if mt.start() > pos:
            lit = str_fmt[pos:mt.start()]; pieces.append((lit, None, None)); width += len(lit)
        name, spec = mt.group(1, 2)
        if name not in f:
            return None
        fw, sg = _orig_field_width(spec, f[name])
        if fw is None:
            return None
        pieces.append((f[name], fw, sg)); width += fw; pos = mt.end()
    if pos < len(str_fmt):
        lit = str_fmt[pos:]; pieces.append((lit, None, None)); width += len(lit)
    buf = np.empty(f["year"].shape + (width,), dtype=np.uint32); col = 0
    for v, fw, sg in pieces:
        if fw is None:
            buf[..., col:col + len(v)] = [ord(c) for c in v]; col += len(v)
        else:
            _orig_write_decimal(buf[..., col:col + fw], v, fw, sg); col += fw
    return buf.view(f"U{width}").reshape(self.jd1.shape)

cur = TimeString._value_fast

def bench(impl, N, reps=5):
    TimeString._value_fast = impl
    best = float("inf")
    for _ in range(reps):
        t = Time("2001-01-01T00:00:00", format="isot") + np.arange(N) * (1.0 / 86400)
        t0 = time.perf_counter()
        _ = t.fits           # first access; .fits is cached afterwards
        best = min(best, time.perf_counter() - t0)
    return best

for N in [100_000, 1_000_000]:
    o = bench(_orig_value_fast, N)
    m = bench(cur, N)
    print(f"N={N:>9,}  orig={o:.4f} s  current={m:.4f} s  ({m / o:.1f}x)")

TimeString._value_fast = cur

@mhvk - I had a quick look at your commit. At this point I share your view that the tests are good so I'm not going to spend too much time confirming that your new code is correct. (I already spent a lot of time on the initial version!).

But I will ask that you revert (or more strongly justify) removing the unit tests on private functions. I disagree with your stated philosophy of not testing those. On the contrary, I often make a point of factoring out smaller (private) functions with the express intent that they can be more easily and transparently tested. That's really the case here, and I was actually quite happy with all those unit tests that confirmed at a low-level that those private functions were working as expected.

for the common uniform-width case — e.g. for a 1M-element iso/isot/fits array, roughly 0.10 s before vs 0.55 s now

Good point @veyron-kairo. (Your comment came in while I was typing mine). I'm hesitant to accept a factor of 5 slowdown for the new code even if it is cleaner. I actually understood your original implementation and it seemed sensible, a bit more like what one would do in C with pre-allocating and filling in digits.

So maybe we do need a real discussion about strategy here.

Thanks @taldcroft — really appreciate you taking the time to dig into this. I'm totally happy to follow whatever you and @mhvk settle on for the strategy.

@veyron-kairo - Ai, that's a horrible slowdown! I had only done the simple test from #19929, and it seemed to be just as fast as your implementation, so thought it was OK. But it sounds like I did not look at the numbers carefully enough.

Given the slowdown, I think it is obvious we should not use my commits. One possibility is to just remove them again. Alternatively, my earlier attempt I am sure was similar in speed to yours - I now pushed this to #19998. But like your original, that will fail the test I added.

Aside: we really should add some new items to the benchmarks.

p.s. @taldcroft - the private function tests I removed were all tests of a function that no longer exists in this PR. My comment reflects mostly that if you test private functions rather than API, then anyone who refactors and removes/changes private stuff will get failing tests and decide they don't matter. To me, that is a reason to avoid testing them if possible, but of course in some of the tests here it is really best to actually check the code even runs!

Thanks @mhvk — and thanks for pushing the alternative up so there's concrete code to compare against.

I ran the same benchmark on #19998 too, just to be thorough (checked out the branch and timed the real .fits/.iso path). On numpy 2.4.6 I'm seeing it come out a few times slower than the original as well — for a 1M-element iso/fits array, very roughly 0.11 s for the original vs ~0.47 s for #19998 (and ~0.55 s for the current np.strings version). From a bit of profiling it looks like the cost is mostly the int → str conversion (the astype to "U"/"T"), which the original happens to sidestep by writing the digits straight into the buffer.

These are of course just my numbers on one machine, so they may well not match what you're seeing — happy to share the exact script if it's useful. And either way I'm glad to go whichever direction you and @taldcroft prefer. If it'd help, I'm happy to take a pass at keeping the fast pre-allocated approach but fixing the large-year overflow so it passes the test you added (and restoring the low-level tests).

@veyron-kairo - thanks for your testing. In the end, it turns out I was simply an idiot, having forgotten about how Time caches its result. Now ensuring the cache is deleted in the timing, I confirm what you have also on my machine:

tm = Time(2451545.0 + np.random.uniform(0, 3650, 100_000), format='jd')
%timeit del tm.cache; tm.isot
# Current master
200 ms ± 1.56 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# PR 19942, your original
19.6 ms ± 97.3 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
# My revision here
71.2 ms ± 517 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
# My alternative (#19998)
60.1 ms ± 176 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

I suggest you simply remove my additional commits, only keeping the new test, and force-push so you are back at what you had. Of course, you'd also need some type of guard that ensures the new test also passes (perhaps just falling back to the slow path, since this is hardly a common case!).

Apologies for leading us all the wrong direction!

p.s. I am actually somewhat surprised that int_array.astype("T") is so slow, but that's for another time!

p.s. Do add a comment to your custom int-to-decimal converted to note that it is faster than astype("T")...

FWIW, astype("T") is so slow because it goes from python string conversion; see https://github.com/numpy/numpy/blob/9721c24ba87e6323f29e0b2f4f73aa22574c2aef/numpy/_core/src/multiarray/stringdtype/casts.cpp#L871-L877. astype("U") is even slower...

Thanks @mhvk — and no worries at all, it was a useful detour and I learned a lot from it!

I've gone ahead and done as you suggested: removed your two commits so it's back to the buffer version, and added a small guard in _field_width so that any zero-padded field that would overflow (like a year past 9999 in the +06d FITS format) now falls back to the per-element path instead of truncating — so your test_string_fits_very_large_years passes (I kept it, just adjusted it to expect the fallback). I also restored the low-level _field_width/_write_decimal tests, and added a comment on the converter noting it's faster than astype, with your numpy source link. Force-pushed, so the branch is linear again.

Thanks @mhvk — and you were right, the subtract-multiply is the fastest of the options. I checked all three end-to-end on numpy 2.4.6 (current %///, np.divmod, and your digits - 10*by10): on the full .isot/.fits path your version came out fastest, ~5% over the current code at high precision, so I've applied it (with you as co-author on the commit). Thanks for the approval and all the help on this one — I learned a lot!

Excellent, happy that I added at least a little! Since @taldcroft approved already, I'll merge. Thanks again!

@mhvk @taldcroft , thanks a lot again got to learn a lot from you 🙏 !