openmmtools nc file parser for alchemlyb analysis

[coparks2012]

Hello,

I am performing ABFE calculation using openmm with the openmmtools toolkit for Hamiltonian replica exchange. I would very much like to be able to process the resulting nc files in a manner that would enable analysis with alchemlyb. In particular, I am interested in plotting the forward and backward convergence.

I am curious regarding whether there is a way to parse the information stored in a MultistateReporter or MultiStateAnalyzer openmmtools variable so that is can be subsequently processed using alchemlyb.

Any help pointing me in the right direction would be greatly appreciated.

[xiki-tempula]

Do you know if the openmmtools is API stable? I have a bunch of old codes that I used to parse the openmmtools but I didn't put it here as it is said that openmmtools is not API stable. I will have a dig and see if I can get it here.

[coparks2012]

Thank you for your response. Any help would be appreciated!

I think openmmtools is under development, but the primary infrastructure needed to perform ABFE calculations seems to be in place. An ABFE implementation was recently benchmarked (https://doi.org/10.1021/acs.jctc.1c01208) and is also available as part of BATPY (https://github.com/GHeinzelmann/BAT.py). I have been performing retrospective tests on the BRD4 system, and have seen quality reproduction thus far.

If others have had a different experience with openmm, I would love to hear more.

[xiki-tempula]

In that case, I think an interface could be supplied. I wonder if I could have some input to have a test? Preferably a set of nc files for different lambdas.

[coparks2012]

Thank you very much for your time.

I have uploaded a relevant nc file, yaml file containing estimated FE estimates at various % values of the run, and rest.py file containing a RestraintComposableState class definition used to modify/set lambda values to github:

https://github.com/coparks2012/openmmtools_alchemlyb_dev

Please let me know if you have any issues accessing the files,, and thank you again.

[xiki-tempula]

Sorry, I'm not too familiar with the NC file generated by OpenMMTools, I could retrieve the reduced potential array but I have trouble finding the all other lambda states.
I wonder if you know how to find a list of all the lambda states? Thanks.

[coparks2012]

Thanks!

Would it be possible to hard specify the lambda values on your end while I work on figuring out how to parse them from the nc file? In general, performing 'grep "lambdas =" rest.py' will print the lambda values to the terminal. For this run, here are those values:

lambdas = [0.00000 , 0.00100 , 0.00240 , 0.00560 , 0.01330 , 0.03160 , 0.07500 , 0.17780 , 0.42170 , 1.00000]

[xiki-tempula]

Asking the openmmtools people if they could include the relevant information in the NC file choderalab/openmmtools#652

[xiki-tempula]

Do you know how to directly read this information from NC file without having openmmtools as a dependency?
I could do

import netCDF4 as nc
df = nc.Dataset('restraint.nc')
df.groups['thermodynamic_states'].variables

But I cannot figure out what to do after this

[coparks2012]

Regarding the error, including the class definition at the beginning of your script should do the trick. Its defined in rest.py:

class RestraintComposableState(GlobalParameterState):
	lambda_restraints=GlobalParameterState.GlobalParameter('lambda_restraints', standard_value=1.0)

I have never tried anything besides openmm to open the file. What calculations are you attempting to perform?

[xiki-tempula]

So the question is

I am performing ABFE calculation using openmm with the openmmtools toolkit for Hamiltonian replica exchange. I would very much like to be able to process the resulting nc files in a manner that would enable analysis with alchemlyb. In particular, I am interested in plotting the forward and backward convergence.

I think the difficulty would be that I will need to make a parser that inputs a NC file generated by openmmtools and output a alchemlyb pandas dataframe. I could extract the reduced potential matrix from the NC file but can't extract the thermodynamic_states from NC file (or I don't know how to extract them).
It seems that one could use openmmtools to extract the thermodynamic_states information but this would need openmmtools, which would be a quite heavy dependency for a light-weighted library like alchemlyb. I was hoping that I could just use netcdf to open a NC file and extract the thermodynamic_states information.

[coparks2012]

Thank you very much for your continued help!

