text: from argparse import ArgumentParser
  import csv
  from functools import partial
  import gzip
  from pathlib import Path
  from typing import Dict, List
  
  from matplotlib import pyplot as plt
  from matplotlib import ticker
  import numpy as np
  from scipy import stats
  from tqdm import tqdm
  
  def extract_smis(library, smiles_col=0, title_line=True) -> List:
      if Path(library).suffix == '.gz':
          open_ = partial(gzip.open, mode='rt')
      else:
          open_ = open
      
      with open_(library) as fid:
          reader = csv.reader(fid)
          if title_line:
              next(reader)
  
          smis = []
          for row in tqdm(reader, desc='Getting smis', leave=False):
              try:
                  smis.append(row[smiles_col])
              except ValueError:
                  continue
  
      return smis
  
  def build_true_dict(true_csv, smiles_col: int = 0, score_col: int = 1,
                      title_line: bool = True,
                      maximize: bool = False) -> Dict[str, float]:
      if Path(true_csv).suffix == '.gz':
          open_ = partial(gzip.open, mode='rt')
      else:
          open_ = open
      
      c = 1 if maximize else -1
  
      with open_(true_csv) as fid:
          reader = csv.reader(fid)
          if title_line:
              next(reader)
  
          d_smi_score = {}
          for row in tqdm(reader, desc='Building dict', leave=False):
              try:
                  d_smi_score[row[smiles_col]] = c * float(row[score_col])
              except ValueError:
                  continue
  
      return d_smi_score
  
  def gather_experiment_predss(experiment) -> List[np.ndarray]:
      chkpts_dir = Path(experiment) / 'chkpts'
  
      chkpt_iter_dirs = sorted(
          chkpts_dir.iterdir(), key=lambda p: int(p.stem.split('_')[-1])
      )[1:]
      try:    # new way
          preds_npzs = [np.load(chkpt_iter_dir / 'preds.npz')
                        for chkpt_iter_dir in chkpt_iter_dirs]
          predss, varss = zip(*[
              (preds_npz['Y_pred'], preds_npz['Y_var'])
              for preds_npz in preds_npzs
          ])
      except FileNotFoundError: # old way
          predss = [np.load(chkpt_iter_dir / 'preds.npy')
                    for chkpt_iter_dir in chkpt_iter_dirs]
  
      return predss
  
  def plot_residualss(Y_true, Y_predss, mask):
      fig, ax = plt.subplots(1, 1, figsize=(6, 4), sharey=True)
  
      for i, Y_pred in enumerate(Y_predss):
          R = np.abs(Y_true - Y_pred)
          ax.hist(
              R, cumulative=True, density=True, bins=np.linspace(0, 10, 20),
              histtype='step', label=f'Iter {i+1}'
          )
  
      ax.xaxis.set_minor_locator(ticker.MultipleLocator(0.5))
      ax.grid(which='both')
      ax.legend()
      ax.set_xlabel('absolute predictive error')
      ax.set_ylabel('density')
      
      return fig
  
  if __name__ == "__main__":
      parser = ArgumentParser()
      parser.add_argument('-e', '--experiments', '--expts', nargs='+',
                          help='the top-level directory generated by the MolPAL run, I.e., the directory with the "data" and "chkpts" directories.')
      parser.add_argument('-l', '--library',
                          help='the library file used for the corresponding MolPAL run.')
      parser.add_argument('--true-csv',
                          help='a pickle file containing a dictionary of the true scoring data')
      parser.add_argument('--smiles-col', type=int, default=0)
      parser.add_argument('--score-col', type=int, default=1)
      parser.add_argument('--no-title-line', action='store_true', default=False)
      parser.add_argument('--maximize', action='store_true', default=False,
                          help='whether the objective for which you are calculating performance should be maximized.')
      parser.add_argument('-n', '--name')
      parser.add_argument('--bins', type=int, default=20)
      
      args = parser.parse_args()
      args.title_line = not args.no_title_line
  
      smis = extract_smis(args.library, args.smiles_col, args.title_line)
      d_smi_score = build_true_dict(
          args.true_csv, args.smiles_col, args.score_col,
          args.title_line, args.maximize
      )
  
      Y_true = np.empty(len(smis))
      mask = np.ones(len(smis), dtype=bool)
      for i, smi in tqdm(enumerate(smis), desc='Building mask', leave=False):
          try:
              Y_true[i] = d_smi_score[smi]
          except KeyError:
              mask[i] = False
  
      for experiment in args.experiments:
          predss = gather_experiment_predss(experiment)
  
          fig1, ax = plt.subplots(1, 1, sharex=True, figsize=(6, 4))
          fig2, axs = plt.subplots(len(predss), 1, figsize=(4, 10),
                                   sharex=True, sharey=True)
  
          for i, preds in enumerate(predss):
              E = (Y_true[mask] - preds[mask])
              mse = (E**2).mean(axis=0)
              pearson, _ = stats.pearsonr(Y_true[mask], preds[mask])
              spearman, _ = stats.spearmanr(Y_true[mask], preds[mask])
  
              print(
                  f'Iteration {i+1}: MSE: {mse:0.3f}, '
                  f'pearson: {pearson:0.3f}, spearman: {spearman:0.3f}',
                  flush=True
              )
  
              R = np.abs(E)
              ax.hist(
                  R, cumulative=True, density=True,
                  bins=np.linspace(0, 2, args.bins),
                  histtype='step', alpha=0.7, label=f'Iter {i+1}'
              )
              axs[i].hist2d(Y_true[mask], R, bins=2*args.bins,
                            range=((5, 12.5), (0, 2.5)),density=True)
              axs[i].set_title(f'Iter {i+1}')
  
          ax.xaxis.set_minor_locator(ticker.MultipleLocator(0.5))
          ax.grid(which='both')
          ax.legend()
          ax.set_xlabel('Absolute predictive error')
          ax.set_ylabel('Density')
  
          ax = fig2.add_subplot(111, frameon=False)
          ax.tick_params(labelcolor='none', which='both', top=False, 
                          bottom=False, left=False, right=False)
          ax.set_xlabel('Docking score')
          ax.set_ylabel('Residual')
  
          fig1.tight_layout()
          fig2.tight_layout()        
          
          fig1.savefig(f'figures/{args.name}_hist.png')
          fig2.savefig(f'figures/{args.name}_2Dhist.png')
  
      exit()