notebooks/Train_single_datapoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This file was used to test models on a single user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "parent_dir = os.path.abspath(os.path.join(os.getcwd(), os.pardir))\n",
    "\n",
    "sys.path.append(parent_dir)\n",
    "\n",
    "from simtrain import SETTINGS_POLIMI as SETTINGS\n",
    "from simtrain.sim_models_new import Toy_intensity_Generator, Toy_intensity_Comparer\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import pytorch_warmup as warmup\n",
    "\n",
    "import paths\n",
    "from os.path import join\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import simtrain.utils as utils\n",
    "from tqdm import tqdm\n",
    "import evotorch\n",
    "import ray\n",
    "from functools import partial\n",
    "import pandas as pd\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from simtrain.Dataset import TimestepFrequencyDataset\n",
    "\n",
    "os.environ['PYTHONPATH'] = \"/home/thahit/github/Recommender_Sim\"\n",
    "#ray.init(ignore_reinit_error=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "state_size = SETTINGS.STATE_SIZE\n",
    "experiment_name = \"toy\"\n",
    "#num_negatives = 100\n",
    "#conditioned=True\n",
    "#kl_weight=.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generated Data\n",
    "Fake data generated by a neural network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter dicts\n",
    "width= 8\n",
    "user_state_dict = {\"model_hyp\": {\"layer_width\": [width, width]}}\n",
    "intensity_state_dict = {\"model_hyp\": {\"layer_width\": [width, width],\n",
    "                                                         \"noise\": 0}\n",
    "                            }\n",
    "\n",
    "\n",
    "hyperparameter_dict = {\"state_size\": state_size, \"state_model\": user_state_dict, \n",
    "                           \"intensity_model\": intensity_state_dict,# \"num_recom\" : num_items_per_recom,\n",
    "                            \"noise\": 0.}\n",
    "gen_model = Toy_intensity_Generator(hyperparameter_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = join(paths.dat, SETTINGS.filepaths_new[\"simulate_intensity_model\"]\n",
    "            )\n",
    "path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#torch.save(gen_model.state_dict(), path)\n",
    "gen_model.load_state_dict(torch.load(path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = torch.zeros((1, state_size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = [gen_model.sample_one(state).detach().numpy() for _ in range(100)]\n",
    "samples = np.array(samples)\n",
    "samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "feed in true values for fake models, \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 6))\n",
    "\n",
    "# Histogram\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.hist(samples, bins=50, density=True, alpha=0.6, color='g', edgecolor='black')\n",
    "plt.title('Histogram of Samples')\n",
    "plt.xlabel('Value')\n",
    "plt.ylabel('Frequency')\n",
    "\n",
    "# Density Plot\n",
    "plt.subplot(1, 2, 2)\n",
    "sns.kdeplot(samples, shade=True, color='g', bw_adjust=0.1)\n",
    "plt.title('Density Plot of Samples')\n",
    "plt.xlabel('Value')\n",
    "plt.ylabel('Density')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_path = gen_model.sample_path(num_samples=5)\n",
    "sample_path =torch.tensor(sample_path)\n",
    "sample_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.eventplot(sample_path, orientation='vertical', colors='r')\n",
    "plt.scatter(sample_path, [1] * len(sample_path), color='blue', label='Time Series 1', s=10, marker='x')\n",
    "\n",
    "plt.yticks([])\n",
    "plt.xlabel('Time')\n",
    "plt.title('Event Plot of Sparse Binary Events')\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(join(paths.dat, SETTINGS.rootpaths['models'],\n",
    "                             \"testing\", \"data.h5\"))\n",
    "list_of_dicts = checkpoint['data']\n",
    "chosen_sample = list_of_dicts[4][\"timestamps\"]\n",
    "sample_path =torch.tensor(chosen_sample)\n",
    "sample_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.eventplot(sample_path, orientation='vertical', colors='r')\n",
    "plt.scatter(sample_path, [1] * len(sample_path), color='blue', label='Time Series 1', s=10, marker='x')\n",
    "\n",
    "plt.yticks([])\n",
    "plt.xlabel('Time')\n",
    "plt.title('Event Plot of Sparse Binary Events')\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Gradient Descent\n",
    "Not working"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_single(model, path, scoring_func,optimizer, num_epochs=100, num_tries=20):\n",
    "    \n",
    "    for iter in tqdm(range(num_epochs)):\n",
    "        last_t = 0\n",
    "        state = torch.zeros((1, state_size))\n",
    "        loss = 0.\n",
    "        #results = []\n",
    "        for timestep in path:\n",
    "            current_pred = []\n",
    "            for _ in range(num_tries):\n",
    "                current_pred.append(last_t + model.sample_one(state))\n",
    "            current_pred = torch.stack(current_pred)\n",
    "            #results.append(torch.mean(current_pred))\n",
    "            loss += torch.log(scoring_func(current_pred, timestep))\n",
    "            delta = timestep-last_t\n",
    "            last_t = timestep\n",
    "            state = model.evolve_state(state, delta)\n",
    "        print(\"loss: \", loss)\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2)\n",
    "        optimizer.step()\n",
    "\n",
    "        #for name, param in model.named_parameters():\n",
    "        #    print(f\"Parameter Name: {name}\")\n",
    "        #    print(f\"Parameter Value: {param}\")\n",
    "        #    print(f\"Gradients: {param.grad}\")\n",
    "        #    print(f\"Parameter Shape: {param.shape}\")\n",
    "        #    print(f\"Requires Gradient: {param.requires_grad}\")\n",
    "        #    print(\"-\" * 40)\n",
    "        #return\n",
    "        optimizer.zero_grad()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter dicts\n",
    "width= 16\n",
    "user_state_dict = {\"model_hyp\": {\"layer_width\": [width, width, width]}}\n",
    "intensity_state_dict = {\"model_hyp\": {\"layer_width\": [width, width, width],\n",
    "                                                         \"noise\": 0}\n",
    "                            }\n",
    "\n",
    "\n",
    "hyperparameter_dict = {\"state_size\": state_size, \"state_model\": user_state_dict, \n",
    "                           \"intensity_model\": intensity_state_dict,# \"num_recom\" : num_items_per_recom,\n",
    "                            \"noise\": 0.}\n",
    "train_model = Toy_intensity_Generator(hyperparameter_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_train = join(paths.dat, SETTINGS.filepaths_new[\"copy_intensity_model\"]\n",
    "            )\n",
    "#torch.save(gen_model.state_dict(), path_train)\n",
    "#train_model.load_state_dict(torch.load(path_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_path = torch.as_tensor(sample_path)\n",
    "optimizer = optim.AdamW(train_model.parameters(), lr=0.1,\n",
    "                        weight_decay=1e-7)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#train_single(train_model, sample_path, scoring_func=utils.energy_score_loss,\n",
    "#            optimizer=optimizer, num_epochs=30, num_tries=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### function approximation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simtrain.sim_models_new import all_in_one_model\n",
    "\n",
    "width=64\n",
    "user_state_dict = {\"model_hyp\": {\"layer_width\": [width for _ in range(4)]}}\n",
    "time_dict = {\"model_hyp\": {\"layer_width\": [width for _ in range(3)]}\n",
    "            }\n",
    "            \n",
    "\n",
    "timecheat = False\n",
    "hyperparameter_dict = {\"state_size\": state_size, \"time_model\": time_dict, \n",
    "                           \"state_model\": user_state_dict}\n",
    "train_model = all_in_one_model(hyperparameter_dict, timecheat=timecheat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simtrain.train import train_single_function_approx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_epoch = 1\n",
    "\n",
    "num_epochs = 1500\n",
    "warmup_period = steps_per_epoch * 10\n",
    "num_steps = num_epochs*steps_per_epoch - warmup_period\n",
    "num_iter_til_first_restart = num_steps + 1\n",
    "user_lr = 0.001\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "optimizer = optim.AdamW(train_model.parameters(), lr=user_lr,\n",
    "                        weight_decay=1e-7)\n",
    "lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(\n",
    "    optimizer, T_0=num_iter_til_first_restart, T_mult=1, eta_min=1e-7)\n",
    "\n",
    "warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_single_function_approx(train_model, sample_path,scoring_func=utils.energy_score_loss,\n",
    "            state_size=state_size, warmup_scheduler=warmup_scheduler, lr_scheduler=lr_scheduler,\n",
    "            optimizer=optimizer, num_epochs=num_epochs, num_tries=30, timecheat=timecheat, loss_print_interval=num_epochs//20,\n",
    "            warmup_period=warmup_period)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.AdamW(train_model.parameters(), lr=0.0001,\n",
    "                        weight_decay=1e-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_single_function_approx(train_model, sample_path,scoring_func=utils.energy_score_loss,\n",
    "            optimizer=optimizer, num_epochs=3000,timecheat=timecheat, num_tries=20, \n",
    "            loss_print_interval=20, state_size=state_size,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train via Evolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from evotorch.algorithms import PGPE, CMAES\n",
    "from evotorch.logging import PandasLogger\n",
    "from functools import partial\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter dicts\n",
    "width= 64\n",
    "user_state_dict = {\"model_hyp\": {\"layer_width\": [width,width, width]}}\n",
    "intensity_state_dict = {\"model_hyp\": {\"layer_width\": [width, width,width],\n",
    "                                                         \"noise\": 0}\n",
    "                            }\n",
    "\n",
    "\n",
    "hyperparameter_dict = {\"state_size\": state_size, \"state_model\": user_state_dict, \n",
    "                           \"intensity_model\": intensity_state_dict,# \"num_recom\" : num_items_per_recom,\n",
    "                            \"noise\": 0.}\n",
    "train_model = Toy_intensity_Generator(hyperparameter_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_train = join(paths.dat, SETTINGS.filepaths_new[\"copy_intensity_model\"]\n",
    "            )\n",
    "#torch.save(gen_model.state_dict(), path_train)\n",
    "#train_model.load_state_dict(torch.load(path_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def eval_single(model, path, scoring_func, num_tries=20):\n",
    "    \n",
    "    last_t = 0\n",
    "    state = torch.zeros((1, state_size))\n",
    "    loss = 0.\n",
    "    #results = []\n",
    "    with torch.no_grad():\n",
    "        for timestep in path:\n",
    "            current_pred = []\n",
    "            for _ in range(num_tries):\n",
    "                out = model.sample_one(state)\n",
    "                current_pred.append(last_t + out)\n",
    "            current_pred = torch.stack(current_pred)\n",
    "            #results.append(torch.mean(current_pred))\n",
    "            loss += scoring_func(current_pred, timestep)\n",
    "            delta = timestep-last_t\n",
    "            last_t = timestep\n",
    "            state = model.evolve_state(state, delta)\n",
    "    \n",
    "    return loss\n",
    "\n",
    "\n",
    "eval_single_partial = partial(eval_single, path=sample_path, scoring_func=utils.energy_score_loss,\n",
    "                              num_tries=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_single_partial(train_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "Single_path_problem = evotorch.neuroevolution.NEProblem(\n",
    "    objective_sense=\"min\",\n",
    "    network= train_model,\n",
    "    #network=Toy_intensity_Generator,\n",
    "    #network_args = {\"hyperparameter_dict\": hyperparameter_dict},\n",
    "    network_eval_func=eval_single_partial,\n",
    "    device=device,\n",
    "    num_actors = 20,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "searcher = PGPE(\n",
    "    Single_path_problem,\n",
    "    popsize=40,\n",
    "    radius_init=5.,\n",
    "    center_learning_rate=0.2,\n",
    "    stdev_learning_rate=0.05,\n",
    ")\n",
    "'''\n",
    "searcher = CMAES(\n",
    "    Single_path_problem,\n",
    "    popsize=20,  \n",
    "    stdev_init= 2.\n",
    ")'''\n",
    "\n",
    "logger = PandasLogger(searcher)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_iterations = 20\n",
    "for _ in tqdm(range(num_iterations), desc=\"Running PGPE\"):\n",
    "    searcher.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1 epoch = 50(population) * 20(samples for expectation) * 1(number of paths) * 10'000(integration/for loop steps, worst case) * (2+1) (neural network calls) *10(number of events)= 300m NN calls per epoch\n",
    "\n",
    "for much worse approximation and still way too much compute:<br>\n",
    "1 epoch = 20(population) * 20(samples for expectation) * 1(number of paths) * 100(integration/for loop steps, worst case) * (2+1) (neural network calls) *10(number of events)= 1.2m NN calls per epoch\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "lstm/rnn  feed last hidden state + actual state to predict probability increase,  stop at  threshhold chosen by uniform \n",
    "\n",
    "\n",
    "neural network(time? + state)  -> time + state\n",
    "\n",
    "noise for time but not state\n",
    "\n",
    "normalizing flows?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.to_dataframe().mean_eval.plot()# test datapoint CMA ES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trained_network = Single_path_problem.parameterize_net(searcher.status[\"center\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_train = join(paths.dat, SETTINGS.filepaths_new[\"copy_intensity_model\"]\n",
    "            )\n",
    "#torch.save(trained_network.state_dict(), path_train)\n",
    "#train_model.load_state_dict(torch.load(path_train))\n",
    "#trained_network = train_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Intensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_sorted = False\n",
    "batchsize = 1  # needs to be 1 for ode\n",
    "model_type = \"simple\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "width=64\n",
    "intensity = {\"model_hyp\": {\"layer_width\": [width for _ in range(3)]}}\n",
    "state_dict = {\"model_hyp\": {\"layer_width\": [width for _ in range(4)],\n",
    "                            \"noise\": 0},\n",
    "            }\n",
    "\n",
    "timecheat = False\n",
    "hyperparameter_dict = {\"state_size\": state_size, \"state_model\": state_dict, \n",
    "        \"intensity_model\": intensity, \"state_model_type\": model_type, # simple\n",
    "        \"time_embedding_size\" :32, \"max_freq\": 20,\n",
    "        }\n",
    "model = Toy_intensity_Comparer(hyperparameter_dict)\n",
    "print(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = TimestepFrequencyDataset(sample_path, num_random_points=150)\n",
    "\n",
    "for i in range(min(len(dataset), 20)):\n",
    "    sample = dataset[i]\n",
    "    print(f\"Timestep: {sample['timestep'].item():.4f}, Frequency: {sample['frequency'].item()}\")\n",
    "\n",
    "dataloader = DataLoader(dataset, batch_size=batchsize, shuffle=not train_sorted\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for batch in dataloader:\n",
    "    print(batch)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps_per_epoch = len(dataset) // batchsize  \n",
    "if len(dataset) % batchsize  != 0:\n",
    "    steps_per_epoch+=1\n",
    "num_epochs = 200\n",
    "warmup_period = steps_per_epoch\n",
    "num_steps = num_epochs*steps_per_epoch - warmup_period\n",
    "num_iter_til_first_restart = num_steps + 1\n",
    "user_lr = 0.001\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "optimizer = optim.AdamW(model.parameters(), lr=user_lr,\n",
    "                        weight_decay=1e-5)\n",
    "lr_scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(\n",
    "    optimizer, T_0=num_iter_til_first_restart, T_mult=1, eta_min=1e-5)\n",
    "\n",
    "warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simtrain.train import train_density\n",
    "\n",
    "def train_density_sorted():\n",
    "    pass\n",
    "\n",
    "def print_res(model, dataloader, criterion,state_size):\n",
    "    for batch in dataloader:\n",
    "        timesteps = batch['timestep'].unsqueeze(1)  # Add batch dimension\n",
    "        frequencies = batch['frequency']\n",
    "        state = torch.zeros((len(timesteps), state_size))\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        outputs = model(state, timesteps)\n",
    "\n",
    "        loss = criterion(outputs, frequencies)  # Remove extra dimension from output\n",
    "        print(f\"loss: {loss} \\tfrequencies: {frequencies} \\t predicted: {outputs}, \\ttime: {timesteps}\")\n",
    "        loss.backward()\n",
    "\n",
    "        #for name, param in model.named_parameters():\n",
    "        #    print(f\"Parameter Name: {name}\")\n",
    "        #    print(f\"Gradients: {param.grad}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simtrain.utils import weighted_mse_loss\n",
    "loss = partial(weighted_mse_loss, weight_pos=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_density(model, dataloader, criterion=weighted_mse_loss, state_size=state_size,\n",
    "                optimizer=optimizer, num_epochs=num_epochs, warmup_scheduler=warmup_scheduler,\n",
    "                loss_print_interval=1, warmup_period=warmup_period, lr_scheduler=lr_scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### visualize training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#dataloader = DataLoader(dataset, batch_size=1, shuffle=not train_sorted) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_res(model, dataloader, weighted_mse_loss, state_size=state_size,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = np.linspace(0, 71, 71*2)  # Adjust the range as needed\n",
    "x_range_tensor = torch.tensor(x_range, dtype=torch.float32).unsqueeze(1)\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    if model_type == \"ode\":\n",
    "        state = torch.zeros((1, state_size))\n",
    "\n",
    "        predictions = []\n",
    "        for el in x_range:\n",
    "            out = model(state, el)\n",
    "            predictions.append(out[0])\n",
    "    else:\n",
    "        state = torch.zeros((len(x_range), state_size))\n",
    "        predictions = model(state, x_range_tensor).numpy()\n",
    "\n",
    "print(f\"area: {np.sum(predictions)*(72/200)}\")\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(x_range, predictions, label='Model Predictions')\n",
    "\n",
    "for i in range(len(dataset)):\n",
    "    sample = dataset[i]\n",
    "    x_pos= sample['timestep'].item()\n",
    "    height = torch.where(sample['frequency']>0, sample['frequency'], .1).item()\n",
    "    plt.plot([x_pos, x_pos], [0, height], linestyle='--', color='red')\n",
    "\n",
    "plt.plot([0, 0], [0, 0], color='red', linestyle='--', alpha=1.0, label='Data Points')\n",
    "\n",
    "#plt.scatter(x_train.numpy(), y_train.numpy(), color='red', label='Training Data')\n",
    "plt.xlabel('Input')\n",
    "plt.ylabel('Output')\n",
    "plt.title('Model Predictions Over Input Range')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sfdsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = np.linspace(-1, 70, 200)  # Adjust the range as needed\n",
    "x_range_tensor = torch.tensor(x_range, dtype=torch.float32).unsqueeze(1)\n",
    "state = torch.zeros((len(x_range), state_size))\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    if model_type == \"ode\":\n",
    "        predictions = []\n",
    "        for el in x_range_tensor:\n",
    "            predictions.append(model(el, x_range_tensor).item())\n",
    "    else:\n",
    "        predictions = model(state, x_range_tensor).numpy()\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(x_range, predictions, label='Model Predictions')\n",
    "\n",
    "for i in range(len(dataset)):\n",
    "    sample = dataset[i]\n",
    "    x_pos= sample['timestep'].item()\n",
    "    height =sample['frequency'].item()+.1\n",
    "    plt.plot([x_pos, x_pos], [0, height], linestyle='--', color='red')\n",
    "\n",
    "#plt.scatter(x_train.numpy(), y_train.numpy(), color='red', label='Training Data')\n",
    "plt.xlabel('Input')\n",
    "plt.ylabel('Output')\n",
    "plt.title('Model Predictions Over Input Range')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = np.linspace(-1, 70, 200)  # Adjust the range as needed\n",
    "x_range_tensor = torch.tensor(x_range, dtype=torch.float32).unsqueeze(1)\n",
    "state = torch.zeros((len(x_range), state_size))\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    predictions = model(state, x_range_tensor).numpy()\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(x_range, predictions, label='Model Predictions')\n",
    "\n",
    "for i in range(len(dataset)):\n",
    "    sample = dataset[i]\n",
    "    x_pos= sample['timestep'].item()\n",
    "    height =sample['frequency'].item()+.1\n",
    "    plt.plot([x_pos, x_pos], [0, height], linestyle='--', color='red')\n",
    "\n",
    "#plt.scatter(x_train.numpy(), y_train.numpy(), color='red', label='Training Data')\n",
    "plt.xlabel('Input')\n",
    "plt.ylabel('Output')\n",
    "plt.title('Model Predictions Over Input Range')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_range = np.linspace(-1, 70, 200)  # Adjust the range as needed\n",
    "x_range_tensor = torch.tensor(x_range, dtype=torch.float32).unsqueeze(1)\n",
    "state = torch.zeros((len(x_range), state_size))\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    predictions = model(state, x_range_tensor).numpy()\n",
    "\n",
    "# Plot the results\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(x_range, predictions, label='Model Predictions')\n",
    "\n",
    "for i in range(len(dataset)):\n",
    "    sample = dataset[i]\n",
    "    x_pos= sample['timestep'].item()\n",
    "    height =sample['frequency'].item()+.1\n",
    "    plt.plot([x_pos, x_pos], [0, height], linestyle='--', color='red')\n",
    "\n",
    "#plt.scatter(x_train.numpy(), y_train.numpy(), color='red', label='Training Data')\n",
    "plt.xlabel('Input')\n",
    "plt.ylabel('Output')\n",
    "plt.title('Model Predictions Over Input Range')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trained_network = train_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def simulate_single_forced(model, path, num_tries=20, state = torch.zeros((1, state_size))):\n",
    "    \n",
    "    last_t = 0\n",
    "    \n",
    "    results = []\n",
    "    with torch.no_grad():\n",
    "        for timestep in path:\n",
    "            current_pred = []\n",
    "            for _ in range(num_tries):\n",
    "                out = model.sample_one(state)\n",
    "                current_pred.append(last_t + out)\n",
    "            current_pred = torch.stack(current_pred)\n",
    "            results.append(torch.mean(current_pred))\n",
    "            print(f\"var(forced): {torch.var(current_pred)}\")\n",
    "            #loss += scoring_func(current_pred, timestep)\n",
    "            delta = timestep-last_t\n",
    "            last_t = timestep\n",
    "            state = model.evolve_state(state, delta)\n",
    "    \n",
    "    return results\n",
    "\n",
    "\n",
    "def simulate_single(model, num_events, num_tries=20, state = torch.zeros((1, state_size))):\n",
    "    \n",
    "    last_t = 0\n",
    "    results = []\n",
    "    with torch.no_grad():\n",
    "        for i in range(num_events):\n",
    "            current_pred = []\n",
    "            for _ in range(num_tries):\n",
    "                out = model.sample_one(state)\n",
    "                current_pred.append(last_t + out)\n",
    "            current_pred = torch.stack(current_pred)\n",
    "            print(f\"var: {torch.var(current_pred)}\")\n",
    "            selected = torch.mean(current_pred)\n",
    "            results.append(selected)\n",
    "            delta = selected-last_t\n",
    "            last_t = selected\n",
    "            state = model.evolve_state(state, delta)\n",
    "    \n",
    "    return results\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simpler nn\n",
    "simulate_single_partial_forced_function_approx = partial(\n",
    "    utils.simulate_single_forced_function_approx, path=sample_path,\n",
    "                              num_tries=1, timecheat=timecheat, state_size=state_size)\n",
    "simulate_single_partial_function_approx =partial(\n",
    "    utils.simulate_single_function_approx, num_events =len(sample_path),\n",
    "                              num_tries=1, timecheat=timecheat, state_size=state_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# integrals\n",
    "simulate_single_partial_forced = partial(simulate_single_forced, path=sample_path, \n",
    "                              num_tries=2)\n",
    "\n",
    "simulate_single_partial = partial(simulate_single, num_events=12,\n",
    "                              num_tries=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# intensity\n",
    "example_out_forced = simulate_single_partial_forced(trained_network)\n",
    "example_out = simulate_single_partial(trained_network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#mlp\n",
    "example_out_forced = simulate_single_partial_forced_function_approx(trained_network)\n",
    "example_out = simulate_single_partial_function_approx(trained_network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Example data: replace these with your actual time series data\n",
    "time_series_1 = sample_path # Timestamps for the first time series\n",
    "time_series_2 = torch.clamp(torch.as_tensor(example_out_forced),0,70).detach().numpy()  # Timestamps for the second time series\n",
    "time_series_3 = torch.clamp(torch.as_tensor(example_out), 0, 70).detach().numpy()  # Timestamps for the second time series\n",
    "\n",
    "\n",
    "# Create a figure and axis\n",
    "fig, ax = plt.subplots(figsize=(8, 5))\n",
    "\n",
    "# Plot the first time series\n",
    "ax.scatter(time_series_1, [1] * len(time_series_1), color='blue', label='Ground Truth', s=10, marker='o')\n",
    "\n",
    "# Plot the second time series\n",
    "ax.scatter(time_series_2, [2] * len(time_series_2), color='red', label='Simmulation Forced', s=10, marker='x')\n",
    "\n",
    "ax.scatter(time_series_3, [3] * len(time_series_3), color='green', label='Simmulation Free', s=10, marker='x')\n",
    "\n",
    "\n",
    "# Add labels, legend, and grid\n",
    "ax.set_xlabel('Time')\n",
    "ax.set_yticks([1, 2, 3])\n",
    "ax.set_yticklabels(['Ground Truth', 'Simmulation Forced', \"Simmulation Free\"])\n",
    "ax.set_title('Comparison of Two Time Series')\n",
    "ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), frameon=False)\n",
    "\n",
    "# Adjust subplot parameters to make room for the legend\n",
    "plt.subplots_adjust(right=0.75)\n",
    "ax.grid(True)\n",
    "\n",
    "# Display the plot\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "WW",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}