<commit_before>__author__ = "<NAME>"

__all__ = ["workflow_compute_cubes"]

# ================> Python Standard  and third-party <==========
# from .initialization import read_time_dependent_coeffs
from collections import namedtuple
from functools import partial
from multiprocessing import (cpu_count, Pool)
from nac.common import (
    Array, Matrix, Tensor3D, Vector, change_mol_units, retrieve_hdf5_data,
    store_arrays_in_hdf5)
from nac.integrals.multipoleIntegrals import calcOrbType_Components
from nac.schedule.components import calculate_mos
from noodles import (gather)
from os.path import join
from qmworks import run
from qmworks.parsers import parse_string_xyz
from typing import (Dict, List, Tuple)
import numpy as np
import logging
import sparse

# Data types
GridCube = namedtuple("GridCube", ("voxel", "shape"))
CGFS = namedtuple("CGFS", ("primitives", "ang_expo"))

# Get logger
logger = logging.getLogger(__name__)


def workflow_compute_cubes(
        package_name: str, project_name: str, package_args: Dict,
        path_time_coeffs: str=None, grid_data: Tuple=None,
        guess_args: Dict=None, geometries: List=None,
        dictCGFs: Dict=None, calc_new_wf_guess_on_points: str=None,
        path_hdf5: str=None, enumerate_from: int=0, package_config: Dict=None,
        traj_folders: List=None, work_dir: str=None, basisname: str=None,
        hdf5_trans_mtx: str=None, nHOMO: int=None, algorithm='levine',
        ignore_warnings=False, **kwargs) -> None:
    """
    :param path_hdf5: Path to the HDF5 file that contains the
    numerical results.
    :param geometries: List of string cotaining the molecular geometries.
    """
    # Molecular orbital calculations calculatios
    mo_paths_hdf5 = calculate_mos(
        package_name, geometries, project_name, path_hdf5, traj_folders,
        package_args, guess_args, calc_new_wf_guess_on_points,
        enumerate_from, package_config=package_config,
        ignore_warnings=ignore_warnings)

    # geometries in atomic units
    molecules_au = [center_molecule_in_cube(center_molecule_in_cube(
        grid_data, change_mol_units(parse_string_xyz(gs))))
        for gs in geometries]

    # Extract the information of the CGFs dictionary
    dictCGFs_array = {l: CGFS(
        get_primitives(dictCGFs, l), get_angular_exponents(dictCGFs, l))
        for l in dictCGFs.keys()}

    # Retrieve the matrix to transform from Cartesian to spherical coordinates
    trans_mtx = sparse.csr_matrix(retrieve_hdf5_data(path_hdf5, hdf5_trans_mtx))

    # Compute the values in the given grid
    promised_grids = [compute_grid_density(
        k, project_name, grid_data, path_hdf5, dictCGFs_array,
        trans_mtx, mol, path_mo)
        for k, mol, path_mo in enumerate(zip(molecules_au, mo_paths_hdf5))]

    # # Compute the density weighted by the population computed with PYXAID
    # # Time-dependent coefficients
    # time_depend_coeffs = read_time_dependent_coeffs(path_time_coeffs)
    # msg = "Reading time_dependent coefficients from: {}".format(path_time_coeffs)
    # logger.info(msg)

    # scheduled_density = schedule(compute_dynamic_density)
    # promised_dynamic_density = scheduled_density(
    #     path_hdf5, promised_grids, time_depend_coeffs)

    # Execute the workflow
    path_grids = run(gather(*promised_grids), folder=work_dir)

    print_grids(path_hdf5, path_grids)


def compute_dynamic_density(
        path_hdf5: str, promised_grids: List, time_depend_coeffs: Matrix):
    pass


