Improve pre-trained channel embeddings handling in Signal-JEPA - #991

Pure function that maps (channel_embedding, chs_info) to the tuple
(effective_chs_info, channel_locations, ch_idxs). All validation lives
here so the model constructor can stay concise and the logic is
unit-testable in isolation.

Case-insensitive name matching for 'pretrain_aligned' mode. Explicit
ValueError with actionable fallback message when a user channel is
outside the pre-training set.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

_PosEncoder now requires ch_idxs at __init__ and stores it via
register_buffer(..., persistent=False). forward() uses this buffer when
its ch_idxs kwarg is None. Renames the init kwarg ch_locs ->
channel_locations for clarity (the locations describe the table rows,
not the input channels).

Non-persistent is deliberate: ch_idxs is derived from chs_info +
channel_embedding at __init__, so saving it in state_dict would
duplicate config.json and conflict with user-overridden chs_info at
from_pretrained time.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

New kwarg channel_embedding: Literal['scratch', 'pretrain_aligned'].
Runs the resolver before super().__init__, so when the user passes
chs_info=None with channel_embedding='pretrain_aligned', the 62
pre-training channels become the effective chs_info (and hence n_chans).

Stores self._channel_embedding so build_model_config serializes it
into config.json for HF round-trips.

The feature_encoder-only path (_init_transformer=False, used by
PreLocal and PostLocal via _BaseSignalJEPA) skips the resolver
entirely.

SignalJEPA.__init__ now forwards the new kwarg to its super. Other
subclasses (Contextual, PreLocal, PostLocal) will be updated in Task 5.

Fixed _PRETRAIN_CHS_INFO locations from all-zeros to realistic 3D
coordinates to avoid division-by-zero in positional encoding.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

Replace the synthetic topographic coordinates (introduced during Task 4
implementation to work around division-by-zero in _ChannelEmbedding
reset_parameters when all locs are identical) with obviously-fake
systematic placeholders: loc = [0.001*i, 0.0, 0.0] for the i-th channel.

These are clearly NOT real EEG coordinates, so there is no risk of the
model being silently used with incorrect locations. The TODO(PRETRAIN_COORDS)
comment continues to mark the block for replacement before publishing
the HuggingFace checkpoint.

Channel names and order are unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

Both public classes now accept channel_embedding at __init__ and
forward it to _BaseSignalJEPA. The ch_idxs argument is removed from
their forward() signatures — it is now resolved once at __init__ and
stored in _PosEncoder.default_ch_idxs (non-persistent buffer).

PreLocal and PostLocal are unchanged: they don't build a pos_encoder.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

The method set_fixed_ch_names did not exist on _PosEncoder; the call
was guarded by 'if chs_info is not None', so it silently produced an
AttributeError whenever a caller used the transfer-learning path of
SignalJEPA_Contextual.from_pretrained with chs_info set.

The functionality it attempted to provide (binding the model to a
specific channel list) is now covered by the channel_embedding
parameter plus the chs_info passed directly to the constructor.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

Verifies that:
  - channel_embedding round-trips through config.json (via
    build_model_config + track_model_init_kwargs)
  - default_ch_idxs is NOT in state_dict (non-persistent buffer)
  - user-overridden chs_info at from_pretrained time correctly rebuilds
    the ch_idxs mapping while still loading the 62-row embedding weights

Gated on huggingface_hub availability.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

Replaces the systematic placeholder locs (0.001 * i on the x axis)
with the authoritative 62 channels used to pre-train Signal-JEPA.

The channel set is derived from the MOABB Lee2019_SSVEP dataset,
subject 1, session '1', run '1train': the first 62 ch_names (the
last 5 of the dataset's 67 channels are 4 EMG + 1 stim). Positions
come from MNE's standard_1005 montage.

The channel NAMES and ORDER in this list are now the definitive
contract with the HuggingFace checkpoint braindecode/signal-jepa;
changing either would silently mismatch the _ChannelEmbedding row
order of the published weights.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

When SignalJEPA_Contextual.from_pretrained(src_instance, chs_info=...)
was called with a pretrain_aligned source and a different destination
chs_info, the deepcopy()'d pos_encoder carried the source's
default_ch_idxs unchanged — mapping the user's channels to the WRONG
rows of the 62-entry embedding table. No runtime error, just silently
incorrect results.

Fix: after the deepcopy, re-resolve default_ch_idxs against the
destination chs_info when the source was built in pretrain_aligned
mode, and propagate _channel_embedding onto the new model so
subsequent save_pretrained round-trips use the right config.

Add regression test test_from_pretrained_transfer_pretrain_aligned_different_subset
that constructs a source on pretrain rows [0, 1, 2] and transfers to
destination rows [5, 0, 10], asserting the recomputed buffer and
preserved embedding weights.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

'hub_id' is not a real kwarg; the first positional arg of from_pretrained
is pretrained_model_name_or_path. Rephrase as a descriptive 'checkpoint'
label to avoid pointing users to a non-existent parameter.

build_model_config captures __init__ kwargs directly, not instance
attributes, so the placement of this assignment relative to super().__init__
is irrelevant for config tracking. The assignment speaks for itself.

The list comprehension was 96 chars, over the 88-char ruff line-length
configured in pyproject.toml. pre-commit's ruff hooks don't currently
run on test/, which is why it slipped through.

Co-authored-by: Copilot <[email protected]>