modelDef.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "h_DsRnn_def.ipynb\n",
    "\n",
    "## Imports for h_DsRnn_v*.ipynb\n",
    "\n",
    "SEE ALSO: `h_DsRnn_def_v1b.ipynb`\n",
    "\n",
    "PREC: `h_DsRnn_v3b8.ipynb`, `h_DsRnn_def.ipynb`\n",
    "\n",
    "AUTH: Hendrik.Mandelkow@icloud.com"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#< os.chdir('/home/mandelkowhc/matlab/htools1/hpy')\n",
    "!jupyter nbconvert h_DsRnn_def.ipynb --output tmp --to python \\\n",
    "    --TemplateExporter.exclude_output=True \\\n",
    "    --TemplateExporter.exclude_raw=True \\\n",
    "    --no-prompt \\\n",
    "    --TemplateExporter.exclude_markdown=False\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----------------------------------------------------------------------------------------\n",
    "#### Biowulf preliminaries:\n",
    "```bash\n",
    "module load python/3.6\n",
    "...OR\n",
    "conda_on\n",
    "conda activate Tf1\n",
    "```\n",
    "\n",
    "##### NB: GPU is required for TF backend!\n",
    "```bash\n",
    "freengpu # free nodes\n",
    "sinteractive -TT -t 24:00:00 -c 14 --mem=16g --gres=gpu:p100:1,lscratch:32'\n",
    "alias sintg='f(){ sinteractive -TT -t 36:00:00 -c $((14*${1##*:})) --mem=64g --gres=gpu:$1,lscratch:10; }; f'\n",
    "sintg 1\n",
    "sintg p100:1\n",
    "\n",
    "sjobs\n",
    "newwall --jobid ??? --time 36:00:00\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensorboard\n",
    "```python\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir ./TB_logs --port $PORT2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# WARNING: Guess this must come before imports!?!\n",
    "# HOWTO enable autoreload for imported modules.\n",
    "%load_ext autoreload\n",
    "%autoreload\n",
    "# %matplotlib auto\n",
    "%matplotlib inline\n",
    "#%aimport ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## IMPORTS\n",
    "import sys, os, re, glob\n",
    "import numpy as np\n",
    "np.set_printoptions(3)\n",
    "%precision 3\n",
    "# import scipy\n",
    "# import tensorflow as tf\n",
    "# import tensorflow.keras as keras\n",
    "import keras\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "#plt.style.use('default')\n",
    "plt.style.use('dark_background') # 'default', 'dark_background', 'seaborn-talk'\n",
    "# plt.style.use('seaborn-talk')\n",
    "from pprint import pprint\n",
    "import scipy\n",
    "from scipy import signal as scsi\n",
    "import hdf5storage as hdf5\n",
    "import nibabel as niba\n",
    "import nilearn as nile\n",
    "import nilearn.plotting as niplt\n",
    "# HOWTO suppress warnings:\n",
    "from warnings import warn, filterwarnings, simplefilter\n",
    "filterwarnings('ignore',module='nilearn',lineno=1569)\n",
    "# simplefilter('ignore')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### import htools"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "if '/hpy:' not in ':'.join(sys.path)+':':\n",
    "    sys.path.insert(0,'/home/mandelkowhc/matlab/htools1/hpy')\n",
    "# pprint(sys.path[:5])\n",
    "os.chdir('/home/mandelkowhc/matlab/htools1/hpy')\n",
    "!jupyter nbconvert htools_v1b.ipynb --to python \\\n",
    "    --TemplateExporter.exclude_output=True \\\n",
    "    --TemplateExporter.exclude_raw=True \\\n",
    "    --TemplateExporter.exclude_markdown=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from htools_v1b import hsys, hcd, hstd, hmovmean, hhline, hvline, hscalez, hreshape, \\\n",
    "    hFpath, hFname, htcode64, htime64, himgtileax, hxcorry, hnormalize, hrescale, \\\n",
    "    hsavefig, ddict, hplotchan\n",
    "# from htools_v1b import hipymagic, hipyshell\n",
    "from htools_v1b import hformatwarnmsg\n",
    "import warnings\n",
    "warnings._formatwarnmsg_impl = hformatwarnmsg"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "import hbiopack as hbp\n",
    "# import hDsCl as hds\n",
    "# from h_BpRnnBw_def import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hrms = lambda X,d: np.sqrt(np.mean(np.abs(X)**2,d))\n",
    "hrss = lambda X,d: np.sqrt(np.sum(np.abs(X)**2,d))\n",
    "hnorm2 = lambda X,d: np.sqrt(np.sum(np.abs(X)**2,d))\n",
    "hcorrxy = lambda x,y: np.corrcoef(x,y,rowvar=False)[:x.shape[1],x.shape[1]:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hplotstyles = lambda : [ plt.style.use('dark_background'), mpl.rcParams.update({'figure.figsize': (18,4)}) ];\n",
    "hplotstyles = lambda : [ plt.style.use('seaborn-talk'), mpl.rcParams.update({'figure.figsize': (18,4)}) ];\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hind2sub = np.unravel_index # lin.index to multi-subscripts\n",
    "hsub2ind = np.ravel_multi_index # to lin.index\n",
    "\n",
    "def hPutMaskSorted(M,X,x=0):\n",
    "    '''[***] Similar to np.putmask but with sorted M.values\n",
    "    '''\n",
    "    # np.put( np.zeros(M.shape, X.dtype), np.unravel_index(np.argsort(M,None)[-X.size:],M.shape), X )\n",
    "    # Img = np.zeros(list(M.shape)+X.shape[1:], X.dtype) + x\n",
    "    # Img = np.zeros(M.shape, X.dtype) + x\n",
    "    Img = np.full(M.shape, x, dtype=X.dtype)\n",
    "    # [+++] HOWTO assign vlues to an ordered mask:\n",
    "    Img[np.unravel_index(np.argsort(M,None)[-X.size:],M.shape)] = X\n",
    "    return Img\n",
    "\n",
    "# Use different name + input sequence?\n",
    "hma2im = lambda X,M,x=0: hPutMaskSorted(M,X,x)\n",
    "\n",
    "# Can't work bc np.put returns None!?\n",
    "# [+--] hPutSortedMask = lambda Y,M,X: np.put( Y=np.zeros(M.shape, X.dtype), np.unravel_index(np.argsort(M,None)[-X.size:],M.shape), X )\n",
    "\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%script _bash\n",
    "# _bash prevents execution.\n",
    "cd $ExId.results\n",
    "set +e # don't exit on error\n",
    "rm -f McPar.1D\n",
    "ln -s dfile.r01.1D McPar.1D || echo Link exists.\n",
    "3dAFNItoNIFTI -overwrite -prefix Epi_mask.nii.gz full_mask.*.BRIK*\n",
    "if $OW || [ ! -e Epi_Mc.nii* ] ; then\n",
    "\t3dAFNItoNIFTI -overwrite -float -prefix Epi_Mc.nii.gz pb01.*.BRIK*\n",
    "fi\n",
    "if $OW || [ ! -e Epi_Mcr.nii* ] ; then\n",
    "\t3dAFNItoNIFTI -overwrite -float -prefix Epi_Mcr.nii.gz errts.*.BRIK*\n",
    "fi\n",
    "if $OW || [ ! -e Epi_Mcr_std.nii* ] ; then\n",
    "\t3dTstat -overwrite -stdev -prefix Epi_Mcr_std.nii.gz Epi_Mcr.nii*\n",
    "fi\n",
    "if $OW || [ ! -e Epi_Mc_mean.nii* ] ; then\n",
    "\t# 3dTstat -overwrite -mean -std -prefix Epi_Mc_mean+std.nii.gz Epi_Mc.nii*\n",
    "\t3dTstat -overwrite -mean -prefix Epi_Mc_mean.nii.gz Epi_Mc.nii*\n",
    "\t3dTstat -overwrite -stdev -prefix Epi_Mc_std.nii.gz Epi_Mc.nii*\n",
    "fi\n",
    "# find . -iname \"*mask*.BRIK\" -print -exec 3dAFNItoNIFTI {} \\;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MASKVAL = +0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "******************************************************************\n",
    "# Batch generator\n",
    "Training a *stateful* RNN on multiple (independent) voxels in parallel requires specially formed batches of training data. These could either be created and fed manually using model.train_on_batch() or by using model.fit_generator() with a custom generator - see below.\n",
    "\n",
    "#### Re: Seq. length\n",
    "In simple Keras the seq. length is the same for training and prediction. It determines e.g. the \"depth\" of backprop in time. If the RNN *stateful* stride should equal seq. length and the batch size Nbatch equal to 1 or, perhaps, the number of inputs (voxels) trained in parallel. If the network is *not* stateful prediction may require longer sequences, but they can overlap i.e. the *stride* can be short e.g. 1TR.\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# NOTE: TSgen will not generate sequence output (Y), just one point.\n",
    "from keras.preprocessing.sequence import TimeseriesGenerator\n",
    "TR = 3; Fs = Bp.Fs; Tseq = 30; Nseq = int(Tseq//TR*Fs)\n",
    "TSgenXY = TimeseriesGenerator(Data[:,[0,1,n]], np.roll(Data[:,n],-1), length=Nseq, stride=Nseq, \\\n",
    "                              start_index=int(TR*Fs), batch_size=Nbatch)\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#%% SCRATCH\n",
    "from scipy.sparse import csc_matrix\n",
    "# Data = csc_matrix((Xtv.flat,(range(TRFS-1,Xtv.shape[0]*TRFS,TRFS),range(NX,NX+NY))),dtype=np.float32)\n",
    "# Data = csc_matrix((Xtv.flat,(range(TRFS-1,Xtv.shape[0]*TRFS,TRFS),range(NX,NX+NY))),dtype=np.float32)\n",
    "#Data = np.concatenate([Bp.Fata[:Data.shape[0],:],Data],axis=1)\n",
    "#Data = csc_matrix([Xtv.shape[0]*TRFS+1, NX+NY])\n",
    "CSC = csc_matrix((Xtv.shape[0]*TRFS+1, NX+NY), dtype=np.float32)\n",
    "CSC.shape"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "def hXtv2Data(X,Y,dt,t0=None):\n",
    "    NX,NY = X.shape[1],Y.shape[1]\n",
    "    t0 = dt-1 if t0 is None else t0\n",
    "    Data = np.zeros((Y.shape[0]*dt+1, NX+NY), np.float32)\n",
    "    Data[:,:NX] = X[:Data.shape[0],:] # +++\n",
    "    Data[t0:Y.shape[0]*dt:dt, NX:] = Y[:,:NY]\n",
    "    return Data\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "def hXtv2Data(X,Y,dt,t0=None,Sparse=False):\n",
    "    '''\n",
    "    Data = hXtv2Data( Bp.Fata, Xtv, Bp.Fs*TR )\n",
    "    '''\n",
    "    NX,NY = X.shape[1],Y.shape[1]\n",
    "    if t0 is None:\n",
    "        t0 = dt-1\n",
    "    if Sparse:\n",
    "        # CSR: Sparse matrix stored in contiguous rows:\n",
    "        Data = csr_matrix((Y.shape[0]*dt+1, NX+NY), np.float32)\n",
    "    else:\n",
    "        Data = np.zeros((Y.shape[0]*dt+1, NX+NY), np.float32)\n",
    "    Data[:,:NX] = X[:Data.shape[0],:] # +++\n",
    "    Data[t0:Y.shape[0]*dt:dt, NX:] = Y[:,:NY]\n",
    "    return Data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "def hXtv2Data( X, Y, dt, t0=None, MaskVal=MASKVAL, MaskCol=False, Sparse=False):\n",
    "    '''[***3ab] Cat time series of unequal sampling rate.\n",
    "    Cat X with Y upsampled by dt with t0 offset and MaskVal between samples.\n",
    "    \n",
    "    USE: Data[:,:NX+NY] = hXtv2Data( Bp.Fata[:,:NX], Xtv[:,:NY], Bp.Fs*TR )\n",
    "    \n",
    "    RETURNS: Data with Data[:,:NX] = X[:,NX] and Data[ t0::dt, NX:] = Y[:,:NY]\n",
    "    \n",
    "    MaskVal: (scalar float) used to interpolate Y\n",
    "    MaskCol: if True prepend Y with a \"boolean\" input mask column like:\n",
    "        Y = np.c_[ Y[:,0]*0+1, Y ] and NX += 1\n",
    "    Sparse: if True return (mem.efficient) sparse type csr_matrix\n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    if MaskCol: # Prepend a mask column of 0/1\n",
    "        Y = np.c_[ Y[:,0]*0+1, Y ]\n",
    "        # Y = np.c_[ np.ones_like(Y[:,0]), Y ]\n",
    "    '''\n",
    "    \n",
    "    NX,NY = X.shape[1],Y.shape[1]\n",
    "    if t0 is None:\n",
    "        t0 = dt-1\n",
    "\n",
    "    Nt = min( Y.shape[0], X.shape[0]//dt)\n",
    "    if Sparse:\n",
    "        # CSR: Sparse matrix stored in contiguous rows:\n",
    "        Data = csr_matrix( (Nt*dt+t0, NX+NY), np.float32)\n",
    "    else:\n",
    "        Data = np.zeros( (Nt*dt, NX+NY), np.float32)\n",
    "\n",
    "    if MaskVal:\n",
    "        Data[...] = MaskVal # Use masking value for Xtv\n",
    "        \n",
    "    # TEST: print(X.shape,end=' '); print(Y.shape,end=' '); print(Data.shape)\n",
    "    Data[:,:NX] = X[:Data.shape[0],:] # +++\n",
    "    # Data[t0:Y.shape[0]*dt:dt, NX:] = Y[:,:NY]\n",
    "    # Data[t0::dt, NX:] = Y[:,:NY]\n",
    "    Data[t0::dt, NX:] = Y[:Nt,:NY]\n",
    "    \n",
    "    if MaskCol:\n",
    "        # numpy.insert(arr, obj, values, axis=None)\n",
    "        Data = np.insert( Data, NX, Data[:,-1]!=MaskVal, axis=1)\n",
    "    \n",
    "    return Data\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "def hXtv2df( X, Y, dt, t0=None, MaskVal=MASKVAL, MaskCol=False, Sparse=False):\n",
    "    '''[***3ab] Cat time series of unequal sampling rate.\n",
    "    Cat X with Y upsampled by dt with t0 offset and MaskVal between samples.\n",
    "    \n",
    "    USE: Data[:,:NX+NY] = hXtv2Data( Bp.Fata[:,:NX], Xtv[:,:NY], Bp.Fs*TR )\n",
    "    \n",
    "    RETURNS: Data with Data[:,:NX] = X[:,NX] and Data[ t0::dt, NX:] = Y[:,:NY]\n",
    "    \n",
    "    MaskVal: (scalar float) used to interpolate Y\n",
    "    MaskCol: if True prepend Y with a \"boolean\" input mask column like:\n",
    "        Y = np.c_[ Y[:,0]*0+1, Y ] and NX += 1\n",
    "    Sparse: if True return (mem.efficient) sparse type csr_matrix\n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    if MaskCol: # Prepend a mask column of 0/1\n",
    "        Y = np.c_[ Y[:,0]*0+1, Y ]\n",
    "        # Y = np.c_[ np.ones_like(Y[:,0]), Y ]\n",
    "    '''\n",
    "    \n",
    "    NX,NY = X.shape[1],Y.shape[1]\n",
    "    if t0 is None:\n",
    "        t0 = dt-1\n",
    "\n",
    "    Nt = min( Y.shape[0], X.shape[0]//dt)\n",
    "    Data = pd.DataFrame( np.zeros( (Nt*dt, NX+NY), np.float32), columns=['X'+str(n) for n in range(NX)]\n",
    "\n",
    "    if MaskVal:\n",
    "        Data[...] = MaskVal # Use masking value for Xtv\n",
    "        \n",
    "    print(X.shape,end=' '); print(Y.shape,end=' '); print(Data.shape)\n",
    "    Data[:,:NX] = X[:Data.shape[0],:] # +++\n",
    "    # Data[t0:Y.shape[0]*dt:dt, NX:] = Y[:,:NY]\n",
    "    Data[t0::dt, NX:] = Y[:,:NY]\n",
    "    \n",
    "    if MaskCol:\n",
    "        # numpy.insert(arr, obj, values, axis=None)\n",
    "        Data = np.insert( Data, NX, Data[:,-1]!=MaskVal, axis=1)\n",
    "    \n",
    "    return Data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## v4b new ValFrac\n",
    "See h_DsRnn_def_v1b.ipynb for older / alternate versions.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TODO:\n",
    " * [x] add Tsh for augmentation\n",
    " * [ ] add suffle + reverse as control\n",
    " * possibly avoid cp of data by reshaping batch instead of data?\n",
    " * Incorporate upsampling of MRI data?\n",
    " * set shift = 0 to test training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Add data augmentation by random time shift.\n",
    "class hBatchSeq1y(keras.utils.Sequence):\n",
    "    '''[**4b++] Batch generator (keras.*.Sequence) for stateful RNN with NX,NY,NB > 1.\n",
    "    Stack NB and NY into a batch \n",
    "    For use with STATEFUL=True\n",
    "    Returns: ( X[NY*NB,NT,NX+1], Y[NY*NB,NT,1]) or ( X, Y, Mask )\n",
    "    \n",
    "    Data[t,:NX+NY]\n",
    "    NT: length of each training sequence [ batch_size=(NY*NB, NT, NX+1) ]\n",
    "    NB: split Data[t,c] into NB sections along t for parallel training\n",
    "    NY= Data.shape[-1]-NX : split Data[t,:NX+NY] for parallel training of each NY\n",
    "    Drop: Dropout 0 < Drop < 1.0, mask Drop*100% of input Y at random [DropMode= 'sample']\n",
    "        DropMode='sequence' # drop Y input for entire samples (sequences - NT) at random\n",
    "        DropMode='odd' # drop Y input for odd samples (sequences - NT)\n",
    "        DropMode='last' # drop Y input for last samples (sequences - NT) in each section\n",
    "    Mask = *sample_weights* = either 1D array of Batch.shape[0] or 2D of Batch.shape[:2]\n",
    "        Return ( X, Y, sample_weights) to serve as a mask for cost functions\n",
    "        Need to set sample_weight_mode = 'temporal' ?!?\n",
    "    ValFrac: if >0 leave out ValFrac*100% at the end of each section\n",
    "    WARNING: This results in *distributed* val.data e.g. NB=4, ValFrac=1/3 -> TTV,TTV,TTV,TTV\n",
    "    ValFrac > 0 : deliver training data\n",
    "    ValFrac < 0 : deliver validation data\n",
    "    \n",
    "    .getX() : retrieve X (full TS)\n",
    "    .getY() : retrieve Y for comparison with Yh\n",
    "    .predict( model ) : compute Yh\n",
    "    .evaluate( model ) : compute losses\n",
    "    .reshapeInput : change model input_shape for prediction on different batch size\n",
    "    ...\n",
    "    \n",
    "    SEE: h_BpRnnBw_def.py\n",
    "    '''\n",
    "    # AUTH: Hendrik.Mandelkow@gmail.com\n",
    "    \n",
    "    def __init__(self, Data, NT, NX=None, NB=1, ValFrac=0, Mask=None, Drop=0, DropMode='sample', Xtrafo=None):\n",
    "        # self.__dict__.update(Data=Data, NT=NT, NY=NY, NB=NB, NX=Data.shape[-1]-NY)\n",
    "        self.__dict__.update(NT=NT, NX=NX, NB=NB, Mask=Mask, Drop=Drop, DropMode=DropMode, MaskVal=-10)\n",
    "        self.__dict__.update(Xtrafo=Xtrafo)\n",
    "        try: self.Drop, self.DropMode = self.Drop[0], self.Drop[1]\n",
    "        except: pass\n",
    "        self.MaskVal = +0.0; warn('+++ TEST +++ MaskVal.')\n",
    "        if NX is None: self.NX = NX = Data.shape[-1]-1\n",
    "        self.NY = Data.shape[-1] - self.NX\n",
    "        self.Xlead = self.RandXlead = 0\n",
    "        self.Yscale = self.RandYscale = 0\n",
    "        assert not (self.RandXlead or self.Xlead and ValFrac), 'Oops! RandXlead and ValFrac are likely incompatible.'\n",
    "        \n",
    "        #< Data[:,self.NX:] = hzscore(Data[:,self.NX:])\n",
    "        \n",
    "        assert Data.ndim <= 3, 'Oops!'\n",
    "        if False and Data.ndim>2:\n",
    "            warn('Using (NB,NT,NXY) = Data.shape')\n",
    "            self.NT = NT = Data.shape[-2]\n",
    "            self.NB = NB = Data.shape[0]\n",
    "            Data = Data.reshape(-1,Data.shape[-1])\n",
    "            \n",
    "        ### Transform X input - very BETA\n",
    "        if isinstance( self.Xtrafo, str): self.Xtrafo = [ self.Xtrafo ]\n",
    "        if self.Xtrafo[0]=='step':\n",
    "            # import pdb; pdb.set_trace()\n",
    "            n = Data[:,-1] != self.MaskVal\n",
    "            tmp = Data[ n, self.NX: ]\n",
    "            tmp = np.diff(tmp,axis=0,prepend=0)\n",
    "            Data[ n, self.NX:] = tmp\n",
    "            Data[ :, self.NX:] = np.cumsum( Data[ :, self.NX:], 0)\n",
    "\n",
    "        ### Reshape Data\n",
    "        Data = Data[:Data.shape[0]//NT//NB*NB*NT,:].reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [B,NB,NT,NX+NY]\n",
    "        if ValFrac:\n",
    "            print('+++ WARNING: *Interleaved* validation data at the end of each block.')\n",
    "        if ValFrac > 0:\n",
    "            print('+ Training data.')\n",
    "            Data = Data[:-round(abs(ValFrac)*Data.shape[0])]\n",
    "        elif ValFrac < 0:\n",
    "            print('+ Validation data.')\n",
    "            Data = Data[-round(abs(ValFrac)*Data.shape[0]):]\n",
    "\n",
    "        print('Batches per epoch: %u, batch size (NY*NB): %u' \\\n",
    "              %( Data.shape[0], Data.shape[1]*self.NY))\n",
    "        \n",
    "        assert Data.size > 0, 'Oops, Data.shape= '+str(Data.shape)\n",
    "        \n",
    "        self.Data = Data\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.Data.shape[0]\n",
    "\n",
    "    def __getitem__(self,idx): # return one *batch*!\n",
    "        NT,NX,NY,NB = map(self.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        X = self.Data[idx] # One batch: X[NB,NT,NX+NY]\n",
    "        if True:\n",
    "            X = np.concatenate([ X[:,:,[*range(NX),NX+y]] for y in range(NY)], axis=0) # X[NY*NB,NT,NX+1]\n",
    "        else:\n",
    "            assert False,'NOT TESTED!'\n",
    "            warn('NOT TESTED!')\n",
    "            X = np.concatenate((np.tile(X[:,:,:NX],[1,1,1,NY]),X[:,:,None,:]),2) # [NB,NT,NX+1,NY]\n",
    "            X = np.moveaxis(X,-1,0).reshape(NY*NB,NT,NX+1) # X[NY*NB,NT,NX+1]\n",
    "\n",
    "        ### Shift Y in time to make prediction non-trivial.\n",
    "        # Circshift each seq. (NT) may be suboptimal but simple and irrelevant.\n",
    "        Y = np.roll(X[:,:,-1:],-1,axis=-2) # +++ Y[NY*NB,NT,1] (out) shifted -1 rel to X\n",
    "        # Y = np.roll(X[:,:,-1:],0,axis=1); warn('TEST TEST TEST!')\n",
    "\n",
    "        ### Transform X input\n",
    "        if isinstance( self.Xtrafo, str): self.Xtrafo = [ self.Xtrafo ]\n",
    "        if not self.Xtrafo:\n",
    "            pass\n",
    "        elif self.Xtrafo[0]=='step': # restore mask values in Y\n",
    "            # if np.all( self.Data[ self.Data[:,NX+1]!=1, NX+1] == 0): # boolean mask col. NX+1\n",
    "            tmp = np.diff( Y, axis=-2, prepend=Y[...,:1,:]) == 0\n",
    "            Y[ tmp ] = self.MaskVal\n",
    "            Y[:,-1,-1] = self.MaskVal\n",
    "        elif self.Xtrafo[0]=='randn':\n",
    "            X[...,self.Xtrafo[1]] = np.random.randn(**X[...,self.Xtrafo[1]].shape)\n",
    "        elif self.Xtrafo[0]=='zeros':\n",
    "            X[...,self.Xtrafo[1]] = 0\n",
    "        elif self.Xtrafo[0]=='const':\n",
    "            X[...,self.Xtrafo[1]] = self.Xtrafo[2]\n",
    "        else:\n",
    "            assert (self.Xtrafo is None), 'Oops! Parameter Xtrafo not recognized.'\n",
    "        \n",
    "        ### Dropout to decrease reliance on BOLD autocorrelations\n",
    "        if self.Drop:\n",
    "            assert (0 <= self.Drop <= 1), 'Oops! Expecting 0 < Drop < 1.'\n",
    "            tmp = self.DropMode[:3].lower()\n",
    "            if tmp in ['sam']: # samples\n",
    "                assert not self.Xtrafo, 'Error, combined Xtrafo and Drop Samples is not implemented.'\n",
    "                X[ np.random.random(X.shape[:-1])<self.Drop, -1] = self.MaskVal # ***\n",
    "            if tmp in ['seq']: # sequences\n",
    "                X[ np.random.random(X.shape[:1])<self.Drop, :, -1] = self.MaskVal # ***\n",
    "            if tmp in ['odd']: # odd sequences 1,3,5,...\n",
    "                if (idx % 2): X[ :, :, -1] = self.MaskVal # ***\n",
    "            if tmp in ['las']: # drop last len()*Drop sequences \n",
    "                if idx >= len(self)*(1-self.Drop): X[ :, :, -1] = self.MaskVal # ***\n",
    "            if tmp in ['lxy']: # drop last len()*Drop sequences in X and all other in Y (no feedback)\n",
    "                if idx >= len(self)*(1-self.Drop): X[ :, :, -1] = self.MaskVal # ***\n",
    "                else: Y[:,:,-1] = self.MaskVal\n",
    "\n",
    "            # Old: if np.all( np.logical_or( X[...,-2]==0, X[...,-2]==1) ): # more complicated\n",
    "            if np.all( np.isin( X[...,-2], [0,1]) ): # If X[...,-2] is a boolean mask...\n",
    "            # ...null any elements that (now) coincide with X[...,-1]==MaskVal\n",
    "                X[...,-2] = np.logical_and( X[...,-2], X[...,-1]!=self.MaskVal)\n",
    "                # Simpler?!: X[ X[...,-1]==self.MaskVal, -2] = False\n",
    "        \n",
    "        ### Return mask?\n",
    "        #< X[X==np.nan] = self.MaskVal\n",
    "        #< X[ np.isnan(X[...,-1]), -1] = self.MaskVal\n",
    "        #< assert not np.any(np.isnan(X)), 'Oops NaNs in X!?!'\n",
    "        if self.Mask is None:\n",
    "            return ( X, Y )\n",
    "        elif isinstance( self.Mask, np.ndarray ):\n",
    "            # assert False, 'Not tested!?!'\n",
    "            return ( X, Y, self.Mask )\n",
    "            # Could use Mask = np.any(Y,-1).astype(float)\n",
    "        else:\n",
    "            assert False, 'TEST! Not this way!'\n",
    "            Mask = np.any( Y != self.Mask, -1).astype(float) # Mask[NY*NB,NT] might be correct?!?\n",
    "            #< Mask = np.any( np.logical_not(np.isnan(Y)), -1).astype(float) # Mask[NY*NB,NT] might be correct?!?\n",
    "            return ( X, Y, Mask ) # (..., sample_weights)\n",
    "        \"\"\"\n",
    "        try:\n",
    "            assert False, 'Not tested!?!'\n",
    "            assert self.Mask.size > 1, 'Not an array!?'\n",
    "            return ( X, Y, self.Mask )\n",
    "            # Could use Mask = np.any(Y,-1).astype(float)\n",
    "        except (AttributeError): # AssertionError\n",
    "            # assert False, 'Error: Work in progress.'\n",
    "            Mask = np.any( Y != self.Mask, 2).astype(float) # Mask[NY*NB,NT] might be correct?!?\n",
    "            return ( X, Y, Mask )\n",
    "        \"\"\"\n",
    "        \n",
    "    def on_epoch_end(self):\n",
    "        '''TODO: Add some data augmentation eg by random t-shift here.\n",
    "        '''\n",
    "        # It should really only matter for card (and resp)?\n",
    "        if self.RandXlead:\n",
    "            self.setXlead( np.random.randint( 1+self.RandXlead ) ) # new random shift\n",
    "    \n",
    "        if self.RandYscale:\n",
    "            # assert False, 'NOT TESTED!'\n",
    "            self.setYscale( (np.random.rand()-0.5)*2*self.RandYscale + 1 ) # new random scale\n",
    "    \n",
    "    \n",
    "    def unstackData(self,Data):\n",
    "        '''\n",
    "        # Data = self.Data # [B,NB,NT,NX+NY]\n",
    "        '''\n",
    "        NT,NX,NY,NB = map(self.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        # Data = self.Data # [B,NB,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [NB,B,NT,NX+NY]\n",
    "        Data = Data.reshape( -1, Data.shape[-1]) # [T,NX+NY]\n",
    "        return Data\n",
    "        \n",
    "    def restackData(self,Data):\n",
    "        '''\n",
    "        '''\n",
    "        NT,NX,NY,NB = map(self.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        #< Data = Data[:Data.shape[0]//NT//NB*NB*NT,:].reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = Data.reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [B,NB,NT,NX+NY]\n",
    "        # self.Data = Data\n",
    "        return Data\n",
    "        \n",
    "    \n",
    "    def setYscale(self, Yscale, Ch=3 ):\n",
    "        '''\n",
    "        '''\n",
    "        # assert False, 'NOT TESTED!'\n",
    "        if not self.Yscale: self.Yscale = 1.0\n",
    "        Data = self.unstackData( self.Data)\n",
    "        Data[:,Ch] = Data[:,Ch] * (Yscale/self.Yscale)\n",
    "        self.Yscale = Yscale\n",
    "        self.Data = self.restackData( Data )\n",
    "        # return # what?\n",
    "    \n",
    "    \n",
    "    def setXlead(self, Xlead, Ch=slice(2) ):\n",
    "        '''Tshift Xchans left to lead Y by Xlead samples.\n",
    "        : self.Xlead= 0: Current (random) time shift (left) in samples np.roll(Data[:,Ch],-Xlead,0)\n",
    "        : self.RandXlead= 0: Max for random self.Xlead\n",
    "        self.setXlead(0) : undo random time shift\n",
    "        '''\n",
    "        ## self.getX, transform, reshape and store as self.Data\n",
    "        Data = self.unstackData( self.Data)        \n",
    "        Data[:,Ch] = np.roll( Data[:,Ch], self.Xlead-Xlead, 0) # apply new shift (left)\n",
    "        self.Xlead = Xlead\n",
    "        self.Data = self.restackData( Data )\n",
    "        # return # what?\n",
    "    \n",
    "    def setXlead0(self, Xlead, Ch=slice(2) ): # OLD: superceded by s.above\n",
    "        '''Tshift Xchans left to lead Y by Xlead samples.\n",
    "        : self.Xlead= 0: Current (random) time shift (left) in samples np.roll(Data[:,Ch],-Xlead,0)\n",
    "        : self.RandXlead= 0: Max for random self.Xlead\n",
    "        self.setXlead(0) : undo random time shift\n",
    "        '''\n",
    "        ## self.getX, transform, reshape and store as self.Data\n",
    "        NT,NX,NY,NB = map(self.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        Data = self.Data # [B,NB,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [NB,B,NT,NX+NY]\n",
    "        Data = Data.reshape( -1, Data.shape[-1]) # [T,NX+NY]\n",
    "        \n",
    "        #< if self.Xlead is None: slef.Xlead = 0 # unnecessary\n",
    "        Data[:,Ch] = np.roll( Data[:,Ch], self.Xlead-Xlead, 0) # apply new shift (left)\n",
    "        self.Xlead = Xlead\n",
    "        \n",
    "        #< Data = Data[:Data.shape[0]//NT//NB*NB*NT,:].reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = Data.reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [B,NB,NT,NX+NY]\n",
    "        self.Data = Data\n",
    "        \n",
    "        # return # what?\n",
    "    \n",
    "    def getX(Bgen):\n",
    "        '''[*1a+]\n",
    "        '''\n",
    "        # PREC: hBatchGen_getX()\n",
    "        # OK this works, according to the test below.\n",
    "        XY = 0 # XY = 0,1 = getX, getY\n",
    "        NT,NX,NY,NB = map(Bgen.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        x = np.stack([ Bgen[n][0] for n in range(len(Bgen)) ],1) # y[NY*NB,B,NT,NX+1]\n",
    "        assert NT == x.shape[-2], 'Oops!'\n",
    "        assert NB == x.shape[0]//NY, 'Oops!'\n",
    "        assert (NX+1) == x.shape[-1], 'Oops!'\n",
    "        x = x.reshape(NY,NB,len(Bgen),NT,NX+1) # [NY,NB,B,NT,NX+1]\n",
    "        x = np.concatenate((x[0,:,:,:,:NX], np.moveaxis(x[:,:,:,:,-1],0,-1)),-1) # [NB,B(PE),NT,NX+NY]\n",
    "        x = x.reshape(-1,x.shape[-1]) \n",
    "        return x\n",
    "\n",
    "    def getY(Bgen,Tsh=False):\n",
    "        '''[*1a+]\n",
    "        Tsh=True : Undo t-shift for training.\n",
    "        Tsh=False : directly comparable to Yh\n",
    "        NB : batch size, nof samples per batch\n",
    "        B : nof batches per epoch\n",
    "        '''\n",
    "        XY = 1 # get Y\n",
    "        NT,NX,NY,NB = map(Bgen.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        y = np.stack([ Bgen[n][XY] for n in range(len(Bgen)) ],1) # y[NY*NB,B,NT,1]\n",
    "        assert NT == y.shape[-2], 'Oops!'\n",
    "        assert NB == y.shape[0]//NY, 'Oops!'\n",
    "        y = np.reshape(y,(NY,NB,len(Bgen),NT)) # [NY,NB,B,NT]\n",
    "        # y = np.transpose(y,(1,2,3,0)) # y[NB,B,NT,NY]\n",
    "        y = np.moveaxis(y,0,-1) # y[NB,B,NT,NY]\n",
    "        if Tsh:\n",
    "            y = np.roll(y,1,-2)\n",
    "        y = y.reshape(-1,y.shape[-1]) # y[NB*N*NT,NY]\n",
    "        return y\n",
    "\n",
    "    def unbatchYh(self,Yh,NY=None,NB=None,Tsh=False):\n",
    "        # This could be a static function Yh2Y()?\n",
    "        '''\n",
    "        # Yh[B*NY*NB,NT,1]\n",
    "        Yh[t,NY] = hUnbatchYh( RNN.predict_generator( hBatchSeq1y(...)))\n",
    "        '''\n",
    "        NY = self.NY if NY is None else NY\n",
    "        NB = self.NB if NB is None else NB\n",
    "        NT = Yh.shape[-2]\n",
    "        yh = np.reshape(Yh,(-1,NY,NB,NT)) # yh[B,NY,NB,NT]\n",
    "        yh = np.transpose(yh,(2,0,3,1)) #yh[NB,B,NT,NY]\n",
    "        if Tsh:\n",
    "            yh = np.roll(yh,1,-2)\n",
    "        yh = yh.reshape(-1,yh.shape[-1]) # yh[NB*B*NT,NY]\n",
    "        return yh\n",
    "    \n",
    "    def predict(Bgen, RNN, Tsh=False, Reset=True):\n",
    "        '''Run generator batches through RNN and reshuffle output into an array.\n",
    "        Bgen = TrainGen or ValidGen\n",
    "        '''\n",
    "        if Reset:\n",
    "            RNN.reset_states()\n",
    "        Yh = RNN.predict_generator(Bgen,verbose=1) # Yh[B*NY*NB,NT,1]\n",
    "        #< print(Yh.shape)\n",
    "        # Yh = hUnbatchYh(Yh,Bgen.NY,Bgen.NB,Tsh)\n",
    "        Yh = Bgen.unbatchYh(Yh,Bgen.NY,Bgen.NB,Tsh)\n",
    "        return Yh\n",
    "\n",
    "    def predict1(Bgen, RNN, Tsh=False, Reset=True):\n",
    "        '''Predict using new RNN with input shape matching Bgen.\n",
    "        E.g. use NB=1 for better stateful prediction.\n",
    "        Bgen = TrainGen or ValidGen\n",
    "        '''\n",
    "        # RNNp = keras.models.clone_model(RNN)\n",
    "        RNNp = keras.models.model_from_json(RNN.to_json())\n",
    "        # RNNp._layers[1].batch_input_shape = (NY,NT,NX+1)\n",
    "        RNNp._layers[1].batch_input_shape = Bgen[0][0].shape\n",
    "        RNNp = keras.models.model_from_json(RNNp.to_json())\n",
    "        RNNp.set_weights(RNN.get_weights())\n",
    "        # [ RNNp.layers[n].set_weights(RNN.layers[n].get_weights()) for n in range(len(RNN.layers))]\n",
    "        # ??? [ L.stateful= True for L in RNNp.layers if hasattr(L,'stateful') ]\n",
    "        # ??? [ L.batch_input_shape= (1,None,3) for L in RNNp.layers if hasattr(L,'stateful') ]\n",
    "        RNNp.summary()\n",
    "\n",
    "        if Reset:\n",
    "            RNNp.reset_states()\n",
    "        Yh = RNNp.predict_generator(Bgen,verbose=1) # Yh[B*NY*NB,NT,1]\n",
    "        #< print(Yh.shape)\n",
    "        # Yh = hUnbatchYh(Yh,Bgen.NY,Bgen.NB,Tsh)\n",
    "        Yh = Bgen.unbatchYh(Yh,Bgen.NY,Bgen.NB,Tsh)\n",
    "        return Yh\n",
    "\n",
    "    def reshapeInput(Bgen, RNN, InputShape=None):\n",
    "        '''Cp RNN weights to new model with batch_input_shape matching generator (for prediction)\n",
    "        Bgen = TrainGen or ValidGen\n",
    "        RNN\n",
    "        InputShape = batch_input_shape = (NB,NT,NY)\n",
    "        '''\n",
    "        # K.clear_session()\n",
    "        # RNNp = keras.models.clone_model(RNN)\n",
    "        RNNp = keras.models.model_from_json(RNN.to_json())\n",
    "        # RNNp._layers[1].batch_input_shape = (NY,NT,NX+1)\n",
    "        if InputShape is None: InputShape =  Bgen[0][0].shape\n",
    "        # RNNp._layers[1].batch_input_shape = Bgen[0][0].shape\n",
    "        RNNp._layers[1].batch_input_shape = InputShape\n",
    "        RNNp = keras.models.model_from_json(RNNp.to_json())\n",
    "        RNNp.set_weights(RNN.get_weights())\n",
    "        # [ RNNp.layers[n].set_weights(RNN.layers[n].get_weights()) for n in range(len(RNN.layers))]\n",
    "        # ??? [ L.stateful= True for L in RNNp.layers if hasattr(L,'stateful') ]\n",
    "        # ??? [ L.batch_input_shape= (1,None,3) for L in RNNp.layers if hasattr(L,'stateful') ]\n",
    "        \n",
    "        try:\n",
    "            RNNp.compile(optimizer=RNN.optimizer, loss=RNN.loss, metrics=RNN.metrics)\n",
    "        except:\n",
    "            warn('Compile failed?!?')\n",
    "            pass\n",
    "        \n",
    "        RNNp.summary()\n",
    "\n",
    "        return RNNp\n",
    "\n",
    "    def evaluate(Bgen, RNN, Reset=True):\n",
    "        '''\n",
    "        Bgen = TrainGen or ValidGen\n",
    "        '''\n",
    "        if Reset:\n",
    "            RNN.reset_states()\n",
    "        Losses = RNN.evaluate_generator(Bgen,verbose=1) # Yh[B*NY*NB,NT,1]\n",
    "        Losses = dict( zip( RNN.metrics_names, Losses))\n",
    "        return Losses\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "    def on_epoch_end(self):\n",
    "        '''TODO: Add some data augmentation eg by random t-shift here.\n",
    "        '''\n",
    "        # It should really only matter for card (and resp)?\n",
    "        if self.RandXlead:\n",
    "            self.setRandXlead()\n",
    "        \n",
    "        return # what?\n",
    "    \n",
    "    \n",
    "    def setRandXlead1(self, RandXlead=None, Ch=slice(2) ):\n",
    "        '''TODO: Add some data augmentation eg by random t-shift here.\n",
    "        : self.Xlead= 0: Current random time shift by rts samples left np.roll(Data[:,Ch=0:2],-rts,0)\n",
    "        : self.RandXlead= 0: Max for self.Xlead\n",
    "        self.setRandXlead() : apply a new random time shift\n",
    "        self.setRandXlead(0) : undo and disable random time shift\n",
    "        '''\n",
    "        assert False, 'Oops! setRandXlead() is untested.'\n",
    "        # It should really only matter for card (and resp)?\n",
    "        # \n",
    "        if RandXlead is not None: self.RandXlead = RandXlead\n",
    "        if not self.RandXlead and not self.Xlead : return\n",
    "        #< assert not self.ValFrac, 'Not compatible with ValFrac'\n",
    "\n",
    "        # self.getX, transform, reshape and store as self.Data\n",
    "        Data = self.getX()\n",
    "        \n",
    "        tmp = self.Xlead # old rts\n",
    "        self.Xlead = np.random.randint( 1+self.RandXlead ) # new random shift\n",
    "        Data[:,Ch] = np.roll( Data[:,Ch], tmp-self.Xlead, 0) # apply new shift (left)        \n",
    "        \n",
    "        NT,NX,NY,NB = map(self.__dict__.get, ['NT','NX','NY','NB'])\n",
    "        #< Data = Data[:Data.shape[0]//NT//NB*NB*NT,:].reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = Data.reshape(NB,-1,NT,Data.shape[-1]) # [NB,B,NT,NX+NY]\n",
    "        Data = np.moveaxis(Data,0,1) # [B,NB,NT,NX+NY]\n",
    "        self.Data = Data\n",
    "        \n",
    "        return # what?\n",
    "    \n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hUnbatchYh(Yh,NY,NB,Tsh=False):\n",
    "    '''\n",
    "    # Yh[B*NY*NB,NT,1]\n",
    "    Yh[t,NY] = hUnbatchYh( RNN.predict_generator( hBatchSeq1y(...)))\n",
    "    '''\n",
    "    NT = Yh.shape[-2]\n",
    "    yh = np.reshape(Yh,(-1,NY,NB,NT)) # yh[B,NY,NB,NT]\n",
    "    yh = np.transpose(yh,(2,0,3,1)) #yh[NB,B,NT,NY]\n",
    "    if Tsh:\n",
    "        yh = np.roll(yh,1,-2)\n",
    "    yh = yh.reshape(-1,yh.shape[-1]) # yh[NB*B*NT,NY]\n",
    "    return yh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hPredictGen(RNN, Bgen, Tsh=False, Reset=True):\n",
    "    '''\n",
    "    Bgen = TrainGen or ValidGen\n",
    "    '''\n",
    "    if Reset:\n",
    "        RNN.reset_states()\n",
    "    Yh = RNN.predict_generator(Bgen,verbose=1) # Yh[B*NY*NB,NT,1]\n",
    "    print(Yh.shape)\n",
    "    Yh = hUnbatchYh(Yh,Bgen.NY,Bgen.NB,Tsh)\n",
    "    return Yh\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hBatchGen_getY(Bgen,Tsh=False):\n",
    "    '''[*1a+]\n",
    "    Tsh=True : Undo t-shift for training.\n",
    "    Tsh=False : directly comparable to Yh\n",
    "    NB : batch size, nof samples per batch\n",
    "    B : nof batches per epoch\n",
    "    '''\n",
    "    XY = 1 # get Y\n",
    "    NT = Bgen.NT\n",
    "    NY = Bgen.NY\n",
    "    NB = Bgen.NB\n",
    "    y = np.stack([ Bgen[n][XY] for n in range(len(Bgen)) ],1) # y[NY*NB,B,NT,1]\n",
    "    assert NT == y.shape[-2], 'Oops!'\n",
    "    assert NB == y.shape[0]//NY, 'Oops!'\n",
    "    y = np.reshape(y,(NY,NB,len(Bgen),NT)) # [NY,NB,B,NT]\n",
    "    # y = np.transpose(y,(1,2,3,0)) # y[NB,B,NT,NY]\n",
    "    y = np.moveaxis(y,0,-1) # y[NB,B,NT,NY]\n",
    "    if Tsh:\n",
    "        y = np.roll(y,1,-2)\n",
    "    y = y.reshape(-1,y.shape[-1]) # y[NB*N*NT,NY]\n",
    "    return y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hBatchGen_getX(Bgen):\n",
    "    '''[*1a+]\n",
    "    '''\n",
    "    # OK this works, according to the test below.\n",
    "    XY = 0 # XY = 0,1 = getX, getY\n",
    "    NT = Bgen.NT\n",
    "    NX = Bgen.NX\n",
    "    NY = Bgen.NY\n",
    "    NB = Bgen.NB\n",
    "    x = np.stack([ Bgen[n][0] for n in range(len(Bgen)) ],1) # y[NY*NB,B,NT,NX+1]\n",
    "    assert NT == x.shape[-2], 'Oops!'\n",
    "    assert NB == x.shape[0]//NY, 'Oops!'\n",
    "    assert (NX+1) == x.shape[-1], 'Oops!'\n",
    "    x = x.reshape(NY,NB,len(Bgen),NT,NX+1) # [NY,NB,B,NT,NX+1]\n",
    "    x = np.concatenate((x[0,:,:,:,:NX], np.moveaxis(x[:,:,:,:,-1],0,-1)),-1)\n",
    "    x = x.reshape(-1,x.shape[-1])\n",
    "    return x\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# TEST: hBatchGen_getX and Y\n",
    "Data = np.linspace(0,4*np.pi,1000)\n",
    "Data = np.stack([np.sin(Data*2),np.sin(Data*4)],1) + Data[:,None]\n",
    "Data = np.c_[Data,Data]\n",
    "# print(Data.shape)\n",
    "\n",
    "Tgen = hBatchSeq1y(Data, 100, NX=2, NB=2, ValFrac=0)\n",
    "\n",
    "x = hBatchGen_getX(Tgen)\n",
    "x.shape\n",
    "\n",
    "np.array_equal(x[:,:2],Data[:,:2])\n",
    "np.array_equal(x,Data)\n",
    "\n",
    "# plt.plot(np.c_[Data[:,2:],x[:,2:]])\n",
    "\n",
    "y = hBatchGen_getY(Tgen,Tsh=True)\n",
    "np.array_equal(y,Data[:,2:])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loss Functions:\n",
    "Missing training data can be handled either by passing sample_weights [with sample_weight_mode='temporal'] for masking the loss function in model.fit(). Alternatively, we may define a custom loss function (with integrated mask) as seen below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NB: Custom objects must be passed to model.compile and also model_load()\n",
    "if 'KCustoms' not in locals(): KCustoms = {}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: Looks like loss functions can return either a scalar or an array that will be summed over.\n",
    "# I think sample_weight should require the array, but there is no error?!\n",
    "# Maybe some erroneous broadcast goint on?!?\n",
    "\n",
    "# https://github.com/keras-team/keras/blob/master/keras/losses.py\n",
    "import keras.backend as K\n",
    "def hMSE(Y, Yh):\n",
    "    if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    Y = K.cast(Y, Yh.dtype)\n",
    "    return K.mean(K.square(Yh - Y),-1) # axis=-1 didn't matter?!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OK\n",
    "import keras.backend as K\n",
    "def hWMSE(Y,Yh):\n",
    "    '''Weighted Mean Square Error (MSE)\n",
    "    '''\n",
    "    Mask= -10.0\n",
    "    Mask= +0.0 # +++ TEST +++\n",
    "    if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    Y = K.cast(Y, Yh.dtype)\n",
    "    mask = K.not_equal(Y,Mask) # OK! could use NaN <-> 0\n",
    "    # TODO: assert there are masked values?!\n",
    "    # return K.mean(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "    mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "    return K.sum(K.square(K.abs(Yh - Y)*mask))/K.sum(mask) # OK!\n",
    "\n",
    "    # return (K.square(K.abs(Yh - Y)*mask))/K.sum(mask)*K.size(Y) # Keras would apply mean implicitly?!\n",
    "    # L = K.sum(K.square(K.abs(Yh - Y)*mask),-1)/K.sum(mask) # Keras would apply mean implicitly?!\n",
    "    # return L * K.size(L)\n",
    "\n",
    "KCustoms['hWMSE'] = hWMSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "def hMSEmask(Mask=-10.0):\n",
    "    def LossFun_(Y,Yh):\n",
    "        '''Weighted Mean Square Error (MSE)\n",
    "        '''\n",
    "        if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "        Y = K.cast(Y, Yh.dtype)\n",
    "        mask = K.not_equal(Y,Mask) # OK! could use NaN <-> 0\n",
    "        # return K.mean(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "        mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "        return K.sum(K.square(K.abs(Yh - Y)*mask))/K.maximum(K.cast(1,K.dtype(mask)),K.sum(mask)) # OK!\n",
    "    return LossFun_\n",
    "\n",
    "# KCustoms['hMSEmask'] = hMSEmask(-10.0)\n",
    "warn('+++ TEST +++ Using alternate MaskVal!')\n",
    "KCustoms['hMSEmask'] = hMSEmask(+0.0)\n",
    "# hMSEmask = hMSEmask(+0.0) # This work better???\n",
    "# KCustoms['hMSEmask'] = hMSEmask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "class hWMSEmaskCl: # Should be a sub-class of LossFunctionWrapper?!\n",
    "    def __init__(self, Mask=-10.0):\n",
    "        self.Mask = Mask\n",
    "        \n",
    "    def __call__(self,Y,Yh):\n",
    "        '''Weighted Mean Square Error (MSE)\n",
    "        '''\n",
    "        if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "        Y = K.cast(Y, Yh.dtype)\n",
    "        mask = K.not_equal(Y,self.Mask) # OK! could use NaN <-> 0\n",
    "        # return K.mean(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "        mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "        return K.sum(K.square(K.abs(Yh - Y)*mask))/K.sum(mask) # OK!\n",
    "\n",
    "# KCustoms['hWMSE'] = hWMSEmaskCl()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "# NOTE: Y,Yh[NB,NT,NY] one batch!\n",
    "# NOTE: If Y is not 0 mean, a constant prediction can yield hWRVF < 1.\n",
    "def hWRVF(Y,Yh):\n",
    "    '''Weighted Residual Variance Fraction\n",
    "    '''\n",
    "    Mask= -10.0\n",
    "    Mask= +0.0 # +++ TEST +++\n",
    "    if not K.is_tensor(Yh): Yh = K.constant(Yh) # need this?? or make variable?\n",
    "    Y = K.cast(Y, Yh.dtype) # need this??\n",
    "    mask = K.not_equal(Y,Mask) # OK! could use NaN <-> 0\n",
    "    # return K.sum(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "    mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "    #<? Y = Y - K.sum( Y*mask, -2, keepdims=True)/K.sum( mask, -2, keepdims=True)\n",
    "    # return K.sum(K.square((Yh - Y)*mask)) / K.sum(K.square(Y*mask))\n",
    "    \n",
    "    # Loss = K.sum(K.square((Yh - Y)*mask)) / K.maximum(0.000001, K.sum(K.square(Y*mask - K.sum(Y*mask)/K.sum(mask))))\n",
    "    Loss = K.sum(K.square((Yh - Y)*mask)) / (0.000001 + K.sum(K.square(Y*mask - K.sum(Y*mask)/K.sum(mask))))\n",
    "    # Loss = K.sum(K.square((Yh - Y)*mask)) / K.sum(K.square(Y*mask - K.sum(Y*mask)/K.sum(mask))) # more appropriate?!!\n",
    "    # assert K.eval( K.all( K.not_equal( Loss, np.nan))), 'Oops! NaN in output.' # TODO: make this work!\n",
    "    # assert not np.any( np.isnan( K.eval( Loss))), 'Oops! NaN in output.' # TODO: make this work!\n",
    "    # Loss = K.switch( Loss, Loss, Loss=0 )\n",
    "    return Loss\n",
    "\n",
    "KCustoms['hWRVF'] = hWRVF"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "\n",
    "a = K.variable(np.nan*np.random.rand(10,100))\n",
    "b = K.variable(np.ones((10,100))-2)\n",
    "K.eval( hWRVF( a, b))\n",
    "#np.all( K.eval( K.not_equal( K.variable( 0*np.ones((10,5))) , 2)))\n",
    "np.isnan( K.eval( K.variable( np.nan*np.ones((3,5)))))\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "import keras.backend as K\n",
    "# NOTE: Y,Yh[NB,NT,NY] one batch!\n",
    "# NOTE: Keras might expect a tensor in return [NB,NT], which it would mean over.\n",
    "def hWRVF(Y,Yh):\n",
    "    '''Weighted Residual Variance Fraction\n",
    "    Proper variance after subtracting the mean.\n",
    "    Actually makes less sense, since a constant offset would eval to zero difference.\n",
    "    Somehow this does not work!?!\n",
    "    '''\n",
    "    Mask= -10.0\n",
    "    Mask= +0.0 # +++ TEST +++\n",
    "    # if not K.is_tensor(Yh): Yh = K.constant(Yh) # need this?? maybe should not be constant?!\n",
    "    # if not K.is_tensor(Yh): Yh = K.variable(Yh) # need this?? maybe should not be constant?!\n",
    "    # Y = K.cast(Y, Yh.dtype) # need this??\n",
    "    mask = K.not_equal(Y,Mask) # OK! could use NaN <-> 0\n",
    "    # return K.sum(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "    # ... Might this work?!: idx = K.tf.where( mask ); K.gather( Yh, idx)\n",
    "    mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "    # assert False, 'Somehow this does not work!?!'\n",
    "    yh = (Yh-Y)*mask # should I reuse Yh?!\n",
    "    yh = yh - K.sum( yh, -2, keepdims=True)/K.sum( mask, -2, keepdims=True)\n",
    "    yh = K.square(yh*mask)\n",
    "    y = Y*mask\n",
    "    y = y - K.sum( y, -2, keepdims=True)/K.sum( mask, -2, keepdims=True)\n",
    "    y = K.square(y*mask)\n",
    "    return K.sum(yh) / K.sum(y) # Suppose I could return K.sum(Yh,-1)/K.sum(Y,-1)\n",
    "\n",
    "KCustoms['hWRVF'] = hWRVF"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# TEST: For some reason this fails half of the time due to rounding?!?\n",
    "tmp = np.float32(np.random.rand(32,64,4)-0.5)\n",
    "tmp[np.abs(tmp)<0.4] = 0\n",
    "a = hWRVF(tmp,tmp/2)\n",
    "tmp[tmp==0] = np.nan\n",
    "tmp = tmp - np.nanmean(tmp,-2,keepdims=True)\n",
    "b = np.nansum((tmp/2-tmp)**2) / np.nansum(tmp**2)\n",
    "K.print_tensor(a, message='a= ')\n",
    "a = K.eval(a)\n",
    "print((a,b))\n",
    "assert a==b, 'Oops'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "def hRVF(Y,Yh):\n",
    "    '''Residual Variance Fraction\n",
    "    return K.sum(K.square(Yh - Y)) / K.sum(K.square(Y))\n",
    "    '''\n",
    "    if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    Y = K.cast(Y, Yh.dtype)\n",
    "    return K.sum(K.square(Yh - Y)) / K.sum(K.square(Y))\n",
    "\n",
    "KCustoms['hRVF'] = hRVF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------------------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Correlation Loss\n",
    "NOTE: Keras defines losses with additional inputs e.g. cosine_similarity in 3 steps: \n",
    " 1. The actual loss fun [cosine_similarity(Y,Yh,axis=-1)] returns an array of losses for each batch of samples.\n",
    " 2. A wrapper sub-class( LossFunctionWrapper ), which is initialized with the extra parameters as well as\n",
    " 3. a \"reduction\" method usually sum over batch to yield a final scalar result.\n",
    " * see: https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/keras/losses.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras.backend as K\n",
    "def hCorrLossMask(Mask=0):\n",
    "    def LossFun_(Y,Yh):\n",
    "        '''\n",
    "        With Mask==0 this is very close to keras.losses.cosine_similarity(y_true, y_pred, axis=-1)\n",
    "        '''\n",
    "        if not K.is_tensor(Yh): Yh = K.constant(Yh) # Why?!\n",
    "        Y = K.cast(Y, Yh.dtype)\n",
    "        mask = K.not_equal(Y,Mask)\n",
    "        # return K.mean(K.square(Yh[mask]-Y[mask])) # BUT: No boolean indexing in Keras!\n",
    "        mask = K.cast(mask,K.dtype(Y)) # OK!\n",
    "        \n",
    "        # Yh[NB,NT,NY] ?!\n",
    "        # Y = K.normalize( Y, -2) # should be scaled appropriately before ?!\n",
    "        # Yh = K.normalize( Yh , -2)\n",
    "        Yh = K.normalize( Yh * mask, -2)\n",
    "        L = -K.sum( Y * Yh * mask, -2)\n",
    "        # Apparently, by default Keras will take the mean, if the result of loss fun is not scalar.\n",
    "        # In order for sample_weights to work Keras would need the array, wouldn't it?!\n",
    "        # Might not even want to sum over time steps for that?!\n",
    "        return L\n",
    "    return LossFun_\n",
    "\n",
    "KCustoms['hCorrLoss'] = hCorrLossMask(+0.0)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## Multiple linear regression loss?!?\n",
    "import keras.backend as K\n",
    "def hMultLinRegMask(Y,Yh):\n",
    "    '''Multiple linear regression loss with mask.\n",
    "    Requires Yh with multiple outputs (and Y to match).\n",
    "    '''\n",
    "    assert False, 'Work in progress!'\n",
    "    #< if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    #< Y = K.cast(Y, Yh.dtype)\n",
    "    Mask = +0.0\n",
    "    mask = K.not_equal(Y,Mask)\n",
    "    mask = K.cast(mask,K.dtype(Y))\n",
    "    ## norm, project, subtract,...\n",
    "    Yh = K.normalize( Yh*mask, -2)\n",
    "    Y[...,0] = Y[...,0] - Yh[...,0]*K.sum(Y[...,0]*Yh[...,0]*mask, -2)\n",
    "    for n in range(1,Y.shape[-1]):\n",
    "        Y[...,n] = Y[:,:,n-1] - Yh[:,:,n]*K.sum(Y[:,:,n]*Yh[:,:,n]*mask, -2)\n",
    "    return K.sum(K.square(K.abs(Y)*mask),-1)/K.sum(mask,-1)\n",
    "\n",
    "# KCustoms['hRVF'] = hRVF"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## Multiple linear regression loss?!?\n",
    "import keras.backend as K\n",
    "def hMultLinRegMask(Y,Yh):\n",
    "    '''Multiple linear regression loss with mask.\n",
    "    '''\n",
    "    assert False, 'Work in progress!'\n",
    "    #< if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    #< Y = K.cast(Y, Yh.dtype)\n",
    "    Mask = +0.0\n",
    "    Y = Y[...,0]\n",
    "    mask = K.not_equal(Y,Mask)\n",
    "    mask = K.cast(mask,K.dtype(Y))\n",
    "    ## norm, project, subtract,...\n",
    "    # Yh = K.normalize( Yh*mask, -2) # necessary?!\n",
    "    for n in range(Yh.shape[-1]):\n",
    "        Y = Y - Yh[:,:,n]*K.sum( Y*Yh[:,:,n]*mask, -2)\n",
    "    return K.sum( K.square(K.abs(Y)*mask),-1)/K.sum(mask,-1) # [NB,]\n",
    "\n",
    "# KCustoms['hRVF'] = hRVF"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## Multiple linear regression loss?!?\n",
    "import keras.backend as K\n",
    "def hMaxXcorr(Y,Yh):\n",
    "    '''Max Xcorr?!?\n",
    "    '''\n",
    "    assert False, 'Work in progress!'\n",
    "    #< if not K.is_tensor(Yh): Yh = K.constant(Yh)\n",
    "    #< Y = K.cast(Y, Yh.dtype)\n",
    "    Mask = +0.0\n",
    "    mask = K.not_equal(Y,Mask)\n",
    "    mask = K.cast(mask,K.dtype(Y))\n",
    "    ## norm, project, subtract,...\n",
    "    # ?!? Yh = K.normalize( Yh*mask, -2)\n",
    "    Yh = K.normalize( Yh, -2)\n",
    "    Y = K.normalize( Y, -2)\n",
    "    K.conv1d( Y, Yh) # padding= same, causal valid\n",
    "    ...\n",
    "    Y[...,0] = Y[...,0] - Yh[...,0]*K.sum(Y[...,0]*Yh[...,0]*mask, -2)\n",
    "    for n in range(Y.shape[-2]):\n",
    "        Y[...,n] = Y[:,:,n-1] - Yh[:,:,n]*K.sum(Y[:,:,n]*Yh[:,:,n]*mask, -2)\n",
    "    return K.sum(K.square(K.abs(Y)*mask),-1)/K.sum(mask,-1)\n",
    "\n",
    "# KCustoms['hRVF'] = hRVF"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "import keras.backend as K\n",
    "a = K.ones((9,7,2)) # [NB, Nt, NX=ChIn]\n",
    "b = K.ones((3,2,1)) # [Nt-Nk+1, Nk*ChIn, ChOut]\n",
    "c = K.conv1d(a,b) # [NB, Nt-Nk+1, NY=ChOut]\n",
    "c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RNN model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Consider activation functions?!\n",
    "# AUTH: HM, 2019-06-26, v3b2: Add GRU, better handling of input layer.\n",
    "# AUTH: HM, 2019-06-26, v3b3: Add Nio=Type\n",
    "# def mkRNN(Nio=[1,1], Nsteps=1, Nbatch=None, Lpar, Cpar):\n",
    "def mkRNN(Nio=[1,1], Nsteps=1, Nbatch=None, **kwarg):\n",
    "    '''[++-] Make simple stateful RNN with final dense layer.\n",
    "    Nsteps: number of time steps in each sample sequence\n",
    "    Nbatch: training batch size (number of sample sequences per batch)\n",
    "    TODO: What about initialization / regularization?\n",
    "    '''\n",
    "    PARS = {} # kwarg for Keras.layer...\n",
    "    # batch_input_shape = (batchsize,timesteps,data_dim) only for input layer\n",
    "    PARS['batch_input_shape'] = (Nbatch, Nsteps, Nio[0])\n",
    "    PARS['return_sequences']=True\n",
    "    PARS['stateful']=True # ***\n",
    "    # PARS['activation'] = 'linear' # BAD!\n",
    "    PARS.update(kwarg)\n",
    "    \n",
    "    RNN = keras.models.Sequential() # +++\n",
    "    # RNN.name = '' # clear generic name sequential_1\n",
    "    # The latest Keras (2.2.4?) does not seem to like name=''.\n",
    "    # Also, the name parameter seems to have some obscure internal uses. Better not touch.\n",
    "    \n",
    "    Type = 'L' # *** default\n",
    "    # RNN.name = str(Nio[0])+'L%u'*(len(Nio)-1)%tuple(Nio[1:])\n",
    "    # RNN.name = '%uL%u'%tuple(Nio[:2])\n",
    "    #< for Nout in Nio[1:]:\n",
    "    for n,Nout in enumerate(Nio[1:]):\n",
    "        if n==len(Nio)-2:\n",
    "            PARS['activation'] = 'linear' # output layer w/ lin. activation\n",
    "        if Nout in ['linear','tanh']:\n",
    "            PARS['activation'] = Nout\n",
    "            continue\n",
    "        if isinstance(Nout,str):\n",
    "            Type = Nout\n",
    "            continue\n",
    "        if Type == 'L':\n",
    "            RNN.add(keras.layers.LSTM( Nout, **PARS))\n",
    "        elif Type == 'G':\n",
    "            # https://keras.io/layers/recurrent/#GRU\n",
    "            # keras.layers.GRU(units, activation='tanh', recurrent_activation='hard_sigmoid',...\n",
    "            RNN.add(keras.layers.GRU( Nout, implementation=2, **PARS))\n",
    "        elif Type == 'D': # Dense\n",
    "            RNN.add(keras.layers.Dense(1, activation='linear'))\n",
    "        else:\n",
    "            raise ValueError('Type must be L,G,D not: '+Type)\n",
    "            \n",
    "        #< RNN.name += Type+'%u'%Nout\n",
    "            \n",
    "    # tmp = [str([x.units for n,x in enumerate(y.layers)]) for y in RNNs]+['line %u'%x for x in range(n+2,len(h)+1)]\n",
    "    PARS.pop('batch_input_shape',None) # remove key, only required for input layer\n",
    "\n",
    "    # TODO: Use only hWMSE\n",
    "    # TODO: Could use MaskVal; hWRVFmask(MaskVal) here\n",
    "    # RNN.compile(loss=hWMSE, optimizer='adam', metrics=[hWRVF]) # +++\n",
    "    RNN.compile(loss=KCustoms['hMSEmask'], optimizer='adam', metrics=[ KCustoms['hWRVF'] ]) # +++\n",
    "    # RNN.compile(loss='MSE', optimizer='adam', metrics=[hWMSE,hWRVF], weighted_metrics=[hRVF,hMSE], sample_weight_mode='temporal') # TEST!!!\n",
    "    # NOTE: Apparently, sample_weights are applied to \"loss\" as well as weighted_metrics.\n",
    "    # TEST! RNN.compile(loss='MSE', optimizer='adam') # TEST!!!\n",
    "    print('+ RNN.name= '+RNN.name)\n",
    "    \n",
    "    #%% Record input parameters: sadly broken!\n",
    "    tmp = locals() # Dunno why this is necessary?!\n",
    "    RNN.mkRNNargs = { n:tmp[n] for n in ['Nio','Nsteps','Nbatch'] }\n",
    "    RNN.mkRNNargs.update(kwarg)\n",
    "    RNN.mkRNNargs['Nio'] = list( Nio ) # Dunno how this can change into a pythetic ListWrapper!?!\n",
    "    # FIXIT: With tf.keras all lists get changed into a pythetic ListWrapper!?!\n",
    "    \n",
    "    return RNN\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from keras.utils import multi_gpu_model\n",
    "\n",
    "def mkRnnGpu(Model,**par):\n",
    "    \"\"\"[+++] Recompile model for multi-GPU training: MModel = multi_gpu_model(Model, **par)\n",
    "    \n",
    "    PAR = { # optional parameters / defaults\n",
    "        'gpus':None, # Nof GPUs or None for all available\n",
    "        'cpu_merge':True # ?!? force merging of weights on CPU\n",
    "        'cpu_relocation':False, # ?!? force transfer of model from GPU to CPU\n",
    "        }\n",
    "        \n",
    "    WARNING: Use Model not MModel for saving:\n",
    "        model.save(Fname)\n",
    "        model.save_weights(Fname)\n",
    "        \n",
    "    MModel.fit() will split batches (evenly) across GPUs.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    try: # Train on mult. GPUs\n",
    "        MModel = multi_gpu_model(Model, **par)\n",
    "    except: # Train on 1 CPU or GPU\n",
    "        MModel = Model\n",
    "\n",
    "    MModel.compile( loss= Model.loss, optimizer= Model.optimizer, metrics= Model.metrics )\n",
    "    \n",
    "    return MModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## hResetStatesCb callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Why reset on epoch end?\n",
    "class hResetStatesCb(keras.callbacks.Callback):\n",
    "    '''\n",
    "    hResetStatesCb(False/True) # reset states at on_epoch_begin (default) / end\n",
    "    '''\n",
    "    def __init__(self,End=False): # required?\n",
    "        self.End = End\n",
    "\n",
    "    def on_epoch_begin(self, epoch, logs={}):\n",
    "        if not self.End:\n",
    "            # HOWTO access model from callback:\n",
    "            self.model.reset_states()\n",
    "\n",
    "    # Doppeltgemoppelt haelt besser - for val_loss also?!\n",
    "    def on_epoch_end(self, epoch, logs={}):\n",
    "        if self.End:\n",
    "            # HOWTO access model from callback:\n",
    "            self.model.reset_states()\n",
    "\n",
    "# FITPAR['callbacks'] += [ hResetStatesCb() ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## hTBoardTextCb: custom TensorBoard callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: Each TensorBoard CB will create a separate events file (for each run).\n",
    "# To avoid duplicate logging of metrics etc., use only one TB callback in model.fit() \n",
    "\n",
    "# hTBoardTextCb v3b2: input {LogTag:LogStr,...} *OR* [(LogTag, LogStr),...]\n",
    "from keras.callbacks import TensorBoard\n",
    "# from keras.callbacks.tensorboard_v1 import TensorBoard # Get older version?!\n",
    "import tensorflow as tf\n",
    "\n",
    "hdict2list = lambda D: list( D.items() )\n",
    "\n",
    "class hTBoardTextCb(TensorBoard):\n",
    "    '''Tensorboard callback extended to log arbitrary text strings.\n",
    "    USE:\n",
    "    callbacks += [ hTBoardTextCb(log_dir, MyLogs, **kwargs)]\n",
    "    with MyLogs = {'Tag1':'Text1',...} or [('Tag1','Text1'),...]\n",
    "    \n",
    "    NOTE: Each TBoard callback creates a separate events.* file (for each run).\n",
    "    To avoid duplicate logging, use only one TB callback in model.fit()\n",
    "    by replacing:\n",
    "    callbacks += [keras.callbacks.TensorBoard(log_dir=TbDir,**kwargs)]\n",
    "    \n",
    "    Logs = [ ('Tag1', 'String1'), ('Tag2', 'String2'), ... ]\n",
    "    Could use an OrderedDict (from collections). But a simple pythetic \"dict\" does not preserve order.\n",
    "    \n",
    "    **{'batch_size':NB, 'histogram_freq':0, 'write_graph':False, 'write_images':False}\n",
    "    '''\n",
    "\n",
    "    def __init__(self, log_dir, MyLogs=None, **kwargs):\n",
    "        super().__init__(log_dir, **kwargs)\n",
    "        self.MyLogs = MyLogs\n",
    "\n",
    "    def on_train_begin(self, logs=None):\n",
    "        super().on_train_begin(logs)\n",
    "        if self.MyLogs is None:\n",
    "            return\n",
    "        try:\n",
    "            MyLogs = list(self.MyLogs.items()) # convert dict to list\n",
    "        except:\n",
    "            MyLogs = self.MyLogs\n",
    "\n",
    "        # Might consider: tf.summary.merge_all()\n",
    "        for TagText in MyLogs:\n",
    "            summary = tf.summary.text( TagText[0], tf.convert_to_tensor(TagText[1]) )\n",
    "            # https://www.tensorflow.org/api_docs/python/tf/summary/text\n",
    "\n",
    "            with  tf.Session() as sess:\n",
    "                # No need for this?: self.writer = tf.summary.FileWriter('./Tensorboard', sess.graph)\n",
    "                s = sess.run(summary)\n",
    "                self.writer.add_summary(s)\n",
    "\n",
    "        # Do we need sth like this?: self.writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# hLoadKm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "\n",
    "def hLoadProcDir(ProcDir, Npool=None):\n",
    "    '''\n",
    "    metapar.mat, model.json, model_BesVal.h5, train_log.csv\n",
    "    '''\n",
    "    if not isinstance( ProcDir,str):\n",
    "        with concurrent.futures.ProcessPoolExecutor(Npool) as pool: # pool size e.g. (8)?\n",
    "            Pars = pool.map( hLoadProcDir, ProcDir)\n",
    "        Pars = list(Pars)\n",
    "        return Pars\n",
    "\n",
    "    par = { 'ProcDir':ProcDir, 'Xtrafo':None }\n",
    "    \n",
    "    hdf5.loadmat( par['ProcDir']+'metapar.mat', par)\n",
    "    par['ProcDir'] = ProcDir\n",
    "    \n",
    "    par['TrainLog'] = pd.read_csv( par['ProcDir']+'train_log.csv')\n",
    "    \n",
    "    hdf5.loadmat( par['ProcDir']+'EvalData.mat', par)\n",
    "    par['ProcDir'] = ProcDir\n",
    "    \n",
    "    print('.',end='')\n",
    "    \n",
    "    return par\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "****************************************\n",
    "# EEG Sleep Scores\n",
    "Just divide the entire EEG duration into 30s intervals and assign them to the sleep scores.\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "ExTb['SscFile'] = '/data/mandelkowhc/Sleep/EEG/'+ExTb.ExId+'_sleepscoringmr.txt'\n",
    "ExTb['VmrkFile'] = '/data/mandelkowhc/Sleep/EEG/'+ExTb.ExId+'.vmrk'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def readVmrkStart( Fname ):\n",
    "    with open(Fname, mode='r') as F: T = F.read() # T is one string.\n",
    "    S = re.findall(r'^Mk1=New.*,(\\d+)$',T,re.M)[0]\n",
    "    return S"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "ExTb['VmrkT0'] = ExTb.VmrkFile.apply( lambda x: readVmrkStart(x) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def readVmrkR128( Fname ):\n",
    "    '''Rean MR triggers from EEG marker file .vmrk'''\n",
    "    with open(Fname, mode='r') as F: T = F.read() # T is one string.\n",
    "    S = re.findall(r'^Mk\\d+=Response,R128,(\\d+)',T,re.M)\n",
    "    S = list( map(int,S)) # int samples at 5kHz, 0.2ms?!\n",
    "    return S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def readSsc( SscFile, VmrkFile=None, Crop=None ):\n",
    "    '''\n",
    "    Returns pd.Series of length NTR with sleep score N+(0.1 if REM) and index= orig. sleepscore intervall number n*30s\n",
    "    '''\n",
    "    # Ssc = pd.read_csv( Fname, '\\t', usecols='Score', squeeze=True)\n",
    "    Ssc = pd.read_csv( SscFile, '\\t')\n",
    "    S = Ssc.Score + (Ssc.REM/10.0)\n",
    "    if VmrkFile is not None:\n",
    "        tmp = ( np.array(readVmrkR128(VmrkFile))/30/5000 ).astype(int) # =floor() to Ssc frames (30s)\n",
    "        S = S.reindex(index=tmp) # =interpolates\n",
    "    if Crop is not None:\n",
    "        S = S[ int(Crop): ]\n",
    "    return S"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "ExTb['SscMr'] = ExTb.apply( lambda x: readSsc(x.SscFile, x.VmrkFile, 3*60/3).values, 'columns')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python/3.8",
   "language": "python",
   "name": "py3.8"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}