# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.
from __future__ import print_function
from pyiron_contrib.protocol.generic import PrimitiveVertex, CompoundVertex, Protocol
from pyiron_contrib.protocol.utils import ensure_iterable
from pyiron_contrib.protocol.primitive.one_state import ExternalHamiltonian, Counter, Norm, Max, GradientDescent
from pyiron_contrib.protocol.primitive.two_state import IsGEq, IsLEq
from pyiron_contrib.protocol.utils import Pointer, IODictionary
import numpy as np
"""
Protocol for running the force-based quantum mechanics/molecular mechanics concurrent coupling scheme described in
Huber et al., Comp. Mat. Sci. 118 (2016) 259-268
"""
__author__ = "Dominik Gehringer, Liam Huber"
__copyright__ = "Copyright 2019, Max-Planck-Institut für Eisenforschung GmbH " \
"- Computational Materials Design (CM) Department"
__version__ = "0.0"
__maintainer__ = "Liam Huber"
__email__ = "huber@mpie.de"
__status__ = "development"
__date__ = "June 6, 2019"
[docs]class QMMM(CompoundVertex):
"""
Relax a QM/MM coupled system.
Needs an MM job (incl. superstructure), a QM job reference, QM target site ids, optional target site chemistry,
optional domain ids, optional QM cell instructions?
During setup, the domain indices are checked -- if None, build them from scratch using shells and the cell...stuff?
Input attributes:
structure (pyiron.atomistics.structure.atoms.Atoms): The full region I+II superstructure.
mm_ref_job_full_path (str): Path to the pyiron job to use for evaluating forces and energies of the MM domains.
qm_ref_job_full_path (str): Path to the pyiron job to use for evaluating forces and energies of the QM domain.
domain_ids (dict): A dictionary of ids, `seed`, `core`, `buffer`, and `filler`, for mapping atoms of the MM
I+II superstructure to the various domains of the QM/MM coupling scheme. (Default is None, gets constructed
using consecutive shells).
seed_ids (int/list): The integer id (or ids) of atom(s) in the MM I+II superstructure to base the QM region on.
These are the only atoms whose species can be changed. (Default is None, which raises an error unless
`domain_ids` was explicitly provided -- else this parameter must be provided.)
shell_cutoff (float): Maximum distance for two atoms to be considered neighbours in the construction of shells
for automatic system partitioning. (Default is None, which will raise an error -- this parameter must be
provided if domain partitioning is done using `seed_ids` and shells. When `domain_ids` are explicitly
provided, `shell_cutoff` is not needed.)
n_core_shells (int): How many neighbour shells around the seed(s) to relax using QM forces. (Default is 2.)
n_buffer_shells (int): How many neighbour shells around the region I core to relax using MM forces. (Default is
2.)
seed_species (str/list): The species for each 'seed' atom in the QM domain. Value(s) should be a atomic symbol,
e.g. 'Mg'. (Default is None, which leaves all seeds the same as they occur in the input structure.)
vacuum_width (float): Minimum vacuum distance between atoms in two periodic images of the QM domain. Influences
the final simulation cell for the QM calculation. (Default is 2.)
filler_width (float): Length added to the bounding box of the region I atoms to create a new box from which to
draw filler atoms using the initial MM I+II superstructure. Influences the final simulation cell for the QM
calculation.If the value is not positive, no filler atoms are used. The second, larger box uses the same
center as the bounding box of the QM region I, so this length is split equally to the positive and negative
sides for each cartesian direction. (Default is 6.)
n_steps (int): The maximum number of minimization steps to make. (Default is 100.)
f_tol (float): The maximum force on any atom below which the calculation terminates. Only atoms which could be
relaxed are considered, i.e. QM forces for region I core, and MM forces for region II and I buffer. (Filler
atoms are not real, so we never care about their forces.) (Default is 1e-4 eV/angstrom.)
Output attributes:
energy_mm (float): The total energy in eV of the region I+II superstructure (without seed species) using the MM
representation.
energy_qm (float): The total energy in eV of the region I+filler structure (with seed species) using QM
representation.
energy_mm_small (float): The total energy in eV of the region I+filler structure (without seed species) using
the representation.
energy_qmmm (float): The composite QM/MM energy: `energy_mm` + `energy_qm` - `energy_mm_small` in eV.
max_force (float): The largest atomic force among QM forces on region I core atoms or MM forces on region I
buffer and region II atoms.
positions (numpy.ndarray): The per-atom vector of cartesian positions for the entire region I+II superstructure.
"""
DefaultWhitelist = {
'calc_static_mm': {
'output': {
'energy_pot': 1,
},
},
'calc_static_qm': {
'output': {
'energy_pot': 1,
},
},
}
def __init__(self, **kwargs):
super(QMMM, self).__init__(**kwargs)
id_ = self.input.default
id_.domain_ids = None
id_.seed_ids = None
id_.shell_cutoff = None
id_.n_core_shells = 2
id_.n_buffer_shells = 2
id_.vacuum_width = 2.
id_.filler_width = 6.
id_.n_steps = 100
id_.f_tol = 1e-4
id_.gamma0 = 0.1
id_.fix_com = True
id_.use_adagrad = False
[docs] def define_vertices(self):
# Components
g = self.graph
g.partition = PartitionStructure()
g.calc_static_mm = ExternalHamiltonian()
g.calc_static_qm = ExternalHamiltonian()
g.clock = Counter()
g.force_norm_mm = Norm()
g.force_norm_qm = Norm()
g.max_force_mm = Max()
g.max_force_qm = Max()
g.check_force_mm = IsLEq()
g.check_force_qm = IsLEq()
g.check_steps = IsGEq()
g.update_buffer_qm = AddDisplacements()
g.update_core_mm = AddDisplacements()
g.gradient_descent_mm = GradientDescent()
g.gradient_descent_qm = GradientDescent()
g.calc_static_small = ExternalHamiltonian()
[docs] def define_execution_flow(self):
g = self.graph
g.make_pipeline(
g.partition,
g.check_steps, 'false',
g.calc_static_mm,
g.calc_static_qm,
g.force_norm_mm,
g.max_force_mm,
g.check_force_mm, 'true',
g.force_norm_qm,
g.max_force_qm,
g.check_force_qm, 'false',
g.gradient_descent_mm,
g.gradient_descent_qm,
g.update_buffer_qm,
g.update_core_mm,
g.clock,
g.check_steps
)
g.make_edge(g.check_force_mm, g.gradient_descent_mm, 'false')
g.make_edge(g.check_force_qm, g.calc_static_small, 'true')
g.make_edge(g.check_steps, g.calc_static_small, 'true')
g.starting_vertex = g.partition
g.restarting_vertex = g.check_steps
[docs] def get_output(self):
gp = Pointer(self.graph)
e_mm = ~gp.calc_static_mm.output.energy_pot[-1]
e_qm = ~gp.calc_static_qm.output.energy_pot[-1]
e_mm_small = ~gp.calc_static_small.output.energy_pot[-1]
try: # If we terminate by step count, the QM force checker may never have gotten called
max_force = max(~gp.max_force_mm.output.amax[-1], ~gp.max_force_qm.output.amax[-1])
except KeyError:
max_force = ~gp.max_force_mm.output.amax[-1]
try: # We might also be converged before ever running a step (e.g. for perfect bulk)
positions = ~gp.update_core_mm.output.positions[-1]
except KeyError:
positions = self.input.structure.positions
return {
'energy_mm': e_mm,
'energy_qm': e_qm,
'energy_mm_small': e_mm_small,
'energy_qmmm': e_mm + e_qm - e_mm_small,
'max_force': max_force,
'positions': positions
}
[docs] def show_mm(self):
try:
mm_full_structure = self.graph.partition.output.mm_full_structure[-1]
mm_full_structure.positions = self.graph.update_core_mm.output.positions[-1]
domain_ids = self.graph.partition.output.domain_ids[-1]
except KeyError:
partitioned = self.partition_input()
mm_full_structure = partitioned['mm_full_structure']
domain_ids = partitioned['domain_ids']
color = 4 * np.ones(len(mm_full_structure)) # This 5th colour makes the balance of atoms
for n, group in enumerate(['seed', 'core', 'buffer', 'filler']):
color[domain_ids[group]] = n
return mm_full_structure.plot3d(scalar_field=color)
[docs] def show_qm(self):
try:
qm_structure = self.graph.partition.output.qm_structure[-1]
qm_structure.positions = self.graph.update_buffer_qm.output.positions[-1]
domain_ids_qm = self.graph.partition.output.domain_ids_qm[-1]
except KeyError:
partitioned = self.partition_input()
qm_structure = partitioned['qm_structure']
domain_ids_qm = partitioned['domain_ids_qm']
color = 4 * np.ones(len(qm_structure)) # If you see this 5th colour, something is wrong
for n, group in enumerate(['seed', 'core', 'buffer', 'filler']):
color[domain_ids_qm[group]] = n
return qm_structure.plot3d(scalar_field=color)
[docs] def show_boxes(self):
try:
structure = self.graph.partition.output.structure[-1]
qm_structure = self.graph.partition.output.qm_structure[-1]
except KeyError:
partitioned = self.partition_input()
structure = partitioned['mm_full_structure']
qm_structure = partitioned['qm_structure']
self._plot_boxes([structure.cell, qm_structure.cell],
colors=['r', 'b'],
titles=['MM Superstructure', 'QM Structure'])
@staticmethod
def _plot_boxes(cells, translate=None, colors=None, titles=None, default_color='b', size=(29, 21)):
"""
Plots one or a list of cells in xy, yz and xt projection
Args:
cells (numpy.ndarray/list): The cells to plot. A list of (3,3) or a single (3,3) matrix
translate (list): list of translations vectors for each cell list of (3,) numpy arrays.
colors (list): list of colors for each cell in the list. e.g ['b', 'g', 'y'] or 'bgky'
titles (list): list of names displayed for each cell. list of str
default_color (str): if colors is not specified this color will be applied to all cells
size (float,float): A tuple of two float specifiyng the size of the plot in centimeters
Returns:
matplotlib.figure: The figure which will be displayed
"""
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
# Make sure that looping over this properties works smoothly
cells = ensure_iterable(cells)
translate = translate or [np.zeros(3,)]*len(cells)
colors = colors or [default_color]*len(cells)
titles = ensure_iterable(titles)
# Scaled coordinates for all edges of the cell
edges = [
[(0, 0, 0), (1, 0, 0)],
[(0, 0, 0), (0, 1, 0)],
[(0, 0, 0), (0, 0, 1)],
[(1, 1, 0), (1, 0, 0)],
[(1, 1, 0), (0, 1, 0)],
[(1, 1, 0), (1, 1, 1)],
[(1, 0, 1), (1, 0, 0)],
[(1, 0, 1), (1, 1, 1)],
[(1, 0, 1), (0, 0, 1)],
[(0, 1, 1), (1, 1, 1)],
[(0, 1, 1), (0, 0, 1)],
[(0, 1, 1), (0, 1, 0)]
]
# Map axis to indices
axis_mapping = {a: i for a, i in zip('xyz', range(3))}
# Planes to plot
planes = ['xy', 'yz', 'xz']
# Get subplots and set the size in inches
fig, axes = plt.subplots(1, len(planes))
fig.set_size_inches(*[s / 2.54 for s in size])
for i, plane in enumerate(planes):
axis = axes[i]
axis.set_aspect('equal')
axis.title.set_text('{} plane'.format(plane))
indices = [axis_mapping[a] for a in plane]
for j, (cell, color, translation) in enumerate(zip(cells, colors, translate)):
calc = lambda mul_, cell_: sum([m_ * vec_ for m_, vec_ in zip(mul_, cell_)])
coords = [(calc(start, cell) + translation, calc(end, cell) + translation) for start, end in edges]
for start, end in coords:
sx, sy = start[indices]
ex, ey = end[indices]
axis.plot([sx, ex], [sy, ey], color=color)
# Create dummy legend outside
legend_lines = [Line2D([0], [0], color=col or default_color, lw=2) for col in colors]
legend_titles = [title or 'Box {}'.format(i + 1) for i, title in enumerate(titles)]
plt.figlegend(legend_lines, legend_titles, loc='lower center', fancybox=True, shadow=True)
return plt
[docs]class AddDisplacements(PrimitiveVertex):
def __init__(self, name=None):
super(AddDisplacements, self).__init__(name=name)
self.input.default.target_mask = None
self.input.default.displacement_mask = None
[docs] def command(self, target, displacement, target_mask, displacement_mask):
result = target.copy()
if target_mask is not None and isinstance(target_mask, list):
target_mask = np.concatenate(target_mask)
if displacement_mask is not None and isinstance(displacement_mask, list):
displacement_mask = np.concatenate(displacement_mask)
if target_mask is not None and displacement_mask is not None:
result[target_mask] += displacement[displacement_mask]
elif target_mask is not None and displacement_mask is None:
result[target_mask] += displacement
elif target_mask is None and displacement_mask is not None:
result += displacement[displacement_mask]
else:
result += displacement
return {
'positions': result
}
[docs]class PartitionStructure(PrimitiveVertex):
"""
"""
def __init__(self, name=None):
super(PartitionStructure, self).__init__(name=name)
id_ = self.input.default
id_.domain_ids = None
id_.seed_ids = None
id_.shell_cutoff = None
id_.n_core_shells = None
id_.n_buffer_shells = None
id_.filler_width = None
[docs] def command(
self,
structure,
domain_ids,
seed_ids, shell_cutoff, n_core_shells, n_buffer_shells,
vacuum_width, filler_width,
seed_species
):
domain_ids, domain_ids_qm, mm_small_structure = self._set_qm_structure(
structure,
domain_ids,
seed_ids, shell_cutoff, n_core_shells, n_buffer_shells,
vacuum_width, filler_width
)
qm_structure = self._change_qm_species(mm_small_structure, domain_ids_qm, seed_species)
domain_ids_qm['only_core'] = self._only_core(domain_ids_qm)
domain_ids['except_core'] = self._except_core(structure, domain_ids)
return {
'mm_full_structure': structure.copy(),
'mm_small_structure': mm_small_structure.copy(),
'qm_structure': qm_structure.copy(),
'domain_ids': domain_ids,
'domain_ids_qm': domain_ids_qm
}
def _set_qm_structure(
self,
superstructure,
domain_ids,
seed_ids, shell_cutoff, n_core_shells, n_buffer_shells,
vacuum_width, filler_width
):
if domain_ids is not None and seed_ids is not None:
raise ValueError('Only *one* of `seed_ids` and `domain_ids` may be provided.')
elif domain_ids is not None:
seed_ids = domain_ids['seed']
core_ids = domain_ids['core']
buffer_ids = domain_ids['buffer']
filler_ids = domain_ids['filler']
region_I_ids = np.concatenate((seed_ids, core_ids, buffer_ids))
bb = self._get_bounding_box(superstructure[np.concatenate((region_I_ids, filler_ids))])
elif seed_ids is not None:
shells = self._build_shells(superstructure, n_core_shells + n_buffer_shells, shell_cutoff, seed_ids)
core_ids = np.concatenate(shells[:n_core_shells])
buffer_ids = np.concatenate(shells[n_core_shells:])
region_I_ids = np.concatenate((seed_ids, core_ids, buffer_ids))
bb = self._get_bounding_box(superstructure[region_I_ids])
extra_box = 0.5 * np.array(filler_width)
bb[:, 0] -= extra_box
bb[:, 1] += extra_box
# Store it because get bounding box return a tight box and is different
filler_ids = self._get_ids_within_box(superstructure, bb)
filler_ids = np.setdiff1d(filler_ids, region_I_ids)
bb = self._get_bounding_box(superstructure[np.concatenate((region_I_ids, filler_ids))])
domain_ids = {'seed': seed_ids, 'core': core_ids, 'buffer': buffer_ids, 'filler': filler_ids}
else:
raise ValueError('At least *one* of `seed_ids` and `domain_ids` must be provided.')
# Build the domain ids in the qm structure
qm_structure = None
domain_ids_qm = {}
offset = 0
for key, ids in domain_ids.items():
if qm_structure is None:
qm_structure = superstructure[ids]
else:
qm_structure += superstructure[ids]
id_length = len(ids)
domain_ids_qm[key] = np.arange(id_length) + offset
offset += id_length
# And put everything in a box near (0,0,0)
extra_vacuum = 0.5 * vacuum_width
bb[:, 0] -= extra_vacuum
bb[:, 1] += extra_vacuum
# If the bounding box is larger than the MM superstructure
bs = np.abs(bb[:, 1] - bb[:, 0])
supercell_lengths = [np.linalg.norm(row) for row in superstructure.cell]
shrinkage = np.array([(box_size - cell_size) / 2.0 if box_size > cell_size else 0.0
for box_size, cell_size in zip(bs, supercell_lengths)])
if np.any(shrinkage > 0):
self.logger.info('The QM box is larger than the MM Box therefore I\'ll shrink it')
bb[:, 0] += shrinkage
bb[:, 1] -= shrinkage
elif any([0.9 < box_size / cell_size < 1.0 for box_size, cell_size in zip(bs, supercell_lengths)]):
# Check if the box is just slightly smaller than the superstructure cell
self.logger.warn(
'Your cell is nearly as large as your supercell. Probably you want to expand it a little bit')
qm_structure.cell = np.identity(3) * np.ptp(bb, axis=1)
box_center = tuple(np.dot(np.linalg.inv(superstructure.cell), np.mean(bb, axis=1)))
qm_structure.wrap(center=box_center)
# Wrap it to the unit cell
qm_structure.positions = np.dot(qm_structure.get_scaled_positions(), qm_structure.cell)
return domain_ids, domain_ids_qm, qm_structure
@staticmethod
def _change_qm_species(qm_structure, domain_ids_qm, seed_species):
qm_structure_alchemized = qm_structure.copy()
for index, species in zip(domain_ids_qm['seed'], seed_species):
qm_structure_alchemized[index] = species
return qm_structure_alchemized
@staticmethod
def _get_ids_within_box(structure, box):
"""
Finds all the atoms in a structure who have a periodic image inside the bounding box.
Args:
structure (Atoms): The structure to search.
box (np.ndarray): A 3x2 array of the x-min through z-max values for the bounding rectangular prism.
Returns:
np.ndarray: The integer ids of the atoms inside the box.
"""
box_center = np.mean(box, axis=1)
box_center_direct = np.dot(np.linalg.inv(structure.cell), box_center)
# Wrap atoms so that they are the closest image to the box center
wrapped_structure = structure.copy()
wrapped_structure.wrap(center=tuple(box_center_direct))
pos = wrapped_structure.positions
# Keep only atoms inside the box limits
masks = []
for d in np.arange(len(box)):
masks.append(pos[:, d] > box[d, 0])
masks.append(pos[:, d] < box[d, 1])
total_mask = np.prod(masks, axis=0).astype(bool)
return np.arange(len(structure), dtype=int)[total_mask]
@staticmethod
def _build_shells(structure, n_shells, shell_cutoff, seed_ids):
indices = [seed_ids]
current_shell_ids = seed_ids
for _ in range(n_shells):
neighbors = structure.get_neighbors(id_list=current_shell_ids, cutoff=shell_cutoff)
new_ids = np.setdiff1d(
np.unique(np.concatenate(neighbors.indices)),
np.concatenate(indices)
)
indices.append(new_ids)
current_shell_ids = new_ids
# Pop seed ids
indices.pop(0)
return indices
@staticmethod
def _get_bounding_box(structure):
"""
Finds the smallest rectangular prism which encloses all atoms in the structure after accounting for periodic
boundary conditions.
So what's the problem?
| ooooo |, easy, wrap by CoM (centre of mass)
|ooo oo|, easy, wrap by CoM
|o ooo o|, uh-oh! Need this function.
Args:
structure (Atoms): The structure to bound.
Returns:
numpy.ndarray: A 3x2 array of the x-min through z-max values for the bounding rectangular prism.
"""
wrapped_structure = structure.copy()
# Take the frist positions and wrap the atoms around there to determine the size of the bounding box
wrap_center = tuple(np.dot(np.linalg.inv(structure.cell), structure.positions[0, :]))
wrapped_structure.wrap(center=wrap_center)
bounding_box = np.vstack([
np.amin(wrapped_structure.positions, axis=0),
np.amax(wrapped_structure.positions, axis=0)
]).T
return bounding_box
@staticmethod
def _except_core(structure, domain_ids):
seed_ids = np.array(domain_ids['seed'])
core_ids = np.array(domain_ids['core'])
return np.setdiff1d(np.arange(len(structure)), np.concatenate([seed_ids, core_ids]))
@staticmethod
def _only_core(domain_ids_qm):
return np.concatenate([domain_ids_qm['seed'], domain_ids_qm['core']])
[docs]class ProtocolQMMM(Protocol, QMMM):
pass