profile_lammps.py

import mbuild as mb
import matplotlib.pyplot as plt
import time
import foyer

def build_mbuild():
    from mbuild.formats.lammpsdata import write_lammpsdata
    from foyer import Forcefield
    
    water = mb.load("files/tip3p.mol2")
    water.name = "SOL"
    
    time_list = list()
    size = [100, 500, 1000, 2000]
    for i in size:
        water_box = mb.fill_box(compound=water,
                n_compounds=i,
                density=1000)
        
        forcefield = Forcefield("files/tip3p.xml")
        
        system = forcefield.apply(water_box)
        
        t0 = time.time() 
        write_lammpsdata(system, 'data.mbuid')
        t1 = time.time()
        time_list.append(t1-t0)

    return time_list

def build_gmso():
    from gmso.formats.lammpsdata import write_lammpsdata
    from gmso import ForceField
    from gmso.external import convert_mbuild
    import gmso
    water = mb.load("files/tip3p.mol2")
    water = water.children[0]
    water.name = "water"
    
    # element_map = which site name corresponds to which atom_type name. 
    # In the future atomtyping will be done through foyer. 
    element_map = {"O": "opls_111",
                   "H": "opls_112"}

    time_list = list()
    size = [100, 500, 1000, 2000]
    for i in size:
        water_box = mb.fill_box(compound=water,
                n_compounds=i,
                density=1000)
        
        forcefield = ForceField("files/gmso_tip3p.xml")

        top = convert_mbuild.from_mbuild(water_box)

        # Assign atom types
        for atom in top.sites:
            atom.atom_type = forcefield.atom_types[element_map[atom.name]]
            
        # Assign bond types
        for bond in top.bonds:
            bond.bond_type = bond.connection_type = forcefield.bond_types["opls_111~opls_112"]
        
        # Create angles with correct atom type and add to top
        for subtop in top.subtops:
            angle = gmso.core.angle.Angle(
                connection_members=[site for site in subtop.sites],
                name="opls_112~opls_111~opls_112",
                connection_type=forcefield.angle_types["opls_112~opls_111~opls_112"]
            )
            top.add_connection(angle, update_types=False)
        
        top.update_topology()

        t0 = time.time()
        write_lammpsdata(top, 'data.gmso')
        t1 = time.time()

        time_list.append(t1-t0)

    return time_list

mbuild_time = build_mbuild()
gmso_time = build_gmso()
size = [100, 500, 1000, 2000]
plt.plot(size, mbuild_time, label='mBuild')
plt.plot(size, gmso_time, label='gmso')
plt.legend()
plt.xlabel('system size')
plt.ylabel('Time (s)')
plt.savefig('times.pdf')