✨ feat: Self-Organizing Map ordering stage (Starkman et al. 2023) - #56
Conversation
github-actions
Bot
added
📝 Documentation
There was a problem hiding this comment.
🟡 Changes recommended
A few correctness/robustness fixes are needed (notably avoiding deprecated JAX APIs and improving deterministic ordering behavior on projection ties).
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds a new optional Self-Organizing Map (SOM) refinement stage to phasecurvefit, implementing Starkman et al. (2023) in pure JAX and integrating it into the orderer pipeline (including chaining via | and an init= contract for refinement).
Changes:
- Introduces
phasecurvefit.somand the pure-JAX SOM core (init_prototypes,fit,densify,chord,SOM1D). - Adds
SOMOrdererplusChainOrderer/AbstractOrderer.__or__and threadsinitthrough the ordering entry points. - Expands docs, citation surfacing, Quantity (
unxt) support, and comprehensive tests (core, orderer behavior, chaining, vmap/jit ensemble-readiness).
File summaries
| File | Description |
|---|---|
| tests/test_som.py | Core SOM functional tests (init/fit/densify/chord, jit, determinism). |
| tests/test_som_orderer.py | Integration tests for SOMOrderer standalone + chained behavior. |
| tests/test_som_ensemble.py | Ensures SOM core remains vmap/jit friendly for ensembles. |
| tests/test_orderers_unxt.py | Verifies Quantity support for SOMOrderer and chained MST→SOM. |
| tests/test_chain.py | Tests ChainOrderer semantics and init= threading contract. |
| src/phasecurvefit/som.py | Public SOM module re-export wrapper. |
| src/phasecurvefit/orderers.py | Re-exports SOMOrderer and ChainOrderer in public orderers API. |
| src/phasecurvefit/_src/som.py | Pure-JAX SOM implementation (training, spline backbone, chord projection). |
| src/phasecurvefit/_src/orderers/som.py | SOMOrderer implementation producing backbone + chord and respecting init. |
| src/phasecurvefit/_src/orderers/result.py | Adds optional chord field to OrderingResult. |
| src/phasecurvefit/_src/orderers/mst.py | Accepts (and ignores) init for chain compatibility. |
| src/phasecurvefit/_src/orderers/localflow.py | Accepts (and ignores) init for chain compatibility. |
| src/phasecurvefit/_src/orderers/chain.py | New chain orderer that forwards init stage-to-stage. |
| src/phasecurvefit/_src/orderers/base.py | Updates abstract contract: init param + ` |
| src/phasecurvefit/_interop/interop_unxt.py | Adds Quantity dispatch for SOMOrderer and threads init through Quantity path. |
| src/phasecurvefit/init.py | Exposes som at top-level package API. |
| README.md | Adds repository-level citation guidance including SOM citation. |
| pyproject.toml | Codespell ignore for “som”; Ruff ignore for conventional matrix names in SOM core. |
| docs/index.md | Adds SOM guide and citation block to docs index. |
| docs/guides/som.md | New SOM user guide (usage, limitations, chord, projection details, ensembles). |
| docs/guides/orderers.md | Documents chaining behavior and adds SOMOrderer to orderer overview. |
| docs/api/index.md | Adds API docs entries for SOM module, SOMOrderer, and ChainOrderer. |
Review details
- Files reviewed: 22/22 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #56 +/- ##
=======================================
Coverage ? 91.97%
=======================================
Files ? 34
Lines ? 1982
Branches ? 109
=======================================
Hits ? 1823
Misses ? 129
Partials ? 30 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Address PR GalacticDynamics#56 review feedback: - `SOMOrderer.order` now passes `stable=True` to `jnp.argsort` explicitly. This is already JAX's default, so behaviour is unchanged, but the tie behaviour is a contract worth stating: observations with equal chord values keep the order they arrived in, i.e. the prior stage's ordering when refining. - `_resample_uniform`'s docstring still claimed `chord` relies on uniform segment lengths. That was true before `chord` was changed to compute true cumulative arc length; it now contradicts `chord`'s own docstring. The uniform spacing is still needed, but by `OrderingResult._interp_backbone`, which interpolates by vertex index. Not applied: the suggestion to replace `jax.ops.segment_sum` with `jax.lax.segment_sum`. That symbol does not exist in JAX 0.8.2 (`jax.ops` is the only home for the segment reductions) and `jax.ops.segment_sum` emits no deprecation warning; the suite runs with `filterwarnings = ["error"]`, so a real deprecation would already be failing it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
🔵 Needs a closer look
Two chaining-contract issues need fixing (ChainOrderer passing metadata=None breaks Quantity dispatch selection, and SOMOrderer assumes init.ordering not guaranteed by AbstractResult).
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
src/phasecurvefit/_src/orderers/chain.py:79
- ChainOrderer always passes
metadata=metadatadown to stages, even whenmetadataisNone. This breaks the docstring claim that arguments are forwarded untouched and can prevent Plum from selecting theunxtQuantity dispatches, since those overloads expectmetadata: StateMetadata(notStateMetadata|None) and rely on their default when omitted.
src/phasecurvefit/_src/orderers/som.py:149 initis typed asAbstractResult, but this code relies oninit.ordering, which is not part ofAbstractResult(it's defined onOrderingResult). If a custom orderer returns a differentAbstractResultimplementation, chaining intoSOMOrdererwill raiseAttributeError. Use the guaranteedindicesfield to derive the visited ordering instead.
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
Address PR GalacticDynamics#56 review feedback. Both issues were real; one had a different root cause than reported. `SOMOrderer.order` accepted `init: AbstractResult` but read `init.ordering`, which is an `OrderingResult` property, not part of the `AbstractResult` contract the signature advertises. Any other `AbstractResult` implementation chained into the SOM would have raised `AttributeError`. It now derives the working set from `indices`, the ordering field the base class does guarantee. Omitting `metadata` on a Quantity run raised `AttributeError: 'NoneType' object has no attribute 'get'` instead of the intended guidance. The cause is not `ChainOrderer`: the `order` facade has always forwarded `metadata` unconditionally, so `MSTOrderer` alone fails identically on `main`. Fixed in `_require_usys`, where every orderer and entry point routes through, rather than special-casing the chain and leaving the direct call broken. Both paths now raise TypeError with the intended guidance, and both fixes carry regression tests -- the `AbstractResult` one using a deliberately bare subclass with no `ordering` property. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
🔵 Needs a closer look
The order() facade and ChainOrderer currently forward init= even when it is None, which can unnecessarily break compatibility with third-party orderers that don’t accept an init keyword.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
src/phasecurvefit/_src/orderers/base.py:101
order()always passesinit=down to the selected orderer even wheninitisNone. This can break third-party/custom orderers that were previously compatible but don't accept aninitkeyword. Passinginitonly when it is non-Nonepreserves backward compatibility while keeping refinement behavior when callers explicitly provide aninitresult.
src/phasecurvefit/_src/orderers/chain.py:78ChainOrdererunconditionally forwardsinit=to every stage, including the first stage whereinitis typicallyNone. This forces all stages (including third-party orderers) to accept aninitkeyword even when chaining is not being used. Consider only passinginitwhen it is non-Noneto avoid needless signature breakage.
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
Address PR GalacticDynamics#56 review feedback. `init` is new in this branch: the released v0.3.1 (which is exactly this PR's merge base, 602b6e2) ships `AbstractOrderer.order` with no such parameter. The `order()` facade and `ChainOrderer` were forwarding `init=` unconditionally, so any third-party orderer written against the released contract would raise TypeError even when no chaining was involved. Both now omit `init` when there is no prior result to forward. A non-conforming orderer keeps working standalone and as the head of a chain, and still fails -- clearly, and correctly -- if placed *after* another stage, which it genuinely cannot participate in. The abstract signature is unchanged: `init` remains part of the contract new orderers should implement. Three regression tests cover it, using a deliberately legacy-shaped orderer whose `order()` has no `init` parameter. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces a substantial new algorithmic stage plus a widened public orderer contract (init/chaining) across multiple layers (core, interop, docs, tests), warranting final human review for API/behavioral compatibility.
Review details
- Files reviewed: 22/22 changed files
- Comments generated: 0 new
- Review effort level: Lite
There was a problem hiding this comment.
🟡 Changes recommended
There is at least one behavior/documentation mismatch (and a related interop forwarding gap) that should be resolved to avoid surprising defaults and ensure chaining semantics remain consistent across the Quantity path.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
src/phasecurvefit/_interop/interop_unxt.py:806
- The Quantity dispatch for
MSTOrderer.orderacceptsinitbut does not forward it to the plain-array implementation, leavinginitunused and potentially diverging from the chaining contract if MST ever starts consuminginit(and may also trip unused-argument lint). Forwardinitto the inner call (it is ignored today).
- Files reviewed: 23/23 changed files
- Comments generated: 1
- Review effort level: Lite
…ty path Address PR GalacticDynamics#56 review feedback. Both were real. `chord` and `fit` still defaulted `metric_scale=1.0` while `SOMOrderer` and `SOM1D` defaulted `0.0` -- and `chord`'s own docstring, rewritten in 47d3298, asserts "the default `metric_scale=0.0`". A caller relying on the documented default with a velocity-aware metric got velocity weighted at 1.0 instead of off, which for `FullPhaseSpaceDistanceMetric` is the difference between pure position distance and velocity dominating on any stream whose speeds exceed its position separations. All four now default to 0.0, so "the default" means one thing across the whole surface. The `MSTOrderer` Quantity dispatch accepted `init` but dropped it when delegating to the plain path, while the `SOMOrderer` one forwards it. MST ignores `init` today so nothing changes, but the asymmetry would become a silent divergence between the Quantity and plain paths the moment MST consumed it. Forwarded. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a substantial new core algorithm + orderer stage touching metrics contracts, unit interop, and multiple user-facing docs/tutorials, warranting final human review despite strong test coverage.
Review effort: Lite
Findings: None
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
The PR introduces a substantial new algorithmic core, new public API surface, and multiple cross-cutting contract changes (metrics/orderers/Quantity interop/docs), which warrants final human review despite strong test coverage.
Review effort: Lite
Findings: None
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
_train_and_project recomputes the SOM backbone twice inside the jitted region (avoidable extra densify work), which should be fixed before approval.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (1)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer currently relies on Python truthiness/membership checks for metric_scale, which can raise unexpectedly for scalar-array inputs and should be made robust/explicit before approval.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a substantial new algorithmic ordering stage and API surface area, so a final human review is recommended despite strong test coverage.
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a large new algorithmic stage affecting public APIs, chaining semantics, unit interop, and documentation/tests, so it warrants final human review despite only minor requested adjustments.
Review effort: Lite
Findings: 1
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a substantial new ordering stage with multiple integration points (chaining semantics, unit interop, public API/docs), which warrants final human review despite strong test coverage.
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer.order() does not reject init.indices < -1, which can silently drop observations despite the documented “-1 for unvisited” contract.
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer.metric_scale’s type annotation is narrower than the accepted runtime inputs (0‑d JAX scalars), creating an API/type-hint mismatch that should be corrected.
Review effort: Lite
Findings: 2
Open (2)
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer currently branches on metric_scale via Python truthiness/membership in ways that can break under jit/grad when metric_scale is traced, so the metric/scale resolution should be made trace-safe (or fail with a clear, intentional error).
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (2)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer currently relies on Python truthiness/containment checks on metric_scale, which is brittle for array-like and traced scalars and can break JAX-transform safety guarantees stated in the implementation comments.
Review effort: Lite
Findings: 2
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a substantial new algorithmic stage plus API/docs/test expansions across many files that warrant final human review for correctness and long-term maintenance implications.
Review effort: Lite
Findings: None
Resolved since last review (1)
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer.order currently uses non-jittable constructs when init is provided (boolean indexing and Python scalar conversions for warnings), which can break jax.jit compilation for chained pipelines.
Review effort: Lite
Findings: 2
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
SOMOrderer.order can raise an unhelpful StopIteration on empty component dict input and should explicitly validate and raise a clear ValueError instead.
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (2)
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
It introduces a substantial new ordering algorithm plus cross-cutting API/interop/docs changes, so a final human review should validate correctness and long-term maintenance impact beyond what can be exhaustively confirmed from the provided diff excerpts.
Review effort: Lite
Findings: None
Resolved since last review (1)
Wraps `phasecurvefit.som` as an orderer. Most useful as a *refinement* after another stage -- the prototypes average over many tracers, so the backbone is far less sensitive to local noise than a greedy walk's step-by-step decisions or an MST's individual graph edges -- but it runs standalone too. It follows the previous stage's velocity-awareness. Ordering on position alone after a stage that used velocity undoes that stage's work: at a crossing the two branches are spatially coincident and differ only in velocity. When `init` reports `velocity_aware` and neither `metric` nor `metric_scale` was set, it resolves to `FullPhaseSpaceDistanceMetric` with a scale of `sigma_phys / (2 |v|)`, which separates anti-parallel branches by about the distance the lattice can already resolve. Measured on a self-intersecting epitrochoid: a velocity-aware MST scoring |rho| = 0.94 drops to 0.59 under a position-only SOM and rises to 0.99 under the default. `metric_scale=0.0` is an explicit choice, not "unset" -- testing truthiness would silently re-derive a non-zero scale for a caller who asked for position only. Chained, it seeds the lattice from the prior ordering and holds sigma at `sigma_end` rather than starting wide, so it refines locally instead of re-deriving a global ordering and discarding what it was given. It warns when its result disagrees wholesale with that input, which is what a lattice too coarse for the curve looks like. Includes the guide, the eight documented deviations from the paper, the Starkman citation surfaces, benchmarks, and the tutorials chained end to end. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GalacticDynamics#109 reorganised tests/ into unit/, smoke/, integration/ and static/ while this branch was open. The merge relocated every file it shared with main, but test_som_orderer.py is new here, so git had no rename to follow and left it at the top level. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`SOMOrderer` traces standalone but not when chained: the working set is whatever the prior stage visited, so `prior[prior >= 0]` has a size that depends on that stage's values. Under a transform `init.indices` is a tracer and no shape fixed at trace time holds the result. That cannot be fixed here -- the mask would have to travel through the whole core so rejected points stay rejected without slicing (GalacticDynamics#125). What it can do is fail legibly: chaining under `jit` now raises a `TypeError` naming the cause and the two ways out, rather than surfacing a `NonConcreteBooleanIndexError` from inside indexing. The limit is stated in the class docstring and the SOM guide, alongside the standalone case that does trace. `_warn_if_disagrees` was a second blocker on the same path, masked by the first: it concretises with `float()` to format its message. It now declines while tracing. A best-effort diagnostic is never worth failing a compile for -- the same reasoning that took `length` out of the track fraction it prints. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three guards caught a JAX error to decide what to do. Catching `TracerBoolConversionError` / `NonConcreteBooleanIndexError` is control flow by exception, and it made the rule look like "behave differently under a transform", which is not the rule. The rule is "this value has to be knowable here". `jax.core.is_concrete` says exactly that and says it up front, so the two guards that remain read as preconditions rather than as recoveries. Behaviour is unchanged: a `grad` trace carries a concrete primal and passes; an abstract `jit` trace does not. The third guard is gone rather than rewritten. `_warn_if_disagrees` returns early when `init is None`, and with an `init` the working-set selection has already refused to trace -- so nothing reached the concretisation it was guarding. It was dead code kept alive by a speculative argument about GalacticDynamics#125, which is the wrong reason to keep anything; GalacticDynamics#125 now records that the warning needs handling when the shape problem is fixed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A 7-angle review of this branch (correctness, removed-behavior, cross-file, reuse, simplification, efficiency, altitude) surfaced ten findings, all fixed here: - Chaining fed `init_prototypes` a synthetic `arange` ordering, so its own duplicate-index guard could never see a repeated `init.indices`. Added a uniqueness check on the real indices in `order()` itself. - `_resolve_metric`/`_warn_if_disagrees` computed arc length with plain `jnp.linalg.norm`, reintroducing the 0/0-gradient bug `_safe_norm` exists to prevent at coincident vertices -- and duplicated the computation between the two methods. Extracted one `_polyline_length` helper using `_safe_norm`. - The disagreement warning could print a self-contradictory pair of numbers (a smoothing length of 0 next to a nonzero percentage) on a degenerate, all-coincident working set. Named the degeneracy in the message instead. - `velocity_aware` was read off `init` via `getattr(..., False)`, duck-typing a contract that only `OrderingResult` actually declared. Promoted it to `AbstractResult` so every result type carries it. - Documented, in `AbstractOrderer.order`'s own docstring, the convention that a stage restricting itself to `init`'s visited subset must not reintroduce what a prior stage rejected -- nothing in `ChainOrderer` enforces this, so it needs to be stated where implementers will read it. - Extracted the near-identical MST/SOM Quantity dispatch bodies in `interop_unxt.py` into one shared `_order_with_backbone_and_chord`. - Added regression tests for two previously-untested paths: chaining with the default `n_prototypes=25` on a too-small visited set, and the `_chord_unit` fallback through the `LocalFlowOrderer` dispatch. - Documented why a uniformly tiny (not zero) `speed` does not blow up the derived `metric_scale`: scale and speed are reciprocal by construction, so a typical velocity's contribution is `sigma_phys / 2` regardless of scale. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Adds an optional Self-Organizing Map ordering stage, implementing the method of
Starkman et al. (2023) in JAX. The SOM learns a
1-D lattice of prototype vectors in phase space, smooths them into a backbone curve,
and orders observations by their arc length ("chord") projected onto it. It is
designed as a refinement stage after an initial ordering, before the learning step.
What's new
phasecurvefit.som— pure-JAX core:init_prototypes,fit,densify,chord, plus a thinSOM1Dwrapper.phasecurvefit.orderers.SOMOrderer— the orderer wrapping that core.ChainOrdererandAbstractOrderer.__or__— orderers now compose:MSTOrderer(k=10) | SOMOrderer(n_prototypes=15).order()gainedinit: AbstractResult | None = None, so a stage can refine a prior ordering; stages that can't simply ignore it.OrderingResult.chord— optional arc-length parameter per observation, in input order,nanwhere unvisited. The natural affine parameter for the downstream fit.SOMOrderer;backboneandchordreturn in the caller's position units.docs/guides/som.mdplus citation surfaces.The N-D projection
The published method defines its projection in 2-D and breaks ties at polyline
vertices with an
arctan2angle sweep, which has no N-D analogue. This implementationreplaces that entirely: prototypes are smoothed with a centripetal Catmull-Rom spline,
which dissolves the tie wedges, and each datum is then assigned to its nearest backbone
vertex under the full phase-space metric and refined to sub-vertex resolution in
position space.
That split is forced rather than stylistic:
AlignedMomentumDistanceMetricis notinduced by an inner product, so "project onto a line segment" is undefined under it.
The metric chooses which segment; Euclidean position geometry chooses where within
it — which also makes the chord a genuine physical arc length.
Deviations from the paper
Both are disclosed in the citation prose, so anyone citing the paper for results from
this code knows which parts are which:
fixed points but not trajectories, so published numbers will not reproduce bit-for-bit.
Known limitations
ordering, prototypes are binned along the first principal axis. A stream winding more
than ~1 turn produces a tangled lattice that training cannot repair (measured:
rho ~ 1.0for <=1 turn, collapsing beyond, at any sigma or epoch count). This is theN-D analogue of the paper's own caveat about equi-frequency binning when the data is
many-to-one in the binning coordinate. Remedy: chain after another orderer, which is
the recommended usage anyway. Documented in the docstring and the guide.
vmap-ability alone does not deliver the planned ensemble. Batch Kohonen iscontractive, so members from jittered initializations converge to bit-identical
results. A real posterior of orderings needs a genuine diversity source — bootstrap
resampling per member, per-member subsampling, or a stochastic update. Documented so
the follow-up isn't built naively.
Ensemble-readiness
The core is shaped so a future Statistically Combined Ensemble of SOMs is a
vmapoverfitandchordwith no change to the module: explicit initial prototypes, explicitkeys, static shapes,
lax.scanover epochs, no host callbacks, no Python branching ontraced values.
tests/test_som_ensemble.pymakes that guarantee executable rather thanaspirational — and it genuinely fails if the property is lost (verified by injecting a
traced-value branch into
fit's scan body).Testing
855 passed, 1 skipped (pre-existing conditional skip), with
filterwarnings = ["error"].Every test that guards a specific regression was verified to actually fail when that
regression is forced — the C1-spline test against a piecewise-linear stub, the chained
seeding test against
ordering=None, the unit assertions against un-reconverted units,and the ensemble distinctness assertions against a broadcast-degenerated
vmap.Citation
Users of this stage are asked to cite:
Surfaced via
__citation__class attributes, module and class docstrings,docs/guides/som.md,docs/index.md, and a new## Citationsection in the README,which previously had none.
🤖 Generated with Claude Code