Skip to content

Model Estimators

centimators.model_estimators

Model estimator abstractions that combine Keras with the scikit-learn API.

This module exposes two estimator classes that conform to scikit-learn's BaseEstimator/TransformerMixin contracts while delegating all heavy‐ lifting to Keras. The goal is to let neural networks participate in classic ML pipelines without boilerplate.

Highlights
  • Drop-in compatibility – works with sklearn.pipeline.Pipeline, GridSearchCV, etc.
  • Distribution strategies – opt-in data-parallel training across multiple devices/GPUs.
  • Sequence support – :class:SequenceEstimator reshapes a flattened lag matrix into the 3-D tensor expected by recurrent or convolutional sequence layers.

BaseKerasEstimator dataclass

Bases: TransformerMixin, BaseEstimator, ABC

Meta-estimator for Keras models following the scikit-learn API.

Parameters:

Name Type Description Default
output_units int, default=1

Dimensionality of the model output. It is forwarded to :meth:build_model and can be used there when constructing the final layer.

 1
optimizer Type[optimizers.Optimizer], default=keras.optimizers.Adam

Optimiser class not instance. The class is instantiated in :meth:fit with the requested learning_rate.

 Adam
learning_rate float, default=1e-3

Learning-rate passed to the optimiser constructor.

 0.001
loss_function str or keras.losses.Loss, default="mse"

Loss forwarded to model.compile.

 'mse'
metrics list[str] | None, default=None

List of metrics forwarded to model.compile.

 None
model keras.Model | None, default=None

Internal Keras model instance. If None it is lazily built on the first call to :meth:fit.

 None
distribution_strategy str | None, default=None

Name of a Keras distribution strategy to activate before training. At the moment only "DataParallel" is recognised.

 None

Attributes:

Name Type Description
_n_features_in_ int | None

Inferred number of features from the data passed to :meth:fit.
Notes

Sub-classes must implement :meth:build_model which should return a compiled (or at least constructed) keras.Model instance.

