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.computeandoutput_filename. The package overview andtracmass.seagrid.debuggot pages. Private helpers (_name) stay out.Land mask and distance to land for every seagrid model.
seagrid.glorys,seagrid.ecco2andseagrid.scrippsgoghavelandmask(k=1)(bool,Truewherekmt < k) anddistance_to_land(k=1)(great-circle km from each cell centre to the nearest land cell centre, 0 on land). Both returnxr.DataArrays on(y, x)with the grid’slon/lat. The shared code is in the newseagrid/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.scipyis 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) taketraj_max_time, read from thetraj_max_timesetting when omitted. It is apd.Timedeltastring ("30D","12h") or calendar months/quarters/years ("1M","6M","1Q","1Y"), which are measured from the run’s start. It sets the existingtimax_secondslimit, counted from each particle’s release.timax_secondsnow defaults to “not given”; passing both raisesValueError. Run metadata getstraj_max_time/timax_seconds. New helpersrunsetup.traj_max_secondsandrunsetup.resolve_timax.Repeat seeding:
seed_freq.time_dependent,time_dependent_patchedandtime_dependent_streamingtakeseed_freq(a pandas frequency, e.g."30D"or"MS"). When the argument is omitted it is read fromseed_freqin the[env.seed]table. Everyseed_freqfromdtm_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, sot0identifies the release andtimax_secondscounts per particle. Particles keep their array index, so existing outputids are unchanged. The window loops no longer stop early while a release is still to come. Run metadata getsseed_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_freqandrunsetup.seed_metadata.seed.from_settingsignores theseed_freqkey.upwelling.movie: movie of thegog_upwellingrun (upwelling.py). Reads the newesttrm_gog_upwellingfile one window at a time (the file holds ~1.4 billion rows, too many to load at once) and scatters the particles on the projmapgogmap, one frame per window.mode="eq"shows particles withka == 12,mode="le"those withka <= 12(the level is set withka=). Writes.mp4when ffmpeg is installed, otherwise.gif.seed.latlon_list: seed one particle per lat/lon point. Takes a DataFrame withlat/loncolumns (optionalkcolumn for the 1-based level, elsekpos), a.csv/.parquetfile, orlat=/lon=lists. Each particle starts at the point’s exact fractional grid position (trajs.coords.lonlat_to_xyon the adapter’slonlat_grid(grid), sosubsetgrids work); points on land move to the nearest ocean cell, and rows below the local bottom (k > kmt) are dropped with a warning. Registered inseed.SETTINGS_APPROACHESasapproach = "latlon_list", withseed_file = ...or inlinelat = [...]/lon = [...]in[env.seed].gcm.scrippsgog/seagrid.scrippsgog: Scripps Gulf of Guinea MITgcm adapter.seagrid.scrippsgog.retrieve(dtm)downloads the dailyGG_i001_YYYY-MM-DD_1day_UVEL_VVEL_top200m.ncfiles from sose.ucsd.edu/GoG/DATA/PRELIM over HTTPS (streamed to a.partfile, 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()buildskmt/dztfrom theDepthfile 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 arelons/lats.gcm.scrippsgog.fluxes(dtm)moves MITgcm’s west/south-faceUVEL/VVELonto the east/north faces the kernels read (the last column’s east face is zero) and builds transports from the files’ owndyG/dxG/drF/hFacW/hFacS. Horizontal-only (noWVELin the files); the regional domain is non-periodic, and its open western/southern boundaries are closed in the tracking. Includeslonlat_grid,lonlat_to_xy,xy_to_lonlat; tests intests/test_scrippsgog.py.Ramp-up saves:
time_dependentandtime_dependent_patchedtakerampup_windows/rampup_substeps(or the settings of the same name; defaults0= off and5). During the firstrampup_windowswindows, positions are also savedrampup_substeps - 1times inside each window, as rows with the newevent_type="rampup"(window= the window they fall in,dtm= their exact time), written by the newParticleRecorder.record_rampup. The run metadata getsrampup_windows,rampup_substepsandsave_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_windownow 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_streamingandscheme="analytical"raise if a ramp-up is requested.tracmass.load(rampup=True)keeps the"rampup"rows along with the snapshots and adds theevent_typecolumn. The default still returns snapshots only, so existing analyses see the same rows.tracmass.run(dtm=...):dtmis accepted (via**kwargs) as an alias fordtm_start, matching thedtmname used byLazyFluxFieldsand the adapters. Passing both raises aTypeError; 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 = truein the active settings environment (also read by the integrators), ortracmass.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_fluxtimes. Uses the standard"tracmass"logger atDEBUG. Added to diagnoserun(env="gog")showing no monitor: before the monitor starts,time_dependentloads 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.computewarns 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. onewindowof 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 incolumns; withclean=Truethe snapshot-row filter is ANDed on. Kept separate fromcolumnsbecause 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_metadatain the recorder’stracmass.runParquet metadata —gcm,gcm_product,gcm_subset,env,integrator,direction,dtm_start,dtm_end,timestep_seconds,tracmass_version,created(previously onlygcm). Keys a caller passes inmetadata=win.tracmass.load(lonlat=True)uses it: with nogcmin the file it looks up[env.gcm] nameof the file’senv(not just the active one), and it converts with the grid of the file’sgcm_product, so a file from another GLORYS/ECCO2 product gets the right lon/lat.Seed by lat/lon from settings:
seed.SETTINGS_APPROACHESgains"latlon_box"and"latlon_point", so[env.seed]can useapproach = "latlon_box"withbox = { lat1, lat2, lon1, lon2, k1, k2 }.seed.latlon_boxwas unusable — it readgrid.lat/grid.lon, whichGridFieldsdoes not have. It now takes the cell-centre lon/lat from the new adapter functionlonlat_grid(grid)(gcm.glorys,gcm.ecco2; the window of the full grid atgrid.i_offset/j_offset, so it follows thesubset), seeds the ocean cells whose centre is in the box (levelsk1..k2, skippingk > 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 fromgcm=or[env.gcm] name.latlon_pointis alatlon_boxoflat ± dlat,lon ± dlon.tracmass.loadfinds trajectory files by environment and date (newoutput.py).load(filename=None, *, env=None, dtm=None, reverse=None, columns=None)reads a recorder file as a DataFrame, with its path andtracmass.runmetadata indf.attrs. Instead of a file name, passenv(read viasettings.from_env, so the active env is not switched) and/ordtmto get the file the settings-driven recorder writes for that run; withoutdtmthe newest file of that env; with no arguments the newest.parquetinbase_output_diror its sub-folders. Lookups also check the flat and{env}sub-folder layouts, so files written under anotheruse_sub_foldersetting are found. Replaces a stub that ignored its argument, read a hard-codedtrm_omi.parquetand returned nothing. The recorder’s path builder (runsetup._auto_recorder_path) now uses the sameoutput_folder/output_stem/output_filenamehelpers.lonlat=Trueaddslon/latcolumns via the GCM adapter’sxy_to_lonlat; the adapter comes fromgcm=, the file’s metadata or[env.gcm] name.clean(defaultTrue) keeps only the snapshot rows and the columnsid,x,y,dtm(pluslon,lat); an explicitcolumnsoverrides the column list, andclean=Falsereturns the whole file. For the lon/lat lookup,manage_run_ionow adds the adapter’s name (gcm) to the recorder’s metadata, andParticleRecorder.metadatais a public attribute (was_metadata).gcm.glorys.xy_to_lonlatreads 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_foldersetting: when true, the settings-built trajectory recorder writes into a sub-directory ofbase_output_dirnamed afteroutput_name, or the active environment whenoutput_nameis 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. Defaultfalse(unchanged layout).tracmass.run: one-call, settings-driven entry point (newdriver.py), generalised fromabstemp.run. It activates an optionalenv, builds grid + flux buffer with a singlenew_flux, seeds (aParticleStateor aseed(grid)callable) and calls the integrator. Omitted arguments come from settings:gcmfrom[env.gcm] name,dtm_start/dtm_endfrom the keys of the same name, elsedtm_end = dtm_start + runtime(negativeruntime→ backward run; with no start date, the adapter’s default date is used), andintegrator(time_dependent/time_dependent_patched/time_dependent_streaming, defaulttime_dependent). Remaining keyword arguments go to the integrator.abstemp.runis now a thin wrapper around it.tracmass.runcan seed from settings; newseed.from_settings:stateis now optional. When omitted, particles come from the active[env.seed]table —approach(a key of the newseed.SETTINGS_APPROACHES, currently"ijk"), its arguments in aboxsub-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 intracmass.toml. The dates come first, so the signature isrun(dtm_start=None, dtm_end=None, state=None, *, env=, gcm=, integrator=, **kwargs)andrun("2007-06-01", env="omi")needs no keyword; aParticleStateor seed function passed in a date slot (the earlierrun(state, dtm_start, ...)order) raises aTypeError.seed.pynow importsParticleStatefromtracmass.particle(wastracmass) sodriver.pycan 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/):LazyFluxFieldsgainstimestepandreverse(copied from the grid bynew_flux),dtm_end(thedtm_endpassed tonew_flux) and a derivedn_windowsproperty,ceil(|dtm_end - dtm| / timestep)— the windows left untildtm_end, decreasing by one perupdate_flux,Nonewithout adtm_end.time_dependent,time_dependent_patchedandtime_dependent_streamingnow defaultn_windowstoNone, resolved (newrunsetup.resolve_n_windows) fromfluxes.n_windows, falling back to the old default of1when nodtm_endwas given; an explicitn_windowsstill wins. Sonew_flux(fn, dtm_start=d0, dtm_end=d1)followed bytime_dependent(state, grid, lazy)replaces thenwin = int((d1 - d0) / tstep)boilerplate — and no longer needs the caller to know the adapter’s timestep, which already lives in eachseagridmodule’ssettings["timestep"]. Adtm_endthat falls inside a window now gets that window integrated (the scripts’int()truncated it).update_flux(lazy): step size, direction and thefluxescallable now come from the buffer itself (lazy.timestep/lazy.reverse/lazy.fluxes_fn);update_flux(lazy, fluxes_fn)overrides the callable. The oldupdate_flux(grid, lazy, fluxes_fn)form still works and is needed only for a buffer built by lower-level helpers without atimestep.update_fluxandLazyFluxFields.preloadnow carry all buffer metadata over viadataclasses.replace(preloadused to dropzarr_chunk_xy).trm_*.tomlfiles in the current directory are now loaded as settings (config.py): any file matchingtrm_*.tomlnext to a script is added to Dynaconf’ssettings_files, sorted by filename, after the current directory’stracmass.toml/.secrets.tomland before theSEAVIEW_SETTINGS_FILE_FOR_DYNACONFoverride. Lets a script layer run-specific settings (e.g.trm_omi.tomldefining an[omi]environment) on top of a sharedtracmass.tomlwithout editing it. NewTRM_FILESmodule attribute holds the resolved, sorted list.Rich monitor title shows the active Dynaconf environment (
monitoring.py): theRunMonitor(backend="rich")table title is nowtracmass · time_dependent · env: <env>(e.g.env: omi), read fromsettings.current_envon each redraw, so it’s clear at a glance whichtracmass.tomlsection a run is using.abstemp.closest_unconnected_db(abstemp.py): for every consecutiveregid_pathleg that no trajectory connects (the legsstitched_traj_dbskips; listed by the newunconnected_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). Addstarget,dist,min_dist,rankand aclosestflag set on every position of then_closestbest trajectories;closest_only=Truereturns just those. The per-pair step is exposed asabstemp.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 (fromclosest_traj_db) instead of being skipped, so the stitched path has no missing legs; aconnectedcolumn marks those stand-in rowsFalse.
Fixed#
seagrid.glorys.rechunk_wvelcrashed a run reading the same days. Agog_upwellingrun caught up with a running rechunk and failed withcannot 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_streamingdid not record the time limit as a stop reason. Particles pasttimax_secondswere stopped but kept reason0, so they looked paused rather than stopped for time. They now getREASON_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.computenow keeps the present slot’s arrays (_computed, keyed by the slot’s DataArray and shared acrossupdate_fluxcopies) 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_dataarraytakemask=(defaultTrue, unchanged for other callers).gcm.glorys.fluxesopens withmask=Falseand applies the new publicseagrid.glorys.mask_fillafter the halo trim and subset.wvelZarrs had one 50-level chunk per tile. The download chunk settings nameddepthtonly, andwvelis ondepthw. New downloads use 10 levels ("depthw": 10). The newseagrid.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 withwflux):compute()went from ~10 s to ~3.5 s per window with the first two fixes, before rechunking. Rechunking cuts thewvelsubset 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 checkedj. On a grid without zonal wrap (iperio = 0, i.e. anygcm.subsetin longitude forglorys/ecco2), particles crossing the east edge kept going pastimtwith clamped flux indices, still active. Particles reaching the west or south edge were held there by a closed wall (um = 0atia == 1), also still active. None of them showed up in the stop tally or asdeactivatedrows. All four kernels (piecewise,analytical,piecewise_local,piecewise_stencil) now stop a particle without_of_bounds_hwhen it enters the outer ring of cells on a non-periodic axis (newflux_grid.edge_mask(grid);LocalFluxFieldsgets anedgefield, 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 rowjmtinstead of leaving it). Regression test:tests/test_domain_edge.py.wfluxis 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_patchandpreloadgaininclude_wflux=True; the integrators passinclude_wflux=not no_wflux.GLORYS velocities sat one cell west of their T-cell:
gcm.glorys.fluxestrimmed the velocity Zarrs (x=4322) withslice(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_wequal the meshglamu/glamv/glamtexactly at velocity columnc + 2= mesh columnc. Souflux[i]was the west face of celli(the kernels read it as the east face), andvflux/wfluxbelonged 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 stayingactive(e.g. 1,489 particles on the Ghana coast in theomi_falkorrun). The velocity trim is nowslice(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 intests/test_gcm_subset.pynow 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), andgrid_definition()wrote every zero — landkmt, landdzt— as that fill value (ds.where(ds, fill)replaces falsy values).open_datasetnow returns 0 for masked velocities (NaN fortemp/salt), andgrid_definitionreturns land askmt == 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 influx_grid.py, but other readers of the grid (e.g.lonlat_to_xy’s land check, which testskmt == 0) were not.seed.ijk(..., grid=grid)skipped the wrong cells as land: it indexed the 0-basedgrid.kmtwith the 1-based TRACMASSi/jit 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-rangei2 = imt/j2 = jmtwent unnoticed because JAX clamps out-of-bounds indices. It now testskmt[i - 1, j - 1]. Runs seeded withseed.ijknext to land (e.g.omi_source_faith.py) seed slightly different cells.time_dependent_patchedcrashed on in-memory float32 GCMs (LazyFluxFields.compute_patch): with the theoretical adapter it raisedValueError: assignment destination is read-onlyfrom the in-place NaN/fill cleanup, because float32 patch slices come back as read-only views and — unlikecompute()—compute_patchhad no copy guard for that case. It now copies non-writeable slices first. Dask-backed (GLORYS/ECCO2) data was unaffected, asdask.computereturns fresh arrays.xy_to_lonlatwas half a cell off, andlonlat_to_xysnapped to cell centres (trajs/coords.py, used bygcm.ecco2/gcm.glorys): the TRACMASS convention is that cellispansi <= x < i + 1with its centre – where the grid’slon/latare defined – atx = i + 0.5, whichseed,lonlat_to_xyandabstempall follow.xy_to_lonlatinstead put centreiatx = i, shifting every converted position by half a cell (0.125° on ECCO2) towards +x/+y. It now interpolates between centres atx - 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_xynow returns the true fractional position – the exact inverse ofxy_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 newperiodic=Noneargument (auto- detected). Positions converted with the oldxy_to_lonlat(e.g.abscarib_30_stitched.parquet) should be regenerated.lonlat_to_xyfailed with currentpyinterp(2026.x has nopyinterp.RTree): the nearest-centre search now usesscipy.spatial.cKDTreeon unit-sphere coordinates (great-circle nearest), sopyinterpis no longer needed.RunMonitor’s rich progress bar showed a misleading0/1(never0/0) for a genuinen_windows=0run (monitoring.py): every place that derived a display total or remaining count fromself._n_windowsusedself._n_windows or <fallback>, which conflates an unset (None) total with an explicitly-zero one —0 or 1evaluates to1. A run with zero windows (e.g. theabstemp.runbug above, or any off-by-typo empty date range) would render a permanently frozen0/1bar 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 checkis not Noneinstead of truthiness, so a genuine zero renders as0/0.abstemp.runsilently 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 —nwinthen computed to0andtime_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 static0/1, which looks like a hang even though the call already returned. Fixed topd.Timedelta("180D").new_grid/new_fluxwithdtm_start=Nonebroke afluxescallable that has no default for itsdtmparameter (tests/test_integration.py): a side effect of thedtm_startdefault change above — three test-locallambda dtm, reverse=False: ...fixtures (test_periodic_east_wrap,test_surface_layer_no_vertical_exit,test_northern_boundary_exit_reason_bounds_h) had no default fordtmand don’t actually use it (time-invariant synthetic fields), so_fetch_slot0’s barefn()call now raisedTypeError: missing 1 required positional argument: 'dtm'. All three shipped GCM adapters (theoretical,fabboni,glorys,ecco2) already givedtma default, so only these test fixtures were affected. Fixed by giving themdtm=Nonetoo, matching the real adapter convention.abstemp.runfetchedwvelover OPeNDAP despitemove_vertically = false, and silently dropped its[abscarib_30.gcm] subset(abstemp.py): it calledtracmass.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()’smove_vertically()/gcm_subset()calls read the globaltracmass.settingssingleton, which was still on[default](move_vertically = true, no subset), soinclude_wfluxresolved toTrueand fetched an uncachedwvelfile over the network instead of using only the cacheduvel/vvel. Changed totracmass.settings.setenv(env), which activates the environment in place — matching the working pattern already used inomi_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 explicittimeout, so it fell back to pydap’s internal 120s default — combined with therequests.Session’sHTTPAdapter(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 toPydapDataStore.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 viaabstemp.runhanging attracmass.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_subcounted from wherever each particle’s clock was, not the time left until the next sub-step boundary. With the defaultiter_sub=1it almost never bound, so the last step of every window was a cell crossing ending past the boundary: snapshots sat exactly on cell faces (integerx/y) withttup to hours past the time in theirdtmcolumn, and the next window started from that overshoot. The integrators now pass a per-particle cap (integrate.full_grid._substep_cap) and snaptt/tsonto the boundary when it was the limiting step (_snap_to_substep), matching Fortran TRACMASS; this coverstime_dependent,time_dependent_patched,time_dependent_streamingandsteady_state.piecewise_stencilaccepts a per-particledsmin_s. Newtests/test_window_boundary.py.GLORYS x-subset crashed with a broadcast error:
gcm.glorys.fluxesappliedsubseti1/i2before 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 withoperands could not be broadcast together with shapes (2, 1100, 1096) (2, 1100, 1098). The halos are now trimmed first, soi1/i2index TRACMASS cells and a subset gives exactlyi2 - i1cells. Each subset cell still holds the same data as before (the start offset is unchanged). An x-subset also now setsiperio = 0(inecco2too), so particles leaving the east edge of a regional box no longer wrap around to its west edge.tracmass.runcatches an environment name passed as a date:tracmass.run("omi_falkor")put the env name in thedtm_startslot and failed deep insidenew_fluxwith a pandasDateParseError. A stringdtm_start/dtm_endthat does not parse as a date now raises aValueErrorup front suggestingenv="omi_falkor".
Changed#
seed.ijktakes absolute (full-grid) cell indices by default, the same frame as saved trajectories andlonlat_to_xy. Withgrid=it subtractsgrid.i_offset/j_offset, drops cells outside the subset, and raisesValueErrorif no cell of the box is inside it. The newuse_subset_coordinates=Truekeyword keeps the old subset-relative behaviour; in a settings table setuse_subset_coordinates = truein[env.seed](defaultfalse). Without a subset (offset 0) nothing changes. Seed boxes written for a subset must be shifted byi1/j1or getuse_subset_coordinates = true.Saved trajectories use full-grid (absolute) x/y on a subset grid.
GridFieldsgainsi_offset/j_offset(default 0), filled from the_i_offset/_j_offsetattrs thatglorys/ecco2now put on their datasets (the subset’si1/j1, via the newflux_grid.subset_offset_attrs).ParticleRecorderhas agrid_offset, set bymanage_run_iofrom the grid, which it adds tox/iaandy/jaon write; a non-zero offset is also stored in thetracmass.runmetadata. 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 returnedParticleStateare unchanged (subset-relative). To match,glorys/ecco2xy_to_lonlat/lonlat_to_xynow always use the full grid; subtractgrid.i_offset/grid.j_offsetto seed fromlonlat_to_xyoutput 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_dbdistances are now great-circle km, and wrap across the periodiciseam (abstemp.py):dist/min_distwere 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_lutnow picks the nearest target-cell centre with a KD-tree on unit-sphere coordinates (replacing the planar distance transform) and stores itstlon/tlat; the DuckDB query computes haversine distance from the particle’s lon/lat, derived fromx/yon the regular global grid. It is still one hash join, so performance is unchanged.stitch_closeststill uses grid-index distances.new_grid/new_fluxno longer hard-code a defaultdtm_start(gcm/flux_grid.py):dtm_startnow defaults toNoneinstead of the literal"2007-06-01". When omitted, thefluxescallable is invoked with no arguments at all, so its own default applies — forgcm.glorys.fluxes/gcm.ecco2.fluxesthis falls through totracmass.seagrid.glorys.open_dataset/tracmass.seagrid.ecco2.open_dataset’s owndtmdefault, which is now the single place a default date lives for those two adapters.gcm.glorys.fluxes/gcm.ecco2.fluxesthemselves also changed theirdtmdefault from"2007-06-01"toNoneto match. New_fetch_slot0helper influx_grid.pyfetches slot 0 and reads back the actual resolved timestamp (from the returned dataset’sdtmvariable) for backward/forward direction inference againstdtm_end, so passing onlydtm_end(nodtm_start) still infers direction correctly.gcm.theoretical.fluxes/gcm.fabboni.fluxesare unaffected — they are synthetic adapters with their own literaldtmdefaults and don’t useseagrid.tracmass.njordrenamed totracmass.seagrid: the low-level Zarr/OPeNDAP access layer (glorys.py,ecco2.py,debug.py) now lives atsrc/tracmass/seagrid/. All internal imports (gcm/glorys.py,gcm/ecco2.py,seed.py), docs, and top-level scripts (io_speed.py) updated to match. Historicalnjord-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.pymoved intotests/and converted to pytest tests: both were previously manual, un-asserting scripts at the project root, run by hand viapixi run python -c "import test_glorys; test_glorys.compare()". They are nowtests/test_glorys.py(test_glorys_full_grid/test_glorys_patched/test_glorys_streaming, sharing a module-scopedglorys_gridfixture) andtests/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 markedpytest.mark.slow(skipped unlesspytest --slowis 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 aGridFields(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_fluxgains a second call form,new_flux(grid, fluxes, dtm_start=..., dtm_end=None) -> LazyFluxFields, which builds only the flux buffer against an existing grid (anydtm_endpassed is checked againstgrid.reverse). The originalnew_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. Stalenew_grid(..., reverse=True)references inintegrate/full_grid.pyand thefluxes_to_griddocstring corrected.n_particle_batchesis now settings-driven (config.py,runsetup.py,integrate/full_grid.py,integrate/streaming.py,tests/conftest.py,tests/test_integration.py): then_particle_batchesargument ofsteady_state,time_dependent,time_dependent_patchedandtime_dependent_streamingnow defaults to a sentinel; when not passed it is resolved byrunsetup.resolve_n_particle_batches()from then_particle_batchessetting (newtracmass.config.n_particle_batches(), default1— no batching, clamped to>= 1). An explicit argument still wins. Both full-grid integrators (steady_state,time_dependent) act on it (particle-batchedlax.while_loopvia_integrate_batched); the streaming integrators resolve it for API symmetry but do not batch.steady_stategained then_particle_batcheskeyword in this release — previously it always ran the whole population through a singlelax.while_loop. Thetracmass.toml[default]n_particle_batches = 8key is now actually consumed instead of being ignored unless every call site passed it by hand.tests/conftest.py’s autouse fixture pins it to1for the test run.
Fixed#
tracmass.trajs.xy_to_lonlat/lonlat_to_xyno longer return garbage for transposed or 1-D lon/lat (trajs/coords.py):_unpack_gridblindly readimt, jmt = lon_arr.shape, so passingnp.meshgrid(lon_1d, lat_1d)(defaultindexing="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_lonlathelper now accepts 1-Dlon(imt,)/lat(jmt,)(broadcast withindexing="ij") and auto-detects a(jmt, imt)2-D layout, transposinglon/lat/kmtto TRACMASS(imt, jmt)with awarnings.warn.tracmass.tomlno longer fails to parse: the shipped settings file definedgcmboth as a string (gcm = "theoretical") and as a table ([default.gcm], twice), which is aTOMLDecodeError. 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.tomlnow carries a[env.gcm]table withname(the adapter —"glorys","ecco2","theoretical", …) and an optionalsubsetinline table.tracmass.config.gcm_name()returns the adapter name (a baregcm = "<name>"string is still accepted).tracmass.config.gcm_subset()readsgcm.subsetand returns a normalised 9-key dict (i1, i2, istep, j1, j2, jstep, k1, k2, kstep— eachintorNone; missing keys and"@none"/falsebecomeNone).tracmass.config.apply_gcm_subset( module)copies the non-Noneentries intomodule.settings["subset"], andnew_fluxcalls it for the adapter it resolves (from a string name or the passed callable’s module) before the firstfluxes()call — sosettings.setenv("omi"); new_flux("glorys", ...)now honours[omi.gcm] subsetwith no manualgcm.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 nosettings["subset"]and are left untouched. Newtests/test_gcm_subset.py.move_verticallysetting drivesno_wflux/include_wflux(config.py,runsetup.py,integrate/full_grid.py,integrate/streaming.py,gcm/ecco2.py): a singletracmass.tomlkey now controls vertical advection end-to-end.tracmass.config.move_vertically()reads it (defaultTrue— particles do move vertically, matching the oldno_wflux=False/include_wflux=Truedefaults). Theno_wfluxargument ofsteady_state,time_dependent,time_dependent_patchedandtime_dependent_streamingnow defaults to a sentinel; when not passed it is resolved byrunsetup.resolve_no_wflux()tonot move_vertically(). An explicitno_wflux=True/Falsestill wins.gcm.ecco2.fluxesgainedinclude_wflux=None(the new default), resolved tomove_vertically()so thewvelOPeNDAP dataset is skipped when vertical motion is off. Settingmove_vertically = falsein atracmass.tomlsection (e.g.[tempvel],[omi]) now turnswfluxoff for both the flux load and the advection kernel with no per-call arguments.tests/conftest.py’s autouse fixture also pinsmove_vertically = Truefor the test run. The[omi]section’s previously-deadno_wflux = truekey is nowmove_vertically = false, andomi_source_faith.pyactivates[omi](settings.setenv("omi")) and drops its explicitno_wflux=False— so omi runs now honour the section’s long-standing “no vertical motion” intent instead of the leftoverno_wflux=Falseargument.monitor/recorderare now optional on the multi-window integrators (runsetup.py— new,integrate/full_grid.py,integrate/streaming.py):time_dependent,time_dependent_patchedandtime_dependent_streamingpreviously required the caller to build aRunMonitor/ParticleRecorderand 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_iodecorator (intracmass/runsetup.py) wraps the three integrators somonitorandrecordereach 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 acontextlib.ExitStackand exited automatically — on normal return and on exception — so thewithblock is no longer needed:tracmass.time_dependent_patched(state, grid, flux, n_windows=nwin)
Settings keys read from
tracmass.toml(viatracmass.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 stemtrm_{active-env}), andbase_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 fromgrid.reverseand the date fromfluxes.dtm, so a forward and a backward run in the same settings environment no longer collide. When neitheroutput_namenorauto_recordis set, an omittedrecorderresolves toNone(no file) — the historical default, so existing call sites are unaffected.tests/conftest.pygained an autouse fixture that neutralises these settings during the test run so a straytracmass.tomlin 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 throughnew_flux→steady_state/time_dependentparticle advection.njord/ecco2.pygained awvelentry invardictandgrid_definition()now storesdxv/dyu(zonal/meridional T-cell face widths) alongsidedxdyso the gcm layer can builduflux/vfluxitself.lonlat_to_xy/xy_to_lonlatreuse the shared curvilinear-grid utilities intrajs/coords.py, same asgcm.glorys. Two ECCO2-specific data-quality issues were found and worked around:the APDRC GDS gateway declares
missing_value/_FillValueas-1e23but actually serves exact0.0for land/below-seaflooruvel/vvelcells —njord.ecco2.calc_kmt_dztderiveskmt/dztfrom exact-zero equality rather than NaN/CF decoding.wvelhas its own invalid-cell pattern that is not implied byuvel’s (e.g. under ice shelveswvelis invalid at levels whereuvel/vvelare valid), and is inconsistently NaN vs. the raw undecoded-1e23sentinel —gcm.ecco2._clean_velocityzeroes anyuvel/vvel/wvelvalue withabs() > 1e10(no real ECCO2 current comes close to that magnitude) before computing fluxes.njord/ecco2.pyis adapted fromnjord/glorys.pyfor 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_datasetselects the nearest day, caches it as a local Zarr, and renames dims to the codebase’s(z, y, x)convention.grid_definitionhas no mesh_hgr/mesh_zgr files to draw on, so it computes cell geometry analytically from the lat/lon spacing. The server’s nativelonruns 0.125..359.875 (Prime Meridian at the array seam, antimeridian at the array centre);_prepare_write_jobnow rolls every cached field through_center_on_prime_meridianto -179.875..179.875 before writing, so 0° sits at the array centre and Atlantic/European regions are contiguous fori1:i2subsetting, matching the conventiongcm.glorysalready used. Verified end-to-end with a Gulf Stream subset (lon -80..-50) expressed as plain contiguous indices.njord/ecco2.pyalso 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 lazytimeCF-decode in a worker thread after a scopedwith warnings.catch_warnings():in the calling thread has already restored the filters, andcatch_warningsis documented as not thread-safe, so a scoped filter intermittently leaked through):xr.SerializationWarning(APDRC declarestimeunits 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 — adap2://URL prefix — makes pydap rewrite the scheme tohttpsinternally, bypassing thesession.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 freshgrid_definitionbuild — no warnings reach the caller in any of them.njord/ecco2.pyalso gained_compute_with_retry: ECCO2 daily fetches are 100-200 MB per variable over a slow remote link and occasionally fail mid-transfer with arequests.exceptions.ChunkedEncodingError(IncompleteRead) from the APDRC server — not covered by theHTTPAdapter/Retryconfig in_prepare_write_job, which only retries failures during connection setup, not a drop while streaming an already-established response deep inside pydap’sr.content. This previously killed afor dtm in dtmvec: ecco2.open_dataset(dtm)loop partway through. Both.compute()call sites (retrieve, and the merged uvel+vvel path inopen_dataset) now route through it (3 retries, linear backoff). Fixing this also removed a latent bug:open_dataset’s localimport dask(used only to reachdask.computearound its owndaskboolean parameter, which shadows the module for the whole function) could leak the dask module into thedask=daskkeyword of its own recursive call whenever a fetch actually ran, silently forcingdask=Trueregardless of what the caller passed.
Changed#
seed.mask()now takes a raw(jmt, imt[, kmt])mask — j (lat) first (seed.py): previouslymask()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:jareaching 100 on a grid withjmt=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 indexgrid.kmtin its own(imt, jmt)convention, which is correct sinceGridFields.kmtreally is built that way by thegcm/*.pyadapters. Existing callers ofmask()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_patchedwith GLORYS12 (njord/glorys.py,rechunk_grid.py): three layered fixes eliminated Dask rechunk overhead that dominated patch reads:grid_definition()was opening the grid zarr withchunks="auto", which collapseddyu/dxv/dxdyinto single 3059×4318 global slabs despite the zarr store having 200×200 chunks. Changed tochunks={}so xarray honours the zarr store’s native chunk layout.The grid zarr was rechunked with
dztat(4, 200, 200)while the velocity zarrs use(50, 200, 200). Multiplyinguvel * dzttriggered a Dask rechunk of the full z-column before every patch read.rechunk_grid.pynow usesz=50fordzt, matching the velocity zarr exactly.rechunk_grid.pyadded 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_fluxAPI:dtmrenamed todtm_start,dtm_endadded;reverseremoved (gcm/flux_grid.py): integration direction is now inferred from the relationship betweendtm_startanddtm_endrather than a separatereverseflag. Whendtm_end < dtm_startthe run is backward:grid.timestepstays positive,lazy.dtmdecrements per window, and velocity fluxes (uflux,vflux,wflux) are lazily negated insidenew_flux/update_flux— no wrapper function needed. This fixes the pre-existing bug wherereverse=Truemadetimestepnegative, causingtseas < 0inside_integrate_windowand completely wrong sub-cycle positions. Old code usingnew_flux(fluxes_fn, dtm=..., reverse=True)must be updated tonew_flux(fluxes_fn, dtm_start=..., dtm_end=<earlier timestamp>).reverseparameter removed from GCM adapterfluxes()functions (gcm/glorys.py,gcm/theoretical.py,gcm/fabboni.py): adapters now always return forward (positive) fluxes; negation is handled by the framework. Thenffsign-multiplier infabboni._velocityis also removed (it evaluated velocities at negative model times rather than negating forward velocities, which was physically incorrect).omi_source_faith.py:_glorys_backwardwrapper 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_BOUNDSsplit intoREASON_BOUNDS_H(3) andREASON_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 forREASON_BOUNDS_H(backward-compatible with existing parquet files where reason=3 was always a lateral exit);REASON_BOUNDS_Vis assigned code 8.Surface-layer particles spuriously exiting through
kb=0(REASON_BOUNDS_V) (advection_scheme/full_grid.py): atka=1,jnp.clip(ka-2, 0, km)clips to0, making the shallower-face fluxum_zidentical to the deeper-face fluxuu_z. When the inter-layer flux is upward (upwelling),um_z < 0yields a finitedsd, pushing the particle above the sea surface (kb=0 < 1). Fixed by zeroingum_z(andvm_zin the analytical scheme) whenka==1; the sea surface is treated as a rigid wall. Eliminates ~140 K spuriousREASON_BOUNDSexits per 60 M particle run.njord.glorys.open_datasetprogress bar garbled during parallel downloads (njord/glorys.py): when all three flux files needed downloading,ThreadPoolExecutorran threeretrievecalls concurrently, each owning its own DaskProgressBar. The bars interleaved on stdout, producing garbled output.retrievenow accepts ashow_progressparameter (defaultTrue);open_datasetpassesFalsefor threaded calls and instead prints a single"Downloading N file(s) … Done."line around the parallel fetch.
Changed#
njord.glorys.retrieveis now significantly faster (njord/glorys.py): depth chunking changed from-1(all 50 levels per chunk) to10, giving Dask ~5× more parallel OPeNDAP requests per tile.write_job.compute()now usesscheduler="threads", num_workers=8for concurrent chunk writes.open_dataset(data_var=None)now pre-fetches uvel, vvel, and wvel concurrently viaThreadPoolExecutor(max_workers=3)before opening the cached files, cutting wall time for a full three-variable download by roughly 3×.time_dependent_patchednow 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 withREASON_BOUNDS. The patch now grows each window when the active-particle bounding box exceeds the current patch size, rounded up topatch_granularity(default 100 cells) to limit JAX recompilations.piecewise_localis decorated with@functools.lru_cacheso the compiled kernel is reused whenever the same patch shape recurs. A newpatch_granularityparameter controls rounding.RunMonitorrich backend: cumulative stopped count + end-of-run reason tally (monitoring.py): theStopped ∑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. Theprintandlogbackends likewise gain a cumulative stopped/reason line in their summary output.
Fixed#
Snapshot rows no longer carry stale non-zero
reasoncodes (integrate/full_grid.py,integrate/streaming.py): the window-boundary reactivation logic usedts >= ts_window_endandkmt > 0as guards, but did not checkreason == 0. Particles stopped for REASON_SLOW, REASON_PROGRESS, REASON_ITERLIMIT, or REASON_DXYZ near a window boundary could be erroneously reactivated withactive=Truewhile retaining their non-zero reason, causing confusing snapshot rows whereevent_type="snapshot"butreason != 0. Fix: all three reactivation blocks now requirereason == 0(paused, not permanently stopped).time_dependent_streamingI/O now O(1) per window (integrate/streaming.py,gcm/flux_grid.py): previouslyextract_stencilcalleddask.compute()once per sub-cycle for every cell crossing, reloading the same large Zarr chunks on each call. Fix: newLazyFluxFields.preload()method issues a single batcheddask.compute()for all flux variables at the start of each window, returning a numpy-backedLazyFluxFields. Allextract_stencilcalls within the window then use pure numpy indexing with no further Zarr I/O. The preload time appears in the monitor’st_iocolumn; sub-cyclet_iocontributions drop to near zero.extract_stenciloverflow on float32 cast (gcm/flux_grid.py): fill values (e.g. GLORYS ~9.97e36) were cast tofloat32before thenan_to_zerocleanup ran, causingRuntimeWarning: overflow encountered in cast. Fix: move.astype(np.float32)to after thenan_to_numand 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-Ddxdyarray (spatially varying), or aGridFieldsobject directly. Returns a DataFrame withid,window,dt,dx_m,dy_m,speed.time_dependent_patchedandtime_dependent_streamingnow acceptmonitorandn_particle_batches—monitoris the sameRunMonitorinterface astime_dependent(defaults toRunMonitor(backend="print"); passNoneto suppress output); per-window timing is split intot_io,t_jax,t_pyandn_sub/reason_tallyare reported toon_window_end.n_particle_batchesis accepted for API compatibility but has no effect (no JAXwhile_loopscratch allocation to reduce).tracmass.trajs.xy_to_lonlat(x, y, grid)andtracmass.trajs.lonlat_to_xy(lon, lat, grid)— generic coordinate-conversion functions extracted intotracmass/trajs/coords.pyand re-exported fromtracmass.trajs. Both accept any object with.lon/.latattributes (or a 2/3-tuple), making them usable with any GCM adapter.gcm.glorys.xy_to_lonlatandgcm.glorys.lonlat_to_xyare 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. Usespyinterp.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 withseed.expand_sub. Applies the same spatial subset and x-halo trim asfluxes()::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 Parquetx/ycolumns) 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. Callsjax.clear_caches()(evicts compiled kernels and thekmt/dxdyXLA constants they pin),gc.collect(), andjax.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=...)— optionalmetadatadict 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 forkmt(bottom-depth index) anddzt(layer thickness). When xarray read the file back without applying fill-value masking, two bugs arose:Seeding in land cells:
seed.ijk(grid=grid)filters ocean cells viagrid.kmt.astype(bool). Withkmt = 9.96921e36for land cells, the bool cast isTrue, so particles were seeded inside land.NaN cascade in particle positions:
dzt = 9.96921e36in the advection kernel causesdxyz_rawto overflow float32 → Inf, makingdsmin_arr = dt / Inf = 0,ds = 0, anddt_particles = 0 × Inf = NaN. On the very next sub-stepintrpg_pp = fmod(NaN, 1) = NaN, which propagates todxyz = NaN,ds = NaN, and finallyx1 = NaN. The integer castfloor(NaN).astype(int32)yields 0 →ia = 1, ja = 1. Becausekmt[0, 0]also contained the fill value (≠ 0), the particle was not detected as land and remained “active” forever withts = NaN, so the window loop ran allmax_sub_iter = 10 000sub-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 todztandzstot(was only applied touflux,vflux,wflux)._static_from_ds(): when loadingkmtintoGridFields, replace any value that is NaN or above_FLUX_FILL_THRESHOLD(1e20) with 0 so that land cells are correctly identified askmt = 0.
After the fix, land cells are excluded from seeding and
dxyz_rawis always finite, eliminating the NaN cascade.Western periodic wall crossing — unbounded negative x positions. When a particle in cell
ia=1crossed the western wall on a periodic grid (iperio=1), the kernel setx1 = iam = imtandib = imt. The subsequent cyclic wrapx1 >= imt → x1 - imtconvertedx1 = imttox1 = 0, leaving stateia=imt, x=0— geometrically inconsistent (cellimtspans[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 >= imttox1 > imt(strict), matching the original Fortranmod_pos_tstep.F90exactly.x1 = imtis left unwrapped intentionally — the already-presentr0_adjspecial case in_calc_pos_1d/_calc_pos_analytical_1dhandlesr0 = imt, ii = 1correctly 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 causedds["lat"] = hgr["nav_lon"]. Users with a cachedtrm_glorys12v1_grid_def.ncmust regenerate it::from tracmass.njord import glorys as _g _g.grid_definition(force=True)
no_wflux=True— wflux still uploaded to GPU. Whenno_wflux=Truethewfluxarray (~212 MB for a global GLORYS12 grid) was still present in the computedFluxFieldsand uploaded to GPU viajax.device_put, contributing toCUDA_ERROR_OUT_OF_MEMORYerrors.wfluxis now stripped (set toNone) before thedevice_putcall in bothsteady_stateandtime_dependent.test_ned_grid.py_allclose— failed on NaN and large fill values._load_single_slotreturns raw xarray values (NaN at land cells; float64 fill ≈ 9.97e36), whileLazyFluxFields.compute()normalises these to zero._allclosenow applies the same normalisation (clip → float32 →nan_to_num→ fill-value zeroing) before comparing.test_speed.py—snapshotcolumn no longer exists. Both speed tests read asnapshotcolumn that was replaced bywindow(int32) in theParticleRecorderschema. Tests now requestwindow, sort by it, and filter onevent_type == "snapshot".test_performance.py— spurioustseasargument. The timed-run call in_timedpassedtseas=tseasto the integration function; the variable was never defined, raisingNameError. 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 intoio,jax, andpyphases plus a totals summary."log"— emits onelogging.inforecord per window with structuredextra=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 byrich: 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_sublinspace fromlinspace(0, 1, …)tolinspace(-1, 0, …)so offsets span(-1, 0), placing the single sub-particle at-0.5relative toivec, i.e. the cell centreivec - 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 |
|---|---|---|
|
0 |
Still running (or paused at window boundary) |
|
1 |
Time limit exceeded ( |
|
2 |
Entered a land cell ( |
|
3 |
Left model domain ( |
|
4 |
Zero float32 progress (ds ≈ 0 / UNDEF — Fortran infLoop/dsCrossError) |
|
5 |
Progressive deactivation (too slow, stuck oscillator) |
|
6 |
Sub-cycle iteration limit reached |
|
7 |
Zero or negative cell volume ( |
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:
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 advancestsby only a fraction of the window width), so windows with 100–200 sub-cycles over 2 M particles became stuck for tens of minutes.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 allN_WINDOWSwindows (snapshot count is the correct completion metric; the finalactiveflag 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 withpytest -sto 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_resultfixture now acceptstmp_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 withpd.read_parquetand filter bysnapshotandntraccolumns instead of indexingrecorder._snapshots.tests/test_speed.py— replacedxarray.open_datasetwithpd.read_parquet; fixture renamedtrajectory_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):
ufluxused−sin(πy)·cos(2πx)→ should be−cos(πy)·sin(2πx)vfluxused+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) accessedwflux[ia-1, ja-1, ka-2, ...]without clamping whenka=1, causing JAX’s-1index wrap-around to read the wrong wflux level. Fixed by usingjnp.maximum(ka - 2, 0), consistent with the analytical scheme’sjnp.clip(ka - 2, 0, km)andextract_stencil’snp.maximum(ka - 2, 0).Code clarity: Cached
jnp.floor()results in_subcycle_one_particleand_subcycle_one_particle_stencilto 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 daskfrom inside theextract_stencilfunction body to module level ingcm/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 Fortranprojects/Theoreticalrun. Step index derived fromdtm - DTM0in units ofTIMESTEP(1 h).gcm.fabboni.initial_state()— single-particle state atx=150.0, y=99.5, z=0.5(isec=1, idir=1seeding).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.
GridFieldsno longer carriesdtm. It now contains only geometry (kmt,dxdy) and fixed metadata (imt,jmt,km,timestep,reverse,iperio,jperio).LazyFluxFieldsgains adtm: pd.Timestamp | Nonefield. Set bynew_fluxandupdate_flux;Nonewhen constructed via lower-level helpers likeswap_fields.new_grid→new_flux— same signature and return type;dtmis now stored on the returnedLazyFluxFieldsinstead ofGridFields.update_grid→update_flux— now returns justLazyFluxFields(not a tuple). Call sites change fromgrid, fluxes = gcm.update_grid(grid, fluxes, fn)tofluxes = 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. UnlikeFluxFields,kmtanddxdyare 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 aLocalFluxFields.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 viada.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 raiseIndexError. All reads are batched into a singledask.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 aFluxStencil(afterjax.vmappeels the N dimension) and computesdxyzfrom per-particle geometry. Returnsdtin seconds (physical time elapsed) rather than normalisedds.advection_scheme.piecewise_stencil(grid, no_wflux)(advection_scheme.py) — factory that closes overgrid(forkmtland checks on the destination cell) and returnsstep_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_endbut still withintimax_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_dependentparticle 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 applyingstate._replace(active=...)before advancing to the next window.time_dependent_patchedparticle reactivation (integrate/streaming.py) — the patched variant had no reactivation logic at all. Added the same numpy-based reactivation astime_dependent_streaming.
Cleanup#
Removed stale example functions (
test_steady_state,example_2–5) fromtracmass/__init__.py; they referenced non-existent APIs from an earlier prototype.Added missing
swap_fieldsimport totracmass/__init__.py.
API additions#
Symbol |
Module |
Description |
|---|---|---|
|
|
Patch-local flux + geometry pytree |
|
|
Per-particle point-extracted flux pytree |
|
|
Batched Dask/NumPy point indexer |
|
|
Patch-based multi-window integration |
|
|
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 |
|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|---|---|
|
Re-exports |
|
|
|
|
stencil.py imports the shared helpers from full_grid.py; no circular
dependency.
integrate/ subpackage#
File |
Contents |
|---|---|
|
Re-exports |
|
|
|
|
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 |
|---|---|---|
|
Full-grid baseline |
Computes entire GLORYS grid |
|
Local-stencil patch |
|
|
Per-particle stencil |
Scales to any particle distribution |
compare() runs all three and prints a side-by-side timing summary with
speedup ratios.