#!/usr/bin/python3

# ex_xval.py
#
# Copyright (C) 2019 S. Koenig, A. Ekstroem, et al.
#
# Last updated: Sep 18, 2019
#
# While the code below has been thoroughly tested, no guarantee is given that
# it is free of errors or will be fit for a desired purpose.
#
# Permission to use the code as basis for further scientific work is granted
# under the conditions that:
#
# (a) An attribution to the original code is made in the published work,
#     preferably by citing:
#     S. Koenig et al., "Eigenvector Continuation as an Efficient and Accurate
#     Emulator for Uncertainty Quantification"
#
# (b) The derived or extended program code is made available together with
#     the new published work, under the same or compatible conditions as given
#     here.

import os
import sys
import argparse
import subprocess
import math
import json

import numpy as np
import scipy.linalg as spla

from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression

from scipy.sparse.linalg import eigsh
from scipy.sparse.linalg import LinearOperator

from pyDOE import lhs

import GPy

import matplotlib.pyplot as plt
import matplotlib.colors as clr

import colorsys

################################################################################
# Arguments

parser = argparse.ArgumentParser()

ag = parser.add_argument_group()

ag.add_argument(
  "-d", "--dim",
  type = int, default = 3,
  help = "set dimension of parameter space, min = 1, max = 16"
)

ag.add_argument(
  "-D", "--domain",
  type = float, default = 2.0,
  dest = "dom",
  help = "set sampling interval, [-dom:dom] for each parameter"
)

ag.add_argument(
  "-n", "--points",
  type = int, default = 6,
  help = "set number of emulator training points"
)

ag.add_argument(
  "-X", "--xval-samples",
  type = int, default = 250,
  dest = "N",
  help = "set number of cross-validation samples"
)

ag.add_argument(
  "-c", "--convex",
  action = "store_true", default = False,
  help = "only evaluate convex combinations of emulator training points, " \
  "i.e., enforce interpolation"
)

ag = parser.add_argument_group()

ag.add_argument(
  "-I", "--input",
  default = "./data/He4_Nmax8_hw36.inp",
  dest = "file",
  help = "select target system via input file"
)

ag.add_argument(
  "-O", "--observable",
  choices = ["energy", "radius"], default = "energy",
  help = "select observable"
)

ag = parser.add_argument_group()

ag.add_argument(
  "-p", "--polynomial",
  action = "store_true", default = False,
  help = "enable polynomial interpolation"
)

ag.add_argument(
  "-m", "--poly-order",
  type = int, default = 3,
  help = "set order for polynomial interpolation"
)

ag = parser.add_argument_group()

ag.add_argument(
  "-g", "--gaussian-process",
  action = "store_true", default = False,
  help = "enable Gaussian process emulation"
)

ag.add_argument(
  "--gp-sigma",
  type = float, default = 1.0,
  help = "choose sigma for Gaussian process"
)

ag.add_argument(
  "--gp-kernel",
  choices = ["rbf", "matern32"], default = "rbf",
  help = "choose kernel for Gaussian process"
)

ag = parser.add_argument_group()

ag.add_argument(
  "-e", "--eigenvector-continuation", action = "store_true",
  default = False,
  help = "enable eigenvector continuation"
)

ag.add_argument(
  "-o", "--orthogonalize",
  action = "store_true", default = False,
  help = "orthogonalize eigenvector continuation basis"
)

ag = parser.add_argument_group()

ag.add_argument(
  "-r", "--random-seed",
  type = int, default = -1,
  help = "set fixed random seed"
)

ag.add_argument(
  "-t", "--lanczos-tol",
  type = float, default = 1e-6,
  help = "set tolerance for Lanczos diagonalization"
)

ag.add_argument(
  "-v", "--verbose",
  action = "count", default = 0,
  help = "increase output verbosity"
)

args = parser.parse_args()

if not args.polynomial and not args.gaussian_process:
  args.eigenvector_continuation = True

################################################################################
# Small helper functions

def newline():
  print('>')

def abort(msg):
  print('ERROR: %s' % msg)
  sys.exit()

def expectation_value(bra, op, ket):
  return np.dot(np.dot(bra.transpose(), op), ket)

def load_matrix(name):
  mat = np.load(name)
  return mat

def parse_system(s):
  element, a = s.split('-')
  return int(a), element

