figure-2.py

#!/usr/bin/env python3
import sys
import os
import argparse
import time
import numpy as np
from scipy.interpolate import interp1d

import matplotlib
#matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme()
from matplotlib.path import Path
from matplotlib.patches import PathPatch

import torch
import torch.nn as nn

parser = argparse.ArgumentParser('IKr NN ODE syn. data plot 1')
parser.add_argument('--method', type=str, choices=['dopri5', 'adams'], default='dopri5')
parser.add_argument('--cached', action='store_true')
args = parser.parse_args()

from torchdiffeq import odeint

#device = torch.device('cuda:' + str(args.gpu) if torch.cuda.is_available() else 'cpu')
device = 'cpu'

# Set random seed
np.random.seed(0)
torch.manual_seed(0)

noise_sigma = 0.1

true_y0s = [torch.tensor([[1., 0.]]).to(device),  # what you get after holding at +40mV
            torch.tensor([[0., 1.]]).to(device)]  # (roughly) what you get after holding at -80mV
gt_true_y0s = [torch.tensor([[0., 0., 1., 0., 0., 0.]]).to(device),  # what you get after holding at +40mV
               torch.tensor([[0., 1., 0., 0., 0., 0.]]).to(device)]  # (roughly) what you get after holding at -80mV

# B1.2 in https://physoc.onlinelibrary.wiley.com/action/downloadSupplement?doi=10.1113%2FJP275733&file=tjp12905-sup-0001-textS1.pdf#page=4
e = torch.tensor([-88.4]).to(device)  # assume we know
# Table F11, Cell #5 GKr, in https://physoc.onlinelibrary.wiley.com/action/downloadSupplement?doi=10.1113%2FJP275733&file=tjp12905-sup-0001-textS1.pdf#page=20
#g = torch.tensor([0.1524]).to(device)  # assume we know
g = torch.tensor([1]).to(device)  # assume we know

g_nn = g #* 1.2  # just because we see a-gate gets to ~1.2 at some point (in prt V=50), so can absorb that into the g.
#e -= 5    # just because in pr4, at -90 mV, a-gates became negative, meaning e < -90mV; and only if adding an extra -5mV, a ~ [0, 1].

#
#
#
def makedirs(dirname):
    if not os.path.exists(dirname):
        os.makedirs(dirname)

makedirs('figure-2')


#
# Load data
#
raw_data1 = np.loadtxt('data/pr3-steady-activation-cell-5.csv', delimiter=',', skiprows=1)
time1 = raw_data1[:, 0]
time1_torch = torch.from_numpy(raw_data1[:, 0]).to(device)
#current1 = raw_data1[:, 1]
voltage1 = raw_data1[:, 2]

#
# Make filters
#
n_ms = 3
dt = 0.1  # ms
n_points = int(n_ms / dt)
change_pt1 = np.append([True], ~(voltage1[1:] != voltage1[:-1]))
cap_mask1 = np.copy(change_pt1)
for i in range(n_points):
    cap_mask1 = cap_mask1 & np.roll(change_pt1, i + 1)
# A bigger/final filter mask
extra_points = 20  # for numerical derivative or smoothing issue
mask1 = np.copy(cap_mask1)
for i in range(extra_points):
    mask1 = mask1 & np.roll(change_pt1, i + n_points + 1)
    mask1 = mask1 & np.roll(change_pt1, -i - 1)



