test_wcst_session_extended_dim.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the current impl of WCST session with an extended feature dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from wcst_session import WcstSession\n",
    "from block_switching_conditions import short_condition, create_bernoulli_condition\n",
    "from card_generators import RandomCardGenerator\n",
    "from rule_generators import RandomRuleGeneratorHuman\n",
    "from constants import EXTENDED_DIM_FEATURE_NAMES, EXTENEDED_DIM_NAMES\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "sess = WcstSession(\n",
    "    correct_value=1, \n",
    "    incorrect_value=0, \n",
    "    block_switching_condition=short_condition, \n",
    "    card_generator=RandomCardGenerator(None, num_cards=4, num_dims=4),\n",
    "    rule_generator=RandomRuleGeneratorHuman(None, num_rules=16, num_dims=4),\n",
    "    feature_names=EXTENDED_DIM_FEATURE_NAMES,\n",
    "    dim_names=EXTENEDED_DIM_NAMES,\n",
    "    random_seed=42\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['SQUARE', 'MAGENTA', 'SWIRL', 'SOLID'],\n",
       "       ['STAR', 'YELLOW', 'POLKADOT', 'DOTTED'],\n",
       "       ['TRIANGLE', 'GREEN', 'ESCHER', 'DASHED'],\n",
       "       ['CIRCLE', 'CYAN', 'RIPPLE', 'DASHDOT']], dtype='<U8')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.get_cards_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(False, 0)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.make_selection(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'SOLID'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.get_current_rule_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['TRIANGLE', 'MAGENTA', 'SWIRL', 'DASHDOT'],\n",
       "       ['STAR', 'YELLOW', 'POLKADOT', 'DASHED'],\n",
       "       ['CIRCLE', 'CYAN', 'RIPPLE', 'DOTTED'],\n",
       "       ['SQUARE', 'GREEN', 'ESCHER', 'SOLID']], dtype='<U8')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.get_cards_text()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(True, 1)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.make_selection(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TrialNumber</th>\n",
       "      <th>BlockNumber</th>\n",
       "      <th>TrialAfterRuleChange</th>\n",
       "      <th>Response</th>\n",
       "      <th>ItemChosen</th>\n",
       "      <th>CurrentRule</th>\n",
       "      <th>Reward</th>\n",
       "      <th>Item0Shape</th>\n",
       "      <th>Item0Color</th>\n",
       "      <th>Item0Pattern</th>\n",
       "      <th>...</th>\n",
       "      <th>Item1Pattern</th>\n",
       "      <th>Item1Outline</th>\n",
       "      <th>Item2Shape</th>\n",
       "      <th>Item2Color</th>\n",
       "      <th>Item2Pattern</th>\n",
       "      <th>Item2Outline</th>\n",
       "      <th>Item3Shape</th>\n",
       "      <th>Item3Color</th>\n",
       "      <th>Item3Pattern</th>\n",
       "      <th>Item3Outline</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>Incorrect</td>\n",
       "      <td>2</td>\n",
       "      <td>SOLID</td>\n",
       "      <td>0</td>\n",
       "      <td>SQUARE</td>\n",
       "      <td>MAGENTA</td>\n",
       "      <td>SWIRL</td>\n",
       "      <td>...</td>\n",
       "      <td>POLKADOT</td>\n",
       "      <td>DOTTED</td>\n",
       "      <td>TRIANGLE</td>\n",
       "      <td>GREEN</td>\n",
       "      <td>ESCHER</td>\n",
       "      <td>DASHED</td>\n",
       "      <td>CIRCLE</td>\n",
       "      <td>CYAN</td>\n",
       "      <td>RIPPLE</td>\n",
       "      <td>DASHDOT</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>Correct</td>\n",
       "      <td>3</td>\n",
       "      <td>SOLID</td>\n",
       "      <td>1</td>\n",
       "      <td>TRIANGLE</td>\n",
       "      <td>MAGENTA</td>\n",
       "      <td>SWIRL</td>\n",
       "      <td>...</td>\n",
       "      <td>POLKADOT</td>\n",
       "      <td>DASHED</td>\n",
       "      <td>CIRCLE</td>\n",
       "      <td>CYAN</td>\n",
       "      <td>RIPPLE</td>\n",
       "      <td>DOTTED</td>\n",
       "      <td>SQUARE</td>\n",
       "      <td>GREEN</td>\n",
       "      <td>ESCHER</td>\n",
       "      <td>SOLID</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>2 rows × 23 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   TrialNumber  BlockNumber  TrialAfterRuleChange   Response  ItemChosen  \\\n",
       "0            0            0                     0  Incorrect           2   \n",
       "1            1            0                     1    Correct           3   \n",
       "\n",
       "  CurrentRule  Reward Item0Shape Item0Color Item0Pattern  ... Item1Pattern  \\\n",
       "0       SOLID       0     SQUARE    MAGENTA        SWIRL  ...     POLKADOT   \n",
       "1       SOLID       1   TRIANGLE    MAGENTA        SWIRL  ...     POLKADOT   \n",
       "\n",
       "  Item1Outline Item2Shape Item2Color Item2Pattern Item2Outline Item3Shape  \\\n",
       "0       DOTTED   TRIANGLE      GREEN       ESCHER       DASHED     CIRCLE   \n",
       "1       DASHED     CIRCLE       CYAN       RIPPLE       DOTTED     SQUARE   \n",
       "\n",
       "  Item3Color Item3Pattern Item3Outline  \n",
       "0       CYAN       RIPPLE      DASHDOT  \n",
       "1      GREEN       ESCHER        SOLID  \n",
       "\n",
       "[2 rows x 23 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.dump_history()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check if gym env works as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "from gymnasium import spaces\n",
    "import wcst_gym\n",
    "from wcst_gym.n_hot_observation import NHotObservation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"WcstSession-v0\", wcst_session=sess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[ 1,  7, 10, 14],\n",
       "        [ 0,  4,  9, 13],\n",
       "        [ 2,  6, 11, 15],\n",
       "        [ 3,  5,  8, 12]]),\n",
       " {'current_rule': 3})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.reset(seed=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 4)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.observation_space.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[ 0,  6, 10, 14],\n",
       "        [ 2,  4,  8, 15],\n",
       "        [ 3,  5, 11, 12],\n",
       "        [ 1,  7,  9, 13]]),\n",
       " 1,\n",
       " False,\n",
       " False,\n",
       " {'current_rule': 3})"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.step(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = NHotObservation(env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0],\n",
       "        [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0],\n",
       "        [0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0],\n",
       "        [1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1]], dtype=int8),\n",
       " {'current_rule': 7})"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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": 2
}