Source code in src/centimators/model_estimators.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@dataclass(kw_only=True)
class BaseKerasEstimator(TransformerMixin, BaseEstimator, ABC):
    """Meta-estimator for Keras models following the scikit-learn API.

    Args:
        output_units (int, default=1): Dimensionality of the model output.
            It is forwarded to :meth:`build_model` and can be used there when
            constructing the final layer.
        optimizer (Type[optimizers.Optimizer], default=keras.optimizers.Adam):
            Optimiser class **not instance**. The class is instantiated in
            :meth:`fit` with the requested ``learning_rate``.
        learning_rate (float, default=1e-3): Learning-rate passed to the
            optimiser constructor.
        loss_function (str or keras.losses.Loss, default="mse"): Loss
            forwarded to ``model.compile``.
        metrics (list[str] | None, default=None): List of metrics forwarded
            to ``model.compile``.
        model (keras.Model | None, default=None): Internal Keras model instance.
            If *None* it is lazily built on the first call to :meth:`fit`.
        distribution_strategy (str | None, default=None): Name of a Keras
            distribution strategy to activate before training. At the moment
            only ``"DataParallel"`` is recognised.

    Attributes:
        _n_features_in_ (int | None): Inferred number of features from the data
            passed to :meth:`fit`.

    Notes:
        Sub-classes **must** implement :meth:`build_model` which should return
        a compiled (or at least constructed) ``keras.Model`` instance.
    """

    output_units: int = 1
    optimizer: Type[optimizers.Optimizer] = optimizers.Adam
    learning_rate: float = 0.001
    loss_function: str = "mse"
    metrics: list[str] | None = None
    model: Any = None
    distribution_strategy: str | None = None

    @abstractmethod
    def build_model(self):
        pass

    def _setup_distribution_strategy(self) -> None:
        """Activate a distribution strategy for multi-device training.

        The current implementation always uses
        ``keras.distribution.DataParallel`` which mirrors the model on all
        available devices and splits the batch.  Support for additional
        strategies can be added later.
        """
        # TODO: allow for different distribution strategies
        strategy = distribution.DataParallel()
        distribution.set_distribution(strategy)

    def fit(
        self,
        X,
        y,
        epochs: int = 100,
        batch_size: int = 32,
        validation_data: tuple[Any, Any] | None = None,
        callbacks: list[Any] | None = None,
        **kwargs: Any,
    ) -> "BaseKerasEstimator":
        """Fit the underlying Keras model.

        The model is **lazily** built and compiled on the first call. All
        extra keyword arguments are forwarded to ``keras.Model.fit``.

        Args:
            X (array-like): Training data of shape (n_samples, n_features).
            y (array-like): Training targets of shape (n_samples,) or (n_samples, n_outputs).
            epochs (int, default=100): Number of training epochs.
            batch_size (int, default=32): Minibatch size.
            validation_data (tuple[Any, Any] | None, default=None): Optional
                validation split forwarded to Keras.
            callbacks (list[Any] | None, default=None): Optional list of callbacks.
            **kwargs: Additional keyword arguments forwarded to ``keras.Model.fit``.

        Returns:
            BaseKerasEstimator: Fitted estimator.
        """
        self._n_features_in_ = X.shape[1]

        if self.distribution_strategy:
            self._setup_distribution_strategy()

        if not self.model:
            self.build_model()

        self.model.fit(
            nw.from_native(X).to_numpy(),
            y=nw.from_native(y, allow_series=True).to_numpy(),
            batch_size=batch_size,
            epochs=epochs,
            validation_data=validation_data,
            callbacks=callbacks,
            **kwargs,
        )
        self._is_fitted = True
        return self

    def predict(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
        """Generate predictions with the trained model.

        Args:
            X (array-like): Input samples of shape (n_samples, n_features).
            batch_size (int, default=512): Batch size used for inference.
            **kwargs: Additional keyword arguments forwarded to ``keras.Model.predict``.

        Returns:
            Any: Model predictions of shape (n_samples, output_units)
                in the same order as *X*.
        """
        if not self.model:
            raise ValueError("Model not built. Call `build_model` first.")

        return self.model.predict(X, batch_size=batch_size, **kwargs)

    def transform(self, X, **kwargs):
        """Alias for :meth:`predict` to comply with scikit-learn pipelines."""
        return self.predict(X, **kwargs)

    def __sklearn_is_fitted__(self) -> bool:
        """Return ``True`` when the estimator has been fitted.

        scikit-learn relies on :func:`sklearn.utils.validation.check_is_fitted`
        to decide whether an estimator is ready for inference.
        """
        return getattr(self, "_is_fitted", False)

fit(X, y, epochs=100, batch_size=32, validation_data=None, callbacks=None, **kwargs)

Fit the underlying Keras model.

The model is lazily built and compiled on the first call. All extra keyword arguments are forwarded to keras.Model.fit.

Parameters:

Name Type Description Default
X array - like

Training data of shape (n_samples, n_features).

 required
y array - like

Training targets of shape (n_samples,) or (n_samples, n_outputs).

 required
epochs int, default=100

Number of training epochs.

 100
batch_size int, default=32

Minibatch size.

 32
validation_data tuple[Any, Any] | None, default=None

Optional validation split forwarded to Keras.

 None
callbacks list[Any] | None, default=None

Optional list of callbacks.

 None
**kwargs Any

Additional keyword arguments forwarded to keras.Model.fit.

 {}

Returns:

Name Type Description
BaseKerasEstimator BaseKerasEstimator

Fitted estimator.
Source code in src/centimators/model_estimators.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def fit(
    self,
    X,
    y,
    epochs: int = 100,
    batch_size: int = 32,
    validation_data: tuple[Any, Any] | None = None,
    callbacks: list[Any] | None = None,
    **kwargs: Any,
) -> "BaseKerasEstimator":
    """Fit the underlying Keras model.

    The model is **lazily** built and compiled on the first call. All
    extra keyword arguments are forwarded to ``keras.Model.fit``.

    Args:
        X (array-like): Training data of shape (n_samples, n_features).
        y (array-like): Training targets of shape (n_samples,) or (n_samples, n_outputs).
        epochs (int, default=100): Number of training epochs.
        batch_size (int, default=32): Minibatch size.
        validation_data (tuple[Any, Any] | None, default=None): Optional
            validation split forwarded to Keras.
        callbacks (list[Any] | None, default=None): Optional list of callbacks.
        **kwargs: Additional keyword arguments forwarded to ``keras.Model.fit``.

    Returns:
        BaseKerasEstimator: Fitted estimator.
    """
    self._n_features_in_ = X.shape[1]

    if self.distribution_strategy:
        self._setup_distribution_strategy()

    if not self.model:
        self.build_model()

    self.model.fit(
        nw.from_native(X).to_numpy(),
        y=nw.from_native(y, allow_series=True).to_numpy(),
        batch_size=batch_size,
        epochs=epochs,
        validation_data=validation_data,
        callbacks=callbacks,
        **kwargs,
    )
    self._is_fitted = True
    return self

predict(X, batch_size=512, **kwargs)

Generate predictions with the trained model.

Parameters:

Name Type Description Default
X array - like

Input samples of shape (n_samples, n_features).

 required
batch_size int, default=512

Batch size used for inference.

 512
**kwargs Any

Additional keyword arguments forwarded to keras.Model.predict.

 {}

Returns:

Name Type Description
Any Any

Model predictions of shape (n_samples, output_units) in the same order as X.
Source code in src/centimators/model_estimators.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def predict(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
    """Generate predictions with the trained model.

    Args:
        X (array-like): Input samples of shape (n_samples, n_features).
        batch_size (int, default=512): Batch size used for inference.
        **kwargs: Additional keyword arguments forwarded to ``keras.Model.predict``.

    Returns:
        Any: Model predictions of shape (n_samples, output_units)
            in the same order as *X*.
    """
    if not self.model:
        raise ValueError("Model not built. Call `build_model` first.")

    return self.model.predict(X, batch_size=batch_size, **kwargs)

transform(X, **kwargs)

Alias for :meth:predict to comply with scikit-learn pipelines.

Source code in src/centimators/model_estimators.py 
154
155
156
def transform(self, X, **kwargs):
    """Alias for :meth:`predict` to comply with scikit-learn pipelines."""
    return self.predict(X, **kwargs)

__sklearn_is_fitted__()

Return True when the estimator has been fitted.

scikit-learn relies on :func:sklearn.utils.validation.check_is_fitted to decide whether an estimator is ready for inference.

Source code in src/centimators/model_estimators.py 
158
159
160
161
162
163
164
def __sklearn_is_fitted__(self) -> bool:
    """Return ``True`` when the estimator has been fitted.

    scikit-learn relies on :func:`sklearn.utils.validation.check_is_fitted`
    to decide whether an estimator is ready for inference.
    """
    return getattr(self, "_is_fitted", False)

SequenceEstimator dataclass

Bases: BaseKerasEstimator

Estimator for models that consume sequential data.

The class assumes that X is a flattened 2-D representation of a sequence built from multiple lagged views of the original signal. The shape transformation performed by :meth:_reshape is visualised below for a concrete example.

Parameters:

Name Type Description Default
lag_windows list[int]

Offsets (in number of timesteps) that have been concatenated to form the flattened design matrix.

 required
n_features_per_timestep int

Number of original features per timestep before creating the lags.

 required

Attributes:

Name Type Description
seq_length int

Inferred sequence length from lag_windows.
Source code in src/centimators/model_estimators.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
@dataclass(kw_only=True)
class SequenceEstimator(BaseKerasEstimator):
    """Estimator for models that consume sequential data.

    The class assumes that *X* is a **flattened** 2-D representation of a
    sequence built from multiple lagged views of the original signal.
    The shape transformation performed by :meth:`_reshape` is visualised
    below for a concrete example.

    Args:
        lag_windows (list[int]): Offsets (in number of timesteps) that have been
            concatenated to form the flattened design matrix.
        n_features_per_timestep (int): Number of *original* features per timestep
            **before** creating the lags.

    Attributes:
        seq_length (int): Inferred sequence length from lag_windows.
    """

    lag_windows: list[int]
    n_features_per_timestep: int

    def __post_init__(self):
        self.seq_length = len(self.lag_windows)

    def _reshape(self, X: IntoFrame, validation_data: tuple[Any, Any] | None = None):
        """Reshape a flattened lag matrix into a 3-D tensor.

        Args:
            X (IntoFrame): Design matrix containing the lagged features.
            validation_data (tuple[Any, Any] | None, default=None): Optional
                validation split; its *X* part will be reshaped in the same way.

        Returns:
            tuple[numpy.ndarray, tuple[Any, Any] | None]:
                A tuple containing the reshaped training data (numpy.ndarray with shape
                ``(n_samples, seq_length, n_features_per_timestep)``) and the
                (potentially reshaped) validation data.
        """
        X = nw.from_native(X).to_numpy()
        X_reshaped = ops.reshape(
            X, (X.shape[0], self.seq_length, self.n_features_per_timestep)
        )

        if validation_data:
            X_val, y_val = validation_data
            X_val = nw.from_native(X_val).to_numpy()
            X_val_reshaped = ops.reshape(
                X_val,
                (X_val.shape[0], self.seq_length, self.n_features_per_timestep),
            )
            validation_data = X_val_reshaped, nw.from_native(y_val).to_numpy()

        return X_reshaped, validation_data

    def fit(
        self, X, y, validation_data: tuple[Any, Any] | None = None, **kwargs: Any
    ) -> "SequenceEstimator":
        """Redefines :meth:`BaseKerasEstimator.fit`
        to include reshaping for sequence data.

        Args:
            X (array-like): Training data.
            y (array-like): Training targets.
            validation_data (tuple[Any, Any] | None, default=None): Optional
                validation split.
            **kwargs: Additional keyword arguments passed to the parent fit method.

        Returns:
            SequenceEstimator: Fitted estimator.
        """
        X_reshaped, validation_data_reshaped = self._reshape(X, validation_data)
        super().fit(
            X_reshaped,
            y=nw.from_native(y).to_numpy(),
            validation_data=validation_data_reshaped,
            **kwargs,
        )
        return self

    def predict(self, X, **kwargs: Any) -> numpy.ndarray:
        """Redefines :meth:`BaseKerasEstimator.predict`
        to include reshaping for sequence data.

        Args:
            X (array-like): Input data.
            **kwargs: Additional keyword arguments passed to the parent predict method.

        Returns:
            numpy.ndarray: Predictions of shape (n_samples, output_units).
        """
        X_reshaped, _ = self._reshape(X)
        return super().predict(X_reshaped, **kwargs)

fit(X, y, validation_data=None, **kwargs)

Redefines :meth:BaseKerasEstimator.fit to include reshaping for sequence data.

Parameters:

Name Type Description Default
X array - like

Training data.

 required
y array - like

Training targets.

 required
validation_data tuple[Any, Any] | None, default=None

Optional validation split.

 None
**kwargs Any

Additional keyword arguments passed to the parent fit method.

 {}

Returns:

Name Type Description
SequenceEstimator SequenceEstimator

Fitted estimator.
Source code in src/centimators/model_estimators.py 
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
def fit(
    self, X, y, validation_data: tuple[Any, Any] | None = None, **kwargs: Any
) -> "SequenceEstimator":
    """Redefines :meth:`BaseKerasEstimator.fit`
    to include reshaping for sequence data.

    Args:
        X (array-like): Training data.
        y (array-like): Training targets.
        validation_data (tuple[Any, Any] | None, default=None): Optional
            validation split.
        **kwargs: Additional keyword arguments passed to the parent fit method.

    Returns:
        SequenceEstimator: Fitted estimator.
    """
    X_reshaped, validation_data_reshaped = self._reshape(X, validation_data)
    super().fit(
        X_reshaped,
        y=nw.from_native(y).to_numpy(),
        validation_data=validation_data_reshaped,
        **kwargs,
    )
    return self

predict(X, **kwargs)

Redefines :meth:BaseKerasEstimator.predict to include reshaping for sequence data.

Parameters:

Name Type Description Default
X array - like

Input data.

 required
**kwargs Any

Additional keyword arguments passed to the parent predict method.

 {}

Returns:

Type Description
ndarray

numpy.ndarray: Predictions of shape (n_samples, output_units).
Source code in src/centimators/model_estimators.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
def predict(self, X, **kwargs: Any) -> numpy.ndarray:
    """Redefines :meth:`BaseKerasEstimator.predict`
    to include reshaping for sequence data.

    Args:
        X (array-like): Input data.
        **kwargs: Additional keyword arguments passed to the parent predict method.

    Returns:
        numpy.ndarray: Predictions of shape (n_samples, output_units).
    """
    X_reshaped, _ = self._reshape(X)
    return super().predict(X_reshaped, **kwargs)

MLPRegressor dataclass

Bases: RegressorMixin, BaseKerasEstimator

A minimal fully-connected multi-layer perceptron for tabular data.

The class follows the scikit-learn estimator interface while delegating the heavy lifting to Keras. It is intended as a sensible baseline model that works out of the box with classic ML workflows such as pipelines or cross-validation.

Parameters:

Name Type Description Default
hidden_units tuple[int, ...], default=(64, 64

Width (number of neurons) for each hidden layer. The length of the tuple defines the depth of the network.

 (64, 64)
activation str, default="relu"

Activation function applied after each hidden Dense layer.

 'relu'
dropout_rate float, default=0.0

Optional dropout applied after each hidden layer. Set to 0 to disable dropout entirely.

 0.0
output_units int, default=1

Copied from :class:BaseKerasEstimator. Defines the dimensionality of the final layer.

 1

Attributes:

Name Type Description
_n_features_in_ int | None

Inferred number of features from the data passed to :meth:fit.
Source code in src/centimators/model_estimators.py 
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
@dataclass(kw_only=True)
class MLPRegressor(RegressorMixin, BaseKerasEstimator):
    """A minimal fully-connected multi-layer perceptron for tabular data.

    The class follows the scikit-learn *estimator* interface while delegating
    the heavy lifting to Keras.  It is intended as a sensible baseline model
    that works *out of the box* with classic ML workflows such as pipelines or
    cross-validation.

    Args:
        hidden_units (tuple[int, ...], default=(64, 64)): Width (number of
            neurons) for each hidden layer.  The length of the tuple defines
            the depth of the network.
        activation (str, default="relu"): Activation function applied after
            each hidden ``Dense`` layer.
        dropout_rate (float, default=0.0): Optional dropout applied **after**
            each hidden layer.  Set to *0* to disable dropout entirely.
        output_units (int, default=1): Copied from :class:`BaseKerasEstimator`.
            Defines the dimensionality of the final layer.

    Attributes:
        _n_features_in_ (int | None): Inferred number of features from the data
            passed to :meth:`fit`.
    """

    hidden_units: tuple[int, ...] = (64, 64)
    activation: str = "relu"
    dropout_rate: float = 0.0
    metrics: list[str] | None = field(default_factory=lambda: ["mse"])

    def build_model(self):
        """Construct a simple MLP with the configured hyper-parameters."""
        inputs = layers.Input(shape=(self._n_features_in_,), name="features")
        x = inputs
        for units in self.hidden_units:
            x = layers.Dense(units, activation=self.activation)(x)
            if self.dropout_rate > 0:
                x = layers.Dropout(self.dropout_rate)(x)
        outputs = layers.Dense(self.output_units, activation="linear")(x)
        self.model = models.Model(inputs=inputs, outputs=outputs, name="mlp_regressor")

        self.model.compile(
            optimizer=self.optimizer(learning_rate=self.learning_rate),
            loss=self.loss_function,
            metrics=self.metrics,
        )

        return self

build_model()

Construct a simple MLP with the configured hyper-parameters.

Source code in src/centimators/model_estimators.py 
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
def build_model(self):
    """Construct a simple MLP with the configured hyper-parameters."""
    inputs = layers.Input(shape=(self._n_features_in_,), name="features")
    x = inputs
    for units in self.hidden_units:
        x = layers.Dense(units, activation=self.activation)(x)
        if self.dropout_rate > 0:
            x = layers.Dropout(self.dropout_rate)(x)
    outputs = layers.Dense(self.output_units, activation="linear")(x)
    self.model = models.Model(inputs=inputs, outputs=outputs, name="mlp_regressor")

    self.model.compile(
        optimizer=self.optimizer(learning_rate=self.learning_rate),
        loss=self.loss_function,
        metrics=self.metrics,
    )

    return self

BottleneckEncoder dataclass

Bases: BaseKerasEstimator

A bottleneck autoencoder that can learn latent representations and predict targets.

This estimator implements a bottleneck autoencoder architecture that: 1. Encodes input features to a lower-dimensional latent space 2. Decodes the latent representation back to reconstruct the input 3. Uses an additional MLP branch to predict targets from the decoded features

The model can be used both as a regressor (via predict) and as a transformer (via transform) to get latent space representations for dimensionality reduction.

Parameters:

Name Type Description Default
gaussian_noise float, default=0.035

Standard deviation of Gaussian noise applied to inputs for regularization.

 0.035
encoder_units list[tuple[int, float]], default=[(1024, 0.1)]

List of (units, dropout_rate) tuples defining the encoder architecture.

 lambda: [(1024, 0.1)]()
latent_units tuple[int, float], default=(256, 0.1

Tuple of (units, dropout_rate) for the latent bottleneck layer.

 (256, 0.1)
ae_units list[tuple[int, float]], default=[(96, 0.4)]

List of (units, dropout_rate) tuples for the autoencoder prediction branch.

 lambda: [(96, 0.4)]()
activation str, default="swish"

Activation function used throughout the network.

 'swish'
reconstruction_loss_weight float, default=1.0

Weight for the reconstruction loss.

 1.0
target_loss_weight float, default=1.0

Weight for the target prediction loss.

 1.0

Attributes:

Name Type Description
encoder Model

The encoder submodel for transforming inputs to latent space.
Source code in src/centimators/model_estimators.py 
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
@dataclass(kw_only=True)
class BottleneckEncoder(BaseKerasEstimator):
    """A bottleneck autoencoder that can learn latent representations and predict targets.

    This estimator implements a bottleneck autoencoder architecture that:
    1. Encodes input features to a lower-dimensional latent space
    2. Decodes the latent representation back to reconstruct the input
    3. Uses an additional MLP branch to predict targets from the decoded features

    The model can be used both as a regressor (via predict) and as a transformer 
    (via transform) to get latent space representations for dimensionality reduction.

    Args:
        gaussian_noise (float, default=0.035): Standard deviation of Gaussian noise
            applied to inputs for regularization.
        encoder_units (list[tuple[int, float]], default=[(1024, 0.1)]): List of 
            (units, dropout_rate) tuples defining the encoder architecture.
        latent_units (tuple[int, float], default=(256, 0.1)): Tuple of 
            (units, dropout_rate) for the latent bottleneck layer.
        ae_units (list[tuple[int, float]], default=[(96, 0.4)]): List of 
            (units, dropout_rate) tuples for the autoencoder prediction branch.
        activation (str, default="swish"): Activation function used throughout the network.
        reconstruction_loss_weight (float, default=1.0): Weight for the reconstruction loss.
        target_loss_weight (float, default=1.0): Weight for the target prediction loss.

    Attributes:
        encoder (keras.Model): The encoder submodel for transforming inputs to latent space.
    """

    gaussian_noise: float = 0.035
    encoder_units: list[tuple[int, float]] = field(default_factory=lambda: [(1024, 0.1)])
    latent_units: tuple[int, float] = (256, 0.1)
    ae_units: list[tuple[int, float]] = field(default_factory=lambda: [(96, 0.4)])
    activation: str = "swish"
    reconstruction_loss_weight: float = 1.0
    target_loss_weight: float = 1.0
    encoder: Any = None

    def build_model(self):
        """Construct the bottleneck autoencoder architecture."""
        if self._n_features_in_ is None:
            raise ValueError("Must call fit() before building the model")

        # Input layer
        inputs = layers.Input(shape=(self._n_features_in_,), name="features")
        x0 = layers.BatchNormalization()(inputs)

        # Encoder path
        encoder = layers.GaussianNoise(self.gaussian_noise)(x0)
        for units, dropout in self.encoder_units:
            encoder = layers.Dense(units)(encoder)
            encoder = layers.BatchNormalization()(encoder)
            encoder = layers.Activation(self.activation)(encoder)
            encoder = layers.Dropout(dropout)(encoder)

        # Latent bottleneck layer
        latent_units, latent_dropout = self.latent_units
        latent = layers.Dense(latent_units)(encoder)
        latent = layers.BatchNormalization()(latent)
        latent = layers.Activation(self.activation)(latent)
        latent_output = layers.Dropout(latent_dropout)(latent)

        # Create separate encoder model for transform method
        self.encoder = models.Model(inputs=inputs, outputs=latent_output, name="encoder")

        # Decoder path (reverse of encoder)
        decoder = latent_output
        for units, dropout in reversed(self.encoder_units):
            decoder = layers.Dense(units)(decoder)
            decoder = layers.BatchNormalization()(decoder)
            decoder = layers.Activation(self.activation)(decoder)
            decoder = layers.Dropout(dropout)(decoder)

        # Reconstruction output
        reconstruction = layers.Dense(self._n_features_in_, name="reconstruction")(decoder)

        # Target prediction branch from decoded features
        target_pred = reconstruction
        for units, dropout in self.ae_units:
            target_pred = layers.Dense(units)(target_pred)
            target_pred = layers.BatchNormalization()(target_pred)
            target_pred = layers.Activation(self.activation)(target_pred)
            target_pred = layers.Dropout(dropout)(target_pred)

        target_output = layers.Dense(self.output_units, activation="linear", name="target_prediction")(target_pred)

        # Create the full model with multiple outputs
        self.model = models.Model(
            inputs=inputs, 
            outputs=[reconstruction, target_output],
            name="bottleneck_encoder"
        )

        # Compile with multiple losses
        self.model.compile(
            optimizer=self.optimizer(learning_rate=self.learning_rate),
            loss={
                "reconstruction": "mse",
                "target_prediction": self.loss_function
            },
            loss_weights={
                "reconstruction": self.reconstruction_loss_weight,
                "target_prediction": self.target_loss_weight
            },
            metrics={
                "target_prediction": self.metrics or ["mse"]
            }
        )

        return self

    def fit(
        self,
        X,
        y,
        epochs: int = 100,
        batch_size: int = 32,
        validation_data: tuple[Any, Any] | None = None,
        callbacks: list[Any] | None = None,
        **kwargs: Any,
    ) -> "BottleneckEncoder":
        """Fit the bottleneck autoencoder.

        Args:
            X (array-like): Training data (features).
            y (array-like): Training targets.
            epochs (int, default=100): Number of training epochs.
            batch_size (int, default=32): Minibatch size.
            validation_data (tuple[Any, Any] | None, default=None): Optional
                validation split.
            callbacks (list[Any] | None, default=None): Optional callbacks.
            **kwargs: Additional arguments passed to keras.Model.fit.

        Returns:
            BottleneckEncoder: Fitted estimator.
        """
        # Store input dimension and build model
        self._n_features_in_ = X.shape[1]

        if self.distribution_strategy:
            self._setup_distribution_strategy()

        if not self.model:
            self.build_model()

        # Convert inputs to numpy arrays
        X_np = nw.from_native(X).to_numpy()
        y_np = nw.from_native(y, allow_series=True).to_numpy()

        # Create target dictionary for multiple outputs
        y_dict = {
            "reconstruction": X_np,
            "target_prediction": y_np
        }

        # Handle validation data
        if validation_data is not None:
            X_val, y_val = validation_data
            X_val_np = nw.from_native(X_val).to_numpy()
            y_val_np = nw.from_native(y_val, allow_series=True).to_numpy()
            validation_data = (
                X_val_np,
                {
                    "reconstruction": X_val_np,
                    "target_prediction": y_val_np
                }
            )

        # Train the model
        self.model.fit(
            X_np,
            y_dict,
            batch_size=batch_size,
            epochs=epochs,
            validation_data=validation_data,
            callbacks=callbacks,
            **kwargs,
        )

        self._is_fitted = True
        return self

    def predict(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
        """Generate target predictions using the fitted model.

        Args:
            X (array-like): Input samples.
            batch_size (int, default=512): Batch size for prediction.
            **kwargs: Additional arguments passed to keras.Model.predict.

        Returns:
            array-like: Target predictions.
        """
        if not self.model:
            raise ValueError("Model not built. Call 'fit' first.")

        X_np = nw.from_native(X).to_numpy()
        predictions = self.model.predict(X_np, batch_size=batch_size, **kwargs)

        # Return only the target predictions (second output)
        return predictions[1] if isinstance(predictions, list) else predictions

    def transform(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
        """Transform input data to latent space representation.

        Args:
            X (array-like): Input samples.
            batch_size (int, default=512): Batch size for transformation.
            **kwargs: Additional arguments passed to keras.Model.predict.

        Returns:
            array-like: Latent space representations.
        """
        if not self.encoder:
            raise ValueError("Encoder not built. Call 'fit' first.")

        X_np = nw.from_native(X).to_numpy()
        return self.encoder.predict(X_np, batch_size=batch_size, **kwargs)

    def fit_transform(self, X, y, **kwargs) -> Any:
        """Fit the model and return latent space representations.

        Args:
            X (array-like): Training data.
            y (array-like): Training targets.
            **kwargs: Additional arguments passed to fit.

        Returns:
            array-like: Latent space representations of X.
        """
        return self.fit(X, y, **kwargs).transform(X)

    def get_feature_names_out(self, input_features=None) -> list[str]:
        """Generate feature names for the latent space output.

        Args:
            input_features (array-like, optional): Ignored. Present for API compatibility.

        Returns:
            list[str]: Feature names for latent dimensions.
        """
        latent_dim = self.latent_units[0]
        return [f"latent_{i}" for i in range(latent_dim)]

build_model()

Construct the bottleneck autoencoder architecture.

Source code in src/centimators/model_estimators.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
def build_model(self):
    """Construct the bottleneck autoencoder architecture."""
    if self._n_features_in_ is None:
        raise ValueError("Must call fit() before building the model")

    # Input layer
    inputs = layers.Input(shape=(self._n_features_in_,), name="features")
    x0 = layers.BatchNormalization()(inputs)

    # Encoder path
    encoder = layers.GaussianNoise(self.gaussian_noise)(x0)
    for units, dropout in self.encoder_units:
        encoder = layers.Dense(units)(encoder)
        encoder = layers.BatchNormalization()(encoder)
        encoder = layers.Activation(self.activation)(encoder)
        encoder = layers.Dropout(dropout)(encoder)

    # Latent bottleneck layer
    latent_units, latent_dropout = self.latent_units
    latent = layers.Dense(latent_units)(encoder)
    latent = layers.BatchNormalization()(latent)
    latent = layers.Activation(self.activation)(latent)
    latent_output = layers.Dropout(latent_dropout)(latent)

    # Create separate encoder model for transform method
    self.encoder = models.Model(inputs=inputs, outputs=latent_output, name="encoder")

    # Decoder path (reverse of encoder)
    decoder = latent_output
    for units, dropout in reversed(self.encoder_units):
        decoder = layers.Dense(units)(decoder)
        decoder = layers.BatchNormalization()(decoder)
        decoder = layers.Activation(self.activation)(decoder)
        decoder = layers.Dropout(dropout)(decoder)

    # Reconstruction output
    reconstruction = layers.Dense(self._n_features_in_, name="reconstruction")(decoder)

    # Target prediction branch from decoded features
    target_pred = reconstruction
    for units, dropout in self.ae_units:
        target_pred = layers.Dense(units)(target_pred)
        target_pred = layers.BatchNormalization()(target_pred)
        target_pred = layers.Activation(self.activation)(target_pred)
        target_pred = layers.Dropout(dropout)(target_pred)

    target_output = layers.Dense(self.output_units, activation="linear", name="target_prediction")(target_pred)

    # Create the full model with multiple outputs
    self.model = models.Model(
        inputs=inputs, 
        outputs=[reconstruction, target_output],
        name="bottleneck_encoder"
    )

    # Compile with multiple losses
    self.model.compile(
        optimizer=self.optimizer(learning_rate=self.learning_rate),
        loss={
            "reconstruction": "mse",
            "target_prediction": self.loss_function
        },
        loss_weights={
            "reconstruction": self.reconstruction_loss_weight,
            "target_prediction": self.target_loss_weight
        },
        metrics={
            "target_prediction": self.metrics or ["mse"]
        }
    )

    return self

fit(X, y, epochs=100, batch_size=32, validation_data=None, callbacks=None, **kwargs)

Fit the bottleneck autoencoder.

Parameters:

Name Type Description Default
X array - like

Training data (features).

 required
y array - like

Training targets.

 required
epochs int, default=100

Number of training epochs.

 100
batch_size int, default=32

Minibatch size.

 32
validation_data tuple[Any, Any] | None, default=None

Optional validation split.

 None
callbacks list[Any] | None, default=None

Optional callbacks.

 None
**kwargs Any

Additional arguments passed to keras.Model.fit.

 {}

Returns:

Name Type Description
BottleneckEncoder BottleneckEncoder

Fitted estimator.
Source code in src/centimators/model_estimators.py 
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
def fit(
    self,
    X,
    y,
    epochs: int = 100,
    batch_size: int = 32,
    validation_data: tuple[Any, Any] | None = None,
    callbacks: list[Any] | None = None,
    **kwargs: Any,
) -> "BottleneckEncoder":
    """Fit the bottleneck autoencoder.

    Args:
        X (array-like): Training data (features).
        y (array-like): Training targets.
        epochs (int, default=100): Number of training epochs.
        batch_size (int, default=32): Minibatch size.
        validation_data (tuple[Any, Any] | None, default=None): Optional
            validation split.
        callbacks (list[Any] | None, default=None): Optional callbacks.
        **kwargs: Additional arguments passed to keras.Model.fit.

    Returns:
        BottleneckEncoder: Fitted estimator.
    """
    # Store input dimension and build model
    self._n_features_in_ = X.shape[1]

    if self.distribution_strategy:
        self._setup_distribution_strategy()

    if not self.model:
        self.build_model()

    # Convert inputs to numpy arrays
    X_np = nw.from_native(X).to_numpy()
    y_np = nw.from_native(y, allow_series=True).to_numpy()

    # Create target dictionary for multiple outputs
    y_dict = {
        "reconstruction": X_np,
        "target_prediction": y_np
    }

    # Handle validation data
    if validation_data is not None:
        X_val, y_val = validation_data
        X_val_np = nw.from_native(X_val).to_numpy()
        y_val_np = nw.from_native(y_val, allow_series=True).to_numpy()
        validation_data = (
            X_val_np,
            {
                "reconstruction": X_val_np,
                "target_prediction": y_val_np
            }
        )

    # Train the model
    self.model.fit(
        X_np,
        y_dict,
        batch_size=batch_size,
        epochs=epochs,
        validation_data=validation_data,
        callbacks=callbacks,
        **kwargs,
    )

    self._is_fitted = True
    return self

predict(X, batch_size=512, **kwargs)

Generate target predictions using the fitted model.

Parameters:

Name Type Description Default
X array - like

Input samples.

 required
batch_size int, default=512

Batch size for prediction.

 512
**kwargs Any

Additional arguments passed to keras.Model.predict.

 {}

Returns:

Type Description
Any

array-like: Target predictions.
Source code in src/centimators/model_estimators.py 
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
def predict(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
    """Generate target predictions using the fitted model.

    Args:
        X (array-like): Input samples.
        batch_size (int, default=512): Batch size for prediction.
        **kwargs: Additional arguments passed to keras.Model.predict.

    Returns:
        array-like: Target predictions.
    """
    if not self.model:
        raise ValueError("Model not built. Call 'fit' first.")

    X_np = nw.from_native(X).to_numpy()
    predictions = self.model.predict(X_np, batch_size=batch_size, **kwargs)

    # Return only the target predictions (second output)
    return predictions[1] if isinstance(predictions, list) else predictions

transform(X, batch_size=512, **kwargs)

Transform input data to latent space representation.

Parameters:

Name Type Description Default
X array - like

Input samples.

 required
batch_size int, default=512

Batch size for transformation.

 512
**kwargs Any

Additional arguments passed to keras.Model.predict.

 {}

Returns:

Type Description
Any

array-like: Latent space representations.
Source code in src/centimators/model_estimators.py 
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
def transform(self, X, batch_size: int = 512, **kwargs: Any) -> Any:
    """Transform input data to latent space representation.

    Args:
        X (array-like): Input samples.
        batch_size (int, default=512): Batch size for transformation.
        **kwargs: Additional arguments passed to keras.Model.predict.

    Returns:
        array-like: Latent space representations.
    """
    if not self.encoder:
        raise ValueError("Encoder not built. Call 'fit' first.")

    X_np = nw.from_native(X).to_numpy()
    return self.encoder.predict(X_np, batch_size=batch_size, **kwargs)

fit_transform(X, y, **kwargs)

Fit the model and return latent space representations.

Parameters:

Name Type Description Default
X array - like

Training data.

 required
y array - like

Training targets.

 required
**kwargs

Additional arguments passed to fit.

 {}

Returns:

Type Description
Any

array-like: Latent space representations of X.
Source code in src/centimators/model_estimators.py 
531
532
533
534
535
536
537
538
539
540
541
542
def fit_transform(self, X, y, **kwargs) -> Any:
    """Fit the model and return latent space representations.

    Args:
        X (array-like): Training data.
        y (array-like): Training targets.
        **kwargs: Additional arguments passed to fit.

    Returns:
        array-like: Latent space representations of X.
    """
    return self.fit(X, y, **kwargs).transform(X)

get_feature_names_out(input_features=None)

Generate feature names for the latent space output.

Parameters:

Name Type Description Default
input_features array - like

Ignored. Present for API compatibility.

 None

Returns:

Type Description
list[str]

list[str]: Feature names for latent dimensions.
Source code in src/centimators/model_estimators.py 
544
545
546
547
548
549
550
551
552
553
554
def get_feature_names_out(self, input_features=None) -> list[str]:
    """Generate feature names for the latent space output.

    Args:
        input_features (array-like, optional): Ignored. Present for API compatibility.

    Returns:
        list[str]: Feature names for latent dimensions.
    """
    latent_dim = self.latent_units[0]
    return [f"latent_{i}" for i in range(latent_dim)]

LSTMRegressor dataclass

Bases: RegressorMixin, SequenceEstimator

LSTM-based regressor for time series prediction.

This estimator uses stacked LSTM layers to model sequential dependencies in time series data. It supports bidirectional processing and various normalization strategies.

Parameters:

Name Type Description Default
lstm_units list[tuple[int, float, float]], default=[(64, 0.01, 0.01)]

List of tuples defining LSTM layers. Each tuple contains: - units: Number of LSTM units - dropout_rate: Dropout rate applied to inputs - recurrent_dropout_rate: Dropout rate applied to recurrent connections

 lambda: [(64, 0.01, 0.01)]()
use_batch_norm bool, default=False

Whether to apply batch normalization after each LSTM layer.

 False
use_layer_norm bool, default=False

Whether to apply layer normalization after each LSTM layer.

 False
bidirectional bool, default=False

Whether to use bidirectional LSTM layers.

 False
lag_windows list[int]

Inherited from SequenceEstimator.

 required
n_features_per_timestep int

Inherited from SequenceEstimator.

 required

Attributes:

Name Type Description
_n_features_in_ int | None

Inferred number of features from training data.
Source code in src/centimators/model_estimators.py 
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
@dataclass(kw_only=True)
class LSTMRegressor(RegressorMixin, SequenceEstimator):
    """LSTM-based regressor for time series prediction.

    This estimator uses stacked LSTM layers to model sequential dependencies
    in time series data. It supports bidirectional processing and various
    normalization strategies.

    Args:
        lstm_units (list[tuple[int, float, float]], default=[(64, 0.01, 0.01)]):
            List of tuples defining LSTM layers. Each tuple contains:
            - units: Number of LSTM units
            - dropout_rate: Dropout rate applied to inputs
            - recurrent_dropout_rate: Dropout rate applied to recurrent connections
        use_batch_norm (bool, default=False): Whether to apply batch normalization
            after each LSTM layer.
        use_layer_norm (bool, default=False): Whether to apply layer normalization
            after each LSTM layer.
        bidirectional (bool, default=False): Whether to use bidirectional LSTM layers.
        lag_windows (list[int]): Inherited from SequenceEstimator.
        n_features_per_timestep (int): Inherited from SequenceEstimator.

    Attributes:
        _n_features_in_ (int | None): Inferred number of features from training data.
    """

    lstm_units: list[tuple[int, float, float]] = field(default_factory=lambda: [(64, 0.01, 0.01)])
    use_batch_norm: bool = False
    use_layer_norm: bool = False
    bidirectional: bool = False
    metrics: list[str] | None = field(default_factory=lambda: ["mse"])

    def build_model(self):
        """Construct the LSTM architecture."""
        if self._n_features_in_ is None:
            raise ValueError("Must call fit() before building the model")

        # Input layer expecting 3D tensor (batch, timesteps, features)
        inputs = layers.Input(
            shape=(self.seq_length, self.n_features_per_timestep), 
            name="sequence_input"
        )

        x = inputs

        # Stack LSTM layers
        for layer_num, (units, dropout, recurrent_dropout) in enumerate(self.lstm_units):
            return_sequences = layer_num < len(self.lstm_units) - 1

            lstm_layer = layers.LSTM(
                units=units,
                activation="tanh",
                return_sequences=return_sequences,
                dropout=dropout,
                recurrent_dropout=recurrent_dropout,
                name=f"lstm_{layer_num}"
            )

            # Apply bidirectional wrapper if requested
            if self.bidirectional:
                x = layers.Bidirectional(lstm_layer, name=f"bidirectional_{layer_num}")(x)
            else:
                x = lstm_layer(x)

            # Apply normalization layers if requested
            if self.use_layer_norm:
                x = layers.LayerNormalization(name=f"layer_norm_{layer_num}")(x)
            if self.use_batch_norm:
                x = layers.BatchNormalization(name=f"batch_norm_{layer_num}")(x)

        # Output layer
        outputs = layers.Dense(self.output_units, activation="linear", name="output")(x)

        # Create and compile model
        self.model = models.Model(inputs=inputs, outputs=outputs, name="lstm_regressor")

        self.model.compile(
            optimizer=self.optimizer(learning_rate=self.learning_rate),
            loss=self.loss_function,
            metrics=self.metrics,
        )

        return self

build_model()

Construct the LSTM architecture.

Source code in src/centimators/model_estimators.py 
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
def build_model(self):
    """Construct the LSTM architecture."""
    if self._n_features_in_ is None:
        raise ValueError("Must call fit() before building the model")

    # Input layer expecting 3D tensor (batch, timesteps, features)
    inputs = layers.Input(
        shape=(self.seq_length, self.n_features_per_timestep), 
        name="sequence_input"
    )

    x = inputs

    # Stack LSTM layers
    for layer_num, (units, dropout, recurrent_dropout) in enumerate(self.lstm_units):
        return_sequences = layer_num < len(self.lstm_units) - 1

        lstm_layer = layers.LSTM(
            units=units,
            activation="tanh",
            return_sequences=return_sequences,
            dropout=dropout,
            recurrent_dropout=recurrent_dropout,
            name=f"lstm_{layer_num}"
        )

        # Apply bidirectional wrapper if requested
        if self.bidirectional:
            x = layers.Bidirectional(lstm_layer, name=f"bidirectional_{layer_num}")(x)
        else:
            x = lstm_layer(x)

        # Apply normalization layers if requested
        if self.use_layer_norm:
            x = layers.LayerNormalization(name=f"layer_norm_{layer_num}")(x)
        if self.use_batch_norm:
            x = layers.BatchNormalization(name=f"batch_norm_{layer_num}")(x)

    # Output layer
    outputs = layers.Dense(self.output_units, activation="linear", name="output")(x)

    # Create and compile model
    self.model = models.Model(inputs=inputs, outputs=outputs, name="lstm_regressor")

    self.model.compile(
        optimizer=self.optimizer(learning_rate=self.learning_rate),
        loss=self.loss_function,
        metrics=self.metrics,
    )

    return self