#
#
#
class GroundTruth(nn.Module):
    def __init__(self):
        super(GroundTruth, self).__init__()

        # Best of 10 fits for data herg25oc1 cell B06 (seed 542811797)
        self.p1 = 5.94625498751561316e-02 * 1e-3
        self.p2 = 1.21417701632850410e+02 * 1e-3
        self.p3 = 4.76436985414236425e+00 * 1e-3
        self.p4 = 3.49383233960778904e-03 * 1e-3
        self.p5 = 9.62243079990877703e+01 * 1e-3
        self.p6 = 2.26404683824047979e+01 * 1e-3
        self.p7 = 8.00924780462999131e+00 * 1e-3
        self.p8 = 2.43749808069009823e+01 * 1e-3
        self.p9 = 2.06822607368134157e+02 * 1e-3
        self.p10 = 3.30791433507312362e+01 * 1e-3
        self.p11 = 1.26069071928587784e+00 * 1e-3
        self.p12 = 2.24844970727316245e+01 * 1e-3

    def set_fixed_form_voltage_protocol(self, t, v):
        # Regular time point voltage protocol time series
        self._t_regular = t
        self._v_regular = v
        self.__v = interp1d(t, v)

    def _v(self, t):
        return torch.from_numpy(self.__v([t.cpu().numpy()])).to(device)

    def voltage(self, t):
        # Return voltage
        return self._v(t).numpy()

    def forward(self, t, y):
        c1, c2, i, ic1, ic2, o = torch.unbind(y[0])

        try:
            v = self._v(t).to(device)
        except ValueError:
            v = torch.tensor([-80]).to(device)

        a1 = self.p1 * torch.exp(self.p2 * v)
        b1 = self.p3 * torch.exp(-self.p4 * v)
        bh = self.p5 * torch.exp(self.p6 * v)
        ah = self.p7 * torch.exp(-self.p8 * v)
        a2 = self.p9 * torch.exp(self.p10 * v)
        b2 = self.p11 * torch.exp(-self.p12 * v)

        dc1dt = a1 * c2 + ah * ic1 + b2 * o - (b1 + bh + a2) * c1
        dc2dt = b1 * c1 + ah * ic2 - (a1 + bh) * c2
        didt = a2 * ic1 + bh * o - (b2 + ah) * i
        dic1dt = a1 * ic2 + bh * c1 + b2 * i - (b1 + ah + a2) * ic1
        dic2dt = b1 * ic1 + bh * c2 - (ah + a1) * ic2
        dodt = a2 * c1 + ah * i - (b2 + bh) * o

        return torch.stack([dc1dt[0], dc2dt[0], didt[0], dic1dt[0], dic2dt[0], dodt[0]])



#
#
#
class Lambda(nn.Module):
    def __init__(self):
        super(Lambda, self).__init__()

        # Fit to GroundTruth model using train-d0.py, in `./d0/model-parameters.txt`.
        self.p1 = 5.694588454735844622e-05
        self.p2 = 1.172955815858964107e-01
        self.p3 = 3.522672347205991382e-05
        self.p4 = 4.972513487995382231e-02
        # Best of 10 fits (M10) for data herg25oc1 cell B06 (seed 542811797) - assume correct.
        self.p5 = 9.62243079990877703e+01 * 1e-3
        self.p6 = 2.26404683824047979e+01 * 1e-3
        self.p7 = 8.00924780462999131e+00 * 1e-3
        self.p8 = 2.43749808069009823e+01 * 1e-3

        self.unity = torch.tensor([1]).to(device)

    def set_fixed_form_voltage_protocol(self, t, v):
        # Regular time point voltage protocol time series
        self._t_regular = t
        self._v_regular = v
        self.__v = interp1d(t, v)

    def _v(self, t):
        return torch.from_numpy(self.__v([t.cpu().numpy()])).to(device)

    def voltage(self, t):
        # Return voltage
        return self._v(t).numpy()

    def forward(self, t, y):
        a, r = torch.unbind(y[0])

        try:
            v = self._v(t).to(device)
        except ValueError:
            v = torch.tensor([-80]).to(device)

        k1 = self.p1 * torch.exp(self.p2 * v)
        k2 = self.p3 * torch.exp(-self.p4 * v)
        k3 = self.p5 * torch.exp(self.p6 * v)
        k4 = self.p7 * torch.exp(-self.p8 * v)

        dadt = k1 * (self.unity - a) - k2 * a
        drdt = -k3 * r + k4 * (self.unity - r)

        return torch.stack([dadt[0], drdt[0]])


class ODEFunc1_6(nn.Module):

    def __init__(self):
        super(ODEFunc1_6, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(2, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 1),
        )

        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0, std=0.1)
                nn.init.constant_(m.bias, val=0)

        self.vrange = torch.tensor([100.]).to(device)
        self.netscale = torch.tensor([1000.]).to(device)

        # Best of 10 fits (M10) for data herg25oc1 cell B06 (seed 542811797)
        self.p5 = 9.62243079990877703e+01 * 1e-3
        self.p6 = 2.26404683824047979e+01 * 1e-3
        self.p7 = 8.00924780462999131e+00 * 1e-3
        self.p8 = 2.43749808069009823e+01 * 1e-3

        self.unity = torch.tensor([1]).to(device)

    def set_fixed_form_voltage_protocol(self, t, v):
        # Regular time point voltage protocol time series
        self._t_regular = t
        self._v_regular = v
        self.__v = interp1d(t, v)

    def _v(self, t):
        #return torch.from_numpy(np.interp([t.cpu().detach().numpy()], self._t_regular,
        #                                  self._v_regular))
        return torch.from_numpy(self.__v([t.cpu().detach().numpy()]))

    def voltage(self, t):
        # Return voltage
        return self._v(t).numpy()

    def forward(self, t, y):
        a, r = torch.unbind(y, dim=1)

        try:
            v = self._v(t).to(device)
        except ValueError:
            v = torch.tensor([-80]).to(device)
        nv = v / self.vrange

        k3 = self.p5 * torch.exp(self.p6 * v)
        k4 = self.p7 * torch.exp(-self.p8 * v)

        drdt = -k3 * r + k4 * (self.unity - r)

        dadt = self.net(torch.stack([nv[0], a[0]]).float()) / self.netscale

        return torch.stack([dadt[0], drdt[0]]).reshape(1, -1)