Regarding the thermodynamic states, did the suggestion from the openmmtools team not work, or is your concern you cannot do it independently of openmmtools?

Would it be possible to parse the NC file with openmmtols, and write a file to storage that alchemlyb would like (i.e. keep the two environments seperate for now)? Perhaps this could be a quick, although hacky, solution for the time being.
This would allow a user to hop into their alchemlyb environment for the analysis without having to load openmmtools. Do you have thoughts regarding how code to perform this function (use openmm to write a file for alchemlyb) would look like? I had a hard time understanding the format of the alchemlyb dataframe.

In the spirit of this, I have been able to hack into the alchemlyb package to generate the forward FE and mixing plot which I can readily share. However, I have been unable to determine how to parse the NC file/prep the array for alchemlyb to create the reverse FE plot. I haven't attempted any of the other plots.

[xiki-tempula]

Regarding the thermodynamic states, did the suggestion from the openmmtools team not work, or is your concern you cannot do it independently of openmmtools?

Yes, I could extract all of the lambda states but the script doesn't extract the current lambda state, which is needed for decorrelation.

Would it be possible to parse the NC file with openmmtols, and write a file to storage that alchemlyb would like (i.e. keep the two environments seperate for now)? Perhaps this could be a quick, although hacky, solution for the time being.

The way that I'm currently using is hracky where I just save the reduced potential array as a numpy array. Then I supply the lambda states and the current lambda myself, since I'm running the simulation myself, so I know the how I schedule the lambdas.

I have a script that I used to generate the alchemlyb style pandas dataframe from a numpy array of reduced potential.

def extract_u_nk(
    array: np.ndarray,
    current_lambda: pd.Series,
    lambda_list: pd.DataFrame,
    time: pd.Series,
    temperature: float = 298.0,
    energy_unit: str = "kT",
) -> pd.DataFrame:
    """Convert the numpy array into a pandas u_nk dataframe for the
    consumption of alchemlyb.

    :param array: The array of reduced potential.
    :param current_lambda: The current lambda Series.
    :param lambda_list: The DataFrame of all lambda.
    :param time: The time Series has "time" as name.
    :return: alchemlyb style u_nk dataframe
    """
    if time.name != "time":
        raise ValueError(f"The name of the time ({time.name}) has to be 'time'.")
    n_samples, n_lambdas = array.shape
    if n_samples != len(time):
        raise ValueError(
            f"The length of the time array ({len(time)}) is not"
            f"the same as the length of the array ({n_samples})."
        )
    if not any((lambda_list == current_lambda).all(axis=1)):
        raise ValueError(
            f"The current_lambda ({current_lambda}) is not in"
            f"the lambda_list ({lambda_list})."
        )

    lambda_dict = lambda_list.to_dict()
    labels = list(lambda_dict.keys())
    if len(labels) == 1:
        columns = [lambda_dict[labels[0]][i] for i in range(len(lambda_list))]
    else:
        columns = [
            tuple([lambda_dict[key][i] for key in labels])
            for i in range(len(lambda_list))
        ]

    u_nk = pd.DataFrame(
        data=array,
        index=time,
        columns=columns,
    )

    for key in labels:
        u_nk[key] = current_lambda[key]

    # set up new multi-index
    newind = ["time"] + list(labels)
    u_nk = u_nk.reset_index().set_index(newind)
    u_nk.name = "u_nk"
    u_nk.attrs = {"temperature": temperature, "energy_unit": energy_unit}
    return u_nk

You could use it as

        array = np.zeros((10, 4))
        time = pd.Series(np.arange(10), name="time")
        lambda_list = pd.DataFrame(
            data={
                "coul-lambda": [1.0, 0.5, 0.0, 0.0],
                "vdw-lambda": [1.0, 1.0, 1.0, 0.0],
            }
        )
        current_lambda = pd.Series(data={"coul-lambda": 0.5, "vdw-lambda": 1.0})
        df = extract_u_nk(array, current_lambda, lambda_list, time)

