Changelog#

[Unreleased]#

Added#

  • API reference covers every public docstring. Sphinx now ignores each module’s __all__, which listed only part of the public names and hid e.g. edge_mask, extract_stencil, LazyFluxFields.compute and output_filename. The package overview and tracmass.seagrid.debug got pages. Private helpers (_name) stay out.

  • Land mask and distance to land for every seagrid model. seagrid.glorys, seagrid.ecco2 and seagrid.scrippsgog have landmask(k=1) (bool, True where kmt < k) and distance_to_land(k=1) (great-circle km from each cell centre to the nearest land cell centre, 0 on land). Both return xr.DataArrays on (y, x) with the grid’s lon / lat. The shared code is in the new seagrid/coast.py. It uses a KD-tree on unit-sphere coordinates holding only coastal land cells, so curvilinear grids and the date line are handled and full GLORYS takes about 5 s. scipy is now a direct dependency (it was already installed through scikit-learn).

  • Trajectory lifetime: traj_max_time. All integrators (steady_state, time_dependent, time_dependent_patched, time_dependent_streaming) take traj_max_time, read from the traj_max_time setting when omitted. It is a pd.Timedelta string ("30D", "12h") or calendar months/quarters/years ("1M", "6M", "1Q", "1Y"), which are measured from the run’s start. It sets the existing timax_seconds limit, counted from each particle’s release. timax_seconds now defaults to “not given”; passing both raises ValueError. Run metadata gets traj_max_time / timax_seconds. New helpers runsetup.traj_max_seconds and runsetup.resolve_timax.

  • Repeat seeding: seed_freq. time_dependent, time_dependent_patched and time_dependent_streaming take seed_freq (a pandas frequency, e.g. "30D" or "MS"). When the argument is omitted it is read from seed_freq in the [env.seed] table. Every seed_freq from dtm_start, a copy of the initial state is appended to the running state at the start of the window the time falls in. The copy’s clock (tt, ts, t0) is moved to that time, so t0 identifies the release and timax_seconds counts per particle. Particles keep their array index, so existing output ids are unchanged. The window loops no longer stop early while a release is still to come. Run metadata gets seed_freq / seed_windows. New helpers: ParticleState.concat(*states) (NumPy in, NumPy out, so the batched path stays on the host), runsetup.seed_windows, runsetup.reseed, runsetup.resolve_seed_freq and runsetup.seed_metadata. seed.from_settings ignores the seed_freq key.

  • upwelling.movie: movie of the gog_upwelling run (upwelling.py). Reads the newest trm_gog_upwelling file one window at a time (the file holds ~1.4 billion rows, too many to load at once) and scatters the particles on the projmap gog map, one frame per window. mode="eq" shows particles with ka == 12, mode="le" those with ka <= 12 (the level is set with ka=). Writes .mp4 when ffmpeg is installed, otherwise .gif.

  • seed.latlon_list: seed one particle per lat/lon point. Takes a DataFrame with lat/lon columns (optional k column for the 1-based level, else kpos), a .csv/.parquet file, or lat=/lon= lists. Each particle starts at the point’s exact fractional grid position (trajs.coords.lonlat_to_xy on the adapter’s lonlat_grid(grid), so subset grids work); points on land move to the nearest ocean cell, and rows below the local bottom (k > kmt) are dropped with a warning. Registered in seed.SETTINGS_APPROACHES as approach = "latlon_list", with seed_file = ... or inline lat = [...] / lon = [...] in [env.seed].

  • gcm.scrippsgog / seagrid.scrippsgog: Scripps Gulf of Guinea MITgcm adapter. seagrid.scrippsgog.retrieve(dtm) downloads the daily GG_i001_YYYY-MM-DD_1day_UVEL_VVEL_top200m.nc files from sose.ucsd.edu/GoG/DATA/PRELIM over HTTPS (streamed to a .part file, retried on dropped transfers) and caches them unchanged; available_dates() lists the server’s dates, retrieve_grid() fetches the binary MITgcm grid files. grid_definition() builds kmt/dzt from the Depth file with MITgcm’s partial-cell rule (hFacMin = 0.2) and raises if that does not reproduce the velocity files’ hFacW/hFacS; its 2-D coordinates are lons/lats. gcm.scrippsgog.fluxes(dtm) moves MITgcm’s west/south-face UVEL/VVEL onto the east/north faces the kernels read (the last column’s east face is zero) and builds transports from the files’ own dyG/dxG/drF/hFacW/hFacS. Horizontal-only (no WVEL in the files); the regional domain is non-periodic, and its open western/southern boundaries are closed in the tracking. Includes lonlat_grid, lonlat_to_xy, xy_to_lonlat; tests in tests/test_scrippsgog.py.

  • Ramp-up saves: time_dependent and time_dependent_patched take rampup_windows / rampup_substeps (or the settings of the same name; defaults 0 = off and 5). During the first rampup_windows windows, positions are also saved rampup_substeps - 1 times inside each window, as rows with the new event_type="rampup" (window = the window they fall in, dtm = their exact time), written by the new ParticleRecorder.record_rampup. The run metadata gets rampup_windows, rampup_substeps and save_dt_seconds, the time between consecutive saved positions (written for every run, ramp-up or not). For resolving the first hours of dispersal, e.g. residence times in boxes particles leave within one daily window. Particles are stopped exactly on each save time: _run_window now caps steps at its (dynamic) end time and takes the piece’s start for the slow-particle pace check, so the pieces reuse the compiled loop. time_dependent_streaming and scheme="analytical" raise if a ramp-up is requested.

  • tracmass.load(rampup=True) keeps the "rampup" rows along with the snapshots and adds the event_type column. The default still returns snapshots only, so existing analyses see the same rows.

  • tracmass.run(dtm=...): dtm is accepted (via **kwargs) as an alias for dtm_start, matching the dtm name used by LazyFluxFields and the adapters. Passing both raises a TypeError; the alias goes through the same date validation and is not forwarded to the integrator.

  • Verbose debug output (new tracmass.verbose): tracmass.run(..., verbose=True), verbose = true in the active settings environment (also read by the integrators), or tracmass.set_verbose() print timestamped messages to stderr for each stage: GCM adapter calls, each window’s flux read (total size, then time per variable), seeding and particle count, the resolved integrator options, monitor start, and per-window I/O / advection (first window: incl. JIT) / update_flux times. Uses the standard "tracmass" logger at DEBUG. Added to diagnose run(env="gog") showing no monitor: before the monitor starts, time_dependent loads the first window, which for global GLORYS with no subset is 42.5 GB (4318 × 3059 × 50, float64, both slots), several minutes per window.

  • Warning for very large flux reads: LazyFluxFields.compute warns when one window reads more than 8 GB, suggesting an [env.gcm] subset.

  • tracmass.load(filters=...) reads only matching rows, filtered inside the Parquet reader so row groups that cannot match are skipped (cheap for e.g. one window of a long run). Takes a dict shorthand ({"window": 0}; a list value means “in”) or pyarrow filters (a list of (column, op, value) tuples, or a list of such lists for OR). Filter columns need not be in columns; with clean=True the snapshot-row filter is ANDed on. Kept separate from columns because column selection and row selection are independent (you often filter on a column you do not want back).

  • Trajectory files describe their run: the integrators now store runsetup.run_metadata in the recorder’s tracmass.run Parquet metadata — gcm, gcm_product, gcm_subset, env, integrator, direction, dtm_start, dtm_end, timestep_seconds, tracmass_version, created (previously only gcm). Keys a caller passes in metadata= win. tracmass.load(lonlat=True) uses it: with no gcm in the file it looks up [env.gcm] name of the file’s env (not just the active one), and it converts with the grid of the file’s gcm_product, so a file from another GLORYS/ECCO2 product gets the right lon/lat.

  • Seed by lat/lon from settings: seed.SETTINGS_APPROACHES gains "latlon_box" and "latlon_point", so [env.seed] can use approach = "latlon_box" with box = { lat1, lat2, lon1, lon2, k1, k2 }. seed.latlon_box was unusable — it read grid.lat / grid.lon, which GridFields does not have. It now takes the cell-centre lon/lat from the new adapter function lonlat_grid(grid) (gcm.glorys, gcm.ecco2; the window of the full grid at grid.i_offset / j_offset, so it follows the subset), seeds the ocean cells whose centre is in the box (levels k1..k2, skipping k > kmt), handles boxes across the dateline (lon1 > lon2), and raises with the grid’s lat/lon extent when the box holds no ocean cell. The adapter comes from gcm= or [env.gcm] name. latlon_point is a latlon_box of lat ± dlat, lon ± dlon.

  • tracmass.load finds trajectory files by environment and date (new output.py). load(filename=None, *, env=None, dtm=None, reverse=None, columns=None) reads a recorder file as a DataFrame, with its path and tracmass.run metadata in df.attrs. Instead of a file name, pass env (read via settings.from_env, so the active env is not switched) and/or dtm to get the file the settings-driven recorder writes for that run; without dtm the newest file of that env; with no arguments the newest .parquet in base_output_dir or its sub-folders. Lookups also check the flat and {env} sub-folder layouts, so files written under another use_sub_folder setting are found. Replaces a stub that ignored its argument, read a hard-coded trm_omi.parquet and returned nothing. The recorder’s path builder (runsetup._auto_recorder_path) now uses the same output_folder / output_stem / output_filename helpers. lonlat=True adds lon / lat columns via the GCM adapter’s xy_to_lonlat; the adapter comes from gcm=, the file’s metadata or [env.gcm] name. clean (default True) keeps only the snapshot rows and the columns id, x, y, dtm (plus lon, lat); an explicit columns overrides the column list, and clean=False returns the whole file. For the lon/lat lookup, manage_run_io now adds the adapter’s name (gcm) to the recorder’s metadata, and ParticleRecorder.metadata is a public attribute (was _metadata).

  • gcm.glorys.xy_to_lonlat reads only the grid window it needs: it used to prepare the whole 4318 × 3059 global lon/lat grid on every call (~28 s here, regardless of how many positions); it now reads the cells spanned by the positions plus a 2-cell margin, giving identical results (33 s → 5 s for a 2.1M-row file). Positions near the zonal edge still use the full grid, so seam handling is unchanged.

  • use_sub_folder setting: when true, the settings-built trajectory recorder writes into a sub-directory of base_output_dir named after output_name, or the active environment when output_name is unset (e.g. /data/raid/trm/trm_tempvel/trm_tempvel_forward_2020-07-01.parquet), so runs from different environments don’t pile up in one directory. The directory is created if missing. Default false (unchanged layout).

  • tracmass.run: one-call, settings-driven entry point (new driver.py), generalised from abstemp.run. It activates an optional env, builds grid + flux buffer with a single new_flux, seeds (a ParticleState or a seed(grid) callable) and calls the integrator. Omitted arguments come from settings: gcm from [env.gcm] name, dtm_start / dtm_end from the keys of the same name, else dtm_end = dtm_start + runtime (negative runtime → backward run; with no start date, the adapter’s default date is used), and integrator (time_dependent / time_dependent_patched / time_dependent_streaming, default time_dependent). Remaining keyword arguments go to the integrator. abstemp.run is now a thin wrapper around it.

  • tracmass.run can seed from settings; new seed.from_settings: state is now optional. When omitted, particles come from the active [env.seed] table — approach (a key of the new seed.SETTINGS_APPROACHES, currently "ijk"), its arguments in a box sub-table, and every other key (sub_dist, i_sub_num, …) passed on as a sub-cell option — so e.g. tracmass.run(env="omi") describes a whole run in tracmass.toml. The dates come first, so the signature is run(dtm_start=None, dtm_end=None, state=None, *, env=, gcm=, integrator=, **kwargs) and run("2007-06-01", env="omi") needs no keyword; a ParticleState or seed function passed in a date slot (the earlier run(state, dtm_start, ...) order) raises a TypeError. seed.py now imports ParticleState from tracmass.particle (was tracmass) so driver.py can import it during package initialisation.

  • The flux buffer carries its own time stepping, so runs no longer need a hand-computed n_windows (gcm/flux_grid.py, runsetup.py, integrate/): LazyFluxFields gains timestep and reverse (copied from the grid by new_flux), dtm_end (the dtm_end passed to new_flux) and a derived n_windows property, ceil(|dtm_end - dtm| / timestep) — the windows left until dtm_end, decreasing by one per update_flux, None without a dtm_end. time_dependent, time_dependent_patched and time_dependent_streaming now default n_windows to None, resolved (new runsetup.resolve_n_windows) from fluxes.n_windows, falling back to the old default of 1 when no dtm_end was given; an explicit n_windows still wins. So new_flux(fn, dtm_start=d0, dtm_end=d1) followed by time_dependent(state, grid, lazy) replaces the nwin = int((d1 - d0) / tstep) boilerplate — and no longer needs the caller to know the adapter’s timestep, which already lives in each seagrid module’s settings["timestep"]. A dtm_end that falls inside a window now gets that window integrated (the scripts’ int() truncated it).

  • update_flux(lazy): step size, direction and the fluxes callable now come from the buffer itself (lazy.timestep / lazy.reverse / lazy.fluxes_fn); update_flux(lazy, fluxes_fn) overrides the callable. The old update_flux(grid, lazy, fluxes_fn) form still works and is needed only for a buffer built by lower-level helpers without a timestep. update_flux and LazyFluxFields.preload now carry all buffer metadata over via dataclasses.replace (preload used to drop zarr_chunk_xy).

  • trm_*.toml files in the current directory are now loaded as settings (config.py): any file matching trm_*.toml next to a script is added to Dynaconf’s settings_files, sorted by filename, after the current directory’s tracmass.toml/.secrets.toml and before the SEAVIEW_SETTINGS_FILE_FOR_DYNACONF override. Lets a script layer run-specific settings (e.g. trm_omi.toml defining an [omi] environment) on top of a shared tracmass.toml without editing it. New TRM_FILES module attribute holds the resolved, sorted list.

  • Rich monitor title shows the active Dynaconf environment (monitoring.py): the RunMonitor(backend="rich") table title is now tracmass · time_dependent · env: <env> (e.g. env: omi), read from settings.current_env on each redraw, so it’s clear at a glance which tracmass.toml section a run is using.

  • abstemp.closest_unconnected_db (abstemp.py): for every consecutive regid_path leg that no trajectory connects (the legs stitched_traj_db skips; listed by the new unconnected_pairs_db), extracts all trajectories starting in the origin region and ranks them by closest approach to the target region (distance to the nearest target-cell centre, via a distance-transform lookup joined inside DuckDB). Adds target, dist, min_dist, rank and a closest flag set on every position of the n_closest best trajectories; closest_only=True returns just those. The per-pair step is exposed as abstemp.closest_traj_db(con, regid1, regid2, ...).

  • abstemp.stitched_traj_db(..., include_unconnected=False): when set, an unconnected leg is filled with its closest-approach trajectory (from closest_traj_db) instead of being skipped, so the stitched path has no missing legs; a connected column marks those stand-in rows False.

