# fmt: off import numpy as np import pytest from numpy.testing import assert_array_almost_equal from scipy.integrate import trapezoid from ase import Atoms from ase.calculators.calculator import CalculationFailed, CalculatorSetupError from ase.calculators.emt import EMT from ase.calculators.harmonic import HarmonicCalculator, HarmonicForceField from ase.calculators.mixing import MixedCalculator from ase.geometry.geometry import get_distances_derivatives from ase.md.andersen import Andersen from ase.md.velocitydistribution import ( MaxwellBoltzmannDistribution, Stationary, ZeroRotation, ) from ase.optimize import BFGS from ase.units import fs from ase.vibrations import Vibrations ref_pos = np.asarray([[8.7161, 7.96276, 8.48206], [8.60594, 8.04985, 9.44464], [8.0154, 8.52264, 8.10545]]) ref_atoms = Atoms('OH2', positions=ref_pos) # relaxed water molecule ref_energy = -14.222189 # y shift of the 'parabola' (harmonic potential) # example Hessian matrix as obtained from DFT hessian_x = np.asarray([[2.82630333e+01, -2.24763667e+01, 7.22478333e+00, -2.96970000e+00, 2.34363333e+00, 2.72788333e+00, -2.52159833e+01, 2.01307833e+01, -9.94651667e+00], [-2.24763667e+01, 1.78621333e+01, -5.77378333e+00, 2.33703333e+00, -1.85276667e+00, -2.15118333e+00, 2.01258667e+01, -1.60350833e+01, 7.93248333e+00], [7.22478333e+00, -5.77378333e+00, 5.72735000e+01, 7.25470000e+00, -5.75313333e+00, -4.69477333e+01, -1.44613000e+01, 1.15504833e+01, -1.03523333e+01], [-2.96970000e+00, 2.33703333e+00, 7.25470000e+00, 2.96170000e+00, -2.36901667e+00, -3.76841667e+00, -2.83833333e-02, 3.06833333e-02, -3.49190000e+00], [2.34363333e+00, -1.85276667e+00, -5.75313333e+00, -2.36901667e+00, 1.89046667e+00, 2.95495000e+00, 2.90666667e-02, -1.80666667e-02, 2.79565000e+00], [2.72788333e+00, -2.15118333e+00, -4.69477333e+01, -3.76841667e+00, 2.95495000e+00, 4.89340000e+01, 1.03146667e+00, -8.18450000e-01, -1.96118333e+00], [-2.52159833e+01, 2.01258667e+01, -1.44613000e+01, -2.83833333e-02, 2.90666667e-02, 1.03146667e+00, 2.52034000e+01, -2.01516833e+01, 1.34293167e+01], [2.01307833e+01, -1.60350833e+01, 1.15504833e+01, 3.06833333e-02, -1.80666667e-02, -8.18450000e-01, -2.01516833e+01, 1.60592333e+01, -1.07369667e+01], [-9.94651667e+00, 7.93248333e+00, -1.03523333e+01, -3.49190000e+00, 2.79565000e+00, -1.96118333e+00, 1.34293167e+01, -1.07369667e+01, 1.23150000e+01]]) def assert_water_is_relaxed(atoms): forces = atoms.get_forces() assert np.allclose(np.zeros(forces.shape), forces) assert np.allclose(ref_energy, atoms.get_potential_energy()) assert np.allclose(atoms.get_angle(1, 0, 2), ref_atoms.get_angle(1, 0, 2)) assert np.allclose(atoms.get_distance(0, 1), ref_atoms.get_distance(0, 1)) assert np.allclose(atoms.get_distance(0, 2), ref_atoms.get_distance(0, 2)) def run_optimize(atoms): opt = BFGS(atoms) opt.run(fmax=1e-9) @pytest.mark.optimize() def test_cartesians(): """In Cartesian coordinates the first 6 trash eigenvalues (translations and rotations) can be slightly different from zero; hence set them to zero using an increased parameter zero_thresh. """ zero_thresh = 0.06 # set eigvals to zero if abs(eigenvalue) < zero_thresh hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy, hessian_x=hessian_x, zero_thresh=zero_thresh) assert np.allclose(hff.hessian_q, hff.hessian_x) atoms = ref_atoms.copy() atoms.calc = HarmonicCalculator(hff) assert_water_is_relaxed(atoms) # atoms has not been distorted run_optimize(atoms) # nothing should happen assert_water_is_relaxed(atoms) # atoms should still be relaxed atoms.set_distance(0, 1, 3.5) # now distort atoms along axis, no rotation run_optimize(atoms) # optimization should recover original assert_water_is_relaxed(atoms) # relaxed geometry with pytest.raises(AssertionError): atoms.rattle() # relaxation should fail to recover the atoms.rotate(90, 'x') # original geometry of the atoms, run_optimize(atoms) # because Cartesian coordinates are assert_water_is_relaxed(atoms) # not rotationally invariant. @pytest.mark.optimize() def test_constraints_with_cartesians(): """Project out forces along x-component of H-atom (index 0 in the q-vector with the Cartesian coordinates (here: x=q)). A change in the x-component of the H-atom should not result in restoring forces, when they were projected out from the Hessian matrix. """ def test_forces(calc): atoms = ref_atoms.copy() atoms.calc = calc pos = ref_pos.copy() pos[0, 0] *= 2 atoms.set_positions(pos) run_optimize(atoms) # (no) restoring force along distorted x-component xdiff = atoms.get_positions() - ref_pos return all(xdiff[xdiff != 0] == pos[0, 0] / 2) zero_thresh = 0.06 # set eigvals to zero if abs(eigenvalue) < zero_thresh parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy, 'hessian_x': hessian_x, 'zero_thresh': zero_thresh} hff = HarmonicForceField(**parameters) calc = HarmonicCalculator(hff) assert not test_forces(calc) # restoring force along distorted x-component parameters['constrained_q'] = [0] # project out the coordinate with index 0 hff = HarmonicForceField(**parameters) calc = HarmonicCalculator(hff) assert test_forces(calc) # no restoring force along distorted x-component def setup_water(calc): atoms = ref_atoms.copy() atoms.calc = calc assert_water_is_relaxed(atoms) atoms.rattle(0.3) atoms.rotate(160, 'x') assert not np.allclose(ref_energy, atoms.get_potential_energy()) return atoms # start doc example 3 dist_defs = [[0, 1], [1, 2], [2, 0]] # define three distances by atom indices def water_get_q_from_x(atoms): """Simple internal coordinates to describe water with three distances.""" q_vec = [atoms.get_distance(i, j) for i, j in dist_defs] return np.asarray(q_vec) def water_get_jacobian(atoms): """Function to return the Jacobian for the water molecule described by three distances.""" pos = atoms.get_positions() dist_vecs = [pos[j] - pos[i] for i, j in dist_defs] derivs = get_distances_derivatives(dist_vecs) jac = [] for i, defin in enumerate(dist_defs): dqi_dxj = np.zeros(ref_pos.shape) for j, deriv in enumerate(derivs[i]): dqi_dxj[defin[j]] = deriv jac.append(dqi_dxj.flatten()) return np.asarray(jac) # end doc example 3 def test_raise_Errors(): with pytest.raises(CalculatorSetupError): HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x, get_q_from_x=lambda x: x) with pytest.raises(CalculatorSetupError): HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x, variable_orientation=True) with pytest.raises(CalculatorSetupError): HarmonicForceField(ref_atoms=ref_atoms, hessian_x=hessian_x, cartesian=False) with pytest.raises(CalculationFailed): hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy, hessian_x=hessian_x, get_q_from_x=water_get_q_from_x, get_jacobian=lambda x: np.ones((3, 9)), cartesian=True, variable_orientation=True) calc = HarmonicCalculator(hff) setup_water(calc) @pytest.mark.optimize() def test_internals(): parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy, 'hessian_x': hessian_x, 'get_q_from_x': water_get_q_from_x, 'get_jacobian': water_get_jacobian, 'cartesian': False} hff = HarmonicForceField(**parameters) # calculation in internals calc = HarmonicCalculator(hff) atoms = setup_water(calc) # distorted copy of ref_atoms run_optimize(atoms) # recover original configuration assert_water_is_relaxed(atoms) parameters['cartesian'] = True # calculation in Cartesian Coordinates hff = HarmonicForceField(**parameters) calc = HarmonicCalculator(hff) atoms = setup_water(calc) # 'variable_orientation' not set to True! # but water has rotational degrees of freedom run_optimize(atoms) with pytest.raises(AssertionError): # hence forces were incorrect assert_water_is_relaxed(atoms) # original configuration not recovered parameters['variable_orientation'] = True hff = HarmonicForceField(**parameters) calc = HarmonicCalculator(hff) atoms = setup_water(calc) run_optimize(atoms) assert_water_is_relaxed(atoms) # relaxation succeeded despite rotation @pytest.mark.optimize() def test_compatible_with_ase_vibrations(): atoms = ref_atoms.copy() atoms.calc = EMT() run_optimize(atoms) opt_atoms = atoms.copy() opt_energy = atoms.get_potential_energy() vib = Vibrations(atoms, nfree=2) vib.run() energies = vib.get_energies() vib_data = vib.get_vibrations() hessian_2d = vib_data.get_hessian_2d() vib.clean() hff = HarmonicForceField(ref_atoms=opt_atoms, ref_energy=opt_energy, hessian_x=hessian_2d) calc_harmonic = HarmonicCalculator(hff) atoms = ref_atoms.copy() atoms.calc = calc_harmonic vib = Vibrations(atoms, nfree=4, delta=1e-5) vib.run() assert np.allclose(energies, vib.get_energies()) vib.clean() hff = HarmonicForceField(ref_atoms=ref_atoms, ref_energy=ref_energy, hessian_x=hessian_2d, get_q_from_x=water_get_q_from_x, get_jacobian=water_get_jacobian, cartesian=True) calc_harmonic = HarmonicCalculator(hff) atoms = ref_atoms.copy() atoms.calc = calc_harmonic vib = Vibrations(atoms, nfree=4, delta=1e-5) vib.run() # 3 transl and 3 rot are removed by internal coordinates assert_array_almost_equal(energies[-3:], vib.get_energies()[-3:], decimal=2) def test_thermodynamic_integration(): parameters = {'ref_atoms': ref_atoms, 'ref_energy': ref_energy, 'hessian_x': hessian_x, 'get_q_from_x': water_get_q_from_x, 'get_jacobian': water_get_jacobian, 'cartesian': True, 'variable_orientation': True} hff_1 = HarmonicForceField(**parameters) calc_harmonic_1 = HarmonicCalculator(hff_1) parameters['cartesian'] = False hff_0 = HarmonicForceField(**parameters) calc_harmonic_0 = HarmonicCalculator(hff_0) ediffs = {} # collect energy difference for varying lambda coupling lambs = [0.00, 0.25, 0.50, 0.75, 1.00] # integration grid for lamb in lambs: ediffs[lamb] = [] calc_linearCombi = MixedCalculator(calc_harmonic_0, calc_harmonic_1, 1 - lamb, lamb) atoms = ref_atoms.copy() atoms.calc = calc_linearCombi MaxwellBoltzmannDistribution(atoms, temperature_K=300, force_temp=True) Stationary(atoms) ZeroRotation(atoms) with Andersen(atoms, 0.5 * fs, temperature_K=300, andersen_prob=0.05, fixcm=False) as dyn: for _ in dyn.irun(50): # should be much longer for production runs e0, e1 = calc_linearCombi.get_energy_contributions(atoms) ediffs[lamb].append(float(e1) - float(e0)) ediffs[lamb] = np.mean(ediffs[lamb]) dA = trapezoid([ediffs[lamb] for lamb in lambs], x=lambs) # anharm. corr. assert -0.005 < dA < 0.005 # the MD run is to short for convergence if dA == 0.0: raise ValueError('there is most likely something wrong, but it could ' 'also be sheer coincidence')