exploding solvate when using the openmm-alchemy tools

[LeoFHall]

Hi all,
hopefully this is the right place to ask about this,
I'm trying to understand how to use the openmm-tools library to do TI and solvate a glycine in water, but as soon as lambda is slightly low (0.8) the glycine coordinate explode all over the place. I seems like all intra-molecular forces are getting turned off rather than just electrostatics and LJ with the rest of the system ?
I've tried gradually turning down lambda during equilibration to see if that help but it doesn't, here's my code bellow
Thanks for the help

#always include at the top of simulation
import openmm.app as app
import openmm as mm
import openmm.unit as unit
import os
os.environ['JAX_ENABLE_X64'] = '1' #set JAX to 64bits by default
import openmmtools as tools
from openmmtools import integrators
from openmmtools import forces #still under construction, check for updates occasionally
from openmmtools import alchemy
from sys import stdout
from time import perf_counter as pf
import MDAnalysis as mda
import numpy as np

start = pf()

pdb = app.PDBFile('ZGY_1p5nmPadWaterCube.pdb')
forcefield = app.ForceField('./amber99sb-ZGY-hack.xml','/home/leo/.conda/envs/openmmEnv/lib/python3.10/site-packages/openmm/app/data/tip3p.xml')
#forcefield = app.ForceField('/home/hall2401/ForceFields/amber99sb-ZGY-hack.xml','/home/hall2401/ForceFields/tip3p.xml')

# System Configuration
nonbondedMethod = app.PME
nonbondedCutoff = 1.2*unit.nanometers
ewaldErrorTolerance = 0.0005
constraints = None
lambda_value = 0.8

# Integration Options
dt = 0.0005*unit.picoseconds # 0.5 femtoseconds
temperature = 298*unit.kelvin 
friction = 1.0/unit.picosecond
pressure = 1.0*unit.atmospheres
barostatInterval = 25

# Simulation Options
steps = 20000#00 #1ns with the 0.5fs timestep
equilibrationSteps = 10000#000 #0.5ns equilibration
CPUplatform = mm.Platform.getPlatformByName('CPU')
GPUplatform = mm.Platform.getPlatformByName('CUDA')
GPUplatformProperties = {'Precision': 'single'}
pdbReporter = app.PDBReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.pdb', 1000)
dataReporter = app.StateDataReporter('trajectory-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.csv', 1000, totalSteps=steps,
    step=True, speed=True, progress=True, potentialEnergy=True, temperature=True, separator=',')
checkpointReporter = app.CheckpointReporter('checkpoint-solvation_CONCGly_CONCNaCl_TEMP_LAMBDA_RUN0p1.chk', 10000)

# Create the system
print('Building system...')
topology = pdb.topology
positions = pdb.positions
system = forcefield.createSystem(topology=topology, nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
    constraints=constraints, ewaldErrorTolerance=ewaldErrorTolerance)

#Get residue of interest
target_residue_index = 0  # Index of the target glycine residue
target_atoms = [atom.index for atom in pdb.topology.atoms() if atom.residue.index == target_residue_index]

#create system where target atoms are alchemically modifiable
alchemical_region = alchemy.AlchemicalRegion(alchemical_atoms=target_atoms)
factory = alchemy.AbsoluteAlchemicalFactory()
alchemical_system = factory.create_alchemical_system(system,alchemical_region)

integrator = integrators.LangevinIntegrator(temperature, friction, dt)
#integrator = mm.LangevinMiddleIntegrator(temperature, friction, dt)
simulation = app.Simulation(topology, alchemical_system, integrator, GPUplatform, GPUplatformProperties)
simulation.context.setPositions(positions)


print("performing energy minimisation...")
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(temperature)

print("equilibrating system...")
simulation.reporters.append(dataReporter)
simulation.step(equilibrationSteps)

print("applying alchemical transformation")
#apply the Lambda parametre to the system
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
alchemical_state.lambda_electrostatics = 0.9
alchemical_state.lambda_sterics = 0.9
#alchemical_state.set_alchemical_parameters(lambda_value)
alchemical_state.apply_to_context(simulation.context)

print("equilibrating system again...")
#simulation.reporters.append(dataReporter)
simulation.currentStep = 0
simulation.step(equilibrationSteps)

print("applying alchemical transformation")
#apply the Lambda parametre to the system
alchemical_state = alchemy.AlchemicalState.from_system(alchemical_system)
alchemical_state.lambda_electrostatics = lambda_value
alchemical_state.lambda_sterics = lambda_value
#alchemical_state.set_alchemical_parameters(lambda_value)
alchemical_state.apply_to_context(simulation.context)

# Simulate
print('Simulating...')
simulation.reporters.append(pdbReporter)
#simulation.reporters.append(dataReporter) #already added during equilibration
simulation.reporters.append(checkpointReporter)
simulation.currentStep = 0
simulation.step(steps)

end = pf()
print('runtime (mn) = ',(end-start)/60)