multimodel.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b07111dc-e5da-47e3-a8eb-56a902b3b39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "threads = 16\n",
    "import os\n",
    "os.environ['OMP_NUM_THREADS']=str(threads)\n",
    "import tensorflow as tf\n",
    "\n",
    "# Tensorflow needs explicit cofig calls\n",
    "tf.config.threading.set_inter_op_parallelism_threads(threads)\n",
    "tf.config.threading.set_intra_op_parallelism_threads(threads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a8d6c9b-166c-4b51-a7cd-5042d2b0d5dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# input conformation\n",
    "#conf = \"alaninedipeptide_H.pdb\"\n",
    "conf = \"trpcage_correct.pdb\"\n",
    "\n",
    "# input trajectory\n",
    "# atom numbering must be consistent with {conf}\n",
    "\n",
    "#traj = \"alaninedipeptide_reduced.xtc\"\n",
    "traj = \"trpcage_red.xtc\"\n",
    "\n",
    "# input topology\n",
    "# expected to be produced with \n",
    "#    gmx pdb2gmx -f {conf} -p {topol} -n {index} -o {gro}\n",
    "\n",
    "# Gromacs changes atom numbering, the index file must be generated and used as well\n",
    "# gro file is used to generate inverse indexing for plumed.dat\n",
    "\n",
    "#topol = \"topol.top\"\n",
    "topol = \"topol_correct.top\"\n",
    "index = 'index_correct.ndx'\n",
    "gro = 'trpcage_correct.gro'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77a36713-f0bf-4225-b8d6-669b96fe7a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mdtraj as md\n",
    "import numpy as np\n",
    "import asmsa\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b5e1882-a39c-46d6-9d88-8538fbfbac91",
   "metadata": {},
   "outputs": [],
   "source": [
    "tr = md.load(traj,top=conf)\n",
    "idx=tr[0].top.select(\"name CA\")\n",
    "#idx=tr[0].top.select(\"element != H\")\n",
    "tr.superpose(tr[0],atom_indices=idx)\n",
    "geom = np.moveaxis(tr.xyz ,0,-1)\n",
    "geom = np.moveaxis(tr.xyz ,0,-1)\n",
    "geom.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "863073c4-e674-4313-9302-4db0e3b1d929",
   "metadata": {},
   "outputs": [],
   "source": [
    "density = 2 # integer in [1, n_atoms-1]\n",
    "\n",
    "sparse_dists = asmsa.NBDistancesSparse(geom.shape[0], density=density)\n",
    "mol = asmsa.Molecule(pdb=conf,top=topol,ndx=index,fms=[sparse_dists])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5699a0f-4a9a-4c10-84cc-e6360b135e90",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = mol.intcoord(geom).T\n",
    "X_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba5fb3d-244c-4d0c-9dd1-d80b68a7e41f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "454e69b0-7db5-4f6d-a547-e36242fac0cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "\n",
    "class AAEMultiModelNaive(keras.models.Model):\n",
    "    \n",
    "    def __init__(self,models):\n",
    "        super().__init__()\n",
    "        self.models = models\n",
    "        \n",
    "    def compile(self,*args,**kwargs):\n",
    "        super().compile()\n",
    "        for m in self.models:\n",
    "            m.compile(args,kwargs)\n",
    "    \n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        ael = [0]\n",
    "        dl = [0]\n",
    "        for m in self.models:\n",
    "            res = m.train_step(batch)\n",
    "#            ael.append(res['ae_loss'].numpy())\n",
    "#            dl.append(res['d_loss'].numpy())\n",
    "        \n",
    "        return { 'ae_loss_min': min(ael), 'ae_loss_max': max(ael), 'd_loss_min': min(dl), 'd_loss_max': max(dl) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e5d1c6-8ca3-4c7d-8af0-88f56d841d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "_default_hp = {\n",
    "        'batch_size' : 64,\n",
    "        'activation' : 'relu', \n",
    "        'ae_number_of_layers': 2,\n",
    "        'disc_number_of_layers': 2,\n",
    "        'ae_neuron_number_seed' : 32,\n",
    "        'disc_neuron_number_seed' : 32,\n",
    "        'ae_loss_fn': 'MeanSquaredError',\n",
    "        'disc_loss_fn': 'BinaryCrossentropy',\n",
    "        'optimizer': 'Adam',\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b79b94f1-f21a-44a9-bbd3-1becec1da75d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mods = []\n",
    "for ael in [2,3]:\n",
    "    for an in range(32,129,8):\n",
    "        hp = _default_hp.copy()\n",
    "        hp['ae_number_of_layers'] = ael\n",
    "        hp['ae_neuron_number_seed'] = an\n",
    "        mods.append(asmsa.AAEModel((X_train.shape[1],),hp=hp))\n",
    "        \n",
    "mmod = AAEMultiModelNaive(mods)\n",
    "mmod.compile()\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90ffd0fc-dcee-4415-a535-07592e74ae8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tf.data.Dataset.from_tensor_slices(X_train).shuffle(2048).batch(_default_hp['batch_size'],drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b706551-900b-4d77-b19c-baca52d0c77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmod.fit(ds,epochs=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15727a69-de7a-4cd6-89cd-999309a9115b",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(mods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178629a8-c3bc-4e20-b2d6-ac20f60b80e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mods[0].fit(ds,epochs=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af4fcd1-8056-44ce-a9ad-261cf58d1b7e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f71d002e-ac42-475c-9079-e8373c5a6443",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _compute_number_of_neurons(params,ae):\n",
    "        tmp = params['ae_neuron_number_seed' if ae else 'disc_neuron_number_seed']\n",
    "        neurons = [ tmp ]\n",
    "        \n",
    "        for _ in range(params['ae_number_of_layers'] if ae else params['disc_number_of_layers']):\n",
    "                tmp = int(tmp / 2)\n",
    "                neurons.append(tmp)\n",
    "        return neurons\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a871edd4-298f-43cd-b224-ff39cce6525a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AAEMultiModel(keras.models.Model):\n",
    "    \n",
    "    def __init__(self,molecule_shape,latent_dim=2,prior='normal',hp=_default_hp):\n",
    "        super().__init__()\n",
    "        self.inp = keras.Input(shape=molecule_shape,name='common.input')\n",
    "        latent = []\n",
    "        out = []\n",
    "        \n",
    "        # XXX\n",
    "        for an in range(32,129,8):\n",
    "            for r in range(4):\n",
    "                m = f\"seed{r}_{an}\"\n",
    "                neurons = _compute_number_of_neurons(hp,ae=True)\n",
    "                l = self.inp\n",
    "                i = 1\n",
    "                for n in neurons:\n",
    "                    l = keras.layers.Dense(n, activation=hp['activation'], name=f'enc.{m}.{i}')(l)\n",
    "                    l = keras.layers.BatchNormalization(momentum=0.8)(l)\n",
    "                    i += 1\n",
    "                l = keras.layers.Dense(latent_dim,activation='linear',name=f'enc.{m}.output')(l)\n",
    "                latent.append(l)\n",
    "\n",
    "                i = 1\n",
    "                for n in reversed(neurons):\n",
    "                    l = keras.layers.Dense(n, activation=hp['activation'], name=f'dec.{m}.{i}')(l)\n",
    "                    l = keras.layers.BatchNormalization(momentum=0.8)(l)\n",
    "                    i += 1\n",
    "                l = keras.layers.Dense(np.prod(molecule_shape), activation=hp['activation'],name=f'dec.{m}.output')(l)\n",
    "                l = keras.layers.Reshape(molecule_shape,name=f'dec.{m}.reshape')(l)\n",
    "                out.append(l)\n",
    "        \n",
    "        self.latent = tf.stack(latent,axis=1,name='all.latent')\n",
    "        self.out = tf.stack(out,axis=1,name='all.output')\n",
    "        self.n_models = len(out)\n",
    "        \n",
    "        self.encs = keras.Model(inputs=self.inp,outputs=self.latent)\n",
    "        self.aes = keras.Model(inputs=self.inp,outputs=self.out)\n",
    "            \n",
    "    def compile(self,*args,**kwargs):\n",
    "        super().compile(*args,**kwargs)\n",
    "        self.encs.compile(*args,**kwargs)\n",
    "        self.aes.compile(*args,**kwargs)\n",
    "        \n",
    "        # XXX\n",
    "        self.ae_loss = keras.losses.MeanSquaredError()\n",
    "        self.optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002)\n",
    "    \n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "            \n",
    "        multibatch = tf.stack([batch]*self.n_models,axis=1)\n",
    "#        print(multibatch)\n",
    "        with tf.GradientTape() as tape:\n",
    "            reconstruct = self.aes(batch)\n",
    "            ae_loss = tf.reduce_sum(keras.metrics.mean_squared_error(multibatch,reconstruct),axis=0)\n",
    "\n",
    "        ae_grad = tape.gradient(ae_loss,self.aes.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.aes.trainable_weights))\n",
    "        \n",
    "        return { 'ae min': tf.reduce_min(ae_loss), 'ae max' : tf.reduce_max(ae_loss) }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d414877-57d7-4e6f-812b-4da0017eb25e",
   "metadata": {},
   "outputs": [],
   "source": [
    "_compute_number_of_neurons(_default_hp,True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2d9d99-e280-49e0-97dc-bdf5daa37f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = AAEMultiModel((X_train.shape[1],))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "112d9ea0-3144-4218-87a6-b739719576a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf09fb82-ef07-42ec-b1f9-79b224e405cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51480d46-f25a-4a2f-b6c5-c860d691b6bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.fit(ds,epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f74f73b0-6513-4355-a628-e55bd299454b",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(list(range(32,129,8)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97c8c814-bf59-47bc-8a3d-0e74e4a4f58c",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.aes.trainable_weights[26]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c9e9a9a-b705-49d1-b7f2-ee6cf7ebd0f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2aebfd59-e34c-4c87-9267-84abd68ae529",
   "metadata": {},
   "outputs": [],
   "source": [
    "364/13"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4deca71-5f41-4908-b9ce-3a6c9b5ebc88",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AAEMultiModel2(keras.models.Model):\n",
    "    \n",
    "    def __init__(self,molecule_shape,latent_dim=2,prior='normal',hp=_default_hp):\n",
    "        super().__init__()\n",
    "        self.inp = keras.Input(shape=molecule_shape,name='common.input')\n",
    "        latent = []\n",
    "        out = []\n",
    "        self.aes = []\n",
    "        \n",
    "        # XXX\n",
    "        for an in range(32,129,8):\n",
    "            for r in range(4):\n",
    "                m = f\"seed{r}_{an}\"\n",
    "                neurons = _compute_number_of_neurons(hp,ae=True)\n",
    "                l = self.inp\n",
    "                i = 1\n",
    "                for n in neurons:\n",
    "                    l = keras.layers.Dense(n, activation=hp['activation'], name=f'enc.{m}.{i}')(l)\n",
    "                    l = keras.layers.BatchNormalization(momentum=0.8)(l)\n",
    "                    i += 1\n",
    "                l = keras.layers.Dense(latent_dim,activation='linear',name=f'enc.{m}.output')(l)\n",
    "                latent.append(l)\n",
    "\n",
    "                i = 1\n",
    "                for n in reversed(neurons):\n",
    "                    l = keras.layers.Dense(n, activation=hp['activation'], name=f'dec.{m}.{i}')(l)\n",
    "                    l = keras.layers.BatchNormalization(momentum=0.8)(l)\n",
    "                    i += 1\n",
    "                l = keras.layers.Dense(np.prod(molecule_shape), activation=hp['activation'],name=f'dec.{m}.output')(l)\n",
    "                l = keras.layers.Reshape(molecule_shape,name=f'dec.{m}.reshape')(l)\n",
    "                out.append(l)\n",
    "                self.aes.append(keras.Model(inputs=self.inp,outputs=l))\n",
    "        \n",
    "        self.latent = tf.stack(latent,axis=1,name='all.latent')\n",
    "        self.out = tf.stack(out,axis=1,name='all.output')        \n",
    "        self.all_aes = keras.Model(inputs=self.inp,outputs=self.out)\n",
    "            \n",
    "    def compile(self,*args,**kwargs):\n",
    "        super().compile(*args,**kwargs)\n",
    "\n",
    "        for m in self.aes:\n",
    "            m.compile(*args,**kwargs)\n",
    "        \n",
    "        # XXX\n",
    "        self.ae_loss = keras.losses.MeanSquaredError()\n",
    "        self.optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002)\n",
    "    \n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "            \n",
    "        multibatch = tf.stack([batch]*len(self.aes),axis=1)\n",
    "#        print(multibatch)\n",
    "        with tf.GradientTape(persistent=True) as tape:\n",
    "            reconstruct = self.all_aes(batch)\n",
    "            ae_loss = tf.reduce_sum(keras.metrics.mean_squared_error(multibatch,reconstruct),axis=0)\n",
    "\n",
    "        for m in self.aes:\n",
    "            ae_grad = tape.gradient(ae_loss,m.trainable_weights)\n",
    "            self.optimizer.apply_gradients(zip(ae_grad,m.trainable_weights))\n",
    "        \n",
    "        return { 'ae min': tf.reduce_min(ae_loss), 'ae max' : tf.reduce_max(ae_loss) }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26538595-5003-4ba9-9cff-3971d14557c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = AAEMultiModel2((X_train.shape[1],)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6c7ba94-adb3-41fd-87d1-3593c9527ca7",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "267e3d48-1654-4626-a39d-108561c18435",
   "metadata": {},
   "outputs": [],
   "source": [
    "m.fit(ds,epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e034cf0-0261-4855-a706-074d010730ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "inp = keras.Input((X_train.shape[1],))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f196f1-b6c1-4f09-abde-5f527d75f2c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = inp\n",
    "for _ in range(4):\n",
    "    out = keras.layers.BatchNormalization(momentum=0.8)(\n",
    "        keras.layers.Dense(850,activation='relu')(out)\n",
    "    )\n",
    "out = keras.layers.Dense(X_train.shape[1],)(out)\n",
    "\n",
    "class TestModel(keras.models.Model):\n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "\n",
    "        with tf.GradientTape(persistent=True) as tape:\n",
    "            reconstruct = self.all_aes(batch)\n",
    "            ae_loss = keras.metrics.mean_squared_error(batch,reconstruct)\n",
    "\n",
    "        ae_grad = tape.gradient(ae_loss,self.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.trainable_weights))\n",
    "        return ae_loss\n",
    "\n",
    "\n",
    "    \n",
    "silim = keras.Model(inputs=inp,outputs=out)\n",
    "silim.compile(loss=keras.losses.MeanSquaredError(),optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ab0f41d-0550-4409-871c-7338bb911513",
   "metadata": {},
   "outputs": [],
   "source": [
    "silim.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb9e79f5-e523-43ea-802c-c2b700a5efe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "silim.fit(X_train,X_train,epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54d75c49-6cac-4ce5-9e32-6b91c4cc10e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "inp=keras.Input((5,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45c4c58-2852-4b93-8944-ed4413c2b381",
   "metadata": {},
   "outputs": [],
   "source": [
    "l1 = keras.layers.Dense(10,activation='relu')(inp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3259f9a8-24be-4992-afe9-0ea72ca954d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "l2 = keras.layers.Dense(10,activation='relu')(l1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d3b63b-3fe2-4c44-8ae7-068f9688e6a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "l2.trainable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdb0f272-e095-47e3-b240-0aa9bc5ceb2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "l2.trainable_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e521a0f-459b-4bdd-9f3e-05a99f3103bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "tm = keras.Model(inputs=inp,outputs=l2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa868633-3870-4a06-80bc-e5870e625949",
   "metadata": {},
   "outputs": [],
   "source": [
    "l2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c1238d6-6655-421b-bffd-25378b38eb99",
   "metadata": {},
   "outputs": [],
   "source": [
    "tm.layers[1].trainable_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469d3682-51c9-4de0-874e-7bd8ddace472",
   "metadata": {},
   "outputs": [],
   "source": [
    "l2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e3134e-b1ff-4a38-b9e0-355cc64885d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tm.layers[2].name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e74806f-418d-4133-9755-0979f2e0c222",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AAEMultiModel3(keras.models.Model):\n",
    "    # enc_neurons: (models,layers)\n",
    "    def __init__(self,inp_shape,enc_neurons,disc_neurons,latent_dim=2):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_models = enc_neurons.shape[0]\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        assert disc_neurons.shape[0] == self.n_models\n",
    "        \n",
    "        inp = keras.Input(shape=inp_shape)\n",
    "        out = inp\n",
    "        # TODO: emulate \"empty\" layers\n",
    "        for num in range(enc_neurons.shape[1]):\n",
    "            print(f'enc_{num} {enc_neurons[:,num]}')\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',name=f'enc_{num}')(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_{num}')(out)\n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*latent_dim,name='enc_out')(out) # \n",
    "        latent = out\n",
    "        \n",
    "        # decoder layers are numbered in reverse so that neuron numbers match with encoder\n",
    "        for num  in reversed(range(enc_neurons.shape[1])):\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',name=f'dec_{num}')(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'dec_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*inp_shape[0],name='dec_out')(out)\n",
    "        out = keras.layers.Reshape((enc_neurons.shape[0],inp_shape[0]))(out)\n",
    "        \n",
    "        self.aes = keras.Model(inputs=inp,outputs=[out,latent])\n",
    "        self.enc = keras.Model(inputs=inp,outputs=latent)\n",
    "        self.dec = keras.Model(inputs=latent,outputs=out)\n",
    "        \n",
    "        inp = keras.Input(shape=(latent_dim * self.n_models,))\n",
    "        disc = inp\n",
    "        for num in range(disc_neurons.shape[1]):\n",
    "            disc = keras.layers.Dense(np.sum(disc_neurons[:,num]),name=f'disc_{num}')(disc)\n",
    "            disc = keras.layers.LeakyReLU(alpha=0.2,name=f'disc_relu_{num}')(disc)\n",
    "            \n",
    "        disc = keras.layers.Dense(disc_neurons.shape[0],name='disc_out')(disc)\n",
    "        \n",
    "        self.disc = keras.Model(inputs=inp,outputs=disc)\n",
    "        \n",
    "        self.masks = {}\n",
    "        \n",
    "        # masks for encoder layers 1:, decoder layers 0: are just .T's\n",
    "        idx = np.concatenate((np.zeros((1,enc_neurons.shape[1]),dtype=np.int32),np.cumsum(enc_neurons,axis=0)))\n",
    "        for layer in range(1, enc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(enc_neurons[:,layer-1]),np.sum(enc_neurons[:,layer])),dtype=np.float32)\n",
    "            for mod in range(enc_neurons.shape[0]):\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'enc_{layer}'] = tf.convert_to_tensor(mask)\n",
    "            self.masks[f'dec_{layer-1}'] = tf.convert_to_tensor(mask.T)\n",
    "        \n",
    "        # mask to and from latent layer\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,-1]),latent_dim * enc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(enc_neurons.shape[0]):\n",
    "            mask[idx[mod,-1]:idx[mod+1,-1],mod*latent_dim:(mod+1)*latent_dim] = 1.\n",
    "        \n",
    "        self.masks['enc_out'] = tf.convert_to_tensor(mask)\n",
    "        self.masks[f'dec_{enc_neurons.shape[1]-1}'] = tf.convert_to_tensor(mask.T)\n",
    "        \n",
    "        # mask for decoder output\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,0]),inp_shape[0]*enc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(enc_neurons.shape[0]):\n",
    "            mask[idx[mod,0]:idx[mod+1,0], mod*inp_shape[0]:(mod+1)*inp_shape[0]] = 1.\n",
    "        self.masks['dec_out'] = tf.convert_to_tensor(mask)\n",
    "             \n",
    "        idx = np.concatenate((np.zeros((1,disc_neurons.shape[1]),dtype=np.int32),np.cumsum(disc_neurons,axis=0)))\n",
    "        # mask for discriminator layer 0:\n",
    "        mask = np.zeros((latent_dim * disc_neurons.shape[0],np.sum(disc_neurons[:,0])),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[mod*latent_dim:(mod+1)*latent_dim, idx[mod,0]:idx[mod+1,0]] = 1.\n",
    "        self.masks['disc_0'] = tf.convert_to_tensor(mask)\n",
    "            \n",
    "        # mask for discriminator layers 1:\n",
    "        for layer in range(1,disc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(disc_neurons[:,layer-1]),np.sum(disc_neurons[:,layer])),dtype=np.float32)\n",
    "            for mod in range(disc_neurons.shape[0]):\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'disc_{layer}'] = tf.convert_to_tensor(mask)\n",
    "            \n",
    "        # mask for discriminator output:\n",
    "        mask = np.zeros((np.sum(disc_neurons[:,-1]),disc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[idx[mod,disc_neurons.shape[1]-1]:idx[mod+1,disc_neurons.shape[1]-1],mod] = 1.\n",
    "        self.masks['disc_out'] = tf.convert_to_tensor(mask)\n",
    "\n",
    "    def compile(self):\n",
    "        # XXX\n",
    "        super().compile(optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002))\n",
    "        self.ae_weights = self.enc.trainable_weights + self.dec.trainable_weights\n",
    "         \n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "        \n",
    "        # multiple models need replicated batch to compute loss simultaneously\n",
    "        multibatch = tf.stack([batch]*self.n_models,axis=1)\n",
    " \n",
    "        with tf.GradientTape() as tape:\n",
    "            reconstruct = self.aes(batch)\n",
    "            ae_multiloss = tf.reduce_mean(keras.metrics.mean_squared_error(multibatch,reconstruct[0]),axis=0)\n",
    "            ae_loss = tf.reduce_sum(ae_multiloss)\n",
    "               \n",
    "        ae_grad = tape.gradient(ae_loss,self.ae_weights)\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.ae_weights))\n",
    "        \n",
    "        for l in self.aes.layers:\n",
    "            if l.name in self.masks:\n",
    "#                print(l.name, self.masks[l.name])\n",
    "                l.weights[0].assign(l.weights[0] * self.masks[l.name])\n",
    "       \n",
    "#        rand_low = tf.random.normal((batch.shape[0],self.latent_dim * self.n_models))\n",
    "        rand_low = tf.random.normal((batch.shape[0],self.latent_dim))\n",
    "#        rand_low = tf.random.uniform((batch.shape[0],self.latent_dim))\n",
    "        rand_low = tf.repeat(rand_low,self.n_models,axis=1)\n",
    "        lows = tf.concat([rand_low, reconstruct[1]],axis=0)\n",
    "        labels = tf.concat([tf.ones((batch.shape[0],self.n_models)), tf.zeros((batch.shape[0],self.n_models))], axis=0)\n",
    "        labels += 0.05 * tf.random.uniform(labels.shape)\n",
    "                        \n",
    "        with tf.GradientTape() as tape:\n",
    "            pred = self.disc(lows)\n",
    "            disc_losses = tf.keras.metrics.binary_crossentropy(labels,pred,from_logits=True,axis=0)\n",
    "            disc_loss = tf.reduce_sum(disc_losses)\n",
    "            \n",
    "        disc_grads = tape.gradient(disc_loss,self.disc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(disc_grads,self.disc.trainable_weights))\n",
    "        \n",
    "        for l in self.disc.layers:\n",
    "            if l.name in self.masks:\n",
    "#                print(l.name, self.masks[l.name])\n",
    "                l.weights[0].assign(l.weights[0] * self.masks[l.name])\n",
    "  \n",
    "        all_true = tf.ones((batch.shape[0],self.n_models))\n",
    "    \n",
    "        with tf.GradientTape() as tape:\n",
    "            cheat = self.disc(self.enc(batch))\n",
    "            cheat_losses = tf.keras.metrics.binary_crossentropy(all_true,cheat,from_logits=True,axis=0)\n",
    "            cheat_loss = tf.reduce_sum(cheat_losses)\n",
    "            \n",
    "        cheat_grads = tape.gradient(cheat_loss,self.enc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(cheat_grads,self.enc.trainable_weights))\n",
    "        \n",
    "        for l in self.enc.layers:\n",
    "            if l.name in self.masks:\n",
    "#                print(l.name, self.masks[l.name])\n",
    "                l.weights[0].assign(l.weights[0] * self.masks[l.name])\n",
    "       \n",
    "\n",
    "        return {\n",
    "            'AE loss min' : tf.reduce_min(ae_multiloss),\n",
    "            'AE loss max' : tf.reduce_max(ae_multiloss),\n",
    "            'disc loss min' : tf.reduce_min(disc_losses),\n",
    "            'disc loss max' : tf.reduce_max(disc_losses),\n",
    "            'cheat loss min' : tf.reduce_min(cheat_losses),\n",
    "            'cheat loss max' : tf.reduce_max(cheat_losses)\n",
    "        }\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self,inp):\n",
    "        return self.aes(inp)\n",
    "                            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36ca5b8-34f5-4b8b-a80a-19afcfc6e747",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch=np.array([[1,2,3],[4,5,6]])\n",
    "multibatch = np.stack([batch]*2,axis=1)\n",
    "multibatch # (batch, model, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "215eff5e-0cf5-4b41-9d3a-92a9d6974a51",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.convert_to_tensor(multibatch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d53c56-e1f7-4c8a-b4d9-2ae702d1b12d",
   "metadata": {},
   "outputs": [],
   "source": [
    "result=np.stack([batch, np.array([[1.1,2.1,2.9],[4.5,5.5,6]])],axis=1)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2eed57-83d3-418b-9f59-eb95280717aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "fr = np.array([[1,2,3,1.1,2.1,2.9],[4,5,6,4.5,5.5,6]])\n",
    "r2 = np.reshape(fr,(2,2,3))\n",
    "r2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c665fd9d-6f6b-4a04-b4c0-158c7605c0ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "metric=keras.metrics.mean_squared_error(multibatch,r2)\n",
    "metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1bcf332-ce9e-4a50-9c36-6af350cf7c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean = tf.reduce_mean(metric,axis=0)\n",
    "mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc1ac549-ad3e-4675-a903-f0b94f5372fb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54746d0-150d-43e0-a7fb-2739ffd91596",
   "metadata": {},
   "outputs": [],
   "source": [
    "en = np.array([[32, 16, 8]]*10)\n",
    "\n",
    "m3 = AAEMultiModel3((X_train.shape[1],),en)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd0b449-b354-4654-995e-3e5bb93fdbd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3(X_train[0:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a98b703-a7b7-4e2e-bd38-14f6f68c69c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c00df80-18c5-4673-be9c-597a6915c2a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3.fit(ds,epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc24b37-614e-41ca-ae10-6eca039976a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3.aes.layers[3].get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0697f5e9-bdaf-4bf5-b0d5-bd5775f9b91b",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93b11be-aed2-474b-a317-3bc7cfa1231a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "048af077-69f1-4fef-b7e8-f4922c84802a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tf.data.Dataset.from_tensor_slices(X_train).shuffle(2048).batch(256,drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0967a2e-64a5-4dae-8f1c-42d0204b0e6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "en = []\n",
    "for an in range(32,129,8):\n",
    "    hp = _default_hp.copy()\n",
    "    hp['ae_neuron_number_seed'] = an\n",
    "    en.append(_compute_number_of_neurons(hp,True))\n",
    "    \n",
    "en = np.array(en)\n",
    "\n",
    "dn = np.array([[ 32, 64, 32 ]] * en.shape[0])\n",
    "dn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58e9c653-6ef2-4b69-af02-b867a65b9f71",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3a = AAEMultiModel3((X_train.shape[1],),en,dn)\n",
    "m3a.compile()\n",
    "m3a.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "964691a2-b762-40c2-a805-e5b5ca5d5f57",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3a.fit(ds,epochs=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d66008a-57c4-4dba-8e3c-d203e126ab0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70dfd69-3549-4ac7-9a24-50a485dcab14",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m3a.enc(X_train).numpy()\n",
    "lows.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd5be4a-c3af-4d72-bc88-df39f9d7d2d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "for mod in range(lows.shape[1]//2):\n",
    "    plows = lows[::200,mod*2:(mod+1)*2]\n",
    "    plt.scatter(plows[:,0],plows[:,1],marker='.',label=str(mod))\n",
    "\n",
    "lim=.05\n",
    "plt.legend()    \n",
    "plt.xlim((-lim,lim))\n",
    "plt.ylim((-lim,lim))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d75c354-5254-4394-8a33-8d448230d35e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "992947aa-a3bb-47bf-b301-a84ed4bd2fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "en = []\n",
    "dn = []\n",
    "hp = _default_hp.copy()\n",
    "\n",
    "for an in range(32,129,32):\n",
    "    hp['ae_neuron_number_seed'] = an\n",
    "    en.append(_compute_number_of_neurons(hp,True))\n",
    "    \n",
    "for an in range(64,257,64):\n",
    "    hp['disc_neuron_number_seed'] = an\n",
    "    dn.append(_compute_number_of_neurons(hp,False))\n",
    "\n",
    "el = len(en)\n",
    "dl = len(dn)\n",
    "en = np.array(en * dl)\n",
    "dn = np.repeat(np.array(dn),el,axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb281b2-c365-46f7-ae1d-9777d8a7a654",
   "metadata": {},
   "outputs": [],
   "source": [
    "en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcc3331d-98ad-4c19-865e-2a69800035b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0751be01-5ba3-4f7a-ba05-6ce1307a878f",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3b = AAEMultiModel3((X_train.shape[1],),en,dn)\n",
    "m3b.compile()\n",
    "m3b.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81f5d8be-3227-4fe6-b739-71e011d8e2b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3b.fit(ds,epochs=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd79b864-472b-4ba6-9d19-efde523777ea",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d357be1-2cbb-47bd-8041-0abd1ec75d71",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m3b.enc(X_train).numpy()\n",
    "lows.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de57449-630d-4d8f-aa69-c631c3ea878d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12,9))\n",
    "for mod in range(lows.shape[1]//2):\n",
    "    plows = lows[::100,mod*2:(mod+1)*2]\n",
    "    plt.scatter(plows[:,0],plows[:,1],marker='.',label=str(mod))\n",
    "\n",
    "lim=.04\n",
    "plt.legend()    \n",
    "plt.xlim((-lim,lim))\n",
    "plt.ylim((-lim,lim))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5ced1ff-8c52-412f-b744-ac0118e44636",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3c = AAEMultiModel3((X_train.shape[1],),np.array([[96,48],[96,48]]),np.array([[224,112],[64,32]]))\n",
    "m3c.compile()\n",
    "m3c.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee5523e1-5271-45fa-921a-e322352b8932",
   "metadata": {},
   "outputs": [],
   "source": [
    "m3c.fit(ds,epochs=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf31030b-b5c6-4a04-b49f-bab7c1eac19a",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m3c.enc(X_train).numpy()\n",
    "lows.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e781e5f-fda7-42b0-954f-14906dd90cdc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aed0af09-6295-4515-9f03-bc281fbbbfbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6182fd2-3735-4e90-9899-d73e881c6589",
   "metadata": {},
   "outputs": [],
   "source": [
    "class _Sparse(keras.constraints.Constraint):\n",
    "    def __init__(self,mask):\n",
    "        self.mask = tf.convert_to_tensor(mask)\n",
    "\n",
    "    @tf.function    \n",
    "    def __call__(self,w):\n",
    "        return w * self.mask\n",
    "    \n",
    "\n",
    "class AAEMultiModel4(keras.models.Model):\n",
    "    # enc_neurons: (models,layers)\n",
    "    \n",
    "    def __init__(self,inp_shape,enc_neurons,disc_neurons,latent_dim=2):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_models = enc_neurons.shape[0]\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        assert disc_neurons.shape[0] == self.n_models\n",
    "        \n",
    "        self.masks = {}\n",
    "        \n",
    "        # masks for encoder layers 1:, decoder layers 0: are just .T's\n",
    "        idx = np.concatenate((np.zeros((1,enc_neurons.shape[1]),dtype=np.int32),np.cumsum(enc_neurons,axis=0)))\n",
    "        for layer in range(1, enc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(enc_neurons[:,layer-1]),np.sum(enc_neurons[:,layer])),dtype=np.float32)\n",
    "            for mod in range(enc_neurons.shape[0]):\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'enc_{layer}'] = _Sparse(mask)\n",
    "            self.masks[f'dec_{layer-1}'] = _Sparse(mask.T)\n",
    "        \n",
    "        # mask to and from latent layer\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,-1]),latent_dim * enc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(enc_neurons.shape[0]):\n",
    "            mask[idx[mod,-1]:idx[mod+1,-1],mod*latent_dim:(mod+1)*latent_dim] = 1.\n",
    "        \n",
    "        self.masks['enc_out'] = _Sparse(mask)\n",
    "        self.masks[f'dec_{enc_neurons.shape[1]-1}'] = _Sparse(mask.T)\n",
    "        \n",
    "        # mask for decoder output\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,0]),inp_shape[0]*enc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(enc_neurons.shape[0]):\n",
    "            mask[idx[mod,0]:idx[mod+1,0], mod*inp_shape[0]:(mod+1)*inp_shape[0]] = 1.\n",
    "        self.masks['dec_out'] = _Sparse(mask)\n",
    "             \n",
    "        idx = np.concatenate((np.zeros((1,disc_neurons.shape[1]),dtype=np.int32),np.cumsum(disc_neurons,axis=0)))\n",
    "        # mask for discriminator layer 0:\n",
    "        mask = np.zeros((latent_dim * disc_neurons.shape[0],np.sum(disc_neurons[:,0])),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[mod*latent_dim:(mod+1)*latent_dim, idx[mod,0]:idx[mod+1,0]] = 1.\n",
    "        self.masks['disc_0'] = _Sparse(mask)\n",
    "#        print('disc_0',mask)\n",
    "            \n",
    "        # mask for discriminator layers 1:\n",
    "        for layer in range(1,disc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(disc_neurons[:,layer-1]),np.sum(disc_neurons[:,layer])),dtype=np.float32)\n",
    "            for mod in range(disc_neurons.shape[0]):\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'disc_{layer}'] = _Sparse(mask)\n",
    "#            print(f'disc_{layer}',mask)\n",
    "            \n",
    "        # mask for discriminator output:\n",
    "        mask = np.zeros((np.sum(disc_neurons[:,-1]),disc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[idx[mod,disc_neurons.shape[1]-1]:idx[mod+1,disc_neurons.shape[1]-1],mod] = 1.\n",
    "        self.masks['disc_out'] = _Sparse(mask)  \n",
    "#        print('disc_out',mask)\n",
    "\n",
    "        \n",
    "        inp = keras.Input(shape=inp_shape)\n",
    "        out = inp\n",
    "\n",
    "        # TODO: emulate \"empty\" layers\n",
    "        for num in range(enc_neurons.shape[1]):\n",
    "            name = f'enc_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',\n",
    "                                     name = name, kernel_constraint = self.masks.get(name))(out)\n",
    "#            out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*latent_dim,name='enc_out', kernel_constraint = self.masks.get('enc_out'))(out) # \n",
    "        latent = out\n",
    "        \n",
    "        # decoder layers are numbered in reverse so that neuron numbers match with encoder\n",
    "        for num  in reversed(range(enc_neurons.shape[1])):\n",
    "            name = f'dec_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',\n",
    "                                     name=name,kernel_constraint=self.masks.get(name))(out)\n",
    "#            out = keras.layers.BatchNormalization(momentum=0.8,name=f'dec_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*inp_shape[0],name='dec_out',kernel_constraint=self.masks.get('dec_out'))(out)\n",
    "        out = keras.layers.Reshape((enc_neurons.shape[0],inp_shape[0]))(out)\n",
    "        \n",
    "        self.aes = keras.Model(inputs=inp,outputs=[out,latent])\n",
    "        self.enc = keras.Model(inputs=inp,outputs=latent)\n",
    "        self.dec = keras.Model(inputs=latent,outputs=out)\n",
    "        \n",
    "        inp = keras.Input(shape=(latent_dim * self.n_models,))\n",
    "        disc = inp\n",
    "        for num in range(disc_neurons.shape[1]):\n",
    "            name = f'disc_{num}'\n",
    "            disc = keras.layers.Dense(np.sum(disc_neurons[:,num]),\n",
    "                                      name=name,kernel_constraint=self.masks.get(name))(disc)\n",
    "            disc = keras.layers.LeakyReLU(alpha=0.2,name=f'disc_relu_{num}')(disc)\n",
    "            \n",
    "        disc = keras.layers.Dense(disc_neurons.shape[0],name='disc_out',kernel_constraint=self.masks.get('disc_out'))(disc)\n",
    "        \n",
    "        self.disc = keras.Model(inputs=inp,outputs=disc)\n",
    "        \n",
    " \n",
    "\n",
    "    def compile(self):\n",
    "        # XXX\n",
    "        super().compile(optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002))\n",
    "        self.ae_weights = self.enc.trainable_weights + self.dec.trainable_weights\n",
    "         \n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "        \n",
    "        # multiple models need replicated batch to compute loss simultaneously\n",
    "        multibatch = tf.stack([batch]*self.n_models,axis=1)\n",
    " \n",
    "        with tf.GradientTape() as tape:\n",
    "            reconstruct = self.aes(batch)\n",
    "            ae_multiloss = tf.reduce_mean(keras.metrics.mean_squared_error(multibatch,reconstruct[0]),axis=0)\n",
    "            ae_loss = tf.reduce_sum(ae_multiloss)\n",
    "               \n",
    "        ae_grad = tape.gradient(ae_loss,self.ae_weights)\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.ae_weights))\n",
    "               \n",
    "#        rand_low = tf.random.normal((batch.shape[0],self.latent_dim * self.n_models))\n",
    "        rand_low = tf.random.normal((batch.shape[0],self.latent_dim))\n",
    "#        rand_low = tf.random.uniform((batch.shape[0],self.latent_dim))\n",
    "        rand_low = tf.tile(rand_low,(1,self.n_models))\n",
    "        lows = tf.concat([rand_low, reconstruct[1]],axis=0)\n",
    "        labels = tf.concat([tf.ones((batch.shape[0],self.n_models)), tf.zeros((batch.shape[0],self.n_models))], axis=0)\n",
    "        labels += 0.05 * tf.random.uniform(labels.shape)\n",
    "                        \n",
    "        with tf.GradientTape() as tape:\n",
    "            \n",
    "            # FIXME: perturbe\n",
    "            neg_pred = self.disc(reconstruct[1])\n",
    "            neg_losses = tf.reduce_sum(neg_pred,axis=0) \n",
    "            pos_pred = self.disc(rand_low)\n",
    "            pos_losses = -tf.reduce_sum(pos_pred,axis=0)\n",
    "            disc_losses = neg_losses + pos_losses\n",
    "            disc_loss = tf.reduce_mean(disc_losses)\n",
    "            \n",
    "        disc_grads = tape.gradient(disc_loss,self.disc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(disc_grads,self.disc.trainable_weights))\n",
    "          \n",
    "        all_true = tf.ones((batch.shape[0],self.n_models))\n",
    "    \n",
    "        with tf.GradientTape() as tape:\n",
    "            cheat = self.disc(self.enc(batch))\n",
    "            cheat_losses = -tf.reduce_sum(cheat,axis=0)\n",
    "            cheat_loss = tf.reduce_mean(cheat_losses)\n",
    "            \n",
    "        cheat_grads = tape.gradient(cheat_loss,self.enc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(cheat_grads,self.enc.trainable_weights))\n",
    "        \n",
    "#        return {\n",
    "#            'AE loss min' : tf.reduce_min(ae_multiloss),\n",
    "#            'AE loss max' : tf.reduce_max(ae_multiloss),\n",
    "#            'disc loss min' : tf.reduce_min(disc_losses),\n",
    "#            'disc loss max' : tf.reduce_max(disc_losses),\n",
    "#            'cheat loss min' : tf.reduce_min(cheat_losses),\n",
    "#            'cheat loss max' : tf.reduce_max(cheat_losses)\n",
    "#        }\n",
    "\n",
    "        return { str(i): disc_losses[i] for i in range(disc_losses.shape[0]) }\n",
    "        #return { '0': disc_losses[0], '1' : disc_losses[1] }\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self,inp):\n",
    "        return self.aes(inp)\n",
    "                            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73add112-7e22-474d-b234-15cb13abc6ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "m4 = AAEMultiModel4((X_train.shape[1],),np.array([[96,48],[96,48]]*10),np.array([[224,112],[64,32]]*10))\n",
    "#m4 = AAEMultiModel4((X_train.shape[1],),np.array([[96,48],[96,48]]),np.array([[224,112],[64,32]]))\n",
    "m4.compile()\n",
    "m4.aes.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c509a1f6-d320-4107-b162-89e1946a1640",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = m4.fit(ds,epochs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a471827-9f54-4218-9fb8-fcaf41ec8f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "f.history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "007035a1-bb5c-423a-9f7b-9fd197d65ffa",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m4.enc(X_train).numpy()\n",
    "\n",
    "plt.figure(figsize=(12,9))\n",
    "for mod in range(lows.shape[1]//2):\n",
    "    plows = lows[::10,mod*2:(mod+1)*2]\n",
    "    plt.scatter(plows[:,0],plows[:,1],marker='.',label=str(mod))\n",
    "\n",
    "lim=5\n",
    "plt.legend()    \n",
    "plt.xlim((-lim,lim))\n",
    "plt.ylim((-lim,lim))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96f07890-be7c-45d4-9e72-91d0ac4e8334",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00a9c4ba-b8ea-492b-bb3a-3962376f2d34",
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_low = tf.random.normal((64,2))\n",
    "rand_low = tf.tile(rand_low,(1,m4.n_models))\n",
    "rand_low"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afa3dfdc-adfa-4a36-beb5-6c41b9d61c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = m4.disc(rand_low)\n",
    "tf.reduce_sum(pred,axis = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2bfe7ed-cc6e-406e-8109-229be6b86880",
   "metadata": {},
   "outputs": [],
   "source": [
    "r2 = tf.tile(tf.random.normal((5000,2),stddev=1.),(1,m4.n_models))\n",
    "p2 = m4.disc(r2)\n",
    "tf.reduce_sum(p2,axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a60b1fd-5e21-4e3c-b1dc-1edcf467364d",
   "metadata": {},
   "outputs": [],
   "source": [
    "d2 = m4.disc(lows[:5000,:])\n",
    "tf.reduce_sum(d2,axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e95a20a-d076-4ff6-8596-83cd7f446ddc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b76c6fdf-9416-4740-8f25-47373f54057d",
   "metadata": {},
   "outputs": [],
   "source": [
    "_random_init = keras.initializers.RandomUnivorm(seed=42,minval=-1.,maxval=1.)\n",
    "\n",
    "class _SparseConstraint(keras.constraints.Constraint):\n",
    "    def __init__(self,mask):\n",
    "        self.mask = mask\n",
    "        \n",
    "    @tf.function    \n",
    "    def __call__(self,w):\n",
    "        return w * self.mask\n",
    "    \n",
    "class _SparseInitializer(keras.initializers.Initializer):\n",
    "    def __init__(self,mask):\n",
    "        self.mask = mask\n",
    "    \n",
    "    def __call__(self,shape,dtype):\n",
    "        return _random_init(shape=shape,dtype=dtype) * self.mask\n",
    "        \n",
    "    \n",
    "class _Sparse():\n",
    "    def __init__(self,mask):\n",
    "        tmask = tf.convert_to_tensor(mask)\n",
    "        self.i = _SparseInitializer(tmask)\n",
    "        self.c = _SparseConstraint(tmask)\n",
    "\n",
    "class AAEMultiModel5(keras.models.Model):\n",
    "    # enc_neurons: (models,layers)\n",
    "    \n",
    "    def __init__(self,inp_shape,enc_neurons,disc_neurons,latent_dim=2):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_models = enc_neurons.shape[0]\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        assert disc_neurons.shape[0] == self.n_models\n",
    "        \n",
    "        self.masks = {}\n",
    "        \n",
    "        # XXX: input to first model only, hack\n",
    "        mask = np.zeros((inp_shape[0],np.sum(enc_neurons[:,0])),dtype=np.float32)\n",
    "        mask[:,:enc_neurons[0,0]] = 1.\n",
    "        self.masks['enc_0'] = _Sparse(mask)\n",
    "        \n",
    "        # masks for encoder layers 1:, decoder layers 0: are just .T's\n",
    "        idx = np.concatenate((np.zeros((1,enc_neurons.shape[1]),dtype=np.int32),np.cumsum(enc_neurons,axis=0)))\n",
    "        for layer in range(1, enc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(enc_neurons[:,layer-1]),np.sum(enc_neurons[:,layer])),dtype=np.float32)\n",
    "#            for mod in range(enc_neurons.shape[0]):\n",
    "            for mod in [0]:\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'enc_{layer}'] = _Sparse(mask)\n",
    "            self.masks[f'dec_{layer-1}'] = _Sparse(mask.T)\n",
    "        \n",
    "        # mask to and from latent layer\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,-1]),latent_dim * enc_neurons.shape[0]),dtype=np.float32)\n",
    "#        for mod in range(enc_neurons.shape[0]):\n",
    "        for mod in [0]:\n",
    "            mask[idx[mod,-1]:idx[mod+1,-1],mod*latent_dim:(mod+1)*latent_dim] = 1.\n",
    "        \n",
    "        print('enc_out:', mask)\n",
    "        self.masks['enc_out'] = _Sparse(mask)\n",
    "        self.masks[f'dec_{enc_neurons.shape[1]-1}'] = _Sparse(mask.T)\n",
    "        \n",
    "        # mask for decoder output\n",
    "        mask = np.zeros((np.sum(enc_neurons[:,0]),inp_shape[0]*enc_neurons.shape[0]),dtype=np.float32)\n",
    "#        for mod in range(enc_neurons.shape[0]):\n",
    "        for mod in [0]:\n",
    "            mask[idx[mod,0]:idx[mod+1,0], mod*inp_shape[0]:(mod+1)*inp_shape[0]] = 1.\n",
    "        self.masks['dec_out'] = _Sparse(mask)\n",
    "             \n",
    "        idx = np.concatenate((np.zeros((1,disc_neurons.shape[1]),dtype=np.int32),np.cumsum(disc_neurons,axis=0)))\n",
    "        # mask for discriminator layer 0:\n",
    "        mask = np.zeros((latent_dim * disc_neurons.shape[0],np.sum(disc_neurons[:,0])),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[mod*latent_dim:(mod+1)*latent_dim, idx[mod,0]:idx[mod+1,0]] = 1.\n",
    "        self.masks['disc_0'] = _Sparse(mask)\n",
    "#        print('disc_0',mask)\n",
    "            \n",
    "        # mask for discriminator layers 1:\n",
    "        for layer in range(1,disc_neurons.shape[1]):\n",
    "            mask = np.zeros((np.sum(disc_neurons[:,layer-1]),np.sum(disc_neurons[:,layer])),dtype=np.float32)\n",
    "            for mod in range(disc_neurons.shape[0]):\n",
    "                mask[idx[mod,layer-1]:idx[mod+1,layer-1],idx[mod,layer]:idx[mod+1,layer]] = 1.\n",
    "            self.masks[f'disc_{layer}'] = _Sparse(mask)\n",
    "#            print(f'disc_{layer}',mask)\n",
    "            \n",
    "        # mask for discriminator output:\n",
    "        mask = np.zeros((np.sum(disc_neurons[:,-1]),disc_neurons.shape[0]),dtype=np.float32)\n",
    "        for mod in range(disc_neurons.shape[0]):\n",
    "            mask[idx[mod,disc_neurons.shape[1]-1]:idx[mod+1,disc_neurons.shape[1]-1],mod] = 1.\n",
    "        self.masks['disc_out'] = _Sparse(mask)  \n",
    "#        print('disc_out',mask)\n",
    "\n",
    "        \n",
    "        inp = keras.Input(shape=inp_shape)\n",
    "        out = inp\n",
    "#        out = keras.layers.RepeatVector(self.n_models,trainable=False,name='inp_repeat')(inp)\n",
    "#        out = keras.layers.Reshape((self.n_models * inp_shape[0],),trainable=False,name='inp_reshape')(out)\n",
    "\n",
    "        \n",
    "        out = keras.layers.Dense(np.sum(enc_neurons[:,0]),\n",
    "                                 activation='relu',name = 'enc_0',\n",
    "                                 kernel_constraint = self.masks['enc_0'].c,\n",
    "                                 kernel_initializer = self.masks['enc_0'].i)(out)\n",
    "        out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_0')(out)\n",
    "        for num in range(1,enc_neurons.shape[1]):\n",
    "            name = f'enc_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',\n",
    "                                     name = name, kernel_constraint = self.masks[name].c,\n",
    "                                     kernel_initializer=self.masks[name].i)(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*latent_dim,name='enc_out',\n",
    "                                 kernel_constraint = self.masks['enc_out'].c,\n",
    "                                 kernel_initializer = self.masks['enc_out'].i)(out) # \n",
    "        latent = out\n",
    "        \n",
    "        # decoder layers are numbered in reverse so that neuron numbers match with encoder\n",
    "        for num  in reversed(range(enc_neurons.shape[1])):\n",
    "            name = f'dec_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',\n",
    "                                     name=name,\n",
    "                                     kernel_constraint=self.masks[name].c,\n",
    "                                     kernel_initializer=self.masks[name].i\n",
    "                                    )(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'dec_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(enc_neurons.shape[0]*inp_shape[0],name='dec_out',\n",
    "                                 kernel_constraint=self.masks['dec_out'].c,\n",
    "                                 kernel_initializer=self.masks['dec_out'].i\n",
    "                                )(out)\n",
    "        out = keras.layers.Reshape((enc_neurons.shape[0],inp_shape[0]))(out)\n",
    "        \n",
    "        self.aes = keras.Model(inputs=inp,outputs=[out,latent])\n",
    "        self.enc = keras.Model(inputs=inp,outputs=latent)\n",
    "        self.dec = keras.Model(inputs=latent,outputs=out)\n",
    "        \n",
    "        inp = keras.Input(shape=(latent_dim * self.n_models,))\n",
    "        disc = inp\n",
    "        for num in range(disc_neurons.shape[1]):\n",
    "            name = f'disc_{num}'\n",
    "            disc = keras.layers.Dense(np.sum(disc_neurons[:,num]),\n",
    "                                      name=name,kernel_constraint=self.masks[name].c,\n",
    "                                      kernel_initializer=self.masks[name].i\n",
    "                                     )(disc)\n",
    "            disc = keras.layers.LeakyReLU(alpha=0.2,name=f'disc_relu_{num}')(disc)\n",
    "            \n",
    "        disc = keras.layers.Dense(disc_neurons.shape[0],name='disc_out',\n",
    "                                  kernel_constraint=self.masks['disc_out'].c,\n",
    "                                  kernel_initializer=self.masks['disc_out'].i,\n",
    "                                  activation='sigmoid')(disc)\n",
    "        \n",
    "        self.disc = keras.Model(inputs=inp,outputs=disc)\n",
    "        \n",
    " \n",
    "\n",
    "    def compile(self):\n",
    "        # XXX\n",
    "        super().compile(optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002))\n",
    "        self.ae_weights = self.enc.trainable_weights + self.dec.trainable_weights\n",
    "\n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "        \n",
    "        # multiple models need replicated batch to compute loss simultaneously\n",
    "        multibatch = tf.stack([batch]*self.n_models,axis=1)\n",
    " \n",
    "        with tf.GradientTape() as aetape:\n",
    "            reconstruct = self.aes(batch)\n",
    "            print(reconstruct[0].shape)\n",
    "            mse = keras.metrics.mean_squared_error(multibatch,reconstruct[0])\n",
    "            print(mse.shape)\n",
    "            ae_multiloss = tf.reduce_mean(mse,axis=0)\n",
    "            print(ae_multiloss.shape)\n",
    "            \n",
    "            #ae_loss = tf.reduce_sum(ae_multiloss)\n",
    "            ae_loss = ae_multiloss[0]\n",
    "               \n",
    "        ae_grad = aetape.gradient(ae_loss,self.ae_weights)\n",
    "        self.last_ae_grad = ae_grad\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.ae_weights))\n",
    "               \n",
    "#        rand_low = tf.random.normal((batch.shape[0],self.latent_dim * self.n_models))\n",
    "        rand_low = tf.random.normal((batch.shape[0],self.latent_dim))\n",
    "#        rand_low = tf.random.uniform((batch.shape[0],self.latent_dim))\n",
    "        rand_low = tf.tile(rand_low,(1,self.n_models))\n",
    "        lows = tf.concat([rand_low, reconstruct[1]],axis=0)\n",
    "        labels = tf.concat([tf.ones((batch.shape[0],self.n_models)), tf.zeros((batch.shape[0],self.n_models))], axis=0)\n",
    "        labels += 0.05 * tf.random.uniform(labels.shape)\n",
    "                        \n",
    "        with tf.GradientTape() as dtape:\n",
    "            \n",
    "            # FIXME: perturbe\n",
    "            neg_pred = self.disc(reconstruct[1])\n",
    "            neg_losses = -tf.reduce_mean(tf.math.log(1-neg_pred),axis=0) \n",
    "            pos_pred = self.disc(rand_low)\n",
    "            pos_losses = -tf.reduce_mean(tf.math.log(pos_pred),axis=0)\n",
    "            disc_losses = neg_losses + pos_losses\n",
    "#            disc_loss = tf.reduce_mean(disc_losses)\n",
    "            disc_loss = disc_losses[0]\n",
    "                \n",
    "        disc_grads = dtape.gradient(disc_loss,self.disc.trainable_weights)\n",
    "#        self.optimizer.apply_gradients(zip(disc_grads,self.disc.trainable_weights))\n",
    "          \n",
    "        all_true = tf.ones((batch.shape[0],self.n_models))\n",
    "    \n",
    "        with tf.GradientTape() as ctape:\n",
    "            cheat = self.disc(self.enc(batch))\n",
    "            cheat_losses = -tf.reduce_mean(tf.math.log(cheat),axis=0)\n",
    "            cheat_loss = cheat_losses[0]\n",
    "            #cheat_loss = tf.reduce_mean(cheat_losses)\n",
    "            \n",
    "        cheat_grads = ctape.gradient(cheat_loss,self.enc.trainable_weights)\n",
    "#        self.optimizer.apply_gradients(zip(cheat_grads,self.enc.trainable_weights))\n",
    "        \n",
    "#        return {\n",
    "#            'AE loss min' : tf.reduce_min(ae_multiloss),\n",
    "#            'AE loss max' : tf.reduce_max(ae_multiloss),\n",
    "#            'disc loss min' : tf.reduce_min(disc_losses),\n",
    "#            'disc loss max' : tf.reduce_max(disc_losses),\n",
    "#            'cheat loss min' : tf.reduce_min(cheat_losses),\n",
    "#            'cheat loss max' : tf.reduce_max(cheat_losses)\n",
    "#        }\n",
    "\n",
    "        return { str(i): ae_multiloss[i] for i in range(ae_multiloss.shape[0]) }\n",
    "        #return { '0': disc_losses[0], '1' : disc_losses[1] }\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self,inp):\n",
    "        return self.aes(inp)\n",
    "                            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0016f19a-db85-4407-b718-8119dd6302bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#m5 = AAEMultiModel5((X_train.shape[1],),np.array([[96,48],[96,48]]*10),np.array([[224,112],[64,32]]*10))\n",
    "#m5 = AAEMultiModel5((X_train.shape[1],),np.array([[96,48],[96,48]]),np.array([[224,112],[64,32]]))\n",
    "m5 = AAEMultiModel5((X_train.shape[1],),np.array([[96,48]]*10),np.array([[224,112]]*10))\n",
    "#m5 = AAEMultiModel5((X_train.shape[1],),np.array([[4]]*10),np.array([[4,3]]*10))\n",
    "m5.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e56a857e-1438-4fb7-9f5f-22e9b44c2926",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tf.data.Dataset.from_tensor_slices(X_train[:256]).shuffle(2048).batch(256,drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ea89fe-02f5-4668-85b9-0c8d08f5344c",
   "metadata": {},
   "outputs": [],
   "source": [
    "m5.fit(ds,epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d8001b-4ded-4861-a484-3b716595ef1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m5.enc(X_train[::500]).numpy()\n",
    "#lows = m5.aes(X_train[::500])[1].numpy()\n",
    "\n",
    "plt.figure(figsize=(12,9))\n",
    "for mod in range(lows.shape[1]//2):\n",
    "    plows = lows[:,mod*2:(mod+1)*2]\n",
    "    plt.scatter(plows[:,0],plows[:,1],marker='.',label=str(mod))\n",
    "\n",
    "lim=2000000\n",
    "plt.legend()    \n",
    "plt.xlim((-lim,lim))\n",
    "plt.ylim((-lim,lim))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "293f24bb-3c5f-4320-a1ed-ada42594e34d",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(m5.last_ae_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a02ca9a6-14aa-48e1-9b47-959c91ae3d10",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in m5.ae_weights:\n",
    "    if 'kernel' in w.name: print(w.name,np.sum(np.abs(w.numpy())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0237a835-d30f-4fa5-81c7-5c66c236b928",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in m5.ae_weights:\n",
    "    if 'kernel' in w.name: print(w.name,np.sum(np.abs(w.numpy())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90a81d2-0f89-494f-994c-ccc3f150817c",
   "metadata": {},
   "outputs": [],
   "source": [
    "m5.ae_weights[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab773843-4e60-4cc6-962d-daa8a36e2ebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "_glorot(shape=(2,5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b135a48-0c4c-40e2-b507-c2544d44ea83",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.sum(np.array([1,2,3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2053c26d-aa52-4034-bcdb-5a46f38cd838",
   "metadata": {},
   "outputs": [],
   "source": [
    "_random_init = keras.initializers.GlorotUniform(seed=42)\n",
    "_random_init((2,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5d2a05b-4657-4312-bc2f-7bc76df2b172",
   "metadata": {},
   "outputs": [],
   "source": [
    "_random_init = keras.initializers.GlorotUniform(seed=42)\n",
    "\n",
    "class _SparseConstraint(keras.constraints.Constraint):\n",
    "    \n",
    "    def __init__(self,left,right):\n",
    "        assert len(left) == len(right)\n",
    "        mask = np.zeros((np.sum(left),np.sum(right)),dtype=np.float32)\n",
    "        idxl = np.concatenate((np.zeros((1,),dtype=np.int32),np.cumsum(left)))\n",
    "        idxr = np.concatenate((np.zeros((1,),dtype=np.int32),np.cumsum(right)))\n",
    "        for mod in range(len(left)):\n",
    "            mask[idxl[mod]:idxl[mod+1],idxr[mod]:idxr[mod+1]] = 1.\n",
    "     \n",
    "        self.mask = tf.convert_to_tensor(mask)\n",
    "        \n",
    "    @tf.function    \n",
    "    def __call__(self,w):\n",
    "        return w * self.mask\n",
    "\n",
    "class _SparseInitializer(keras.initializers.Initializer):\n",
    "    def __init__(self,left,right):\n",
    "        assert len(left) == len(right)\n",
    "        self.left = left\n",
    "        self.right = right\n",
    "        self.idxl = np.concatenate((np.zeros((1,),dtype=np.int32),np.cumsum(left)))\n",
    "        self.idxr = np.concatenate((np.zeros((1,),dtype=np.int32),np.cumsum(right)))\n",
    "    \n",
    "    def __call__(self,shape,dtype=None):\n",
    "#        print(shape,self.left,self.right)\n",
    "        assert shape == [np.sum(self.left),np.sum(self.right)]\n",
    "        \n",
    "        init = np.zeros((np.sum(self.left),np.sum(self.right)),dtype=dtype.as_numpy_dtype)\n",
    "        for mod in range(len(self.left)):\n",
    "            init[self.idxl[mod]:self.idxl[mod+1],self.idxr[mod]:self.idxr[mod+1]] = _random_init((self.left[mod],self.right[mod])).numpy()\n",
    "        \n",
    "        return tf.convert_to_tensor(init)\n",
    "    \n",
    "    \n",
    "def _masks(left,right):\n",
    "    return { 'kernel_initializer': _SparseInitializer(left,right),\n",
    "            'kernel_constraint': _SparseConstraint(left,right) }\n",
    "\n",
    "\n",
    "\n",
    "class AAEMultiModel6(keras.models.Model):\n",
    "    # enc_neurons: (models,layers)\n",
    "    \n",
    "    def __init__(self,inp_shape,enc_neurons,disc_neurons,latent_dim=2):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_models = enc_neurons.shape[0]\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        assert disc_neurons.shape[0] == self.n_models\n",
    "                \n",
    "        inp = keras.Input(shape=inp_shape)\n",
    "        out = inp\n",
    "\n",
    "        out = keras.layers.Dense(np.sum(enc_neurons[:,0]),activation='relu',name = 'enc_0')(out)\n",
    "        out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_0')(out)\n",
    "        \n",
    "        for num in range(1,enc_neurons.shape[1]):\n",
    "            name = f'enc_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',\n",
    "                                     name = name, **_masks(enc_neurons[:,num-1],enc_neurons[:,num]))(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'enc_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(self.n_models*latent_dim,name='enc_out',\n",
    "                                 **_masks(enc_neurons[:,-1],[latent_dim]*self.n_models))(out) \n",
    "        latent = out\n",
    "        \n",
    "        out = keras.layers.Dense(np.sum(enc_neurons[:,-1]),activation='relu',name=f'dec_{enc_neurons.shape[1]}',\n",
    "                                 **_masks([latent_dim]*self.n_models,enc_neurons[:,-1]))(out)\n",
    "        out = keras.layers.BatchNormalization(momentum=0.8,name=f'dec_bn_{enc_neurons.shape[1]}')(out)\n",
    "        \n",
    "        # decoder layers are numbered in reverse so that neuron numbers match with encoder\n",
    "        for num in reversed(range(enc_neurons.shape[1]-1)):\n",
    "            name = f'dec_{num}'\n",
    "            out = keras.layers.Dense(np.sum(enc_neurons[:,num]),activation='relu',name=name,\n",
    "                                     **_masks(enc_neurons[:,num+1],enc_neurons[:,num]))(out)\n",
    "            out = keras.layers.BatchNormalization(momentum=0.8,name=f'dec_bn_{num}')(out)\n",
    "            \n",
    "        out = keras.layers.Dense(self.n_models*inp_shape[0],name='dec_out',activation='relu',\n",
    "                                    **_masks(enc_neurons[:,0],[inp_shape[0]]*self.n_models))(out)\n",
    "        out = keras.layers.Reshape((self.n_models,inp_shape[0]))(out)\n",
    "        \n",
    "        self.aes = keras.Model(inputs=inp,outputs=[out,latent])\n",
    "        self.enc = keras.Model(inputs=inp,outputs=latent)\n",
    "        self.dec = keras.Model(inputs=latent,outputs=out)\n",
    "        \n",
    "        inp = keras.Input(shape=(latent_dim * self.n_models,))\n",
    "        disc = inp\n",
    "        disc = keras.layers.Dense(np.sum(disc_neurons[:,0]),name='disc_0',\n",
    "                                  **_masks([latent_dim]*self.n_models,disc_neurons[:,0]))(disc)\n",
    "        disc = keras.layers.LeakyReLU(alpha=0.2,name=f'disc_relu_{num}')(disc)\n",
    "        \n",
    "        for num in range(1,disc_neurons.shape[1]):\n",
    "            name = f'disc_{num}'\n",
    "            disc = keras.layers.Dense(np.sum(disc_neurons[:,num]),name=name,\n",
    "                                      **_masks(disc_neurons[:,num-1],disc_neurons[:,num]))(disc)\n",
    "            disc = keras.layers.LeakyReLU(alpha=0.2,name=f'disc_relu_{num}')(disc)\n",
    "            \n",
    "        disc = keras.layers.Dense(self.n_models,name='disc_out',\n",
    "                                  **_masks(disc_neurons[:,-1],[1]*self.n_models))(disc)\n",
    "        \n",
    "        self.disc = keras.Model(inputs=inp,outputs=disc)\n",
    "        \n",
    " \n",
    "\n",
    "    def compile(self):\n",
    "        # XXX\n",
    "        super().compile(optimizer = keras.optimizers.legacy.Adagrad(learning_rate=0.0002))\n",
    "        self.ae_weights = self.enc.trainable_weights + self.dec.trainable_weights\n",
    "\n",
    "    @tf.function\n",
    "    def train_step(self,batch):\n",
    "        if isinstance(batch,tuple):\n",
    "            batch = batch[0]\n",
    "            \n",
    "        # multiple models need replicated batch to compute loss simultaneously\n",
    "        multibatch = tf.stack([batch]*self.n_models,axis=1)\n",
    " \n",
    "        with tf.GradientTape() as aetape:\n",
    "            reconstruct = self.aes(batch)\n",
    "            mse = keras.metrics.mean_squared_error(multibatch,reconstruct[0])\n",
    "            ae_multiloss = tf.reduce_mean(mse,axis=0)\n",
    "            \n",
    "            ae_loss = tf.reduce_sum(ae_multiloss)\n",
    "               \n",
    "        ae_grad = aetape.gradient(ae_loss,self.ae_weights)\n",
    "        self.last_ae_grad = ae_grad\n",
    "        self.optimizer.apply_gradients(zip(ae_grad,self.ae_weights))\n",
    "               \n",
    "#        labels += 0.05 * tf.random.uniform(labels.shape)\n",
    "        rand_low = tf.random.normal((tf.shape(batch)[0],self.latent_dim))\n",
    "#        rand_low = tf.random.uniform((batch.shape[0],self.latent_dim))\n",
    "        rand_low = tf.tile(rand_low,(1,self.n_models))\n",
    "                        \n",
    "        with tf.GradientTape() as dtape:\n",
    "            # perturbe\n",
    "            neg_pred = self.disc(reconstruct[1])\n",
    "            neg_losses = tf.reduce_mean(neg_pred*tf.random.uniform(tf.shape(neg_pred),1.,1.05),axis=0) \n",
    "            pos_pred = self.disc(rand_low)\n",
    "            pos_losses = -tf.reduce_mean(pos_pred*tf.random.uniform(tf.shape(pos_pred),1.,1.05),axis=0)\n",
    "            disc_losses = neg_losses + pos_losses\n",
    "            disc_loss = tf.reduce_mean(disc_losses)\n",
    "                \n",
    "        disc_grads = dtape.gradient(disc_loss,self.disc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(disc_grads,self.disc.trainable_weights))\n",
    "          \n",
    "        with tf.GradientTape() as ctape:\n",
    "            # perturbe\n",
    "            cheat = self.disc(self.enc(batch))\n",
    "            cheat_losses = -tf.reduce_mean(cheat*tf.random.uniform(tf.shape(cheat),1.,1.05),axis=0)\n",
    "            cheat_loss = tf.reduce_mean(cheat_losses)\n",
    "            \n",
    "        cheat_grads = ctape.gradient(cheat_loss,self.enc.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(cheat_grads,self.enc.trainable_weights))\n",
    "        \n",
    "#        return {\n",
    "#            'AE loss min' : tf.reduce_min(ae_multiloss),\n",
    "#            'AE loss max' : tf.reduce_max(ae_multiloss),\n",
    "#            'disc loss min' : tf.reduce_min(disc_losses),\n",
    "#            'disc loss max' : tf.reduce_max(disc_losses),\n",
    "#            'cheat loss min' : tf.reduce_min(cheat_losses),\n",
    "#            'cheat loss max' : tf.reduce_max(cheat_losses)\n",
    "#        }\n",
    "\n",
    "        return { str(i): ae_multiloss[i] for i in range(ae_multiloss.shape[0]) }\n",
    "        #return { '0': disc_losses[0], '1' : disc_losses[1] }\n",
    "    \n",
    "    @tf.function\n",
    "    def call(self,inp):\n",
    "        return self.aes(inp)\n",
    "                            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b5ef81-7797-4b0d-bc58-661b2328a644",
   "metadata": {},
   "outputs": [],
   "source": [
    "m=_masks([2,3],[4,5])\n",
    "m['kernel_constraint'].mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108bf3d9-1996-4f37-af05-4f97359bddfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "m['kernel_initializer']([5,9],tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95bed5cd-0bba-40ce-a303-c6d94a244580",
   "metadata": {},
   "outputs": [],
   "source": [
    "#m6 = AAEMultiModel6((X_train.shape[1],),np.array([[96,48],[96,48]]*10),np.array([[224,112],[64,32]]*10))\n",
    "#m6 = AAEMultiModel6((X_train.shape[1],),np.array([[96,48],[96,48]]),np.array([[224,112],[64,32]]))\n",
    "#m6 = AAEMultiModel6((X_train.shape[1],),np.array([[96,48]]*10),np.array([[224,112]]*10))\n",
    "#m6 = AAEMultiModel6((X_train.shape[1],),np.array([[4]]*10),np.array([[4,3]]*10))\n",
    "m6 = AAEMultiModel6((X_train.shape[1],),en,dn)\n",
    "m6.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d02c160-7778-4764-9a52-2ee3def229a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in m6.ae_weights:\n",
    "    if 'kernel' in w.name: print(w.name,np.sum(np.abs(w.numpy())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbd704d-127b-42f3-bcdc-7b916333dec0",
   "metadata": {},
   "outputs": [],
   "source": [
    "lows = m6.enc(X_train[::200]).numpy()\n",
    "#lows = m5.aes(X_train[::500])[1].numpy()\n",
    "\n",
    "plt.figure(figsize=(12,9))\n",
    "for mod in range(lows.shape[1]//2):\n",
    "    plows = lows[:,mod*2:(mod+1)*2]\n",
    "    plt.scatter(plows[:,0],plows[:,1],marker='.',label=str(mod))\n",
    "\n",
    "lim=2\n",
    "plt.legend()    \n",
    "plt.xlim((-lim,lim))\n",
    "plt.ylim((-lim,lim))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db33490-0cb7-44cf-bd0a-32008f81d33d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = tf.data.Dataset.from_tensor_slices(X_train).shuffle(2048).batch(256,drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a62c27f8-447a-48ef-89b6-fcea4221f9c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "m6.fit(ds,epochs=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c575f84b-6429-4dc8-bfa6-6d99edf7c987",
   "metadata": {},
   "outputs": [],
   "source": [
    "en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c6bce1-de24-4f8c-ac7f-bf5be007725a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8570a052-e4f8-43a6-9665-90209040f675",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}