Source code for ai2_kit.algorithm.aos_analysis

from MDAnalysis.lib.distances import distance_array
import MDAnalysis as mda
import numpy as np
import fire
import os


def _ecn_analysis_per_frame(frame_index, reference_group, configuration_group, universe):
    l_avg, ecn = 0, 0
    universe.trajectory[frame_index]
    distance_mat = distance_array(reference=reference_group, configuration=configuration_group, box=universe.dimensions)
    # Calculate l_av (average bond length).
    n_in, n_o = distance_mat.shape
    l_min_array = np.min(distance_mat, axis=1)
    l_av = np.divide(distance_mat,l_min_array[:, np.newaxis])
    l_av = np.exp(1.0 - np.power(l_av, 6))
    denominator = np.sum(l_av, axis=1)
    l_av = np.divide(np.multiply(distance_mat, l_av), denominator[:, np.newaxis])
    l_av = np.sum(l_av, axis=1)
    l_av = np.mean(l_av)
    # Calculate ECN (effective coordination number)
    ecn_part = l_av.copy()
    ecn_part = np.exp(1.0- np.power(np.divide(distance_mat, ecn_part), 6))
    ecn_part = np.sum(ecn_part)/n_in
    l_avg += l_av
    ecn += ecn_part
    return frame_index, l_avg, ecn


[docs] def ecn_analysis(input_traj: str, out_dir: str, center: str, ligand: str, cell: list[float]): """ Calculate ECN (effective coordination number) and l_av (average bond length). :param input_traj: path of input trajectory file :param out_dir: path of output directory :param center: central atoms, in the format of MDA atom selection string, e.g. 'name In' :param ligand: atoms that surround the central atoms, in the format of MDA atom selection string, e.g. 'name O' :param cell: cell dimensions """ universe = mda.Universe(input_traj) universe.dimensions = np.array(cell) atomgroup = universe.atoms reference_group = atomgroup.select_atoms(center) configuration_group = atomgroup.select_atoms(ligand) # TODO: this can be parallelized results = [] for k in range(universe.trajectory.n_frames): result = _ecn_analysis_per_frame(frame_index=k, reference_group=reference_group, configuration_group=configuration_group, universe=universe) results.append(result) results_arr = np.array(results) avr_arr = np.mean(results_arr, axis=0) # dump results os.makedirs(out_dir, exist_ok=True) np.savetxt(os.path.join(out_dir, 'raw.txt'), results_arr, fmt=' %d %.4f %.4f', delimiter=' ', header='frame_index l_av ECN') with open(os.path.join(out_dir, 'stats.txt'), mode='w') as fp: fp.write('\n'.join([f'l_av = {avr_arr[1]:.4f}', f'ECN = {avr_arr[2]:.4f}']))
def _count_shared_polyhedra_per_frame(frame_index, reference_group, configuration_group, universe, cutoff, coordination_num = -1): universe.trajectory[frame_index] distance_mat = distance_array(reference=reference_group, configuration=configuration_group, box=universe.dimensions) neighbor_mat = distance_mat < cutoff neighbor_mat = neighbor_mat.astype(int) # common_neighbor_mat: a matrix of coordination and shared oxygen numbers. # The former numbers are diagonal and the later are other elements. common_neighbor_mat = np.dot(neighbor_mat, neighbor_mat.T) # Counting polyhedra with particular coordination number. if coordination_num > 0: for i in range(reference_group.n_atoms): if common_neighbor_mat[i][i] != coordination_num: common_neighbor_mat[i] = 0 common_neighbor_mat[:, i] = 0 # counting common_neighbor_mat = common_neighbor_mat[np.triu_indices_from(common_neighbor_mat, k=1)] corner = np.count_nonzero(common_neighbor_mat == 1) edge = np.count_nonzero(common_neighbor_mat == 2) face = np.count_nonzero(common_neighbor_mat>=3) return frame_index, corner, edge, face
[docs] def count_shared_polyhedra(input_traj: str, out_dir: str, center: str, ligand: str, cell: list[float], cutoff: float, coord_num: int = -1): """ Count the number of shared polyhedra in each frame of the trajectory. :param input_traj: path of input trajectory file :param out_dir: path of output directory :param center: central atoms, in the format of MDA atom selection string, e.g. 'name In' :param ligand: atoms that surround the central atoms, in the format of MDA atom selection string, e.g. 'name O' :param cell: cell dimensions :param cutoff: cutoff distance :param coord_num: coordination number """ universe = mda.Universe(input_traj) universe.dimensions = np.array(cell) atomgroup = universe.atoms reference_group = atomgroup.select_atoms(center) configuration_group = atomgroup.select_atoms(ligand) # TODO: this can be parallelized results = [] for k in range(universe.trajectory.n_frames): result = _count_shared_polyhedra_per_frame(frame_index=k, reference_group=reference_group, configuration_group=configuration_group, universe=universe, cutoff=cutoff, coordination_num=coord_num) results.append(result) results_arr = np.array(results) avg_arr = np.mean(results_arr, axis=0) # dump results os.makedirs(out_dir, exist_ok=True) np.savetxt(os.path.join(out_dir, 'raw.txt'), results_arr, fmt=' %d %.4f %.4f %.4f', delimiter=' ', header='frame_index corner edge face') with open(os.path.join(out_dir, 'stats.txt'), mode='w') as fp: fp.write('\n'.join([f'Corner-share = {avg_arr[1]:.4f}', f'Edge-share = {avg_arr[2]:.4f}', f'Face-share = {avg_arr[3]:.4f}', ]))
cli_entry = { 'ecn-analysis': ecn_analysis, 'count-shared-polyhedra': count_shared_polyhedra, } if __name__ == '__main__': fire.Fire(cli_entry)