ctm_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "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 scipy.signal import convolve2d\n",
    "from matplotlib.animation import FuncAnimation\n",
    "import matplotlib.animation as animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#edge padding \n",
    "def edge_pad_3x3(arr):\n",
    "    \"\"\"\n",
    "    Pad an array for a 3x3 kernel. This function adds one layer of padding around the array.\n",
    "    \"\"\"\n",
    "    # Pad top and bottom with 1 row\n",
    "    top_bottom_padded = np.vstack([arr[0, :], arr, arr[-1, :]])\n",
    "    # Pad left and right with 1 column\n",
    "    fully_padded = np.hstack([top_bottom_padded[:, 0:1], top_bottom_padded, top_bottom_padded[:, -1:]])\n",
    "    return fully_padded\n",
    "\n",
    "def edge_pad_5x5(arr):\n",
    "    \"\"\"\n",
    "    Pad an array for a 5x5 kernel. This function adds two layers of padding around the array.\n",
    "    \"\"\"\n",
    "    # Pad top and bottom with 2 rows\n",
    "    top_bottom_padded = np.vstack([arr[0, :], arr[0, :], arr, arr[-1, :], arr[-1, :]])\n",
    "    # Pad left and right with 2 columns\n",
    "    fully_padded = np.hstack([top_bottom_padded[:, 0:1], top_bottom_padded[:, 0:1], \n",
    "                              top_bottom_padded, top_bottom_padded[:, -1:], top_bottom_padded[:, -1:]])\n",
    "    return fully_padded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CTM:\n",
    "    def __init__(self):\n",
    "        # Configurations\n",
    "        self.dim = (20, 20)\n",
    "        self.max_steps = 100  # Maximum number of steps in the simulation\n",
    "        self.state = np.zeros(self.dim)  # Initialize the state grid\n",
    "        self.kernel = np.random.rand(3, 3)*2-1  # Initialize the kernel\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 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",
    "                padded_state = edge_pad_3x3(self.state)\n",
    "                ds = convolve2d(padded_state, self.kernel, mode='valid')\n",
    "                self.state += ds * self.dt\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=-3, vmax=3)\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage:\n",
    "# Initialize the model\n",
    "model = CTM()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[114474579.86280113, -1451049535.348981]\n"
     ]
    }
   ],
   "source": [
    "# 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
}