Fixed#

  • seagrid.glorys.rechunk_wvel crashed a run reading the same days. A gog_upwelling run caught up with a running rechunk and failed with cannot reshape array of size 10000000 into shape (50,1000,1000): it had opened a store with 50-level chunks one window before reading it, and the rechunk swapped in the 10-level store in between. The docstring and README now say not to rechunk days a run is reading. The swap now renames the original away and the new store in before deleting the original, so the store is never missing while the original is deleted.

  • time_dependent_streaming did not record the time limit as a stop reason. Particles past timax_seconds were stopped but kept reason 0, so they looked paused rather than stopped for time. They now get REASON_TIME, as in the other integrators.

  • Slow flux reads on small GLORYS subsets with vertical motion (e.g. gog_upwelling, ~10 s per window). There were three causes:

    • Every snapshot was read twice, once as the present slot and again as the next window’s past slot. LazyFluxFields.compute now keeps the present slot’s arrays (_computed, keyed by the slot’s DataArray and shared across update_flux copies) and reuses them for the past slot. This applies to all adapters with the full-grid integrators.

    • Land fill values were masked before the subset was cut out (seagrid.glorys.open_dataset). Every whole 1000 × 1000 Zarr chunk the subset touched was masked, about 7 times slower than masking only the subset. open_dataset / open_dataarray take mask= (default True, unchanged for other callers). gcm.glorys.fluxes opens with mask=False and applies the new public seagrid.glorys.mask_fill after the halo trim and subset.

    • wvel Zarrs had one 50-level chunk per tile. The download chunk settings named deptht only, and wvel is on depthw. New downloads use 10 levels ("depthw": 10). The new seagrid.glorys.rechunk_wvel(dtm1, dtm2) rewrites existing stores. It checks the values against the original before swapping, and skips stores that are already rechunked.

    Measured on gog_upwelling (300 × 400 × 25 with wflux): compute() went from ~10 s to ~3.5 s per window with the first two fixes, before rechunking. Rechunking cuts the wvel subset read from 0.23 s to 0.05 s. Tests: tests/test_flux_io.py.

  • Particles did not leave a subset domain (advection_scheme, gcm/flux_grid.py). The “left the domain” test (REASON_BOUNDS_H) only checked j. On a grid without zonal wrap (iperio = 0, i.e. any gcm.subset in longitude for glorys/ecco2), particles crossing the east edge kept going past imt with clamped flux indices, still active. Particles reaching the west or south edge were held there by a closed wall (um = 0 at ia == 1), also still active. None of them showed up in the stop tally or as deactivated rows. All four kernels (piecewise, analytical, piecewise_local, piecewise_stencil) now stop a particle with out_of_bounds_h when it enters the outer ring of cells on a non-periodic axis (new flux_grid.edge_mask(grid); LocalFluxFields gets an edge field, the mask sliced to the patch, so patch edges inside the domain are not treated as domain edges). This changes results of subset runs: particles now stop at the edge. The north edge stops one cell earlier than before (entering row jmt instead of leaving it). Regression test: tests/test_domain_edge.py.

  • wflux is no longer read when particles do not move vertically (no_wflux=True / move_vertically = false). All three integrators read it every window and then dropped it: for global GLORYS that was a quarter of the per-window read (10.6 of 42.5 GB). LazyFluxFields.compute, compute_patch and preload gain include_wflux=True; the integrators pass include_wflux=not no_wflux.

  • GLORYS velocities sat one cell west of their T-cell: gcm.glorys.fluxes trimmed the velocity Zarrs (x=4322) with slice(1, -1) twice but the mesh files (x=4320) once, assuming the two extra velocity columns are one halo per side. They are both on the west: the Zarrs’ lon_u/lon_v/lon_w equal the mesh glamu/glamv/glamt exactly at velocity column c + 2 = mesh column c. So uflux[i] was the west face of cell i (the kernels read it as the east face), and vflux/wflux belonged to the cell to the west. Along coastlines this gave cells with inflow and no outflow, where particles froze for the rest of the run while staying active (e.g. 1,489 particles on the Ghana coast in the omi_falkor run). The velocity trim is now slice(3, -1) (_VEL_XTRIM); every surface face now has zero flux exactly when one of its cells is land (tests/test_glorys_alignment.py). The fake GLORYS source in tests/test_gcm_subset.py now has the real column layout (it encoded the old one-halo-per-side assumption). GLORYS runs made before this fix used shifted velocities everywhere and should be rerun.

  • GLORYS fill values are masked when read (seagrid.glorys): the cached velocity Zarrs hold the NetCDF fill value 9.96921e36 at land points (it is not decoded on download), and grid_definition() wrote every zero — land kmt, land dzt — as that fill value (ds.where(ds, fill) replaces falsy values). open_dataset now returns 0 for masked velocities (NaN for temp/salt), and grid_definition returns land as kmt == 0 / dzt == 0, also for grid Zarrs built before; new builds no longer write the fill value. The kernels were already protected by the fill clean-up in flux_grid.py, but other readers of the grid (e.g. lonlat_to_xy’s land check, which tests kmt == 0) were not.

  • seed.ijk(..., grid=grid) skipped the wrong cells as land: it indexed the 0-based grid.kmt with the 1-based TRACMASS i/j it seeds, so it tested the land mask of the cell diagonally up-right of each seeded one — seeding particles onto land next to coasts (99 of them over the theoretical grid) and dropping ocean cells beside it. Out-of-range i2 = imt / j2 = jmt went unnoticed because JAX clamps out-of-bounds indices. It now tests kmt[i - 1, j - 1]. Runs seeded with seed.ijk next to land (e.g. omi_source_faith.py) seed slightly different cells.

  • time_dependent_patched crashed on in-memory float32 GCMs (LazyFluxFields.compute_patch): with the theoretical adapter it raised ValueError: assignment destination is read-only from the in-place NaN/fill cleanup, because float32 patch slices come back as read-only views and — unlike compute() — compute_patch had no copy guard for that case. It now copies non-writeable slices first. Dask-backed (GLORYS/ECCO2) data was unaffected, as dask.compute returns fresh arrays.

  • xy_to_lonlat was half a cell off, and lonlat_to_xy snapped to cell centres (trajs/coords.py, used by gcm.ecco2 / gcm.glorys): the TRACMASS convention is that cell i spans i <= x < i + 1 with its centre – where the grid’s lon/lat are defined – at x = i + 0.5, which seed, lonlat_to_xy and abstemp all follow. xy_to_lonlat instead put centre i at x = i, shifting every converted position by half a cell (0.125° on ECCO2) towards +x/+y. It now interpolates between centres at x - 0.5, wraps across the zonal seam on global grids, and linearly extrapolates in the outer half cell elsewhere (it used to clamp). lonlat_to_xy now returns the true fractional position – the exact inverse of xy_to_lonlat – via Newton refinement from the nearest centre, falling back to the nearest ocean-cell centre for points on land or outside the grid. Both take a new periodic=None argument (auto- detected). Positions converted with the old xy_to_lonlat (e.g. abscarib_30_stitched.parquet) should be regenerated.

  • lonlat_to_xy failed with current pyinterp (2026.x has no pyinterp.RTree): the nearest-centre search now uses scipy.spatial.cKDTree on unit-sphere coordinates (great-circle nearest), so pyinterp is no longer needed.

  • RunMonitor’s rich progress bar showed a misleading 0/1 (never 0/0) for a genuine n_windows=0 run (monitoring.py): every place that derived a display total or remaining count from self._n_windows used self._n_windows or <fallback>, which conflates an unset (None) total with an explicitly-zero one — 0 or 1 evaluates to 1. A run with zero windows (e.g. the abstemp.run bug above, or any off-by-typo empty date range) would render a permanently frozen 0/1 bar even though the call had already returned successfully, looking exactly like a hang. All five occurrences (on_window_start, _print_window, _log_window, _start_rich, _rich_window) now check is not None instead of truthiness, so a genuine zero renders as 0/0.

  • abstemp.run silently integrated zero windows (abstemp.py): dtm2 = dtm1 + pd.Timedelta("180") was missing its unit suffix, so pandas parsed it as 180 nanoseconds rather than 180 days — nwin then computed to 0 and time_dependent_patched(..., n_windows=0) returned almost instantly having done no integration at all. The rich monitor’s progress bar (total = n_windows or 1, see above) still rendered a static 0/1, which looks like a hang even though the call already returned. Fixed to pd.Timedelta("180D").

  • new_grid/new_flux with dtm_start=None broke a fluxes callable that has no default for its dtm parameter (tests/test_integration.py): a side effect of the dtm_start default change above — three test-local lambda dtm, reverse=False: ... fixtures (test_periodic_east_wrap, test_surface_layer_no_vertical_exit, test_northern_boundary_exit_reason_bounds_h) had no default for dtm and don’t actually use it (time-invariant synthetic fields), so _fetch_slot0’s bare fn() call now raised TypeError: missing 1 required positional argument: 'dtm'. All three shipped GCM adapters (theoretical, fabboni, glorys, ecco2) already give dtm a default, so only these test fixtures were affected. Fixed by giving them dtm=None too, matching the real adapter convention.

  • abstemp.run fetched wvel over OPeNDAP despite move_vertically = false, and silently dropped its [abscarib_30.gcm] subset (abstemp.py): it called tracmass.settings.from_env(env), which — per Dynaconf’s own docstring, “the existing settings object remains the same” — returns an isolated copy of the settings without activating that environment globally. gcm.ecco2.fluxes()’s move_vertically()/gcm_subset() calls read the global tracmass.settings singleton, which was still on [default] (move_vertically = true, no subset), so include_wflux resolved to True and fetched an uncached wvel file over the network instead of using only the cached uvel/vvel. Changed to tracmass.settings.setenv(env), which activates the environment in place — matching the working pattern already used in omi_source_faith.py.

  • OPeNDAP fetches could hang for 10+ minutes per chunk against an unreachable/dead server (seagrid/glorys.py, seagrid/ecco2.py, _prepare_write_job): PydapDataStore.open(url, session=session) was called with no explicit timeout, so it fell back to pydap’s internal 120s default — combined with the requests.Session’s HTTPAdapter(max_retries=Retry(total=5, ...)), a single stalled chunk request (each dask chunk of a downloaded field makes its own OPeNDAP GET) could retry up to 6 times at ~120s each plus backoff, appearing to hang rather than fail. New module constant _OPENDAP_TIMEOUT = (10, 60) (connect, read seconds) is now passed explicitly to PydapDataStore.open, and the retry count dropped from 5 to 2, so a dead server now fails within roughly a minute or two instead of 10+ minutes per chunk. Reported via abstemp.run hanging at tracmass.new_flux(grid, "ecco2", dtm_start=dtm1, dtm_end=dtm2).

  • Particles overshot window boundaries (piecewise scheme): the step cap passed to the kernel was a fixed tseas / iter_sub counted from wherever each particle’s clock was, not the time left until the next sub-step boundary. With the default iter_sub=1 it almost never bound, so the last step of every window was a cell crossing ending past the boundary: snapshots sat exactly on cell faces (integer x / y) with tt up to hours past the time in their dtm column, and the next window started from that overshoot. The integrators now pass a per-particle cap (integrate.full_grid._substep_cap) and snap tt / ts onto the boundary when it was the limiting step (_snap_to_substep), matching Fortran TRACMASS; this covers time_dependent, time_dependent_patched, time_dependent_streaming and steady_state. piecewise_stencil accepts a per-particle dsmin_s. New tests/test_window_boundary.py.

  • GLORYS x-subset crashed with a broadcast error: gcm.glorys.fluxes applied subset i1/i2 before trimming the NEMO halo columns, which are trimmed twice from the velocity files (x=4322) but once from the mesh files (x=4320). With an x-subset the two came out 2 cells apart in width, so e.g. [omi] / [omi_falkor] (i1 = 2600, i2 = 3700) failed with operands could not be broadcast together with shapes (2, 1100, 1096) (2, 1100, 1098). The halos are now trimmed first, so i1/i2 index TRACMASS cells and a subset gives exactly i2 - i1 cells. Each subset cell still holds the same data as before (the start offset is unchanged). An x-subset also now sets iperio = 0 (in ecco2 too), so particles leaving the east edge of a regional box no longer wrap around to its west edge.

  • tracmass.run catches an environment name passed as a date: tracmass.run("omi_falkor") put the env name in the dtm_start slot and failed deep inside new_flux with a pandas DateParseError. A string dtm_start / dtm_end that does not parse as a date now raises a ValueError up front suggesting env="omi_falkor".

