# Get Started

## Installation

To install GMMVI (optionally) create a virtual environment and run

    (.venv) $ pip install .

## Usage

For performing the optimization, you can directly instantiate a `GMMVI` and run `GMMVI.train_iter()` in a loop, or, for adding basic logging capability and easier integration, for example with WandB, you can instantiate a `GmmviRunner` and run `GmmviRunner.iterate_and_log(n)` in a loop.

### Directly Using GMMVI

Before instantiating the `GMMVI`, we need to create several other objects, namely:

1. A `wrapped model` which stores the parameters of the GMM, as well as component-specific meta-information (reward histories, learning-rates, etc.)

2. A `SampleDB` for storing samples.

3. One object for each of the seven Design Choices.

Fortunately, each of these classes and also `GMMVI` itself, have a static method called build_from_config(), which allows to create the object from a common config dictionary (which can be created from a YAML file). Using a common dictionary is recommended, to ensure that the parameters passed to the different constructors are consistent (e.g. the sample dimensions needs to be the same).

It is easiest to directly use `GMMVI.build_from_config`, which will automatically construct most of the required objects. However, you still need to pass

1. the dictionary containing the hyperparameters,

2. the `target distribution`,

3. and the `initial model`.

The following example script directly uses `GMMVI` using the hyperparameters from the following YAML file: examples/example_config.yml.

    import os
    import logging
    # Tensorflow may give warnings when the Cholesky decomposition fails.
    # However, these warning can usually be ignored because the NgBasedOptimizer
    # will handle them by rejecting the update and decreasing the stepsize for
    # the failing component. To keep the console uncluttered, we suppress warnings.
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # ERROR
    logging.getLogger('tensorflow').setLevel(logging.ERROR)
    import tensorflow as tf

    from gmmvi.optimization.gmmvi import GMMVI
    from gmmvi.configs import load_yaml
    from gmmvi.experiments.target_distributions.logistic_regression import make_breast_cancer
    from gmmvi.models.full_cov_gmm import FullCovGMM
    from gmmvi.models.gmm_wrapper import GmmWrapper
    from gmmvi.experiments.setup_experiment import construct_initial_mixture

    #For creating a GMMVI object using GMMVI.build_from_config, we need:
    # 1. A dictionary containing the hyperparameters
    my_path = os.path.dirname(os.path.realpath(__file__))
    config = load_yaml(os.path.join(my_path, "example_config.yml"))

    # 2. A target distribution
    target_distribution = make_breast_cancer()

    # 3. An (wrapped) initial model
    dims = target_distribution.get_num_dimensions()
    initial_weights = tf.ones(1, tf.float32)
    initial_means = tf.zeros((1, dims), tf.float32)
    initial_covs = tf.reshape(100 * tf.eye(dims), [1, dims, dims])
    model = FullCovGMM(initial_weights, initial_means, initial_covs)
    # Above config contains a section model_initialization, and, therefore,
    # we could also create the initial model using:
    # model = construct_initial_mixture(dims, **config["model_initialization"])
    wrapped_model = GmmWrapper.build_from_config(model=model, config=config)

    # Now we can create the GMMVI object and start optimizing
    gmmvi = GMMVI.build_from_config(config=config,
                                    target_distribution=target_distribution,
                                    model=wrapped_model)
    max_iter = 1001
    for n in range(max_iter):
        gmmvi.train_iter()

        if n % 100  == 0:
            samples = gmmvi.model.sample(1000)[0]
            elbo = tf.reduce_mean(target_distribution.log_density(samples)
                                  - model.log_density(samples))
            print(f"{n}/{max_iter}: "
                  f"The model now has {gmmvi.model.num_components} components "
                  f"and an elbo of {elbo}.")

The script can be found under examples/1_directly_using_gmmvi.py.

### Using the GmmviRunner

The `GmmviRunner` wraps around `GMMVI` to add logging capabilities. Furthermore, the `GmmviRunner` takes care of initializing the model and the target distribution (when using one of the provided target distributions). Hence, we only need to provide the config to create it, as shown by the following script:

    import os
    from gmmvi.gmmvi_runner import GmmviRunner
    from gmmvi.configs import load_yaml

    my_path = os.path.dirname(os.path.realpath(__file__))
    config = load_yaml(os.path.join(my_path, "example_config.yml"))
    gmmvi_runner = GmmviRunner.build_from_config(config)

    for n in range(10001):
        gmmvi_runner.iterate_and_log(n)

The script can be found under examples/2_using_the_gmmvi_runner.py.

### Using the GmmviRunner with Default Configs

