add transformer models and predictive maintenance simulator

tests/test_simulator.py

@@ -32,3 +32,21 @@ def test_simulator_base():
 
     s.fit()
     MockDriver.fit.assert_called_once_with(s._data.X, s._data.y)
+
+
+def test_pdm_simulator():
+    n_samples = 10
+    data = tsgm.dataset.DatasetProperties(N=100, T=12, D=23)
+    pdm_simulator = tsgm.simulator.PredictiveMaintenanceSimulator(data)
+    syn_dataset, equipment = pdm_simulator.generate(n_samples)
+    assert len(equipment) == 10
+    assert len(syn_dataset) == 10
+    for d in equipment:
+        assert isinstance(d, dict)
+
+    new_sim = pdm_simulator.clone()
+    params1 = pdm_simulator.params()
+    params2 = new_sim.params()
+    assert params1["switches"] == params2["switches"]
+    assert params1["m_norms"] == params2["m_norms"]
+    assert params1["sigma_norms"] == params2["sigma_norms"]

tests/test_zoo.py

@@ -1,9 +1,6 @@
 import pytest
 
-import functools
 import numpy as np
-import random
-import tensorflow as tf
 from tensorflow.keras import layers
 
 import sklearn.metrics.pairwise
@@ -48,7 +45,8 @@ def test_zoo_cgan(model_type):
 @pytest.mark.parametrize("model_type_name", [
     "clf_cn",
     "clf_cl_n",
-    "clf_block"],
+    "clf_block",
+    "clf_transformer",]
 )
 def test_zoo_clf(model_type_name):
     seq_len = 10

tsgm/metrics/metrics.py

@@ -306,10 +306,6 @@ class EntropyMetric(Metric):
     """
     Calculates the spectral entropy of a dataset or tensor.
 
-    This metric measures the randomness or disorder in a dataset or tensor
-    using spectral entropy, which is a measure of the distribution of energy
-    in the frequency domain.
-
     Args:
         d (tsgm.dataset.DatasetOrTensor): The input dataset or tensor.
 

tsgm/models/architectures/zoo.py

@@ -653,7 +653,7 @@ def _build_model(self) -> keras.Model:
 
 class BasicRecurrentArchitecture(Architecture):
     """
-    Base class for basic recurrent neural network architectures.
+    Base class for recurrent neural network architectures.
 
     Inherits from Architecture.
     """
@@ -726,6 +726,76 @@ def build(self, activation: str = "sigmoid", return_sequences: bool = True) -> k
         return model
 
 
+class TransformerClfArchitecture(BaseClassificationArchitecture):
+    """
+    Base class for transformer architectures.
+
+    Inherits from BaseClassificationArchitecture.
+    """
+
+    arch_type = "downstream:classification"
+
+    def __init__(self, seq_len: int, feat_dim: int, num_heads: int = 2, ff_dim: int = 64, n_blocks: int = 1, dropout_rate=0.5, output_dim: int = 2) -> None:
+        """
+        Initializes the TransformerClfArchitecture.
+
+        :param seq_len: Length of input sequences.
+        :type seq_len: int
+        :param feat_dim: Dimensionality of input features.
+        :type feat_dim: int
+        :param num_heads: Number of attention heads (default is 2).
+        :type num_heads: int
+        :param ff_dim: Feed forward dimension in the attention block (default is 64).
+        :type ff_dim: int
+        :param output_dim: Dimensionality of the output.
+        :type output_dim: int
+        :param dropout_rate: Dropout probability (default is 0.5).
+        :type dropout_rate: float, optional
+        :param n_blocks: Number of transformer blocks (default is 1).
+        :type n_blocks: int, optional
+        :param output_dim: Number of classes (default is 2).
+        :type output_dim: int, optional
+        """
+
+        self._num_heads = num_heads
+        self._ff_dim = ff_dim
+        self._n_blocks = n_blocks
+        self._dropout_rate = dropout_rate
+
+        super().__init__(seq_len, feat_dim, output_dim)
+
+    def transformer_block(self, inputs):
+        # Multi-Head Attention
+        attention_output = layers.MultiHeadAttention(
+            num_heads=self._num_heads,
+            key_dim=inputs.shape[-1]
+        )(inputs, inputs)
+        attention_output = layers.Dropout(self._dropout_rate)(attention_output)
+        attention_output = layers.LayerNormalization(epsilon=1e-6)(attention_output + inputs)
+
+        # Feed-Forward Network
+        ff_output = layers.Dense(self._ff_dim, activation="relu")(attention_output)
+        ff_output = layers.Dense(inputs.shape[-1])(ff_output)
+        ff_output = layers.Dropout(self._dropout_rate)(ff_output)
+        ff_output = layers.LayerNormalization(epsilon=1e-6)(ff_output + attention_output)
+
+        return ff_output
+
+    def _build_model(self) -> keras.Model:
+        inputs = layers.Input(shape=(self._seq_len, self._feat_dim))
+
+        x = inputs
+        for _ in range(self._n_blocks):
+            x = self.transformer_block(x)
+
+        x = layers.GlobalAveragePooling1D()(x)
+        x = layers.Dropout(self._dropout_rate)(x)
+        outputs = layers.Dense(self._output_dim, activation="softmax")(x)
+
+        model = keras.Model(inputs, outputs)
+        return model
+
+
 class cGAN_LSTMnArchitecture(BaseGANArchitecture):
     """
     Conditional Generative Adversarial Network (cGAN) with LSTM-based architecture.
@@ -837,5 +907,6 @@ def summary(self) -> None:
         "clf_cl_n": ConvnLSTMnArchitecture,
         "clf_block": BlockClfArchitecture,
         "recurrent": BasicRecurrentArchitecture,
+        "clf_transformer": TransformerClfArchitecture
     }
 )