Changed#

  • seed.ijk takes absolute (full-grid) cell indices by default, the same frame as saved trajectories and lonlat_to_xy. With grid= it subtracts grid.i_offset / j_offset, drops cells outside the subset, and raises ValueError if no cell of the box is inside it. The new use_subset_coordinates=True keyword keeps the old subset-relative behaviour; in a settings table set use_subset_coordinates = true in [env.seed] (default false). Without a subset (offset 0) nothing changes. Seed boxes written for a subset must be shifted by i1 / j1 or get use_subset_coordinates = true.

  • Saved trajectories use full-grid (absolute) x/y on a subset grid. GridFields gains i_offset / j_offset (default 0), filled from the _i_offset / _j_offset attrs that glorys / ecco2 now put on their datasets (the subset’s i1 / j1, via the new flux_grid.subset_offset_attrs). ParticleRecorder has a grid_offset, set by manage_run_io from the grid, which it adds to x / ia and y / ja on write; a non-zero offset is also stored in the tracmass.run metadata. Previously a subset run (e.g. [omi_falkor], i1 = 2600) saved positions relative to the subset, so they could not be compared across subsets or mapped to lon/lat without knowing the subset. Integration, seeding and the returned ParticleState are unchanged (subset-relative). To match, glorys / ecco2 xy_to_lonlat / lonlat_to_xy now always use the full grid; subtract grid.i_offset / grid.j_offset to seed from lonlat_to_xy output on a subset grid. A strided subset (istep / jstep) warns, since its positions cannot be made absolute by an offset alone.

  • abstemp.closest_traj_db / closest_unconnected_db distances are now great-circle km, and wrap across the periodic i seam (abstemp.py): dist / min_dist were Euclidean grid-index distances, which overstate east-west separation away from the equator (0.25° cells narrow towards the poles) and measured ~1440 cells between points either side of the ±180° seam. _nearest_target_lut now picks the nearest target-cell centre with a KD-tree on unit-sphere coordinates (replacing the planar distance transform) and stores its tlon/tlat; the DuckDB query computes haversine distance from the particle’s lon/lat, derived from x/y on the regular global grid. It is still one hash join, so performance is unchanged. stitch_closest still uses grid-index distances.

  • new_grid/new_flux no longer hard-code a default dtm_start (gcm/flux_grid.py): dtm_start now defaults to None instead of the literal "2007-06-01". When omitted, the fluxes callable is invoked with no arguments at all, so its own default applies — for gcm.glorys.fluxes/gcm.ecco2.fluxes this falls through to tracmass.seagrid.glorys.open_dataset/tracmass.seagrid.ecco2.open_dataset’s own dtm default, which is now the single place a default date lives for those two adapters. gcm.glorys.fluxes/gcm.ecco2.fluxes themselves also changed their dtm default from "2007-06-01" to None to match. New _fetch_slot0 helper in flux_grid.py fetches slot 0 and reads back the actual resolved timestamp (from the returned dataset’s dtm variable) for backward/forward direction inference against dtm_end, so passing only dtm_end (no dtm_start) still infers direction correctly. gcm.theoretical.fluxes/gcm.fabboni.fluxes are unaffected — they are synthetic adapters with their own literal dtm defaults and don’t use seagrid.

  • tracmass.njord renamed to tracmass.seagrid: the low-level Zarr/OPeNDAP access layer (glorys.py, ecco2.py, debug.py) now lives at src/tracmass/seagrid/. All internal imports (gcm/glorys.py, gcm/ecco2.py, seed.py), docs, and top-level scripts (io_speed.py) updated to match. Historical njord-prefixed paths in earlier changelog entries above are left as-written since they describe the module under its old name at the time.

  • test_glorys.py / test_time_increase.py moved into tests/ and converted to pytest tests: both were previously manual, un-asserting scripts at the project root, run by hand via pixi run python -c "import test_glorys; test_glorys.compare()". They are now tests/test_glorys.py (test_glorys_full_grid / test_glorys_patched / test_glorys_streaming, sharing a module-scoped glorys_grid fixture) and tests/test_time_increase.py (test_theoretical_backward_run_completes / test_glorys_backward_run_completes), each with real assertions (particle count preserved, positions finite, snapshot counts, iteration bounds) in place of the old print-only output. Both files are marked pytest.mark.slow (skipped unless pytest --slow is passed) since each case runs many GCM windows or hundreds of thousands of particles; the GLORYS-dependent tests skip cleanly when the data is unreachable.

Added#

  • tracmass.new_grid() — build the static grid on its own (gcm/flux_grid.py, gcm/__init__.py, __init__.py): new_grid(fluxes, dtm_start=..., dtm_end=None) returns just a GridFields (geometry + timestep + reverse), calling the adapter once for geometry and loading no flux data. The grid can then be reused across many flux buffers on the same grid. new_flux gains a second call form, new_flux(grid, fluxes, dtm_start=..., dtm_end=None) -> LazyFluxFields, which builds only the flux buffer against an existing grid (any dtm_end passed is checked against grid.reverse). The original new_flux(fluxes, dtm_start=..., dtm_end=None) -> (GridFields, LazyFluxFields) form is unchanged and still fetches slot 0 only once internally. Shared internals extracted as _resolve_fluxes, _grid_from_ds, _lazy_from_ds, _infer_backward. Stale new_grid(..., reverse=True) references in integrate/full_grid.py and the fluxes_to_grid docstring corrected.

  • n_particle_batches is now settings-driven (config.py, runsetup.py, integrate/full_grid.py, integrate/streaming.py, tests/conftest.py, tests/test_integration.py): the n_particle_batches argument of steady_state, time_dependent, time_dependent_patched and time_dependent_streaming now defaults to a sentinel; when not passed it is resolved by runsetup.resolve_n_particle_batches() from the n_particle_batches setting (new tracmass.config.n_particle_batches(), default 1 — no batching, clamped to >= 1). An explicit argument still wins. Both full-grid integrators (steady_state, time_dependent) act on it (particle-batched lax.while_loop via _integrate_batched); the streaming integrators resolve it for API symmetry but do not batch. steady_state gained the n_particle_batches keyword in this release — previously it always ran the whole population through a single lax.while_loop. The tracmass.toml [default] n_particle_batches = 8 key is now actually consumed instead of being ignored unless every call site passed it by hand. tests/conftest.py’s autouse fixture pins it to 1 for the test run.

Fixed#

  • tracmass.trajs.xy_to_lonlat / lonlat_to_xy no longer return garbage for transposed or 1-D lon/lat (trajs/coords.py): _unpack_grid blindly read imt, jmt = lon_arr.shape, so passing np.meshgrid(lon_1d, lat_1d) (default indexing="xy", shape (jmt, imt)) or 1-D coordinate vectors silently scrambled the axes — e.g. an ECCO2 path came back with longitude driven by the latitude index and clipped latitudes. A new _orient_lonlat helper now accepts 1-D lon (imt,) / lat (jmt,) (broadcast with indexing="ij") and auto-detects a (jmt, imt) 2-D layout, transposing lon/lat/kmt to TRACMASS (imt, jmt) with a warnings.warn.

  • tracmass.toml no longer fails to parse: the shipped settings file defined gcm both as a string (gcm = "theoretical") and as a table ([default.gcm], twice), which is a TOMLDecodeError. Adapter selection and the spatial subset are now sibling keys of one [env.gcm] table (gcm.name / gcm.subset).

