{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2058498/2740422639.py:11: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  data = torch.load(file_path)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 25912 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import os\n",
    "os.sys.path.append(\"../src\")\n",
    "from fair_PCA import FairPCA\n",
    "from main import *\n",
    "\n",
    "import torch\n",
    "import os\n",
    "exp_path = \"../output/exps/protect=[gender]\"\n",
    "file_path = f\"{exp_path}/extracted_features.pt\"\n",
    "data = torch.load(file_path)\n",
    "protect = list(data.keys())[0]\n",
    "hdim = 800\n",
    "fpca = FairPCA(target_dim = hdim, standardize = False, tradeoff_param = 0)\n",
    "\n",
    "if \"UserStableDiffusion3Pipeline\" not in locals():\n",
    "    pipe = UserStableDiffusion3Pipeline.from_pretrained(\"stabilityai/stable-diffusion-3-medium-diffusers\", torch_dtype=torch.float16)\n",
    "pipe.exp_dir = exp_path\n",
    "pipe.calc_projection_matrix_sg(data[protect], fpca)\n",
    "\n",
    "projections = {}\n",
    "for group in data[protect]:\n",
    "    projections[group] = fpca.nzTXT.matmul(data[protect][group].T)\n",
    "\n",
    "# plot projections with bin plot, each group is a different color. Write code woth matplotlib\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Example data in projections dict\n",
    "\n",
    "# Plot histogram for each group with different colors\n",
    "colors = ['blue', 'orange', 'green', 'red', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan']\n",
    "for i, group in enumerate(projections):\n",
    "    plt.hist(projections[group], bins=20, alpha=0.5, color=colors[i], label=group)\n",
    "\n",
    "# Labeling the plot\n",
    "plt.xlabel(\"Values\")\n",
    "plt.ylabel(\"Frequency\")\n",
    "plt.legend()\n",
    "plt.title(\"Bin Plot of Projections for Each Group\")\n",
    "\n",
    "# Display the plot\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting matplotlib\n",
      "  Using cached matplotlib-3.7.5-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl.metadata (5.7 kB)\n",
      "Collecting contourpy>=1.0.1 (from matplotlib)\n",
      "  Using cached contourpy-1.1.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (5.9 kB)\n",
      "Collecting cycler>=0.10 (from matplotlib)\n",
      "  Using cached cycler-0.12.1-py3-none-any.whl.metadata (3.8 kB)\n",
      "Collecting fonttools>=4.22.0 (from matplotlib)\n",
      "  Downloading fonttools-4.54.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (163 kB)\n",
      "Collecting kiwisolver>=1.0.1 (from matplotlib)\n",
      "  Downloading kiwisolver-1.4.7-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl.metadata (6.3 kB)\n",
      "Requirement already satisfied: numpy<2,>=1.20 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from matplotlib) (1.24.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from matplotlib) (24.1)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from matplotlib) (9.5.0)\n",
      "Collecting pyparsing>=2.3.1 (from matplotlib)\n",
      "  Downloading pyparsing-3.1.4-py3-none-any.whl.metadata (5.1 kB)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from matplotlib) (2.9.0)\n",
      "Collecting importlib-resources>=3.2.0 (from matplotlib)\n",
      "  Downloading importlib_resources-6.4.5-py3-none-any.whl.metadata (4.0 kB)\n",
      "Requirement already satisfied: zipp>=3.1.0 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from importlib-resources>=3.2.0->matplotlib) (3.20.2)\n",
      "Requirement already satisfied: six>=1.5 in /data/zihao.fu/programs/miniconda/envs/fairimg/lib/python3.8/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n",
      "Using cached matplotlib-3.7.5-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (9.2 MB)\n",
      "Using cached contourpy-1.1.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (301 kB)\n",
      "Using cached cycler-0.12.1-py3-none-any.whl (8.3 kB)\n",
      "Downloading fonttools-4.54.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.7/4.7 MB\u001b[0m \u001b[31m37.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading importlib_resources-6.4.5-py3-none-any.whl (36 kB)\n",
      "Downloading kiwisolver-1.4.7-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.2 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.2/1.2 MB\u001b[0m \u001b[31m35.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading pyparsing-3.1.4-py3-none-any.whl (104 kB)\n",
      "Installing collected packages: pyparsing, kiwisolver, importlib-resources, fonttools, cycler, contourpy, matplotlib\n",
      "Successfully installed contourpy-1.1.1 cycler-0.12.1 fonttools-4.54.1 importlib-resources-6.4.5 kiwisolver-1.4.7 matplotlib-3.7.5 pyparsing-3.1.4\n"
     ]
    }
   ],
   "source": [
    "! pip install matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([635, 2048])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[protect][group].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fairimg",
   "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.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
