import matplotlib.pyplot as plt
import MDAnalysis as mda
from ase.neighborlist import NeighborList
import numpy as np
from multiprocessing import Pool
from multiprocessing import Manager
[docs]
class Scan:
""" This is a class for performing cans on a molecule using a calculator of your choice.
Parameters
----------
calculator : object
The calculator to use for the calculations. This should be an instance of a class that inherits from AbstractIO.
Attributes
----------
calculator : object
The calculator to use for the calculations. This should be an instance of a class that inherits from AbstractIO.
torsion_angles : list
A list of torsion angles for the scan.
energies : list
A list of energies for the scan.
forces : list
A list of forces for the scan.
coordinates : list
A list of coordinates for the scan.
dihedrals : list
A list of dihedrals for the scan.
unit : str
The unit of the dihedral angle. Default is None.
initial_structure : ase.Atoms
The initial structure of the molecule.
atom_types : list
A list of atom types in the molecule.
mask : list
A list of atom indices that should be rotated.
See Also
--------
ase.neighborlist.NeighborList : The ASE neighbor list object used to determine connectivity.
MDAnalysis : The MDAnalysis library used for molecular dynamics analysis.
multiprocessing : The multiprocessing library used for parallel processing.
"""
def __init__(self, interface):
self.calculator = interface
self.torsion_angles = []
self.energies = []
self.forces = []
self.coordinates = []
self.dihedrals = []
self.schedule="None"
self.unit = None
return
[docs]
def add_initial_coordinates(self, structure):
""" Add the initial coordinates of the molecule to the torsional scan.
Parameters
----------
structure : ase.Atoms
The ASE atoms object containing the initial coordinates of the molecule.
"""
self.initial_structure = structure
self.atom_types = structure.get_chemical_symbols()
return
[docs]
def scan(self, atoms=None, scan_range=(0,10,1), mask=None, verbose=False, num_workers=1, charge=0, schedule=None, constrain=None, optimize=True):
""" Perform a scan on the molecule.
Parameters
----------
atoms : list
A list of atoms indices to scan. Must be length 2, 3, or 4.
scan_range : tuple
A tuple containing the start, end, and step size of the scan. Default is (0,10,1).
mask : list
A list of atom indices that should be rotated. If None, the function will automatically select the mask.
verbose : bool
If True, print the dihedral angle being scanned.
num_workers : int
The number of workers to use for parallel processing. Default is 1.
charge : int
The net charge of the system. Default is 0.
schedule : str
The schedule for the scan. Can be "linear" or "iterative". Default is None, which means no schedule is applied.
Returns
-------
energies : list
A list of energies for each scan point.
forces : list
A list of forces for each scan point.
coordinates : list
A list of coordinates for each scan point.
Raises
------
ValueError
If the number of atoms is not 2, 3, or 4.
"""
if schedule is not None:
if schedule not in ["linear", "iterative", "bothways", "iterative-bothways"]:
raise ValueError("Invalid schedule. Must be 'linear', 'iterative', 'bothways', or 'iterative-bothways'.")
else:
self.schedule = schedule
if constrain is None:
self.constrain = atoms
else:
self.constrain = constrain
if num_workers < 1:
raise ValueError("num_workers must be greater than 0. Please set num_workers to 1 or greater.")
if len(atoms) == 2:
if num_workers > 1 and self.calculator.allow_parallel:
self._bond_scan_parallel()
else:
self._bond_scan_serial()
elif len(atoms) == 3:
if num_workers > 1 and self.calculator.allow_parallel:
self._angle_scan_parallel()
else:
self._angle_scan()
elif len(atoms) == 4:
if num_workers > 1 and self.calculator.allow_parallel:
self._dihedral_scan_parallel(atoms, scan_range, mask, verbose=verbose, num_workers=num_workers, optimize=optimize)
else:
self._dihedral_scan_serial(atoms, scan_range, mask, verbose=verbose, charge=charge, optimize=optimize)
else:
raise ValueError("Invalid number of atoms. Must be 2, 3, or 4.")
print("Scan complete.")
print(f"Total calculations performed: {len(self.energies)}")
print("Energies: ", self.energies)
return
[docs]
@staticmethod
def process_bond_scan(args):
raise NotImplementedError("Parallel dihedral scan is not implemented yet.")
[docs]
@staticmethod
def process_angle_scan(args):
raise NotImplementedError("Parallel dihedral scan is not implemented yet.")
def _bond_scan_serial(self):
pass
def _bond_scan_paralell(self):
pass
def _angle_scan_serial(self):
pass
def _angle_scan_parallel(self):
pass
[docs]
@staticmethod
def process_structure(args):
structure, calculator = args
calc = calculator._calculate(structure)
e, f = calculator._report_data(calc)
return e, f
def _set_dihedral_range(self, structure, scan_range, atoms, schedule=None):
""" Generate a range of dihedral angles based on the scan range and the current dihedral angle.
There are several options for the schedule of the scan:
- "linear": The dihedral angle is scanned linearly from the start to the end of the scan range.
- "iterative": The dihedral angle is scanned iteratively from the current dihedral angle to the end of the scan range.
- "bothways": The dihedral angle is scanned from the start to the end of the scan range and then from the end to the start of the scan range.
- "iterative-bothways": The dihedral angle is scanned iteratively from the current dihedral angle to the end of the scan range and then from the end to the start of the scan range.
If the schedule is None, the dihedral angle is scanned linearly from the start to the end of the scan range.
The dihedral angles are rounded to 2 decimal places.
Note - bothways and iterative-bothways are more effective for scans, as they can help find actual minima in the PES; howeever, they require effectively double the number of calculations, so they are slower.
Parameters
----------
structure : ase.Atoms
The ASE atoms object containing the system data.
scan_range : tuple
A tuple containing the start, end, and step size of the scan. Default is (0,10,1).
atoms : list
A list of atom indices that define the dihedral angle. Must be length 4.
schedule : str, optional
The schedule for the scan. Can be "linear", "iterative", "bothways", or "iterative-bothways". Default is None, which means no schedule is applied.
Returns
-------
dihedral_range : list
A list of dihedral angles to scan. The angles are in degrees and are rounded to 2 decimal places.
"""
if schedule == "linear" or schedule is None:
dihedral_range = range(scan_range[0], scan_range[1], scan_range[2])
elif schedule == "iterative":
current_dihedral = structure.get_dihedral(atoms[0], atoms[1], atoms[2], atoms[3])
dihedral_range = [current_dihedral + angle for angle in range(scan_range[0], scan_range[1], scan_range[2])]
elif schedule == "bothways":
dihedral_range1 = range(scan_range[0], scan_range[1], scan_range[2])
dihedral_range2 = range(scan_range[1], scan_range[0], -scan_range[2])
dihedral_range = list(dihedral_range1) + list(dihedral_range2)
elif schedule == "iterative-bothways":
current_dihedral = structure.get_dihedral(atoms[0], atoms[1], atoms[2], atoms[3])
dihedral_range1 = [current_dihedral + angle for angle in range(scan_range[0], scan_range[1], scan_range[2])]
dihedral_range2 = [current_dihedral - angle for angle in range(scan_range[0], scan_range[1], scan_range[2])]
tmp_dihedral_range = list(dihedral_range1) + list(dihedral_range2)
dihedral_range = [np.round(angle, 2) for angle in tmp_dihedral_range]
else:
raise ValueError("Invalid schedule. Must be 'linear', 'iterative', 'bothways', or 'iterative-bothways'.")
print("Dihedral range set to:", dihedral_range)
print("Current dihedral angle:", structure.get_dihedral(atoms[0], atoms[1], atoms[2], atoms[3]))
return dihedral_range
def _dihedral_iter(self, structure, charge, optimize=True):
"""Helper function to iterate over dihedral angles and calculate energies and forces.
Parameters
----------
structure : ase.Atoms
The ASE atoms object containing the system data.
charge : int
The net charge of the system.
optimize : bool
If True, perform a geometry optimization before the scan. Default is True.
Returns
-------
e : float
The energy of the system.
f : list
The forces acting on the atoms in the system.
"""
structure = self.calculator._calculate(structure, charge=charge)
if optimize: structure = self.calculator._optimize(structure, constrain_mask=self.constrain)
e, f = self.calculator._report_data(structure)
structure.calc.reset()
return e, f
def _sort_dihedrals(self):
"""Sort the dihedrals in ascending order."""
sorted_indices = sorted(range(len(self.dihedrals)), key=lambda k: self.dihedrals[k]%360.0)
self.dihedrals = [self.dihedrals[i]%360 for i in sorted_indices]
self.energies = [self.energies[i] for i in sorted_indices]
self.forces = [self.forces[i] for i in sorted_indices]
self.coordinates = [self.coordinates[i] for i in sorted_indices]
# Filter dihedrals to keep only the lowest energy for duplicates
unique_dihedrals = {}
for i, dihedral in enumerate(self.dihedrals):
normalized_dihedral = dihedral % 360.0
if normalized_dihedral not in unique_dihedrals or self.energies[i] < unique_dihedrals[normalized_dihedral][0]:
unique_dihedrals[normalized_dihedral] = (self.energies[i], self.forces[i], self.coordinates[i])
# Update the lists with filtered values
tmp_dih = list(unique_dihedrals.keys())
self.dihedrals = [ np.round(d, 2) for d in tmp_dih]
self.energies = [unique_dihedrals[d][0] for d in self.dihedrals]
self.forces = [unique_dihedrals[d][1] for d in self.dihedrals]
self.coordinates = [unique_dihedrals[d][2] for d in self.dihedrals]
def _dihedral_scan_serial(self, atoms, scan_range, mask, verbose=False, charge=0, optimize=True):
""" Perform a dihedral scan on the molecule.
This is a serial implementation of the dihedral scan. It rotates the dihedral angle
defined by the four atoms in the list `atoms` from `scan_range[0]` to `scan_range[1]`
in steps of `scan_range[2]`. The energies and forces are calculated for each step
and stored in the `energies`, `forces`, and `coordinates` lists. The dihedral angles
are stored in the `dihedrals` list. The unit of the dihedral angle is degrees.
Parameters
----------
atoms : list
A list of atom indices to scan. Must be length 4.
scan_range : tuple
A tuple containing the start, end, and step size of the scan. Default is (0,10,1).
mask : list
A list of atom indices that should be rotated. If None, the function will automatically select the mask.
verbose : bool
If True, print the dihedral angle being scanned.
charge : int
The net charge of the system. Default is 0.
optimize : bool
If True, perform a geometry optimization before the scan. Default is True.
Returns
-------
None
Raises
------
ValueError
If the number of atoms is not 4.
"""
self.unit = "Angle (degrees)"
if mask is None:
mask = self._select_mask(self.initial_structure, atoms[2], atoms[3])
structure = self.initial_structure.copy()
structure = self.calculator._calculate(structure, charge=charge)
if optimize: structure = self.calculator._optimize(structure)
for dihedral in self._set_dihedral_range(structure, scan_range, atoms, schedule=self.schedule):
if verbose:
print(f"Dihedral angle: {dihedral} degrees")
structure.set_dihedral(atoms[0], atoms[1], atoms[2], atoms[3], dihedral, indices=mask)
try:
e, f = self._dihedral_iter(structure, charge, optimize=optimize)
self.energies.append(e); self.forces.append(f)
self.dihedrals.append(dihedral)
self.coordinates.append(structure.get_positions())
except Exception as e:
print(f"Error calculating dihedral angle {dihedral} degrees: {e}")
continue
self._sort_dihedrals()
def _dihedral_scan_parallel(self, atoms, scan_range, mask=None, verbose=False, num_workers=1, optimize=True):
""" Perform a dihedral scan on the molecule in parallel."""
raise NotImplementedError("Parallel dihedral scans are temporarily disabled due to changes in the approach to dihedral scans. Please use the serial implementation for now.")
manager = Manager()
energies_mp = manager.list()
forces_mp = manager.list()
self.unit = "Angle (degrees)"
if mask is None:
mask = self._select_mask(self.initial_structure, atoms[1], atoms[2])
tasks = []
for dihedral in range(scan_range[0], scan_range[1], scan_range[2]):
self.dihedrals.append(dihedral)
structure = self.initial_structure.copy()
structure.set_dihedral(atoms[0], atoms[1], atoms[2], atoms[3], dihedral, indices=mask)
tasks.append((structure, self.calculator))
self.coordinates.append(structure.get_positions())
# Use multiprocessing to perform the dihedral scan in parallel
with Pool(processes=num_workers) as pool:
results = pool.map(Scan.process_structure, tasks)
print("Mapping complete.")
for result in results:
e, f= result
energies_mp.append(e)
forces_mp.append(f)
self.energies = list(energies_mp)
self.forces = list(forces_mp)
[docs]
def write_xyz(self, filename="scan.xyz"):
""" Write the scan results to an XYZ file.
Parameters
----------
filename : str
The filename to save the XYZ file. Default is "scan.xyz".
"""
with open(filename, "w") as f:
for i, coord in enumerate(self.coordinates):
f.write(f"{len(coord)}\n")
f.write("Scan\n")
for j, atom in enumerate(coord):
f.write(f"{self.atom_types[j]} {atom[0]} {atom[1]} {atom[2]}\n")
return
def _select_mask(self, structure, atom1, atom2):
"""Automatically select the mask for atoms on either side of the bond.
Parameters
----------
structure : ase.Atoms
The ASE atoms object.
atom1 : int
The index of the first atom in the bond.
atom2 : int
The index of the second atom in the bond.
Returns
-------
mask : list
A list of atom indices that should be rotated.
"""
# Create a neighbor list to determine connectivity
nl = NeighborList([0.8] * len(structure), self_interaction=False, bothways=True)
nl.update(structure)
# Perform a breadth-first search (BFS) to find atoms connected to atom2
visited = set()
queue = [atom2]
mask = []
while queue:
current = queue.pop(0)
if current in visited or current == atom1:
continue
visited.add(current)
mask.append(current)
neighbors = nl.get_neighbors(current)[0]
queue.extend(neighbors)
print(f"Mask for dihedral scan: {mask}")
return mask
[docs]
def plot_scan(self, save=None):
"""Plot the scan results.
Parameters
----------
save : str
The filename to save the plot. If None, the plot will be shown but not saved.
"""
fig = plt.figure(dpi=300, figsize=(5,5))
plt.plot(self.dihedrals, self.energies)
plt.xlabel(f"Scan {self.unit}")
plt.ylabel("Energy (eV)")
plt.title("Torsional Scan Results")
plt.tight_layout()
if save:
plt.savefig(save)
else:
plt.show()