<commit_before>__author__ = "<NAME>"

# ================> Python Standard  and third-party <==========
from os.path import join
from scipy.optimize import linear_sum_assignment

import logging
import numpy as np
import os
# ==================> Internal modules <==========
from nac.integrals import (
    calculate_couplings_levine, calculate_couplings_3points,
    compute_overlaps_for_coupling, correct_phases)
from nac.common import (
    Matrix, Vector, Tensor3D,
    change_mol_units, femtosec2au, retrieve_hdf5_data,
    search_data_in_hdf5, store_arrays_in_hdf5)

from noodles import (gather, schedule)
from qmworks.parsers import parse_string_xyz

# Types hint
from typing import (Callable, Dict, List, Tuple)

# Starting logger
logger = logging.getLogger(__name__)

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


def lazy_couplings(paths_overlaps: List, path_hdf5: str, project_name: str,
                   enumerate_from: int, nHOMO: int, dt: float, tracking: bool,
                   write_overlaps: bool, algorithm='levine') -> List:
    """
    Compute the Nonadibatic coupling using a 3 point approximation. See:
    The Journal of Chemical Physics 137, 22A514 (2012); doi: 10.1063/1.4738960

    or a 2Point approximation using an smoothing function:
    J. Phys. Chem. Lett. 2014, 5, 2351−2356; doi: 10.1021/jz5009449

    Notice that the states can cross frequently due to unavoided crossing and
    such crossing must be track. see:
    J. Chem. Phys. 137, 014512 (2012); doi: 10.1063/1.4732536
    """
    if tracking:
        fixed_phase_overlaps, swaps = compute_the_fixed_phase_overlaps(
            paths_overlaps, path_hdf5, project_name, enumerate_from, nHOMO)
    else:
        # Do not track the crossings
        fixed_phase_overlaps = np.stack(
            retrieve_hdf5_data(path_hdf5, paths_overlaps))
        nOverlaps, nOrbitals, _ = fixed_phase_overlaps.shape
        swaps = np.tile(np.arange(nOrbitals), (nOverlaps + 1, 1))

    # Write the overlaps in text format
    if write_overlaps: 
       logger.debug("Writing down the overlaps in ascii format")
       write_overlaps_in_ascii(fixed_phase_overlaps)

    # Compute the couplings using either the levine method
    # or the 3Points approximation
    coupling_algorithms = {'levine': (calculate_couplings_levine, 1),
                           '3points': (calculate_couplings_3points, 2)}
    # Choose an algorithm to compute the couplings
    fun_coupling, step = coupling_algorithms[algorithm]

    # Number of couplings to compute
    nCouplings = fixed_phase_overlaps.shape[0] - step + 1
<commit_msg>Fix phases also when coupling is not tracked<commit_after>__author__ = "<NAME>"

# ================> Python Standard  and third-party <==========
from os.path import join
from scipy.optimize import linear_sum_assignment

import logging
import numpy as np
import os
# ==================> Internal modules <==========
from nac.integrals import (
    calculate_couplings_levine, calculate_couplings_3points,
    compute_overlaps_for_coupling, correct_phases)
from nac.common import (
    Matrix, Vector, Tensor3D,
    change_mol_units, femtosec2au, retrieve_hdf5_data,
    search_data_in_hdf5, store_arrays_in_hdf5)

from noodles import (gather, schedule)
from qmworks.parsers import parse_string_xyz

# Types hint
from typing import (Callable, Dict, List, Tuple)

# Starting logger
logger = logging.getLogger(__name__)

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


def lazy_couplings(paths_overlaps: List, path_hdf5: str, project_name: str,
                   enumerate_from: int, nHOMO: int, dt: float, tracking: bool,
                   write_overlaps: bool, algorithm='levine') -> List:
    """
    Compute the Nonadibatic coupling using a 3 point approximation. See:
    The Journal of Chemical Physics 137, 22A514 (2012); doi: 10.1063/1.4738960

    or a 2Point approximation using an smoothing function:
    J. Phys. Chem. Lett. 2014, 5, 2351−2356; doi: 10.1021/jz5009449

    Notice that the states can cross frequently due to unavoided crossing and
    such crossing must be track. see:
    J. Chem. Phys. 137, 014512 (2012); doi: 10.1063/1.4732536
    """
    if tracking:
        fixed_phase_overlaps, swaps = compute_the_fixed_phase_overlaps(
            paths_overlaps, path_hdf5, project_name, enumerate_from, nHOMO)
    else:
        # Do not track the crossings
        mtx_0 = retrieve_hdf5_data(path_hdf5, paths_overlaps[0])
        _, dim = mtx_0.shape
        overlaps = np.stack(
            retrieve_hdf5_data(path_hdf5, paths_overlaps))
        nOverlaps, nOrbitals, _ = overlaps.shape
        swaps = np.tile(np.arange(nOrbitals), (nOverlaps + 1, 1))
        mtx_phases = compute_phases(overlaps, nOverlaps, dim)
        fixed_phase_overlaps = correct_phases(overlaps, mtx_phases)

    # Write the overlaps in text format
    if write_overlaps: 
       logger.debug("Writing down the overlaps in ascii format")
       write_overlaps_in_ascii(fixed_phase_overlaps)

    # Compute the couplings using either the levine method
    # or the 3Points approximation
    coupling_algorithms = {'levine': (calculate_couplings_levine, 1),
                           '3points': (calculate_couplings_3points, 2)}
    # Choose an algorithm to compute the couplings
    fun_coupling, step = coupling_algorithms[algorithm]

    # Number of couplings to compute
    nCouplings = fixed_phase_overlaps.shape[0] - step + 1

{'fraction_non_alphanumeric': 0.11323502639723285, 'fraction_numerical': 0.037320225741853265, 'mean_word_length': 3.703767123287671, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3691630', 'n_tokens_mistral': 1910, 'n_tokens_neox': 1692, 'n_words': 479}