def compute_grid_density(
        k: int, project_name: str, grid_data: Tuple, path_hdf5: str,
        time_depend_coeffs: Matrix, dictCGFs_array: Dict, trans_mtx: Matrix,
        mol: List, path_mo: str) -> str:
    """
    Compute the grid density for a given geometry and store it in the HDF5

    :param grid_data: Tuple containing the grid specification
    :param path_hdf5: Path to the HDF5 file that contains the
    numerical results.
    :paramter dictCGFS: Dictionary from Atomic Label to basis set
    :type     dictCGFS: Dict String [CGF],
              CGF = ([Primitives], AngularMomentum),
              Primitive = (Coefficient, Exponent)
    :param mol: List of atomic labels and cartesian coordinates
    :param path_mo: Path to the MO in the HDF5.
    :returns: path where the grid is stored in the HDF5
    """
    path_grid = join(project_name, 'density_grid_{}'.format(k))
    # Nuclear coordinates of the grid
    grid_coordinates = create_grid_nuclear_coordinates(grid_data, mol)

    # Atomic symbols
    symbols = [at.symbol for at in mol]

    # Compute the values of the orbital using all the Avialable CPUs
    # Before multiplying for the MO coefficient
    orbital_grid = distribute_grid_computation(
        symbols, grid_coordinates, dictCGFs_array)

    # |phi| ^ 2
    orbital_grid *= orbital_grid

    # Read the molecular orbitals from the HDF5
    css = retrieve_hdf5_data(path_hdf5, path_mo)

    # Ci ^ 2
    css *= css

    # Ci^2 * |phi| ^ 2. Matrix shape: points ** 3, number_of_orbitals
    density_grid = np.dot(orbital_grid, css)

    # Store the density grid in the HDF5
    store_arrays_in_hdf5(path_hdf5, path_grid, density_grid)

    return path_grid


def distribute_grid_computation(
        symbols: List, grid_coordinates: Tensor3D,
        dictCGFs_array: Tensor3D) -> Tensor3D:
    """
    Use all the available CPUs to compute the grid of density for a given
    molecular geometry.
    """<commit_msg>used th sparse representation of the trans_mtx to compute the grid_sphericals<commit_after>__author__ = "<NAME>"

__all__ = ["workflow_compute_cubes"]

# ================> Python Standard  and third-party <==========
# from .initialization import read_time_dependent_coeffs
from collections import namedtuple
from functools import partial
from multiprocessing import (cpu_count, Pool)
from nac.common import (
    Array, Matrix, Tensor3D, Vector, change_mol_units, retrieve_hdf5_data,
    store_arrays_in_hdf5)
from nac.integrals.multipoleIntegrals import calcOrbType_Components
from nac.schedule.components import calculate_mos
from noodles import (gather)
from os.path import join
from qmworks import run
from qmworks.parsers import parse_string_xyz
from typing import (Dict, List, Tuple)
import numpy as np
import logging
import sparse

# Data types
GridCube = namedtuple("GridCube", ("voxel", "shape"))
CGFS = namedtuple("CGFS", ("primitives", "ang_expo"))

# Get logger
logger = logging.getLogger(__name__)


def workflow_compute_cubes(
        package_name: str, project_name: str, package_args: Dict,
        path_time_coeffs: str=None, grid_data: Tuple=None,
        guess_args: Dict=None, geometries: List=None,
        dictCGFs: Dict=None, calc_new_wf_guess_on_points: str=None,
        path_hdf5: str=None, enumerate_from: int=0, package_config: Dict=None,
        traj_folders: List=None, work_dir: str=None, basisname: str=None,
        hdf5_trans_mtx: str=None, nHOMO: int=None, algorithm='levine',
        ignore_warnings=False, **kwargs) -> None:
    """
    :param path_hdf5: Path to the HDF5 file that contains the
    numerical results.
    :param geometries: List of string cotaining the molecular geometries.
    """
    # Molecular orbital calculations calculatios
    mo_paths_hdf5 = calculate_mos(
        package_name, geometries, project_name, path_hdf5, traj_folders,
        package_args, guess_args, calc_new_wf_guess_on_points,
        enumerate_from, package_config=package_config,
        ignore_warnings=ignore_warnings)

    # geometries in atomic units
    molecules_au = [center_molecule_in_cube(center_molecule_in_cube(
        grid_data, change_mol_units(parse_string_xyz(gs))))
        for gs in geometries]

    # Extract the information of the CGFs dictionary
    dictCGFs_array = {l: CGFS(
        get_primitives(dictCGFs, l), get_angular_exponents(dictCGFs, l))
        for l in dictCGFs.keys()}

    # Retrieve the matrix to transform from Cartesian to spherical coordinates
    trans_mtx = sparse.csr_matrix(
        retrieve_hdf5_data(path_hdf5, hdf5_trans_mtx))

    # Compute the values in the given grid
    promised_grids = [compute_grid_density(
        k, project_name, grid_data, path_hdf5, dictCGFs_array,
        trans_mtx, mol, path_mo)
        for k, mol, path_mo in enumerate(zip(molecules_au, mo_paths_hdf5))]

    # # Compute the density weighted by the population computed with PYXAID
    # # Time-dependent coefficients
    # time_depend_coeffs = read_time_dependent_coeffs(path_time_coeffs)
    # msg = "Reading time_dependent coefficients from: {}".format(path_time_coeffs)
    # logger.info(msg)

    # scheduled_density = schedule(compute_dynamic_density)
    # promised_dynamic_density = scheduled_density(
    #     path_hdf5, promised_grids, time_depend_coeffs)

    # Execute the workflow
    path_grids = run(gather(*promised_grids), folder=work_dir)

    print_grids(path_hdf5, path_grids)


