.. visualisation:

=============
Visualisation
=============

Here we detail the methods available for outputting the learned RAFFLE descriptor and RAFFLE fingerprints (distribution functions) for individual structures.


Visualising the learned RAFFLE descriptor
-----------------------------------------

RAFFLE generates 2-body, 3-body, and 4-body distribution functions for each atomic species in the system (element pairs for the 2-body function).
These are generated by combining formation energy-weighted (or convex hull-weighted) `n`-body distribution functions for each structure provided in the learning database.

The ``raffle_generator`` object has a method ``get_descriptor()`` that returns the 2-, 3-, and 4-body forms of the learned RAFFLE generalised descriptor.
The output is a list of three numpy arrays, with the first array containing the 2-body descriptor, the second array containing the 3-body descriptor, and the third array containing the 4-body descriptor.
Each `n`-body descriptor is a 2D array, with the first column containing the species index (or element pair index for the 2-body descriptor) and the second column containing the binned descriptor value.
The bin lengths are (``nbins`` component) set either explicitly or determined by the ``cutoff_min``, ``cutoff_max``, and ``width`` components of the generator.

Here is an example of how to use the ``get_descriptor()`` method:

.. code-block:: python

    # Initialise RAFFLE generator
    from raffle.generator import raffle_generator

    generator = raffle_generator()

    # Set the host structure
    host = Atoms(
        # Host structure for the generator
    )
    generator.set_host(host)

    # Set the reference energies (i.e. chemical potential references)
    generator.distributions.set_element_energies(
        {
            # reference energies for all elements in the systems
        }
    )

    # Optional parameters
    generator.distributions.set_kBT(0.2)
    generator.distributions.set_width([0.04, np.pi/160.0, np.pi/160.0])
    generator.distributions.set_cutoff_min([0.5, 0.0, 0.0])
    generator.distributions.set_cutoff_max([6.0, np.pi, np.pi])

    # Set and learn from the initial database
    database = [
        # List of structures in the learning database
    ]
    generator.distributions.create(database)

    # Retrieve the descriptor
    descriptor_init = generator.get_descriptor()

    # Print the 2-body descriptor
    print("2-body descriptor:")
    print(descriptor_init[0])

    # Print the 3-body descriptor
    print("3-body descriptor:")
    print(descriptor_init[1])

    # Print the 4-body descriptor
    print("4-body descriptor:")
    print(descriptor_init[2])

With this, we can now plot the descriptor using any plotting library of your choice.

.. code-block:: python

    import matplotlib.pyplot as plt
    import numpy as np

    # Create a figure with 3 subplots side by side
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    # Plot for each n-body descriptor (2-body, 3-body, 4-body)
    for j in range(3):
        # Calculate x-axis values
        x = np.arange(generator.distributions.cutoff_min[j],
                    generator.distributions.cutoff_max[j] + generator.distributions.width[j],
                    generator.distributions.width[j])

        # Plot on the respective subplot
        for idx in range(len(descriptor_init[j])):
            axes[j].plot(x, descriptor_init[j][idx,:])

        # Set labels and title for each subplot
        axes[j].set_ylabel('Descriptor value')
        axes[j].set_title(f'{j+2}-body descriptor')

    axes[0].set_xlabel('Distance (Å)')
    axes[1].set_xlabel('3-body angle (radians)')
    axes[2].set_xlabel('Improper dihedral angle (radians)')
    plt.tight_layout()
    plt.show()

An example python notebook is provided in :git:`examples/python_pkg/visualisation/descriptor.ipynb <examples/python_pkg/visualisation/descriptor.ipynb>`

We can now use this to compare the initial descriptor with the updated descriptor after generating new structures.

.. code-block:: python

    # Generate new structures and update the descriptor
    structures = [
        # List of structures to be generated
    ]
    generator.distributions.update(structures)

    # Retrieve the updated descriptor
    descriptor_new = generator.get_descriptor()

    # Print the updated descriptor on the plots and compare
    ...


Visualising a RAFFLE fingerprint
--------------------------------