def in_range(x, sigma, x0):
  if x0 >= x - sigma and x0 <= x + sigma:
    return True
  else:
    return False

def adjust(x, lo, hi):
  return x * (hi - lo) + lo

################################################################################
# Hamiltonian
#
# H = H0 + H1 + H2 + ... + HN

split_lecs = [
  'c1','c3','c4','Ct_1S0np','Ct_1S0nn','Ct_1S0pp', \
  'Ct_3S1','C_1S0','C_3P0','C_1P1','C_3P1','C_3S1', \
  'C_3S1-3D1','C_3P2','c_D','c_E'
]

nnlo_sat_lecs = [ \
  -1.12152119963259, -3.92500585648682,  3.76568715858592, -0.15982244957832, \
  -0.15915026828018, -0.15814937937011, -0.17767436449900,  2.53936778505038, \
   1.39836559187614,  0.55595876513335, -1.13609526332782,  1.00289267348351, \
   0.60071604833596, -0.80230029533846,  0.81680589148271, -0.03957471270351, \
]

def read_hamiltonian():
  print('--> loading Hamiltonian %s_%s.npy' % (H_name, 'const'))

  global H_const, H_parts
  global obs_mat

  H_const = load_matrix('%s_%s.npy' % (H_name, 'const'))

  H_parts = []

  for idx, lec in enumerate(split_lecs):
    print('--> loading Hamiltonian part %s_%s.npy' % (H_name, lec))
    H_parts.append(load_matrix('%s_%s.npy' % (H_name, lec)))

  H_parts = np.asarray(H_parts)

  if args.observable == "radius":
    print('-> loading <r^2> observable matrix %s.npy' % R_name)
    obs_mat= load_matrix('%s.npy' % (R_name))
  else:
    obs_mat = None

################################################################################
# Basic configuration

# Make sure we can find the data relative to the script:
os.chdir(os.path.dirname(sys.argv[0]))

inp = json.load(open(args.file))
path = os.path.dirname(args.file)

H_name = path + '/' + inp['h_name']
R_name = path + '/' + inp['r_name']

inp['n_body'], inp['element'] = parse_system(inp['system'])

full_dim = len(nnlo_sat_lecs)

################################################################################
# Bookkeeping

# Error:
#########

error = {}

if args.polynomial:
  error['poly'] = 0

if args.gaussian_process:
  error['gp'] = 0

if args.eigenvector_continuation:
  error['evc'] = 0

def add_error(tag, exact, actual):
  error[tag] += (exact - actual) ** 2

# Lanczos counting:
####################

mv_total = 0
n_lanczos = 0

# Summary output:
##################

def print_summary():
  newline()

  print("> Performance summary:")
  print("> --------------------")

  print("> Lanczos average = %d" % (mv_total / n_lanczos))

  print("> Dim(exact) = %d" % H_const.shape[0])

  if args.orthogonalize:
    print(
      "> Dim(evc) = %d (%d)"
        % (args.points, subspace_info['subspace_basis'].shape[1])
    )
  else:
    print("> Dim(evc) = %d" % args.points)

  for k, v in sorted(error.items()):
    print("> Error(%s) = %g" % (k, v))

################################################################################
# Hamiltonian setup and diagonalization

def setup(lecs, const, parts):
  H_mtx = np.copy(const)

  for idx, lec in enumerate(lecs):
    H_mtx += lec * parts[idx]

  return H_mtx

def diagonalize(mat):
  def do_mv(mat, x):
    do_mv.count += 1
    return np.dot(mat, x)

  do_mv.count = 0

  w, v = eigsh(
    LinearOperator(mat.shape, matvec = lambda x : do_mv(mat, x)),
    k = 1, which = 'SA', tol = args.lanczos_tol
  )

  s = np.argsort(w)

  return w[s[0]], v[:,s[0]], do_mv.count

################################################################################
# Generates exact data in the subspace,
# used for both EC and polynomial interpolation
#
# Input: list of args.points points (each a vector args.dim long)
#        in the LEC domain
#
# Output: - a dictionary with exact eigenvectors at each point in the LEC domain
#         - the list of points
#         - exact energy eigenvalues for each point

