pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.metric_dataloader import MetricDataPreprocessor\n",
    "from utils.tshae_utils import load_tshae_model\n",
    "from models.ddpm_models import ContextUnet, DDPM\n",
    "from omegaconf import OmegaConf\n",
    "import torch\n",
    "import numpy as np\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best_models/FD003/tshae/tshae_best_model.pt\n",
      "{'train_ids': array([ 1,  2,  3,  5,  6,  7,  8,  9, 11, 13, 14, 15, 16, 17, 19, 20, 21,\n",
      "       23, 24, 25, 26, 27, 28, 29, 32, 34, 35, 36, 37, 38, 40, 41, 42, 43,\n",
      "       46, 47, 48, 49, 50, 51, 52, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,\n",
      "       64, 65, 66, 67, 68, 69, 71, 72, 74, 75, 78, 79, 81, 82, 84, 85, 86,\n",
      "       87, 88, 89, 91, 92, 93, 94, 95, 96, 97, 98, 99]), 'test_ids': array([  1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,\n",
      "        14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,\n",
      "        27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,\n",
      "        40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,\n",
      "        53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,\n",
      "        66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,\n",
      "        79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,\n",
      "        92,  93,  94,  95,  96,  97,  98,  99, 100]), 'val_ids': array([ 4, 10, 12, 18, 22, 30, 31, 33, 39, 44, 45, 53, 70, 73, 76, 77, 80,\n",
      "       83, 90])}\n"
     ]
    }
   ],
   "source": [
    "class LatentEncoder:\n",
    "    def __init__(self, checkpoint_path: str = 'best_models/FD003/tshae/', device: str ='cuda'):\n",
    "        config_path = checkpoint_path + \".hydra/config.yaml\"\n",
    "        model_path = checkpoint_path + \"tshae_best_model.pt\"\n",
    "        self.device = device if torch.cuda.is_available() else \"cpu\"\n",
    "        self.config = self._get_congig(config_path=config_path)\n",
    "        print(model_path)\n",
    "        self.model = load_tshae_model(model_path=model_path).to(self.device)\n",
    "        train_dataset, test_dataset, val_dataset = self._get_datasets()\n",
    "        self.train_dataset = train_dataset\n",
    "        self.test_dataset = test_dataset\n",
    "        self.val_dataset = val_dataset\n",
    "    \n",
    "    def _get_congig(self, config_path: str):\n",
    "        config = OmegaConf.load(config_path)\n",
    "        return config\n",
    "\n",
    "    def _get_datasets(self):\n",
    "        preproc = MetricDataPreprocessor(**self.config.data_preprocessor)\n",
    "        train_dataset, test_dataset, val_dataset = preproc.get_datasets()\n",
    "        return train_dataset, test_dataset, val_dataset\n",
    "    \n",
    "    def get_run_ids(self):\n",
    "        run_ids = {\n",
    "            'train_ids': self.train_dataset.ids,\n",
    "            'test_ids': self.test_dataset.ids,\n",
    "            'val_ids': self.val_dataset.ids\n",
    "            }\n",
    "        return run_ids\n",
    "\n",
    "    def encode_run_id(self, run_id: int) -> np.ndarray:\n",
    "        if run_id in self.train_dataset.ids:\n",
    "            x, true_rul = self.train_dataset.get_run(run_id)\n",
    "        elif run_id in self.test_dataset.ids:\n",
    "            x, true_rul = self.test_dataset.get_run(run_id)\n",
    "        elif run_id in self.val_dataset.ids:\n",
    "            x, true_rul = self.val_dataset.get_run(run_id)\n",
    "        else:\n",
    "            raise KeyError('No such run_id in datasets!')\n",
    "        rul_hat, z, *_ = self.model(x.to(self.device))\n",
    "        return z.detach().cpu().numpy()\n",
    "\n",
    "le = LatentEncoder()\n",
    "print(le.get_run_ids())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = le.encode_run_id(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "sampling timestep 100\r"
     ]
    }
   ],
   "source": [
    "class DiffusionGen:\n",
    "    def __init__(self, checkpoint_path = 'best_models/FD003/ddpm/', tshae_model_path='best_models/FD003/tshae/tshae_best_model.pt', device='cuda'):\n",
    "        config_path = checkpoint_path + \".hydra/config.yaml\"\n",
    "        model_path = checkpoint_path + \"ddpm_best_model.pt\"\n",
    "        self.device = device if torch.cuda.is_available() else \"cpu\"\n",
    "        self.config = self._get_congig(config_path=config_path)\n",
    "        self.ddpm_model = self._get_model().to(self.device)\n",
    "        self.tshae_model = load_tshae_model(model_path=tshae_model_path).to(self.device)\n",
    "    \n",
    "    def _get_congig(self, config_path):\n",
    "        config = OmegaConf.load(config_path)\n",
    "        return config\n",
    "    \n",
    "    def _get_model(self):\n",
    "        n_T = self.config.diffusion.ddpm_train.n_T \n",
    "        z_dim   = self.config.diffusion.ddpm_train.z_dim\n",
    "        n_feat = self.config.diffusion.ddpm_train.n_feat\n",
    "        drop_prob = self.config.diffusion.ddpm_model.drop_prob\n",
    "        \n",
    "        ddpm = DDPM(\n",
    "            nn_model=ContextUnet(\n",
    "            in_channels=1, \n",
    "            n_feat=n_feat, \n",
    "            z_dim=z_dim), \n",
    "            betas=(1e-4, 0.02), \n",
    "            n_T= n_T, \n",
    "            device=self.device, \n",
    "            drop_prob=drop_prob)\n",
    "        ddpm.load_state_dict(torch.load(self.config.diffusion.checkpoint_ddpm.path))\n",
    "        ddpm.eval()\n",
    "        return ddpm\n",
    "    \n",
    "    def generate_from_latent(self, z_space, num_samples=4, w=0.5, quantile=0.25, mode='best'):\n",
    "\n",
    "        history = defaultdict(dict)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            z = torch.FloatTensor(z_space).to(self.device)\n",
    "            x_tshae_samples = self.tshae_model.decoder(z)\n",
    "            x_hat_diffusion, _ = self.ddpm_model.sample_cmapss(n_sample=num_samples, size=(1,32,32), device=z.device, z_space_contexts=z, guide_w=w)\n",
    "\n",
    "            # Calculate distances between the passed latent trajectory and generated trajectories\n",
    "\n",
    "            num_z = z.shape[0]\n",
    "            x_hat_diffusion = x_hat_diffusion.squeeze(1)[:,:,:21]  # Adjusting shape to match the desired sensor reconstruction\n",
    "\n",
    "\n",
    "            #================================================\n",
    "            z = z.unsqueeze(1)\n",
    "\n",
    "            x_hat_diffusion = x_hat_diffusion.squeeze(1)[:,:,:21]\n",
    "            with torch.no_grad(): \n",
    "                rul_hat_diff, z_diff, *_  = self.tshae_model(x_hat_diffusion)\n",
    "            z_diff = z_diff.reshape(num_samples, num_z, 2).permute(1,0,2)\n",
    " \n",
    "            rul_hat_diff = rul_hat_diff.reshape(num_samples, num_z,  1).permute(1,0,2)\n",
    "\n",
    "\n",
    "            distances = torch.norm(z - z_diff, dim=-1)\n",
    "\n",
    "            if mode == \"quantile\":\n",
    "                limits = torch.quantile(distances.squeeze(1), quantile, interpolation='linear', dim=1, keepdim=True)\n",
    "                choices = []\n",
    "                for i in range(distances.shape[0]):\n",
    "                    choices.append(\n",
    "                        np.random.choice(\n",
    "                            np.flatnonzero(\n",
    "                                np.where(\n",
    "                                    distances.squeeze(1).detach().cpu().numpy()[i] < limits.detach().cpu().numpy()[i], 1, 0\n",
    "                                )\n",
    "                            ), \n",
    "                        1)\n",
    "                    )\n",
    "                best_samples = torch.from_numpy(np.array(choices).squeeze(1))\n",
    "            else:\n",
    "                best_samples = torch.argmin(distances.squeeze(1), dim=-1)\n",
    "            \n",
    "\n",
    "            x_diff_samples = x_hat_diffusion.reshape(num_samples, num_z, 32, 21).permute(1,0,2,3) #[num_z, num_samples, 32, 21]\n",
    "\n",
    "            x_diff_samples = x_diff_samples[range(x_diff_samples.shape[0]), best_samples]\n",
    "\n",
    "            z_diff = z_diff[range(z_diff.shape[0]), best_samples]\n",
    "            rul_hat_diff = rul_hat_diff[range(rul_hat_diff.shape[0]), best_samples]\n",
    "            \n",
    "            x_diff_samples = x_diff_samples.to(\"cpu\")\n",
    "            x_tshae_samples = x_tshae_samples.to(\"cpu\")\n",
    "            \n",
    "            sensors_diff_reconstructed = []\n",
    "            sensors_tshae_reconstructed = []\n",
    "\n",
    "            for ind in range(num_z):\n",
    "                if ind == 0:\n",
    "                    sensors_diff_reconstructed.append(x_diff_samples[ind])\n",
    "                    sensors_tshae_reconstructed.append(x_tshae_samples[ind])\n",
    "                else:\n",
    "                    sensors_diff_reconstructed.append(np.expand_dims(x_diff_samples[ind, -1, :], axis=0))\n",
    "                    sensors_tshae_reconstructed.append(np.expand_dims(x_tshae_samples[ind, -1, :], axis=0))\n",
    "\n",
    "            sensors_diff_reconstructed = np.concatenate(sensors_diff_reconstructed, axis=0)\n",
    "            sensors_tshae_reconstructed = np.concatenate(sensors_tshae_reconstructed, axis=0)\n",
    "\n",
    "\n",
    "            history['z'] = z.squeeze().detach().cpu().numpy()\n",
    "            history[\"z_diff\"] = z_diff.detach().cpu().numpy()\n",
    "            \n",
    "            history[\"x_diff_samples\"] = x_diff_samples.detach().cpu().numpy()\n",
    "            \n",
    "            history[\"rul_hat_diff\"] = rul_hat_diff.detach().cpu().numpy()\n",
    "            \n",
    "            history[\"sensors_diff_reconstructed\"] =  sensors_diff_reconstructed\n",
    "            history[\"sensors_tshae_reconstructed\"] =  sensors_tshae_reconstructed\n",
    "        return history\n",
    "\n",
    "dg = DiffusionGen()\n",
    "h = dg.generate_from_latent(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(228, 2)\n",
      "(259, 21)\n",
      "(259, 21)\n"
     ]
    }
   ],
   "source": [
    "print(z.shape)\n",
    "print(h['sensors_diff_reconstructed'].shape)\n",
    "print(h['sensors_tshae_reconstructed'].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_ml",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}