Examples

This page collects examples included in the different modules of the VAE project.

Models

VAE.models.example_VAE

example_VAE()

Example of a VAE model.

This function demonstrates how to build a VAE model using the method :func:VAE.models.VAE.

Source code in VAE/models.py

def example_VAE():
    """Example of a VAE model.

    This function demonstrates how to build a VAE model using the method :func:`VAE.models.VAE`.
    """

    # We first define the parameters of the model:
    params = {
        'encoder_blocks': 1,
        'cond_size': 12,
        'fc_units': 48,
        'filters': 16,
        'input_shape': [16, 7],
        'latent_dim': 10,
        'trainable': ['*bn*'],
    }

    # Then we build the different parts of the model. We start with the encoder:
    encoder = Encoder(**params, name='encoder')

    # and the latent sampling layer:
    latent_sampling = LatentSampling(**params, name='latent')

    # and finally the decoder:
    decoder = Decoder(output_shape=params['input_shape'],
                      decoder_blocks=params['encoder_blocks'],
                      output_reverse=True,
                      **params,
                      name='decoder')

    # Once we have the different parts of the model, we can build the full model:
    model = VAE(encoder, decoder, latent_sampling, **params, name='VAE')

    # Let's have a look at the model:
    model.summary()

    # We can also have a look at the trainable parameters:
    collection.summary_trainable(model)

    # and plot the model:
    ks.utils.plot_model(model, show_shapes=True, dpi=75, rankdir='LR', to_file='example_VAE.png')

    return model

VAE.models.example_VAEp

example_VAEp()

Example of a VAEp model.

This function demonstrates how to build a VAEp model using the method :func:VAE.models.VAEp.

Source code in VAE/models.py

def example_VAEp():
    """Example of a VAEp model.

    This function demonstrates how to build a VAEp model using the method :func:`VAE.models.VAEp`.
    """

    # We first define the parameters of the model:
    params = {
        'encoder_blocks': 1,
        'cond_size': 12,
        'fc_units': 48,
        'filters': 16,
        'input_shape': [16, 7],
        'latent_dim': 10,
        'trainable': ['*bn*'],
        'prediction_shape': [16, 1],
    }

    # Then we build the different parts of the model. We start with the encoder:
    encoder = Encoder(**params, name='encoder')

    # and the latent sampling layer:
    latent_sampling = LatentSampling(**params, name='latent')

    # Then we build the decoder:
    decoder = Decoder(output_shape=params['input_shape'],
                      decoder_blocks=params['encoder_blocks'],
                      output_reverse=True,
                      **params,
                      name='decoder')

    # and a second decoder for the prediction:
    prediction = Decoder(output_shape=params['prediction_shape'],
                         decoder_blocks=params['encoder_blocks'],
                         output_reverse=False,
                         **params,
                         name='prediction')

    # Once we have the different parts of the model, we can build the full model:
    model = VAEp(encoder, decoder, latent_sampling, prediction, **params, name='VAEp')

    # Let's have a look at the model:
    model.summary()

    # We can also have a look at the trainable parameters:
    collection.summary_trainable(model)

    # and plot the model:
    ks.utils.plot_model(model, show_shapes=True, dpi=75, rankdir='LR', to_file='example_VAEp.png')

    return model

Layers

VAE.layers.example_Film

example_Film()

Example of Film layer.

Source code in VAE/layers.py

def example_Film():
    """Example of Film layer."""
    input_shape = (1, 16, 12)
    cond_shape = (10, )
    x = tf.constant(1., shape=(32, ) + input_shape)
    c = tf.constant(1., shape=(32, ) + cond_shape)
    x_in = ks.layers.Input(shape=input_shape)
    cond_in = ks.layers.Input(shape=cond_shape)
    out = Film(
        use_scale=True,
        use_offset=True,
        use_bias=True,
        shape=(1, None, None),
    )([x_in, cond_in])
    model = ks.Model(inputs=[x_in, cond_in], outputs=out)
    model.summary()
    _ = model.predict([x, c])
    for w in model.weights:
        print(w.name, ':', w.shape)

VAE.layers.example_GumbelSoftmax

example_GumbelSoftmax()

Example of GumbelSoftmax layer.

Source code in VAE/layers.py

def example_GumbelSoftmax():
    """Example of GumbelSoftmax layer."""
    input_shape = (5, 4)
    x = tf.zeros((2, ) + input_shape)
    x_in = ks.layers.Input(shape=input_shape)
    out = GumbelSoftmax(axis=-1, hard=True, noise_shape=(None, 1, None))(x_in)
    model = ks.Model(inputs=x_in, outputs=out)
    y = model.predict(x)
    print(y)