class ODEFunc1_6_2(nn.Module):

    def __init__(self):
        super(ODEFunc1_6_2, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(2, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 200),
            nn.LeakyReLU(),
            nn.Linear(200, 1),
        )

        for m in self.net.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0, std=1e-3)
                nn.init.constant_(m.bias, val=0)

        self.vrange = torch.tensor([100.]).to(device)
        self.netscale = torch.tensor([1000.]).to(device)

        # https://github.com/CardiacModelling/hERGRapidCharacterisation/blob/master/room-temperature-only/out/herg25oc1/herg25oc1-staircaseramp-B06-solution-542811797.txt
        self.p1 = 1.12592345582957387e-01 * 1e-3
        self.p2 = 8.26751134920666146e+01 * 1e-3
        self.p3 = 3.38768033864048357e-02 * 1e-3
        self.p4 = 4.67106147665183542e+01 * 1e-3
        # 
        # Best of 10 fits (M10) for data herg25oc1 cell B06 (seed 542811797)
        self.p5 = 9.62243079990877703e+01 * 1e-3
        self.p6 = 2.26404683824047979e+01 * 1e-3
        self.p7 = 8.00924780462999131e+00 * 1e-3
        self.p8 = 2.43749808069009823e+01 * 1e-3

        self.unity = torch.tensor([1]).to(device)

    def set_fixed_form_voltage_protocol(self, t, v):
        # Regular time point voltage protocol time series
        self._t_regular = t
        self._v_regular = v
        self.__v = interp1d(t, v)

    def _v(self, t):
        #return torch.from_numpy(np.interp([t.cpu().detach().numpy()], self._t_regular,
        #                                  self._v_regular))
        return torch.from_numpy(self.__v([t.cpu().detach().numpy()]))

    def voltage(self, t):
        # Return voltage
        return self._v(t).numpy()

    def _dadt(self, a, v):
        k1 = self.p1 * torch.exp(self.p2 * v)
        k2 = self.p3 * torch.exp(-self.p4 * v)
        return k1 * (self.unity - a) - k2 * a

    def _drdt(self, r, v):
        k3 = self.p5 * torch.exp(self.p6 * v)
        k4 = self.p7 * torch.exp(-self.p8 * v)
        return -k3 * r + k4 * (self.unity - r)

    def forward(self, t, y):
        a, r = torch.unbind(y, dim=1)

        try:
            v = self._v(t).to(device)
        except ValueError:
            v = torch.tensor([-80]).to(device)
        nv = v / self.vrange

        drdt = self._drdt(r, v)

        dadt = self._dadt(a, v).reshape(-1)
        ddadt = self.net(torch.stack([nv[0], a[0]]).float()) / self.netscale
        dadt += ddadt.reshape(-1)

        return torch.stack([dadt[0], drdt[0]]).reshape(1, -1)
#
#
#

#
#
#
true_model = GroundTruth().to(device)

func_o = Lambda().to(device)
func_o.eval()

func_1 = ODEFunc1_6().to(device)
func_1.load_state_dict(torch.load('d1/model-state-dict.pt'))
func_1.eval()

func_2 = ODEFunc1_6_2().to(device)
func_2.load_state_dict(torch.load('d2/model-state-dict.pt', map_location=torch.device('cpu')))
func_2.eval()

prediction3 = np.loadtxt('data/pr5-deactivation-cell-5.csv', delimiter=',', skiprows=1)
timep3 = prediction3[:, 0]
timep3_torch = torch.from_numpy(prediction3[:, 0]).to(device)
#currentp3 = prediction3[:, 1]
voltagep3 = prediction3[:, 2]

true_gy0 = gt_true_y0s[1]  # (roughly holding at -80mV)
true_y0 = true_y0s[1]  # (roughly holding at -80mV)

