Skip to content

Make spatiotemporal autoencoders reconstruct sequences - #154

Open
sergioald wants to merge 1 commit into
GeoOcean:developfrom
sergioald:add-full-sequence-reconstruction
Open

Make spatiotemporal autoencoders reconstruct sequences#154
sergioald wants to merge 1 commit into
GeoOcean:developfrom
sergioald:add-full-sequence-reconstruction

Conversation

@sergioald

@sergioald sergioald commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR changes the two spatiotemporal autoencoders to reconstruct complete
input sequences.

The public mappings become:

ConvLSTMAutoencoder:
(B, T, C, H, W) -> (B, k) -> (B, T, C, H, W)

HybridConvLSTMTransformerAutoencoder:
(B, T, C, H, W) -> (B, k) -> (B, T, C, H, W)

There is no reconstruction_mode option. Both classes now have one clear
autoencoder objective: reconstruct the complete sample supplied to the
encoder.

Rationale

The previous implementation encoded a complete sequence but reconstructed
only its final input frame. That sequence-to-frame behaviour was useful, but
it was surprising for classes named autoencoders and prevented a direct
full-window comparison with PCA and the existing LSTMAutoencoder.

A new user would normally expect:

X_hat = model.predict(X)
assert X_hat.shape == X.shape

This PR makes that expectation true for all reconstruction-oriented sequence
models.

Architecture

ConvLSTM autoencoder

  • Encodes the complete sequence with the existing ConvLSTM stack.
  • Maps the final recurrent representation to one latent vector of size k.
  • Adds learned temporal embeddings to the latent vector.
  • Uses an LSTM temporal decoder to create one code per output timestep.
  • Applies one shared spatial decoder to every timestep.

Hybrid ConvLSTM-Transformer autoencoder

  • Preserves the existing ConvLSTM and temporal-attention encoder.
  • Maps the full input window to one latent vector of size k.
  • Uses latent-conditioned learned temporal queries.
  • Refines the queries with the configured linear or standard attention blocks.
  • Applies one shared spatial decoder to every timestep.

API behaviour

model = ConvLSTMAutoencoder(k=32)
model.fit(X)

Z = model.encode(X)
X_hat = model.predict(X)
X_from_z = model.decode(Z)

assert Z.shape == (len(X), 32)
assert X_hat.shape == X.shape
assert X_from_z.shape == X.shape

The same behaviour applies to
HybridConvLSTMTransformerAutoencoder.

fit(X) and the reconstruction metric helpers use the complete sequence X
as their default target. A custom y must have the same shape as X.

Breaking change

This intentionally changes the output shape of both spatiotemporal classes:

before: (B, C, H, W)
after:  (B, T, C, H, W)

Existing final-frame checkpoints are not compatible with the new decoder
architecture. They should remain tied to the earlier BlueMath_tk version that
created them. Loading them into the new classes is expected to fail rather
than silently initialise missing sequence-decoder parameters.

The removed final-frame operation is better represented in future by a
separately named sequence-to-frame model, not by an autoencoder mode.

Tests

The PR adds or updates coverage for:

  • complete-sequence output shapes for ConvLSTM and Hybrid models;
  • predict(X) == decode(encode(X));
  • automatic full-sequence reconstruction targets;
  • manual reconstruction-error agreement;
  • frame-target rejection;
  • fixed sequence-length validation;
  • sequence checkpoint round trips;
  • one-vector decoding;
  • non-divisible spatial dimensions and output cropping;
  • linear and standard Hybrid attention;
  • finite non-zero gradients in sequence-specific decoder parameters;
  • timestep-dependent decoder outputs;
  • rejection of the removed reconstruction_mode argument;
  • updated public API and smoke-test expectations.

Local validation

Replace these placeholders with the results from the complete checkout:

compileall: passed
focused Ruff checks: passed
sequence reconstruction tests: 16 passed
autoencoder API tests: 20 passed
autoencoder smoke tests: 8 passed
metrics tests: 8 passed
full suite: 257 passed, 113 warnings

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant