2.5. ReaxFF: Gaseous H₂O¶

def run_reference():
    # choose a reference engine. If you do not have a license for ADF, use DFTB, ForceField, or another engine for which you have the license.
    # from_smiles('O') creates an H2O molecule
    job = AMSJob.from_input(name='H2O_pesscan', molecule=from_smiles('O'), text_input='''
    Task PESScan
    PESScan
        # Scan O-H bond length between 0.85 and 1.35 angstrom (11 points)
        ScanCoordinate
            nPoints 11
            Distance 1 2 0.85 1.35
        End
    End
    Engine ADF
        XC
            GGA PBE
        End
    EndEngine
    #Engine ForceField
    #EndEngine
    #Engine DFTB
    #EndEngine
    ''')
    job.run()
    return job

def create_job_collection():
    # Create a Job Collection, holding all jobs needed for the parametrization:
    jc = JobCollection()
    # Create a single Job Collection Entry with a H2O system:
    jce = JCEntry(molecule=from_smiles('O'))
    # The task is a PESScan with a H-O bond stretch, ranged [0.85, 1.35] angstrom:
    jce.settings.input.ams.task = 'pesscan'
    pes_points = np.linspace(0.85, 1.35, 11)
    jce.settings.input.ams.pesscan.scancoordinate.npoints = len(pes_points)
    jce.settings.input.ams.pesscan.scancoordinate.distance = f"1 2 {pes_points[0]} {pes_points[-1]}" # atom1 atom2 start end
    # Add the Entry to the Collection:
    jc.add_entry('water', jce)

    return jc, pes_points

def run_reference_calculations(jc : JobCollection):
    # The data set is missing reference values. Lets calculate them with PBE:
    pbe = Settings()
    pbe.input.adf.xc.gga = 'pbe'
    # Run all jobs in the Job Collection with the same settings:
    print('Running the reference calculation. This will take a moment ...')
    results = jc.run(pbe)

    return results

def import_results():
    # If you ran the PESScan job via the graphical user interface, set
    #pesscan_job = '/path/to/ams.rkf'
    # Otherwise, the below line runs the reference job
    pesscan_job = run_reference()
    ri = ResultsImporter(settings={'units': {'energy': 'kcal/mol'}})
    ri.add_singlejob(pesscan_job, task='PESScan', properties={
        'pes': {
            # apply a weights scheme (optional). 
            # This sets the weight for the larger energies (very short or very long bond distance) to smaller numbers.
            # stdev is the standard deviation of the Gaussian in energy units (here kcal/mol)
            'weights_scheme': WeightsSchemeGaussian(normalization=11.0, stdev=20.0)
        }
    })
    ri.store(binary=False, backup=False, store_settings=False)

def create_training_set():
    # Define the training set. Here, we will just add one entry: The PES of the H-O bond scan of water:
    ds = DataSet()
    ds.add_entry("pes('water')")
    return ds

def update_training_set_with_reference_results(ds : DataSet, results : dict):
    # Extract the relevant results and populate all reference values in the Data Set:
    ds.calculate_reference(results)

def create_parameter_interface():
    # start from the Water2017.ff force field
    parameter_interface = ReaxFFParameters('Water2017.ff', bounds_scale=1.2)
    parameter_interface.header['head'] = "Reparametrization of Water2017.ff"
    # only activate the "Standard" bond parameters for O-H
    for p in parameter_interface:
        p.is_active = (p.metadata['block'] == 'BND') and \
                      (p.metadata['category'] == 'Standard') and \
                      (p.metadata['atoms'] == ['H', 'O'] or p.metadata['atoms'] == ['O', 'H'])
    parameter_interface.yaml_store('parameter_interface.yaml')

def create_parameter_interface():
    # Load the ReaxFF force field we want to optimize:
    parameter_interface = ReaxFFParameters('Water2017.ff')
    # First line of the output ffield.ff file
    parameter_interface.header['head'] = "Reparametrization of Water2017.ff"
    # Mark only parameters from the H-O Bonds block for optimization:
    for parameter in parameter_interface:
        parameter.is_active = parameter.metadata['block'] == 'BND' and \
                              (parameter.metadata['atoms'] == ['H', 'O'] or parameter.metadata['atoms'] == ['O', 'H'])

    return parameter_interface

optimizer = CMAOptimizer(popsize=10, sigma=0.1, minsigma=1e-4)

max_evaluations = 100

logger_every = 5

def run_optimization(jc : JobCollection, ds : DataSet, pi : ReaxFFParameters):
    # Choose an optimization algorthm:
    o = CMAOptimizer(sigma=0.1) # sigma range: (0, 1). smaller sigma -> more local sampling

    # Piece it all together and run:
    O = Optimization(
        job_collection=jc,
        data_sets=ds,
        optimizer=o,
        parameter_interface=pi,
        callbacks=[MaxIter(100)], # The MaxIter callback stops an optimization after N evaluations
        )
    O.optimize()
    return O

def create_plots(workdir, pes_points):
    # At this point the optimization is done.
    # All relevant optimization data is stored in `O.workdir`
    # Data with the best parameters is stored in `training_set_results/best`
    # For more information on what is logged, refer to the `Logger()` callback documentation
    # Lets visualize some of it:
    from os.path import join as opj
    import matplotlib.pyplot as plt

    initial_ffield_ds     = DataSet( opj(workdir, 'training_set_results', 'initial', 'data_set_predictions.yaml') )
    initial_predictions   = initial_ffield_ds["pes('water')"].metadata['Prediction']
    optimized_ffield_ds   = DataSet( opj(workdir, 'training_set_results', 'best', 'data_set_predictions.yaml') )
    optimized_predictions = optimized_ffield_ds["pes('water')"].metadata['Prediction']
    reference             = initial_ffield_ds["pes('water')"].reference

    plt.rcParams['font.size'] = 16
    plt.figure(figsize=(8,6))
    plt.plot(pes_points, 2.6255e+03*reference, color='k', label='Reference') # convert hartree to kj/mol
    plt.plot(pes_points, 2.6255e+03*initial_predictions, label='Initial FF')
    plt.plot(pes_points, 2.6255e+03*optimized_predictions, label='Optimized FF')
    plt.ylabel('ΔE [kJ/mol]')
    plt.xlabel('H-O distance [Å]')
    plt.legend()
    plt.savefig('pes_results.png')