def sim_data(func, time, voltage, time_torch, gg, y0, noise=None):
    func.set_fixed_form_voltage_protocol(time, voltage)
    with torch.no_grad():
        pred_y = odeint(func, y0, time_torch).to(device)
    sim = gg * pred_y[:, 0, -1] * (func._v(time_torch).to(device) - e)
    if noise:
        sim += torch.from_numpy(np.random.normal(0, noise, sim.cpu().numpy().shape)).to(device)
    return sim.reshape(-1).cpu().numpy()

def predict(func, time, voltage, time_torch, data, gg, y0, name):
    func.set_fixed_form_voltage_protocol(time, voltage)
    with torch.no_grad():
        pred_y = odeint(func, y0, time_torch).to(device)
    pred_yo = gg * pred_y[:, 0, 0] * pred_y[:, 0, 1] * (func._v(time_torch).to(device) - e)
    loss = torch.mean(torch.abs(pred_yo - torch.from_numpy(data).to(device)))
    print('{:s} prediction | Total Loss {:.6f}'.format(name, loss.item()))
    return pred_yo


if args.cached:
    current1 = torch.load('figure-2/yc-pr3.pt')
    pred_y_o_pr3 = torch.load('figure-2/yo-pr3.pt')
    pred_y_1_pr3 = torch.load('figure-2/y1-pr3.pt')
    pred_y_2_pr3 = torch.load('figure-2/y2-pr3.pt')
    currentp3 = torch.load('figure-2/yc-pr5.pt')
    pred_y_o_pr5 = torch.load('figure-2/yo-pr5.pt')
    pred_y_1_pr5 = torch.load('figure-2/y1-pr5.pt')
    pred_y_2_pr5 = torch.load('figure-2/y2-pr5.pt')