Sorry, I'm quite busy at the moment so cannot answer the question as fast as I should but if you would like to improve alchemlyb, I'm happy to review it.

[coparks2012]

No need to apologize. Your responses are quite helpful, and I believe I am nearly there. I have a remaining question that I think resolve the issue fully. For any future readers, here is how I envision the process:

step 1: read in openmm written NC file, and extract replica_energies and state_indices

output_data = './trajectory/restraint.nc'
reporter    = openmmtools.multistate.MultiStateReporter(output_data, open_mode='r')
states      = reporter.read_thermodynamic_states()[0]
analyzer = ReplicaExchangeAnalyzer(reporter)
replica_energies, _, _, replica_state_indices = analyzer.read_energies()

replica_energies is an [i,j,k] array where i = replica index, j = state index, and k = step index. So, replica_energies[0,1,3] would give the reduced potential energy of replica 0 evaluated at state 1 at step 3.

replica_state_indices gives is an [i,j] array where replica_state_indices[i,j] gives the state index of replica i and step j. For example, if replica 3 is in state 4 at step 3, replica_state_indcies[3,3] = 4

step 2: obtain state vector
I can obtain all of the states with the following code. Since this particular run was calculating the FE of attaching restraints, I will grab the restraint_lambdas:

#===restraint calculation --> need restrant lambdas
lambdas = []
for i in range(0, len(states)):
        print("replica, lambda restraints, lambda sterics, lambda electrostatics:", \
              i, states[i]._composable_states[1].lambda_restraints, \
              states[i]._composable_states[0].lambda_sterics_A, \
              states[i]._composable_states[0].lambda_electrostatics_A)
        lambdas += [states[i]._composable_states[1].lambda_restraints]

step 3: extract u_nk for each lambda i, and create data_list array.
this is the step that I am having a hard time grasping. For each lambda, I need to grab a subset of the replica_energies array and feed it to your extract_unk function. But which subset is not clear to me. Since I am running replica exchange, I may need to access the information in replica_state_indices.

In the code snippet below, I am grabbing a subset that has replica i evaluated at all states at all timesteps. I do not think this is the right subset to grab, but I just put something in to test if the whole code would work end to end.

Could you please help me determine how to parse the replica_energies/replica_state_indices to create the right data_slice for a given lambda i? This would be greatly appreciated.

lambda_list = pd.DataFrame( data = {"lambda_restraint":lambdas})
data_list = []
for i,lam in enumerate(lambdas):
    current_lambda = pd.Series(data={"lambda_restraint": lam})
    # *** what code needs to go here to create the right data slice
    data_slice = replica_energies[i,:,:].copy().transpose()
    # ***
    time    = pd.Series(np.arange(len(data_slice)), name="time")
    df      = extract_u_nk(data_slice, current_lambda, lambda_list, time)
    data_list += [df]

Step4: plot
If I can get the right data_list above, now we can plot!

from alchemlyb.visualisation import plot_convergence
from alchemlyb.convergence import forward_backward_convergence
dff =  forward_backward_convergence(data_list, 'mbar')
ax = plot_convergence(dff)
ax.figure.savefig('dF_t.pdf')

I confirmed that dF_t.pdf was succesfully generated

[xiki-tempula]

step 3: extract u_nk for each lambda i, and create data_list array.
this is the step that I am having a hard time grasping. For each lambda, I need to grab a subset of the replica_energies array and feed it to your extract_unk function. But which subset is not clear to me. Since I am running replica exchange, I may need to access the information in replica_state_indices.

This is the only step that I have figured out actually. Here is my code snippet, where i will be your for i,lam in enumerate(lambdas):

        dataset = nc.Dataset(self._work_dir + "/" + self._nc_file)
        eng = np.array(dataset["energies"])
        eng = eng[:, i, :]
        eng -= eng[:, i].reshape(len(eng), 1)

In this case, you could see that I tried to read the NC file without relying on the openmmtools

I think your approach works for the case where people are running all the windows at the same time but for the case where one is only running a single window at a time, one might still have the problem of grabbing the current lambda state

[coparks2012]