def generate_subspace_basis(lecs_list, **kwargs):
  newline()
  print('> Generating subspace basis...')

  subspace_basis = []
  parameter_values = []
  eigenvalues = []

  subspace_info = {}

  if 'observable_matrix' in kwargs:
    obsmtx = kwargs.get('observable_matrix')
    observable_values = []
  else:
    obsmtx = None

  # Compute exact eigenvalues and eigenvectors:

  for lecs in lecs_list:
    e, psi, _ = diagonalize(setup(lecs, H_const, H_parts))

    if args.verbose > 0:
      print('--> lecs =', *lecs)
      print('--> lowest eigenvalue = %.6f' % e)

    subspace_basis.append(psi)

    parameter_values.append(lecs[0 : args.dim])
    eigenvalues.append(e)

    if obsmtx is not None:
      observable_values.append(expectation_value(psi, obsmtx, psi))

  subspace_info['subspace_basis']   = np.column_stack(subspace_basis)
  subspace_info['parameter_values'] = parameter_values
  subspace_info['eigenvalues']      = eigenvalues

  if args.eigenvector_continuation:
    vs = np.column_stack(subspace_basis)

    if args.orthogonalize:
      vs = spla.orth(vs)
      n_mat = None
    else:
      n_mat = np.transpose(vs).dot(vs)

    h_const = np.transpose(vs).dot(H_const.dot(vs))
    h_parts = []

    for idx, mat in enumerate(H_parts):
      h_parts.append(np.transpose(vs).dot(mat.dot(vs)))

    if obsmtx is not None:
      obs_mat = np.transpose(vs).dot(obsmtx.dot(vs))
    else:
      obs_mat = None

    subspace_info['subspace_basis'] = vs

    if args.orthogonalize:
      print('-> effective rank = %d  ' % vs.shape[1])
    else:
      wn = np.sort(np.real(spla.eig(n_mat)[0]))

      print('-> norm matrix determinant  = %.6e' % np.linalg.det(n_mat))
      print('-> norm matrix rank         = %d  ' % np.linalg.matrix_rank(n_mat))
      print('-> smallest norm eigenvalue = %.6e' % wn[0])

      if wn[0] < 0:
        abort('Norm matrix is not positive definite!')

    subspace_info['norm_matrix'] = n_mat

    subspace_info['h_const'] = np.asarray(h_const)
    subspace_info['h_parts'] = np.asarray(h_parts)

    if obsmtx is not None:
      subspace_info['obs_mat'] = obs_mat
      subspace_info['observable_values'] = observable_values

  return subspace_info

################################################################################
# Gaussian process

def train_gaussian_process(train_X, train_Y):
  newline()
  print('> Training a Gaussian process...')
  print('-> kernel = %s' % args.gp_kernel)
  print('-> sigma = %g' % args.gp_sigma)

  d = args.dim

  if args.gp_kernel == 'matern32':
    kernel = GPy.kern.Matern32(input_dim = d, variance = 1.0, lengthscale = 1.0)
  if args.gp_kernel == 'rbf':
    kernel = GPy.kern.RBF(input_dim = d, variance = 1.0, lengthscale = 1.0)

  GP = GPy.models.GPRegression(train_X, train_Y, kernel, normalizer = True)

  GP.optimize(messages = False)

  if args.verbose > 0:
    print(GP)

  return GP

################################################################################
# Eigenvector continuation

def eigenvector_continuation(subspace_info, lecs_target):
  vs = subspace_info['subspace_basis']
  N = subspace_info['norm_matrix']

  if N is not None:
    N = N.copy()

  h_const = subspace_info['h_const']
  h_parts = subspace_info['h_parts']

  def run_ec():
    H = setup(lecs_target, h_const, h_parts)

    if N is not None:
      return spla.eigh(H, N)
    else:
      return spla.eigh(H)

  w, v = run_ec()

  s = np.argsort(np.real(w))

  return w[s[0]], v[:, s[0]]

################################################################################
# Compare EC method & polynomial-type interpolation with exact calculation
#
# Input: - set of target LECs
#        - subspace basis for the EC
#        - set of parameters and eigenvalues for the polynomial interpolation
#
# Output: comparison results for energy and observable

