Source code for moldesign.min.base

# Copyright 2016 Autodesk Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
import sys

import numpy as np

import moldesign as mdt
from moldesign import data
from moldesign import units as u

[docs]class MinimizerBase(object): _strip_units = True # callbacks expect and return dimensionless quantities scaled to default unit system def __init__(self, mol, nsteps=20, force_tolerance=data.DEFAULT_FORCE_TOLERANCE, frame_interval=None, _restart_from=0, _restart_energy=None): self.mol = mol self.nsteps = nsteps - _restart_from self.force_tolerance = force_tolerance self.frame_interval = mdt.utils.if_not_none(frame_interval, max(nsteps/10, 1)) self._restart_from = _restart_from # Set up the trajectory to track the minimization self.traj = mdt.Trajectory(mol) self.current_step = _restart_from if _restart_energy is None: self.traj.new_frame(minimization_step=0, annotation='minimization steps:%d (energy=%s)' % (0, mol.calc_potential_energy())) self._initial_energy = _restart_energy self._last_energy = None self._last_grad = None # Figure out whether we'll use gradients self.request_list = ['potential_energy'] if 'forces' in mol.energy_model.ALL_PROPERTIES: self.gradtype = 'analytical' self.request_list.append('forces') else: self.gradtype = 'approximate' assert len(mol.constraints) == 0, \ 'Constrained minimization only available with analytical gradients' def _sync_positions(self, vector): """ Set the molecule's position """ c = vector.reshape((self.mol.num_atoms, 3)) if self._strip_units: self.mol.positions = c*u.default.length else: self.mol.positions = c def _coords_to_vector(self, coords): """ Convert position array to a flat vector """ vec = coords.reshape(self.mol.num_atoms * 3).copy() if self._strip_units: return vec.magnitude else: return vec
[docs] def objective(self, vector): """ Callback function to calculate the objective function """ self._sync_positions(vector) try: self.mol.calculate(requests=self.request_list) except mdt.QMConvergenceError: # returning infinity can help rescue some line searches return np.inf pot = self.mol.potential_energy if self._initial_energy is None: self._initial_energy = pot self._last_energy = pot if self._strip_units: return pot.defunits().magnitude else: return pot.defunits()
[docs] def grad(self, vector): """ Callback function to calculate the objective's gradient """ self._sync_positions(vector) self.mol.calculate(requests=self.request_list) grad = -self.mol.forces grad = grad.reshape(self.mol.num_atoms * 3) self._last_grad = grad if self._strip_units: return grad.defunits().magnitude else: return grad.defunits()
def __call__(self): """ Run the minimization Returns: moldesign.Trajectory: the minimization trajectory """ if self.traj.num_frames == 0 or self.traj[-1].minimization_step != self.current_step: self.traj.new_frame(minimization_step=self.current_step, annotation='minimization result (%d steps) (energy=%s)'% (self.current_step, self.mol.potential_energy)) return self.traj
[docs] def run(self): raise NotImplementedError('This is an abstract base class')
[docs] def callback(self, *args): """ To be called after each minimization step Args: *args: ignored """ self.current_step += 1 if self.current_step % self.frame_interval != 0: return self.mol.calculate(self.request_list) self.traj.new_frame(minimization_step=self.current_step, annotation='minimization steps:%d (energy=%s)'% (self.current_step, self.mol.potential_energy)) if self.nsteps is None: message = ['Minimization step %d' % self.current_step] else: message = ['Step %d/%d' % (self.current_step, self.nsteps + self._restart_from)] if self._last_energy is not None: message.append(u'\u0394E={x.magnitude:.3e} {x.units}'.format( x=self._last_energy - self._initial_energy)) if self.gradtype == 'analytical' and self._last_grad is not None: force = self._last_grad message.append(u'RMS \u2207E={rmsf.magnitude:.3e}, ' u'max \u2207E={mf.magnitude:.3e} {mf.units}'.format( rmsf=np.sqrt( / self.mol.ndims), mf=np.abs(force).max())) if self.mol.constraints: nsatisfied = 0 for c in self.mol.constraints: if c.satisfied(): nsatisfied += 1 message.append('constraints:%d/%d' % (nsatisfied, len(self.mol.constraints))) print(', '.join(message)) sys.stdout.flush()
@classmethod def _as_function(cls, newname): """ Create a function that runs this minimization """ @mdt.utils.args_from(cls, allexcept=['self']) def asfn(*args, **kwargs): obj = cls(*args, **kwargs) return obj() asfn.__name__ = newname asfn.__doc__ = cls.__doc__ return asfn