VAE.layers.example_RandomSampling

example_RandomSampling()

Example of RandomSampling layer.

Source code in VAE/layers.py

def example_RandomSampling():
    """Example of RandomSampling layer."""
    input_shape = (5, 4)
    z_mean = tf.zeros((2, ) + input_shape)
    z_log_var = tf.zeros((2, ) + input_shape)
    z_mean_in = ks.layers.Input(shape=input_shape)
    z_log_var_in = ks.layers.Input(shape=input_shape)
    out = RandomSampling(noise_shape=(None, 1, None))([z_mean_in, z_log_var_in])
    model = ks.Model(inputs=[z_mean_in, z_log_var_in], outputs=out)
    z = model.predict([z_mean, z_log_var])
    print(z)

Losses

VAE.losses.example_total_correlation_losses

example_total_correlation_losses()

Example of total correlation loss functions.

Source code in VAE/losses.py

def example_total_correlation_losses():
    """Example of total correlation loss functions."""
    batch_size = 32
    repeat_samples = 20
    shape = (batch_size * repeat_samples, 8)
    set_size = 7

    z_mean = tf.random.normal(shape)
    z_log_var = tf.random.normal(shape) * 0.1 - 1.
    z = z_mean + tf.exp(z_log_var * 0.5) * tf.random.normal(shape)
    z = tf.expand_dims(z, axis=1)
    z = tf.repeat(z, repeats=set_size, axis=1)

    fcns = {
        'TC loss': TotalCorrelation(z, z_mean, z_log_var),
        'TC loss between': TotalCorrelationBetween(z, z_mean, z_log_var, repeat_samples=repeat_samples),
        'TC loss within': TotalCorrelationWithin(z, z_mean, z_log_var, repeat_samples=repeat_samples),
    }

    print(f'{"Batch size":<20} {batch_size} * {repeat_samples} = {batch_size * repeat_samples}')

    for name, fcn in fcns.items():
        tc_loss = fcn(None, None)
        tc_mean = tf.reduce_mean(tc_loss)
        tc_std = tf.math.reduce_std(tc_loss)
        print(f'{name:<20} mean={tc_mean:.2f}  std={tc_std:.2f}  shape={tc_loss.shape}')

VAE.losses.example_similarity_losses

example_similarity_losses()

Example of similarity loss functions.

Source code in VAE/losses.py

def example_similarity_losses():
    """Example of similarity loss functions."""
    batch_size = 32
    repeat_samples = 5
    shape = (batch_size * repeat_samples, 1, 160, 3)

    inputs = tf.random.normal(shape)

    fcns = {
        'Sim loss': Similarity(),
        'Sim loss between': _SimilarityBetween(repeat_samples=repeat_samples),
        'Sim loss between (fast)': SimilarityBetween(repeat_samples=repeat_samples),
    }

    print(f'{"Batch size":<25} {batch_size} * {repeat_samples} = {batch_size * repeat_samples}')

    losses = []
    for name, fcn in fcns.items():
        loss = fcn(None, inputs)
        losses.append(loss)
        mean_loss = tf.reduce_mean(loss)
        std_loss = tf.math.reduce_std(loss)
        print(f'{name:<25} mean={mean_loss:.2f}  std={std_loss:.2f}  shape={loss.shape}')

    return losses

Generators

VAE.generators.example_FitGenerator

example_FitGenerator()

Example of :class:FitGenerator.

This example shows how to use the :class:FitGenerator class.

Source code in VAE/generators.py

def example_FitGenerator():
    """Example of :class:`FitGenerator`.

    This example shows how to use the :class:`FitGenerator` class.

    """

    # first we create some dummy data
    shape = (1, 32, 3)  # (set_size, data_length, channels)
    dataset = np.reshape(np.arange(np.prod(shape)), shape)
    datasets = [dataset] * 3

    # the corresponding time values
    time = range(shape[1])

    # and the corresponding conditions
    # the encoder and decoder conditions are different
    encoder_cond = np.linspace(-1, 1, 32)
    decoder_cond = np.linspace(1, -1, 32)

    # then we create the generator
    fit_gen = FitGenerator(datasets,
                           condition={
                               'encoder': encoder_cond,
                               'decoder': decoder_cond
                           },
                           input_length=1,
                           prediction_length=4,
                           batch_size=128,
                           ensemble_size=len(datasets),
                           ensemble_type='index',
                           tp_period=12,
                           time=time,
                           shuffle=False)

    # we can see the summary of the generator
    fit_gen.summary()

    # we can now use the generator to get the inputs for the model
    inputs, *_ = fit_gen[0]

    # we can plot the inputs, to see what the model will get
    # we show the encoder and decoder conditions
    fig, (lax, rax) = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(16, 5))
    lax.pcolormesh(inputs['encoder_cond'][:, 0, :])
    lax.set_title("inputs['encoder_cond']")
    mp = rax.pcolormesh(inputs['decoder_cond'][:, 0, :])
    rax.set_title("inputs['decoder_cond']")

    fig.colorbar(mp, ax=(lax, rax))