Hmmm, in case anyone else reads this, I suspect our approach is buggy somehow. I am testing the alchemlyb implementation which produced the following plot:

dF_t.pdf

This differs substantially from the plot produced by plotting the openmmtools FE estimate (which is also using mbar) using the same nc file. I am fairly confident in the openmm code, so not sure where the issue is with the openmm/alchemlyb interface.

time, fe, std = [],[],[]
for i in range(1, iterations_to_analyze+1):
    samples_to_analyze = i *10
    analyzer.max_n_iterations = samples_to_analyze
    Delta_f_ij, dDelta_f_ij = analyzer.get_free_energy()
    end_pt_f, end_pt_std = Delta_f_ij[0, nstates - 1], dDelta_f_ij[0, nstates - 1]
    print("relaive FE:", i, iterations_to_analyze+1, end_pt_f, end_pt_std)
    time += [i]
    fe   += [end_pt_f]
    std  += [end_pt_std]

#===forward convergence                                                                                                                                                                                                                            
spacing = int(len(fe)/10)
ax = plot_convergence(fe[0:-1:spacing], fe[0:-1:spacing], forward_error=std[0:-1:spacing], backward_error=std[0:-1:spacing], units = 'kT')

openmmtools_fe.pdf

Please note that since I am unable to calculate the reverse FE using openmmtools, I have just been passing in fe to plot_convergence to produce a plot.

[xiki-tempula]

Thanks for keeping trying. I kind of got the reading thermodynamic state to work previously but when I tried to investigate this problem today. I cannot get it to work. My code is just

from openmmtools.states import GlobalParameterState
class RestraintComposableState(GlobalParameterState):
    lambda_restraints=GlobalParameterState.GlobalParameter('lambda_restraints', standard_value=1.0)
from openmmtools.multistate import MultiStateReporter
reporter = MultiStateReporter('restraint.nc', open_mode='r')
states = reporter.read_thermodynamic_states()[0]

But this would give me

