src/forge_lab/models/tensorflow_model.py

from __future__ import annotations

import numpy as np
import tensorflow as tf

from .numpy_runtime import DEFAULT_CENTER, DEFAULT_SCALE, WINDOW_SIZE


class TensorFlowDriftModel(tf.Module):
    """TensorFlow implementation with weights shared from the PyTorch head."""

    def __init__(self, linear_weight: np.ndarray, linear_bias: float) -> None:
        super().__init__()
        x = np.arange(WINDOW_SIZE, dtype=np.float32)
        x_centered = x - x.mean()
        self.x_centered = tf.constant(x_centered)
        self.slope_denominator = tf.constant(np.square(x_centered).sum(), dtype=tf.float32)
        self.feature_center = tf.constant(DEFAULT_CENTER)
        self.feature_scale = tf.constant(DEFAULT_SCALE)
        self.linear_weight = tf.constant(np.asarray(linear_weight, dtype=np.float32))
        self.linear_bias = tf.constant(float(linear_bias), dtype=tf.float32)

    @tf.function(
        input_signature=[tf.TensorSpec([None, WINDOW_SIZE, 2], tf.float32, name="telemetry")]
    )
    def __call__(self, telemetry: tf.Tensor) -> dict[str, tf.Tensor]:
        force = telemetry[:, :, 0]
        deviation = telemetry[:, :, 1]
        slope = tf.reduce_sum(force * self.x_centered, axis=1) / self.slope_denominator
        shift = tf.reduce_mean(force[:, -10:], axis=1) - tf.reduce_mean(force[:, :10], axis=1)
        std = tf.math.reduce_std(force, axis=1)
        max_deviation = tf.reduce_max(deviation, axis=1)
        last_deviation = deviation[:, -1]
        force_range = tf.reduce_max(force, axis=1) - tf.reduce_min(force, axis=1)
        features = tf.stack(
            [slope, shift, std, max_deviation, last_deviation, force_range],
            axis=1,
        )
        normalized = (features - self.feature_center) / self.feature_scale
        logits = tf.linalg.matvec(normalized, self.linear_weight) + self.linear_bias
        return {"probability": tf.math.sigmoid(logits), "features": features}