FileIO

VAE.utils.fileio.example_read_climexp_raw_data_multi

example_read_climexp_raw_data_multi()

Example of how to use the function read_climexp_raw_data_multi.

The function reads multiple files of raw data from the example_data/ folder.

example_data/icmip5_tos_Omon_one_rcp45_pc01.txt
example_data/icmip5_tos_Omon_one_rcp45_pc02.txt

Source code in VAE/utils/fileio.py

def example_read_climexp_raw_data_multi():
    """Example of how to use the function `read_climexp_raw_data_multi`.

    The function reads multiple files of raw data from the `example_data/` folder.

    ```
    example_data/icmip5_tos_Omon_one_rcp45_pc01.txt
    example_data/icmip5_tos_Omon_one_rcp45_pc02.txt
    ```

    """

    filename = [
        'example_data/icmip5_tos_Omon_one_rcp45_pc01.txt',
        'example_data/icmip5_tos_Omon_one_rcp45_pc02.txt',
    ]

    # read data
    df, metadata = read_climexp_raw_data_multi(filename, ensemble_members=[0, 1, 2, 3, 4, 5], join='outer')

    with pd.option_context('display.precision', 3):
        print(df)

    _, ax = plt.subplots(1, 1, figsize=(10, 5))
    time = df.index.to_numpy()

    # access data by level-zero index
    for idx in df.columns.levels[0]:
        x = df[idx]
        x = x.to_numpy()
        x_mean = x.mean(axis=1)
        x_std = x.std(axis=1) * 3

        ax.plot(time, x_mean, label=metadata[idx].get('description'), zorder=2.2)
        ax.fill_between(time, x_mean - x_std, x_mean + x_std, alpha=0.5, zorder=2.1)

    ax.legend(loc='upper left')
    ax.grid(linestyle=':')

VAE.utils.fileio.example1_read_netcdf

example1_read_netcdf()

Example of how to use the function read_netcdf.

This function reads EOFs in a netCDF file from the Climate explorer from the example_data/ folder:

example_data/eofs_icmip5_tos_Omon_one_rcp45.nc

Source code in VAE/utils/fileio.py

def example1_read_netcdf():
    """Example of how to use the function `read_netcdf`.

    This function reads EOFs in a netCDF file from the Climate explorer from the `example_data/` folder:

    ```
    example_data/eofs_icmip5_tos_Omon_one_rcp45.nc
    ```

    """

    filename = 'example_data/eofs_icmip5_tos_Omon_one_rcp45.nc'

    # read data
    variables, dimensions, attributes = read_netcdf(filename)

    print('variables')
    for key, value in variables.items():
        print('  ', key, value.shape)

    print('dimensions')
    for key, value in dimensions.items():
        print('  ', key, value.shape)

    print('attributes')
    pprint(attributes)

    # remove singleton dimensions
    squeeze_variables = {key: np.squeeze(value) for key, value in variables.items()}
    # plot variables
    vplt.map_plot(dimensions['lat'], dimensions['lon'], squeeze_variables, ncols=5, figwidth=20, cmap='seismic')

VAE.utils.fileio.example2_read_netcdf

example2_read_netcdf()

Example of how to use the function read_netcdf.

This function reads EOFs in a netCDF file from the output of the climate data operators (CDO). The file is from the example_data/ folder:

example_data/eofs_anom_gpcc_v2020_1dgr.nc

:material-github: For the calculation of the EOFs and PCs with CDO see the CDO scripts.

Source code in VAE/utils/fileio.py

