alchemiscale_gen_tyk2_network.py

import os
import copy
import json
import openfe
from openfe import ChemicalSystem
from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol as RHTP
from gufe import AlchemicalNetwork, tokenization
from openff.units import unit
from rdkit import Chem
from openfe_benchmarks import tyk2
from alchemiscale import AlchemiscaleClient, Scope, ScopedKey


settings = RHTP.default_settings()
settings.forcefield_settings.hydrogen_mass = 3.0
settings.simulation_settings.equilibration_length = 1000 * unit.picosecond
settings.simulation_settings.production_length = 5000 * unit.picosecond
settings.integrator_settings.timestep = 4.0 * unit.femtosecond
settings.integrator_settings.n_steps = 250 * unit.timestep
settings.alchemical_sampler_settings.n_repeats = 1
settings.alchemical_sampler_settings.online_analysis_interval = 200
settings.alchemical_sampler_settings.online_analysis_target_error = 0.2 * unit.boltzmann_constant * unit.kelvin
settings.simulation_settings.output_filename = f"transform.nc"
settings.simulation_settings.checkpoint_storage = f"transform_checkpoint.nc"
settings.simulation_settings.output_indices = "not water"
settings.engine_settings.compute_platform = "CUDA"
settings.system_settings.nonbonded_cutoff = 0.9 * unit.nanometer
solvent_protocol = RHTP(settings=settings)

system = tyk2.get_system()
edges = list(system.ligand_network.edges)

solvent_transformations = []
complex_transformations = []
alchemical_nodes = []

for edge in edges:
    stateA_solv = ChemicalSystem({'ligand': edge.componentA,
                                  'solvent': system.solvent_component})
    stateA_prot = ChemicalSystem({'ligand': edge.componentA,
                                  'protein': system.protein_component,
                                  'solvent': system.solvent_component})
    stateB_solv = ChemicalSystem({'ligand': edge.componentB,
                                  'solvent': system.solvent_component})
    stateB_prot = ChemicalSystem({'ligand': edge.componentB,
                                  'protein': system.protein_component,
                                  'solvent': system.solvent_component})
    alchemical_nodes.extend([stateA_prot, stateA_solv, stateB_prot, stateB_solv])

    solvent_transformations.append(
        openfe.Transformation(
            stateA=stateA_solv, stateB=stateB_solv,
            mapping={'ligand': edge},
            protocol=solvent_protocol,
        )
    )

    complex_transformations.append(
        openfe.Transformation(
            stateA=stateA_prot, stateB=stateB_prot,
            mapping={'ligand': edge},
            protocol=solvent_protocol,
        )
    )

network = AlchemicalNetwork(
    edges=(solvent_transformations + complex_transformations),
    nodes=alchemical_nodes,
    name="tyk2_off2.0_benchmark")

print("generated network: ", network)
with open('network.json', 'w') as f:
    json.dump(network.to_dict(), f, cls=tokenization.JSON_HANDLER.encoder)

user_id = os.environ['ALCHEMISCALE_ID']
user_key = os.environ['ALCHEMISCALE_KEY']

asc =  AlchemiscaleClient('https://api.alchemiscale.org', user_id, user_key)
scope = Scope('openfe', 'v0_7_3', 'tyk2_off2_0')
network_sk = asc.create_network(network, scope)
print(network_sk)

# store the scoped key
with open('scoped-key.dat', 'w') as f:
    f.write(str(network_sk))


# action out tasks
for transform in network.edges:
    transform_sk = asc.get_scoped_key(transform, scope)
    tasks = asc.create_tasks(transform_sk, count=3)
    asc.action_tasks(tasks, network_sk)