def run_comparison(idx, lecs_target, subspace_info, **kwargs):
  if args.verbose > 0:
    print('--> number of training points = %i' % args.points)

  global mv_total, n_lanczos

  # Exact calculation for set of target LECs:
  ############################################

  e, psi, mv_count = diagonalize(setup(lecs_target, H_const, H_parts))

  mv_total += mv_count
  n_lanczos += 1

  exact_en_value = e

  print('-> exact lowest eigenvalue = %.6f' % exact_en_value)

  if args.verbose > 0:
    print("--> Lanczos matrix-vector calls = %d" % mv_count)

  if 'observable_matrix' in kwargs:
    obsmtx = kwargs.get('observable_matrix')
  else:
    obsmtx = None

  def run_exact_obs():
    if obsmtx is not None:
      return expectation_value(psi, obsmtx, psi)
    else:
      return None

  exact_obs_value = run_exact_obs()

  # Polynomial interpolation:
  ############################

  interpolated = [0, 0]

  cs = np.asarray(subspace_info['parameter_values'])
  es = np.asarray(subspace_info['eigenvalues']).reshape(-1, 1)

  if obsmtx is not None:
    os = np.asarray(subspace_info['observable_values']).reshape(-1, 1)

  c0 = np.asarray(lecs_target[0 : args.dim]).reshape(1, -1)

  if args.polynomial:
    if args.verbose > 0:
      print('--> using multivariate polynomial regression')
      print('--> order = %i' % args.poly_order)

    poly = PolynomialFeatures(degree = args.poly_order)
    clf = LinearRegression()

    cs = poly.fit_transform(cs)
    c0 = poly.fit_transform(c0)

    clf.fit(cs, es)

    interpolated_en_value = clf.predict(c0)[0][0]

    if obsmtx is not None:
      clf.fit(cs, os)
      interpolated_obs_value = clf.predict(c0)[0][0]
    else:
      interpolated_obs_value = None

    print('-> interpolated lowest eigenvalue = %.6f' % interpolated_en_value)

    add_error('poly', exact_en_value, interpolated_en_value)

    interpolated = [interpolated_en_value, interpolated_obs_value]

  # Gaussian process:
  ####################

  gp = [0, 0]

  if 'gaussian_process_model' in kwargs:
    # NOTE: The predict function returns [mean, variance].

    gp_en = kwargs.get('gaussian_process_model')[0]

    gp_en_value = gp_en.predict(
      np.asarray(lecs_target[0 : args.dim]).reshape(1, -1)
    )
    gp_en_value = (gp_en_value[0][0][0], gp_en_value[1][0][0])

    if obsmtx is not None:
      gp_obs = kwargs.get('gaussian_process_model')[1]

      gp_obs_value = gp_obs.predict(
        np.asarray(lecs_target[0 : args.dim]).reshape(1, -1)
      )
      gp_obs_value = (gp_obs_value[0][0][0], gp_obs_value[1][0][0])
    else:
      gp_obs_value = (None, None)

    print('-> gp-predicted eigenvalue = %.6f +- %.6f'
      % (gp_en_value[0], args.gp_sigma * np.sqrt(gp_en_value[1]))
    )

    gp = [gp_en_value, gp_obs_value]

    add_error('gp', exact_en_value, gp_en_value[0])

  # Eigenvector continuation:
  ############################

  ec = [0, 0]

  if args.eigenvector_continuation:

    if obsmtx is not None:
      obsmtx_sub = subspace_info['obs_mat']

    def get_ec_obs(obsmtx, v):
      if obsmtx is not None:
        return expectation_value(v, obsmtx, v)
      else:
        return None

    w, v = eigenvector_continuation(subspace_info, lecs_target)

    ec_en_value = (w.real, None, None)

    if obsmtx is not None:
      ec_obs_value = (get_ec_obs(obsmtx_sub, v), None, None)
    else:
      ec_obs_value = (None, None, None)

    print('-> continued lowest eigenvalue = {:.6f}'.format(ec_en_value[0]))

    ec = [ec_en_value, ec_obs_value]

    add_error('evc', exact_en_value, ec_en_value[0])

  # More output:
  ###############

  if obsmtx is not None:
    print('-> exact observable value = %.6f' % exact_obs_value)

    if args.polynomial:
      print('-> interpolated observable value = %.6f' % interpolated_obs_value)

    if 'gaussian_process_model' in kwargs:
      print('-> gp-predicted observable = %.6f +- %.6f'
        % (gp_obs_value[0], args.gp_sigma * np.sqrt(gp_obs_value[1]))
      )

    if args.eigenvector_continuation:
      print('-> continued observable value = %.6f' % ec_obs_value[0])

      if ec_obs_value[1] is not None and args.verbose > 0:
        print(
          '--> uncertainty = +%.6f -%.6f' % (ec_obs_value[1], ec_obs_value[2])
        )

  # Combine and return results:
  ##############################

  exact = [exact_en_value, exact_obs_value]

  return exact, ec, interpolated, gp