else:
    with torch.no_grad():
        ###
        ### Training protocols
        ###

        #
        # Pr3
        #
        # True model
        current1 = sim_data(true_model, time1, voltage1, time1_torch, g, true_gy0, noise_sigma)
        # Trained Neural ODE
        makedirs('figure-2/pr3')
        pred_y_o = predict(func_o, time1, voltage1, time1_torch, current1, g, true_y0, 'Pr3 (Mo)')
        pred_y_1 = predict(func_1, time1, voltage1, time1_torch, current1, g_nn, true_y0, 'Pr3 (M1)')
        pred_y_2 = predict(func_2, time1, voltage1, time1_torch, current1, g_nn, true_y0, 'Pr3 (M2)')

        l = int(len(time1) / 7)  # 7 steps

        fig1, ax1 = plt.subplots(1, 1, figsize=(6, 4))
        ax1.set_xlabel('Time (ms)')
        ax1.set_ylabel('Current (nA)')
        for i in range(7):
            ax1.plot(time1[:l], current1[l*i:l*(i+1)], c='#7f7f7f', label='__nolegend__' if i else 'Data')
            ax1.plot(time1[:l], pred_y_o.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C0', label='__nolegend__' if i else 'Original')
            ax1.plot(time1[:l], pred_y_1.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C1', label='__nolegend__' if i else 'Full NN')
            ax1.plot(time1[:l], pred_y_2.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '-.', c='C2', label='__nolegend__' if i else 'NN discrepancy')

            fig2, ax2 = plt.subplots(1, 1, figsize=(6, 4))
            ax2.set_xlabel('Time (ms)')
            ax2.set_ylabel('Current (nA)')
            ax2.plot(time1[:l], current1[l*i:l*(i+1)], c='#7f7f7f', label='__nolegend__' if i else 'Data')
            ax2.plot(time1[:l], pred_y_o.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C0', label='Original')
            ax2.plot(time1[:l], pred_y_1.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C1', label='Full NN')
            ax2.plot(time1[:l], pred_y_2.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '-.', c='C2', label='NN discrepancy')
            ax2.set_xlim(time1[:l].min(), time1[:l].max())
            #ax2.set_ylim(-4, 1.9)
            ax2.legend()
            fig2.tight_layout()
            fig2.savefig('figure-2/pr3/s%s' % i, dpi=200)
            plt.close(fig2)
        ax1.set_xlim(time1[:l].min(), time1[:l].max())
        #ax1.set_ylim(-4, 1.9)
        ax1.legend()
        fig1.tight_layout()
        fig1.savefig('figure-2/pr3', dpi=200)
        # do another one with zooms
        ax1.set_xlim(5000, 7000)
        #ax1.set_ylim(-2, 1.7)
        fig1.tight_layout()
        fig1.savefig('figure-2/pr3-z', dpi=200)
        #plt.show()
        plt.close(fig1)

        # Cache it
        torch.save(current1, 'figure-2/yc-pr3.pt')
        torch.save(pred_y_o, 'figure-2/yo-pr3.pt')
        torch.save(pred_y_1, 'figure-2/y1-pr3.pt')
        torch.save(pred_y_2, 'figure-2/y2-pr3.pt')
        pred_y_o_pr3 = pred_y_o
        pred_y_1_pr3 = pred_y_1
        pred_y_2_pr3 = pred_y_2
        del(pred_y_o, pred_y_1, pred_y_2)

        #
        # Pr5
        #
        makedirs('figure-2/pr5')
        # True model
        currentp3 = sim_data(true_model, timep3, voltagep3, timep3_torch, g, true_gy0, noise_sigma)
        # Trained Neural ODE
        pred_y_o = predict(func_o, timep3, voltagep3, timep3_torch, currentp3, g, true_y0, 'Pr5 (Mo)')
        pred_y_1 = predict(func_1, timep3, voltagep3, timep3_torch, currentp3, g_nn, true_y0, 'Pr5 (M1)')
        pred_y_2 = predict(func_2, timep3, voltagep3, timep3_torch, currentp3, g_nn, true_y0, 'Pr5 (M2)')

        l = int(len(timep3) / 9)  # 9 steps

        fig1, ax1 = plt.subplots(1, 1, figsize=(6, 4))
        ax1.set_xlabel('Time (ms)')
        ax1.set_ylabel('Current (nA)')
        for i in range(9):
            ax1.plot(timep3[:l], currentp3[l*i:l*(i+1)], c='#7f7f7f', label='__nolegend__' if i else 'Data')
            ax1.plot(timep3[:l], pred_y_o.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C0', label='__nolegend__' if i else 'Original')
            ax1.plot(timep3[:l], pred_y_1.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C1', label='__nolegend__' if i else 'Full NN')
            ax1.plot(timep3[:l], pred_y_2.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '-.', c='C2', label='__nolegend__' if i else 'NN discrepancy')

            fig2, ax2 = plt.subplots(1, 1, figsize=(6, 4))
            ax2.set_xlabel('Time (ms)')
            ax2.set_ylabel('Current (nA)')
            ax2.plot(timep3[:l], currentp3[l*i:l*(i+1)], c='#7f7f7f', label='__nolegend__' if i else 'Data')
            ax2.plot(timep3[:l], pred_y_o.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C0', label='Original')
            ax2.plot(timep3[:l], pred_y_1.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '--', c='C1', label='Full NN')
            ax2.plot(timep3[:l], pred_y_2.reshape(-1).cpu().numpy()[l*i:l*(i+1)], '-.', c='C2', label='NN discrepancy')
            ax2.set_xlim(timep3[:l].min(), timep3[:l].max())
            #ax2.set_ylim(-3, 7.5)
            ax2.legend()
            fig2.tight_layout()
            fig2.savefig('figure-2/pr5/s%s' % i, dpi=200)
            plt.close(fig2)
        ax1.set_xlim(timep3[:l].min(), timep3[:l].max())
        #ax1.set_ylim(-3, 7.5)
        ax1.legend()
        fig1.tight_layout()
        fig1.savefig('figure-2/pr5', dpi=300)
        # do another one with zooms
        #ax1.set_xlim(1175, 1475)
        #ax1.set_ylim(-2.5, 7)
        #fig1.tight_layout()
        #fig1.savefig('figure-2/pr5-z', dpi=300)
        #plt.show()
        plt.close(fig1)

        # Cache it
        torch.save(currentp3, 'figure-2/yc-pr5.pt')
        torch.save(pred_y_o, 'figure-2/yo-pr5.pt')
        torch.save(pred_y_1, 'figure-2/y1-pr5.pt')
        torch.save(pred_y_2, 'figure-2/y2-pr5.pt')
        pred_y_o_pr5 = pred_y_o
        pred_y_1_pr5 = pred_y_1
        pred_y_2_pr5 = pred_y_2
        del(pred_y_o, pred_y_1, pred_y_2)


#
# Settings
#
zoom_in_win = {
    0: [(1000, 5000), (6600, 7100)],  # pr3
    1: [(2600, 3000), (8650, 9100)],  # pr5
}
zoom_in_y = {
    0: [(-1, 4), (-19, 1)],  # pr3
    1: [(-24, 18), (-6, 1)],  # pr5
}
facecolors = [
    [sns.color_palette("Set2")[0], sns.color_palette("Set2")[1]],  # pr3
    [sns.color_palette("Set2")[2], sns.color_palette("Set2")[3]],  # pr5
]


#
# Plot
#
ds = 20
fig = plt.figure(figsize=(11, 5))
n_maxzoom = 2
grid = plt.GridSpec(4 + 1 + 12 + 5 + 9, 4,
                    hspace=0.0, wspace=0.0)
axes = np.empty([3, 2], dtype=object)
for i in range(2):
    i_grid = i * (2 + 0)
    f_grid = (i + 1) * 2 + i * 0

    axes[0, i] = fig.add_subplot(grid[:4, i_grid:f_grid])
    axes[0, i].set_xticklabels([])
    axes[1, i] = fig.add_subplot(grid[5:17, i_grid:f_grid])
    axes[2, i] = np.empty(n_maxzoom, dtype=object)

    n_zoom = len(zoom_in_win[i])
    for ii in range(n_zoom):
        axes[2, i][ii] = fig.add_subplot(
                grid[-9:, i_grid+ii*1:i_grid+(ii+1)*1])
        axes[2, i][ii].set_xticklabels([])
        axes[2, i][ii].set_xticks([])
        axes[2, i][ii].set_yticklabels([])
        axes[2, i][ii].set_yticks([])

# Set labels
axes[0, 0].set_ylabel('Voltage\n(mV)', fontsize=12)
axes[1, 0].set_ylabel('Current\n(nA)', fontsize=12)
axes[2, 0][0].set_ylabel('Zoom in', fontsize=12)
axes[1, 0].set_xlabel('Time (ms)', fontsize=12)
axes[1, 1].set_xlabel('Time (ms)', fontsize=12)

# Plot!
l = int(len(time1) / 7)  # 7 steps
y_mins = [np.inf] * len(zoom_in_win[0])
y_maxs = [-np.inf] * len(zoom_in_win[0])
for i in range(7):
    axes[0, 0].plot(time1[:l], voltage1[l*i:l*(i+1)], c='#7f7f7f', ds='steps')

    axes[1, 0].plot(time1[:l:ds], current1.reshape(-1)[l*i:l*(i+1):ds], c='#7f7f7f', label='__nolegend__' if i else 'Data')
    axes[1, 0].plot(time1[:l:ds], pred_y_o_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '--', c='C0', label='__nolegend__' if i else 'Original')
    axes[1, 0].plot(time1[:l:ds], pred_y_1_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '--', c='C1', label='__nolegend__' if i else r'$a$-gate as NN (NN-f)')
    axes[1, 0].plot(time1[:l:ds], pred_y_2_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '-.', c='C2', label='__nolegend__' if i else r'NN as discrepancy term (NN-d)')

    axes[0, 0].set_xlim([time1[:l:ds][0], time1[:l:ds][-1]])
    axes[1, 0].set_xlim([time1[:l:ds][0], time1[:l:ds][-1]])

    # Zooms
    for i_z, (t_i, t_f) in enumerate(zoom_in_win[0]):
        # Find closest time
        idx_i = np.argmin(np.abs(time1[:l:ds] - t_i))
        idx_f = np.argmin(np.abs(time1[:l:ds] - t_f))
        # Data
        t = time1[:l:ds][idx_i:idx_f]
        c = current1.reshape(-1)[l*i:l*(i+1):ds][idx_i:idx_f]
        y0 = pred_y_o_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        y1 = pred_y_1_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        y2 = pred_y_2_pr3.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        # Work out the max and min
        y_mins[i_z] = min(np.min(c), y_mins[i_z])
        y_maxs[i_z] = max(np.max(c), y_maxs[i_z])
        # Work out third panel plot
        axes[2, 0][i_z].plot(t, c, c='#7f7f7f')
        axes[2, 0][i_z].plot(t, y0, '--', c='C0')
        axes[2, 0][i_z].plot(t, y1, '--', c='C1')
        axes[2, 0][i_z].plot(t, y2, '-.', c='C2')

        if i == 6:
            axes[2, 0][i_z].set_xlim([t[0], t[-1]])
            # Re-adjust the max and min
            if False:
                y_min, y_max = y_mins[i_z], y_maxs[i_z]
                y_amp = y_max - y_min
                y_min -=  0.2 * y_amp
                y_max +=  0.2 * y_amp
                y_amp = y_max - y_min
            else:
                y_min, y_max = zoom_in_y[0][i_z]
            axes[2, 0][i_z].set_ylim([y_min, y_max])
            # And plot shading over second panels
            codes = [Path.MOVETO] + [Path.LINETO] * 3 + [Path.CLOSEPOLY]
            vertices = np.array([(t[0], y_min),
                                 (t[0], y_max),
                                 (t[-1], y_max),
                                 (t[-1], y_min),
                                 (0, 0)], float)
            pathpatch = PathPatch(Path(vertices, codes),
                    facecolor=facecolors[0][i_z],
                    edgecolor=facecolors[0][i_z],
                    #edgecolor=None,
                    alpha=0.25)
            plt.sca(axes[1, 0])
            pyplot_axes = plt.gca()
            pyplot_axes.add_patch(pathpatch)
            # Set background color to match shading color
            vertices = np.array([(t[0], y_min),
                                 (t[0], y_max),
                                 (t[-1], y_max),
                                 (t[-1], y_min),
                                 (t[0], y_min)], float)
            pathpatch = PathPatch(Path(vertices, codes),
                    facecolor=facecolors[0][i_z],
                    #edgecolor=facecolors[0][i_z],
                    edgecolor=None,
                    alpha=0.25)
            plt.sca(axes[2, 0][i_z])
            pyplot_axes = plt.gca()
            pyplot_axes.add_patch(pathpatch)
            # Set arrow and time duration
            axes[2, 0][i_z].arrow(1, -0.05, -1, 0,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 0][i_z].transAxes)
            axes[2, 0][i_z].arrow(0, -0.05, 1, 0,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 0][i_z].transAxes)
            axes[2, 0][i_z].text(0.5, -0.15,
                    '%s ms' % np.around(t_f - t_i, decimals=0),
                    transform=axes[2, 0][i_z].transAxes,
                    horizontalalignment='center',
                    verticalalignment='center')
            # Set arrow and current range
            axes[2, 0][i_z].arrow(-0.05, 1, 0, -1,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 0][i_z].transAxes)
            axes[2, 0][i_z].arrow(-0.05, 0, 0, 1,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 0][i_z].transAxes)
            axes[2, 0][i_z].text(-0.15, 0.5,
                    '%s nA' % np.around(y_max - y_min, decimals=0),
                    rotation=90,
                    transform=axes[2, 0][i_z].transAxes,
                    horizontalalignment='center',
                    verticalalignment='center')

l = int(len(timep3) / 9)  # 9 steps
y_mins = [np.inf] * len(zoom_in_win[1])
y_maxs = [-np.inf] * len(zoom_in_win[1])
for i in range(9):
    axes[0, 1].plot(timep3[:l], voltagep3[l*i:l*(i+1)], c='#7f7f7f', ds='steps')

    axes[1, 1].plot(timep3[:l:ds], currentp3.reshape(-1)[l*i:l*(i+1):ds], c='#7f7f7f', label='__nolegend__' if i else 'Data')
    axes[1, 1].plot(timep3[:l:ds], pred_y_o_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '--', c='C0', label='__nolegend__' if i else 'Original')
    axes[1, 1].plot(timep3[:l:ds], pred_y_1_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '--', c='C1', label='__nolegend__' if i else r'$a$-gate as NN (NN-f)')
    axes[1, 1].plot(timep3[:l:ds], pred_y_2_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds], '-.', c='C2', label='__nolegend__' if i else 'NN as discrepancy term (NN-d)')

    axes[0, 1].set_xlim([timep3[:l:ds][0], timep3[:l:ds][-1]])
    axes[1, 1].set_xlim([timep3[:l:ds][0], timep3[:l:ds][-1]])

    # Zooms
    for i_z, (t_i, t_f) in enumerate(zoom_in_win[1]):
        # Find closest time
        idx_i = np.argmin(np.abs(timep3[:l:ds] - t_i))
        idx_f = np.argmin(np.abs(timep3[:l:ds] - t_f))
        # Data
        t = timep3[:l:ds][idx_i:idx_f]
        c = currentp3.reshape(-1)[l*i:l*(i+1):ds][idx_i:idx_f]
        y0 = pred_y_o_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        y1 = pred_y_1_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        y2 = pred_y_2_pr5.reshape(-1).cpu().numpy()[l*i:l*(i+1):ds][idx_i:idx_f]
        # Work out the max and min
        y_mins[i_z] = min(np.min(c), y_mins[i_z])
        y_maxs[i_z] = max(np.max(c), y_maxs[i_z])
        # Work out third panel plot
        axes[2, 1][i_z].plot(t, c, c='#7f7f7f')
        axes[2, 1][i_z].plot(t, y0, '--', c='C0')
        axes[2, 1][i_z].plot(t, y1, '--', c='C1')
        axes[2, 1][i_z].plot(t, y2, '-.', c='C2')

        if i == 6:
            axes[2, 1][i_z].set_xlim([t[0], t[-1]])
            # Re-adjust the max and min
            if False:
                y_min, y_max = y_mins[i_z], y_maxs[i_z]
                y_amp = y_max - y_min
                y_min -=  0.2 * y_amp
                y_max +=  0.2 * y_amp
                y_amp = y_max - y_min
            else:
                y_min, y_max = zoom_in_y[1][i_z]
            axes[2, 1][i_z].set_ylim([y_min, y_max])
            # And plot shading over second panels
            codes = [Path.MOVETO] + [Path.LINETO] * 3 + [Path.CLOSEPOLY]
            vertices = np.array([(t[0], y_min),
                                 (t[0], y_max),
                                 (t[-1], y_max),
                                 (t[-1], y_min),
                                 (0, 0)], float)
            pathpatch = PathPatch(Path(vertices, codes),
                    facecolor=facecolors[1][i_z],
                    edgecolor=facecolors[1][i_z],
                    #edgecolor=None,
                    alpha=0.25)
            plt.sca(axes[1, 1])
            pyplot_axes = plt.gca()
            pyplot_axes.add_patch(pathpatch)
            # Set background color to match shading color
            vertices = np.array([(t[0], y_min),
                                 (t[0], y_max),
                                 (t[-1], y_max),
                                 (t[-1], y_min),
                                 (t[0], y_min)], float)
            pathpatch = PathPatch(Path(vertices, codes),
                    facecolor=facecolors[1][i_z],
                    #edgecolor=facecolors[1][i_z],
                    edgecolor=None,
                    alpha=0.25)
            plt.sca(axes[2, 1][i_z])
            pyplot_axes = plt.gca()
            pyplot_axes.add_patch(pathpatch)
            # Set arrow and time duration
            axes[2, 1][i_z].arrow(1, -0.05, -1, 0,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 1][i_z].transAxes)
            axes[2, 1][i_z].arrow(0, -0.05, 1, 0,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 1][i_z].transAxes)
            axes[2, 1][i_z].text(0.5, -0.15,
                    '%s ms' % np.around(t_f - t_i, decimals=0),
                    transform=axes[2, 1][i_z].transAxes,
                    horizontalalignment='center',
                    verticalalignment='center')
            # Set arrow and current range
            axes[2, 1][i_z].arrow(-0.05, 1, 0, -1,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 1][i_z].transAxes)
            axes[2, 1][i_z].arrow(-0.05, 0, 0, 1,
                                  length_includes_head=True,
                                  head_width=0.03, head_length=0.05,
                                  clip_on=False, fc='k', ec='k',
                                  transform=axes[2, 1][i_z].transAxes)
            axes[2, 1][i_z].text(-0.15, 0.5,
                    '%s nA' % np.around(y_max - y_min, decimals=0),
                    rotation=90,
                    transform=axes[2, 1][i_z].transAxes,
                    horizontalalignment='center',
                    verticalalignment='center')

axes[1, 0].legend(loc='lower left', bbox_to_anchor=(-0.02, 1.55), ncol=4,
                  columnspacing=4, #handletextpad=1,
                  bbox_transform=axes[1, 0].transAxes)

axes[0, 0].text(-0.05, 1.05, '(A)', size=12, weight='bold',
                va='bottom', ha='right', transform=axes[0, 0].transAxes)
axes[0, 1].text(-0.05, 1.05, '(B)', size=12, weight='bold',
                va='bottom', ha='right', transform=axes[0, 1].transAxes)

fig.align_ylabels([axes[0, 0], axes[1, 0], axes[2, 0][0]])
#grid.tight_layout(fig, pad=0.1, rect=(0, 0, -0.8, 0.95))
grid.update(wspace=0.4, hspace=0)

plt.savefig('figure-2/fig2.pdf', format='pdf', pad_inches=0.02,
            bbox_inches='tight')
fig.canvas.start_event_loop(sys.float_info.min)  # Silence Tkinter callback
plt.savefig('figure-2/fig2', pad_inches=0.02, dpi=300, bbox_inches='tight')
#plt.show()
plt.close()