RAFFLE fingerprints are the distribution functions for each structure in the learning database.
These are then weighted by energy (formation or convex hull) to form the RAFFLE descriptor.

However, the individual fingerprints can also be extracted and visualised.

The `raffle_generator` object has a method `generate_fingerprint()` that returns the distribution functions for a provided structure.
The output is a list of three numpy arrays, with the first array containing the 2-body fingerprint, the second array containing the 3-body fingerprint, and the third array containing the 4-body fingerprint.
Each `n`-body fingerprint is a 2D array, with the first column containing the species index (or element pair index for the 2-body fingerprint) and the second column containing the binned fingerprint value.
Like above, the bin lengths are set either explicitly or determined by the `cutoff_min`, `cutoff_max`, and `width` components of the generator.
Here is an example of how to use the `generate_fingerprint()` method:

.. code-block:: python

    # Initialise RAFFLE generator
    from raffle.generator import raffle_generator

    generator = raffle_generator()

    # Optional parameters
    generator.distributions.set_width([0.04, np.pi/160.0, np.pi/160.0])
    generator.distributions.set_cutoff_min([0.5, 0.0, 0.0])
    generator.distributions.set_cutoff_max([6.0, np.pi, np.pi])

    # Structure to obtain the fingerprint for
    structure = Atoms(
        # Structure to be used for the fingerprint
    )

    fingerprint = generator.distributions.generate_fingerprint(structure)

This can then be visualised in a similar way to the descriptor.

.. code-block:: python

    # Create a figure with 3 subplots side by side
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))

    # Plot for each n-body function (2-body, 3-body, 4-body)
    for j in range(3):
        # Calculate x-axis values
        x = np.arange(generator.distributions.cutoff_min[j],
                    generator.distributions.cutoff_max[j] + generator.distributions.width[j],
                    generator.distributions.width[j])

        # Plot on the respective subplot
        for idx in range(len(fingerprint[j])):
            axes[j].plot(x, fingerprint[j][idx,:])

        # Set labels and title for each subplot
        axes[j].set_ylabel('Fingerprint value')
        axes[j].set_title(f'{j+2}-body fingerprint')

    axes[0].set_xlabel('Distance (Å)')
    axes[1].set_xlabel('3-body angle (radians)')
    axes[2].set_xlabel('Improper dihedral angle (radians)')
    plt.tight_layout()
    plt.show()

An example python notebook is provided in  :git:`examples/python_pkg/visualisation/fingerprint.ipynb <examples/python_pkg/visualisation/fingerprint.ipynb>`.


Visualising RAFFLE probability density
--------------------------------------

RAFFLE probability density is the probability of finding a given element in a given position in the system.
This is calculated by the RAFFLE generator and can be visualised using the `get_probability_density()` method.
The output is a 2D array, with the first column containing the coordinates (spatial and species) the second column containing the binned probability density value.

A structure is provided to the `get_probability_density()`, along with a list of elements to calculate the probability density for.
These elements must be present in the RAFFLE descriptor.

An example of how to use the `get_probability_density()` method is shown below:

.. code-block:: python

    # Initialise RAFFLE generator
    from raffle.generator import raffle_generator

    generator = raffle_generator()

    generator.distributions.set_element_energies(
        {
            # reference energies for all elements in the systems
        }
    )

    database = [
        # List of structures in the learning database
    ]
    generator.distributions.create(database)

    # Structure to obtain the probability density for
    structure = Atoms(
        # Structure to be used for the probability density
    )

    species = 'SiGe'

    probability_density, grid = generator.get_probability_density(structure, species, return_grid=True)

The first index of the first column of `probability_density` is the x-coordinate, the second index is the y-coordinate, and the third index is the z-coordinate.
The fourth index is the distance between the position and the nearest atom (i.e. the void value).
The fifth index onwards is the species index (in order of the species list provided).
The second column is the site index.

For a more extensive example, see the `examples/python_pkg/visualisation/probability_density.ipynb` notebook.
This also provides a visualisation of the probability density using the `matplotlib` library.
