<commit_before>__author__ = "<NAME>"

# ================> Python Standard  and third-party <==========
from noodles import schedule  # Workflow Engine
from os.path import join

import h5py
import numpy as np
# ==================> Internal modules <==========
from nac.integrals import calculateCoupling3Points
from nac.common import (femtosec2au, retrieve_hdf5_data)
from qmworks.hdf5.quantumHDF5 import StoreasHDF5

# ==============================> Schedule Tasks <=============================


@schedule
def lazy_schedule_couplings(i, path_hdf5, dictCGFs, geometries, mo_paths, dt=1,
                            hdf5_trans_mtx=None, output_folder=None,
                            enumerate_from=0, nCouplings=None):
    """
    Calculate the non-adiabatic coupling using 3 consecutive set of MOs in
    a dynamics, using 3 consecutive geometries in atomic units.

    :param i: nth coupling calculation
    :type i: int
    :paramter dictCGFS: Dictionary from Atomic Label to basis set
    :type     dictCGFS: Dict String [CGF],
              CGF = ([Primitives], AngularMomentum),
              Primitive = (Coefficient, Exponent)
    :parameter geometries: molecular geometries stored as list of
                           namedtuples.
    :type      geometries: ([AtomXYZ], [AtomXYZ], [AtomXYZ])
    :parameter coefficients: Tuple of Molecular Orbital coefficients.
    :type      coefficients: (Matrix, Matrix, Matrix)
    :parameter dt: dynamic integration time
    :type      dt: Float (Femtoseconds)
    :param hdf5_trans_mtx: path to the transformation matrix in the HDF5 file.
    :type hdf5_trans_mtx: String
    :param output_folder: name of the path inside the HDF5 where the coupling
    is stored.
    :type output_folder: String
    :param enumerate_from: Number from where to start enumerating the folders
    create for each point in the MD
    :type enumerate_from: Int
    :returns: path to the Coupling inside the HDF5
    """
    def calc_coupling(output_path, dt):

        if hdf5_trans_mtx is not None:
            trans_mtx = retrieve_hdf5_data(path_hdf5, hdf5_trans_mtx)
        else:
            trans_mtx = None

        mos = tuple(map(lambda j:
                        retrieve_hdf5_data(path_hdf5,
                                           mo_paths[i + j][1]), range(3)))

        # Calculate the coupling among nCouplings/2 HOMOs and nCouplings/2 LUMOs.
        if nCouplings is not None:
            _, nStates = mos[0].shape
            middle, ncs = [n // 2 for n in [nStates, nCouplings]]
            lower, upper = middle - ncs, middle + ncs
            # Extrract a subset of nCouplings coefficients
            mos = tuple(map(lambda xs: xs[:, lower: upper], mos))
<commit_msg>return to previous coupling calculation scheme<commit_after>__author__ = "<NAME>"

# ================> Python Standard  and third-party <==========
from noodles import schedule  # Workflow Engine
from os.path import join

import h5py
import numpy as np
# ==================> Internal modules <==========
from nac.integrals import calculateCoupling3Points
from nac.common import (femtosec2au, retrieve_hdf5_data)
from qmworks.hdf5.quantumHDF5 import StoreasHDF5

# ==============================> Schedule Tasks <=============================


@schedule
def lazy_schedule_couplings(i, path_hdf5, dictCGFs, geometries, mo_paths, dt=1,
                            hdf5_trans_mtx=None, output_folder=None,
                            enumerate_from=0, nCouplings=None):
    """
    Calculate the non-adiabatic coupling using 3 consecutive set of MOs in
    a dynamics, using 3 consecutive geometries in atomic units.

    :param i: nth coupling calculation
    :type i: int
    :paramter dictCGFS: Dictionary from Atomic Label to basis set
    :type     dictCGFS: Dict String [CGF],
              CGF = ([Primitives], AngularMomentum),
              Primitive = (Coefficient, Exponent)
    :parameter geometries: molecular geometries stored as list of
                           namedtuples.
    :type      geometries: ([AtomXYZ], [AtomXYZ], [AtomXYZ])
    :parameter coefficients: Tuple of Molecular Orbital coefficients.
    :type      coefficients: (Matrix, Matrix, Matrix)
    :parameter dt: dynamic integration time
    :type      dt: Float (Femtoseconds)
    :param hdf5_trans_mtx: path to the transformation matrix in the HDF5 file.
    :type hdf5_trans_mtx: String
    :param output_folder: name of the path inside the HDF5 where the coupling
    is stored.
    :type output_folder: String
    :param enumerate_from: Number from where to start enumerating the folders
    create for each point in the MD
    :type enumerate_from: Int
    :returns: path to the Coupling inside the HDF5
    """
    def calc_coupling(output_path, dt):

        if hdf5_trans_mtx is not None:
            trans_mtx = retrieve_hdf5_data(path_hdf5, hdf5_trans_mtx)
        else:
            trans_mtx = None

        mos = tuple(map(lambda j:
                        retrieve_hdf5_data(path_hdf5,
                                           mo_paths[i + j][1]), range(3)))

        #  Calculate the coupling among nCouplings/2 HOMOs
        #  and nCouplings/2 LUMOs.
        if nCouplings is not None:
            _, nStates = mos[0].shape
            middle, ncs = [n // 2 for n in [nStates, nCouplings]]
            lower, upper = middle - ncs, middle + ncs
            # Extrract a subset of nCouplings coefficients
            mos = tuple(map(lambda xs: xs[:, lower: upper], mos))

{'fraction_non_alphanumeric': 0.1094518302677108, 'fraction_numerical': 0.010926971407758149, 'mean_word_length': 2.6983164983164984, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '797177', 'n_tokens_mistral': 1659, 'n_tokens_neox': 1574, 'n_words': 521}