{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d6f169d2",
   "metadata": {},
   "source": [
    "# Vector store properties and metrics\n",
    "Quantify vector store properties and performance\n",
    "\n",
    "E.M.F. August 2025\n",
    "\n",
    "1. Read in vector store\n",
    "2. What are properties of the vector store?\n",
    "\n",
    "    - number of vectors\n",
    "    - distribution of distances\n",
    "    - clustering?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6711a068",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preliminaries\n",
    "\n",
    "from llama_index.core import (\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    "    ServiceContext,\n",
    "    Settings,\n",
    ")\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "# Load embedding model (same as used for vector store)\n",
    "# The problem for offline work is that this attempts to download\n",
    "# Need to find a local storage\n",
    "\n",
    "# For online use:\n",
    "# embed_model = HuggingFaceEmbedding(model_name=\"all-mpnet-base-v2\")\n",
    "# embed_model = HuggingFaceEmbedding(model_name=\"BAAI/bge-large-en-v1.5\")\n",
    "\n",
    "# For offline use (run once with local_files_only=False to cache the model)\n",
    "#embed_model = HuggingFaceEmbedding(model_name_or_path=\"/home/furst/.cache/huggingface/hub/sentence-transformers/all-mpnet-base-v2\")\n",
    "#embed_model = HuggingFaceEmbedding(cache_folder=\"./models\",model_name=\"all-mpnet-base-v2\",local_files_only=True)\n",
    "embed_model = HuggingFaceEmbedding(cache_folder=\"./models\",model_name=\"BAAI/bge-large-en-v1.5\",local_files_only=False)\n",
    "\n",
    "Settings.embed_model = embed_model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "77c2ba93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading llama_index.core.storage.kvstore.simple_kvstore from ./storage_exp/docstore.json.\n",
      "Loading llama_index.core.storage.kvstore.simple_kvstore from ./storage_exp/index_store.json.\n"
     ]
    }
   ],
   "source": [
    "# Load persisted vector store + metadata\n",
    "storage_context = StorageContext.from_defaults(persist_dir=\"./storage_exp\")\n",
    "index = load_index_from_storage(storage_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fd7e554",
   "metadata": {},
   "source": [
    "## Vector store details\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93acab2e",
   "metadata": {},
   "source": [
    "Show the index type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84eb17b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "llama_index.core.indices.vector_store.base.VectorStoreIndex"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(index)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a274c1d",
   "metadata": {},
   "source": [
    "These are all of the attributes and methods of the class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6f429f2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__abstractmethods__',\n",
       " '__annotations__',\n",
       " '__class__',\n",
       " '__class_getitem__',\n",
       " '__delattr__',\n",
       " '__dict__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__eq__',\n",
       " '__format__',\n",
       " '__ge__',\n",
       " '__getattribute__',\n",
       " '__getstate__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__le__',\n",
       " '__lt__',\n",
       " '__module__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__orig_bases__',\n",
       " '__parameters__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__repr__',\n",
       " '__setattr__',\n",
       " '__sizeof__',\n",
       " '__slots__',\n",
       " '__str__',\n",
       " '__subclasshook__',\n",
       " '__weakref__',\n",
       " '_abc_impl',\n",
       " '_add_nodes_to_index',\n",
       " '_adelete_from_docstore',\n",
       " '_adelete_from_index_struct',\n",
       " '_aget_node_with_embedding',\n",
       " '_async_add_nodes_to_index',\n",
       " '_build_index_from_nodes',\n",
       " '_callback_manager',\n",
       " '_delete_from_docstore',\n",
       " '_delete_from_index_struct',\n",
       " '_delete_node',\n",
       " '_docstore',\n",
       " '_embed_model',\n",
       " '_get_node_with_embedding',\n",
       " '_graph_store',\n",
       " '_index_struct',\n",
       " '_insert',\n",
       " '_insert_batch_size',\n",
       " '_object_map',\n",
       " '_show_progress',\n",
       " '_storage_context',\n",
       " '_store_nodes_override',\n",
       " '_transformations',\n",
       " '_use_async',\n",
       " '_validate_serializable',\n",
       " '_vector_store',\n",
       " 'adelete_nodes',\n",
       " 'adelete_ref_doc',\n",
       " 'ainsert',\n",
       " 'ainsert_nodes',\n",
       " 'arefresh_ref_docs',\n",
       " 'as_chat_engine',\n",
       " 'as_query_engine',\n",
       " 'as_retriever',\n",
       " 'aupdate_ref_doc',\n",
       " 'build_index_from_nodes',\n",
       " 'delete',\n",
       " 'delete_nodes',\n",
       " 'delete_ref_doc',\n",
       " 'docstore',\n",
       " 'from_documents',\n",
       " 'from_vector_store',\n",
       " 'index_id',\n",
       " 'index_struct',\n",
       " 'index_struct_cls',\n",
       " 'insert',\n",
       " 'insert_nodes',\n",
       " 'ref_doc_info',\n",
       " 'refresh',\n",
       " 'refresh_ref_docs',\n",
       " 'set_index_id',\n",
       " 'storage_context',\n",
       " 'summary',\n",
       " 'update',\n",
       " 'update_ref_doc',\n",
       " 'vector_store']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# list all attributes and methods of the class\n",
    "dir(index)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a607abff",
   "metadata": {},
   "source": [
    "Here are some of the attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "be8205a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "55cdda60-5f8d-4f57-b1cc-8be97f4c4e30\n",
      "None\n"
     ]
    }
   ],
   "source": [
    "print(index.index_id)\n",
    "print(index.summary)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35d11dbe",
   "metadata": {},
   "source": [
    "The vector store is a part of the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8e9da9fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__abstractmethods__',\n",
       " '__annotations__',\n",
       " '__class__',\n",
       " '__class_getitem__',\n",
       " '__class_vars__',\n",
       " '__copy__',\n",
       " '__deepcopy__',\n",
       " '__delattr__',\n",
       " '__dict__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__eq__',\n",
       " '__fields__',\n",
       " '__fields_set__',\n",
       " '__format__',\n",
       " '__ge__',\n",
       " '__get_pydantic_core_schema__',\n",
       " '__get_pydantic_json_schema__',\n",
       " '__getattr__',\n",
       " '__getattribute__',\n",
       " '__getstate__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__iter__',\n",
       " '__le__',\n",
       " '__lt__',\n",
       " '__module__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__pretty__',\n",
       " '__private_attributes__',\n",
       " '__pydantic_complete__',\n",
       " '__pydantic_computed_fields__',\n",
       " '__pydantic_core_schema__',\n",
       " '__pydantic_custom_init__',\n",
       " '__pydantic_decorators__',\n",
       " '__pydantic_extra__',\n",
       " '__pydantic_fields__',\n",
       " '__pydantic_fields_set__',\n",
       " '__pydantic_generic_metadata__',\n",
       " '__pydantic_init_subclass__',\n",
       " '__pydantic_parent_namespace__',\n",
       " '__pydantic_post_init__',\n",
       " '__pydantic_private__',\n",
       " '__pydantic_root_model__',\n",
       " '__pydantic_serializer__',\n",
       " '__pydantic_setattr_handlers__',\n",
       " '__pydantic_validator__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__replace__',\n",
       " '__repr__',\n",
       " '__repr_args__',\n",
       " '__repr_name__',\n",
       " '__repr_recursion__',\n",
       " '__repr_str__',\n",
       " '__rich_repr__',\n",
       " '__setattr__',\n",
       " '__setstate__',\n",
       " '__signature__',\n",
       " '__sizeof__',\n",
       " '__slots__',\n",
       " '__str__',\n",
       " '__subclasshook__',\n",
       " '__weakref__',\n",
       " '_abc_impl',\n",
       " '_calculate_keys',\n",
       " '_copy_and_set_values',\n",
       " '_data',\n",
       " '_get_value',\n",
       " '_iter',\n",
       " '_setattr_handler',\n",
       " 'aclear',\n",
       " 'add',\n",
       " 'adelete',\n",
       " 'adelete_nodes',\n",
       " 'aget_nodes',\n",
       " 'aquery',\n",
       " 'async_add',\n",
       " 'class_name',\n",
       " 'clear',\n",
       " 'client',\n",
       " 'construct',\n",
       " 'copy',\n",
       " 'custom_model_dump',\n",
       " 'data',\n",
       " 'delete',\n",
       " 'delete_nodes',\n",
       " 'dict',\n",
       " 'from_dict',\n",
       " 'from_json',\n",
       " 'from_namespaced_persist_dir',\n",
       " 'from_orm',\n",
       " 'from_persist_dir',\n",
       " 'from_persist_path',\n",
       " 'get',\n",
       " 'get_nodes',\n",
       " 'is_embedding_query',\n",
       " 'json',\n",
       " 'model_computed_fields',\n",
       " 'model_config',\n",
       " 'model_construct',\n",
       " 'model_copy',\n",
       " 'model_dump',\n",
       " 'model_dump_json',\n",
       " 'model_extra',\n",
       " 'model_fields',\n",
       " 'model_fields_set',\n",
       " 'model_json_schema',\n",
       " 'model_parametrized_name',\n",
       " 'model_post_init',\n",
       " 'model_rebuild',\n",
       " 'model_validate',\n",
       " 'model_validate_json',\n",
       " 'model_validate_strings',\n",
       " 'parse_file',\n",
       " 'parse_obj',\n",
       " 'parse_raw',\n",
       " 'persist',\n",
       " 'query',\n",
       " 'schema',\n",
       " 'schema_json',\n",
       " 'stores_text',\n",
       " 'to_dict',\n",
       " 'to_json',\n",
       " 'update_forward_refs',\n",
       " 'validate']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Access the vector store\n",
    "\n",
    "vector_store = index.vector_store\n",
    "dir(vector_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942b96f4",
   "metadata": {},
   "source": [
    "In this case, it is a SimpleVectorStore. Other vector stores that use databases are possible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b699c703",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "llama_index.core.vector_stores.simple.SimpleVectorStore"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(vector_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb522c1",
   "metadata": {},
   "source": [
    "The actual vector store is part of the SimpleVectorStore.data attribute. It has attributes and methods that are listed. Among these aret the embedded_dict, which is the raw vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "776e1e60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['__abstractmethods__', '__annotations__', '__class__', '__dataclass_fields__', '__dataclass_params__', '__delattr__', '__dict__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__getstate__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__match_args__', '__module__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__slots__', '__str__', '__subclasshook__', '__weakref__', '_abc_impl', 'dataclass_json_config', 'embedding_dict', 'from_dict', 'from_json', 'metadata_dict', 'schema', 'text_id_to_ref_doc_id', 'to_dict', 'to_json']\n"
     ]
    }
   ],
   "source": [
    "vs_data = vector_store.data\n",
    "print(dir(vs_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76657adb",
   "metadata": {},
   "source": [
    "We can find the number of embedded vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d1b9d1d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of embedded vectors: 4907\n"
     ]
    }
   ],
   "source": [
    "# Get the embedding dictionary: node_id -> embedding vector\n",
    "embedded_node_dict = vs_data.embedding_dict\n",
    "\n",
    "# Count the embedded vectors (i.e., nodes)\n",
    "print(f\"Number of embedded vectors: {len(embedded_node_dict)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85addeb7",
   "metadata": {},
   "source": [
    "List the first five nodes in the embedding_node_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ef8e24e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node ID: e9d9d86c-bce1-4d09-9d4e-5adc2cb64c39, Embedding dim: 1024\n",
      "Node ID: b9e70a22-8a2b-48dd-97af-6f838952de0f, Embedding dim: 1024\n",
      "Node ID: 347dba75-6ad8-4cae-a7cf-3f0edf347954, Embedding dim: 1024\n",
      "Node ID: afe2889c-b5b7-45ba-b3c8-28c26bdb0b2c, Embedding dim: 1024\n",
      "Node ID: 7428857b-8eef-43e2-82a1-e17c86402aef, Embedding dim: 1024\n"
     ]
    }
   ],
   "source": [
    "for node_id, embedding in list(embedded_node_dict.items())[:5]:\n",
    "    print(f\"Node ID: {node_id}, Embedding dim: {len(embedding)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5edf3f2f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(embedded_node_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aa5dd08",
   "metadata": {},
   "source": [
    "Since embedded_node_dict is a standard Python dict, I can use typical dictionary methods to inspect keys, values, or items:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5aa0d3",
   "metadata": {},
   "source": [
    "The keys are the node IDs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7df6236d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node ID: e9d9d86c-bce1-4d09-9d4e-5adc2cb64c39\n",
      "Node ID: b9e70a22-8a2b-48dd-97af-6f838952de0f\n",
      "Node ID: 347dba75-6ad8-4cae-a7cf-3f0edf347954\n",
      "Node ID: afe2889c-b5b7-45ba-b3c8-28c26bdb0b2c\n",
      "Node ID: 7428857b-8eef-43e2-82a1-e17c86402aef\n"
     ]
    }
   ],
   "source": [
    "for node_id in list(embedded_node_dict.keys())[:5]:\n",
    "    print(f\"Node ID: {node_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fc01dd7",
   "metadata": {},
   "source": [
    "The values are the vectors. Here we see they have a length 1024!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "969a05c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding (dim=1024): [0.015274446457624435, -0.005610795691609383, 0.011983891017735004, 0.020918143913149834, -0.03160902485251427]...\n",
      "Embedding (dim=1024): [-0.006667741108685732, -0.04432462900876999, -0.0077421520836651325, 0.03053889609873295, -0.0067572263069450855]...\n",
      "Embedding (dim=1024): [0.018348904326558113, 0.02425410784780979, 0.00039119378197938204, 0.06972303986549377, -0.02752634882926941]...\n",
      "Embedding (dim=1024): [0.0035999149549752474, 0.03229404613375664, 0.015970565378665924, 0.03491900861263275, -0.018849430605769157]...\n",
      "Embedding (dim=1024): [0.002958822064101696, 0.0470910407602787, -0.0005777133628726006, 0.030806731432676315, -0.03956307843327522]...\n"
     ]
    }
   ],
   "source": [
    "for embedding in list(embedded_node_dict.values())[:5]:\n",
    "    print(f\"Embedding (dim={len(embedding)}): {embedding[:5]}...\")  # Show first 5 values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4572d162",
   "metadata": {},
   "source": [
    "Some code to return previews of the keys and values. Each key and value is an item in the dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3da19f68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Node ID: e9d9d86c-bce1-4d09-9d4e-5adc2cb64c39\n",
      "Embedding length: 1024\n",
      "Embedding preview: [0.015274446457624435, -0.005610795691609383, 0.011983891017735004, 0.020918143913149834, -0.03160902485251427]\n",
      "---\n",
      "Node ID: b9e70a22-8a2b-48dd-97af-6f838952de0f\n",
      "Embedding length: 1024\n",
      "Embedding preview: [-0.006667741108685732, -0.04432462900876999, -0.0077421520836651325, 0.03053889609873295, -0.0067572263069450855]\n",
      "---\n",
      "Node ID: 347dba75-6ad8-4cae-a7cf-3f0edf347954\n",
      "Embedding length: 1024\n",
      "Embedding preview: [0.018348904326558113, 0.02425410784780979, 0.00039119378197938204, 0.06972303986549377, -0.02752634882926941]\n",
      "---\n",
      "Node ID: afe2889c-b5b7-45ba-b3c8-28c26bdb0b2c\n",
      "Embedding length: 1024\n",
      "Embedding preview: [0.0035999149549752474, 0.03229404613375664, 0.015970565378665924, 0.03491900861263275, -0.018849430605769157]\n",
      "---\n",
      "Node ID: 7428857b-8eef-43e2-82a1-e17c86402aef\n",
      "Embedding length: 1024\n",
      "Embedding preview: [0.002958822064101696, 0.0470910407602787, -0.0005777133628726006, 0.030806731432676315, -0.03956307843327522]\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "for node_id, embedding in list(embedded_node_dict.items())[:5]:\n",
    "    print(f\"Node ID: {node_id}\")\n",
    "    print(f\"Embedding length: {len(embedding)}\")\n",
    "    print(f\"Embedding preview: {embedding[:5]}\")  # First few dims\n",
    "    print(\"---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dad51092",
   "metadata": {},
   "source": [
    "So now I have access to the raw embedding vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "690db919",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Visualize the distribution of embedding values for a sample of nodes\n",
    "sample_embeddings = list(embedded_node_dict.values())[:]  # Sample all embeddings\n",
    "flattened_values = [value for emb in sample_embeddings for value in emb]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8013e5f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23352 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "plt.hist(flattened_values, bins=50)\n",
    "plt.title(\"Distribution of Embedding Values\")\n",
    "plt.xlabel(\"Embedding Value\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.show() \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5ca7aaf7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.015274446457624435, -0.005610795691609383, 0.011983891017735004, 0.020918143913149834, -0.03160902485251427, -0.009153444319963455, -0.017094973474740982, -0.011836066842079163, 0.006088125519454479, 0.037640053778886795]\n"
     ]
    }
   ],
   "source": [
    "print(sample_embeddings[0][:10])  # Print first 10 dimensions of the first embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3434d687",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1160728 characters]",
      "text/plain": [
       "<Figure size 1000x800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Make a heat map of the embedding vectors\n",
    "\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "embedding_matrix = np.array(sample_embeddings)\n",
    "plt.figure(figsize=(10, 8))\n",
    "sns.heatmap(embedding_matrix, cmap=\"plasma\", cbar=True)\n",
    "plt.title(\"Heatmap of Embedding Vectors (4816 total)\")\n",
    "plt.xlabel(\"Embedding Dimension (1024 total, BAAI/bge-large-en-v1.5)\")\n",
    "plt.ylabel(\"Node Index\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "23375ecd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4907, 1024)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Shape of the embedding matrix\n",
    "embedding_matrix.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "54682af1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22952 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Magnitude of each embedding vector\n",
    "\n",
    "magnitudes = np.linalg.norm(embedding_matrix, axis=1)\n",
    "plt.hist(magnitudes, bins=100)\n",
    "plt.title(\"Distribution of Embedding Vector Magnitudes\")\n",
    "plt.xlabel(\"Magnitude\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2f383b7d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Variance across embedding dimensions')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 54264 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Mean and variance across dimensions\n",
    "# shows if some dimensions carry more signal than others\n",
    "\n",
    "mean_per_dim = embedding_matrix.mean(axis=0)\n",
    "var_per_dim = embedding_matrix.var(axis=0)\n",
    "plt.plot(var_per_dim)\n",
    "plt.title(\"Variance across embedding dimensions\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b737ddb9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Axes: ylabel='Count'>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16108 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Cosine similarity distribution\n",
    "# reveals how clustered versus spread out embeddings are\n",
    "\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "sims = cosine_similarity(embedding_matrix[:]) \n",
    "sns.histplot(sims.flatten(), bins=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "092c8797",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 482332 characters]",
      "text/plain": [
       "<Figure size 1200x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot heatmap\n",
    "# NOTE that this is only a subset of the vectors\n",
    "plt.figure(figsize=(12, 10))\n",
    "sns.heatmap(sims[200:500,200:500], cmap=\"plasma\", cbar=True)\n",
    "plt.title(\"Cosine Similarity Matrix of Embeddings, subset of 300\")\n",
    "plt.xlabel(\"Node Index\")\n",
    "plt.ylabel(\"Node Index\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b484d6e7",
   "metadata": {},
   "source": [
    "### Hierarchical clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b45beeff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19088 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy.cluster.hierarchy import linkage, dendrogram\n",
    "linkage_matrix = linkage(embedding_matrix, method='ward')  # sample\n",
    "# Note: use semicolon to suppress output in Jupyter\n",
    "dendrogram(linkage_matrix);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b9cfbfa",
   "metadata": {},
   "source": [
    "### PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f66490b5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'PCA projection of embeddings')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 83260 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "pca = PCA(n_components=2)\n",
    "reduced = pca.fit_transform(embedding_matrix)\n",
    "plt.scatter(reduced[:,0], reduced[:,1], s=5)\n",
    "plt.title(\"PCA projection of embeddings\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "544f9fa0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 't-SNE projection of embeddings')"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 100288 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sklearn.manifold import TSNE\n",
    "tsne = TSNE(n_components=2, perplexity=30, random_state=42)\n",
    "reduced = tsne.fit_transform(embedding_matrix)\n",
    "plt.scatter(reduced[:,0], reduced[:,1], s=5)\n",
    "plt.title(\"t-SNE projection of embeddings\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ab8208",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
