def run_simulation(smiles: str, n_nodes: int, spec_name: str = 'small_basis', solvent: Optional[str] = None)\
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Run the ionization potential computation
Args:
smiles: SMILES string to evaluate
n_nodes: Number of nodes to use
spec_name: Name of the quantum chemistry specification
solvent: Name of the solvent to use
Returns:
Relax records for the neutral and ionized geometry
"""
from moldesign.simulate.functions import generate_inchi_and_xyz, relax_structure, run_single_point
from moldesign.simulate.specs import get_qcinput_specification
from moldesign.utils.chemistry import get_baseline_charge
# Make the initial geometry
inchi, xyz = generate_inchi_and_xyz(smiles)
neu_charge = get_baseline_charge(smiles)
chg_charge = neu_charge + 1
# Make the compute spec
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2, 'ncores': 64}
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
if code == "nwchem":
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
# Compute the neutral geometry and hessian
neutral_xyz, _, neutral_relax = relax_structure(xyz, spec, compute_config=compute_config, charge=neu_charge, code=code)
# Compute the relaxed geometry
oxidized_xyz, _, oxidized_relax = relax_structure(neutral_xyz, spec, compute_config=compute_config, charge=chg_charge, code=code)
# If desired, compute the solvent energies
if solvent is None:
return [neutral_relax, oxidized_relax], []
spec, code = get_qcinput_specification(spec_name, solvent=solvent)
if code == "nwchem":
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
neutral_solvent = run_single_point(neutral_xyz, 'energy', spec, compute_config=compute_config, charge=neu_charge, code=code)
charged_solvent = run_single_point(oxidized_xyz, 'energy', spec, compute_config=compute_config, charge=chg_charge, code=code)
return [neutral_relax, oxidized_relax], [neutral_solvent, charged_solvent]
def _run_simulation(smiles: str, solvent: Optional[str], spec_name: str = 'xtb')\
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Run the ionization potential computation
Args:
smiles: SMILES string to evaluate
solvent: Name of the solvent
spec: Quantum chemistry specification for the molecule
Returns:
Relax records for the neutral and ionized geometry
"""
from moldesign.simulate.functions import generate_inchi_and_xyz, relax_structure, run_single_point
from moldesign.simulate.specs import get_qcinput_specification
from moldesign.utils.chemistry import get_baseline_charge
from qcelemental.models import DriverEnum
# Make the initial geometry
inchi, xyz = generate_inchi_and_xyz(smiles)
init_charge = get_baseline_charge(smiles)
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
# Compute the geometries
neutral_xyz, _, neutral_relax = relax_structure(xyz,
spec,
charge=init_charge,
code=code)
oxidized_xyz, _, oxidized_relax = relax_structure(neutral_xyz,
spec,
charge=init_charge + 1,
code=code)
# Perform the solvation energy computations, if desired
if solvent is None:
return [neutral_relax, oxidized_relax], []
solv_spec, code = get_qcinput_specification(spec_name, solvent=solvent)
neutral_solv = run_single_point(neutral_xyz,
DriverEnum.energy,
solv_spec,
charge=init_charge,
code=code)
oxidized_solv = run_single_point(oxidized_xyz,
DriverEnum.energy,
solv_spec,
charge=init_charge + 1,
code=code)
return [neutral_relax, oxidized_relax], [neutral_solv, oxidized_solv]
def _compute_adiabatic(xyz: str, init_charge: int, solvent: Optional[str], spec_name: str = 'xtb') \
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
# Compute the geometries
oxid_xyz, _, oxidized_relaxed = relax_structure(xyz,
spec,
charge=init_charge + 1,
code=code)
# Perform the solvation energy computations, if desired
if solvent is None:
return [oxidized_relaxed], []
solv_spec, code = get_qcinput_specification(spec_name, solvent=solvent)
oxidized_solv = run_single_point(oxid_xyz,
DriverEnum.energy,
solv_spec,
charge=init_charge + 1,
code=code)
return [oxidized_relaxed], [oxidized_solv]
def run_simulation(
smiles: str,
n_nodes: int) -> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Run the ionization potential computation
Args:
smiles: SMILES string to evaluate
n_nodes: Number of nodes to use
Returns:
Relax records for the neutral and ionized geometry
"""
from moldesign.simulate.functions import generate_inchi_and_xyz, relax_structure, run_single_point
from moldesign.simulate.specs import get_qcinput_specification
from qcelemental.models import DriverEnum
# Make the initial geometry
inchi, xyz = generate_inchi_and_xyz(smiles)
# Make the compute spec
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
# Get the specification and make it more resilient
spec, code = get_qcinput_specification('small_basis')
if code == "nwchem":
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
# Compute the neutral geometry and hessian
neutral_xyz, _, neutral_relax = relax_structure(
xyz, spec, compute_config=compute_config, charge=0, code=code)
neutral_hessian = run_single_point(neutral_xyz,
DriverEnum.hessian,
spec,
charge=0,
compute_config=compute_config,
code=code)
# Compute the relaxed geometry
oxidized_xyz, _, oxidized_relax = relax_structure(
neutral_xyz, spec, compute_config=compute_config, charge=1, code=code)
oxidized_hessian = run_single_point(oxidized_xyz,
DriverEnum.hessian,
spec,
charge=1,
compute_config=compute_config,
code=code)
return [neutral_relax, oxidized_relax], [neutral_hessian, oxidized_hessian]
def compute_adiabatic(xyz: str, init_charge: int, oxidize: bool,
spec_name: str = 'small_basis', n_nodes: int = 2) \
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Compute the adiabatic ionization potential starting from a neutral geometry
Just relaxes the structure in the oxidized state
Args:
xyz: XYZ of the neutral geometry
oxidize: Compute the oxidation potential instead of reduction
init_charge: Charge of the nuetral molecule
spec_name: Name of the computation spec to perform
n_nodes: Number of nodes for the computation
Returns:
- Relaxation record
- Not used, but to make the same interface for compute_vertical and
"""
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
if code == "nwchem":
spec.keywords['dft__convergence__energy'] = 1e-7
spec.keywords['dft__convergence__fast'] = True
spec.keywords["dft__iterations"] = 150
spec.keywords["driver__maxiter"] = 150
spec.keywords["geometry__noautoz"] = True
# Make sure to allow restarting
spec.extras["allow_restarts"] = True
runhash = hashlib.sha256(
f'{xyz}_{oxidize}_{spec_name}'.encode()).hexdigest()[:12]
spec.extras["scratch_name"] = f'nwc_{runhash}'
# Compute the geometries
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
new_charge = init_charge + 1 if oxidize else init_charge - 1
_, _, oxidized_relaxed = relax_structure(xyz,
spec,
charge=new_charge,
code=code,
compute_config=compute_config)
# Perform the solvation energy computations, if desired
return [oxidized_relaxed], []
def compute_single_point(xyz: str, charge: int, solvent: Optional[str] = None,
spec_name: str = 'normal_basis', n_nodes: int = 2) \
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Perform a single point energy computation, return results in same format as other assays
Args:
xyz: Molecular geometry
charge: Molecular charge
solvent: Name of the solvent, if desired
spec_name: Name of the QC specification
n_nodes: Number of nodes for the computation
Returns:
- Not used
- Single point energy computation
"""
# Get the specification and make it more resilient
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
if spec_name == 'diffuse_basis':
compute_config['cores_per_rank'] = 8
spec, code = get_qcinput_specification(spec_name, solvent)
if code == "nwchem":
# Reduce the accuracy needed to 1e-7
spec.keywords['dft__convergence__energy'] = 1e-7
spec.keywords['dft__convergence__fast'] = True
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
# Make sure to allow restarting
spec.extras["allow_restarts"] = True
runhash = hashlib.sha256(
f'{xyz}_{charge}_{spec_name}_{solvent}'.encode()).hexdigest()[:12]
spec.extras["scratch_name"] = f'nwc_{runhash}'
# Run the computation
spe_record = run_single_point(xyz,
DriverEnum.energy,
spec,
charge=charge,
code=code,
compute_config=compute_config)
return [], [spe_record]
def run_simulation(smiles: str, n_nodes: int,
mode: str) -> Union[OptimizationResult, AtomicResult]:
"""Run a single-point or relaxation computation computation
Args:
smiles: SMILES string to evaluate
n_nodes: Number of nodes to use
mode: What computation to perform: single, gradient, hessian, relax
Returns:
Result of the energy computation
"""
from moldesign.simulate.functions import generate_inchi_and_xyz, relax_structure, run_single_point
from moldesign.simulate.specs import get_qcinput_specification
from qcelemental.models import DriverEnum
# Make the initial geometry
inchi, xyz = generate_inchi_and_xyz(smiles)
# Make the compute spec
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
# Get the specification and make it more resilient
spec, code = get_qcinput_specification('small_basis')
if code == "nwchem":
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
# Compute the neutral geometry and hessian
if mode == 'relax':
_, _, neutral_relax = relax_structure(xyz,
spec,
compute_config=compute_config,
charge=0,
code=code)
return neutral_relax
else:
return run_single_point(xyz,
mode,
spec,
charge=0,
compute_config=compute_config,
code=code)
def get_relaxation_args(xyz: str, charge: int, spec_name: str = 'small_basis', n_nodes: int = 2) \
-> Tuple[List[Any], Dict[str, Any]]:
"""Get the function inputs to a relaxation computation
Args:
xyz: XYZ of the starting geometry
charge: Charge of the molecule
spec_name: Name of the computation spec to perform
n_nodes: Number of nodes for the computation
Returns:
- Relaxation record
- Not used, but to make the same interface for compute_vertical and
"""
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
if code == "nwchem":
spec.keywords['dft__convergence__energy'] = 1e-7
spec.keywords['dft__convergence__fast'] = True
spec.keywords["dft__iterations"] = 150
spec.keywords["driver__maxiter"] = 150
spec.keywords["geometry__noautoz"] = True
# Make sure to allow restarting
spec.extras["allow_restarts"] = True
runhash = hashlib.sha256(
f'{xyz}_{charge}_{spec_name}'.encode()).hexdigest()[:12]
spec.extras["scratch_name"] = f'nwc_{runhash}'
# Set up compute configuration
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
return [xyz, spec], dict(charge=charge,
code=code,
compute_config=compute_config)
compute_config["ncores"] = 64 // args.qc_parallelism
else:
qc_workers = nnodes // args.qc_parallelism
run_params["nnodes"] = nnodes
run_params["qc_workers"] = qc_workers
# Load in the models, initial dataset, agent and search space
with open(os.path.join(args.mpnn_config_directory,
'atom_types.json')) as fp:
atom_types = json.load(fp)
with open(os.path.join(args.mpnn_config_directory,
'bond_types.json')) as fp:
bond_types = json.load(fp)
# Get QC specification
qc_spec, code = get_qcinput_specification(args.qc_spec)
if args.qc_spec != "xtb":
qc_spec.keywords["dft__iterations"] = 150
qc_spec.keywords["geometry__noautoz"] = True
ref_energies = lookup_reference_energies(args.qc_spec)
# Create an output directory with the time and run parameters
start_time = datetime.utcnow()
params_hash = hashlib.sha256(
json.dumps(run_params).encode()).hexdigest()[:6]
out_dir = os.path.join(
'runs',
f'{args.qc_spec}-{start_time.strftime("%d%b%y-%H%M%S")}-{params_hash}')
os.makedirs(out_dir, exist_ok=True)
# Save the run parameters to disk
def compute_vertical(smiles: str, oxidize: bool, solvent: Optional[str] = None,
spec_name: str = 'small_basis', n_nodes: int = 2) \
-> Tuple[List[OptimizationResult], List[AtomicResult]]:
"""Perform the initial ionization potential computation of the vertical
First relaxes the structure and then runs a single-point energy at the
Args:
smiles: SMILES string to evaluate
oxidize: Whether to perform an oxidation or reduction
solvent: Name of the solvent, if desired. Runs on the neutral geometry after relaxation
spec_name: Quantum chemistry specification for the molecule
n_nodes: Number of nodes per computation
Returns:
- Relax records for the neutral
- Single point energy in oxidized or reduced state
"""
# Make the initial geometry
inchi, xyz = generate_inchi_and_xyz(smiles)
init_charge = get_baseline_charge(smiles)
# Make the compute spec
compute_config = {'nnodes': n_nodes, 'cores_per_rank': 2}
# Get the specification and make it more resilient
spec, code = get_qcinput_specification(spec_name)
if code == "nwchem":
spec.keywords['dft__convergence__energy'] = 1e-7
spec.keywords['dft__convergence__fast'] = True
spec.keywords["dft__iterations"] = 150
spec.keywords["driver__maxiter"] = 150
spec.keywords["geometry__noautoz"] = True
# Make a repeatably-named scratch directory.
# We cannot base it off a hash of the input file,
# because the XYZ file generator is stochastic.
runhash = hashlib.sha256(
f'{smiles}_{oxidize}_{spec_name}'.encode()).hexdigest()[:12]
spec.extras["scratch_name"] = f'nwc_{runhash}'
spec.extras["allow_restarts"] = True
# Compute the geometries
neutral_xyz, _, neutral_relax = relax_structure(
xyz,
spec,
charge=init_charge,
code=code,
compute_config=compute_config)
# Perform the single-point energy for the ionized geometry
new_charge = init_charge + 1 if oxidize else init_charge - 1
oxid_spe = run_single_point(neutral_xyz,
DriverEnum.energy,
spec,
charge=new_charge,
code=code,
compute_config=compute_config)
# If desired, submit a solvent computation as well
if solvent is None:
return [neutral_relax], [oxid_spe]
spec, code = get_qcinput_specification(spec_name, solvent)
if code == "nwchem":
# Reduce the accuracy needed to 1e-7
spec.keywords['dft__convergence__energy'] = 1e-7
spec.keywords['dft__convergence__fast'] = True
spec.keywords["dft__iterations"] = 150
spec.keywords["geometry__noautoz"] = True
# Make sure to allow restarting
spec.extras["allow_restarts"] = True
solv_spe = run_single_point(neutral_xyz,
DriverEnum.energy,
spec,
charge=init_charge,
code=code,
compute_config=compute_config)
return [neutral_relax], [oxid_spe, solv_spe]
args = parser.parse_args()
# Get the number of cores per node
with open('qcengine.yaml') as fp:
qce_settings = yaml.load(fp)
proc_per_node = qce_settings['all']['ncores']
# Load in the molecules
mols = pd.read_csv("example_molecules.csv")
mols = mols[mols["num_electrons"].apply(lambda x: x in args.sizes)]
logging.info(
f"Pulled {len(mols)} for testing with sizes: {mols['num_electrons'].to_list()}"
)
# Get the run config
spec, program = get_qcinput_specification(args.config)
logging.info(f"Pulled spec for '{args.config}'")
# Loop over the molecules
n_nodes = int(os.environ.get("COBALT_JOBSIZE", "1"))
for _, mol in mols.iterrows():
# Parse the molecule
mol_obj = qcel.models.Molecule.from_data(mol["xyz"],
molecular_charge=args.charge)
# Make the input
input_obj = qcel.models.AtomicInput(molecule=mol_obj,
driver=args.driver,
**spec.dict(exclude={'driver'}))
# Run the calculation
"""Generate the data used by the tests"""
from moldesign.simulate.functions import generate_inchi_and_xyz, relax_structure, run_single_point
from moldesign.simulate.specs import get_qcinput_specification
if __name__ == "__main__":
# Make a water molecule
inchi, xyz = generate_inchi_and_xyz('O')
# Generate the neutral geometry with XTB
xtb_spec, xtb = get_qcinput_specification("xtb")
xtb_neutral_xyz, _, record = relax_structure(xyz, xtb_spec, code=xtb)
with open('records/xtb-neutral.json', 'w') as fp:
print(record.json(), file=fp)
# Compute the vertical oxidation potential
record = run_single_point(xtb_neutral_xyz,
"energy",
xtb_spec,
code=xtb,
charge=1)
with open('records/xtb-neutral_xtb-oxidized-energy.json', 'w') as fp:
print(record.json(), file=fp)
# Compute the adiabatic oxidation potential
xtb_oxidized_xyz, _, record = relax_structure(xtb_neutral_xyz,
xtb_spec,
code=xtb,
charge=1)
with open('records/xtb-oxidized.json', 'w') as fp:
print(record.json(), file=fp)
