Upload 33 files

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+denoising_example_16.gif filter=lfs diff=lfs merge=lfs -text
+denoising_example_25.gif filter=lfs diff=lfs merge=lfs -text
+denoising_example_4.gif filter=lfs diff=lfs merge=lfs -text
+denoising_example_9.gif filter=lfs diff=lfs merge=lfs -text
+noising_example_16.gif filter=lfs diff=lfs merge=lfs -text
+noising_example_25.gif filter=lfs diff=lfs merge=lfs -text
+noising_example_4.gif filter=lfs diff=lfs merge=lfs -text
+noising_example_9.gif filter=lfs diff=lfs merge=lfs -text

GAN/diffusion.py

@@ -0,0 +1,965 @@
+import math
+import shutil
+import numpy as np
+from tqdm.auto import tqdm
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+from GAN.utils import linear_beta_schedule, cosine_beta_schedule
+import matplotlib.pyplot as plt
+import os
+from GAN.utils import TSFeatureScaler
+
+class GaussianDiffusion:
+
+    def __init__(
+            self,
+            beta_schedule='cosine',
+            timesteps=10,
+            clip_min=-1.0,
+            clip_max=1.0,
+    ):
+
+        self.timesteps = timesteps
+        self.clip_min = clip_min
+        self.clip_max = clip_max
+
+        if beta_schedule == 'linear':
+            betas = linear_beta_schedule(timesteps)
+        elif beta_schedule == 'cosine':
+            betas = cosine_beta_schedule(timesteps)
+        else:
+            raise ValueError(f'unknown beta schedule {beta_schedule}')
+
+        alphas = 1. - betas
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+        self.betas = tf.constant(betas, dtype=tf.float32)
+        self.alphas_cumprod = tf.constant(alphas_cumprod, dtype=tf.float32)
+        self.alphas_cumprod_prev = tf.constant(alphas_cumprod_prev, dtype=tf.float32)
+        self.sqrt_recip_alphas = tf.constant(np.sqrt(1. / alphas), dtype=tf.float32)
+
+        self.sqrt_alphas_cumprod = tf.constant(np.sqrt(self.alphas_cumprod), dtype=tf.float32)
+        self.sqrt_one_minus_alphas_cumprod = tf.constant(np.sqrt(1.0 - self.alphas_cumprod), dtype=tf.float32)
+        self.log_one_minus_alphas_cumprod = tf.constant(np.log(1. - alphas_cumprod), dtype=tf.float32)
+        self.sqrt_recip_alphas_cumprod = tf.constant(np.sqrt(1. / alphas_cumprod), dtype=tf.float32)
+        self.sqrt_recipm1_alphas_cumprod = tf.constant(np.sqrt(1.0 / alphas_cumprod - 1), dtype=tf.float32)
+        self.posterior_variance = (betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod))
+
+        self.posterior_log_variance_clipped = tf.constant(
+            np.log(np.maximum(self.posterior_variance, 1e-20)), dtype=tf.float32
+        )
+
+        self.posterior_mean_coef1 = tf.constant(
+            betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
+            dtype=tf.float32,
+        )
+
+        self.posterior_mean_coef2 = tf.constant(
+            (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod),
+            dtype=tf.float32,
+        )
+
+    def _extract(self, a, t, x_shape):
+
+        batch_size = x_shape[0]
+        out = tf.gather(a, t)
+        return tf.reshape(out, [batch_size, 1, 1])
+
+    def q_sample(self, x_start, t):
+
+        x_start_shape = tf.shape(x_start)
+        samp = self._extract(self.sqrt_alphas_cumprod, t, x_start_shape) * x_start
+        noise = tf.random.normal(shape=tf.shape(x_start), dtype='float32')
+        weight_noise = self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start_shape) * noise * 0.5
+        # diffused_sample = self._extract(self.sqrt_alphas_cumprod, t, x_start_shape) * x_start + self._extract(
+        #     self.sqrt_one_minus_alphas_cumprod, t, x_start_shape) * noise #* 0.1
+        diffused_sample = x_start + weight_noise 
+        diffused_sample = tf.clip_by_value(diffused_sample , -0.99, 0.99)
+        weight_noise = diffused_sample - x_start
+        return samp, weight_noise, diffused_sample
+
+    def predict_start_from_noise(self, x_t, t, noise):
+
+        x_t_shape = tf.shape(x_t)
+        return (
+                self._extract(self.sqrt_recip_alphas_cumprod, t, x_t_shape) * x_t
+                - self._extract(self.sqrt_recipm1_alphas_cumprod, t, x_t_shape) * noise
+        )
+
+    def q_posterior(self, x_start, x_t, t):
+
+        x_t_shape = tf.shape(x_t)
+        posterior_mean = (
+                self._extract(self.posterior_mean_coef1, t, x_t_shape) * x_start
+                + self._extract(self.posterior_mean_coef2, t, x_t_shape) * x_t
+        )
+        posterior_variance = self._extract(self.posterior_variance, t, x_t_shape)
+        posterior_log_variance_clipped = self._extract(
+            self.posterior_log_variance_clipped, t, x_t_shape
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, pred_noise, x, t, clip_denoised=False):
+
+        x_recon = self.predict_start_from_noise(x, t=t, noise=pred_noise)
+
+        if clip_denoised:
+            x_recon = tf.clip_by_value(x_recon, self.clip_min, self.clip_max)
+
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
+            x_start=x_recon, x_t=x, t=t
+        )
+        return model_mean, posterior_variance, posterior_log_variance
+
+    def p_sample(self, pred_noise, x, t, clip_denoised=False):
+
+        model_mean, _, model_log_variance = self.p_mean_variance(
+        pred_noise, x=x, t=t, clip_denoised=clip_denoised
+    )
+        variance_term = tf.exp(0.5 * model_log_variance)
+        noise = tf.random.normal(shape=tf.shape(x), dtype=x.dtype)
+        nonzero_mask = tf.reshape(
+        1 - tf.cast(tf.equal(t, 0), tf.float32), [tf.shape(x)[0], 1, 1]
+    )
+        noise_term = variance_term * nonzero_mask * noise
+        sample = model_mean + noise_term
+        return sample
+
+
+class TimeEmbedding(layers.Layer):
+    def __init__(self, dim, **kwargs):
+        super().__init__(**kwargs)
+        self.dim = dim
+        self.half_dim = dim // 2
+        self.emb = math.log(10000) / (self.half_dim - 1)
+        self.emb = tf.exp(tf.range(self.half_dim, dtype=tf.float32) * -self.emb)
+
+    def call(self, inputs):
+        inputs = tf.cast(inputs, dtype=tf.float32)
+        emb = inputs[:, None] * self.emb[None, :]
+        emb = tf.concat([tf.sin(emb), tf.cos(emb)], axis=-1)
+        return emb
+
+def TimeMLP(units, activation_fn=keras.activations.swish):
+    def apply(inputs):
+        temb = layers.Dense(
+            units, activation=activation_fn, kernel_initializer=kernel_init(1.0)
+        )(inputs)
+        # temb = layers.Dense(units, kernel_initializer=kernel_init(1.0))(temb)
+        return temb
+
+    return apply
+
+# Kernel initializer to use
+def kernel_init(scale):
+    scale = max(scale, 1e-10)
+    return keras.initializers.VarianceScaling(
+        scale, mode="fan_avg", distribution="uniform"
+    )
+
+def build_encoder_time(embed_dim=16, num_heads=2, ff_dim=32):
+    def apply(inputs):
+        x, t = inputs
+        position_embedding_layer = layers.Embedding(x.shape[1], embed_dim)
+        pos_encoding = position_embedding_layer(tf.range(x.shape[1]))
+        embeddings = x + pos_encoding + t
+
+        # Encoder blocks
+        for _ in range(2):  # Repeat twice
+            # Multi-head self-attention mechanism
+            attention_output, attention_score = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)(
+                embeddings, embeddings, return_attention_scores=True)
+
+            # Add residual connection and layer normalization
+            x = layers.Add()([embeddings, attention_output])
+            x = layers.LayerNormalization(epsilon=1e-6)(x)
+
+            # Feed-forward network
+            ff_output = layers.Dense(ff_dim, activation="relu")(x)
+            ff_output = layers.Dense(embed_dim)(ff_output)
+
+            # Add residual connection and layer normalization
+            x = layers.Add()([x, ff_output])
+            x = layers.LayerNormalization(epsilon=1e-6)(x)
+
+        return x, attention_score
+
+    return apply
+
+def build_encoder_variales(embed_dim=16, num_heads=2, ff_dim=32):
+    def apply(inputs):
+        x, t = inputs
+        x = layers.Conv1D(16, kernel_size=3, padding='same')(x)
+        embeddings = x + t
+
+        # Encoder blocks
+        for _ in range(2):  # Repeat twice
+            # Multi-head self-attention mechanism
+            attention_output, attention_score = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)(
+                embeddings, embeddings, return_attention_scores=True)
+
+            # Add residual connection and layer normalization
+            x = layers.Add()([embeddings, attention_output])
+            x = layers.LayerNormalization(epsilon=1e-6)(x)
+
+            # Feed-forward network
+            ff_output = layers.Dense(ff_dim, activation="relu")(x)
+            ff_output = layers.Dense(embed_dim)(ff_output)
+
+            # Add residual connection and layer normalization
+            x = layers.Add()([x, ff_output])
+            x = layers.LayerNormalization(epsilon=1e-6)(x)
+
+        return x, attention_score
+
+    return apply
+
+def build_decoder(embed_dim=16, num_heads=2, ff_dim=32):
+    def apply(inputs):
+        encoder_outputs, t = inputs
+        position_embedding_layer = layers.Embedding(encoder_outputs.shape[1], embed_dim)
+        pos_encoding = position_embedding_layer(tf.range(encoder_outputs.shape[1]))
+        dec_embeddings = encoder_outputs + pos_encoding + t
+
+        # Decoder blocks
+        dec_output = dec_embeddings
+        for _ in range(2):  # Repeat twice
+
+            # Multi-head attention over encoder outputs
+            attention2_output = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)(
+                dec_output, encoder_outputs)
+
+            # Add residual connection and layer normalization
+            dec_output = layers.Add()([dec_output, attention2_output])
+            dec_output = layers.LayerNormalization(epsilon=1e-6)(dec_output)
+
+            # Feed-forward network
+            ff_output = layers.Dense(ff_dim, activation="relu")(dec_output)
+            ff_output = layers.Dense(embed_dim)(ff_output)
+
+            # Add residual connection and layer normalization
+            dec_output = layers.Add()([dec_output, ff_output])
+            dec_output = layers.LayerNormalization(epsilon=1e-6)(dec_output)
+
+        return dec_output
+
+    return apply
+
+def build_model(time_len, fea_num, d_model=16, n_heads=2, encoder_type='dual'):
+    """
+    Build the transformer-based diffusion model.
+    """
+    print(f"\nBuilding model with encoder type: {encoder_type}")
+    print(f"Input shape: time_len={time_len}, features={fea_num}, d_model={d_model}")
+    
+    # Input layers
+    x_input = layers.Input(shape=(time_len, fea_num))
+    time_input = layers.Input(shape=())
+    
+    # Time step embeddings
+    time_emb = get_time_embedding(time_input, d_model)
+    
+    encoded_features = []
+    
+    if encoder_type in ['time', 'dual']:
+        print("→ Using Time Transformer Encoder")
+        # Time Transformer
+        time_encoded = time_transformer_encoder(
+            x_input, 
+            time_emb,
+            d_model=d_model,
+            n_heads=n_heads
+        )
+        print(f"  Time encoder output shape: {time_encoded.shape}")
+        encoded_features.append(time_encoded)
+    
+    if encoder_type in ['pairwise', 'dual']:
+        print("→ Using Pairwise Correlation Encoder")
+        # Pairwise Correlation Transformer
+        pairwise_encoded = pairwise_transformer_encoder(
+            x_input,
+            time_emb,
+            d_model=d_model,
+            n_heads=n_heads
+        )
+        print(f"  Pairwise encoder output shape: {pairwise_encoded.shape}")
+        encoded_features.append(pairwise_encoded)
+    
+    # Combine encodings based on encoder type
+    if encoder_type == 'dual':
+        print("→ Combining both encoders")
+        encoded = layers.Concatenate(axis=-1)(encoded_features)
+        print(f"  Combined shape before projection: {encoded.shape}")
+        encoded = layers.Dense(d_model)(encoded)
+        print(f"  Final encoded shape after projection: {encoded.shape}")
+    else:
+        encoded = encoded_features[0]
+        print(f"→ Using single encoder output shape: {encoded.shape}")
+    
+    # Add residual connection
+    if encoder_type != 'dual':
+        print("→ Adding residual connection")
+        encoded = layers.Add()([encoded, layers.Dense(d_model)(x_input)])
+    
+    # Decoder
+    decoded = decoder_module(encoded, time_emb)
+    print(f"→ Decoder output shape: {decoded.shape}")
+    
+    # Final output layer
+    output = layers.Dense(fea_num)(decoded)
+    print(f"→ Final output shape: {output.shape}")
+    print("Model building completed!\n")
+    
+    return keras.Model(inputs=[x_input, time_input], outputs=output)
+
+def pairwise_transformer_encoder(x, time_emb, d_model=16, n_heads=2):
+    """Pairwise Correlation Transformer encoder implementation"""
+    # Get input shape
+    input_shape = x.shape
+    time_len = input_shape[1]
+    
+    # Transpose input to treat features as sequence
+    x_transposed = tf.transpose(x, perm=[0, 2, 1])  # [batch, features, time]
+    
+    # Project input
+    x_proj = layers.Dense(d_model)(x_transposed)
+    
+    # Expand time embeddings
+    time_emb = tf.expand_dims(time_emb, 1)
+    time_emb = tf.tile(time_emb, [1, tf.shape(x_proj)[1], 1])
+    
+    # Add time embeddings
+    x = x_proj + time_emb
+    
+    # Transformer encoder layers
+    for _ in range(2):
+        x = transformer_encoder_layer(x, d_model, n_heads)
+    
+    # Project to correct time dimension and transpose back
+    x = layers.Dense(time_len)(x)  # Project to original time dimension
+    x = tf.transpose(x, perm=[0, 2, 1])  # [batch, time, features]
+    
+    # Final projection to match d_model dimension
+    x = layers.Dense(d_model)(x)  # [batch, time_len, d_model]
+    
+    return x
+
+def time_transformer_encoder(x, time_emb, d_model=16, n_heads=2):
+    """Time Transformer encoder implementation"""
+    # Position embeddings
+    pos_emb = get_positional_embedding(tf.shape(x)[1], d_model)
+    
+    # Project input
+    x = layers.Dense(d_model)(x)
+    
+    # Expand time embeddings
+    time_emb = tf.expand_dims(time_emb, 1)
+    time_emb = tf.tile(time_emb, [1, tf.shape(x)[1], 1])
+    
+    # Add embeddings
+    x = x + pos_emb + time_emb
+    
+    # Transformer encoder layers
+    for _ in range(2):
+        x = transformer_encoder_layer(x, d_model, n_heads)
+    
+    return x  # [batch, time_len, d_model]
+
+def transformer_encoder_layer(x, d_model, n_heads):
+    """Single transformer encoder layer with multi-head attention"""
+    # Multi-head self attention
+    attention_output = layers.MultiHeadAttention(
+        num_heads=n_heads,
+        key_dim=d_model // n_heads
+    )(x, x)
+    
+    # Add & Norm
+    x = layers.Add()([x, attention_output])
+    x = layers.LayerNormalization()(x)
+    
+    # Feed Forward Network
+    ffn_output = layers.Dense(d_model * 4, activation='relu')(x)
+    ffn_output = layers.Dense(d_model)(ffn_output)
+    
+    # Add & Norm
+    x = layers.Add()([x, ffn_output])
+    x = layers.LayerNormalization()(x)
+    
+    return x
+
+def decoder_module(encoded, time_emb):
+    """Decoder implementation"""
+    # Expand time embeddings
+    time_emb = tf.expand_dims(time_emb, 1)
+    time_emb = tf.tile(time_emb, [1, tf.shape(encoded)[1], 1])
+    
+    # Concatenate along feature dimension
+    x = layers.Concatenate(axis=-1)([encoded, time_emb])
+    
+    # Decoder layers
+    x = layers.Dense(256, activation='relu')(x)
+    x = layers.Dense(128, activation='relu')(x)
+    
+    return x
+
+def get_time_embedding(timesteps, embedding_dim):
+    """
+    Create sinusoidal time embeddings.
+    
+    Args:
+        timesteps: Tensor of shape [batch_size] containing timesteps
+        embedding_dim: Dimension of the embeddings to create
+    """
+    # Ensure timesteps is a 2D tensor
+    timesteps = tf.expand_dims(timesteps, -1)
+    
+    # Calculate positions and dimensions
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = tf.exp(tf.range(half_dim, dtype=tf.float32) * -emb)
+    
+    # Create embeddings
+    emb = tf.cast(timesteps, dtype=tf.float32) * emb[None, :]
+    emb = tf.concat([tf.sin(emb), tf.cos(emb)], axis=-1)
+    
+    # Handle odd embedding dimensions
+    if embedding_dim % 2 == 1:
+        emb = tf.pad(emb, [[0, 0], [0, 1]])
+    
+    return emb  # Shape: [batch_size, embedding_dim]
+
+def get_positional_embedding(sequence_length, embedding_dim):
+    """
+    Create sinusoidal position embeddings.
+    
+    Args:
+        sequence_length: Length of the sequence
+        embedding_dim: Dimension of the embeddings to create
+    """
+    # Create position indices
+    positions = tf.range(sequence_length, dtype=tf.float32)[:, tf.newaxis]
+    
+    # Create dimension indices
+    dimensions = tf.range(0, embedding_dim, 2, dtype=tf.float32)[tf.newaxis, :]
+    
+    # Calculate angle rates
+    angle_rates = 1 / tf.pow(10000.0, (2 * dimensions) / tf.cast(embedding_dim, tf.float32))
+    
+    # Calculate angle rads
+    angle_rads = positions * angle_rates
+    
+    # Apply sin and cos
+    pos_encoding = tf.concat(
+        [tf.sin(angle_rads), tf.cos(angle_rads)],
+        axis=-1
+    )
+    
+    # Handle odd embedding dimensions
+    if embedding_dim % 2 == 1:
+        pos_encoding = tf.pad(pos_encoding, [[0, 0], [0, 1]])
+    
+    # Add batch dimension
+    pos_encoding = tf.expand_dims(pos_encoding, 0)
+    
+    return pos_encoding
+
+class DiffusionModel(keras.Model):
+    def __init__(self, network, ema_network, timesteps, gdf_util, data, ema=0.999):
+        super().__init__()
+        self.network = network
+        self.ema_network = ema_network
+        self.timesteps = timesteps
+        self.gdf_util = gdf_util
+        self.data = data 
+        self.ema = ema
+        
+    def train_step(self, data):
+
+        batch_size = tf.shape(data)[0]
+
+        t = tf.random.uniform(
+            minval=0, 
+            maxval=self.timesteps, 
+            shape=(batch_size,), 
+            dtype=tf.int32
+        )
+
+        old_weights = [tf.identity(w) for w in self.network.trainable_weights]
+        
+        with tf.GradientTape() as tape:
+            _, noise, x_t = self.gdf_util.q_sample(data, t)
+            pred_noise = self.network([x_t, t], training=True)
+            loss = self.loss(noise, pred_noise)
+        
+        gradients = tape.gradient(loss, self.network.trainable_weights)
+        self.optimizer.apply_gradients(zip(gradients, self.network.trainable_weights))
+
+        for weight, ema_weight in zip(self.network.weights, self.ema_network.weights):
+            ema_weight.assign(self.ema * ema_weight + (1 - self.ema) * weight)
+
+        new_weights = self.network.trainable_weights
+        weight_changes = []
+        for old_w, new_w in zip(old_weights, new_weights):
+            diff = tf.reduce_max(tf.abs(old_w - new_w))
+            weight_changes.append(diff)
+        max_change = tf.reduce_max(weight_changes)
+
+        return {
+        "loss": loss,
+        "weight_max_change": max_change,
+        "has_weight_changed": max_change > 0
+    }
+
+    def check_noise_levels(self):
+
+        x_0 = tf.cast(self.data, tf.float32)  
+        
+        print("\n=== Noise Level Analysis ===")
+        print(f"Using all {len(self.data)} samples")
+        print("Checking noise levels at different timesteps:")
+
+        timesteps_to_check = [
+            0,  
+            self.timesteps//4,  
+            self.timesteps//2,  
+            3*self.timesteps//4,  
+            self.timesteps-1 
+        ]
+        
+        for t in timesteps_to_check:
+
+            sqrt_alphas = tf.sqrt(self.gdf_util.alphas_cumprod[t])
+            sqrt_one_minus_alphas = tf.sqrt(1 - self.gdf_util.alphas_cumprod[t])
+
+            _, noise, x_t = self.gdf_util.q_sample(x_0, tf.fill([len(x_0)], t))
+            pred_noise = self.ema_network.predict([x_t, tf.fill([len(x_0)], t)], verbose=0)
+
+            print(f"\nTimestep {t}:")
+            print(f"Signal scaling factor (sqrt_alphas): {sqrt_alphas:.4f}")
+            print(f"Noise scaling factor (sqrt_1-alphas): {sqrt_one_minus_alphas:.4f}")
+            print(f"Original data range: [{tf.reduce_min(x_0):.4f}, {tf.reduce_max(x_0):.4f}]")
+            print(f"Noisy data range: [{tf.reduce_min(x_t):.4f}, {tf.reduce_max(x_t):.4f}]")
+            print(f"Added noise - Mean: {tf.reduce_mean(noise):.4f}, Std: {tf.math.reduce_std(noise):.4f}")
+            print(f"Scaled noise - Mean: {tf.reduce_mean(sqrt_one_minus_alphas * noise):.4f}, Std: {tf.math.reduce_std(sqrt_one_minus_alphas * noise):.4f}")
+            print(f"Predicted noise - Mean: {tf.reduce_mean(pred_noise):.4f}, Std: {tf.math.reduce_std(pred_noise):.4f}")
+            print(f"Noise prediction error: {tf.reduce_mean(tf.abs(noise - pred_noise)):.4f}")
+
+    def generate_ts(self, num_ts=16):
+       
+        if num_ts > len(self.data):
+            indices = tf.random.uniform(
+                shape=[num_ts],
+                minval=0,
+                maxval=len(self.data),
+                dtype=tf.int32
+            )
+            initial_samples = tf.cast(
+            tf.gather(self.data, indices),
+            tf.float32
+        )
+        else:
+            initial_samples = self.data
+
+        _, _, samples = self.gdf_util.q_sample(initial_samples,  tf.fill([num_ts], self.timesteps-1))
+        samples0 = samples 
+
+        for i in reversed(range(0, self.timesteps)):
+            tt = tf.fill([num_ts], i)
+            pred_noise = self.ema_network.predict([samples0, tt], verbose=0, batch_size=num_ts
+            )
+            # print(f"\nStep {i}:")
+            # print(f"Predicted noise - Mean: {tf.reduce_mean(pred_noise):.4f}")
+            # print(f"Predicted noise - Std: {tf.math.reduce_std(pred_noise):.4f}")
+            # print(f"Predicted noise - Min: {tf.reduce_min(pred_noise):.4f}")
+            # print(f"Predicted noise - Max: {tf.reduce_max(pred_noise):.4f}")
+
+            samples = self.gdf_util.p_sample(
+                pred_noise, samples0, tt, clip_denoised=False
+            )
+            # scaler = TSFeatureScaler()
+            # print(f"Generated samples - Mean: {tf.reduce_mean(scaler.fit_transform(samples.numpy())):.4f}")
+            # print(f"Generated samples- Std: {tf.math.reduce_std(scaler.fit_transform(samples.numpy())):.4f}")
+            # print(f"Generated samples - Min: {tf.reduce_min(scaler.fit_transform(samples.numpy())):.4f}")
+            # print(f"Generated samples - Max: {tf.reduce_max(scaler.fit_transform(samples.numpy())):.4f}")
+        
+        return samples
+
+    def plot_denoise_process(self, save_dir='denoise_process'):
+
+        if os.path.exists(save_dir):
+            shutil.rmtree(save_dir)
+        os.makedirs(save_dir, exist_ok=True)
+        
+        indices = tf.random.uniform(
+        shape=[16], 
+        minval=0, 
+        maxval=len(self.data), 
+        dtype=tf.int32
+        )
+        x_0 = tf.cast(tf.gather(self.data, indices), tf.float32) 
+        _, _, samples = self.gdf_util.q_sample(x_0,  tf.fill([16], self.timesteps-1))
+        samples0 = samples 
+
+        time_steps = np.arange(self.data.shape[1])
+
+        self._plot_step_grid(x_0.numpy(), self.timesteps, time_steps, save_dir)
+        print(f"Denoising: Generate 16 samples for visualization")
+        
+        for i in reversed(range(0, self.timesteps)):
+            print(f"Processing step {i}/{self.timesteps}")
+            tt = tf.fill([16], i)
+            pred_noise = self.ema_network.predict([samples0, tt], verbose=0, batch_size=1
+            )
+            samples = self.gdf_util.p_sample(
+                pred_noise, samples0, tt, clip_denoised=False
+            )
+            scaler = TSFeatureScaler()
+            scaled_samples = scaler.fit_transform(samples.numpy())
+            self._plot_step_grid(scaled_samples, i, time_steps, save_dir)
+
+        print(f"Saved {self.timesteps + 1} plots to {save_dir}/")
+
+    def _plot_step_grid(self, samples, step, time_steps, save_dir):
+
+        fig, axes = plt.subplots(4, 4, figsize=(20, 20))
+        fig.suptitle(f'Generated Samples at Step {step}', fontsize=16)
+
+        for idx in range(16):
+            row = idx // 4
+            col = idx % 4
+
+            for feature_idx in range(samples.shape[-1]):
+                axes[row, col].plot(
+                time_steps, 
+                samples[idx, :, feature_idx],
+                label=f'Feature {feature_idx+1}',
+                alpha=0.8
+            )
+        
+            axes[row, col].set_title(f'Sample {idx+1}')
+            axes[row, col].grid(True)
+            if idx % 4 == 0:  
+                axes[row, col].set_ylabel('Value')
+            if idx >= 12:  
+                axes[row, col].set_xlabel('Time Steps')
+            if idx == 15:  
+                axes[row, col].legend()
+    
+        plt.tight_layout()
+        plt.savefig(os.path.join(save_dir, f'step_{step:04d}.png'), 
+                dpi=300, bbox_inches='tight')
+        plt.close()
+
+    def plot_denoise_detailed_process(self, save_dir='denoise_detailed_process'):
+
+        if os.path.exists(save_dir):
+            shutil.rmtree(save_dir)
+        os.makedirs(save_dir, exist_ok=True)
+
+        indices = tf.random.uniform(
+        shape=[8], 
+        minval=0, 
+        maxval=len(self.data), 
+        dtype=tf.int32
+    )
+        x_0 = tf.cast(tf.gather(self.data, indices), tf.float32)
+        _, _, samples = self.gdf_util.q_sample(x_0, tf.fill([8], self.timesteps-1))
+        samples0 = samples 
+    
+        self._plot_step_comparison(
+        x_0.numpy(),
+        x_0.numpy(),  
+        -1,  
+        time_steps=np.arange(self.data.shape[1]),
+        save_dir=save_dir,
+        title="Original Data"
+    )
+    
+        self._plot_step_comparison(
+        samples.numpy(),
+        samples.numpy(), 
+        self.timesteps,
+        time_steps=np.arange(self.data.shape[1]),
+        save_dir=save_dir,
+        title="Initial Noisy Data"
+    )
+    
+        for i in reversed(range(0, self.timesteps)):
+            print(f"Processing step {i}/{self.timesteps}")
+            tt = tf.fill([8], i)
+            pred_noise = self.ema_network.predict([samples0, tt], verbose=0, batch_size=8)
+            samples = self.gdf_util.p_sample(pred_noise, samples0, tt, clip_denoised=False)
+            scaler = TSFeatureScaler()
+
+            self._plot_step_comparison(
+            pred_noise,
+            scaler.fit_transform(samples.numpy()),
+            i,
+            time_steps=np.arange(self.data.shape[1]),
+            save_dir=save_dir,
+            title=f"Denoising Step {i}"
+        )
+
+        print(f"Visualization completed. Check {save_dir}/ for results.")
+
+    def _plot_step_comparison(self, noise_data, sample_data, step, time_steps, save_dir, title):
+
+        fig, axes = plt.subplots(2, 8, figsize=(24, 8))
+        fig.suptitle(f'{title}', fontsize=16)
+    
+        for idx in range(8):
+            ax = axes[0, idx]
+            for feature_idx in range(noise_data.shape[-1]):
+                ax.plot(
+                time_steps, 
+                noise_data[idx, :, feature_idx],
+                label=f'Feature {feature_idx+1}',
+                alpha=0.8
+            )
+            ax.set_title(f'Noise {idx+1}')
+            ax.grid(True)
+            if idx == 0: 
+                ax.set_ylabel('Value')
+            if idx == 7:  
+                ax.legend()
+
+        for idx in range(8):
+            ax = axes[1, idx]
+            for feature_idx in range(sample_data.shape[-1]):
+                ax.plot(
+                time_steps, 
+                sample_data[idx, :, feature_idx],
+                label=f'Feature {feature_idx+1}',
+                alpha=0.8
+            )
+            ax.set_title(f'Sample {idx+1}')
+            ax.grid(True)
+            if idx == 0:  
+                ax.set_ylabel('Value')
+            ax.set_xlabel('Time Steps')
+    
+        plt.tight_layout()
+        plt.savefig(
+        os.path.join(save_dir, f'step_{step:04d}.png'), 
+        dpi=300,
+        bbox_inches='tight'
+    )
+        plt.close()
+
+    def plot_noise_process(self, save_dir='noise_process'): 
+
+        if os.path.exists(save_dir):
+            shutil.rmtree(save_dir)
+        os.makedirs(save_dir, exist_ok=True)
+
+        indices = tf.random.uniform(
+        shape=[16], 
+        minval=0, 
+        maxval=len(self.data), 
+        dtype=tf.int32
+        )
+        x_start = tf.cast(tf.gather(self.data, indices), tf.float32)
+        print(f"Noising: Select 16 samples for visualization")
+    
+        for t in range(self.timesteps):
+            print(f"Processing step {t}/{self.timesteps}")
+            _, _, x_noisy = self.gdf_util.q_sample(
+            x_start, 
+            tf.fill([16], t))
+        
+            fig, axes = plt.subplots(4, 4, figsize=(20, 20))
+            fig.suptitle(f'Noise Process at Step {t}', fontsize=16)
+        
+            time_steps = np.arange(x_start.shape[1])
+
+            for idx in range(16):
+                row = idx // 4
+                col = idx % 4
+            
+                for feature_idx in range(x_noisy.shape[-1]):
+                    axes[row, col].plot(
+                    time_steps, 
+                    x_noisy[idx, :, feature_idx],
+                    label=f'Feature {feature_idx+1}',
+                    alpha=0.8
+                )
+            
+                axes[row, col].set_title(f'Sample {idx+1}')
+                axes[row, col].grid(True)
+                if idx % 4 == 0:  
+                    axes[row, col].set_ylabel('Value')
+                if idx >= 12:  
+                    axes[row, col].set_xlabel('Time Steps')
+                if idx == 15:  
+                    axes[row, col].legend()
+        
+            plt.tight_layout()
+            plt.savefig(os.path.join(save_dir, f'step_{t:04d}.png'), 
+                   dpi=300, bbox_inches='tight')
+            plt.close()
+    
+        print(f"Saved {self.timesteps} plots to {save_dir}/")
+
+    def plot_noise_process_app(self, num_samples=16):
+        """为应用程序创建加噪过程的动态可视化"""
+        # 验证num_samples是否为有效值
+        valid_sizes = [4, 9, 16, 25]
+        if num_samples not in valid_sizes:
+            raise ValueError(f"num_samples must be one of {valid_sizes}")
+        
+        # 计算网格大小
+        grid_size = int(np.sqrt(num_samples))
+        
+        # 随机选择样本
+        indices = tf.random.uniform(
+            shape=[num_samples],
+            minval=0,
+            maxval=len(self.data),
+            dtype=tf.int32
+        )
+        x_start = tf.cast(tf.gather(self.data, indices), tf.float32)
+        
+        # 存储每一步的图像
+        frames = []
+        time_steps = np.arange(x_start.shape[1])
+        
+        # 逐步添加噪声
+        for t in range(0, self.timesteps, max(1, self.timesteps // 10)):
+            # 添加噪声
+            _, _, x_noisy = self.gdf_util.q_sample(
+                x_start,
+                tf.fill([num_samples], t)
+            )
+            
+            # 创建图形
+            fig, axes = plt.subplots(grid_size, grid_size, figsize=(5*grid_size, 5*grid_size))
+            fig.suptitle(f'Noise Process at Step {t}', fontsize=16)
+
+            colors = ['#0a9396', '#ee9b00', '#9b2226']
+            
+            # 确保axes是二维数组
+            if grid_size == 2:
+                axes = axes.reshape(2, 2)
+            
+            # 绘制每个样本
+            for idx in range(num_samples):
+                row = idx // grid_size
+                col = idx % grid_size
+                
+                for feature_idx in range(x_noisy.shape[-1]):
+                    axes[row, col].plot(
+                        time_steps,
+                        x_noisy[idx, :, feature_idx],
+                        label=f'Feature {feature_idx+1}',
+                        color=colors[feature_idx], 
+                        alpha=0.8
+                    )
+                
+                axes[row, col].set_title(f'Sample {idx+1}')
+                axes[row, col].grid(True)
+                if idx % grid_size == 0:
+                    axes[row, col].set_ylabel('Value')
+                if idx >= num_samples - grid_size:
+                    axes[row, col].set_xlabel('Time Steps')
+                if idx == num_samples - 1:
+                    axes[row, col].legend()
+            
+            plt.tight_layout()
+            
+            # 将图形转换为图像
+            fig.canvas.draw()
+            image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+            frames.append(image)
+            
+            plt.close()
+        
+        return frames
+
+    def plot_denoise_process_app(self, num_samples=16):
+        """为应用程序创建去噪过程的动态可视化"""
+        # 验证num_samples是否为有效值
+        valid_sizes = [4, 9, 16, 25]
+        if num_samples not in valid_sizes:
+            raise ValueError(f"num_samples must be one of {valid_sizes}")
+        
+        # 计算网格大小
+        grid_size = int(np.sqrt(num_samples))
+        
+        # 随机选择样本
+        indices = tf.random.uniform(
+            shape=[num_samples],
+            minval=0,
+            maxval=len(self.data),
+            dtype=tf.int32
+        )
+        x_0 = tf.cast(tf.gather(self.data, indices), tf.float32)
+        _, _, samples = self.gdf_util.q_sample(x_0, tf.fill([num_samples], self.timesteps-1))
+        samples0 = samples
+        
+        # 存储每一步的图像
+        frames = []
+        time_steps = np.arange(self.data.shape[1])
+        
+        # 添加原始数据的图像
+        frames.append(self._plot_step_grid_app(x_0.numpy(), self.timesteps, time_steps, grid_size))
+        
+        # 逐步去噪
+        for i in reversed(range(0, self.timesteps)):
+            if i % max(1, self.timesteps // 10) == 0:
+                tt = tf.fill([num_samples], i)
+                pred_noise = self.ema_network.predict([samples0, tt], verbose=0, batch_size=num_samples)
+                samples = self.gdf_util.p_sample(pred_noise, samples0, tt, clip_denoised=False)
+                scaler = TSFeatureScaler()
+                scaled_samples = scaler.fit_transform(samples.numpy())
+                frames.append(self._plot_step_grid_app(scaled_samples, i, time_steps, grid_size))
+        
+        return frames
+
+    def _plot_step_grid_app(self, samples, step, time_steps, grid_size):
+        """辅助函数：为应用程序创建单个时间步的网格图"""
+        fig, axes = plt.subplots(grid_size, grid_size, figsize=(5*grid_size, 5*grid_size))
+        fig.suptitle(f'Generated Samples at Step {step}', fontsize=16)
+
+        colors = ['#0a9396', '#ee9b00', '#9b2226']
+        
+        # 确保axes是二维数组
+        if grid_size == 2:
+            axes = axes.reshape(2, 2)
+        
+        # 绘制每个样本
+        num_samples = grid_size * grid_size
+        for idx in range(num_samples):
+            row = idx // grid_size
+            col = idx % grid_size
+            
+            for feature_idx in range(samples.shape[-1]):
+                axes[row, col].plot(
+                    time_steps,
+                    samples[idx, :, feature_idx],
+                    label=f'Feature {feature_idx+1}',
+                    color=colors[feature_idx], 
+                    alpha=0.8
+                )
+            
+            axes[row, col].set_title(f'Sample {idx+1}')
+            axes[row, col].grid(True)
+            if idx % grid_size == 0:
+                axes[row, col].set_ylabel('Value')
+            if idx >= num_samples - grid_size:
+                axes[row, col].set_xlabel('Time Steps')
+            if idx == num_samples - 1:
+                axes[row, col].legend()
+        
+        plt.tight_layout()
+        
+        # 将图形转换为图像
+        fig.canvas.draw()
+        image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        
+        plt.close()
+        return image
+

GAN/timegan.py

@@ -0,0 +1,601 @@
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.python.types.core import TensorLike
+import numpy as np
+import numpy.typing as npt
+from tqdm import tqdm, trange
+from collections import OrderedDict
+import typing as T
+
+import logging
+
+from GAN.zoo import BasicRecurrentArchitecture
+
+logger = logging.getLogger("models")
+logger.setLevel(logging.DEBUG)
+
+
+class LossTracker(OrderedDict):
+    """
+    Dictionary of lists, extends python OrderedDict.
+    Example: Given {'loss_a': [1], 'loss_b': [2]}, adding key='loss_a' with value=0.7
+            gives {'loss_a': [1, 0.7], 'loss_b': [2]}, and adding key='loss_c' with value=1.2
+            gives {'loss_a': [1, 0.7], 'loss_b': [2], 'loss_c': [1.2]}
+    """
+
+    def __setitem__(self, key: T.Any, value: T.Any) -> None:
+        try:
+            # Assumes the key already exists
+            # and the value is a list [oldest_value, another_old, ...]
+            # key -> [oldest_value, another_old, ..., new_value]
+            self[key].append(value)
+        # If there is no key, add key -> [new_value]
+        except KeyError:
+            # key -> [new_value]
+            super(LossTracker, self).__setitem__(key, [value])
+
+    def to_numpy(self) -> npt.NDArray:
+        """
+        :return 2d vector of losses
+        """
+        _losses = np.array([np.array(v) for v in self.values() if isinstance(v, list)])
+        return _losses
+
+    def labels(self) -> T.List:
+        """
+        :return list of keys
+        """
+        return list(self.keys())
+
+
+
+class TimeGAN(keras.Model):
+    """
+    Time-series Generative Adversarial Networks (TimeGAN)
+
+    Reference: Jinsung Yoon, Daniel Jarrett, Mihaela van der Schaar,
+    "Time-series Generative Adversarial Networks,"
+    Neural Information Processing Systems (NeurIPS), 2019.
+
+    Paper link: https://papers.nips.cc/paper/8789-time-series-generative-adversarial-networks
+    """
+
+    def __init__(
+        self,
+        seq_len: int = 24,
+        module: str = "gru",
+        hidden_dim: int = 24,
+        n_features: int = 6,
+        n_layers: int = 3,
+        batch_size: int = 256,
+        gamma: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.seq_len = seq_len
+        self.hidden_dim = hidden_dim
+        self.dim = n_features
+
+        assert module in ["gru", "lstm", "lstmLN"]
+        self.module = module
+
+        self.n_layers = n_layers
+
+        self.batch_size = batch_size
+
+        self.gamma = gamma
+
+        # ----------------------------
+        # Basic Architectures
+        # ----------------------------
+        self.embedder = BasicRecurrentArchitecture(
+            hidden_dim=self.hidden_dim,
+            output_dim=self.hidden_dim,
+            n_layers=self.n_layers,
+            network_type=self.module,
+            name="Embedder",
+        ).build()
+
+        self.recovery = BasicRecurrentArchitecture(
+            hidden_dim=self.hidden_dim,
+            output_dim=self.dim,
+            n_layers=self.n_layers,
+            network_type=self.module,
+            name="Recovery",
+        ).build()
+
+        self.supervisor = BasicRecurrentArchitecture(
+            hidden_dim=self.hidden_dim,
+            output_dim=self.hidden_dim,
+            n_layers=self.n_layers,
+            network_type=self.module,
+            name="Supervisor",
+        ).build()
+
+        self.discriminator = BasicRecurrentArchitecture(
+            hidden_dim=self.hidden_dim,
+            output_dim=1,
+            n_layers=self.n_layers,
+            network_type=self.module,
+            name="Discriminator",
+        ).build()
+
+        self.generator_aux = BasicRecurrentArchitecture(
+            hidden_dim=self.hidden_dim,
+            output_dim=self.hidden_dim,
+            n_layers=self.n_layers,
+            network_type=self.module,
+            name="Generator",
+        ).build()
+
+        # ----------------------------
+        # Optimizers: call .compile() to set them
+        # ----------------------------
+        self.autoencoder_opt = keras.optimizers.Adam()
+        self.adversarialsup_opt = keras.optimizers.Adam()
+        self.generator_opt = keras.optimizers.Adam()
+        self.embedder_opt = keras.optimizers.Adam()
+        self.discriminator_opt = keras.optimizers.Adam()
+        # ----------------------------
+        # Loss functions: call .compile() to set them
+        # ----------------------------
+        self._mse = keras.losses.MeanSquaredError()
+        self._bce = keras.losses.BinaryCrossentropy()
+
+        # --------------------------
+        # All losses: will be populated in .fit()
+        # --------------------------
+        self.training_losses_history = LossTracker()
+
+        # --------------------------
+        # Synthetic data generation during training: will be populated in .fit()
+        # --------------------------
+        self.synthetic_data_generated_in_training = dict()
+
+    def compile(
+        self,
+        d_optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(), # keras.optimizers.legacy.Adam()
+        g_optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(),
+        emb_optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(),
+        supgan_optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(),
+        ae_optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam(),
+        emb_loss: keras.losses.Loss = keras.losses.MeanSquaredError(),
+        clf_loss: keras.losses.Loss = keras.losses.BinaryCrossentropy(),
+    ) -> None:
+        """
+        Assign optimizers and loss functions.
+
+        :param d_optimizer: An optimizer for the GAN's discriminator
+        :param g_optimizer: An optimizer for the GAN's generator
+        :param emb_optimizer: An optimizer for the GAN's embedder
+        :param supgan_optimizer: An optimizer for the adversarial supervised network
+        :param ae_optimizer: An optimizer for the autoencoder network
+        :param emb_loss: A loss function for the embedding recovery
+        :param clf_loss: A loss function for the discriminator task
+        :return: None
+        """
+        # ----------------------------
+        # Optimizers
+        # ----------------------------
+        self.autoencoder_opt = ae_optimizer
+        self.adversarialsup_opt = supgan_optimizer
+        self.generator_opt = g_optimizer
+        self.embedder_opt = emb_optimizer
+        self.discriminator_opt = d_optimizer
+        # ----------------------------
+        # Loss functions
+        # ----------------------------
+        self._mse = emb_loss
+        self._bce = clf_loss
+
+    def _define_timegan(self) -> None:
+        # --------------------------------
+        # Data and Noise Inputs
+        # --------------------------------
+        X = keras.layers.Input(
+            shape=[self.seq_len, self.dim], batch_size=self.batch_size, name="RealData"
+        )
+
+        Z = keras.layers.Input(
+            shape=[self.seq_len, self.dim],
+            batch_size=self.batch_size,
+            name="RandomNoise",
+        )
+
+        # --------------------------------
+        # Autoencoder: Embedder + Recovery
+        # --------------------------------
+        H = self.embedder(X)
+        X_tilde = self.recovery(H)
+
+        self.autoencoder = keras.models.Model(
+            inputs=X, outputs=X_tilde, name="Autoencoder"
+        )
+        self.autoencoder.summary()
+
+        # ---------------------------------
+        # Adversarial Supervised
+        # ---------------------------------
+        E_Hat = self.generator_aux(Z)
+        H_hat = self.supervisor(E_Hat)
+        Y_fake = self.discriminator(H_hat)
+
+        self.adversarial_supervised = keras.models.Model(
+            inputs=Z, outputs=Y_fake, name="AdversarialSupervised"
+        )
+        self.adversarial_supervised.summary()
+
+        # ---------------------------------
+        # Adversarial embedded in latent space
+        # ---------------------------------
+        Y_fake_e = self.discriminator(E_Hat)
+
+        self.adversarial_embedded = keras.models.Model(
+            inputs=Z, outputs=Y_fake_e, name="AdversarialEmbedded"
+        )
+        self.adversarial_embedded.summary()
+
+        # ---------------------------------
+        # Synthetic data generator
+        # ---------------------------------
+        X_hat = self.recovery(H_hat)
+        self.generator = keras.models.Model(
+            inputs=Z, outputs=X_hat, name="FinalGenerator"
+        )
+        self.generator.summary()
+
+        # --------------------------------
+        # Discriminator
+        # --------------------------------
+        Y_real = self.discriminator(H)
+        self.discriminator_model = keras.models.Model(
+            inputs=X, outputs=Y_real, name="FinalDiscriminator"
+        )
+        self.discriminator_model.summary()
+
+    @tf.function
+    def _train_autoencoder(
+        self, X: TensorLike, optimizer: keras.optimizers.Optimizer
+    ) -> float:
+        """
+        1. Embedding network training: minimize E_loss0
+        """
+        with tf.GradientTape() as tape:
+            X_tilde = self.autoencoder(X)
+            E_loss_T0 = self._mse(X, X_tilde)
+            E_loss0 = 10.0 * tf.sqrt(E_loss_T0)
+
+        e_vars = self.embedder.trainable_variables
+        r_vars = self.recovery.trainable_variables
+        all_trainable = e_vars + r_vars
+
+        gradients = tape.gradient(E_loss0, all_trainable)
+        optimizer.apply_gradients(zip(gradients, all_trainable))
+        return E_loss0
+
+    @tf.function
+    def _train_supervisor(
+        self, X: TensorLike, optimizer: keras.optimizers.Optimizer
+    ) -> float:
+        """
+        2. Training with supervised loss only: minimize G_loss_S
+        """
+        with tf.GradientTape() as tape:
+            H = self.embedder(X)
+            H_hat_supervised = self.supervisor(H)
+            G_loss_S = self._mse(H[:, 1:, :], H_hat_supervised[:, :-1, :])
+
+        g_vars = self.generator.trainable_variables
+        s_vars = self.supervisor.trainable_variables
+        all_trainable = g_vars + s_vars
+        gradients = tape.gradient(G_loss_S, all_trainable)
+        apply_grads = [
+            (grad, var)
+            for (grad, var) in zip(gradients, all_trainable)
+            if grad is not None
+        ]
+        optimizer.apply_gradients(apply_grads)
+        return G_loss_S
+
+    @tf.function
+    def _train_generator(
+        self, X: TensorLike, Z: TensorLike, optimizer: keras.optimizers.Optimizer
+    ) -> T.Tuple[float, float, float, float, float]:
+        """
+        3. Joint training (Generator training twice more than discriminator training): minimize G_loss
+        """
+        with tf.GradientTape() as tape:
+            # 1. Adversarial loss
+            Y_fake = self.adversarial_supervised(Z)
+            G_loss_U = self._bce(y_true=tf.ones_like(Y_fake), y_pred=Y_fake)
+
+            Y_fake_e = self.adversarial_embedded(Z)
+            G_loss_U_e = self._bce(y_true=tf.ones_like(Y_fake_e), y_pred=Y_fake_e)
+            # 2. Supervised loss
+            H = self.embedder(X)
+            H_hat_supervised = self.supervisor(H)
+            G_loss_S = self._mse(H[:, 1:, :], H_hat_supervised[:, :-1, :])
+
+            # 3. Two Moments
+            X_hat = self.generator(Z)
+            G_loss_V = self._compute_generator_moments_loss(X, X_hat)
+
+            # 4. Summation
+            G_loss = (
+                G_loss_U
+                + self.gamma * G_loss_U_e
+                + 100 * tf.sqrt(G_loss_S)
+                + 100 * G_loss_V
+            )
+
+        g_vars = self.generator_aux.trainable_variables
+        s_vars = self.supervisor.trainable_variables
+        all_trainable = g_vars + s_vars
+        gradients = tape.gradient(G_loss, all_trainable)
+        apply_grads = [
+            (grad, var)
+            for (grad, var) in zip(gradients, all_trainable)
+            if grad is not None
+        ]
+        optimizer.apply_gradients(apply_grads)
+        return G_loss_U, G_loss_U_e, G_loss_S, G_loss_V, G_loss
+
+    @tf.function
+    def _train_embedder(
+        self, X: TensorLike, optimizer: keras.optimizers.Optimizer
+    ) -> T.Tuple[float, float]:
+        """
+        Train embedder during joint training: minimize E_loss
+        """
+        with tf.GradientTape() as tape:
+            # Supervised Loss
+            H = self.embedder(X)
+            H_hat_supervised = self.supervisor(H)
+            G_loss_S = self._mse(H[:, 1:, :], H_hat_supervised[:, :-1, :])
+
+            # Reconstruction Loss
+            X_tilde = self.autoencoder(X)
+            E_loss_T0 = self._mse(X, X_tilde)
+            E_loss0 = 10 * tf.sqrt(E_loss_T0)
+
+            E_loss = E_loss0 + 0.1 * G_loss_S
+
+        e_vars = self.embedder.trainable_variables
+        r_vars = self.recovery.trainable_variables
+        all_trainable = e_vars + r_vars
+        gradients = tape.gradient(E_loss, all_trainable)
+        optimizer.apply_gradients(zip(gradients, all_trainable))
+        return E_loss, E_loss_T0
+
+    @tf.function
+    def _train_discriminator(
+        self, X: TensorLike, Z: TensorLike, optimizer: keras.optimizers.Optimizer
+    ) -> float:
+        """
+        minimize D_loss
+        """
+        with tf.GradientTape() as tape:
+            D_loss = self._check_discriminator_loss(X, Z)
+
+        d_vars = self.discriminator.trainable_variables
+        gradients = tape.gradient(D_loss, d_vars)
+        optimizer.apply_gradients(zip(gradients, d_vars))
+        return D_loss
+
+    @staticmethod
+    def _compute_generator_moments_loss(
+        y_true: TensorLike, y_pred: TensorLike
+    ) -> float:
+        """
+        :param y_true: TensorLike
+        :param y_pred: TensorLike
+        :return G_loss_V: float
+        """
+        _eps = 1e-6
+        y_true_mean, y_true_var = tf.nn.moments(x=y_true, axes=[0])
+        y_pred_mean, y_pred_var = tf.nn.moments(x=y_pred, axes=[0])
+        # G_loss_V2
+        g_loss_mean = tf.reduce_mean(abs(y_true_mean - y_pred_mean))
+        # G_loss_V1
+        g_loss_var = tf.reduce_mean(
+            abs(tf.sqrt(y_true_var + _eps) - tf.sqrt(y_pred_var + _eps))
+        )
+        # G_loss_V = G_loss_V1 + G_loss_V2
+        return g_loss_mean + g_loss_var
+
+    def _check_discriminator_loss(self, X: TensorLike, Z: TensorLike) -> float:
+        """
+        :param X: TensorLike
+        :param Z: TensorLike
+        :return D_loss: float
+        """
+        # Loss on false negatives
+        Y_real = self.discriminator_model(X)
+        D_loss_real = self._bce(y_true=tf.ones_like(Y_real), y_pred=Y_real)
+
+        # Loss on false positives
+        Y_fake = self.adversarial_supervised(Z)
+        D_loss_fake = self._bce(y_true=tf.zeros_like(Y_fake), y_pred=Y_fake)
+
+        Y_fake_e = self.adversarial_embedded(Z)
+        D_loss_fake_e = self._bce(y_true=tf.zeros_like(Y_fake_e), y_pred=Y_fake_e)
+
+        D_loss = D_loss_real + D_loss_fake + self.gamma * D_loss_fake_e
+        return D_loss
+
+    def _generate_noise(self) -> TensorLike:
+        """
+        Random vector generation
+        :return Z, generated random vector
+        """
+        while True:
+            yield np.random.uniform(low=0, high=1, size=(self.seq_len, self.dim))
+
+    def get_noise_batch(self) -> T.Iterator:
+        """
+        Return an iterator of random noise vectors
+        """
+        return iter(
+            tf.data.Dataset.from_generator(
+                self._generate_noise, output_types=tf.float32
+            )
+            .batch(self.batch_size)
+            .repeat()
+        )
+
+    def _get_data_batch(self, data: TensorLike, n_windows: int) -> T.Iterator:
+        """
+        Return an iterator of shuffled input data
+        """
+        data = tf.convert_to_tensor(data, dtype=tf.float32)
+        return iter(
+            tf.data.Dataset.from_tensor_slices(data)
+            .shuffle(buffer_size=n_windows)
+            .batch(self.batch_size)
+            .repeat()
+        )
+
+    def fit(
+        self,
+        data: T.Union[TensorLike, tf.data.Dataset],
+        epochs: int,
+        checkpoints_interval: T.Optional[int] = None,
+        generate_synthetic: T.Tuple = (),
+        *args,
+        **kwargs,
+    ):
+        """
+        :param data: TensorLike, the training data
+        :param epochs: int, the number of epochs for the training loops
+        :param checkpoints_interval: int, the interval for printing out loss values
+            (loss values will be print out every 'checkpoints_interval' epochs)
+            Default: None (no print out)
+        :param generate_synthetic: list of int, a list of epoch numbers when synthetic data samples are generated
+            Default: [] (no generation)
+        :return None
+        """
+        assert not (
+            self.autoencoder_opt is None
+            or self.adversarialsup_opt is None
+            or self.generator_opt is None
+            or self.embedder_opt is None
+            or self.discriminator_opt is None
+        ), "One of the optimizers is not defined. Please call .compile() to set them"
+        assert not (
+            self._mse is None or self._bce is None
+        ), "One of the loss functions is not defined. Please call .compile() to set them"
+
+        # take tf.data.Dataset | TensorLike
+        if isinstance(data, tf.data.Dataset):
+            batches = iter(data.repeat())
+        else:
+            batches = self._get_data_batch(data, n_windows=len(data))
+
+        # Define the model
+        self._define_timegan()
+
+        # 1. Embedding network training
+        logger.info("Start Embedding Network Training")
+
+        for epoch in tqdm(range(epochs), desc="Autoencoder - training"):
+            X_ = next(batches)
+            step_e_loss_0 = self._train_autoencoder(X_, self.autoencoder_opt)
+
+            # Checkpoint
+            if checkpoints_interval is not None and epoch % checkpoints_interval == 0:
+                logger.info(f"step: {epoch}/{epochs}, e_loss: {step_e_loss_0}")
+            self.training_losses_history["autoencoder"] = float(step_e_loss_0)
+
+        logger.info("Finished Embedding Network Training")
+
+        # 2. Training only with supervised loss
+        logger.info("Start Training with Supervised Loss Only")
+
+        # Adversarial Supervised network training
+        for epoch in tqdm(range(epochs), desc="Adversarial Supervised - training"):
+            X_ = next(batches)
+            step_g_loss_s = self._train_supervisor(X_, self.adversarialsup_opt)
+
+            # Checkpoint
+            if checkpoints_interval is not None and epoch % checkpoints_interval == 0:
+                logger.info(
+                    f"step: {epoch}/{epochs}, s_loss: {np.round(np.sqrt(step_g_loss_s), 4)}"
+                )
+            self.training_losses_history["adversarial_supervised"] = float(
+                np.sqrt(step_g_loss_s)
+            )
+
+        logger.info("Finished Training with Supervised Loss Only")
+
+        # 3. Joint Training
+        logger.info("Start Joint Training")
+
+        # GAN with embedding network training
+        for epoch in tqdm(range(epochs), desc="GAN with embedding - training"):
+            # Generator training (twice more than discriminator training)
+            for kk in range(2):
+                X_ = next(batches)
+                Z_ = next(self.get_noise_batch())
+                # --------------------------
+                # Train the generator
+                # --------------------------
+                (
+                    step_g_loss_u,
+                    step_g_loss_u_e,
+                    step_g_loss_s,
+                    step_g_loss_v,
+                    step_g_loss,
+                ) = self._train_generator(X_, Z_, self.generator_opt)
+
+                # --------------------------
+                # Train the embedder
+                # --------------------------
+                _, step_e_loss_t0 = self._train_embedder(X_, self.embedder_opt)
+
+            X_ = next(batches)
+            Z_ = next(self.get_noise_batch())
+            step_d_loss = self._check_discriminator_loss(X_, Z_)
+            if step_d_loss > 0.15:
+                logger.info(
+                    "Train Discriminator (discriminator does not work well yet)"
+                )
+                step_d_loss = self._train_discriminator(X_, Z_, self.discriminator_opt)
+
+            # Print multiple checkpoints
+            if checkpoints_interval is not None and epoch % checkpoints_interval == 0:
+                logger.info(
+                    f"""step: {epoch}/{epochs},
+                    d_loss: {np.round(step_d_loss, 4)},
+                    g_loss_u: {np.round(step_g_loss_u, 4)},
+                    g_loss_u_e: {np.round(step_g_loss_u_e, 4)},
+                    g_loss_s: {np.round(np.sqrt(step_g_loss_s), 4)},
+                    g_loss_v: {np.round(step_g_loss_v, 4)},
+                    g_loss_v: {np.round(step_g_loss, 4)},
+                    e_loss_t0: {np.round(np.sqrt(step_e_loss_t0), 4)}"""
+                )
+            self.training_losses_history["discriminator"] = float(step_d_loss)
+            self.training_losses_history["generator_u"] = float(step_g_loss_u)
+            self.training_losses_history["generator_u_e"] = float(step_g_loss_u_e)
+            self.training_losses_history["generator_v"] = float(step_g_loss_v)
+            self.training_losses_history["generator_s"] = float(np.sqrt(step_g_loss_s))
+            self.training_losses_history["generator"] = float(step_g_loss)
+            self.training_losses_history["embedder"] = float(np.sqrt(step_e_loss_t0))
+
+            # Synthetic data generation
+            if epoch in generate_synthetic:
+                _sample = self.generate(n_samples=len(data))
+                self.synthetic_data_generated_in_training[epoch] = _sample
+
+        logger.info("Finished Joint Training")
+        return
+
+    def generate(self, n_samples: int) -> TensorLike:
+        """
+        Generate synthetic time series
+        """
+        steps = n_samples // self.batch_size + 1
+        data = []
+        for _ in trange(steps, desc="Synthetic data generation"):
+            Z_ = next(self.get_noise_batch())
+            records = self.generator(Z_)
+            data.append(records)
+        return np.array(np.vstack(data))[:n_samples]

GAN/timevae.py

@@ -0,0 +1,430 @@
+import os, warnings
+warnings.filterwarnings('ignore')
+
+from abc import ABC, abstractmethod
+import numpy as np
+import joblib
+import tensorflow as tf
+from tensorflow.keras import backend as K
+from tensorflow.keras.layers import Conv1D,  Flatten, Dense, Conv1DTranspose, Reshape, Input, Layer
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.metrics import Mean
+from tensorflow.keras.backend import random_normal
+
+class Sampling(Layer):
+    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""
+
+    def call(self, inputs):
+        z_mean, z_log_var = inputs
+        batch = tf.shape(z_mean)[0]
+        dim = tf.shape(z_mean)[1]
+        epsilon = random_normal(shape=(batch, dim))
+        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
+
+class BaseVariationalAutoencoder(Model, ABC):
+    def __init__(self,
+                 seq_len,
+                 feat_dim,
+                 latent_dim,
+                 reconstruction_wt=3.0,
+                 **kwargs):
+        super(BaseVariationalAutoencoder, self).__init__(**kwargs)
+        self.seq_len = seq_len
+        self.feat_dim = feat_dim
+        self.latent_dim = latent_dim
+        self.reconstruction_wt = reconstruction_wt
+        self.total_loss_tracker = Mean(name="total_loss")
+        self.reconstruction_loss_tracker = Mean(name="reconstruction_loss")
+        self.kl_loss_tracker = Mean(name="kl_loss")
+
+        self.encoder = None
+        self.decoder = None
+
+    def call(self, X):
+        z_mean, _, _ = self.encoder(X)
+        x_decoded = self.decoder(z_mean)
+        if len(x_decoded.shape) == 1: x_decoded = x_decoded.reshape((1, -1))
+        return x_decoded
+
+    def get_num_trainable_variables(self):
+        trainableParams = int(np.sum([np.prod(v.get_shape()) for v in self.trainable_weights]))
+        nonTrainableParams = int(np.sum([np.prod(v.get_shape()) for v in self.non_trainable_weights]))
+        totalParams = trainableParams + nonTrainableParams
+        return trainableParams, nonTrainableParams, totalParams
+
+    def get_prior_samples(self, num_samples):
+        Z = np.random.randn(num_samples, self.latent_dim)
+        samples = self.decoder.predict(Z)
+        return samples
+
+    def get_prior_samples_given_Z(self, Z):
+        samples = self.decoder.predict(Z)
+        return samples
+
+    @abstractmethod
+    def _get_encoder(self, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_decoder(self, **kwargs):
+        raise NotImplementedError
+
+    def summary(self):
+        self.encoder.summary()
+        self.decoder.summary()
+
+    def _get_reconstruction_loss(self, X, X_recons):
+        def get_reconst_loss_by_axis(X, X_c, axis):
+            x_r = tf.reduce_mean(X, axis=axis)
+            x_c_r = tf.reduce_mean(X_recons, axis=axis)
+            err = tf.math.squared_difference(x_r, x_c_r)
+            loss = tf.reduce_sum(err)
+            return loss
+
+        # overall
+        err = tf.math.squared_difference(X, X_recons)
+        reconst_loss = tf.reduce_sum(err)
+
+        reconst_loss += get_reconst_loss_by_axis(X, X_recons, axis=[2])  # by time axis
+        # reconst_loss += get_reconst_loss_by_axis(X, X_recons, axis=[1])    # by feature axis
+        return reconst_loss
+
+    def train_step(self, X):
+        with tf.GradientTape() as tape:
+            z_mean, z_log_var, z = self.encoder(X)
+
+            reconstruction = self.decoder(z)
+
+            reconstruction_loss = self._get_reconstruction_loss(X, reconstruction)
+
+            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
+            kl_loss = tf.reduce_sum(tf.reduce_sum(kl_loss, axis=1))
+            # kl_loss = kl_loss / self.latent_dim
+
+            total_loss = self.reconstruction_wt * reconstruction_loss + kl_loss
+
+        grads = tape.gradient(total_loss, self.trainable_weights)
+
+        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
+
+        self.total_loss_tracker.update_state(total_loss)
+        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
+        self.kl_loss_tracker.update_state(kl_loss)
+
+        return {
+            "loss": self.total_loss_tracker.result(),
+            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
+            "kl_loss": self.kl_loss_tracker.result(),
+        }
+
+    def test_step(self, X):
+        z_mean, z_log_var, z = self.encoder(X)
+        reconstruction = self.decoder(z)
+        reconstruction_loss = self._get_reconstruction_loss(X, reconstruction)
+
+        kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
+        kl_loss = tf.reduce_sum(tf.reduce_sum(kl_loss, axis=1))
+        # kl_loss = kl_loss / self.latent_dim
+
+        total_loss = self.reconstruction_wt * reconstruction_loss + kl_loss
+
+        self.total_loss_tracker.update_state(total_loss)
+        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
+        self.kl_loss_tracker.update_state(kl_loss)
+
+        return {
+            "loss": self.total_loss_tracker.result(),
+            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
+            "kl_loss": self.kl_loss_tracker.result(),
+        }
+
+    def save_weights(self, model_dir, file_pref):
+        encoder_wts = self.encoder.get_weights()
+        decoder_wts = self.decoder.get_weights()
+        joblib.dump(encoder_wts, os.path.join(model_dir, f'{file_pref}encoder_wts.h5'))
+        joblib.dump(decoder_wts, os.path.join(model_dir, f'{file_pref}decoder_wts.h5'))
+
+    def load_weights(self, model_dir, file_pref):
+        encoder_wts = joblib.load(os.path.join(model_dir, f'{file_pref}encoder_wts.h5'))
+        decoder_wts = joblib.load(os.path.join(model_dir, f'{file_pref}decoder_wts.h5'))
+
+        self.encoder.set_weights(encoder_wts)
+        self.decoder.set_weights(decoder_wts)
+
+    def save(self, model_dir, file_pref):
+        self.save_weights(model_dir, file_pref)
+        dict_params = {
+
+            'seq_len': self.seq_len,
+            'feat_dim': self.feat_dim,
+            'latent_dim': self.latent_dim,
+            'reconstruction_wt': self.reconstruction_wt,
+            'hidden_layer_sizes': self.hidden_layer_sizes,
+        }
+        params_file = os.path.join(model_dir, f'{file_pref}parameters.pkl')
+        joblib.dump(dict_params, params_file)
+
+class TimeVAE(BaseVariationalAutoencoder):
+
+    def __init__(self, hidden_layer_sizes, trend_poly=0, num_gen_seas=0, custom_seas=None,
+                 use_scaler=False, use_residual_conn=True, **kwargs):
+        '''
+            hidden_layer_sizes: list of number of filters in convolutional layers in encoder and residual connection of decoder.
+            trend_poly: integer for number of orders for trend component. e.g. setting trend_poly = 2 will include linear and quadratic term.
+            num_gen_seas: Number of sine-waves to use to model seasonalities. Each sine wae will have its own amplitude, frequency and phase.
+            custom_seas: list of tuples of (num_seasons, len_per_season).
+                num_seasons: number of seasons per cycle.
+                len_per_season: number of epochs (time-steps) per season.
+            use_residual_conn: boolean value indicating whether to use a residual connection for reconstruction in addition to
+            trend, generic and custom seasonalities.
+        '''
+
+        super(TimeVAE, self).__init__(**kwargs)
+
+        self.hidden_layer_sizes = hidden_layer_sizes
+        self.trend_poly = trend_poly
+        self.num_gen_seas = num_gen_seas
+        self.custom_seas = custom_seas
+        self.use_scaler = use_scaler
+        self.use_residual_conn = use_residual_conn
+        self.encoder = self._get_encoder()
+        self.decoder = self._get_decoder()
+
+    def _get_encoder(self):
+        encoder_inputs = Input(shape=(self.seq_len, self.feat_dim), name='encoder_input')
+        x = encoder_inputs
+        for i, num_filters in enumerate(self.hidden_layer_sizes):
+            x = Conv1D(
+                filters=num_filters,
+                kernel_size=3,
+                strides=2,
+                activation='relu',
+                padding='same',
+                name=f'enc_conv_{i}')(x)
+
+        x = Flatten(name='enc_flatten')(x)
+
+        # save the dimensionality of this last dense layer before the hidden state layer. We need it in the decoder.
+        self.encoder_last_dense_dim = x.get_shape()[-1]
+
+        z_mean = Dense(self.latent_dim, name="z_mean")(x)
+        z_log_var = Dense(self.latent_dim, name="z_log_var")(x)
+
+        encoder_output = Sampling()([z_mean, z_log_var])
+        self.encoder_output = encoder_output
+
+        encoder = Model(encoder_inputs, [z_mean, z_log_var, encoder_output], name="encoder")
+        return encoder
+
+    def _get_decoder(self):
+        decoder_inputs = Input(shape=(int(self.latent_dim)), name='decoder_input')
+
+        outputs = None
+        outputs = self.level_model(decoder_inputs)
+
+        # trend polynomials
+        if self.trend_poly is not None and self.trend_poly > 0:
+            trend_vals = self.trend_model(decoder_inputs)
+            outputs = trend_vals if outputs is None else outputs + trend_vals
+
+            # # generic seasonalities
+        # if self.num_gen_seas is not None and self.num_gen_seas > 0:
+        #     gen_seas_vals, freq, phase, amplitude = self.generic_seasonal_model(decoder_inputs)
+        #     # gen_seas_vals = self.generic_seasonal_model2(decoder_inputs)
+        #     outputs = gen_seas_vals if outputs is None else outputs + gen_seas_vals
+
+        # custom seasons
+        if self.custom_seas is not None and len(self.custom_seas) > 0:
+            cust_seas_vals = self.custom_seasonal_model(decoder_inputs)
+            outputs = cust_seas_vals if outputs is None else outputs + cust_seas_vals
+
+        if self.use_residual_conn:
+            residuals = self._get_decoder_residual(decoder_inputs)
+            outputs = residuals if outputs is None else outputs + residuals
+
+        if self.use_scaler and outputs is not None:
+            scale = self.scale_model(decoder_inputs)
+            outputs *= scale
+
+        # outputs = Activation(activation='sigmoid')(outputs)
+
+        if outputs is None:
+            raise Exception('''Error: No decoder model to use. 
+            You must use one or more of:
+            trend, generic seasonality(ies), custom seasonality(ies), and/or residual connection. ''')
+
+        decoder = Model(decoder_inputs, [outputs], name="decoder")
+        return decoder
+
+    def level_model(self, z):
+        level_params = Dense(self.feat_dim, name="level_params", activation='relu')(z)
+        level_params = Dense(self.feat_dim, name="level_params2")(level_params)
+        level_params = Reshape(target_shape=(1, self.feat_dim))(level_params)  # shape: (N, 1, D)
+
+        ones_tensor = tf.ones(shape=[1, self.seq_len, 1], dtype=tf.float32)  # shape: (1, T, D)
+
+        level_vals = level_params * ones_tensor
+        return level_vals
+
+    def scale_model(self, z):
+        scale_params = Dense(self.feat_dim, name="scale_params", activation='relu')(z)
+        scale_params = Dense(self.feat_dim, name="scale_params2")(scale_params)
+        scale_params = Reshape(target_shape=(1, self.feat_dim))(scale_params)  # shape: (N, 1, D)
+
+        scale_vals = tf.repeat(scale_params, repeats=self.seq_len, axis=1)  # shape: (N, T, D)
+        return scale_vals
+
+    def trend_model(self, z):
+        trend_params = Dense(self.feat_dim * self.trend_poly, name="trend_params", activation='relu')(z)
+        trend_params = Dense(self.feat_dim * self.trend_poly, name="trend_params2")(trend_params)
+        trend_params = Reshape(target_shape=(self.feat_dim, self.trend_poly))(trend_params)  # shape: N x D x P
+
+        lin_space = K.arange(0, float(self.seq_len), 1) / self.seq_len  # shape of lin_space : 1d tensor of length T
+        poly_space = K.stack([lin_space ** float(p + 1) for p in range(self.trend_poly)], axis=0)  # shape: P x T
+
+        trend_vals = K.dot(trend_params, poly_space)  # shape (N, D, T)
+        trend_vals = tf.transpose(trend_vals, perm=[0, 2, 1])  # shape: (N, T, D)
+        trend_vals = K.cast(trend_vals, tf.float32)
+        return trend_vals
+
+    def custom_seasonal_model(self, z):
+
+        N = tf.shape(z)[0]
+        ones_tensor = tf.ones(shape=[N, self.feat_dim, self.seq_len], dtype=tf.int32)
+
+        all_seas_vals = []
+        for i, season_tup in enumerate(self.custom_seas):
+            num_seasons, len_per_season = season_tup
+
+            season_params = Dense(self.feat_dim * num_seasons, name=f"season_params_{i}")(z)  # shape: (N, D * S)
+            season_params = Reshape(target_shape=(self.feat_dim, num_seasons))(season_params)  # shape: (N, D, S)
+
+            season_indexes_over_time = self._get_season_indexes_over_seq(num_seasons, len_per_season)  # shape: (T, )
+
+            dim2_idxes = ones_tensor * tf.reshape(season_indexes_over_time, shape=(1, 1, -1))  # shape: (1, 1, T)
+
+            season_vals = tf.gather(season_params, dim2_idxes, batch_dims=-1)  # shape (N, D, T)
+
+            all_seas_vals.append(season_vals)
+
+        all_seas_vals = K.stack(all_seas_vals, axis=-1)  # shape: (N, D, T, S)
+        all_seas_vals = tf.reduce_sum(all_seas_vals, axis=-1)  # shape (N, D, T)
+        all_seas_vals = tf.transpose(all_seas_vals, perm=[0, 2, 1])  # shape (N, T, D)
+        return all_seas_vals
+
+    def _get_season_indexes_over_seq(self, num_seasons, len_per_season):
+        curr_len = 0
+        season_idx = []
+        curr_idx = 0
+        while curr_len < self.seq_len:
+            reps = len_per_season if curr_len + len_per_season <= self.seq_len else self.seq_len - curr_len
+            season_idx.extend([curr_idx] * reps)
+            curr_idx += 1
+            if curr_idx == num_seasons: curr_idx = 0
+            curr_len += reps
+        return season_idx
+
+    def generic_seasonal_model(self, z):
+
+        freq = Dense(self.feat_dim * self.num_gen_seas, name="g_season_freq", activation='sigmoid')(z)
+        freq = Reshape(target_shape=(1, self.feat_dim, self.num_gen_seas))(freq)  # shape: (N, 1, D, S)
+
+        phase = Dense(self.feat_dim * self.num_gen_seas, name="g_season_phase")(z)
+        phase = Reshape(target_shape=(1, self.feat_dim, self.num_gen_seas))(phase)  # shape: (N, 1, D, S)
+
+        amplitude = Dense(self.feat_dim * self.num_gen_seas, name="g_season_amplitude")(z)
+        amplitude = Reshape(target_shape=(1, self.feat_dim, self.num_gen_seas))(amplitude)  # shape: (N, 1, D, S)
+
+        lin_space = K.arange(0, float(self.seq_len), 1) / self.seq_len  # shape of lin_space : 1d tensor of length T
+        lin_space = tf.reshape(lin_space, shape=(1, self.seq_len, 1, 1))  # shape: 1, T, 1, 1
+
+        seas_vals = amplitude * K.sin(2. * np.pi * freq * lin_space + phase)  # shape: N, T, D, S
+        seas_vals = tf.math.reduce_sum(seas_vals, axis=-1)  # shape: N, T, D
+
+        return seas_vals
+
+    def generic_seasonal_model2(self, z):
+
+        season_params = Dense(self.feat_dim * self.num_gen_seas, name="g_season_params")(z)
+        season_params = Reshape(target_shape=(self.feat_dim, self.num_gen_seas))(season_params)  # shape: (D, S)
+
+        p = self.num_gen_seas
+        p1, p2 = (p // 2, p // 2) if p % 2 == 0 else (p // 2, p // 2 + 1)
+
+        ls = K.arange(0, float(self.seq_len), 1) / self.seq_len  # shape of ls : 1d tensor of length T
+
+        s1 = K.stack([K.cos(2 * np.pi * i * ls) for i in range(p1)], axis=0)
+        s2 = K.stack([K.sin(2 * np.pi * i * ls) for i in range(p2)], axis=0)
+        if p == 1:
+            s = s2
+        else:
+            s = K.concatenate([s1, s2], axis=0)
+        s = K.cast(s, np.float32)
+
+        seas_vals = K.dot(season_params, s, name='g_seasonal_vals')
+        seas_vals = tf.transpose(seas_vals, perm=[0, 2, 1])  # shape: (N, T, D)
+        seas_vals = K.cast(seas_vals, np.float32)
+        print('seas_vals shape', tf.shape(seas_vals))
+
+        return seas_vals
+
+    def _get_decoder_residual(self, x):
+
+        x = Dense(self.encoder_last_dense_dim, name="dec_dense", activation='relu')(x)
+        x = Reshape(target_shape=(-1, self.hidden_layer_sizes[-1]), name="dec_reshape")(x)
+
+        for i, num_filters in enumerate(reversed(self.hidden_layer_sizes[:-1])):
+            x = Conv1DTranspose(
+                filters=num_filters,
+                kernel_size=3,
+                strides=2,
+                padding='same',
+                activation='relu',
+                name=f'dec_deconv_{i}')(x)
+
+        # last de-convolution
+        x = Conv1DTranspose(
+            filters=self.feat_dim,
+            kernel_size=3,
+            strides=2,
+            padding='same',
+            activation='relu',
+            name=f'dec_deconv__{i + 1}')(x)
+
+        x = Flatten(name='dec_flatten')(x)
+        x = Dense(self.seq_len * self.feat_dim, name="decoder_dense_final")(x)
+        residuals = Reshape(target_shape=(self.seq_len, self.feat_dim))(x)
+        return residuals
+
+    def save(self, model_dir, file_pref):
+
+        super().save_weights(model_dir, file_pref)
+        dict_params = {
+            'seq_len': self.seq_len,
+            'feat_dim': self.feat_dim,
+            'latent_dim': self.latent_dim,
+            'reconstruction_wt': self.reconstruction_wt,
+
+            'hidden_layer_sizes': self.hidden_layer_sizes,
+            'trend_poly': self.trend_poly,
+            'num_gen_seas': self.num_gen_seas,
+            'custom_seas': self.custom_seas,
+            'use_scaler': self.use_scaler,
+            'use_residual_conn': self.use_residual_conn,
+        }
+        params_file = os.path.join(model_dir, f'{file_pref}parameters.pkl')
+        joblib.dump(dict_params, params_file)
+
+    @staticmethod
+    def load(model_dir, file_pref):
+        params_file = os.path.join(model_dir, f'{file_pref}parameters.pkl')
+        dict_params = joblib.load(params_file)
+
+        vae_model = TimeVAE(**dict_params)
+
+        vae_model.load_weights(model_dir, file_pref)
+
+        vae_model.compile(optimizer=Adam())
+
+        return vae_model

GAN/utils.py

@@ -0,0 +1,315 @@
+import math
+import numpy as np
+import typing as T
+import seaborn as sns
+import matplotlib.pyplot as plt
+import sklearn
+import sklearn.manifold
+import tensorflow as tf
+import numpy.typing as npt
+
+from tensorflow import keras
+from tensorflow.python.types.core import TensorLike
+
+Tensor = T.Union[tf.Tensor, npt.NDArray]
+OptTensor = T.Optional[Tensor]
+
+
+EPS = 1e-18
+class TSFeatureScaler:
+    """Global time series scaler that scales all features to [0,1] then normalizes to [-1,1]"""
+    
+    def __init__(self) -> None:
+        self.min_val = None
+        self.max_val = None
+        
+    def fit(self, X: TensorLike) -> "TSFeatureScaler":
+        """
+        Fit scaler to data
+        
+        Args:
+            X: Input tensor of shape [N, T, D] 
+               (N: samples, T: timesteps, D: features)
+        """
+        # 计算整个数据集的全局最大最小值
+        self.min_val = np.min(X)
+        self.max_val = np.max(X)
+        return self
+    
+    def transform(self, X: TensorLike) -> TensorLike:
+        """
+        Transform data in two steps:
+        1. Scale to [0,1] using min-max scaling
+        2. Normalize to [-1,1]
+        """
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("Scaler must be fitted before transform")
+            
+        # 1. 缩放到[0,1]
+        X_scaled = (X - self.min_val) / (self.max_val - self.min_val + EPS)
+        
+        # 2. 归一化到[-1,1]
+        X_normalized = 2.0 * X_scaled - 1.0
+        
+        return X_normalized
+    
+    def inverse_transform(self, X: TensorLike) -> TensorLike:
+        """
+        Inverse transform data:
+        1. From [-1,1] back to [0,1]
+        2. From [0,1] back to original range
+        """
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("Scaler must be fitted before inverse_transform")
+            
+        # 1. 从[-1,1]转回[0,1]
+        X_scaled = (X + 1.0) / 2.0
+        
+        # 2. 从[0,1]转回原始范围
+        X_original = X_scaled * (self.max_val - self.min_val + EPS) + self.min_val
+        
+        return X_original
+    
+    def fit_transform(self, X: TensorLike) -> TensorLike:
+        """Fit to data, then transform it"""
+        return self.fit(X).transform(X)
+    
+    def get_range(self) -> T.Tuple[float, float]:
+        """获取原始数据的范围"""
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("Scaler must be fitted first")
+        return (self.min_val, self.max_val)
+
+
+EPS = 1e-18
+class TSFeatureWiseScaler():
+    def __init__(self, feature_range: T.Tuple[float, float] = (0, 1)) -> None:
+        assert len(feature_range) == 2
+
+        self._min_v, self._max_v = feature_range
+
+    # X: N x T x D
+    def fit(self, X: TensorLike) -> "TSFeatureWiseScaler":
+        D = X.shape[2]
+        self.mins = np.zeros(D)
+        self.maxs = np.zeros(D)
+
+        for i in range(D):
+            self.mins[i] = np.min(X[:, :, i])
+            self.maxs[i] = np.max(X[:, :, i])
+
+        return self
+
+    def transform(self, X: TensorLike) -> TensorLike:
+        return ((X - self.mins) / (self.maxs - self.mins + EPS)) * (self._max_v - self._min_v) + self._min_v
+
+    def inverse_transform(self, X: TensorLike) -> TensorLike:
+        X -= self._min_v
+        X /= self._max_v - self._min_v
+        X *= (self.maxs - self.mins + EPS)
+        X += self.mins
+        return X
+
+    def fit_transform(self, X: TensorLike) -> TensorLike:
+        self.fit(X)
+        return self.transform(X)
+ 
+
+def linear_beta_schedule(timesteps, beta_start=1e-4, beta_end=0.99): # beta_end=0.99
+    betas = np.linspace(beta_start, beta_end, timesteps, dtype=np.float32)
+    return betas
+
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    steps = timesteps + 1
+    x = np.linspace(0, timesteps, steps, dtype=np.float64)
+    alphas_cumprod = np.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    betas = np.clip(betas, 0, 0.999)
+    return betas
+
+
+def reconstruction_loss_by_axis(original: tf.Tensor, reconstructed: tf.Tensor, axis: int = 0) -> tf.Tensor:
+    """
+    Calculate the reconstruction loss based on a specified axis.
+
+    This function computes the reconstruction loss between the original data and
+    the reconstructed data along a specified axis. The loss can be computed in
+    two ways depending on the chosen axis:
+
+    - When `axis` is 0, it computes the loss as the sum of squared differences
+      between the original and reconstructed data for all elements.
+    - When `axis` is 1 or 2, it computes the mean squared error (MSE) between the
+      mean values along the chosen axis for the original and reconstructed data.
+
+    Parameters:
+    ----------
+    original : tf.Tensor
+        The original data tensor.
+
+    reconstructed : tf.Tensor
+        The reconstructed data tensor, typically produced by an autoencoder.
+
+    axis : int, optional (default=0)
+        The axis along which to compute the reconstruction loss:
+        - 0: All elements (sum of squared differences).
+        - 1: Along features (MSE).
+        - 2: Along time steps (MSE).
+
+    Returns:
+    -------
+    tf.Tensor
+        The computed reconstruction loss as a TensorFlow tensor.
+    Notes:
+    ------
+    - This function is commonly used in the context of autoencoders and other
+      reconstruction-based models to assess the quality of the reconstruction.
+    - The choice of `axis` determines how the loss is calculated, and it should
+      align with the data's structure.
+    """
+
+    # axis=0 all (sum of squared diffs)
+    # axis=1 features (MSE)
+    # axis=2 times (MSE)
+    if axis == 0:
+        return tf.reduce_sum(tf.math.squared_difference(original, reconstructed))
+    else:
+        return tf.losses.mean_squared_error(tf.reduce_mean(original, axis=axis), tf.reduce_mean(reconstructed, axis=axis))
+
+
+def gen_sine_dataset(N: int, T: int, D: int, max_value: int = 10) -> npt.NDArray:
+    result = []
+    for i in range(N):
+        result.append([])
+        a = np.random.random() * max_value
+        shift = np.random.random() * max_value + 1
+        ts = np.arange(0, T, 1)
+        for d in range(1, D + 1):
+            result[-1].append((a * np.sin((d + 3) * ts / 25. + shift)).T)
+
+    return np.transpose(np.array(result), [0, 2, 1])
+
+
+def gen_sine_vs_const_dataset(N: int, T: int, D: int, max_value: int = 10, const: int = 0) -> T.Tuple[TensorLike, TensorLike]:
+    result_X, result_y = [], []
+    for i in range(N):
+        scales = np.random.random(D) * max_value
+        consts = np.random.random(D) * const
+        shifts = np.random.random(D) * 2
+        alpha = np.random.random()
+        if np.random.random() < 0.5:
+            times = np.repeat(np.arange(0, T, 1)[:, None], D, axis=1) / 10
+            result_X.append(np.sin(alpha * times + shifts) * scales)
+            result_y.append(0)
+        else:
+            result_X.append(np.tile(consts, (T, 1)))
+            result_y.append(1)
+    return np.array(result_X), np.array(result_y)
+
+
+def visualize_ts_lineplot(
+        ts: Tensor,
+        ys: OptTensor = None,
+        num: int = 5,
+        unite_features: bool = True,
+) -> None:
+    assert len(ts.shape) == 3
+
+    fig, axs = plt.subplots(num, 1, figsize=(14, 10))
+    if num == 1:
+        axs = [axs]
+
+    ids = np.random.choice(ts.shape[0], size=num, replace=False)
+    for i, sample_id in enumerate(ids):
+        if not unite_features:
+            feature_id = np.random.randint(ts.shape[2])
+            sns.lineplot(
+                x=range(ts.shape[1]),
+                y=ts[sample_id, :, feature_id],
+                ax=axs[i],
+                label=rf"feature \#{feature_id}",
+            )
+        else:
+            for feat_id in range(ts.shape[2]):
+                sns.lineplot(
+                    x=range(ts.shape[1]), y=ts[sample_id, :, feat_id], ax=axs[i]
+                )
+        if ys is not None:
+            if len(ys.shape) == 1:
+                axs[i].set_title(ys[sample_id])
+            elif len(ys.shape) == 2:
+                sns.lineplot(
+                    x=range(ts.shape[1]),
+                    y=ys[sample_id],
+                    ax=axs[i].twinx(),
+                    color="g",
+                    label="Target variable",
+                )
+            else:
+                raise ValueError("ys contains too many dimensions")
+    #plt.show()
+
+def visualize_tsne(
+    X: Tensor,
+    y: Tensor,
+    X_gen: Tensor,
+    y_gen: Tensor,
+    path: str = "/tmp/tsne_embeddings.pdf",
+    feature_averaging: bool = False,
+    perplexity=30.0
+) -> None:
+    """
+    Visualizes t-SNE embeddings of real and synthetic data.
+
+    This function generates a scatter plot of t-SNE embeddings for real and synthetic data.
+    Each data point is represented by a marker on the plot, and the colors of the markers
+    correspond to the corresponding class labels of the data points.
+
+    :param X: The original real data tensor of shape (num_samples, num_features).
+    :type X: tsgm.types.Tensor
+    :param y: The labels of the original real data tensor of shape (num_samples,).
+    :type y: tsgm.types.Tensor
+    :param X_gen: The generated synthetic data tensor of shape (num_samples, num_features).
+    :type X_gen: tsgm.types.Tensor
+    :param y_gen: The labels of the generated synthetic data tensor of shape (num_samples,).
+    :type y_gen: tsgm.types.Tensor
+    :param path: The path to save the visualization as a PDF file. Defaults to "/tmp/tsne_embeddings.pdf".
+    :type path: str, optional
+    :param feature_averaging: Whether to compute the average features for each class. Defaults to False.
+    :type feature_averaging: bool, optional
+    """
+    tsne = sklearn.manifold.TSNE(n_components=2, perplexity=perplexity, learning_rate="auto", init="random")
+
+    if feature_averaging:
+        X_all = np.concatenate((np.mean(X, axis=2), np.mean(X_gen, axis=2)))
+
+        X_emb = tsne.fit_transform(np.resize(X_all, (X_all.shape[0], X_all.shape[1])))
+    else:
+        X_all = np.concatenate((X, X_gen))
+
+        X_emb = tsne.fit_transform(
+            np.resize(X_all, (X_all.shape[0], X_all.shape[1] * X_all.shape[2]))
+        )
+
+    y_all = np.concatenate((y, y_gen))
+
+    c = np.argmax(y_all, axis=1)
+    colors = {0: "class 0", 1: "class 1"}
+    c = [colors[el] for el in c]
+    point_styles = ["hist"] * X.shape[0] + ["gen"] * X_gen.shape[0]
+
+    plt.figure(figsize=(8, 6), dpi=80)
+    sns.scatterplot(
+        x=X_emb[:, 0],
+        y=X_emb[:, 1],
+        hue=c,
+        style=point_styles,
+        markers={"hist": "<", "gen": "H"},
+        alpha=0.7,
+    )
+    plt.legend()
+    plt.box(False)
+    plt.axis("off")
+    plt.savefig(path)
+    plt.show()

GAN/zoo.py

@@ -0,0 +1,517 @@
+import abc
+import math
+import typing as T
+import tensorflow as tf
+import numpy.typing as npt
+from tensorflow import keras
+from tensorflow.keras import layers
+
+from prettytable import PrettyTable
+
+Tensor = T.Union[tf.Tensor, npt.NDArray]
+OptTensor = T.Optional[Tensor]
+
+class Sampling(tf.keras.layers.Layer):
+    def call(self, inputs: Tensor) -> Tensor:
+        z_mean, z_log_var = inputs
+        epsilon = tf.keras.backend.random_normal(shape=tf.shape(z_mean))
+        return z_mean + tf.exp(0.5 * z_log_var) * epsilon
+
+
+class Architecture(abc.ABC):
+    @abc.abstractproperty
+    def arch_type(self):
+        raise NotImplementedError
+
+class BaseGANArchitecture(Architecture):
+    @property
+    def discriminator(self) -> keras.models.Model:
+        if hasattr(self, "_discriminator"):
+            return self._discriminator
+        else:
+            raise NotImplementedError
+
+    @property
+    def generator(self) -> keras.models.Model:
+        if hasattr(self, "_generator"):
+            return self._generator
+        else:
+            raise NotImplementedError
+
+    def get(self) -> T.Dict:
+        if hasattr(self, "_discriminator") and hasattr(self, "_generator"):
+            return {"discriminator": self._discriminator, "generator": self._generator}
+        else:
+            raise NotImplementedError
+
+class BaseVAEArchitecture(Architecture):
+    @property
+    def encoder(self) -> keras.models.Model:
+        if hasattr(self, "_encoder"):
+            return self._encoder
+        else:
+            raise NotImplementedError
+
+    @property
+    def decoder(self) -> keras.models.Model:
+        if hasattr(self, "_decoder"):
+            return self._decoder
+        else:
+            raise NotImplementedError
+
+    def get(self) -> T.Dict:
+        if hasattr(self, "_encoder") and hasattr(self, "_decoder"):
+            return {"encoder": self._encoder, "decoder": self._decoder}
+        else:
+            raise NotImplementedError
+
+class VAE_CONV5Architecture(BaseVAEArchitecture):
+    arch_type = "vae: conv"
+
+    def __init__(self, seq_len: int, feat_dim: int, latent_dim: int) -> None:
+        super().__init__()
+        self._seq_len = seq_len
+        self._feat_dim = feat_dim
+        self._latent_dim = latent_dim
+        self._encoder = self._build_encoder()
+        self._decoder = self._build_decoder()
+
+    def _build_encoder(self) -> keras.models.Model:
+        encoder_inputs = keras.Input(shape=(self._seq_len, self._feat_dim))
+        x = layers.Conv1D(64, 3, activation="relu", strides=1, padding="same")(
+            encoder_inputs
+        )
+        x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 4, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        x = layers.Flatten()(x)
+        #x = layers.Dense(512, activation="relu")(x)
+        x = layers.Dense(64, activation="relu")(x)
+        z_mean = layers.Dense(self._latent_dim, name="z_mean")(x)
+        z_log_var = layers.Dense(self._latent_dim, name="z_log_var")(x)
+        z = Sampling()([z_mean, z_log_var])
+        encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
+        return encoder
+
+    def _build_decoder(self) -> keras.models.Model:
+        latent_inputs = keras.Input(shape=(self._latent_dim,))
+        x = layers.Dense(64, activation="relu")(latent_inputs)
+        # x = layers.Dense(512, activation="relu")(x)
+        # x = layers.Dense(64, activation="relu")(x)
+
+        dense_shape = self._encoder.layers[-6].output_shape[1] * self._seq_len
+
+        x = layers.Dense(dense_shape, activation="relu")(x)
+
+        x = layers.Reshape((self._seq_len, dense_shape // self._seq_len))(x)
+        x = layers.Conv1DTranspose(64, 2, activation="relu", strides=1, padding="same")(
+            x
+        )
+        x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(64, 2, activation="relu", strides=1, padding="same")(
+        #     x
+        # )
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(64, 2, activation="relu", strides=1, padding="same")(
+        #     x
+        # )
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(64, 2, activation="relu", strides=1, padding="same")(
+        #     x
+        # )
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(
+        #     64, 10, activation="relu", strides=1, padding="same"
+        # )(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        decoder_outputs = layers.Conv1DTranspose(
+            self._feat_dim, 3, activation="sigmoid", padding="same"
+        )(x)
+        decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
+        return decoder
+
+class cVAE_CONV5Architecture(BaseVAEArchitecture):
+    arch_type = "vae:conditional"
+
+    def __init__(self, seq_len: int, feat_dim: int, latent_dim: int, output_dim: int = 2) -> None:
+        self._seq_len = seq_len
+        self._feat_dim = feat_dim
+        self._latent_dim = latent_dim
+        self._output_dim = output_dim
+
+        self._encoder = self._build_encoder()
+        self._decoder = self._build_decoder()
+
+    def _build_encoder(self) -> keras.models.Model:
+        encoder_inputs = keras.Input(
+            shape=(self._seq_len, self._feat_dim + self._output_dim)
+        )
+
+        x = layers.Conv1D(64, 3, activation="relu", strides=1, padding="same")(
+            encoder_inputs
+        )
+        x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 2, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(64, 4, activation="relu", strides=1, padding="same")(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        x = layers.Flatten()(x)
+        #x = layers.Dense(512, activation="relu")(x)
+        x = layers.Dense(64, activation="relu")(x)
+        z_mean = layers.Dense(self._latent_dim * self._seq_len, name="z_mean")(x)
+        z_log_var = layers.Dense(self._latent_dim * self._seq_len, name="z_log_var")(x)
+        z = Sampling()([z_mean, z_log_var])
+        encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
+        return encoder
+
+    def _build_decoder(self) -> keras.models.Model:
+        inputs = keras.Input(
+            shape=(
+                self._seq_len,
+                self._latent_dim + self._output_dim,
+            )
+        )
+        x = layers.Conv1DTranspose(64, 2, strides=2, padding="same")(inputs)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(64, 2, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1DTranspose(64, 2, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Dropout(rate=0.2)(x)
+
+        pool_and_stride = round((x.shape[1] + 1) / (self._seq_len + 1))
+        x = layers.AveragePooling1D(pool_size=pool_and_stride, strides=pool_and_stride)(
+            x
+        )
+        d_output = layers.LocallyConnected1D(self._feat_dim, 1, activation="sigmoid")(x)
+
+        decoder = keras.Model(inputs, d_output, name="decoder")
+        return decoder
+class GAN_ConvLSTM4Architecture(BaseGANArchitecture):
+    arch_type = "gan: conv_lstm"
+
+    def __init__(self, seq_len: int, feat_dim: int, latent_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self._seq_len = seq_len
+        self._feat_dim = feat_dim
+        self._latent_dim = latent_dim
+        self._output_dim = output_dim
+        self.generator_in_channels = latent_dim + output_dim
+        self.discriminator_in_channels = feat_dim + output_dim
+
+        self._discriminator = self._build_discriminator()
+        self._generator = self._build_generator()
+
+    def _build_discriminator(self) -> keras.models.Model:
+        d_input = keras.Input((self._seq_len, self.discriminator_in_channels))
+        x = layers.Conv1D(64, 3, strides=2, padding="same")(d_input)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        # x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Dropout(rate=0.2)(x)
+        x = layers.GlobalAvgPool1D()(x)
+        d_output = layers.Dense(1, activation="sigmoid")(x)
+        discriminator = keras.Model(d_input, d_output, name="discriminator")
+        return discriminator
+
+    def _build_generator(self) -> keras.models.Model:
+        g_input = keras.Input((self.generator_in_channels,))
+        x = layers.Dense(8 * 8 * self._seq_len)(g_input)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Reshape((self._seq_len, 64))(x)
+        x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        # x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        # x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Conv1D(1, 8, padding="same")(x)
+        x = layers.LSTM(256, return_sequences=True)(x)
+
+        pool_and_stride = math.ceil((x.shape[1] + 1) / (self._seq_len + 1))
+
+        x = layers.AveragePooling1D(pool_size=pool_and_stride, strides=pool_and_stride)(
+            x
+        )
+        g_output = layers.LocallyConnected1D(self._feat_dim, 1, activation="tanh")(x)
+        generator = keras.Model(g_input, g_output, name="generator")
+        return generator
+
+
+
+class GAN_Conv2LSTM4Architecture(BaseGANArchitecture):
+    arch_type = "gan: conv_lstm_2"
+
+    def __init__(self, seq_len: int, feat_dim: int, latent_dim: int, output_dim: int, n_blocks: int = 1, output_activation: str = "tanh") -> None:
+        super().__init__()
+        self._seq_len = seq_len
+        self._feat_dim = feat_dim
+        self._latent_dim = latent_dim
+        self._output_dim = output_dim
+        self._n_blocks = n_blocks
+        self._output_activation = output_activation
+
+        self.generator_in_channels = latent_dim + output_dim
+        self.discriminator_in_channels = feat_dim + output_dim
+
+        self._discriminator = self._build_discriminator()
+        self._generator = self._build_generator(output_activation=output_activation)
+
+    def _build_discriminator(self) -> keras.Model:
+        d_input = keras.Input((self._seq_len, self.discriminator_in_channels))
+        x = d_input
+        for i in range(self._n_blocks - 1):
+            x = layers.LSTM(64, return_sequences=True)(x)
+            x = layers.Dropout(rate=0.2)(x)
+
+        x = layers.LSTM(64, return_sequences=True)(x)
+        x = layers.Dropout(rate=0.2)(x)
+
+        x = layers.GlobalAvgPool1D()(x)
+        d_output = layers.Dense(1, activation="sigmoid")(x)
+        discriminator = keras.Model(d_input, d_output, name="discriminator")
+        return discriminator
+
+    def _build_generator(self, output_activation: str) -> keras.Model:
+        g_input = keras.Input((self.generator_in_channels,))
+
+        x = layers.Dense(8 * 8 * self._seq_len)(g_input)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Reshape((self._seq_len, 64))(x)
+
+        for i in range(self._n_blocks - 1):
+            x = layers.LSTM(64, return_sequences=True)(x)
+            x = layers.Dropout(rate=0.2)(x)
+        x = layers.LSTM(256, return_sequences=True)(x)
+
+        pool_and_stride = round((x.shape[1] + 1) / (self._seq_len + 1))
+
+        x = layers.AveragePooling1D(pool_size=pool_and_stride, strides=pool_and_stride)(x)
+        g_output = layers.LocallyConnected1D(self._feat_dim, 1, activation=output_activation)(x)
+        generator = keras.Model(g_input, g_output, name="generator")
+        return generator
+
+class GAN_Conv3LSTM4Architecture(BaseGANArchitecture):
+    arch_type = "gan: conv_lstm_3"
+
+    def __init__(self, seq_len: int, feat_dim: int, latent_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self._seq_len = seq_len
+        self._feat_dim = feat_dim
+        self._latent_dim = latent_dim
+        self._output_dim = output_dim
+
+        self.generator_in_channels = latent_dim + output_dim
+        self.discriminator_in_channels = feat_dim + output_dim
+
+        self._discriminator = self._build_discriminator()
+        self._generator = self._build_generator()
+
+    def _build_discriminator(self) -> keras.models.Model:
+        d_input = keras.Input((self._seq_len, self.discriminator_in_channels))
+        x = layers.LSTM(64, return_sequences=True)(d_input)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        x = layers.Conv1D(128, 3, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Dropout(rate=0.2)(x)
+        x = layers.GlobalAvgPool1D()(x)
+        d_output = layers.Dense(1, activation="sigmoid")(x)
+        discriminator = keras.Model(d_input, d_output, name="discriminator")
+        return discriminator
+
+    def _build_generator(self) -> keras.models.Model:
+        g_input = keras.Input((self.generator_in_channels,))
+        x = layers.Dense(8 * 8 * self._seq_len)(g_input)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Reshape((self._seq_len, 64))(x)
+        x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Conv1DTranspose(128, 4, strides=2, padding="same")(x)
+        x = layers.LeakyReLU(alpha=0.2)(x)
+        x = layers.Conv1D(1, 8, padding="same")(x)
+        x = layers.LSTM(256, return_sequences=True)(x)
+
+        pool_and_stride = round((x.shape[1] + 1) / (self._seq_len + 1))
+
+        x = layers.AveragePooling1D(pool_size=pool_and_stride, strides=pool_and_stride)(x)
+        g_output = layers.LocallyConnected1D(self._feat_dim, 1, activation="tanh")(x)
+        generator = keras.Model(g_input, g_output, name="generator")
+        return generator
+
+
+# class BaseClassificationArchitecture(Architecture):
+#     arch_type = "downstream:classification"
+#
+#     def __init__(self, seq_len: int, feat_dim: int, output_dim: int) -> None:
+#         self._seq_len = seq_len
+#         self._feat_dim = feat_dim
+#         self._output_dim = output_dim
+#         self._model = self._build_model()
+#
+#     @property
+#     def model(self) -> keras.models.Model:
+#         return self._model
+#
+#     def get(self) -> T.Dict:
+#         return {"model": self.model}
+#
+#     def _build_model(self) -> None:
+#         raise NotImplementedError
+
+
+# class ConvnArchitecture(BaseClassificationArchitecture):
+#     def __init__(
+#         self, seq_len: int, feat_dim: int, output_dim: int, n_conv_blocks: int = 1
+#     ) -> None:
+#         self._n_conv_blocks = n_conv_blocks
+#         super().__init__(seq_len, feat_dim, output_dim)
+#
+#     def _build_model(self) -> keras.models.Model:
+#         m_input = keras.Input((self._seq_len, self._feat_dim))
+#         x = m_input
+#         for _ in range(self._n_conv_blocks):
+#             x = layers.Conv1D(filters=64, kernel_size=3, activation="relu")(x)
+#             x = layers.Dropout(0.2)(x)
+#         x = layers.Flatten()(x)
+#         x = layers.Dense(128, activation="relu")(x)
+#         m_output = layers.Dense(self._output_dim, activation="softmax")(x)
+#         return keras.Model(m_input, m_output, name="classification_model")
+
+# class ConvnLSTMnArchitecture(BaseClassificationArchitecture):
+#     def __init__(
+#         self, seq_len: int, feat_dim: int, output_dim: int, n_conv_lstm_blocks: int = 1
+#     ) -> None:
+#         self._n_conv_lstm_blocks = n_conv_lstm_blocks
+#         super().__init__(seq_len, feat_dim, output_dim)
+#
+#     def _build_model(self) -> keras.models.Model:
+#         m_input = keras.Input((self._seq_len, self._feat_dim))
+#         x = m_input
+#         for _ in range(self._n_conv_lstm_blocks):
+#             x = layers.Conv1D(filters=64, kernel_size=3, activation="relu")(x)
+#             x = layers.Dropout(0.2)(x)
+#             x = layers.LSTM(128, activation="relu", return_sequences=True)(x)
+#             x = layers.Dropout(0.2)(x)
+#         x = layers.Flatten()(x)
+#         x = layers.Dense(128, activation="relu")(x)
+#         m_output = layers.Dense(self._output_dim, activation="softmax")(x)
+#         return keras.Model(m_input, m_output, name="classification_model")
+
+class BasicRecurrentArchitecture(Architecture):
+    arch_type = "rnn_architecture"
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        output_dim: int,
+        n_layers: int,
+        network_type: str,
+        name: str = "Sequential",
+    ) -> None:
+        """
+        :param hidden_dim: int, the number of units (e.g. 24)
+        :param output_dim: int, the number of output units (e.g. 1)
+        :param n_layers: int, the number of layers (e.g. 3)
+        :param network_type: str, one of 'gru', 'lstm', or 'lstmLN'
+        :param name: str, model name
+            Default: "Sequential"
+        """
+        self.hidden_dim = hidden_dim
+        self.output_dim = output_dim
+        self.n_layers = n_layers
+
+        self.network_type = network_type.lower()
+        assert self.network_type in ["gru", "lstm"]
+
+        self._name = name
+
+    def _rnn_cell(self) -> keras.layers.Layer:
+        """
+        Basic RNN Cell
+        :return cell: keras.layers.Layer
+        """
+        cell = None
+        # GRU
+        if self.network_type == "gru":
+            cell = keras.layers.GRUCell(self.hidden_dim, activation="tanh")
+        # LSTM
+        elif self.network_type == "lstm":
+            cell = keras.layers.LSTMCell(self.hidden_dim, activation="tanh")
+        return cell
+
+    def _make_network(self, model: keras.models.Model, activation: str, return_sequences: bool) -> keras.models.Model:
+        _cells = tf.keras.layers.StackedRNNCells(
+            [self._rnn_cell() for _ in range(self.n_layers)],
+            name=f"{self.network_type}_x{self.n_layers}",
+        )
+        model.add(keras.layers.RNN(_cells, return_sequences=return_sequences))
+        model.add(
+            keras.layers.Dense(units=self.output_dim, activation=activation, name="OUT")
+        )
+        return model
+
+    def build(self, activation: str = "sigmoid", return_sequences: bool = True) -> keras.models.Model:
+        model = keras.models.Sequential(name=f"{self._name}")
+        model = self._make_network(model, activation=activation, return_sequences=return_sequences)
+        return model
+
+
+class Zoo(dict):
+    def __init__(self, *arg, **kwargs) -> None:
+        super(Zoo, self).__init__(*arg, **kwargs)
+
+    def summary(self) -> None:
+        summary_table = PrettyTable()
+        summary_table.field_names = ["id", "type"]
+        for k, v in self.items():
+            summary_table.add_row([k, v.arch_type])
+        print(summary_table)
+
+
+zoo = Zoo(
+    {
+        # Generative models
+        "vae_conv5": VAE_CONV5Architecture,
+        "cvae_conv5": cVAE_CONV5Architecture,
+        "gan_conv_lstm": GAN_ConvLSTM4Architecture,
+        "gan_conv_lstm_2": GAN_Conv2LSTM4Architecture,
+        "gan_conv_lstm_3": GAN_Conv3LSTM4Architecture
+
+        # # Downstream models
+        # "clf_cn": ConvnArchitecture,
+        # "clf_cl_n": ConvnLSTMnArchitecture,
+       # "recurrent": BasicRecurrentArchitecture,
+    }
+)

app.py

@@ -0,0 +1,226 @@
+import gradio as gr
+import numpy as np
+import matplotlib.pyplot as plt
+from GAN.diffusion import build_model, GaussianDiffusion, DiffusionModel
+import tensorflow as tf
+from tensorflow.python.types.core import TensorLike
+import imageio
+import tempfile
+import os
+
+EPS = 1e-18
+class TSFeatureScaler:
+    """Global time series scaler that scales all features to [0,1] then normalizes to [-1,1]"""
+    
+    def __init__(self) -> None:
+        self.min_val = None
+        self.max_val = None
+        
+    def fit(self, X: TensorLike) -> "TSFeatureScaler":
+        """
+        Fit scaler to data
+        
+        Args:
+            X: Input tensor of shape [N, T, D] 
+               (N: samples, T: timesteps, D: features)
+        """
+        # 计算整个数据集的全局最大最小值
+        self.min_val = np.min(X)
+        self.max_val = np.max(X)
+        return self
+    
+    def transform(self, X: TensorLike) -> TensorLike:
+        """
+        Transform data in two steps:
+        1. Scale to [0,1] using min-max scaling
+        2. Normalize to [-1,1]
+        """
+        if self.min_val is None or self.max_val is None:
+            raise ValueError("Scaler must be fitted before transform")
+            
+        # 1. 缩放到[0,1]
+        X_scaled = (X - self.min_val) / (self.max_val - self.min_val + EPS)
+        
+        # 2. 归一化到[-1,1]
+        X_normalized = 2.0 * X_scaled - 1.0
+        
+        return X_normalized
+    
+    def fit_transform(self, X: TensorLike) -> TensorLike:
+        """Fit to data, then transform it"""
+        return self.fit(X).transform(X)
+    
+
+def create_animation(frames, fps=1):
+    """将帧列表转换为GIF动画数据"""
+    import tempfile
+    import os
+    
+    temp_dir = tempfile.gettempdir()
+    temp_path = os.path.join(temp_dir, f"temp_{id(frames)}.gif")
+    
+    # 将fps转换为duration (毫秒)
+    duration = int(1000 / fps)  # 1000ms = 1s
+    
+    # 保存为GIF文件，设置循环播放
+    imageio.mimsave(temp_path, frames, format='GIF', duration=duration, loop=0)  # loop=0 表示无限循环
+    
+    return temp_path
+
+def generate_timeseries(input_file, num_samples=16):
+    try:
+        # 加载数据
+        real_data = np.load(input_file.name)
+        scaler = TSFeatureScaler()
+        real_data = scaler.fit_transform(real_data)
+        print(f"Loaded data shape: {real_data.shape}")
+        
+        # 确保数据形状正确
+        expected_shape = (None, 96, 3)
+        if len(real_data.shape) != 3 or real_data.shape[1:] != expected_shape[1:]:
+            return None, None
+        
+        # 创建模型和必要的组件
+        network = build_model(
+            time_len=96, 
+            fea_num=3,
+            d_model=16,
+            n_heads=2,
+            encoder_type='dual'
+        )
+        ema_network = build_model(
+            time_len=96, 
+            fea_num=3,
+            d_model=16,
+            n_heads=2,
+            encoder_type='dual'
+        )
+        ema_network.set_weights(network.get_weights())  
+        noise_util = GaussianDiffusion(timesteps=10)
+        
+        print("Creating model...")
+        model = DiffusionModel(
+            network=network,
+            ema_network=ema_network,
+            timesteps=10,
+            gdf_util=noise_util,
+            data=real_data[:num_samples]
+        )
+        
+        # 加载预训练权重
+        checkpoint_path = "/Users/lindan/Dropbox/PhD/Projects/PLF/GAN/code_github/checkpoint/cp.ckpt"
+        print(f"Loading weights from {checkpoint_path}")
+        model.load_weights(checkpoint_path)
+
+        
+        # 生成加噪过程的动画
+        print("Generating noising animation...")
+        noise_frames = model.plot_noise_process_app(num_samples)
+        noise_gif = create_animation(noise_frames)
+        
+        # 生成去噪过程的动画
+        print("Generating denoising animation...")
+        denoise_frames = model.plot_denoise_process_app(num_samples)[1:]
+        denoise_gif = create_animation(denoise_frames)
+        
+        return noise_gif, denoise_gif
+        
+    except Exception as e:
+        import traceback
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_msg)
+        return None, None
+
+def update_example_gifs(num_samples):
+    """根据选择的样本数更新示例GIF"""
+    return f"noising_example_{num_samples}.gif", f"denoising_example_{num_samples}.gif"
+
+# 创建Gradio界面
+with gr.Blocks(title="Wearable Sensors Time-Series Generation") as demo:
+    with gr.Column(elem_id="container"):
+        # Logo
+        gr.Image("logo.webp", elem_id="logo", show_label=False, container=False)
+        
+        # 标题和副标题
+        gr.Markdown(
+            """
+            # Wearable Sensors Time-Series Generation
+            
+            <h3 style='font-weight: normal; color: #666;'>-- mainly targeted at livestock wearables sensors data</h3>
+            """)
+        
+    with gr.Row():
+        with gr.Column():
+            noise_gif = gr.Image(value="noising_example_16.gif", label="Noising Process", show_label=True)
+        with gr.Column():
+            denoise_gif = gr.Image(value="denoising_example_16.gif", label="Denoising Process", show_label=True)
+    
+    with gr.Row():
+        with gr.Column():
+
+            num_samples = gr.Radio(
+                choices=[4, 9, 16, 25],
+                value=16,
+                label="Number of samples to generate"
+            )
+            generate_btn = gr.Button("Generate")
+
+            # 将File组件改为Examples组件
+            input_file = gr.File(label="Select example data")
+            gr.Examples(
+                examples=[
+                    ["app_examples/example1.npy"],
+                    ["app_examples/example2.npy"],
+                    ["app_examples/example3.npy"],
+                    ["app_examples/example4.npy"]
+                ],
+                inputs=input_file,
+                label="Example Datasets"
+            )
+
+
+    
+    # 添加按钮事件处理
+    generate_btn.click(
+        fn=generate_timeseries,
+        inputs=[input_file, num_samples],
+        outputs=[noise_gif, denoise_gif]
+    )
+    
+    # 添加样本数选择的事件处理
+    num_samples.change(
+        fn=update_example_gifs,
+        inputs=[num_samples],
+        outputs=[noise_gif, denoise_gif]
+    )
+    
+    # 添加CSS样式
+    gr.HTML(
+        """
+        <style>
+        #container {
+            text-align: center;
+            padding: 2rem 0;
+        }
+        #logo {
+            width: 120px;
+            height: 120px;
+            margin: 0 auto;
+            margin-bottom: 1rem;
+        }
+        h1 {
+            font-size: 3.5rem;
+            margin-bottom: 0.5rem;
+        }
+        h3 {
+            font-size: 1.8rem;
+            margin-top: 0;
+            color: #666;
+        }
+        </style>
+        """
+    )
+
+# 启动应用
+if __name__ == "__main__":
+    demo.launch(share=True) 

app_examples/example1.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b82cbcf8248355a0c78e9bd5cf48d580363fb70e7b36fdf0bc1ec4ab96c77fe2
+size 20864

app_examples/example2.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c81e71805ffaa6b0f7c74a6046ae2d55c86136ac668fe2517a4e1646f450773
+size 18560

app_examples/example3.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:620459ceb06b59e8b4c9e24b3b4b99a06cab624a632e9f1f4de1166f0f49a1b6
+size 41600

app_examples/example4.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1291113ea1aaea7c8bc9e138abcc9bc1213741960ca23a77f6f41ab069318499
+size 20864

checkpoint/checkpoint

@@ -0,0 +1,2 @@
+model_checkpoint_path: "cp.ckpt"
+all_model_checkpoint_paths: "cp.ckpt"