Added#

  • [env.gcm] table wires adapter + spatial subset into the GCM modules (config.py, gcm/flux_grid.py, tracmass.toml): tracmass.toml now carries a [env.gcm] table with name (the adapter — "glorys", "ecco2", "theoretical", …) and an optional subset inline table. tracmass.config.gcm_name() returns the adapter name (a bare gcm = "<name>" string is still accepted). tracmass.config.gcm_subset() reads gcm.subset and returns a normalised 9-key dict (i1, i2, istep, j1, j2, jstep, k1, k2, kstep — each int or None; missing keys and "@none" / false become None). tracmass.config.apply_gcm_subset( module) copies the non-None entries into module.settings["subset"], and new_flux calls it for the adapter it resolves (from a string name or the passed callable’s module) before the first fluxes() call — so settings.setenv("omi"); new_flux("glorys", ...) now honours [omi.gcm] subset with no manual gcm.glorys.settings["subset"][...] = .... Only the keys present are written, so a partial table composes with limits set directly on the module. Generated-grid adapters (theoretical, fabboni) expose no settings["subset"] and are left untouched. New tests/test_gcm_subset.py.

  • move_vertically setting drives no_wflux / include_wflux (config.py, runsetup.py, integrate/full_grid.py, integrate/streaming.py, gcm/ecco2.py): a single tracmass.toml key now controls vertical advection end-to-end. tracmass.config.move_vertically() reads it (default True — particles do move vertically, matching the old no_wflux=False / include_wflux=True defaults). The no_wflux argument of steady_state, time_dependent, time_dependent_patched and time_dependent_streaming now defaults to a sentinel; when not passed it is resolved by runsetup.resolve_no_wflux() to not move_vertically(). An explicit no_wflux=True/False still wins. gcm.ecco2.fluxes gained include_wflux=None (the new default), resolved to move_vertically() so the wvel OPeNDAP dataset is skipped when vertical motion is off. Setting move_vertically = false in a tracmass.toml section (e.g. [tempvel], [omi]) now turns wflux off for both the flux load and the advection kernel with no per-call arguments. tests/conftest.py’s autouse fixture also pins move_vertically = True for the test run. The [omi] section’s previously-dead no_wflux = true key is now move_vertically = false, and omi_source_faith.py activates [omi] (settings.setenv("omi")) and drops its explicit no_wflux=False — so omi runs now honour the section’s long-standing “no vertical motion” intent instead of the leftover no_wflux=False argument.

  • monitor / recorder are now optional on the multi-window integrators (runsetup.py — new, integrate/full_grid.py, integrate/streaming.py): time_dependent, time_dependent_patched and time_dependent_streaming previously required the caller to build a RunMonitor / ParticleRecorder and drive its context manager:

    with (RunMonitor(backend="rich") as mon,
          ParticleRecorder(f"trm_{dtm.date()}.parquet") as rec):
        tracmass.time_dependent_patched(..., monitor=mon, recorder=rec)
    

    A new @manage_run_io decorator (in tracmass/runsetup.py) wraps the three integrators so monitor and recorder each accept an instance (used as-is, caller keeps ownership), a string ("rich" → RunMonitor(backend=...), "run.parquet" → ParticleRecorder(...), built and closed for you), None (feature disabled), or omitted (built from settings). Anything the integrator constructs is entered on a contextlib.ExitStack and exited automatically — on normal return and on exception — so the with block is no longer needed:

    tracmass.time_dependent_patched(state, grid, flux, n_windows=nwin)
    

    Settings keys read from tracmass.toml (via tracmass.config) for the omitted case: runmonitor_backend (monitor backend, default "print"), output_name (recorder filename stem — its presence also enables auto-recording), auto_record (enable auto-recording with the default stem trm_{active-env}), and base_output_dir (output directory). The auto-built recorder path is {base_output_dir}/{stem}_{forward|reverse}_{dtm_start.date()}.parquet — the direction token is taken from grid.reverse and the date from fluxes.dtm, so a forward and a backward run in the same settings environment no longer collide. When neither output_name nor auto_record is set, an omitted recorder resolves to None (no file) — the historical default, so existing call sites are unaffected. tests/conftest.py gained an autouse fixture that neutralises these settings during the test run so a stray tracmass.toml in the working directory cannot redirect test output or switch on the rich display.

  • gcm.ecco2: full ECCO2 cube92 GCM adapter (gcm/ecco2.py, njord/ecco2.py): gcm.ecco2.fluxes(dtm, include_wflux=True) now returns 3-D fluxes (uflux, vflux, wflux) plus geometry, verified end-to-end through new_flux → steady_state/time_dependent particle advection. njord/ecco2.py gained a wvel entry in vardict and grid_definition() now stores dxv/dyu (zonal/meridional T-cell face widths) alongside dxdy so the gcm layer can build uflux/vflux itself. lonlat_to_xy/xy_to_lonlat reuse the shared curvilinear-grid utilities in trajs/coords.py, same as gcm.glorys. Two ECCO2-specific data-quality issues were found and worked around:

    1. the APDRC GDS gateway declares missing_value/_FillValue as -1e23 but actually serves exact 0.0 for land/below-seafloor uvel/vvel cells — njord.ecco2.calc_kmt_dzt derives kmt/dzt from exact-zero equality rather than NaN/CF decoding.

    2. wvel has its own invalid-cell pattern that is not implied by uvel’s (e.g. under ice shelves wvel is invalid at levels where uvel/vvel are valid), and is inconsistently NaN vs. the raw undecoded -1e23 sentinel — gcm.ecco2._clean_velocity zeroes any uvel/vvel/wvel value with abs() > 1e10 (no real ECCO2 current comes close to that magnitude) before computing fluxes. njord/ecco2.py is adapted from njord/glorys.py for the APDRC GDS ECCO2 cube92 daily reanalysis (uvel, vvel, wvel, theta, salt), a global 0.25° regular lat/lon product served as one long OPeNDAP time series per variable (1992-01-04 – 2025-12-31) rather than per-day files. open_dataset selects the nearest day, caches it as a local Zarr, and renames dims to the codebase’s (z, y, x) convention. grid_definition has no mesh_hgr/mesh_zgr files to draw on, so it computes cell geometry analytically from the lat/lon spacing. The server’s native lon runs 0.125..359.875 (Prime Meridian at the array seam, antimeridian at the array centre); _prepare_write_job now rolls every cached field through _center_on_prime_meridian to -179.875..179.875 before writing, so 0° sits at the array centre and Atlantic/European regions are contiguous for i1:i2 subsetting, matching the convention gcm.glorys already used. Verified end-to-end with a Gulf Stream subset (lon -80..-50) expressed as plain contiguous indices. njord/ecco2.py also silences two more benign, unfixable warnings from this data source at module import time (not per call site — dask’s threaded scheduler can run the lazy time CF-decode in a worker thread after a scoped with warnings.catch_warnings(): in the calling thread has already restored the filters, and catch_warnings is documented as not thread-safe, so a scoped filter intermittently leaked through): xr.SerializationWarning (APDRC declares time units as "days since 1-1-1 00:00:0.0", an unpadded year xarray auto-corrects but warns about) and pydap’s "PyDAP was unable to determine the DAP protocol" (its own suggested fix — a dap2:// URL prefix — makes pydap rewrite the scheme to https internally, bypassing the session.mount("http://", ...) retry adapter, so the warning is silenced instead of following that suggestion). Verified across a fresh single-variable fetch, a cached-Zarr reopen, the threaded merged uvel+vvel fetch, and a fresh grid_definition build — no warnings reach the caller in any of them. njord/ecco2.py also gained _compute_with_retry: ECCO2 daily fetches are 100-200 MB per variable over a slow remote link and occasionally fail mid-transfer with a requests.exceptions.ChunkedEncodingError (IncompleteRead) from the APDRC server — not covered by the HTTPAdapter/Retry config in _prepare_write_job, which only retries failures during connection setup, not a drop while streaming an already-established response deep inside pydap’s r.content. This previously killed a for dtm in dtmvec: ecco2.open_dataset(dtm) loop partway through. Both .compute() call sites (retrieve, and the merged uvel+vvel path in open_dataset) now route through it (3 retries, linear backoff). Fixing this also removed a latent bug: open_dataset’s local import dask (used only to reach dask.compute around its own dask boolean parameter, which shadows the module for the whole function) could leak the dask module into the dask=dask keyword of its own recursive call whenever a fetch actually ran, silently forcing dask=True regardless of what the caller passed.

Changed#

  • seed.mask() now takes a raw (jmt, imt[, kmt]) mask — j (lat) first (seed.py): previously mask() silently assumed (imt, jmt) — the opposite of how raw grid data actually comes out of xarray/NetCDF sources (e.g. njord.ecco2.grid_definition().kmt.values, shape (720, 1440) = (jmt, imt)). Passing such an array in directly did not error; it silently swapped the i/j axes (confirmed: ja reaching 100 on a grid with jmt=80). mask() now interprets axis 0 as j and axis 1 as i, matching the native orientation of grid fields, and warns (mask.shape [0] > mask.shape[1]) when the array looks like it was passed already transposed to the old (imt, jmt) convention instead. Added a docstring (mask() previously had none). ijk()/from_df() are unchanged — they still index grid.kmt in its own (imt, jmt) convention, which is correct since GridFields.kmt really is built that way by the gcm/*.py adapters. Existing callers of mask() that were passing an already (imt, jmt)-transposed array must transpose their input the other way (or stop transposing, if they were compensating for the old bug).

Fixed#

  • 8.7× IO speedup for time_dependent_patched with GLORYS12 (njord/glorys.py, rechunk_grid.py): three layered fixes eliminated Dask rechunk overhead that dominated patch reads:

    1. grid_definition() was opening the grid zarr with chunks="auto", which collapsed dyu/dxv/dxdy into single 3059×4318 global slabs despite the zarr store having 200×200 chunks. Changed to chunks={} so xarray honours the zarr store’s native chunk layout.

    2. The grid zarr was rechunked with dzt at (4, 200, 200) while the velocity zarrs use (50, 200, 200). Multiplying uvel * dzt triggered a Dask rechunk of the full z-column before every patch read. rechunk_grid.py now uses z=50 for dzt, matching the velocity zarr exactly.

    3. rechunk_grid.py added for in-place rechunking of the GLORYS12 grid-definition zarr on /data/fast. Total wall time for a 10-window run dropped from ~1m49s to ~12.5s (IO from ~9s/window to ~700ms/window).

Changed#

  • new_flux API: dtm renamed to dtm_start, dtm_end added; reverse removed (gcm/flux_grid.py): integration direction is now inferred from the relationship between dtm_start and dtm_end rather than a separate reverse flag. When dtm_end < dtm_start the run is backward: grid.timestep stays positive, lazy.dtm decrements per window, and velocity fluxes (uflux, vflux, wflux) are lazily negated inside new_flux/update_flux — no wrapper function needed. This fixes the pre-existing bug where reverse=True made timestep negative, causing tseas < 0 inside _integrate_window and completely wrong sub-cycle positions. Old code using new_flux(fluxes_fn, dtm=..., reverse=True) must be updated to new_flux(fluxes_fn, dtm_start=..., dtm_end=<earlier timestamp>).

  • reverse parameter removed from GCM adapter fluxes() functions (gcm/glorys.py, gcm/theoretical.py, gcm/fabboni.py): adapters now always return forward (positive) fluxes; negation is handled by the framework. The nff sign-multiplier in fabboni._velocity is also removed (it evaluated velocities at negative model times rather than negating forward velocities, which was physically incorrect).

  • omi_source_faith.py: _glorys_backward wrapper removed The backward GLORYS run is now configured directly: new_flux(gcm.glorys.fluxes, dtm_start="2007-12-31", dtm_end="2007-01-01").

Fixed#

  • REASON_BOUNDS split into REASON_BOUNDS_H (3) and REASON_BOUNDS_V (8) (particle.py, advection_scheme/full_grid.py, advection_scheme/stencil.py, __init__.py): horizontal exits (j out of range) and vertical exits (k out of range) now carry distinct reason codes, making it easier to diagnose whether particles leave through open lateral boundaries or through the surface/bottom. Numeric value 3 is preserved for REASON_BOUNDS_H (backward-compatible with existing parquet files where reason=3 was always a lateral exit); REASON_BOUNDS_V is assigned code 8.

  • Surface-layer particles spuriously exiting through kb=0 (REASON_BOUNDS_V) (advection_scheme/full_grid.py): at ka=1, jnp.clip(ka-2, 0, km) clips to 0, making the shallower-face flux um_z identical to the deeper-face flux uu_z. When the inter-layer flux is upward (upwelling), um_z < 0 yields a finite dsd, pushing the particle above the sea surface (kb=0 < 1). Fixed by zeroing um_z (and vm_z in the analytical scheme) when ka==1; the sea surface is treated as a rigid wall. Eliminates ~140 K spurious REASON_BOUNDS exits per 60 M particle run.

  • njord.glorys.open_dataset progress bar garbled during parallel downloads (njord/glorys.py): when all three flux files needed downloading, ThreadPoolExecutor ran three retrieve calls concurrently, each owning its own Dask ProgressBar. The bars interleaved on stdout, producing garbled output. retrieve now accepts a show_progress parameter (default True); open_dataset passes False for threaded calls and instead prints a single "Downloading N file(s) … Done." line around the parallel fetch.

Changed#

  • njord.glorys.retrieve is now significantly faster (njord/glorys.py): depth chunking changed from -1 (all 50 levels per chunk) to 10, giving Dask ~5× more parallel OPeNDAP requests per tile. write_job.compute() now uses scheduler="threads", num_workers=8 for concurrent chunk writes. open_dataset(data_var=None) now pre-fetches uvel, vvel, and wvel concurrently via ThreadPoolExecutor(max_workers=3) before opening the cached files, cutting wall time for a full three-variable download by roughly 3×.

  • time_dependent_patched now grows its patch dynamically (integrate/streaming.py, advection_scheme/stencil.py): previously the patch shape was fixed at initialisation; as particles spread over long runs (e.g. year-long GLORYS integrations) they hit the artificial patch boundary and were killed with REASON_BOUNDS. The patch now grows each window when the active-particle bounding box exceeds the current patch size, rounded up to patch_granularity (default 100 cells) to limit JAX recompilations. piecewise_local is decorated with @functools.lru_cache so the compiled kernel is reused whenever the same patch shape recurs. A new patch_granularity parameter controls rounding.

  • RunMonitor rich backend: cumulative stopped count + end-of-run reason tally (monitoring.py): the Stopped ∑ column now shows the running total of permanently stopped particles rather than the per-window count, making it easier to track overall particle loss at a glance. The summary row now also shows the full reason breakdown (e.g. land:12, time:5) so the dominant stop cause is visible at the end of the run. The print and log backends likewise gain a cumulative stopped/reason line in their summary output.

Fixed#

  • Snapshot rows no longer carry stale non-zero reason codes (integrate/full_grid.py, integrate/streaming.py): the window-boundary reactivation logic used ts >= ts_window_end and kmt > 0 as guards, but did not check reason == 0. Particles stopped for REASON_SLOW, REASON_PROGRESS, REASON_ITERLIMIT, or REASON_DXYZ near a window boundary could be erroneously reactivated with active=True while retaining their non-zero reason, causing confusing snapshot rows where event_type="snapshot" but reason != 0. Fix: all three reactivation blocks now require reason == 0 (paused, not permanently stopped).

  • time_dependent_streaming I/O now O(1) per window (integrate/streaming.py, gcm/flux_grid.py): previously extract_stencil called dask.compute() once per sub-cycle for every cell crossing, reloading the same large Zarr chunks on each call. Fix: new LazyFluxFields.preload() method issues a single batched dask.compute() for all flux variables at the start of each window, returning a numpy-backed LazyFluxFields. All extract_stencil calls within the window then use pure numpy indexing with no further Zarr I/O. The preload time appears in the monitor’s t_io column; sub-cycle t_io contributions drop to near zero.

  • extract_stencil overflow on float32 cast (gcm/flux_grid.py): fill values (e.g. GLORYS ~9.97e36) were cast to float32 before the nan_to_zero cleanup ran, causing RuntimeWarning: overflow encountered in cast. Fix: move .astype(np.float32) to after the nan_to_num and threshold-zeroing pass.

Added#

  • tracmass.trajs.speeds(df, dxdy) — computes inter-snapshot speeds (m/s) for all particles in a trajectory DataFrame. Accepts a scalar cell area (uniform grid), a 2-D dxdy array (spatially varying), or a GridFields object directly. Returns a DataFrame with id, window, dt, dx_m, dy_m, speed.

  • time_dependent_patched and time_dependent_streaming now accept monitor and n_particle_batches — monitor is the same RunMonitor interface as time_dependent (defaults to RunMonitor(backend="print"); pass None to suppress output); per-window timing is split into t_io, t_jax, t_py and n_sub/reason_tally are reported to on_window_end. n_particle_batches is accepted for API compatibility but has no effect (no JAX while_loop scratch allocation to reduce).

  • tracmass.trajs.xy_to_lonlat(x, y, grid) and tracmass.trajs.lonlat_to_xy(lon, lat, grid) — generic coordinate-conversion functions extracted into tracmass/trajs/coords.py and re-exported from tracmass.trajs. Both accept any object with .lon/.lat attributes (or a 2/3-tuple), making them usable with any GCM adapter. gcm.glorys.xy_to_lonlat and gcm.glorys.lonlat_to_xy are now thin wrappers that pass _glorys_grid() to the generic implementations.

  • gcm.glorys.lonlat_to_xy(lon, lat) — converts geographic coordinates (degrees) to fractional 0-based TRACMASS (x, y) grid positions for particle seeding. Uses pyinterp.RTree (pip install pyinterp) to build a spherical spatial index from ocean T-cell centres (kmt > 0) and performs a nearest-neighbour query. Particles are placed at the cell centre (offset = 0.5), consistent with seed.expand_sub. Applies the same spatial subset and x-halo trim as fluxes()::

    import numpy as np
    from tracmass.gcm import glorys as gcm
    
    lon = np.array([-30.0, 10.0])
    lat = np.array([45.0, 55.0])
    x, y = gcm.lonlat_to_xy(lon, lat)
    
  • gcm.glorys.xy_to_lonlat(x, y) — converts fractional 0-based TRACMASS (x, y) positions (as stored in the trajectory Parquet x/y columns) back to longitude and latitude in degrees. Uses bilinear interpolation from the four surrounding T-cell centres; interpolation is performed in Cartesian (X, Y, Z) space on the unit sphere so it is correct for curvilinear (NEMO ORCA) grids and handles the ±180° seam transparently::

    import pandas as pd
    from tracmass.gcm import glorys as gcm
    
    df = pd.read_parquet("trm_omi.parquet", columns=["x", "y"])
    lon, lat = gcm.xy_to_lonlat(df["x"], df["y"])
    
  • tracmass.free_device_memory() — releases GPU memory between runs in an interactive session. Calls jax.clear_caches() (evicts compiled kernels and the kmt/dxdy XLA constants they pin), gc.collect(), and jax.effects_barrier() (waits for async ops to complete). Safe on CPU-only installations::

    result = tracmass.time_dependent(state, grid, lazy, ...)
    tracmass.free_device_memory()   # GPU back to baseline
    
  • ParticleRecorder(filepath, metadata=...) — optional metadata dict of JSON-serialisable key/value pairs is now embedded in the Parquet schema metadata under the key "tracmass.run" and can be read back without scanning any rows::

    import json, pyarrow.parquet as pq
    meta = json.loads(pq.read_schema("traj.parquet").metadata[b"tracmass.run"])
    
    with ParticleRecorder("traj.parquet", metadata=dict(
        gcm="glorys", integration="time_dependent",
        n_windows=50, no_wflux=True,
    )) as rec:
        tracmass.time_dependent(..., recorder=rec)
    

Fixed#

  • GLORYS fill values in dzt/kmt — NaN particle positions and false ocean mask. The GLORYS grid definition file stores land cells using the NEMO fill value (9.96921e36) rather than 0/NaN for kmt (bottom-depth index) and dzt (layer thickness). When xarray read the file back without applying fill-value masking, two bugs arose:

    1. Seeding in land cells: seed.ijk(grid=grid) filters ocean cells via grid.kmt.astype(bool). With kmt = 9.96921e36 for land cells, the bool cast is True, so particles were seeded inside land.

    2. NaN cascade in particle positions: dzt = 9.96921e36 in the advection kernel causes dxyz_raw to overflow float32 → Inf, making dsmin_arr = dt / Inf = 0, ds = 0, and dt_particles = 0 × Inf = NaN. On the very next sub-step intrpg_pp = fmod(NaN, 1) = NaN, which propagates to dxyz = NaN, ds = NaN, and finally x1 = NaN. The integer cast floor(NaN).astype(int32) yields 0 → ia = 1, ja = 1. Because kmt[0, 0] also contained the fill value (≠ 0), the particle was not detected as land and remained “active” forever with ts = NaN, so the window loop ran all max_sub_iter = 10 000 sub-cycles every window while the monitor showed 0 stopped particles.

    Fix (gcm/flux_grid.py):

    • LazyFluxFields.compute() / LocalFluxFields.compute(): extend the NaN→0 and fill-value→0 zeroing to dzt and zstot (was only applied to uflux, vflux, wflux).

    • _static_from_ds(): when loading kmt into GridFields, replace any value that is NaN or above _FLUX_FILL_THRESHOLD (1e20) with 0 so that land cells are correctly identified as kmt = 0.

    After the fix, land cells are excluded from seeding and dxyz_raw is always finite, eliminating the NaN cascade.

  • Western periodic wall crossing — unbounded negative x positions. When a particle in cell ia=1 crossed the western wall on a periodic grid (iperio=1), the kernel set x1 = iam = imt and ib = imt. The subsequent cyclic wrap x1 >= imt → x1 - imt converted x1 = imt to x1 = 0, leaving state ia=imt, x=0 — geometrically inconsistent (cell imt spans [imt-1, imt]; x=0 is 4317+ cells away). On every following sub-step the crossing-time formulae produced large negative displacements that accumulated into the enormous negative x values (e.g. −15 269) visible in trajectory Parquet files.

    Fix: change the cyclic wrap condition from x1 >= imt to x1 > imt (strict), matching the original Fortran mod_pos_tstep.F90 exactly. x1 = imt is left unwrapped intentionally — the already-present r0_adj special case in _calc_pos_1d / _calc_pos_analytical_1d handles r0 = imt, ii = 1 correctly on the next sub-step. Applied to both the piecewise-linear and time-analytical advection schemes.

  • njord/glorys.py — grid_definition() stored longitude in the latitude field. A copy-paste error caused ds["lat"] = hgr["nav_lon"]. Users with a cached trm_glorys12v1_grid_def.nc must regenerate it::

    from tracmass.njord import glorys as _g
    _g.grid_definition(force=True)
    
  • no_wflux=True — wflux still uploaded to GPU. When no_wflux=True the wflux array (~212 MB for a global GLORYS12 grid) was still present in the computed FluxFields and uploaded to GPU via jax.device_put, contributing to CUDA_ERROR_OUT_OF_MEMORY errors. wflux is now stripped (set to None) before the device_put call in both steady_state and time_dependent.

  • test_ned_grid.py _allclose — failed on NaN and large fill values. _load_single_slot returns raw xarray values (NaN at land cells; float64 fill ≈ 9.97e36), while LazyFluxFields.compute() normalises these to zero. _allclose now applies the same normalisation (clip → float32 → nan_to_num → fill-value zeroing) before comparing.

  • test_speed.py — snapshot column no longer exists. Both speed tests read a snapshot column that was replaced by window (int32) in the ParticleRecorder schema. Tests now request window, sort by it, and filter on event_type == "snapshot".

  • test_performance.py — spurious tseas argument. The timed-run call in _timed passed tseas=tseas to the integration function; the variable was never defined, raising NameError. Argument removed.

Fixed — _allclose in test_ned_grid.py now handles NaN and large fill values#

test_new_grid_slot0_matches_dtm[glorys], test_new_grid_slot1_matches_dtm_plus_timestep[glorys], and test_new_grid_vflux_slot0_matches_dtm[glorys] were failing because _load_single_slot returns raw xarray values (NaN at masked land cells, possible float64 fill values ≈ 9.96921e+36), while LazyFluxFields.compute() applies NaN→0 and fill-value→0 normalisation before stacking. np.allclose with equal_nan=False (the default) then returned False even when the ocean- cell values agreed exactly.

Fixed _allclose to apply the same normalisation (clip → float32 → nan_to_num → fill-value zeroing) to both arrays before comparing, mirroring what compute() does. The overflow RuntimeWarning from casting float64 fill values directly to float32 is also eliminated by the clip step.

Fixed — snapshot column renamed to window in test_speed.py#

test_speed_never_exceeds_3ms and test_initial_speed_matches_fabboni were reading a snapshot column from the trajectory Parquet file, which no longer exists. The ParticleRecorder schema uses window (int32) for the GCM window index. Fixed both tests to request window instead of snapshot, sort by window, and filter on event_type == "snapshot" to exclude deactivation rows from the speed calculation.

Fixed — spurious tseas argument in test_performance.py _timed helper#

The timed run call in _timed passed tseas=tseas to the integration function, but tseas was never defined in the helper and none of the integration functions (time_dependent, time_dependent_patched, time_dependent_streaming) accept that argument. Removed the stray keyword argument; all 15 parametrised performance tests now pass.

Fixed — corner-crossing into land cells (REASON_LAND) in piecewise and analytical schemes#

When a particle simultaneously exits through two cell walls (e.g. north and east), the corner logic in _subcycle_one_particle assigned the diagonal neighbour cell unconditionally, without checking whether that cell is ocean or land. In forward mode this is rare because boundary currents flow along coasts. In reverse mode the negated circulation pushes particles toward coastlines, making diagonal corner hits far more frequent and causing spurious REASON_LAND deactivations.

Fix: added three pre-computed ocean masks (_ne_ocean, _nw_ocean, _se_ocean) using the existing closed-over grid.kmt constant. Each corner assignment is now guarded by & _*_ocean; if the diagonal cell is land the ib/jb update is skipped and the particle remains at the cardinal crossing (north or east face). The same fix is applied identically to both the piecewise-linear and analytical schemes in advection_scheme/full_grid.py. Performance cost is negligible: three extra kmt index lookups (XLA constants) and four conditional selects per particle per sub-cycle.

Fixed — ValueError: assignment destination is read-only in LazyFluxFields.compute()#

When flux arrays were already float32, the astype() path was skipped and np.nan_to_num(copy=False, ...) attempted to modify the raw numpy array from xarray .values in-place. For numpy-backed (non-Dask) DataArrays this array is read-only. Fix: added an elif not t0.flags.writeable branch that copies the arrays before in-place NaN/fill-value replacement.

Added — RunMonitor class for structured run monitoring#

Replaced all inline print() calls in time_dependent with a RunMonitor context-manager that supports three output backends:

  • "print" (default) — reproduces the original stdout output, now including per-window timing split into io, jax, and py phases plus a totals summary.

  • "log" — emits one logging.info record per window with structured extra= fields (win, n_active_before/after, n_stopped, t_io, t_jax, t_py, …) suitable for log aggregators, plus a summary record on completion.

  • "rich" — live terminal dashboard powered by rich: scrolling table of window results with timing columns, progress bar with spinner, and a totals row.

time_dependent gains a monitor keyword argument (default RunMonitor() for backward-compatible print output; pass None to suppress all output). RunMonitor is exported from the top-level tracmass namespace.

Fixed — GPU OOM (RESOURCE_EXHAUSTED) when using n_particle_batches on GPU#

When lax.while_loop was called outside jax.jit with closed-over fluxes tensors, JAX embedded the concrete flux data as jaxpr constants, producing a distinct HLO (and thus a new ~5 MB compiled CUBIN) for every _integrate_window call. After ~8 calls the accumulated CUBINs consumed enough VRAM that the next XLA scratch allocation failed with RESOURCE_EXHAUSTED: Failed to allocate 44–89 MiB. The failure always occurred on the 9th call regardless of the number of batches or windows.

Fix: extract the entire lax.while_loop body into a module-level _run_window function decorated with @jax.jit and static_argnames=(step_fn, tseas, iter_sub, max_sub_iter, timax_seconds, pace_factor). Passing fluxes as a dynamic (non-static) JIT argument makes JAX trace it by abstract type (shape + dtype) only, guaranteeing a single compiled XLA program for all windows and batches. Additionally, jax.effects_barrier() is called after each batch in _integrate_batched to flush deferred cudaFree callbacks before the next batch allocates. State is kept as numpy (CPU) between windows when n_particle_batches > 1, so GPU holds nothing outside of _integrate_window calls.

Fixed — residual boundary fluxes causing spurious REASON_LAND in analytical scheme#

_subcycle_one_particle_analytical in advection_scheme/full_grid.py read uflux/vflux values toward land-adjacent faces without checking whether the neighbor cell is ocean or land. GLORYS/NEMO no-flow boundary faces should be exactly 0.0, but small residual values (below the _FLUX_FILL_THRESHOLD = 1e20 fill-value threshold) can survive in practice due to numerical noise or imperfect masking. A non-zero flux toward a land cell produces a finite crossing time, moving the particle into the land cell and triggering REASON_LAND.

Fix: after reading uu_x/um_x/vv_x/vm_x and uu_y/um_y/vv_y/ vm_y, zero all four fluxes for any direction whose neighbor cell has kmt == 0. The same land-neighbor kmt-check pattern already applied to the piecewise scheme (_subcycle_one_particle) is now also applied here.

Fixed — NaN flux values causing spurious REASON_LAND deactivations#

GLORYS/NEMO Zarr files store land-masked cells as NaN (via xarray’s _FillValue mechanism) rather than as 0.0. Some ocean cells adjacent to land have NaN rather than 0.0 stored at their boundary face, because NEMO’s output mask treats those U/V-grid faces as land even though the physical no-flow boundary condition is 0.

When land-masked flux values are not decoded to NaN by xarray (e.g., when the Zarr fill-value round-trip leaves some cells with the raw GLORYS fill value 9.96921e+36), those huge values enter _calc_pos_1d and cause exp(fill_value × dt) to overflow to +inf. The particle position becomes inf, and jnp.floor(inf).astype(int32) produces a garbage cell index that may point to a kmt=0 (land) cell, causing spurious REASON_LAND deactivations. Verified by inspecting the parquet trajectory file: max(y) = inf for deactivated particles.

Fix: in LazyFluxFields.compute(), LazyFluxFields.compute_patch(), and extract_stencil(), apply np.nan_to_num(..., nan=0.0, posinf=0.0, neginf=0.0) followed by zeroing any value whose absolute value exceeds _FLUX_FILL_THRESHOLD = 1e20, for uflux, vflux, and wflux arrays before converting to JAX arrays. Zero is the correct physical value (no flow through a masked face), so this substitution is both safe and correct.

Fixed — seed.mask placing particles one cell outside intended position#

seed.mask builds 1-based cell indices with np.mgrid[1:imt+1, 1:jmt+1] but then used them as 0-based array indices to look up the mask array: mask[ivec, jvec]. On a 0-based numpy/JAX array this silently shifts every cell one step, and the last row/column clips via JAX’s out-of-bounds clamping.

Additionally, expand_sub with sub_dist="rect_xyz" placed particles at ivec + 0.5 (cell ia+1) instead of ivec - 0.5 (cell ia). Cell ia in TRACMASS coordinates spans [ia-1, ia], so the centre is ia - 0.5. With linspace(0, 1, …)[1:-1] the single sub-particle offset was +0.5, pushing every particle one full cell east/north of its intended grid cell.

For particles seeded in the last ocean column before a land boundary, both bugs combined to place them inside the land cell, causing immediate REASON_LAND deactivations in window 1 (thousands per window).

Fix:

  • Changed mask array lookups to mask[ivec-1, jvec-1] (0-based, correct).

  • Changed expand_sub linspace from linspace(0, 1, …) to linspace(-1, 0, …) so offsets span (-1, 0), placing the single sub-particle at -0.5 relative to ivec, i.e. the cell centre ivec - 0.5.

Fixed — njord/glorys.open_dataset corrupting zero velocities at land boundaries#

open_dataset called ds.where(ds, 9.96921e36) after loading velocity data from Zarr. In xarray/numpy, where(cond, other) treats the data values as a boolean condition: zero is falsy, so any cell with exactly 0.0 velocity was replaced with the fill value 9.96921e36.

In NEMO/GLORYS12, the U/V-velocity at the face between an ocean cell and a land cell is exactly 0.0 (no-flow boundary condition stored correctly in the Zarr file). After the bad where call that cell’s velocity became ~10^37, making the east/north wall crossing time essentially zero for any particle adjacent to land. The particle would cross into the land cell and be deactivated with REASON_LAND every window, causing a steady loss of thousands of particles per window near coastlines.

Fix: remove the ds.where(ds, 9.96921e36) call. Zarr files already encode masked (land) cells as NaN via xarray’s _FillValue mechanism; the explicit where was both redundant for NaN cells and actively harmful for zero-velocity boundary cells.

Fixed — GLORYS geometry fields off-by-one in x, causing wrong land checks#

gcm.glorys.fluxes builds uflux and vflux by applying isel(x=slice(1, -1)) to the raw NEMO ORCA velocity arrays. This trim removes the two halo columns so that TRACMASS flux cell i=0 corresponds to NEMO T-cell at original x-position 1. However, the geometry fields kmt, dxdy, and dzt were stored without the same trim, so kmt[i] described NEMO T-cell i instead of the correct T-cell i+1.

Consequence: the land-mask check at the east boundary was off by one. A particle whose eastern neighbour was a land cell could cross into it because kmt[i] (the halo cell, which is ocean) was checked instead of kmt[i+1] (the actual eastern neighbour, which is land). Over many windows, particles accumulated on land despite the GLORYS velocity field being physically land-masked.

Fix: apply the same isel(x=slice(1, -1)) trim to kmt, dxdy, dzt, and zstot so all geometry arrays share the same x-indexing as the flux arrays.

Added — stencil caching and inactive-particle skipping in time_dependent_streaming#

Two I/O optimisations reduce the number of Dask chunk reads per sub-cycle in time_dependent_streaming (the per-particle stencil integration mode used for widely scattered particles on large grids such as GLORYS12):

Skip inactive particles — extract_stencil now accepts an optional mask parameter (boolean array of shape (N,)). When mask is None (the default), only active particles (state.active) are read from the Zarr/Dask backend; deactivated particles are replaced with zeros in the returned FluxStencil. This avoids loading Zarr chunks for particles that have already stopped.

Stencil cache — time_dependent_streaming now tracks the per-particle cell indices (ia, ja, ka) between consecutive sub-cycles. Before each sub-cycle it computes a moved mask (active particles that crossed a cell wall since the previous step) and calls extract_stencil(mask=moved) to reload stencil values only for those particles. Particles that remain in the same cell reuse their previously extracted stencil via the new _update_stencil helper (a jnp.where merge, no Dask I/O). The cache is reset at the start of every GCM window when new flux data is loaded.

Together, these changes can eliminate the majority of dask.compute calls within a window when particle speeds are moderate (< 1 cell crossing per sub-cycle on average) or when a significant fraction of particles have been deactivated.

Changed — TrajectoryRecorder + DeactivationRecorder merged into ParticleRecorder#

TrajectoryRecorder and DeactivationRecorder have been replaced by a single ParticleRecorder class that writes both trajectory snapshots and deactivation events to one Parquet file, distinguished by the event_type column ("snapshot" or "deactivated").

Unified schema (13 columns): event_type, window, dtm, id, x, y, z, ia, ja, ka, tt, t0, reason.

All three multi-window integration functions (time_dependent, time_dependent_patched, time_dependent_streaming) now accept a single recorder: ParticleRecorder | None parameter.

# Before
with (
    tracmass.TrajectoryRecorder("traj.parquet") as rec,
    tracmass.DeactivationRecorder("deact.parquet") as drec,
):
    tracmass.time_dependent(..., recorder=rec, deactivation_recorder=drec)

# After
with tracmass.ParticleRecorder("traj.parquet") as rec:
    tracmass.time_dependent(..., recorder=rec)

Filter events when analysing:

df = pd.read_parquet("traj.parquet")
snapshots    = df[df["event_type"] == "snapshot"]
deactivations = df[df["event_type"] == "deactivated"]

Fixed — GLORYS k-subset off-by-one: settings["subset"]["k1"] now 1-based#

gcm.glorys.fluxes previously passed settings["subset"]["k1"] directly as the start of a Python slice, treating it as a 0-based index. This was inconsistent with the rest of the TRACMASS convention (1-based layer numbers) and with seed.ijk, which also uses 1-based k indices.

Concretely, settings["subset"]["k1"] = 1, k2 = 2 produced slice(1, 2) → one depth level (the second global GLORYS level) → km = 1. Meanwhile, seed.ijk(..., k1=1) placed particles at z = 1.5 → ka = 2, which immediately exceeded km = 1, deactivating every particle in the very first sub-cycle with REASON_BOUNDS. The failure was invisible before the per-window tally was added; it manifested as “515 100 → 0 active, subcycles:1”.

The fix converts the 1-based k1 to a 0-based slice start: k_start = k1 - 1 (or None when unset), so k1=1, k2=2 → slice(0, 2) → two depth levels → km = 2, and ka = 2 ≤ km = 2 is valid.

Added — per-particle deactivation reasons and per-window tally#

ParticleState now carries a reason field (int16, shape (N,)) that records why each particle was permanently deactivated. The reason codes are exported as module-level constants from tracmass:

Constant

Value

Meaning

REASON_ACTIVE

0

Still running (or paused at window boundary)

REASON_TIME

1

Time limit exceeded (tt − t0 > timax)

REASON_LAND

2

Entered a land cell (kmt == 0)

REASON_BOUNDS

3

Left model domain (j or k out of range)

REASON_PROGRESS

4

Zero float32 progress (ds ≈ 0 / UNDEF — Fortran infLoop/dsCrossError)

REASON_SLOW

5

Progressive deactivation (too slow, stuck oscillator)

REASON_ITERLIMIT

6

Sub-cycle iteration limit reached

REASON_DXYZ

7

Zero or negative cell volume (dxyz ≤ 0 — Fortran dxyzError)

Reasons are set inside JIT-compiled kernels (land/bounds/dxyz) and in _integrate_window (time/progress/slow/iter-limit), so there is no performance overhead beyond a few extra jnp.where calls per sub-cycle.

time_dependent now prints a per-window line showing the number of particles that entered and exited each window, how many stopped and why:

Window 1/50: 10000 → 9990 active  stopped:10 (land:8, out_of_bounds:2), subcycles:6, wall:1.1s

ParticleState.__repr__ and write_trajectories_parquet also reflect the new field.

Fortran comparison: the implementation now matches all TRACMASS v7 errorCheck conditions — infLoopError (→ REASON_PROGRESS + REASON_ITERLIMIT), dxyzError (→ REASON_DXYZ), boundError (→ REASON_BOUNDS), landError (→ REASON_LAND), dsCrossError (→ REASON_PROGRESS). The coordboxError debug checks (errCodes 5–7) and user-defined kill_zones are not implemented.

Fixed — periodic east-wall wrap: x1 > imt → x1 >= imt#

In all three advection kernels (piecewise in full_grid.py, analytical in full_grid.py, and piecewise_stencil in stencil.py), the east-wall cyclic wrap used a strict > comparison:

# Before (buggy)
x1 = jnp.where((iperio != 0) & (x1 > imt), x1 - imt, x1)

When a particle reached the eastern wall at exactly x1 == float(imt), the wrap was skipped. The index ib was already wrapped to 1 (via the separate ib > imt guard), but x1 stayed at float(imt). On the next sub-cycle, the kernel saw ia=1 but x=imt, producing negative crossing-time estimates → UNDEF crossing times → REASON_PROGRESS deactivation.

The fix changes all three sites to x1 >= imt, matching the Fortran reference (x1 >= DBLE(IMT) in mod_pos_tstep.F90).

A regression test test_periodic_east_wrap was added to tests/test_integration.py: it seeds a single particle 0.2 cells from the eastern boundary of an 8-cell-wide periodic grid with uniform eastward flux, runs 30 sub-cycles, and asserts that the particle does not receive REASON_BOUNDS or REASON_PROGRESS and its final position satisfies 0.0 ≤ x < imt and 1 ≤ ia ≤ imt.

Fixed — RuntimeWarning: overflow encountered in cast from GLORYS float64 data#

LazyFluxFields.compute() and compute_patch() now explicitly clip float64 numpy arrays to float32 range before passing them to jnp.stack. GLORYS (and similar reanalysis products) provide velocity transports as float64; derived quantities or dataset-specific fill values can occasionally exceed the float32 maximum (~3.4×10³⁸), causing JAX’s implicit float64→float32 cast to emit a RuntimeWarning: overflow encountered in cast. The clip is a no-op for adapters that already produce float32 data (e.g. the theoretical GCM), so there is no performance impact on those paths.

Changed — fluxes_fn is now optional in time_dependent and variants#

LazyFluxFields now carries a fluxes_fn attribute set by :func:new_flux and propagated by :func:update_flux. All three multi-window integration functions (time_dependent, time_dependent_patched, time_dependent_streaming) fall back to fluxes.fluxes_fn when no explicit fluxes_fn argument is supplied.

This eliminates the redundancy where the same GCM callable had to be passed twice — once to new_flux and again to time_dependent. The explicit form still works for custom cases where the update callable differs from the one used to create the flux buffer.

Before:

grid, lazy = tracmass.new_flux(gcm.glorys.fluxes)
result = tracmass.time_dependent(state, grid, lazy, gcm.glorys.fluxes, n_windows=50)

After:

grid, lazy = tracmass.new_flux(gcm.glorys.fluxes)
result = tracmass.time_dependent(state, grid, lazy, n_windows=50)

Fixed — corner-trap oscillation and convergence-zone subcycle inflation#

Two complementary fixes reduce excessive sub-cycle counts (1500–9000+) that occur when particles enter convergence zones in the piecewise advection scheme.

EPS wall-clamping in _subcycle_one_particle (advection_scheme/full_grid.py): When _calc_pos_1d returns a continuous position (x1c, y1c, z1c) that is exactly equal to a cell-wall integer value — which happens when ds is too small to move the particle in float32 — the position is now nudged by EPS_WALL = 1e-5 off the wall. Without this, the update_bounce logic fires again on the next sub-cycle, locking the particle at the wall in an infinite near-zero-progress loop. This mirrors the clamping already present in _calc_pos_analytical_1d for the analytical scheme.

pace_factor progressive deactivation (integrate/full_grid.py, integrate/streaming.py): Adds a pace_factor keyword argument to time_dependent, time_dependent_patched, and time_dependent_streaming (default 5). This parameter is forwarded to _integrate_window, where it controls the progressive-deactivation denominator: _pace_denom = (max_sub_iter_per_window - 2) // pace_factor. Particles that would need more than max_sub_iter_per_window / pace_factor sub-cycles to complete the window (convergence-zone oscillators) are detected at step ~3–10 and deactivated early. Particles completing in fewer sub-cycles are never affected. With max_sub_iter_per_window=10000 and pace_factor=5, the soft cap is 2000 sub-cycles. In a 20-window, 10 000-particle theoretical test, window 12 dropped from 9034 → 1007 sub-cycles and overall run time fell ~22%. Use pace_factor=1 to disable early deactivation in diagnostic runs.

EPS wall-clamping also added to _subcycle_one_particle_stencil (advection_scheme/stencil.py): the same EPS_WALL = 1e-5 nudge applied in full_grid.py is now also applied in the stencil (streaming) path, ensuring consistent behaviour across all advection scheme variants.

Fixed — bounce-chain oscillation from per-particle intrpg#

Switched piecewise.step_all to use the global time-interpolation weight (intrpg, computed from mean particle ts in _integrate_window.body_fn) for all flux reads inside the per-particle kernel, instead of the per-particle weight mod(ts_i, 1.0).

The per-particle weight changes by Δintrpg = Δts between sub-cycles. In GLORYS (and similar monthly reanalysis), the flux can swing by orders of magnitude between slot 0 and slot 1, so even Δintrpg ≈ 1e-4 (one fast cell crossing) can flip the interpolated flux sign at a wall. The update_bounce logic then fires in alternating directions on consecutive sub-cycles — east bounce in step N, west bounce in step N+1 — producing a “bounce-chain” oscillation with thousands of sub-cycles per window in eddy-active seasons.

The Fortran TRACMASS reference implementation uses global (mean-ts) intrpg; this change restores that behaviour. intrpg_pp = mod(ts_i, 1.0) is still used for the per-particle layer-thickness interpolation (dxyz), which affects only the dsmin_arr time-step cap rather than the crossing direction.

Fixed — oscillating particles in piecewise scheme (update_bounce)#

Added the update_bounce logic to _subcycle_one_particle in advection_scheme/full_grid.py, matching mod_pos_tstep.F90::update_bounce() in the Fortran reference implementation.

When a particle sits exactly on a cell wall (e.g. x0 == float(ia)) and the time-interpolated flux at that wall pushes the particle further into the adjacent cell, the cell indices (ia/iam, ja, ka) are reassigned to the adjacent cell before the crossing-time calculation. Without this step, a particle can oscillate indefinitely between two cells whenever the flux changes sign across a window boundary, causing thousands of sub-cycles per window and making runs with per-particle intrpg infeasibly slow.

Implemented for all three spatial directions (x, y, z); the z-direction branch is compiled away when no_wflux=True.

The implementation uses a two-phase read: phase 1 reads fluxes at the pre-bounce cell indices to determine bounce direction; phase 2 re-reads fluxes at the post-bounce indices for the crossing-time calculation. Using pre-bounce fluxes with post-bounce indices (a single-read shortcut) gives incorrect crossing times that cause full-length (dsmin) steps on every sub-cycle, making ts overshoot the window boundary and permanently deactivating particles via the timax exceeded check. The two-phase approach avoids this at the cost of doubling the flux reads for bouncing particles (12 vs. 6), which is acceptable because only a small fraction of particles bounce per sub-cycle.

Boundary guards are applied to all bounce conditions: south bounce requires ja > 1, north bounce requires ja < jmt, down bounce requires ka > 1, up bounce requires ka < km, and east bounce on a non-periodic grid requires ia < imt. These prevent spurious bounces from wrapped Python indices at the domain boundaries. The south-wall flux read uses max(ja-2, 0) with an explicit zero-out when ja == 1 to avoid the same wrap-around in the crossing-time computation.

Fixed — _integrate_window uses Python loop instead of lax.while_loop#

Replaced lax.while_loop with a plain Python for loop in _integrate_window. Each sub-cycle is now a single JIT-compiled step_fn dispatch rather than part of a monolithic XLA while-loop kernel.

The previous design caused two problems:

  1. Uninterruptible hangs. XLA compiles the full while-loop (up to max_sub_iter_per_window = 100 000 iterations) as a single kernel. Once it started it could not be stopped with Ctrl+C. Particles in high-flux regions legitimately need many sub-cycles (each wall crossing advances ts by only a fraction of the window width), so windows with 100–200 sub-cycles over 2 M particles became stuck for tens of minutes.

  2. Occasional runaway. Some windows randomly hit the 100 k limit, producing wall times orders of magnitude longer than typical windows.

With the Python loop each sub-cycle returns to the interpreter. Ctrl+C works between sub-cycles, the per-window cost is proportional to the actual number of sub-cycles needed, and the loop can exit early as soon as all particles finish their window. Per-sub-cycle overhead (two np.asarray syncs) is negligible on CPU.

Fixed — lax.while_loop recompilation when flux data changes between windows#

_integrate_window passed the FluxFields to the loop body via a functools.partial closure. JAX treats closed-over arrays as XLA constants, so every window with new flux data triggered a full recompilation of the while loop body. This caused periodic ~4-window-interval slowdowns as JAX’s JIT cache (default size 4) cycled through evictions, and eventually a hang once compilation time grew too large.

Fixed by threading fluxes through the lax.while_loop carry instead of closing over it. The compiled XLA program is now keyed on shape/dtype only and reused for all windows regardless of flux values. Removed the now-unused functools.partial imports from both integration modules.

Fixed — gcm.theoretical.fluxes JAX recompilation on each window#

The time-evolution branch in theoretical.fluxes used jax.vmap over jnp.arange(n_days), whose shape grows by 1 every window. JAX recompiles vmapped functions when input shapes change, so each successive update_flux call triggered a full recompile — causing wall time to blow up roughly linearly after window 10.

Replaced with a single fixed-shape call using jax.random.fold_in(key, n_days) (amplitude scaled by sqrt(n_days) to preserve random-walk variance growth). No variable-shape JAX array is created, so the compiled graph is reused for all windows.

Fixed — time_dependent_streaming particle reactivation at window boundary#

The early-exit check (if not np.any(active): break) ran before the window-boundary reactivation logic, causing all particles to be treated as permanently dead after the first window in any multi-window run. The fix moves the early-exit to after reactivation, matching the behaviour of time_dependent_patched and time_dependent (full-grid).

Added — tests/test_strategy_comparison.py#

New test module that runs both time_dependent (full-grid JAX while_loop) and time_dependent_streaming (Python for-loop with per-particle stencil) on the identical Fabboni initial state for all 99 windows, then compares them against each other and against the Fortran TRACMASS CSV reference.

  • test_both_strategies_complete_full_run — asserts both strategies record snapshots for all N_WINDOWS windows (snapshot count is the correct completion metric; the final active flag differs by design between the two strategies).

  • test_streaming_x/y_matches_fortran — streaming positions at windows 10, 20, 50, 80 agree with the Fortran CSV to within 1 grid cell.

  • test_strategies_agree_at_window — full-grid and streaming positions agree to within 0.05 cells at windows 10, 20, 50, 80, 98.

  • test_divergence_profile — computes and prints a per-window divergence table (run with pytest -s to see it); asserts max strategy divergence stays below 0.05 cells. On the Fabboni single-particle case the two strategies produce bit-identical trajectories (divergence = 0.00000 at every window).


Changed — trajectory I/O switched from NetCDF to Parquet; TrajectoryRecorder now streams directly to disk#

particle.py no longer depends on xarray or NetCDF for trajectory output. Both the single-snapshot writer and the multi-window recorder now use pandas / PyArrow Parquet.

write_trajectories_parquet(filepath, state, nff=1)#

Replaces write_trajectories_nc. Writes all fields of a ParticleState (all particles, active and inactive) to a single .parquet file via pandas.DataFrame.to_parquet(). The tt and ts columns are stored as nff × value so backward trajectories have negative times, consistent with the previous NetCDF convention.

TrajectoryRecorder(filepath)#

The recorder now requires a filepath argument and streams each snapshot as a PyArrow row group directly to the Parquet file via pyarrow.parquet.ParquetWriter. The in-memory self._snapshots list is eliminated — memory usage is constant regardless of the number of windows.

The to_netcdf() method is removed. To finalise the file, call recorder.close() or use the recorder as a context manager:

with tracmass.TrajectoryRecorder("trajectories.parquet") as recorder:
    result = tracmass.time_dependent(..., recorder=recorder)

Reading snapshots back:

import pandas as pd
df = pd.read_parquet("trajectories.parquet")
snap10 = df[df["snapshot"] == 10]

Schema written by record(): snapshot (int32), label (str), ntrac (int32), x, y, z, tt (float64). Only active particles are written per snapshot.

Updated call sites#

  • tests/conftest.py — simulation_result fixture now accepts tmp_path_factory, creates the recorder with a session-scoped temp path, and returns (final_state, parquet_path) instead of (final_state, recorder).

  • tests/test_fortran_comparison.py — intermediate-snapshot tests read the Parquet file with pd.read_parquet and filter by snapshot and ntrac columns instead of indexing recorder._snapshots.

  • tests/test_speed.py — replaced xarray.open_dataset with pd.read_parquet; fixture renamed trajectory_nc → trajectory_parquet.


[Unreleased] — branch feature/local-stencil#

Changed — gcm.theoretical layer thickness and flux amplitudes rescaled to realistic speeds#

Added DZT = 10.0 m constant (matching gcm.fabboni) and a private _FLUX_AMP = 0.1 × DYU × DZT = 250.0 m³/s that targets a max horizontal speed of ~0.1 m/s. All hardcoded flux magnitudes were replaced:

  • Double-gyre amplitude: 0.05 → _FLUX_AMP (250 m³/s)

  • Initial noise: 0.005 → 0.1 × _FLUX_AMP (25 m³/s, 10 % of mean)

  • Day-to-day perturbation: 0.002 → 0.04 × _FLUX_AMP (10 m³/s, 4 % per day)

  • Vertical flux: 0.001 → 1e-4 × DXDY (6.25 m³/s ≈ 1×10⁻⁴ m/s)

  • dzt: 1.0 → DZT (10 m)

Result: max horizontal speed ~0.14 m/s, mean ~0.07 m/s (previously ~0.3 mm/s).


Changed — gcm.theoretical cell size set to 250 m × 250 m (matching gcm.fabboni)#

Added module-level constants DXV = DYU = 250.0 (m) and DXDY = DXV * DYU = 62500.0 (m²) to gcm/theoretical.py, consistent with gcm.fabboni. The fluxes() function now uses DXDY instead of the previous hard-coded 100.0.


Fixed — gcm.theoretical boundary-leaking double-gyre (caused mass deactivation by window 3)#

The synthetic double-gyre had sin/cos swapped relative to the correct streamfunction ψ = A·sin(πy/H)·sin(2πx/L):

  • uflux used −sin(πy)·cos(2πx) → should be −cos(πy)·sin(2πx)

  • vflux used +cos(πy)·sin(2πx) → should be +sin(πy)·cos(2πx)

The wrong cos(πy) in vflux evaluates to ±1 at y=0 and y=H, producing non-zero meridional flux through the north and south walls. Particles drifted across those boundaries into land cells (kmt[:,0]=0, kmt[:,-1]=0) and were permanently deactivated, emptying the domain within 2–3 windows.

The correct sin(πy) is zero at y=0 and y=H, giving a properly closed domain. uflux was updated simultaneously to keep the flow consistent with the streamfunction.


Fixed — seed.ijk coordinate system bug (caused mass particle deactivation)#

seed.ijk placed particles at x0 = ivec + 0.5 with ia = floor(x0) = ivec. In TRACMASS, cell ia spans the continuous range [ia-1, ia], so the centre of cell ia is at ia - 0.5, not ia + 0.5. With the wrong offset every particle started outside its assigned cell (past the east wall at x = ia). The crossing-time solver found no valid wall crossings, fell back to dsmin, and the subsequent position update placed particles out-of-domain, causing immediate deactivation of nearly all particles.

Fix: changed offset to -0.5 and added the missing +1 to the cell index computation (ia = floor(x0) + 1), consistent with gcm.theoretical.initial_state.


Fixed — JAX correctness and code quality#

  • Bug fix: _subcycle_one_particle (piecewise full-grid scheme) accessed wflux[ia-1, ja-1, ka-2, ...] without clamping when ka=1, causing JAX’s -1 index wrap-around to read the wrong wflux level. Fixed by using jnp.maximum(ka - 2, 0), consistent with the analytical scheme’s jnp.clip(ka - 2, 0, km) and extract_stencil’s np.maximum(ka - 2, 0).

  • Code clarity: Cached jnp.floor() results in _subcycle_one_particle and _subcycle_one_particle_stencil to avoid double computation of the same floor value (matching the analytical scheme’s existing pattern). XLA CSE already deduplicates these, but caching makes the intent explicit.

  • Style: Moved import dask from inside the extract_stencil function body to module level in gcm/flux_grid.py.


Two new memory-efficient integration modes for the GLORYS12 global grid, motivated by the memory cost of computing the full (4320, 2041, 50) GLORYS velocity arrays for every GCM window.


New module — tracmass.gcm.fabboni#

Extracted the Fabboni analytical test-case from tests/conftest.py into a first-class GCM adapter at src/tracmass/gcm/fabboni.py, mirroring the interface of :mod:tracmass.gcm.theoretical.

  • gcm.fabboni.fluxes(dtm, reverse=False) — spatially-uniform, analytically-defined flux fields matching the Fortran projects/Theoretical run. Step index derived from dtm - DTM0 in units of TIMESTEP (1 h).

  • gcm.fabboni.initial_state() — single-particle state at x=150.0, y=99.5, z=0.5 (isec=1, idir=1 seeding).

  • Module-level constants: IMT, JMT, KM, DXV, DYU, DZT, DXDY, DTM0, TIMESTEP, UG, U0, FCOR, GAMMA, GAMMAG.

tests/conftest.py is simplified: all physics constants and helper functions are removed; the simulation_result fixture now uses gcm.fabboni.fluxes and gcm.fabboni.initial_state() directly. gcm.fabboni is exported from tracmass.gcm.__all__.


Refactoring — new_grid/update_grid renamed to new_flux/update_flux; dtm moved to LazyFluxFields#

dtm (the timestamp of the past time slot) describes the flux buffer, not the static grid geometry. Moving it to LazyFluxFields makes GridFields truly immutable for the lifetime of a run, and eliminates the need to return an updated GridFields from update_flux.

  • GridFields no longer carries dtm. It now contains only geometry (kmt, dxdy) and fixed metadata (imt, jmt, km, timestep, reverse, iperio, jperio).

  • LazyFluxFields gains a dtm: pd.Timestamp | None field. Set by new_flux and update_flux; None when constructed via lower-level helpers like swap_fields.

  • new_grid → new_flux — same signature and return type; dtm is now stored on the returned LazyFluxFields instead of GridFields.

  • update_grid → update_flux — now returns just LazyFluxFields (not a tuple). Call sites change from grid, fluxes = gcm.update_grid(grid, fluxes, fn) to fluxes = gcm.update_flux(grid, fluxes, fn).

All public exports in tracmass and tracmass.gcm are updated.


New features#

Local-stencil patch-based advection (time_dependent_patched)#

integrate.time_dependent_patched() loads only the rectangular flux patch that covers the current particle bounding box, plus a configurable buffer of extra cells on each side. This yields a ~4–5× speedup when particles are clustered in a small region of the global domain.

  • LocalFluxFields (gcm/flux_grid.py) — new JAX pytree dataclass carrying flux, geometry (kmt, dxdy), and layer-thickness arrays for a spatial sub-patch. Unlike FluxFields, kmt and dxdy are dynamic arguments rather than XLA constants, so a single compiled kernel can be reused across windows as the patch content changes.

  • LazyFluxFields.compute_patch(i0, i1, j0, j1, grid) (gcm/flux_grid.py) — computes only the Zarr chunks intersecting the requested slice, returning a LocalFluxFields.

  • advection_scheme.piecewise_local(imt_p, jmt_p, km, iperio) (advection_scheme.py) — kernel factory for patch-based advection. Closes over the fixed patch dimensions so the JAX graph is compiled once for the initial patch shape and reused for all subsequent windows.

  • integrate._active_bbox(state, buffer, imt_grid, jmt_grid) — helper that returns the 0-based (i0, i1, j0, j1) bounding box of active particles plus a buffer, clamped to the full-grid extent.

  • integrate._shift_state(state, di, dj) — helper that converts particle grid indices and fractional positions between global and patch-local coordinates.

Limitation: patch shape is fixed at initialisation. If particles disperse beyond the initial bbox + buffer, those particles reach the patch boundary and are deactivated. Not suitable for widely scattered particles or very long runs.


Streaming per-particle stencil advection (time_dependent_streaming)#

integrate.time_dependent_streaming() replaces the JAX while_loop with a Python for-loop. Before each sub-cycle, the 6 flux face values needed at each particle’s current cell are extracted from the lazy Dask arrays via a single batched dask.compute() call. No rectangular patch is allocated.

This approach is designed for:

  • Widely scattered particles — no bounding-box blowup.

  • Very long runs — 30,000 days ≈ 1,000 monthly windows; the fixed-shape patch constraint never applies.

  • FluxStencil (gcm/flux_grid.py) — new JAX pytree dataclass holding the 6 flux face values (u_east, u_west, v_north, v_south, w_top, w_bot), layer thickness (dzt), sea-surface height (zstot), bottom-level index (kmt), and cell area (dxdy) for all N particles, with a trailing slot dimension of size 2 (past / present).

  • extract_stencil(lazy, static, state, no_wflux=False) (gcm/flux_grid.py) — vectorised point indexing via da.data[i, j, k] fancy indexing (works for both Dask-backed GLORYS and plain-NumPy theoretical grids). All particle indices are clamped to valid grid bounds before use so that deactivated particles with out-of-domain indices do not raise IndexError. All reads are batched into a single dask.compute() call to minimise Dask scheduler overhead.

  • advection_scheme._subcycle_one_particle_stencil() (advection_scheme.py) — per-particle sub-cycle kernel that reads flux values from a FluxStencil (after jax.vmap peels the N dimension) and computes dxyz from per-particle geometry. Returns dt in seconds (physical time elapsed) rather than normalised ds.

  • advection_scheme.piecewise_stencil(grid, no_wflux) (advection_scheme.py) — factory that closes over grid (for kmt land checks on the destination cell) and returns step_all(stencil, state, intrpg, dsmin_s) -> ParticleState.

  • Window-boundary particle reactivation — particles that were paused at the end of the previous window (deactivated because ts >= ts_window_end but still within timax_seconds) are correctly reactivated at the start of each new window.

Performance characteristics: total Python iterations ≈ n_windows × (subcycles per window). For monthly GLORYS windows with typical current speeds, each window requires ~10–20 sub-cycles, giving ~10,000–20,000 iterations for a 30,000-day run. Each iteration performs one Dask scheduler call (loads the Zarr chunks containing the requested cells) and one small JAX JIT dispatch.

Numerical correctness: verified against the baseline time_dependent implementation. Produces an identical active-particle mask and near-identical trajectories; maximum position difference < 0.01 cells, arising from float64 rounding in intrpg between the Python-loop and JAX-while_loop code paths.


Bug fixes#

  • time_dependent particle reactivation (integrate/full_grid.py) — the window-boundary reactivation mask (paused & still_valid) was computed but never applied, so in multi-window runs all particles became permanently inactive after the first window. Fixed by applying state._replace(active=...) before advancing to the next window.

  • time_dependent_patched particle reactivation (integrate/streaming.py) — the patched variant had no reactivation logic at all. Added the same numpy-based reactivation as time_dependent_streaming.


Cleanup#

  • Removed stale example functions (test_steady_state, example_2–5) from tracmass/__init__.py; they referenced non-existent APIs from an earlier prototype.

  • Added missing swap_fields import to tracmass/__init__.py.


API additions#

Symbol

Module

Description

LocalFluxFields

tracmass.gcm

Patch-local flux + geometry pytree

FluxStencil

tracmass.gcm

Per-particle point-extracted flux pytree

extract_stencil

tracmass.gcm

Batched Dask/NumPy point indexer

time_dependent_patched

tracmass

Patch-based multi-window integration

time_dependent_streaming

tracmass

Streaming per-particle multi-window integration

All new symbols are exported from tracmass.__all__ and tracmass.gcm.__all__.


Refactoring — class and parameter renaming#

Classes and parameter names have been made more consistent across the codebase:

Old name

New name

Location

FluxSlice

FluxFields

tracmass.gcm

LazyFluxSlice

LazyFluxFields

tracmass.gcm

LocalFluxSlice

LocalFluxFields

tracmass.gcm

ParticleStencil

FluxStencil

tracmass.gcm

The first positional parameter of _integrate_window, steady_state, time_dependent, and update_grid was renamed from static to grid to better reflect that it is a GridFields object.

All public imports from tracmass and tracmass.gcm are updated accordingly.


Refactoring — subpackage split#

advection_scheme.py and integrate.py have been converted from single flat modules into subpackages. All public imports are unchanged.

advection_scheme/ subpackage#

File

Contents

advection_scheme/__init__.py

Re-exports piecewise, analytical, piecewise_local, piecewise_stencil

advection_scheme/full_grid.py

piecewise(), analytical(), and all shared private helpers (_cross_time_1d, _calc_pos_1d, _subcycle_one_particle, …)

advection_scheme/stencil.py

piecewise_local(), piecewise_stencil(), and the stencil sub-cycle kernel

stencil.py imports the shared helpers from full_grid.py; no circular dependency.

integrate/ subpackage#

File

Contents

integrate/__init__.py

Re-exports steady_state, time_dependent, time_dependent_patched, time_dependent_streaming

integrate/full_grid.py

_ensure_computed, _integrate_window (lax.while_loop), steady_state, time_dependent

integrate/streaming.py

_active_bbox, _shift_state, time_dependent_patched, time_dependent_streaming

streaming.py imports _integrate_window from full_grid.py.

Backward compatibility: from tracmass.integrate import time_dependent, from tracmass import advection_scheme, and all other existing import paths continue to work without modification.


Performance test suite (tests/test_performance.py)#

New pytest module that benchmarks the three integration strategies across multiple particle counts, with JIT warm-up excluded from measurements.

pixi run -e test pytest tests/test_performance.py -v -s --tb=short

Test group

Particle counts

Windows

Theoretical (always runs)

1 000, 10 000, 100 000

3

GLORYS12 (skipped if data unavailable)

1 000, 5 000

2

Each test calls the integration function twice: once with n_windows=1, max_sub_iter_per_window=2 to prime the JAX JIT cache, then a second timed run. A formatted summary table is printed by test_zz_print_summary after all tests complete.


Benchmark script (test_glorys.py)#

Added three benchmark functions exercisable without a Python test runner:

pixi run python -c "import test_glorys; test_glorys.compare()"

Function

Strategy

Notes

box()

Full-grid baseline

Computes entire GLORYS grid

box_patched(patch_buffer)

Local-stencil patch

patch_buffer=100 by default

box_streaming()

Per-particle stencil

Scales to any particle distribution

compare() runs all three and prints a side-by-side timing summary with speedup ratios.