def compute_dynamic_density(
        path_hdf5: str, promised_grids: List, time_depend_coeffs: Matrix):
    pass


def compute_grid_density(
        k: int, project_name: str, grid_data: Tuple, path_hdf5: str,
        time_depend_coeffs: Matrix, dictCGFs_array: Dict, trans_mtx: Matrix,
        mol: List, path_mo: str) -> str:
    """
    Compute the grid density for a given geometry and store it in the HDF5

    :param grid_data: Tuple containing the grid specification
    :param path_hdf5: Path to the HDF5 file that contains the
    numerical results.
    :paramter dictCGFS: Dictionary from Atomic Label to basis set
    :type     dictCGFS: Dict String [CGF],
              CGF = ([Primitives], AngularMomentum),
              Primitive = (Coefficient, Exponent)
    :param mol: List of atomic labels and cartesian coordinates
    :param path_mo: Path to the MO in the HDF5.
    :returns: path where the grid is stored in the HDF5
    """
    path_grid = join(project_name, 'density_grid_{}'.format(k))
    # Nuclear coordinates of the grid
    grid_coordinates = create_grid_nuclear_coordinates(grid_data, mol)

    # Atomic symbols
    symbols = [at.symbol for at in mol]

    # Compute the values of the orbital using all the Avialable CPUs
    # Before multiplying for the MO coefficient
    orbital_grid_cartesian = distribute_grid_computation(
        symbols, grid_coordinates, dictCGFs_array)

    # Transform to spherical coordinate
    transpose = trans_mtx.transpose()
    orbital_grid_sphericals = sparse.csr_matrix.dot(
        orbital_grid_cartesian, transpose)

    # |phi| ^ 2
    orbital_grid_sphericals *= orbital_grid_sphericals

    # Read the molecular orbitals from the HDF5
    css = retrieve_hdf5_data(path_hdf5, path_mo)

    # Ci ^ 2
    css *= css

    # Ci^2 * |phi| ^ 2. Matrix shape: points ** 3, number_of_orbitals
    density_grid = np.dot(orbital_grid_sphericals, css)

    # Store the density grid in the HDF5
    store_arrays_in_hdf5(path_hdf5, path_grid, density_grid)

    return path_grid


def distribute_grid_computation(
        symbols: List, grid_coordinates: Tensor3D,
        dictCGFs_array: Tensor3D) -> Tensor3D:
    """
    Use all the available CPUs to compute the grid of density for a given
    molecular geometry.
    """
{'fraction_non_alphanumeric': 0.0797392624341459, 'fraction_numerical': 0.006786320207161353, 'mean_word_length': 3.923076923076923, 'pattern_counts': {'":': 0, '<': 7, '<?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': '3372601', 'n_tokens_mistral': 3720, 'n_tokens_neox': 3565, 'n_words': 993}