################################################################################
# Plotting

def generate_plot():
  nucleus = "$^%d$%s" % (inp['n_body'], inp['element'])

  red = (0.9, 0.2, 0.1)
  green = (0.2, 0.6, 0.3)
  blue = (0.1, 0.4, 0.8)

  if args.observable == 'radius':
    title = '%s ground state radius squared' % nucleus
    prefix = 'EC_r_xval'
    unit = 'fm$^2$'

    x_vals = tr_obs_vals
    y_vals_pf = pf_obs_vals

    if args.gaussian_process:
      y_vals_gp_mu = gp_obs_vals_mu
      y_vals_gp_er = gp_obs_vals_er

    if args.eigenvector_continuation:
      y_vals_ec_mu = ec_obs_vals_mu
      y_vals_ec_er = ec_obs_vals_er

  else:
    title = '%s ground state energy' % nucleus
    prefix = 'EC_E_xval'
    unit = 'MeV'

    x_vals = tr_en_vals
    y_vals_pf = pf_en_vals

    if args.gaussian_process:
      y_vals_gp_mu = gp_en_vals_mu
      y_vals_gp_er = gp_en_vals_er

    if args.eigenvector_continuation:
      y_vals_ec_mu = ec_en_vals_mu
      y_vals_ec_er = ec_en_vals_er

    plt.rcParams["text.latex.preamble"] = [
      r"\usepackage{amsmath}"
    ]

  plt.ylabel('Regression prediction (%s)' % unit)
  plt.xlabel('Exact value (%s)' % unit)

  def plot(x_vals, y_vals, **kwargs):
    base = clr.to_rgba(kwargs['color'])
    trans = (base[0], base[1], base[2], 0.25)

    if 'yerr' in kwargs:
      kwargs['markerfacecolor'] = trans
      plt.errorbar(np.asarray(x_vals), np.asarray(y_vals), **kwargs)
    else:
      kwargs['facecolors'] = trans
      plt.scatter(np.asarray(x_vals), np.asarray(y_vals), **kwargs)

  size = 60
  line = 0.75

  if args.polynomial:
    plot(
      x_vals, y_vals_pf,
      label = 'Polynomial Interpolation',
      color = blue, s = size, marker = "s", linewidths = line
    )

  if args.gaussian_process:
    plot(
      x_vals, y_vals_gp_mu,
      label = 'Gaussian Process',
      yerr = y_vals_gp_er,
      color = green, ms = np.sqrt(size * 1.03), fmt = 'v', linewidth = line
    )

  if args.eigenvector_continuation:
    plot(
      x_vals, y_vals_ec_mu,
      label = 'Eigenvector Continuation',
      color = red, s = size * 1.2, marker = 'o', zorder = 5, linewidths = line
    )

  plt.plot(
    [plt.xlim()[0], plt.xlim()[1]],
    [plt.xlim()[0], plt.xlim()[1]],
    'k-',
    zorder = 10
  )

  setup = '%d dimensions, %d training data' % (args.dim, args.points)
  params = '%s = %g, %s = %d' % ('hw', inp['hw'], 'N_max', inp['n_max'])

  plt.title(
    title + ' [%d train-data, %d dims, seed: %d]'
      % (args.points, args.dim, seed)
  )

  plt.tight_layout()
  plt.show()

################################################################################
# Main program

# Run the actual calculation:
#########################################

print('> Initializing...')

read_hamiltonian()

# Generate LHS points for desired number of dimensions:

if args.random_seed >= 0:
  seed = args.random_seed
else:
  seed = np.random.randint(0, 2 ** 32 - 1)

print("-> random seed = %d" % seed)

np.random.seed(seed)

# Generate points using latin hypercube sampling:

lim_hi = args.dom * np.ones(full_dim)
lim_lo = -args.dom * np.ones(full_dim)

