deepkoopman_model_definitions.py

import tensorflow as tf

# Hyperparameters
INPUT_DIM = 2  # Same as decoder output dimension
HIDDEN_DIM = 30
LATENT_DIM = 2

# Encodes input to low-dimensional code
encoder = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(HIDDEN_DIM, activation=tf.nn.relu, name='encoder_hidden'),
        tf.keras.layers.Dense(LATENT_DIM, activation=tf.nn.sigmoid, name='code')
    ],
    name='encoder'
)

# Decodes from low-dimensional code to output
decoder = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(HIDDEN_DIM, activation=tf.nn.relu, name='decoder_hidden'),
        tf.keras.layers.Dense(INPUT_DIM, activation=tf.keras.activations.linear, name='decoder_output_flat'),
    ],
    name='decoder'
)

# Linear dynamics `K` 
linear_dynamics = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(LATENT_DIM, activation='linear', name='linear_dynamics')
    ],
    name='linear_dynamics'
)