snm_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "import time\n",
    "from scipy import optimize\n",
    "from matplotlib.animation import FuncAnimation\n",
    "import matplotlib.animation as animation\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KernelNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(KernelNetwork, self).__init__()\n",
    "        # Define the layers\n",
    "        self.fc1 = nn.Linear(9, 32)  # Input layer\n",
    "        self.fc2 = nn.Linear(32, 9)  # Output layer\n",
    "\n",
    "        # Custom weight initialization\n",
    "        self.init_weights()\n",
    "\n",
    "    def init_weights(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                # Initialize weights to a higher value, e.g., using a uniform distribution\n",
    "                nn.init.normal_(m.weight, mean=0, std=0.3)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Flatten the input 3x3 grid\n",
    "        x = x.view(-1, 9)\n",
    "        # Apply layers\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        # Reshape output to 3x3 grid\n",
    "        x = x.view(-1, 3, 3)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SNM:\n",
    "    def __init__(self):\n",
    "        # Configurations\n",
    "        self.dim = (20, 20)\n",
    "        self.max_steps = 200  # Maximum number of steps in the simulation\n",
    "        self.kernel_net=KernelNetwork()\n",
    "        self.kernel_pos=(1,1)\n",
    "        self.state = np.zeros(self.dim)  # Initialize the state grid\n",
    "        self.halting_threshold = 0.7\n",
    "        self.dt=0.1\n",
    "        self.time_idx=0\n",
    "\n",
    "        self.state_history=[]\n",
    "\n",
    "        # Significant cells\n",
    "        self.input_indexes = [(3, 0), (6, 0)]\n",
    "        self.output_indexes = [(3, 19), (6, 19)]\n",
    "        self.halting_index = (10, 19)\n",
    "        self.reward_index = (15, 19)\n",
    "\n",
    "    def set_input(self,input_values):\n",
    "        for idx, value in zip(self.input_indexes, input_values):\n",
    "            self.state[idx] = value\n",
    "\n",
    "    def set_reward(self,reward):\n",
    "        self.state[self.reward_index] = reward\n",
    "\n",
    "    def get_output(self):\n",
    "        return [self.state[idx] for idx in self.output_indexes]\n",
    "    \n",
    "    def get_halting(self):\n",
    "        return self.state[self.halting_index]\n",
    "    \n",
    "    def reset_halting(self):\n",
    "        # Reset the halting cell\n",
    "        self.state[self.halting_index] = 0\n",
    "\n",
    "        # Define the halving kernel\n",
    "        halving_kernel = np.array([[0.3, 0.3, 0.3],\n",
    "                                   [0.3, 0.0, 0.3],\n",
    "                                   [0.3, 0.3, 0.3]])\n",
    "\n",
    "        # Get the coordinates of the halting cell\n",
    "        y, x = self.halting_index\n",
    "\n",
    "        # Determine the slicing bounds, ensuring they are within the grid boundaries\n",
    "        y_start, y_end = max(y - 1, 0), min(y + 2, self.dim[0])\n",
    "        x_start, x_end = max(x - 1, 0), min(x + 2, self.dim[1])\n",
    "\n",
    "        # Adjust the kernel size if slicing bounds are at the edges\n",
    "        kernel_y_start, kernel_x_start = max(1 - y, 0), max(1 - x, 0)\n",
    "        kernel_y_end, kernel_x_end = 3 - max(y + 2 - self.dim[0], 0), 3 - max(x + 2 - self.dim[1], 0)\n",
    "\n",
    "        # Apply the halving kernel to the surrounding cells\n",
    "        self.state[y_start:y_end, x_start:x_end] *= halving_kernel[kernel_y_start:kernel_y_end, kernel_x_start:kernel_x_end]\n",
    "\n",
    "\n",
    "    def reset_recording(self):\n",
    "        self.state_history=[]\n",
    "\n",
    "    def save_state(self):\n",
    "        self.state_history.append(self.state.copy())\n",
    "\n",
    "    \n",
    "    def step_kernel(self):\n",
    "        # Get current kernel position\n",
    "        y,x = self.kernel_pos\n",
    "\n",
    "        # Extract the current 3x3 region\n",
    "        kernel_region = self.state[y-1:y+2, x-1:x+2]\n",
    "\n",
    "        # Convert to PyTorch tensor and process through the kernel network\n",
    "        kernel_input = torch.tensor(kernel_region, dtype=torch.float32).unsqueeze(0)\n",
    "        kernel_output = self.kernel_net(kernel_input).squeeze(0).detach() #warning: maybe needed to modify \"detach\" for RL\n",
    "\n",
    "        # Assuming the mapping: 'u' (1), 'd' (7), 'l' (3), 'r' (5), 's' (4)\n",
    "        movement_indices = torch.tensor([1, 7, 3, 5, 4])\n",
    "        movement_values = kernel_output.flatten()[movement_indices]\n",
    "\n",
    "        # Apply softmax to get movement probabilities\n",
    "        movement_probabilities = F.softmax(movement_values, dim=0).numpy()\n",
    "\n",
    "        # Movement directions corresponding to indices\n",
    "        directions = ['u', 'd', 'l', 'r', 's']\n",
    "        movement = np.random.choice(directions, p=movement_probabilities)\n",
    "\n",
    "        # Update kernel position based on selected movement\n",
    "        if movement == 'u' and y > 1:\n",
    "            y -= 1\n",
    "        elif movement == 'd' and y < self.dim[0] - 2:\n",
    "            y += 1\n",
    "        elif movement == 'l' and x > 1:\n",
    "            x -= 1\n",
    "        elif movement == 'r' and x < self.dim[1] - 2:\n",
    "            x += 1\n",
    "        # No movement for 's'\n",
    "\n",
    "        self.kernel_pos = (y, x)\n",
    "\n",
    "        # Update the state space with the output of the kernel\n",
    "        self.state[y-1:y+2, x-1:x+2] = kernel_output.numpy()\n",
    "\n",
    "\n",
    "    def forward(self, input_values, reward):\n",
    "            if len(input_values) != len(self.input_indexes):\n",
    "                raise ValueError(\"Input values size does not match input indexes size.\")\n",
    "\n",
    "            #reset halting\n",
    "            self.reset_halting()\n",
    "\n",
    "        \n",
    "            for step in range(self.max_steps):\n",
    "                \n",
    "                self.set_input(input_values)\n",
    "                self.set_reward(reward)\n",
    "\n",
    "                # Add the current state\n",
    "                self.save_state()\n",
    "\n",
    "                # Apply convolution with edge padding\n",
    "                self.step_kernel()\n",
    "\n",
    "                # Check halting condition\n",
    "                if self.get_halting() >= self.halting_threshold:\n",
    "                    self.save_state()\n",
    "                    break\n",
    "\n",
    "            # Retrieve output values\n",
    "            outputs = self.get_output()\n",
    "            \n",
    "            return outputs\n",
    "    \n",
    "    def visualize(self):\n",
    "        # Set up the figure for animation\n",
    "        fig, ax = plt.subplots()\n",
    "        ims = []\n",
    "\n",
    "        for state in self.state_history:\n",
    "            # Create the heatmap from the state\n",
    "            im = ax.imshow(state, animated=True, cmap='jet', vmin=-1, vmax=1)\n",
    "\n",
    "            # Function to draw border around a cell\n",
    "            def draw_border(y, x, color):\n",
    "                rect = plt.Rectangle((x-0.5, y-0.5), 1, 1, fill=False, edgecolor=color, lw=4)\n",
    "                ax.add_patch(rect)\n",
    "\n",
    "            # Draw borders around the special cells\n",
    "            for y, x in self.input_indexes:\n",
    "                draw_border(y, x, 'green')  # Green for input\n",
    "            for y, x in self.output_indexes:\n",
    "                draw_border(y, x, 'blue')  # Blue for output\n",
    "            y, x = self.halting_index\n",
    "            draw_border(y, x, 'red')  # Red for halting\n",
    "            y, x = self.reward_index\n",
    "            draw_border(y, x, 'orange')  # Orange for reward\n",
    "\n",
    "            ims.append([im])\n",
    "\n",
    "        # Create the animation\n",
    "        ani = animation.ArtistAnimation(fig, ims, interval=200, blit=True)\n",
    "\n",
    "        # Add a colorbar\n",
    "        fig.colorbar(im, ax=ax)\n",
    "\n",
    "        # Save the animation to a file\n",
    "        ani.save('simulation.mp4', writer='ffmpeg')\n",
    "\n",
    "        # Close the plot to prevent it from displaying inline or in a new window\n",
    "        plt.close(fig)\n",
    "\n",
    "    def visualize2(self): #for visualizing also the kernel, not implemented yet\n",
    "        # Set up the figure for animation\n",
    "        fig, ax = plt.subplots()\n",
    "        ims = []\n",
    "\n",
    "        for state in self.state_history:\n",
    "            # Create the heatmap from the state\n",
    "            im = ax.imshow(state, animated=True, cmap='jet', vmin=-1, vmax=1)\n",
    "\n",
    "            # Function to draw a large square around the kernel\n",
    "            def draw_kernel_border(y, x, color, width=2):\n",
    "                # Draw a square border of 3x3 around the kernel's center\n",
    "                rect = plt.Rectangle((x-1, y-1), 3, 3, fill=False, edgecolor=color, lw=width)\n",
    "                ax.add_patch(rect)\n",
    "\n",
    "            # Draw borders around the special cells\n",
    "            for y, x in self.input_indexes:\n",
    "                draw_kernel_border(y, x, 'green')  # Green for input\n",
    "            for y, x in self.output_indexes:\n",
    "                draw_kernel_border(y, x, 'blue')  # Blue for output\n",
    "            y, x = self.halting_index\n",
    "            draw_kernel_border(y, x, 'red')  # Red for halting\n",
    "            y, x = self.reward_index\n",
    "            draw_kernel_border(y, x, 'orange')  # Orange for reward\n",
    "\n",
    "            # Get the kernel's current position and draw the border\n",
    "            kernel_y, kernel_x = self.kernel_pos\n",
    "            draw_kernel_border(kernel_y, kernel_x, 'white')\n",
    "\n",
    "            ims.append([im])\n",
    "\n",
    "        # Create the animation\n",
    "        ani = animation.ArtistAnimation(fig, ims, interval=200, blit=True)\n",
    "\n",
    "        # Add a colorbar\n",
    "        fig.colorbar(im, ax=ax)\n",
    "\n",
    "        # Save the animation to a file\n",
    "        ani.save('simulation.mp4', writer='ffmpeg')\n",
    "\n",
    "        # Close the plot to prevent it from displaying inline or in a new window\n",
    "        plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0, 0.0]\n"
     ]
    }
   ],
   "source": [
    "model = SNM()\n",
    "# Example input and reward values\n",
    "input_values = [1.0, -1.0]  # Example input values\n",
    "reward = 0.5  # Example reward value\n",
    "\n",
    "# Perform the forward pass\n",
    "model.reset_recording()\n",
    "output_values = model.forward(input_values, reward)\n",
    "model.visualize()\n",
    "print(output_values)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}