ai2_kit.feat.spectrum.algorithm module

ai2_kit.feat.spectrum.algorithm module#

ai2_kit.feat.spectrum.algorithm.FT(DT: float, C: ndarray, M: Optional[int] = None) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.FT_sin(DT: float, C: ndarray, M: Optional[int] = None) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.apply_gussian_filter(corr: ndarray, width: float)[source]#
ai2_kit.feat.spectrum.algorithm.apply_lorenz_filter(corr: ndarray, width: float, dt)[source]#
ai2_kit.feat.spectrum.algorithm.box_shift(dx: ndarray, cells: ndarray) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.cal_corr_sfg_method1(atomic_polar: ndarray, atomic_dipole: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, window: int, rc: float = 6.75)[source]#
ai2_kit.feat.spectrum.algorithm.cal_corr_sfg_method2(atomic_polar: ndarray, atomic_dipole: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, window: int, rc: float = 6.75)[source]#
ai2_kit.feat.spectrum.algorithm.cal_range_dipole_polar(z, atomic_dipole, z_lo, z_hi, r_smth)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_corr(A: ndarray, B: ndarray, NMAX: int, window: Optional[int] = None)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_corr_polar(atomic_polar: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, window: int, rc: float = 6.75)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_corr_polar_atomic(atomic_polar: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, window: int, rc: float = 6.75)#
ai2_kit.feat.spectrum.algorithm.calculate_corr_vdipole(atomic_dipole: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, dt_ps: float, window: int, rc: float = 6.75)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_corr_vdipole_atomic(atomic_dipole: ndarray, weight: ndarray, coords: ndarray, cells: ndarray, dt_ps: float, window: int, rc: float = 6.75)#
ai2_kit.feat.spectrum.algorithm.calculate_dipole(coords_O: ndarray, coords_H: ndarray, cells: ndarray, wannier: ndarray) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.calculate_dipole_H(coords_O: ndarray, coords_H: ndarray, cells: ndarray)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_dipole_OH(coords_O: ndarray, coords_H: ndarray, coords_Al: ndarray, cells: ndarray, wannier: ndarray) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.calculate_dipole_OH_H(coords_O: ndarray, coords_H: ndarray, coords_Al: ndarray, cells: ndarray)[source]#
ai2_kit.feat.spectrum.algorithm.calculate_ir(corr: ndarray, width: float, dt_ps: float, temperature: float, M: Optional[int] = None, filter_type: str = 'gaussian')[source]#
ai2_kit.feat.spectrum.algorithm.calculate_ir_atomic(corr: ndarray, width: float, dt_ps: float, temperature: float, M: Optional[int] = None, filter_type: str = 'gaussian')#
ai2_kit.feat.spectrum.algorithm.calculate_raman(corr: ndarray, width: float, dt_ps: float, temperature: float, M: Optional[int] = None, filter_type: str = 'gaussian')[source]#
ai2_kit.feat.spectrum.algorithm.calculate_raman_atomic(corr: ndarray, width: float, dt_ps: float, temperature: float, M: Optional[int] = None, filter_type: str = 'gaussian')#
ai2_kit.feat.spectrum.algorithm.calculate_sfg(corr: ndarray, width: int, dt_ps: float, temperature: float, M: Optional[int] = None, filter_type: str = 'gaussian')[source]#
ai2_kit.feat.spectrum.algorithm.change_unit_ir(freq_ps, CHAT: ndarray, temperature: float)[source]#
ai2_kit.feat.spectrum.algorithm.change_unit_raman(freq_ps, CHAT: ndarray, temperature: float)[source]#
ai2_kit.feat.spectrum.algorithm.change_unit_sfg(freq_ps, CHAT: ndarray, temperature: float)[source]#
ai2_kit.feat.spectrum.algorithm.compute_atomic_dipole_h2o(traj: System, wannier: ndarray, type_O: int = 0, type_H: int = 1, r_bond: float = 1.2, save_h2o: Optional[str] = None, save_atomic_dipole: Optional[str] = None)[source]#

Compute atomic dipole moments for water molecules from trajectory and Wannier center data.

