def initialize_from_api(cls, storage_socket, queue_socket, meta, molecule):
# Grab initial molecule
meta["initial_molecule"] = molecule["id"]
# Copy initial intial input and build out a torsiondrive_state
meta = copy.deepcopy(meta)
# Remove identity info from template
molecule_template = copy.deepcopy(molecule)
del molecule_template["id"]
del molecule_template["identifiers"]
# Initiate torsiondrive meta
meta["torsiondrive_state"] = td_api.create_initial_state(
dihedrals=meta["torsiondrive_meta"]["dihedrals"],
grid_spacing=meta["torsiondrive_meta"]["grid_spacing"],
elements=molecule_template["symbols"],
init_coords=[molecule_template["geometry"]])
# Save initial molecule and add hash
meta["status"] = "READY"
meta["required_jobs"] = False
meta["remaining_jobs"] = False
meta["molecule_template"] = molecule_template
meta["optimization_history"] = {}
dihedral_template = []
for idx in meta["torsiondrive_meta"]["dihedrals"]:
tmp = ('dihedral', ) + tuple(str(z + 1) for z in idx)
dihedral_template.append(tmp)
meta["torsiondrive_meta"]["dihedral_template"] = dihedral_template
# Move around geometric data
meta["optimization_program"] = meta["optimization_meta"].pop("program")
# Temporary hash index
single_keys = copy.deepcopy(meta["qc_meta"])
single_keys["molecule_id"] = meta["initial_molecule"]
keys = {
"type": "torsiondrive",
"program": "torsiondrive",
"keywords": meta["torsiondrive_meta"],
"optimization_keys": {
"procedure": meta["optimization_program"],
"keywords": meta["optimization_meta"],
},
"single_keys": schema.format_result_indices(single_keys)
}
meta["success"] = False
meta["procedure"] = "torsiondrive"
meta["program"] = "torsiondrive"
meta["hash_index"] = procedures.procedures_util.hash_procedure_keys(
keys),
meta["hash_keys"] = keys
meta["tag"] = None
return cls(storage_socket, queue_socket, meta)
def run_torsiondrive_scan(self):
"""
Run torsiondrive scan in the following steps:
1. Create json input for torsiondrive
2. Send the json input dictionary to td_api.next_jobs_from_state(), get the next set of jobs
3. If there are no jobs needed, finish and return the lowest energy on each dihedral grid
4. If there are new jobs, run them with geomeTRIC.run_json.
5. Collect the results and put them into new json input dictionary
6. Go back to step 2.
"""
# step 1
td_state = td_api.create_initial_state(
dihedrals=self.dihedrals,
grid_spacing=self.grid_spacing,
elements=self.elements,
init_coords=self.init_coords,
dihedral_ranges=self.dihedral_ranges,
energy_decrease_thresh=self.energy_decrease_thresh,
energy_upper_limit=self.energy_upper_limit,
)
while True:
# step 2
next_jobs = td_api.next_jobs_from_state(td_state, verbose=True)
# step 3
if len(next_jobs) == 0:
print("torsiondrive Scan Finished")
return td_api.collect_lowest_energies(td_state)
# step 4
job_results = collections.defaultdict(list)
for grid_id_str, job_geo_list in next_jobs.items():
for job_geo in job_geo_list:
dihedral_values = td_api.grid_id_from_string(grid_id_str)
# Run geometric
geometric_input_dict = self.make_geomeTRIC_input(
dihedral_values, job_geo)
geometric_output_dict = geometric.run_json.geometric_run_json(
geometric_input_dict)
# Pull out relevevant data
final_result = geometric_output_dict['trajectory'][-1]
final_geo = final_result['molecule']['geometry']
final_energy = final_result['properties']['return_energy']
# Note: the results should be appended in the same order as in the inputs
# It's not a problem here when running serial for loop
job_results[grid_id_str].append(
(job_geo, final_geo, final_energy))
# step 5
td_api.update_state(td_state, job_results)
def create_state(self):
"""
:return:
"""
self.td_state = td_api.create_initial_state(
dihedrals=self.dihedrals,
grid_spacing=self.grid_spacing,
init_coords=self.init_coords,
elements=self.elements,
dihedral_ranges=self.dihedral_ranges,
energy_decrease_thresh=self.energy_decrease_thresh,
energy_upper_limit=self.energy_upper_limit,
)
def _create_initial_state(
self, molecule: "Ligand", dihedral_data: TorsionScan
) -> Dict[str, Any]:
"""
Create the initial state for the torsion drive using the input settings.
Note:
We also put the running spec into the state.
Args:
molecule: The molecule we want to create our initial state for.
dihedral_data: The dihedral and scan range information.
Returns:
The torsiondrive dict or the initial state.
"""
coords = copy.deepcopy(molecule.coordinates)
td_state = td_api.create_initial_state(
dihedrals=[
dihedral_data.torsion,
],
grid_spacing=[
self.grid_spacing,
],
elements=[atom.atomic_symbol for atom in molecule.atoms],
init_coords=[(coords * constants.ANGS_TO_BOHR)],
dihedral_ranges=[
dihedral_data.scan_range,
],
energy_decrease_thresh=self.energy_decrease_thresh,
energy_upper_limit=self.energy_upper_limit,
)
td_state["spec"] = {
"program": self.program.lower(),
"method": self.method.lower(),
"basis": self.basis.lower() if self.basis is not None else self.basis,
}
return td_state
def initialize_from_api(cls,
storage_socket,
logger,
service_input,
tag=None,
priority=None):
_check_td()
import torsiondrive
from torsiondrive import td_api
# Build the record
# TODO: This removes the task_id which may exist on old records, but does not exist
# in newer TorsionDriveRecords.
# If a proper migration is ever done,
output = TorsionDriveRecord(
**service_input.dict(exclude={"initial_molecule", "task_id"}),
initial_molecule=[x.id for x in service_input.initial_molecule],
provenance={
"creator": "torsiondrive",
"version": torsiondrive.__version__,
"routine": "torsiondrive.td_api",
},
final_energy_dict={},
minimum_positions={},
optimization_history={},
)
meta = {"output": output}
# Remove identity info from molecule template
molecule_template = copy.deepcopy(
service_input.initial_molecule[0].dict(encoding="json"))
molecule_template.pop("id", None)
molecule_template.pop("identifiers", None)
meta["molecule_template"] = json.dumps(molecule_template)
# Initiate torsiondrive meta
# The torsiondrive package uses print, so capture that using
# contextlib
td_stdout = io.StringIO()
with contextlib.redirect_stdout(td_stdout):
meta["torsiondrive_state"] = td_api.create_initial_state(
dihedrals=output.keywords.dihedrals,
grid_spacing=output.keywords.grid_spacing,
elements=molecule_template["symbols"],
init_coords=[
x.geometry for x in service_input.initial_molecule
],
dihedral_ranges=output.keywords.dihedral_ranges,
energy_decrease_thresh=output.keywords.energy_decrease_thresh,
energy_upper_limit=output.keywords.energy_upper_limit,
)
meta["stdout"] = td_stdout.getvalue()
# Build dihedral template
dihedral_template = []
for idx in output.keywords.dihedrals:
tmp = {"type": "dihedral", "indices": idx}
dihedral_template.append(tmp)
meta["dihedral_template"] = json.dumps(dihedral_template)
# Build optimization template
opt_template = {
"meta": {
"procedure": "optimization",
"qc_spec": output.qc_spec.dict(),
"tag": meta.pop("tag", None)
}
}
opt_template["meta"].update(output.optimization_spec.dict())
meta["optimization_template"] = json.dumps(opt_template)
# Move around geometric data
meta["optimization_program"] = output.optimization_spec.program
meta["hash_index"] = output.get_hash_index()
meta["task_tag"] = tag
meta["task_priority"] = priority
return cls(**meta, storage_socket=storage_socket, logger=logger)
def initialize_from_api(cls,
storage_socket,
logger,
service_input,
tag=None,
priority=None):
_check_td()
import torsiondrive
from torsiondrive import td_api
# Build the record
output = TorsionDriveRecord(**service_input.dict(
exclude={"initial_molecule"}),
initial_molecule=[
x.id
for x in service_input.initial_molecule
],
provenance={
"creator": "torsiondrive",
"version": torsiondrive.__version__,
"routine": "torsiondrive.td_api"
},
final_energy_dict={},
minimum_positions={},
optimization_history={})
meta = {"output": output}
# Remove identity info from molecule template
molecule_template = copy.deepcopy(
service_input.initial_molecule[0].json_dict())
molecule_template.pop("id", None)
molecule_template.pop("identifiers", None)
meta["molecule_template"] = json.dumps(molecule_template)
# Initiate torsiondrive meta
meta["torsiondrive_state"] = td_api.create_initial_state(
dihedrals=output.keywords.dihedrals,
grid_spacing=output.keywords.grid_spacing,
elements=molecule_template["symbols"],
init_coords=[x.geometry for x in service_input.initial_molecule],
dihedral_ranges=output.keywords.dihedral_ranges,
energy_decrease_thresh=output.keywords.energy_decrease_thresh,
energy_upper_limit=output.keywords.energy_upper_limit)
# Build dihedral template
dihedral_template = []
for idx in output.keywords.dihedrals:
tmp = {"type": "dihedral", "indices": idx}
dihedral_template.append(tmp)
meta["dihedral_template"] = json.dumps(dihedral_template)
# Build optimization template
meta["optimization_template"] = json.dumps({
"meta": {
"procedure": "optimization",
"keywords": output.optimization_spec.keywords,
"program": output.optimization_spec.program,
"qc_spec": output.qc_spec.dict(),
"tag": meta.pop("tag", None)
},
})
# Move around geometric data
meta["optimization_program"] = output.optimization_spec.program
meta["hash_index"] = output.get_hash_index()
meta["task_tag"] = tag
meta["task_priority"] = priority
return cls(**meta, storage_socket=storage_socket, logger=logger)
def _create_initial_state(
self,
molecule: "Ligand",
dihedral_data: TorsionScan,
qc_spec: "QCOptions",
seed_coordinates: Optional[List[np.ndarray]] = None,
) -> Dict[str, Any]:
"""
Create the initial state for the torsion drive using the input settings.
Note:
We also put the running spec into the state.
Args:
molecule: The molecule we want to create our initial state for.
dihedral_data: The dihedral and scan range information.
Returns:
The torsiondrive dict or the initial state.
"""
if seed_coordinates is not None:
coords = seed_coordinates
elif self.starting_conformations > 1:
coords = molecule.generate_conformers(
n_conformers=self.starting_conformations)
else:
coords = [
copy.deepcopy(molecule.coordinates),
]
# write out the starting geometries
molecule.to_multiconformer_file(file_name="starting_coords.xyz",
positions=coords)
td_state = td_api.create_initial_state(
dihedrals=[
dihedral_data.torsion,
],
grid_spacing=[
self.grid_spacing,
],
elements=[atom.atomic_symbol for atom in molecule.atoms],
init_coords=[(coord * constants.ANGS_TO_BOHR) for coord in coords],
dihedral_ranges=[
dihedral_data.scan_range,
],
energy_decrease_thresh=self.energy_decrease_thresh,
energy_upper_limit=self.energy_upper_limit,
)
td_dependent_settings = (None if qc_spec.td_settings is None else
qc_spec.td_settings.dict())
td_state["spec"] = {
"program":
qc_spec.program.lower(),
"method":
qc_spec.method.lower(),
"basis":
qc_spec.basis.lower()
if qc_spec.basis is not None else qc_spec.basis,
"td":
td_dependent_settings,
}
return td_state