def example2_read_netcdf():
    """Example of how to use the function `read_netcdf`.

    This function reads EOFs in a netCDF file from the output of the climate data operators (CDO). The file is from the
    `example_data/` folder:

    ```
    example_data/eofs_anom_gpcc_v2020_1dgr.nc
    ```

    :material-github: For the calculation of the EOFs and PCs with CDO see the [CDO
        scripts](https://github.com/andr-groth/cdo-scripts).

    """

    filename = 'example_data/eofs_anom_gpcc_v2020_1dgr.nc'

    # read data
    variables, dimensions, attributes = read_netcdf(filename)

    print('variables')
    for key, value in variables.items():
        print('  ', key, value.shape)

    print('dimensions')
    for key, value in dimensions.items():
        print('  ', key, value.shape)

    print('attributes')
    pprint(attributes)

    # plot first variable
    key, *_ = list(variables)
    variable = variables[key]
    vplt.map_plot(dimensions['lat'], dimensions['lon'], variable, ncols=10, figwidth=20, cmap='seismic')

VAE.utils.fileio.example3_read_netcdf

example3_read_netcdf()

Example of how to use the function read_netcdf.

This function reads PCs in a netCDF file from the output of the climate data operators (CDO). The file is from the example_data/ folder:

example_data/pcs_anom_gpcc_v2020_1dgr.nc

:material-github: For the calculation of the EOFs and PCs with CDO see the CDO scripts.

Source code in VAE/utils/fileio.py

def example3_read_netcdf():
    """Example of how to use the function `read_netcdf`.

    This function reads PCs in a netCDF file from the output of the climate data operators (CDO). The file is from the
    `example_data/` folder:

    ```
    example_data/pcs_anom_gpcc_v2020_1dgr.nc
    ```

    :material-github: For the calculation of the EOFs and PCs with CDO see the [CDO
        scripts](https://github.com/andr-groth/cdo-scripts).

    """

    filename = 'example_data/pcs_anom_gpcc_v2020_1dgr.nc'

    # read data
    variables, dimensions, attributes = read_netcdf(filename, num2date=True)

    print('variables')
    for key, value in variables.items():
        print('  ', key, value.shape)

    print('dimensions')
    for key, value in dimensions.items():
        print('  ', key, value.shape)

    print('attribtutes')
    pprint(attributes)

    # plot first variable
    key, *_ = list(variables)
    variable = np.squeeze(variables[key])  # remove singleton spatial dimensions
    variable = variable.T
    variable = np.atleast_2d(variable)
    cols = min(len(variable), 5)
    rows = -(-len(variable) // cols)
    _, axs = plt.subplots(rows, cols, figsize=(4 * cols, 2 * rows), sharex=True, sharey=True, squeeze=False)
    for n, (ax, value) in enumerate(zip(axs.flatten(), variable)):
        ax.plot(dimensions['time'], value)
        ax.set_title(n)

VAE.utils.fileio.example_read_netcdf_multi

example_read_netcdf_multi(filename)

Example of how to use the function read_netcdf_multi.

This function reads PCs in multiple netCDF files from the output of the climate data operators (CDO). The files are from the example_data/ folder:

example_data/pcs_anom_pr_*.nc

:material-github: For the calculation of the EOFs and PCs with CDO see the CDO scripts.

Source code in VAE/utils/fileio.py

def example_read_netcdf_multi(filename: str):
    """Example of how to use the function `read_netcdf_multi`.

    This function reads PCs in multiple netCDF files from the output of the climate data operators (CDO). The files are
    from the `example_data/` folder:

    ```
    example_data/pcs_anom_pr_*.nc
    ```

    :material-github: For the calculation of the EOFs and PCs with CDO see the [CDO
        scripts](https://github.com/andr-groth/cdo-scripts).

    """

    filename = 'example_data/pcs_anom_pr_*.nc'

    # read data
    variables, dimensions, attributes = read_netcdf_multi(filename, num2date=True)

    print('variables')
    for key, value in variables.items():
        print('  ', key)
        for k, v in value.items():
            print('    ', k, v.shape)

    print('dimensions')
    for key, value in dimensions.items():
        print('  ', key)
        for k, v in value.items():
            print('    ', k, v.shape)

    rows = 5
    cols = 2
    _, axs = plt.subplots(rows, cols, figsize=(8 * cols, 3 * rows), sharex=True, sharey=True, squeeze=False)
    # cycle through different files
    for key, variable in variables.items():
        var_name, *_ = list(variable)  # extract first variable
        values = np.squeeze(variable[var_name])  # remove singleton spatial dimensions
        values = values.T
        values = np.atleast_2d(values)
        for n, (ax, value) in enumerate(zip(axs.flatten(), values)):
            ax.plot(dimensions[key]['time'], value, label=key)
            ax.set_title(n)

    axs.flat[0].legend()