Traceback (most recent call last):
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/utils.py", line 607, in deserialize
    instance.__setstate__(serialization)
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/states.py", line 2973, in __setstate__
    self._composable_states = [utils.deserialize(state)
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/states.py", line 2973, in <listcomp>
    self._composable_states = [utils.deserialize(state)
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/utils.py", line 604, in deserialize
    cls = getattr(module, class_name)
AttributeError: module '__main__' has no attribute 'RestraintComposableState'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/Users/zhiyiwu/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.7571.203/PyCharm.app/Contents/plugins/python/helpers/pydev/pydevconsole.py", line 364, in runcode
    coro = func()
  File "<input>", line 1, in <module>
  File "/Users/zhiyiwu/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.7571.203/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 198, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/Users/zhiyiwu/Library/Application Support/JetBrains/Toolbox/apps/PyCharm-P/ch-0/223.7571.203/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/zhiyiwu/Desktop/openmmtools_alchemlyb_dev-main/rest.py", line 8, in <module>
    states = reporter.read_thermodynamic_states()[0]
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/utils.py", line 87, in _wrapper
    return func(*args, **kwargs)
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/multistate/multistatereporter.py", line 562, in read_thermodynamic_states
    states[state_type].append(deserialize(serialized_state))
  File "/Users/zhiyiwu/miniconda3_x86/envs/alchemlyb_omm/lib/python3.9/site-packages/openmmtools/utils.py", line 609, in deserialize
    raise ValueError('Cannot deserialize class {} without a __setstate__ method'.format(class_name))
ValueError: Cannot deserialize class CompoundThermodynamicState without a __setstate__ method

[xiki-tempula]

I got this code snippet to work though

import matplotlib.pyplot as plt
import numpy as np
from alchemlyb.visualisation import plot_convergence
from alchemlyb.convergence import forward_backward_convergence
import netCDF4 as nc
import pandas as pd

def extract_u_nk(
    array: np.ndarray,
    current_lambda: pd.Series,
    lambda_list: pd.DataFrame,
    time: pd.Series,
    temperature: float = 298.0,
    energy_unit: str = "kT",
) -> pd.DataFrame:
    """Convert the numpy array into a pandas u_nk dataframe for the
    consumption of alchemlyb.

    :param array: The array of reduced potential.
    :param current_lambda: The current lambda Series.
    :param lambda_list: The DataFrame of all lambda.
    :param time: The time Series has "time" as name.
    :return: alchemlyb style u_nk dataframe
    """
    if time.name != "time":
        raise ValueError(f"The name of the time ({time.name}) has to be 'time'.")
    n_samples, n_lambdas = array.shape
    if n_samples != len(time):
        raise ValueError(
            f"The length of the time array ({len(time)}) is not"
            f"the same as the length of the array ({n_samples})."
        )
    if not any((lambda_list == current_lambda).all(axis=1)):
        raise ValueError(
            f"The current_lambda ({current_lambda}) is not in"
            f"the lambda_list ({lambda_list})."
        )

    lambda_dict = lambda_list.to_dict()
    labels = list(lambda_dict.keys())
    if len(labels) == 1:
        columns = [lambda_dict[labels[0]][i] for i in range(len(lambda_list))]
    else:
        columns = [
            tuple([lambda_dict[key][i] for key in labels])
            for i in range(len(lambda_list))
        ]

    u_nk = pd.DataFrame(
        data=array,
        index=time,
        columns=columns,
    )

    for key in labels:
        u_nk[key] = current_lambda[key]

    # set up new multi-index
    newind = ["time"] + list(labels)
    u_nk = u_nk.reset_index().set_index(newind)
    u_nk.name = "u_nk"
    u_nk.attrs = {"temperature": temperature, "energy_unit": energy_unit}
    return u_nk

lambdas = pd.DataFrame(data={'fep-lambda': [0.00000 , 0.00100 , 0.00240 , 0.00560 , 0.01330 , 0.03160 , 0.07500 , 0.17780 , 0.42170 , 1.00000]})

dataset = nc.Dataset('restraint.nc')
eng = np.array(dataset["energies"])
u_nk_list = []
for i, lam in lambdas.iterrows():
    current = eng[:, i, :]
    time = pd.Series(data=np.arange(len(current)),name='time')
    current -= current[:, i].reshape(len(current), 1)
    u_nk_list.append(extract_u_nk(current, current_lambda=lam, lambda_list=lambdas, time=time))
dff = forward_backward_convergence(u_nk_list, 'mbar')
ax = plot_convergence(dff)
plt.show()

Gives

[schuhmc]

Sorry for digging this out, but your code above seems to work for me (as in 'produces plausible results' and the forward average is consistent with the real-time analysis output from openmmtools). Are there any plans to go forward with this?

From what I got from the discussion, the main thing you are missing is the lambda values, but are these even necessary for the use with bar/mbar? Passing something like lambdas = pd.DataFrame(data={'smth arbitraty': range(41)}) seems to work for me as long as the number of states is correct and that you should be able to get from dataset['energies'].shape[1] I think. (Note that my original lambda values here were np.linspace(0,1,41)).

For TI, things are different, but openmmtools does not calculate dU/dL values anyways (at least not directly).

[xiki-tempula]

Sorry for digging this out, but your code above seems to work for me (as in 'produces plausible results' and the forward average is consistent with the real-time analysis output from openmmtools). Are there any plans to go forward with this?

So it depends on whether there is people need it. I made this code because I need it but I don't know whether other people would need it and what are the API that they have in mind. Do you mind open an issue explain where you would use it and how would you hog it up with OpenMM, so we could have a discussion of how to go forward.

From what I got from the discussion, the main thing you are missing is the lambda values, but are these even necessary for the use with bar/mbar? Passing something like lambdas = pd.DataFrame(data={'smth arbitraty': range(41)}) seems to work for me as long as the number of states is correct and that you should be able to get from dataset['energies'].shape[1] I think. (Note that my original lambda values here were np.linspace(0,1,41)).

This only works when one has linear spaced lambda, but you do have a point. I think it would be good to open an issue to show how would you use this function.