def make_sample():
  cs_list = lhs(args.dim, args.points)

  # Map LHS points to physical domain:

  for j in range(args.points):
    for i in range(args.dim):
      cs_list[j, i] = adjust(cs_list[j, i], lim_lo[i], lim_hi[i])

  # Fill up with the remaining LECs that are held fixed:

  lecs_list = []

  for cs in cs_list:
    lecs = np.copy(nnlo_sat_lecs)
    lecs[0 : args.dim] = cs

    lecs_list.append(lecs)

  return cs_list, lecs_list

cs_list, lecs_list = make_sample()

# Generate target set of random LEC vectors:

if args.convex:
  target_lecs = []

  for j in range(args.N):
    coeffs = np.random.rand(args.points)
    coeffs /= np.sum(coeffs)

    target_lecs.append(np.dot(coeffs, cs_list))

else:
  target_lecs = np.random.rand(args.N, args.dim)

  for j in range(args.N):
    for i in range(args.dim):
      target_lecs[j, i] = adjust(target_lecs[j, i], lim_lo[i], lim_hi[i])

subspace_info = generate_subspace_basis(lecs_list, observable_matrix = obs_mat)

if args.gaussian_process:
  gaussian_process_energy = train_gaussian_process(
    np.asarray(subspace_info['parameter_values']),
    np.asarray(subspace_info['eigenvalues']).reshape(-1, 1)
  )

  if obs_mat is not None:
    gaussian_process_obs = train_gaussian_process(
      np.asarray(subspace_info['parameter_values']),
      np.asarray(subspace_info['observable_values']).reshape(-1, 1)
    )
  else:
    gaussian_process_obs = None

newline()
print("> Running cross validation with %d samples..." % (args.N))

# Initialization:
##################

target_lecs_list = []

for tl in target_lecs:
  lecs = np.copy(nnlo_sat_lecs)
  lecs[0 : args.dim] = tl

  target_lecs_list.append(lecs)

target_lecs_list = np.asarray(target_lecs_list)

tr_en_vals = [] ; tr_obs_vals = []
ec_en_vals = [] ; ec_obs_vals = []
pf_en_vals = [] ; pf_obs_vals = []
gp_en_vals = [] ; gp_obs_vals = []

int_vals = []

ec_error = []

for i in range(args.N):
  newline()
  print('> Xval %d of %d' % (i + 1, args.N))

  if args.gaussian_process:
    tr, ec, pf, gp = run_comparison(
      i, target_lecs_list[i],
      subspace_info,
      observable_matrix = obs_mat,
      gaussian_process_model = [
        gaussian_process_energy, gaussian_process_obs
      ]
    )

    gp_en_vals.append(gp[0])
    gp_obs_vals.append(gp[1])

  else:
    tr, ec, pf, gp = run_comparison(
      i, target_lecs_list[i],
      subspace_info,
      observable_matrix = obs_mat
    )

  tr_en_vals.append(tr[0]) ; tr_obs_vals.append(tr[1])
  pf_en_vals.append(pf[0]) ; pf_obs_vals.append(pf[1])
  ec_en_vals.append(ec[0]) ; ec_obs_vals.append(ec[1])

if args.gaussian_process:
  # Extract the central values and sigma from the gaussian process data:

  gp_en_vals_mu  = [en[0] for en in gp_en_vals]
  gp_en_vals_er  = [args.gp_sigma * np.sqrt(en[1]) for en in gp_en_vals]

  if obs_mat is not None:
    gp_obs_vals_mu = [o[0] for o in gp_obs_vals]
    gp_obs_vals_er = [args.gp_sigma * np.sqrt(o[1]) for o in gp_obs_vals]
  else:
    gp_obs_vals_mu = [None for _ in gp_en_vals]
    gp_obs_vals_er = [None for _ in gp_en_vals]

if args.eigenvector_continuation:
  # Likewise for the eigenvector continuation data:

  ec_en_vals_mu  = [en[0] for en in ec_en_vals]
  ec_en_vals_er  = [[en[1], en[2]] for en in ec_en_vals]

  if obs_mat is not None:
    ec_obs_vals_mu = [o[0] for o in ec_obs_vals]
    ec_obs_vals_er = [[o[1], o[2]] for o in ec_obs_vals]
  else:
    ec_obs_vals_mu = [None for _ in ec_obs_vals]
    ec_obs_vals_er = [[None, None] for _ in ec_obs_vals]

# Print performance summary and generate plot:
#############################
print_summary()
generate_plot()