Parameters:
  • traj (dpdata.System) – System trajectory containing coordinates, cells, and atom types.

  • wannier (np.ndarray) – Predicted Wannier center coordinates array.

  • type_O (int, optional) – Atom type index for Oxygen. Defaults to 0.

  • type_H (int, optional) – Atom type index for Hydrogen. Defaults to 1.

  • r_bond (float, optional) – Cutoff distance to define an O-H bond. Defaults to 1.2.

  • save_h2o (Optional[str], optional) – File path to save the extracted water oxygen coordinates. Defaults to None.

  • save_atomic_dipole (Optional[str], optional) – File path to save the computed atomic dipoles. Defaults to None.

Returns:

  • h2o_O: Coordinates of the Oxygen atoms of water molecules.

  • atomic_dipole: The computed atomic dipole array for each water molecule.

Return type:

tuple[np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.compute_bulk_ir_h2o(h2o: ndarray, cells: ndarray, atomic_dipole: ndarray, dt: float = 0.0005, window: int = 2000, z0: float = 25.0, zc: float = 7.5, zw: float = 0.5, rc: float = 6.0, width: int = 240, temperature: float = 330.0, M: int = 20000, filter_type: str = 'lorenz', nuclear_quantum_factor: float = 0.96, save_ir_plot: Optional[str] = None, save_ir_spectrum: Optional[str] = None)[source]#

Compute and optionally save/plot the bulk IR spectrum from atomic dipole data.

Parameters:
  • h2o (np.ndarray) – Water molecule coordinates.

  • cells (np.ndarray) – Simulation box array.

  • atomic_dipole (np.ndarray) – Atomic dipole array.

  • dt (float, optional) – Time step between frames in ps. Defaults to 0.0005.

  • window (int, optional) – Window size for correlation/frequency interval. Defaults to 2000.

  • z0 (float, optional) – Z cutoff center for the weighting function. Defaults to 25.0.

  • zc (float, optional) – Z cutoff width for the bulk region. Defaults to 7.5.

  • zw (float, optional) – Z cutoff smoothing parameter. Defaults to 0.5.

  • rc (float, optional) – Cutoff radius for correlation function. Defaults to 6.0.

  • width (int, optional) – Smoothing width parameter for the IR spectrum. Defaults to 240.

  • temperature (float, optional) – Simulation temperature in Kelvin. Defaults to 330.0.

  • M (int, optional) – Number of points for IR Fourier transform. Defaults to 20000.

  • filter_type (str, optional) – Filter type (“lorenz” or “gaussian”). Defaults to “lorenz”.

  • nuclear_quantum_factor (float, optional) – Correction factor for nuclear quantum effects. Defaults to 0.96.

  • save_ir_plot (Optional[str], optional) – Path to save the IR plot figure. Defaults to None.

  • save_ir_spectrum (Optional[str], optional) – Path to save the IR spectrum data. Defaults to None.

Returns:

  • wavenumber: Array of evaluated frequencies (cm^-1).

  • ir: The computed total IR intensity.

Return type:

tuple[np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.compute_bulk_raman_h2o(h2o: ndarray, cells: ndarray, atomic_polar: ndarray, dt: float = 0.0005, window: int = 2000, z0: float = 25.0, zc: float = 7.5, zw: float = 0.5, rc: float = 6.0, width: int = 240, temperature: float = 330.0, M: int = 20000, filter_type: str = 'lorenz', nuclear_quantum_factor: float = 0.96, save_iso_plot: Optional[str] = None, save_aniso_plot: Optional[str] = None, save_low_freq_plot: Optional[str] = None, save_raman_spectrum: Optional[str] = None)[source]#

Compute and optionally save/plot the bulk Raman spectra (isotropic, anisotropic, and low-frequency).

Parameters:
  • h2o (np.ndarray) – Water molecule coordinates.

  • cells (np.ndarray) – Simulation box array.

  • atomic_polar (np.ndarray) – Atomic polarizability tensor array.

  • dt (float, optional) – Time step between frames in ps. Defaults to 0.0005.

  • window (int, optional) – Window size for correlation calculation. Defaults to 2000.

  • z0 (float, optional) – Z cutoff center for the weighting function. Defaults to 25.0.

  • zc (float, optional) – Z cutoff width for the bulk region. Defaults to 7.5.

  • zw (float, optional) – Z cutoff smoothing parameter. Defaults to 0.5.

  • rc (float, optional) – Cutoff radius for correlation function. Defaults to 6.0.

  • width (int, optional) – Smoothing width parameter for the Raman spectrum. Defaults to 240.

  • temperature (float, optional) – Simulation temperature in Kelvin. Defaults to 330.0.

  • M (int, optional) – Number of points for Raman Fourier transform. Defaults to 20000.

  • filter_type (str, optional) – Filter type (“lorenz” or “gaussian”). Defaults to “lorenz”.

  • nuclear_quantum_factor (float, optional) – Correction factor for nuclear quantum effects. Defaults to 0.96.

  • save_iso_plot (Optional[str], optional) – Path to save the isotropic Raman plot. Defaults to None.

  • save_aniso_plot (Optional[str], optional) – Path to save the anisotropic Raman plot. Defaults to None.

  • save_low_freq_plot (Optional[str], optional) – Path to save the low-frequency Raman plot. Defaults to None.

  • save_raman_spectrum (Optional[str], optional) – Path to save the Raman spectrum data. Defaults to None.

Returns:

  • wavenumber: Array of evaluated frequencies (cm^-1).

  • total_iso: Isotropic Raman intensity.

  • total_aniso: Anisotropic Raman intensity.

  • low_range: Low-frequency Raman intensity.

Return type:

tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.compute_surface_ir_spectra_h2o(h2o: ndarray, cells: ndarray, atomic_dipole: ndarray, dt: float = 0.0005, window: int = 2000, z0: float = 25.0, zc: float = 7.5, zw: float = 0.5, rc: float = 6.0, width: int = 240, temperature: float = 330.0, M: int = 20000, filter_type: str = 'lorenz', nuclear_quantum_factor: float = 0.96, save_ir_plot: Optional[str] = None, save_ir_spectrum: Optional[str] = None)[source]#

Compute and optionally plot/save surface IR spectra for interfacial water.

Parameters:
  • h2o (np.ndarray) – Interfacial water molecule coordinates.

  • cells (np.ndarray) – Simulation box array.

  • atomic_dipole (np.ndarray) – Atomic dipole array.

  • dt (float, optional) – Time step between frames in ps. Defaults to 0.0005.

  • window (int, optional) – Window size for correlation function. Defaults to 2000.

  • z0 (float, optional) – Z cutoff center to isolate surface water. Defaults to 25.0.

  • zc (float, optional) – Z cutoff width to isolate surface water. Defaults to 7.5.

  • zw (float, optional) – Z cutoff smoothing parameter. Defaults to 0.5.

  • rc (float, optional) – Cutoff radius for correlation function. Defaults to 6.0.

  • width (int, optional) – Smoothing width parameter for IR spectrum. Defaults to 240.

  • temperature (float, optional) – Simulation temperature in Kelvin. Defaults to 330.0.

  • M (int, optional) – Number of points for Fourier transform. Defaults to 20000.

  • filter_type (str, optional) – Filter type (“lorenz” or “gaussian”). Defaults to “lorenz”.

  • nuclear_quantum_factor (float, optional) – Correction factor for nuclear quantum effects. Defaults to 0.96.

  • save_ir_plot (Optional[str], optional) – Path to save the IR plot figure. Defaults to None.

  • save_ir_spectrum (Optional[str], optional) – Path to save the IR spectrum data. Defaults to None.

Returns:

  • wavenumber: Array of evaluated frequencies (cm^-1).

  • ir_h2o: The computed IR intensity for surface water.

Return type:

tuple[np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.compute_surface_raman_h2o(h2o: ndarray, cells: ndarray, atomic_polar: ndarray, dt: float = 0.0005, window: int = 2000, z0: float = 25.0, zc: float = 7.5, zw: float = 0.5, rc: float = 6.0, width: int = 240, temperature: float = 330.0, M: int = 20000, filter_type: str = 'lorenz', nuclear_quantum_factor: float = 0.96, save_iso_plot: Optional[str] = None, save_aniso_plot: Optional[str] = None, save_low_freq_plot: Optional[str] = None, save_raman_spectrum: Optional[str] = None)[source]#

Compute and optionally save the surface Raman spectra for interfacial water.

Parameters:
  • h2o (np.ndarray) – Interfacial water molecule coordinates.

  • cells (np.ndarray) – Simulation box array.

  • atomic_polar (np.ndarray) – Atomic polarizability tensor array.

  • dt (float, optional) – Time step between frames in ps. Defaults to 0.0005.

  • window (int, optional) – Window size for correlation calculation. Defaults to 2000.

  • z0 (float, optional) – Z cutoff center to isolate surface water. Defaults to 25.0.

  • zc (float, optional) – Z cutoff width to isolate surface water. Defaults to 7.5.

  • zw (float, optional) – Z cutoff smoothing parameter. Defaults to 0.5.

  • rc (float, optional) – Cutoff radius for correlation function. Defaults to 6.0.

  • width (int, optional) – Smoothing width parameter for the Raman spectrum. Defaults to 240.

  • temperature (float, optional) – Simulation temperature in Kelvin. Defaults to 330.0.

  • M (int, optional) – Number of points for Fourier transform. Defaults to 20000.

  • filter_type (str, optional) – Filter type (“lorenz” or “gaussian”). Defaults to “lorenz”.

  • nuclear_quantum_factor (float, optional) – Correction factor for nuclear quantum effects. Defaults to 0.96.

  • save_iso_plot (Optional[str], optional) – Path to save the isotropic Raman plot. Defaults to None.

  • save_aniso_plot (Optional[str], optional) – Path to save the anisotropic Raman plot. Defaults to None.

  • save_low_freq_plot (Optional[str], optional) – Path to save the low-frequency Raman plot. Defaults to None.

  • save_raman_spectrum (Optional[str], optional) – Path to save the Raman spectrum data. Defaults to None.

Returns:

  • wavenumber: Array of evaluated frequencies (cm^-1).

  • total_iso: Isotropic Raman intensity.

  • total_aniso: Anisotropic Raman intensity.

  • low_range: Low-frequency Raman intensity.

Return type:

tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.compute_surface_sfg_h2o(h2o: ndarray, cells: ndarray, atomic_dipole: ndarray, atomic_polar: ndarray, dt: float = 0.0005, window: int = 2000, z0: float = 25.0, zc: float = 7.5, zw: float = 0.5, rc: float = 6.0, width: int = 240, temperature: float = 330.0, M: int = 20000, filter_type: str = 'lorenz', nuclear_quantum_factor: float = 0.96, save_sfg_plot: Optional[str] = None, save_sfg_spectrum: Optional[str] = None)[source]#

Compute and optionally plot/save the surface SFG spectrum (xxz/yyz).

Parameters:
  • h2o (np.ndarray) – Interfacial water molecule coordinates.

  • cells (np.ndarray) – Simulation box array.

  • atomic_dipole (np.ndarray) – Atomic dipole array.

  • atomic_polar (np.ndarray) – Atomic polarizability tensor array.

  • dt (float, optional) – Time step between frames in ps. Defaults to 0.0005.

  • window (int, optional) – Window size for correlation calculation. Defaults to 2000.

  • z0 (float, optional) – Z cutoff center to isolate surface water. Defaults to 25.0.

  • zc (float, optional) – Z cutoff width to isolate surface water. Defaults to 7.5.

  • zw (float, optional) – Z cutoff smoothing parameter. Defaults to 0.5.

  • rc (float, optional) – Cutoff radius for correlation function. Defaults to 6.0.

  • width (int, optional) – Smoothing width parameter for the SFG spectrum. Defaults to 240.

  • temperature (float, optional) – Simulation temperature in Kelvin. Defaults to 330.0.

  • M (int, optional) – Number of points for Fourier transform. Defaults to 20000.

  • filter_type (str, optional) – Filter type (“lorenz” or “gaussian”). Defaults to “lorenz”.

  • nuclear_quantum_factor (float, optional) – Correction factor for nuclear quantum effects. Defaults to 0.96.

  • save_sfg_plot (Optional[str], optional) – Path to save the SFG plot figure. Defaults to None.

  • save_sfg_spectrum (Optional[str], optional) – Path to save the SFG spectrum data. Defaults to None.

Returns:

  • wavenumber: Array of evaluated frequencies (cm^-1).

  • sfg_imag_h2o: The computed imaginary part of the SFG susceptibility (Im[chi^(2)]).

Return type:

tuple[np.ndarray, np.ndarray]

ai2_kit.feat.spectrum.algorithm.cutoff_ir_raman(arr, low, high, smooth_width)[source]#
ai2_kit.feat.spectrum.algorithm.cutoff_z(z_arr, z0, zc, zw)[source]#
ai2_kit.feat.spectrum.algorithm.do_pbc(coords: ndarray, cells: ndarray) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.extract_atomic_polar_from_traj_h2o(traj: System, polar: ndarray, type_O: int = 0, type_H: int = 1, r_bond: float = 1.2, save_atomic_polar: Optional[str] = None)[source]#

Extract atomic polarizability tensors specifically for water molecules.

Parameters:
  • traj (dpdata.System) – System trajectory containing coordinates, cells, and atom types.

  • polar (np.ndarray) – Raw polarizability tensor array for all components.

  • type_O (int, optional) – Atom type index for Oxygen. Defaults to 0.

  • type_H (int, optional) – Atom type index for Hydrogen. Defaults to 1.

  • r_bond (float, optional) – Cutoff distance to define an O-H bond. Defaults to 1.2.

  • save_atomic_polar (Optional[str], optional) – File path to save the water polarizability array. Defaults to None.

Returns:

The extracted atomic polarizability array for water molecules.

Return type:

np.ndarray

ai2_kit.feat.spectrum.algorithm.find_h2o(coords_sel: ndarray, coords_oth: ndarray, cells: ndarray, r_bond: float)[source]#
ai2_kit.feat.spectrum.algorithm.get_distance(coords_A: ndarray, coords_B: Optional[ndarray], cells: ndarray, remove_diag: bool = False, offset: int = 0)[source]#
ai2_kit.feat.spectrum.algorithm.get_lumped_wacent_poses_rel_h2o(stc: Atoms, elem_symbol: str, cutoff: float = 1.0, expected_cn: int = 4)[source]#

Determine the positions of the Wannier centers around a target element (e.g., O) and lump them into a mean relative displacement centroid.

Parameters:
  • stc (Atoms) – ASE Atoms object for a specific simulation frame.

  • elem_symbol (str) – Chemical symbol of the element to analyze.

  • cutoff (float, optional) – Neighbor cutoff radius in angstroms. Defaults to 1.0.

  • expected_cn (int, optional) – Expected coordination number of Wannier centers. Defaults to 4.

Returns:

Matrix of relative lumped Wannier centroid vectors for each targeted atom.

Return type:

np.ndarray

ai2_kit.feat.spectrum.algorithm.inv_cells(cells: ndarray)[source]#
ai2_kit.feat.spectrum.algorithm.k_nearest(coords_A: ndarray, coords_B: Optional[ndarray], cells: ndarray, k: int, batch_size: int = - 1)[source]#
ai2_kit.feat.spectrum.algorithm.set_cells_h2o(stc_list: List[Atoms], cell: List[float])[source]#

Write the length and pbc conditions to the bounding box.

Parameters:
  • stc_list (List[Atoms]) – A list of ASE Atoms objects representing simulation frames.

  • cell (List[float]) – Defining the dimensions/parameters of the simulation box.

Returns:

The modified list of ASE Atoms objects with updated cell/PBC info.

Return type:

List[Atoms]

ai2_kit.feat.spectrum.algorithm.set_lumped_wfc_h2o(stc_list: List[Atoms], lumped_dict: Dict[str, int])[source]#

Use get_lumped_wacent_poses_rel every frame to compute and return the absolute coordinates of Wannier centroids.

Parameters:
  • stc_list (List[Atoms]) – List of ASE Atoms objects for simulation frames.

  • lumped_dict (Dict[str, int]) – Dictionary where keys are elemental symbols and values are their expected coordination numbers.

Returns:

Evaluated coordinates of the lumped Wannier centroids across all frames.

Return type:

np.ndarray

ai2_kit.feat.spectrum.algorithm.to_frac(coords: ndarray, cells: ndarray) ndarray[source]#
ai2_kit.feat.spectrum.algorithm.write_to_diagonal(a: ndarray, diag: ndarray, offset: int = 0, axis1: int = 0, axis2: int = 1)[source]#