.. _tutorial_es_data:

Data with excited states
========================

MLatom has a special data format for excited-state information. The general formats are the same as usual, i.e., everything revolves around Python classes :class:`molecule <mlatom.data.molecule>`, :class:`molecular_database <mlatom.data.molecular_database>`, :class:`molecular_trajectory <mlatom.data.molecular_trajectory>` which can be dumped, loaded, and have many useful properties (see the :ref:`dedicated tutorial <tutorial_data>`). Here we explain how to access and add excited-state proporties in such data formats, which is essential to know, e.g., when you want to train an ML model. Interfaces to third-party electronic-structure programsa nd ML models for excited states must also conform with this data format.

First of all, ``molecule.electronic_states`` stores information about different electronic state properties. The advantage of this format is that it allows to access a property of any state the same way as accessing a property of a 'normal' molecule object. For example:

.. code-block:: python

    # assume we obtained electronic state properties for some mlatom.data.molecule() instance mol
    # let's get S1 state
    S1 = mol.electronic_states[1]
    print(f'Energy of S1 state is {S1.energy} Hartree')
    print(f'Its forces are (in Hartree/Angstrom)')
    print(-S1.energy_gradients)
    print(f'checking multiplicity: {S1.multiplicity}')

Often, we want to extract or add information involving a pair of states, e.g., excitation energy (from the ground to excited state) or oscillator strengths. These special properties can be accessed or added as:

.. code-block:: python

    mol.excitation_energies = np.array([0.01, 0.003]) * mlatom.constants.eV2hartree # excitation energies must be in hartree
    mol.oscillator_strengths = np.array([0.5, 0.001])

We can also access and add couplings, e.g., nonadiabatic couplings as ``molecule.nonadiabatic_coupling_vectors`` where they are matrices corresponding to each pair of state and each atom:

If you need to transform data into MLatom format, the most difficult part might be how to create electronic states. Here is a simple example:

.. code-block:: python

    import mlatom as ml
    mol = ml.molecule()
    # get some coordinates, etc. (can be done later too)
    ...
    mol.electronic_states = [mol.copy(atomic_labels=[], molecular_labels=[]) for ii in range(mol.nstates)]
    # example of getting some properties
    ground_state_energy = ...
    first_excitation_energy = ...
    mol.electronic_states[0].energy = ground_state_energy
    mol.electronic_states[1].energy = ground_state_energy + first_excitation_energy # all in hartree!
    # or
    first_excited_state_energy = ground_state_energy + first_excitation_energy
    mol.electronic_states[1].energy = first_excited_state_energy
    # this scheme is particular useful if you have different spin states for different states:
    mol.electronic_states[0].multiplicity = 1 # ground state is S0, singlet
    mol.electronic_states[1].multiplicity = 3 # first excited states is T1, triplet.
    # a small trick is that mlatom will populate electronic state energies
    # if mol.electronic_states[0].energy is available and we assign excitation_energies (in hartree)
    mol.excitation_energies = np.array([0.01, 0.003]) * mlatom.constants.eV2hartree
    # you can check if they are equal
    mol.electronic_state[2].energy == ground_state_energy + mol.excitation_energies[1]