mixle.utils.parallel.lightning_data module

Lightning-backed encoded-data handle for mini-batch / stochastic EM (WS-C2).

LightningEncodedData plugs PyTorch Lightning’s data tooling into mixle’s encoded-data backend registry (planner.encoded_data(..., backend="lightning")). Full-data EM operations delegate to a resident LocalEncodedData (identical results to backend="local"); the Lightning-specific value is mini-batch iteration via a lightning.pytorch.LightningDataModule + DataLoader (shuffling, batching, multi-worker collation), which drives stochastic / mini-batch EM through StreamingEstimator.

Lightning is an optional dependency: this module is imported only when the "lightning" backend is requested, so the rest of mixle (and CI without Lightning installed) is unaffected.

class LightningEncodedData(data, estimator=None, model=None, encoder=None, batch_size=None, shuffle=True, seed=0, sub_chunks=1, **_)[source]

Bases: EncodedDataHandle

Encoded-data handle that mini-batches via a Lightning DataModule for stochastic EM.

Parameters:

  • data (Any)

  • estimator (Any | None)

  • model (Any | None)

  • encoder (DataSequenceEncoder | None)

  • batch_size (int | None)

  • shuffle (bool)

  • seed (int)

  • sub_chunks (int)

  • _ (Any)

pysp_seq_log_density_sum(estimate)[source]

Return (num_observations, summed_log_density) for estimate.

Parameters:

estimate (Any)

Return type:

tuple[float, float]

pysp_seq_estimate(estimator, prev_estimate)[source]

Run one distributed/local sufficient-statistic fold and M-step.

Parameters:

  • estimator (Any)

  • prev_estimate (Any)

Return type:

Any

pysp_seq_initialize(estimator, rng, p)[source]

Initialize a model through the handle’s resident encoded data.

Parameters:
Return type:

Any

pysp_stream_accumulate(estimator, model)[source]

Return folded sufficient statistics for streaming/incremental EM.

Parameters:
Return type:

tuple[float, Any]

property datamodule

Return the underlying lightning.pytorch.LightningDataModule.

minibatches()[source]

Yield one epoch of raw-observation mini-batches via the Lightning DataLoader.

Return type:

Iterator[list[Any]]

stochastic_em(estimator, *, epochs=5, schedule=None, init_p=0.2, seed=0)[source]

Fit estimator by mini-batch stochastic EM over the Lightning DataLoader batches.

Runs epochs passes, feeding each DataLoader mini-batch to a StreamingEstimator (decayed accumulator + M-step). Returns the fitted model.

Parameters:
Return type:

Any

close()[source]

Release worker resources owned by this handle, if any.

Return type:

None