We can also directly create a default config based on the 7-letter Codeword to specify the design choices, thereby, not requiring an external YAML file:

    from gmmvi.gmmvi_runner import GmmviRunner
    import gmmvi.configs

    # In this example, we will create the config for a GmmviRunner using default configs
    # for a given Codename (we weill use SAMYROX) and an and an environment name
    # (we will use GMM20).
    # Let's first get the default config for SAMYROX
    algorithm_config = gmmvi.configs.get_default_algorithm_config("SAMYROX")

    # Internally, this loaded the yaml files in gmmvi/configs/module_configs corresponding
    # to the chosen design choices and stored them in a single dict "algorithm_config".
    # Note that these default values were chosen independently for every design choice,
    # and, thus, may not always be sensible. For example, the initial_stepsize defined in
    # gmmvi/configs/module_configs/component_stepsize_adaptation/improvement_based.yml
    # (Codeletter "R") is suitable if the stepsize is treated as a trust-region
    # (Codeletter "T"), but not if it directly corresponds to the stepsize
    # (Codeletter "I" or "Y")! Hence, we will overwrite the stepsize to something more
    # suitable for SAMYROX:
    better_stepsize_config = {
       'initial_stepsize': 0.0001,
       'min_stepsize': 0.0001,
       'max_stepsize': 0.001
    }
    algorithm_config = gmmvi.configs.update_config(algorithm_config, better_stepsize_config)

    # We will use a target distribution that was shipped with the framework, namely "gmm20":
    environment_config = gmmvi.configs.get_default_experiment_config("gmm20")

    # The last call searched configs/experiment_configs for a corresponding yml-file and found
    # gmm20.yml and stored the config in the dictionary "environment_config". We now just need
    # to merge both config files:
    config = gmmvi.configs.update_config(algorithm_config, environment_config)

    # Create the GmmviRunner and start optimizing.
    gmmvi_runner = GmmviRunner.build_from_config(config=config)
    for n in range(1500):
        gmmvi_runner.iterate_and_log(n)

The script can be found under examples/3_gmmvi_runner_with_default_configs.py.

### Using the GmmviRunner with Custom Environments

We can still use the `GmmviRunner` with custom environments, but we need to store the `target distribution object` in the config:

    from gmmvi.gmmvi_runner import GmmviRunner
    from gmmvi.configs import get_default_algorithm_config, update_config
    import tensorflow as tf
    import numpy as np
    import matplotlib
    matplotlib.use("tkAgg")
    import matplotlib.pyplot as plt

    # For creating a custom environment, we need to extend
    # gmmvi.experiments.target_distributions.lnpdf.LNPDF:
    from gmmvi.experiments.target_distributions.lnpdf import LNPDF
    class Rosenbrock(LNPDF):
        """ We treat the negative Rosenbrock function as unnormalized target distribution.
        We implement it in numpy and do not allow GMMVI to backpropagate through log_density().
        As we want to use Stein's Lemma for estimating the natural gradient (Codeletter "S"),
        we need to implement the gradient ourselves, and, therefore, we set
        use_log_density_and_grad=True and implement the corresponding method.
        """
        def __init__(self):
            super(Rosenbrock, self).__init__(use_log_density_and_grad=True,
                                             safe_for_tf_graph=False)
            self.a = 1
            self.b = 100

        def get_num_dimensions(self) -> int:
            return 2

        def log_density(self, samples: tf.Tensor) -> tf.Tensor:
            x = samples[:, 0].numpy().astype(np.float32)
            y = samples[:, 1].numpy().astype(np.float32)
            my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
            return tf.convert_to_tensor(my_log_density, dtype=tf.float32)

        def log_density_and_grad(self, samples: tf.Tensor) -> tf.Tensor:
            x = samples[:, 0].numpy().astype(np.float32)
            y = samples[:, 1].numpy().astype(np.float32)
            my_log_density = -((self.a - x)**2 + self.b * (y - x**2)**2)
            my_grad_x = -(-2 * (self.a - x) - 4 * self.b * (y - x**2) * x)
            my_grad_y = -(2 * self.b * (y - x**2))
            my_grad = np.vstack((my_grad_x, my_grad_y)).T
            return [tf.convert_to_tensor(my_log_density, dtype=tf.float32),
                   tf.convert_to_tensor(my_grad, dtype=tf.float32)]

    # We can also use the GmmviRunner, when using custom environments, but we have
    # to put the LNPDF object into the dict. Furthermore, we need to define the other
    # environment-specific settings that would otherwise be defined in
    # the corresponding config in gmmvi/config/experiment_configs:
    environment_config = {
        "target_fn": Rosenbrock(),
        "start_seed": 0,
        "environment_name": "Rosenbrock",
        "model_initialization": {
            "use_diagonal_covs": False,
            "num_initial_components": 1,
            "prior_mean": 0.,
            "prior_scale": 1.,
            "initial_cov": 1.,
        },
        "gmmvi_runner_config": {
            "log_metrics_interval": 100
        },
        "use_sample_database": True,
        "max_database_size": int(1e6),
        "temperature": 1.
    }

    # We will again use the automatically generated config for the algorithm,
    # but this time, we will use "SAMTRUX". The default settings are reasonable for
    # SAMTRUX, so we do not make any modifications to the hyperparameters.
    algorithm_config = get_default_algorithm_config("SAMTRUX")

    # Now we just need to merge the configs and use GmmviRunner as before:
    merged_config = update_config(algorithm_config, environment_config)
    gmmvi_runner = GmmviRunner.build_from_config(merged_config)

    for n in range(500):
        gmmvi_runner.iterate_and_log(n)

    # Plot samples from our "Rosenbrock-distribution"
    test_samples = gmmvi_runner.gmmvi.model.sample(10000)[0]
    plt.plot(test_samples[:, 0], test_samples[:, 1], 'x')
    plt.show()
    plt.pause(0.1)

The script can be found under examples/4_gmmvi_runner_with_custom_environments.py.