text: #!/usr/bin/env python
  
  # =============================================================================
  # GLOBAL IMPORTS
  # =============================================================================
  
  import os
  
  import numpy as np
  import pandas
  import scipy.stats
  import seaborn as sns
  from matplotlib import pyplot as plt
  
  
  # =============================================================================
  # CONSTANTS
  # =============================================================================
  
  # The column names in the Pandas Dataframe associated to the DG prediction and their uncertainties.
  DG_KEY = '$\Delta$G [kcal/mol]'
  DDG_KEY = 'd$\Delta$G [kcal/mol]'
  
  # Path to the directory containing the Pandas Dataframe.
  DATA_DIR_PATH = 'figure7_data'
  
  # Thermal energy at 300K in kcal/mol.
  KT = 0.5961612775922496
  
  
  # =============================================================================
  # UTILITY FUNCTIONS
  # =============================================================================
  
  def to_numpy_free_energy_trajectory(system_data):
      """Convert the free energy trajectories of the replicate calculations from Dataframe to array format.
  
      Parameters
      ----------
      system_data : pandas.Dataframe
          The data with the replicate free energy trajectories for a system.
  
      Returns
      -------
      free_energy_trajectories : numpy.ndarray
          free_energy_trajectories[replicate_idx][t] is the free energy of the
          replicate_idx-th replicate calculation evaluated after
          n_energy_evaluations[t] energy/force evaluations.
      n_energy_evaluations : numpy.ndarray
          The number of force/energy evaluations for each time point.
  
      """
      # Find all the name of the replicates.
      replicate_ids = sorted(system_data['System ID'].unique())
  
      # Initialize the free energy array.
      n_replicates = len(replicate_ids)
      n_time_points = int(len(system_data) / n_replicates)
      free_energy_trajectories = np.empty(shape=(n_replicates, n_time_points))
  
      # Extract data for each replicate.
      for replicate_idx, replicate_id in enumerate(replicate_ids):
          replicate_data = system_data[system_data['System ID'] == replicate_id]
          free_energy_trajectories[replicate_idx] = replicate_data[DG_KEY].values
  
      # The number of energy evaluations is identical for all replicates.
      n_energy_evaluations = np.array(sorted(system_data['N energy evaluations'].unique()))
      assert len(n_energy_evaluations) == n_time_points
  
      return free_energy_trajectories, n_energy_evaluations
  
  
  def compute_mean_free_energy_trajectory(free_energy_trajectories):
      """Compute mean free energy trajectory and CI from replicate trajectories.
  
      Parameters
      ----------
      free_energy_trajectories : numpy.ndarray
          free_energy_trajectories[replicate_idx][t] is the free energy of the
          replicate_idx-th replicate calculation evaluated at the
          t-th number of energy/force evaluations.
  
      Returns
      -------
      mean_free_energy_trajectory : numpy.ndarray
          free_energy_trajectory[t] is the mean free energy evaluated at the
          t-th number of energy/force evaluations.
      ci_free_energy_trajectory : numpy.ndarray
          ci_free_energy_trajectory[t] is the t-based 95% confidence interval
          so that the final estimate at t is
  
              mean_free_energy_trajectory[t] +- ci_free_energy_trajectory[t]
  
      """
      mean_free_energy_trajectory = free_energy_trajectories.mean(axis=0)
  
      # The t-based confidence interval depends on the number of independent replicates.
      n_replicates = free_energy_trajectories.shape[0]
      t_degrees_of_freedom = n_replicates - 1
      t_statistics = scipy.stats.t.interval(alpha=0.95, df=t_degrees_of_freedom)[1]
      sem = scipy.stats.sem(free_energy_trajectories, axis=0)
  
      return mean_free_energy_trajectory, t_statistics * sem
  
  
  # =============================================================================
  # PLOTTING FUNCTIONS
  # =============================================================================
  
  def plot_mean_free_energy(t, mean_free_energy_trajectory, ci_trajectory, ax,
                            color_mean=None, color_ci=None, **plot_kwargs):
      """Plot mean trajectory with confidence intervals."""
      # Plot mean trajectory confidence intervals.
      upper_bound = mean_free_energy_trajectory + ci_trajectory
      lower_bound = mean_free_energy_trajectory - ci_trajectory
      ax.fill_between(t, upper_bound, lower_bound, alpha=0.15, color=color_ci)
  
      # Plot the mean free energy trajectory.
      ax.plot(t, mean_free_energy_trajectory, color=color_mean, alpha=1.0, **plot_kwargs)
  
  
  def plot_yank_trajectories(data, forward_reverse_data):
      """Plot the mean free energy trajectory and uncertainty of the YANK calculations."""
      sns.set_style('whitegrid')
      sns.set_context('paper')
      system_colors = {'CB8-G3': 'C0', 'OA-G6': 'C2'}
  
      # -------------------- #
      # Plot submission data #
      # -------------------- #
  
      fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(3.8, 3.4))
  
      # for system_name in unique_system_names:
      for system_name, system_data in data.items():
          color = system_colors[system_name]
  
          # Obtain the replicate free energy trajectoreis as numpy arrays.
          free_energy_trajectories, n_energy_evaluations = to_numpy_free_energy_trajectory(system_data)
  
          # Compute the mean and confidence interval.
          mean_free_energy_trajectory, ci_trajectory = compute_mean_free_energy_trajectory(free_energy_trajectories)
  
          # Convert number of energy evaluations into an equivalent number of nanoseconds of MD. Timestep is 2fs.
          t = n_energy_evaluations / 500 / 1000
  
          plot_mean_free_energy(t, mean_free_energy_trajectory, ci_trajectory, ax,
                                color_mean=color, color_ci=color, label=system_name)
  
          # Plot asymptotic free energy.
          asymptotic_DG = mean_free_energy_trajectory[-1]
          print(system_name, t[-1], 'ns')
          ax.plot(t, [asymptotic_DG for _ in t], color='black', ls='-.', alpha=0.8, zorder=1)
  
      # Plot forward/reverse estimates.
      for system_name in forward_reverse_data:
          # First build the overall free energy from complex/solvent data.
          system_fr_data = {k: np.array(v) for k, v in forward_reverse_data[system_name]['solvent'].items()}
          for k in system_fr_data:
              system_fr_data[k] -= np.array(forward_reverse_data[system_name]['complex'][k])
              system_fr_data[k] -= forward_reverse_data[system_name]['standard_state_correction']
          t = np.array(forward_reverse_data[system_name]['n_energy_evaluations']) / 500 / 1000
  
          # Plot forward/reverse free energy trajectories.
          for direction, color in zip(['forward', 'reverse'], ['blue', 'red']):
              # Add only one time the reverse/forward labels for the legend.
              if 'CB8' in system_name:
                  label = direction
              else:
                  label = None
  
              # plt.plot(system_fr_data['simulation_percentage'], system_fr_data[phase_name][direction], color=color)
              ax.plot(
                  t, system_fr_data[direction + '_f'],
                  # yerr=system_fr_data[direction + '_df'],
                  color=color, ls='--', alpha=0.65, label=label
              )
  
      # Set limits.
      ax.set_xlim((0, 2500))
      ax.set_ylim((-13, -6))
      ax.set_xlabel('total simulation time [ns]')
      ax.set_ylabel(DG_KEY)
  
      plt.tight_layout(rect=[0, 0, 1, 0.9])
  
      # Create legend.
      ax.legend(loc='upper left', ncol=4, borderpad=0.5, borderaxespad=0.2,
                bbox_to_anchor=(-0.13, 1.2),
                labelspacing=0.8, handletextpad=0.5, columnspacing=1.2)
  
      # Save figure.
      plt.savefig('free_energy_trajectories.pdf')
      # plt.show()
  
  
  def load_data():
      """Load the pandas dataframes containing the data used for plotting.
  
      This is a subset of the data that was generated for the SAMPLing challenge.
      See the file reference_free_energies.json in:
  
          https://github.com/samplchallenges/SAMPL6/tree/master/host_guest/Analysis/SAMPLing/Data
  
      Returns
      -------
      data : Dict[str, pandas.Dataframe]
          data[system_name] is the dataframe containing the replicate free
          energy trajectory for the system.
      """
      system_names = ['CB8-G3', 'OA-G6']
  
      data = {system_name: None for system_name in system_names}
      for system_name in data:
          json_file_path = os.path.join(DATA_DIR_PATH, system_name+'_pandas_dataframe.json')
          dataframe = pandas.read_json(json_file_path)
  
          # Sort the records by key as the JSON changes the ordering.
          data[system_name] = dataframe.sort_index()
  
      return data
  
  
  # =============================================================================
  # MAIN
  # =============================================================================
  
  if __name__ == '__main__':
      # Read the data.
      data = load_data()
  
      import json
      with open('figure7_data/forward_reverse_data.json', 'r') as f:
          forward_reverse_data = json.load(f)
  
      # Plot YANK calculations on a single graphics.
      plot_yank_trajectories(data, forward_reverse_data)
  