{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "61e8ae81-7b3f-435f-80fe-c536dd2e4352",
   "metadata": {},
   "source": [
    "# Color space conversion\n",
    "\n",
    "Let try the [online demo app](https://cutcutcodec-colorspace.streamlit.app)!\n",
    "\n",
    "* The terminology and the way equations are constructed are detailed in [the documentation](https://cutcutcodec.readthedocs.io/latest/build/api/cutcutcodec.core.colorspace.html#module-cutcutcodec.core.colorspace).\n",
    "* The available **gammas** and **gamuts** are also detailed in [the documentation](https://cutcutcodec.readthedocs.io/latest/build/api/cutcutcodec.core.colorspace.cst.html). They come from International Telecommunication Union reports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e7491441-4e33-439b-8ecd-0f075e9460f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import sympy\n",
    "import torch\n",
    "from IPython.display import Math, display\n",
    "\n",
    "import cutcutcodec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "172249d4-63bd-4061-bf85-be0868549597",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Source Colorspace: Colorspace(\"y'pbpr\", 'ntsc', 'smpte240m')\n",
      "Target Colorspace: Colorspace(\"y'pbpr\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    }
   ],
   "source": [
    "src = cutcutcodec.Colorspace(\"y'pbpr\", \"smpte240m\", \"smpte240m\")\n",
    "dst = cutcutcodec.Colorspace(\"y'pbpr\", \"bt2020\", \"bt2020\")\n",
    "print(f\"Source Colorspace: {src}\")\n",
    "print(f\"Target Colorspace: {dst}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0024c53-a356-4ff8-b810-a29c222341e8",
   "metadata": {},
   "source": [
    "## Symbolic conversions\n",
    "\n",
    "This allows you to retrieve the passage equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "843203a0-c1f4-43e3-a066-569f91e9430d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"y'pbpr\", 'ntsc', 'smpte240m') -> Colorspace(\"r'g'b'\", 'ntsc', 'smpte240m')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{430238494 p_{r}}{273127803} + y', \\  - \\frac{7623978355136 p_{b}}{33806856032607} - \\frac{209727006492056 p_{r}}{439489128423891} + y', \\  \\frac{38383246 p_{b}}{21009831} + y'\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"r'g'b'\", 'ntsc', 'smpte240m') -> Colorspace('rgb', 'ntsc')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\begin{cases} 0.25 r' & \\text{for}\\: r' \\leq 0.09128634211778008882172239358977671184 \\\\\\left(0.8996266762239280431763196027708491089 r' + 0.1003733237760719568236803972291508911\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 0.25 g' & \\text{for}\\: g' \\leq 0.09128634211778008882172239358977671184 \\\\\\left(0.8996266762239280431763196027708491089 g' + 0.1003733237760719568236803972291508911\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 0.25 b' & \\text{for}\\: b' \\leq 0.09128634211778008882172239358977671184 \\\\\\left(0.8996266762239280431763196027708491089 b' + 0.1003733237760719568236803972291508911\\right)^{2.222222222222222222222222222222222222} & \\text{otherwise} \\end{cases}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('rgb', 'ntsc') -> Colorspace('xyz')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{4026032 b}{21009831} + \\frac{99764014 g}{273127803} + \\frac{35828814 r}{91042601}, \\  \\frac{1818208 b}{21009831} + \\frac{191482543 g}{273127803} + \\frac{58008556 r}{273127803}, \\  \\frac{20130160 b}{21009831} + \\frac{30572843 g}{273127803} + \\frac{1706134 r}{91042601}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('xyz') -> Colorspace('rgb', 'bt2020')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{21532150939 x}{12543355000} - \\frac{4461219061 y}{12543355000} - \\frac{3178002061 z}{12543355000}, \\  - \\frac{1976779337 x}{2965129750} + \\frac{4793012913 y}{2965129750} + \\frac{46755163 z}{2965129750}, \\  \\frac{79263327 x}{4492356500} - \\frac{192186423 y}{4492356500} + \\frac{4233267077 z}{4492356500}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace('rgb', 'bt2020') -> Colorspace(\"r'g'b'\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\begin{cases} 4.5 r & \\text{for}\\: r \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 r^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 4.5 g & \\text{for}\\: g \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 g^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}, \\  \\begin{cases} 4.5 b & \\text{for}\\: b \\leq 0.01805396851080780733586959258468773494 \\\\1.099296826809442940347282759215782542 b^{0.45} - 0.099296826809442940347282759215782542 & \\text{otherwise} \\end{cases}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Colorspace(\"r'g'b'\", 'bt2020', 'bt1361e, bt1361') -> Colorspace(\"y'pbpr\", 'bt2020', 'bt1361e, bt1361')\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left( \\frac{826593596 b'}{13942086899} + \\frac{9452833643 g'}{13942086899} + \\frac{3662659660 r'}{13942086899}, \\  \\frac{b'}{2} - \\frac{859348513 g'}{2384635146} - \\frac{166484530 r'}{1192317573}, \\  - \\frac{413296798 b'}{10279427239} - \\frac{9452833643 g'}{20558854478} + \\frac{r'}{2}\\right)$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "steps = [  # details of each transition stage\n",
    "    src,\n",
    "    cutcutcodec.Colorspace(\"r'g'b'\", src.primaries, src.transfer),\n",
    "    cutcutcodec.Colorspace(\"rgb\", src.primaries),\n",
    "    cutcutcodec.Colorspace(\"xyz\"),\n",
    "    cutcutcodec.Colorspace(\"rgb\", dst.primaries),\n",
    "    cutcutcodec.Colorspace(\"r'g'b'\", dst.primaries, dst.transfer),\n",
    "    dst,\n",
    "]\n",
    "for step_n, step_np1 in zip(steps[:-1], steps[1:]):\n",
    "    print(f\"{step_n} -> {step_np1}\")\n",
    "    eq = step_n.to_equation(step_np1)\n",
    "    display(Math(sympy.latex(eq)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0e4adfa-0851-47e4-b18a-f2fd9c035f0e",
   "metadata": {},
   "source": [
    "## Numerical conversion\n",
    "\n",
    "* Offers more flexibility than the FilterVideoColorspace.\n",
    "* Compiled in C, this function optimizes cache and overhead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f66adfe1-fbb7-4236-82aa-d1c6c0c0dc29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\begin{align}\n",
       "\\textbf{def}~f\\left(u, v, y\\right): \\\\\n",
       "\\quad x_0 \\leftarrow 0.25 y \\\\\n",
       "\\quad x_1 \\leftarrow 1.575227747868641553126687728674769884 v \\\\\n",
       "\\quad x_1 \\leftarrow x_{1} + y \\\\\n",
       "\\quad x_1 \\leftarrow x_{1} \\leq 0.09128634211778008882172239358977671184 \\\\\n",
       "\\quad x_2 \\leftarrow 0.8996266762239280431763196027708491089 y \\\\\n",
       "\\quad x_2 \\leftarrow x_{2} + 0.1003733237760719568236803972291508911 \\\\\n",
       "\\quad x_3 \\leftarrow 1.417116903110769752212105343075630963 v \\\\\n",
       "\\quad x_3 \\leftarrow x_{2} + x_{3} \\\\\n",
       "\\quad x_3 \\leftarrow x_{3}^{2.222222222222222222222222222222222222} \\\\\n",
       "\\quad x_4 \\leftarrow 0.393806936967160388281671932168692471 v \\\\\n",
       "\\quad x_4 \\leftarrow x_{0} + x_{4} \\\\\n",
       "\\quad x_4 \\leftarrow \\begin{cases} x_{4} & \\text{for}\\: x_{1} \\\\x_{3} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_5 \\leftarrow - x_{0} \\\\\n",
       "\\quad x_6 \\leftarrow - y \\\\\n",
       "\\quad x_7 \\leftarrow 0.4772063583101252952258757092198383568 v \\\\\n",
       "\\quad x_8 \\leftarrow 0.2255157459120897865523043115836921751 u \\\\\n",
       "\\quad x_6 \\leftarrow x_{6} + x_{7} + x_{8} \\\\\n",
       "\\quad x_6 \\leftarrow x_{6} \\geq -0.09128634211778008882172239358977671184 \\\\\n",
       "\\quad x_7 \\leftarrow 0.8996266762239280431763196027708491089 y \\\\\n",
       "\\quad x_8 \\leftarrow - 0.2028799809310532225771413906512419189 u \\\\\n",
       "\\quad x_9 \\leftarrow - 0.4293075699994628824952572134562781054 v \\\\\n",
       "\\quad x_7 \\leftarrow x_{7} + x_{8} + x_{9} + 0.1003733237760719568236803972291508911 \\\\\n",
       "\\quad x_7 \\leftarrow x_{7}^{2.222222222222222222222222222222222222} \\\\\n",
       "\\quad x_8 \\leftarrow - x_{5} \\\\\n",
       "\\quad x_9 \\leftarrow - 0.1193015895775313238064689273049595892 v \\\\\n",
       "\\quad x_{10} \\leftarrow - 0.05637893647802244663807607789592304378 u \\\\\n",
       "\\quad x_8 \\leftarrow x_{8} + x_{9} + x_{10} \\\\\n",
       "\\quad x_8 \\leftarrow \\begin{cases} x_{8} & \\text{for}\\: x_{6} \\\\x_{7} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_9 \\leftarrow 1.826918360266677061800259126310916066 u \\\\\n",
       "\\quad x_9 \\leftarrow x_{9} + y \\\\\n",
       "\\quad x_9 \\leftarrow x_{9} \\leq 0.09128634211778008882172239358977671184 \\\\\n",
       "\\quad x_{10} \\leftarrow 1.643544492179179412130221165880667149 u \\\\\n",
       "\\quad x_{10} \\leftarrow x_{2} + x_{10} \\\\\n",
       "\\quad x_{10} \\leftarrow x_{10}^{2.222222222222222222222222222222222222} \\\\\n",
       "\\quad x_2 \\leftarrow 0.4567295900666692654500647815777290165 u \\\\\n",
       "\\quad x_2 \\leftarrow x_{0} + x_{2} \\\\\n",
       "\\quad x_2 \\leftarrow \\begin{cases} x_{2} & \\text{for}\\: x_{9} \\\\x_{10} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_{11} \\leftarrow x_{1} \\wedge x_{6} \\wedge x_{9} \\\\\n",
       "\\quad x_{12} \\leftarrow 0.005616235079712466335963477384737015402 u \\\\\n",
       "\\quad x_{13} \\leftarrow 0.1927478152272564917376197799690692823 v \\\\\n",
       "\\quad x_{14} \\leftarrow x_{6} \\wedge x_{9} \\\\\n",
       "\\quad x_{15} \\leftarrow 0.5952701227710405783213474436823374812 x_{3} \\\\\n",
       "\\quad x_{16} \\leftarrow - 0.04167368848927251101449620675513716215 v \\\\\n",
       "\\quad x_{16} \\leftarrow x_{15} + x_{16} \\\\\n",
       "\\quad x_{17} \\leftarrow x_{1} \\wedge x_{9} \\\\\n",
       "\\quad x_{18} \\leftarrow 0.02531017415011096215842454741477759892 u \\\\\n",
       "\\quad x_{19} \\leftarrow 0.3493137739140495835481416592001061499 x_{7} \\\\\n",
       "\\quad x_{20} \\leftarrow 0.2344215037165290027521159867242064444 v \\\\\n",
       "\\quad x_{20} \\leftarrow x_{19} + x_{20} \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} + x_{19} \\\\\n",
       "\\quad x_{19} \\leftarrow x_{1} \\wedge x_{6} \\\\\n",
       "\\quad x_{21} \\leftarrow 0.05541610331490983813051089711755636883 x_{10} \\\\\n",
       "\\quad x_{22} \\leftarrow - 0.01969393907039849582246107003004058352 u \\\\\n",
       "\\quad x_{22} \\leftarrow x_{21} + x_{22} \\\\\n",
       "\\quad x_{23} \\leftarrow 2.678715552469682602446063496570518666 x_{4} \\\\\n",
       "\\quad x_{24} \\leftarrow 1.571911982613223125966637466400477675 x_{8} \\\\\n",
       "\\quad x_{25} \\leftarrow 0.2493724649170942715872990370290036597 x_{2} \\\\\n",
       "\\quad x_{23} \\leftarrow x_{23} + x_{24} + x_{25} \\\\\n",
       "\\quad x_{24} \\leftarrow x_{0} + x_{12} + x_{13} \\\\\n",
       "\\quad x_{24} \\leftarrow x_{24} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{25} \\leftarrow 0.1011824693072398554196631390794156297 y \\\\\n",
       "\\quad x_{12} \\leftarrow x_{12} + x_{16} + x_{25} \\\\\n",
       "\\quad x_{12} \\leftarrow x_{12} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{25} \\leftarrow 0.1626715565214876041129645851999734625 y \\\\\n",
       "\\quad x_{25} \\leftarrow x_{18} + x_{20} + x_{25} \\\\\n",
       "\\quad x_{25} \\leftarrow x_{25} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{26} \\leftarrow 0.01385402582872745953262772427938909221 y \\\\\n",
       "\\quad x_{18} \\leftarrow x_{15} + x_{18} + x_{26} \\\\\n",
       "\\quad x_{18} \\leftarrow x_{18} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{26} \\leftarrow 0.2361459741712725404673722757206109078 y \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} + x_{22} + x_{26} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{26} \\leftarrow 0.08732844347851239588703541480002653748 y \\\\\n",
       "\\quad x_{16} \\leftarrow x_{16} + x_{22} + x_{26} \\\\\n",
       "\\quad x_{16} \\leftarrow x_{16} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{22} \\leftarrow 0.1488175306927601445803368609205843703 y \\\\\n",
       "\\quad x_{20} \\leftarrow x_{20} + x_{21} + x_{22} \\\\\n",
       "\\quad x_{20} \\leftarrow x_{20} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} + x_{21} \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{15} \\leftarrow \\operatorname{ITE}\\left(x_{1}, x_{20}, x_{15}\\right) \\\\\n",
       "\\quad x_{15} \\leftarrow \\operatorname{ITE}\\left(x_{6}, x_{16}, x_{15}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow \\operatorname{ITE}\\left(x_{19}, x_{13}, x_{15}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow \\operatorname{ITE}\\left(x_{9}, x_{18}, x_{13}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow \\operatorname{ITE}\\left(x_{17}, x_{25}, x_{13}\\right) \\\\\n",
       "\\quad x_{12} \\leftarrow \\operatorname{ITE}\\left(x_{14}, x_{12}, x_{13}\\right) \\\\\n",
       "\\quad x_{12} \\leftarrow \\operatorname{ITE}\\left(x_{11}, x_{24}, x_{12}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow 0.3493137739140495835481416592001061499 x_{8} \\\\\n",
       "\\quad x_{15} \\leftarrow 0.05541610331490983813051089711755636883 x_{2} \\\\\n",
       "\\quad x_{16} \\leftarrow 0.5952701227710405783213474436823374812 x_{4} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} + x_{15} + x_{16} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13}^{0.45} \\\\\n",
       "\\quad x_{13} \\leftarrow 1.099296826809442940347282759215782542 x_{13} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} - 0.099296826809442940347282759215782542 \\\\\n",
       "\\quad x_{12} \\leftarrow \\begin{cases} x_{23} & \\text{for}\\: x_{12} \\\\x_{13} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_{13} \\leftarrow 0.07436298549573965448441799964934607902 v \\\\\n",
       "\\quad x_{15} \\leftarrow 0.03781848215812902175715342131049288056 u \\\\\n",
       "\\quad x_{16} \\leftarrow 0.0124438037631045686968790177159063483 u \\\\\n",
       "\\quad x_{18} \\leftarrow 0.8915082309299459767548066740119863528 x_{7} \\\\\n",
       "\\quad x_{20} \\leftarrow 0.03199536357565577701992240690343537089 v \\\\\n",
       "\\quad x_{20} \\leftarrow x_{18} + x_{20} \\\\\n",
       "\\quad x_{21} \\leftarrow 0.1063583490713954315043404065527814499 v \\\\\n",
       "\\quad x_{22} \\leftarrow 0.08124631785834660030870500609373494137 x_{3} \\\\\n",
       "\\quad x_{18} \\leftarrow x_{18} + x_{22} \\\\\n",
       "\\quad x_{23} \\leftarrow 0.05026228592123359045403243902639922886 u \\\\\n",
       "\\quad x_{24} \\leftarrow 0.02724545121170742293648831989427870584 x_{10} \\\\\n",
       "\\quad x_{25} \\leftarrow 4.011787039184756895396630033053938588 x_{8} \\\\\n",
       "\\quad x_{26} \\leftarrow 0.3656084303625597013891725274218072362 x_{4} \\\\\n",
       "\\quad x_{27} \\leftarrow 0.1226045304526834032141974395242541763 x_{2} \\\\\n",
       "\\quad x_{27} \\leftarrow x_{25} + x_{26} + x_{27} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} + x_{13} + x_{15} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} \\geq -0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{25} \\leftarrow 0.0271229422675135058112983314970034118 y \\\\\n",
       "\\quad x_{25} \\leftarrow x_{16} + x_{20} + x_{25} \\\\\n",
       "\\quad x_{25} \\leftarrow x_{25} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{26} \\leftarrow - x_{22} \\\\\n",
       "\\quad x_{28} \\leftarrow - 0.2296884205354133499228237484765662647 y \\\\\n",
       "\\quad x_{28} \\leftarrow x_{15} + x_{21} + x_{26} + x_{28} \\\\\n",
       "\\quad x_{28} \\leftarrow x_{28} \\geq -0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{15} \\leftarrow 0.006811362802926855734122079973569676459 y \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} + x_{16} + x_{18} \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{16} \\leftarrow - x_{24} \\\\\n",
       "\\quad x_{26} \\leftarrow - 0.2431886371970731442658779200264303235 y \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} + x_{16} + x_{23} + x_{26} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} \\geq -0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{16} \\leftarrow 0.02031157946458665007717625152343373534 y \\\\\n",
       "\\quad x_{16} \\leftarrow x_{16} + x_{20} + x_{24} \\\\\n",
       "\\quad x_{16} \\leftarrow x_{16} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{20} \\leftarrow - x_{21} \\\\\n",
       "\\quad x_{21} \\leftarrow - x_{23} \\\\\n",
       "\\quad x_{23} \\leftarrow 0.2228770577324864941887016685029965882 y \\\\\n",
       "\\quad x_{20} \\leftarrow x_{20} + x_{21} + x_{22} + x_{23} + x_{24} \\\\\n",
       "\\quad x_{20} \\leftarrow x_{20} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{18} \\leftarrow x_{18} + x_{24} \\\\\n",
       "\\quad x_{18} \\leftarrow x_{18} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{18} \\leftarrow \\operatorname{ITE}\\left(x_{6}, x_{20}, x_{18}\\right) \\\\\n",
       "\\quad x_{16} \\leftarrow \\operatorname{ITE}\\left(x_{1}, x_{16}, x_{18}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow \\operatorname{ITE}\\left(x_{19}, x_{13}, x_{16}\\right) \\\\\n",
       "\\quad x_{13} \\leftarrow \\operatorname{ITE}\\left(x_{9}, x_{15}, x_{13}\\right) \\\\\n",
       "\\quad x_{28} \\leftarrow \\operatorname{ITE}\\left(x_{14}, x_{28}, x_{13}\\right) \\\\\n",
       "\\quad x_{28} \\leftarrow \\operatorname{ITE}\\left(x_{17}, x_{25}, x_{28}\\right) \\\\\n",
       "\\quad x_{28} \\leftarrow \\operatorname{ITE}\\left(x_{11}, x_{5}, x_{28}\\right) \\\\\n",
       "\\quad x_5 \\leftarrow 0.8915082309299459767548066740119863528 x_{8} \\\\\n",
       "\\quad x_{13} \\leftarrow 0.08124631785834660030870500609373494137 x_{4} \\\\\n",
       "\\quad x_{15} \\leftarrow 0.02724545121170742293648831989427870584 x_{2} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} + x_{13} + x_{15} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5}^{0.45} \\\\\n",
       "\\quad x_5 \\leftarrow 1.099296826809442940347282759215782542 x_{5} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} - 0.099296826809442940347282759215782542 \\\\\n",
       "\\quad x_{27} \\leftarrow \\begin{cases} x_{27} & \\text{for}\\: x_{28} \\\\x_{5} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_{28} \\leftarrow 0.4076023144074993873444640152111481931 u \\\\\n",
       "\\quad x_5 \\leftarrow 0.003664093591595216860355383337237092768 v \\\\\n",
       "\\quad x_{13} \\leftarrow 0.004619273747905038815867694158849998107 u \\\\\n",
       "\\quad x_{10} \\leftarrow 0.9025506494887533722687808235712819565 x_{10} \\\\\n",
       "\\quad x_{15} \\leftarrow 0.006110597724791278555250393601747546687 v \\\\\n",
       "\\quad x_7 \\leftarrow 0.0819326159106551655887679794018842299 x_{7} \\\\\n",
       "\\quad x_{16} \\leftarrow 0.4122215881554044261603317093699981912 u \\\\\n",
       "\\quad x_{16} \\leftarrow x_{7} + x_{16} \\\\\n",
       "\\quad x_7 \\leftarrow x_{7} + x_{10} \\\\\n",
       "\\quad x_{18} \\leftarrow 0.01551673460059146214245119702683381359 x_{3} \\\\\n",
       "\\quad x_{20} \\leftarrow - 0.009774691316386495415605776938984639455 v \\\\\n",
       "\\quad x_{20} \\leftarrow x_{18} + x_{20} \\\\\n",
       "\\quad x_{21} \\leftarrow 4.061477922699390175209513706070768804 x_{2} \\\\\n",
       "\\quad x_{22} \\leftarrow 0.06982530570266157964103038662075216118 x_{4} \\\\\n",
       "\\quad x_{23} \\leftarrow 0.3686967715979482451494559073084790346 x_{8} \\\\\n",
       "\\quad x_{21} \\leftarrow x_{21} + x_{22} + x_{23} \\\\\n",
       "\\quad x_{22} \\leftarrow - x_{5} \\\\\n",
       "\\quad x_0 \\leftarrow x_{0} + x_{22} + x_{28} \\\\\n",
       "\\quad x_0 \\leftarrow x_{0} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{22} \\leftarrow - x_{10} \\\\\n",
       "\\quad x_{23} \\leftarrow - 0.02436233762781165693280479410717951087 y \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} + x_{13} + x_{22} + x_{23} \\\\\n",
       "\\quad x_5 \\leftarrow x_{5} \\geq -0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{22} \\leftarrow 0.2295168460223362086028080051495289425 y \\\\\n",
       "\\quad x_{22} \\leftarrow x_{15} + x_{16} + x_{22} \\\\\n",
       "\\quad x_{22} \\leftarrow x_{22} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{23} \\leftarrow 0.003879183650147865535612799256708453399 y \\\\\n",
       "\\quad x_{15} \\leftarrow x_{7} + x_{15} + x_{23} \\\\\n",
       "\\quad x_{15} \\leftarrow x_{15} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{23} \\leftarrow 0.2461208163498521344643872007432915466 y \\\\\n",
       "\\quad x_{28} \\leftarrow x_{20} + x_{23} + x_{28} \\\\\n",
       "\\quad x_{28} \\leftarrow x_{28} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{13} \\leftarrow - x_{13} \\\\\n",
       "\\quad x_{23} \\leftarrow 0.02048315397766379139719199485047105748 y \\\\\n",
       "\\quad x_{10} \\leftarrow x_{10} + x_{13} + x_{20} + x_{23} \\\\\n",
       "\\quad x_{10} \\leftarrow x_{10} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_{13} \\leftarrow 0.2256376623721883430671952058928204891 y \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} + x_{16} + x_{18} \\\\\n",
       "\\quad x_{13} \\leftarrow x_{13} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_7 \\leftarrow x_{7} + x_{18} \\\\\n",
       "\\quad x_7 \\leftarrow x_{7} \\leq 0.01805396851080780733586959258468773494 \\\\\n",
       "\\quad x_7 \\leftarrow \\operatorname{ITE}\\left(x_{9}, x_{13}, x_{7}\\right) \\\\\n",
       "\\quad x_6 \\leftarrow \\operatorname{ITE}\\left(x_{6}, x_{10}, x_{7}\\right) \\\\\n",
       "\\quad x_{28} \\leftarrow \\operatorname{ITE}\\left(x_{14}, x_{28}, x_{6}\\right) \\\\\n",
       "\\quad x_1 \\leftarrow \\operatorname{ITE}\\left(x_{1}, x_{15}, x_{28}\\right) \\\\\n",
       "\\quad x_1 \\leftarrow \\operatorname{ITE}\\left(x_{17}, x_{22}, x_{1}\\right) \\\\\n",
       "\\quad x_1 \\leftarrow \\operatorname{ITE}\\left(x_{19}, x_{5}, x_{1}\\right) \\\\\n",
       "\\quad x_0 \\leftarrow \\operatorname{ITE}\\left(x_{11}, x_{0}, x_{1}\\right) \\\\\n",
       "\\quad x_1 \\leftarrow 0.0819326159106551655887679794018842299 x_{8} \\\\\n",
       "\\quad x_{28} \\leftarrow 0.9025506494887533722687808235712819565 x_{2} \\\\\n",
       "\\quad x_5 \\leftarrow 0.01551673460059146214245119702683381359 x_{4} \\\\\n",
       "\\quad x_1 \\leftarrow x_{1} + x_{5} + x_{28} \\\\\n",
       "\\quad x_1 \\leftarrow x_{1}^{0.45} \\\\\n",
       "\\quad x_1 \\leftarrow 1.099296826809442940347282759215782542 x_{1} \\\\\n",
       "\\quad x_1 \\leftarrow x_{1} - 0.099296826809442940347282759215782542 \\\\\n",
       "\\quad x_0 \\leftarrow \\begin{cases} x_{21} & \\text{for}\\: x_{0} \\\\x_{1} & \\text{otherwise} \\end{cases} \\\\\n",
       "\\quad x_1 \\leftarrow 0.2627052669039600712074144417509988725 x_{12} \\\\\n",
       "\\quad x_{28} \\leftarrow 0.6780070811119393525736767106632850431 x_{27} \\\\\n",
       "\\quad x_5 \\leftarrow 0.05928765198410057621890884758571608441 x_{0} \\\\\n",
       "\\quad x_1 \\leftarrow x_{1} + x_{5} + x_{28} \\\\\n",
       "\\quad x_{28} \\leftarrow 0.5 x_{0} \\\\\n",
       "\\quad x_5 \\leftarrow - 0.1396310293247686621154916082076398282 x_{12} \\\\\n",
       "\\quad x_6 \\leftarrow - 0.3603689706752313378845083917923601718 x_{27} \\\\\n",
       "\\quad x_{28} \\leftarrow x_{5} + x_{6} + x_{28} \\\\\n",
       "\\quad x_5 \\leftarrow 0.5 x_{12} \\\\\n",
       "\\quad x_{27} \\leftarrow - 0.4597937911917934633089336856193296705 x_{27} \\\\\n",
       "\\quad x_0 \\leftarrow - 0.04020620880820653669106631438067032949 x_{0} \\\\\n",
       "\\quad x_0 \\leftarrow x_{0} + x_{5} + x_{27} \\\\\n",
       "\\quad x\\_ \\leftarrow \\left( x_{1}, \\  x_{28}, \\  x_{0}\\right) \\\\\n",
       "\\quad \\textbf{return}~x\\_\n",
       "\\end{align}$"
      ],
      "text/plain": [
       "<IPython.core.display.Math object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "func = src.to_function(dst)  # optimize the full chain\n",
    "display(Math(str(func)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc4e62b5-99b5-4e36-8277-dcaaabb2bf89",
   "metadata": {},
   "outputs": [],
   "source": [
    "src_yuv = torch.rand(2160, 3840, 3)\n",
    "src_yuv[..., 1:] -= 0.5\n",
    "src_yuv *= 0.2  # valid range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b15c50c4-821a-4eae-b5ab-c92548c44a85",
   "metadata": {},
   "outputs": [],
   "source": [
    "dst_y, dst_u, dst_v = func(y=src_yuv[..., 0], u=src_yuv[..., 1], v=src_yuv[..., 2])\n",
    "dst_yuv = torch.cat([dst_y[..., None], dst_u[..., None], dst_v[..., None]], dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5caba756-82c1-4173-8326-06cecf08d730",
   "metadata": {},
   "source": [
    "## Alternatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "93ad2ab6-d4ab-4325-bed5-1de1e2d117e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit\n",
    "\n",
    "import colour  # pip install colour-science"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6e628f5b-a15a-40f3-ad67-b1d275732bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def colourscience_convert(src_yuv):\n",
    "    src_colourspace = colour.models.RGB_COLOURSPACE_SMPTE_240M\n",
    "    src_k = colour.WEIGHTS_YCBCR[\"SMPTE-240M\"]\n",
    "    dst_colourspace = colour.models.RGB_COLOURSPACE_BT2020\n",
    "    dst_k = colour.WEIGHTS_YCBCR[\"ITU-R BT.2020\"]\n",
    "    src_rgb = colour.YCbCr_to_RGB(src_yuv, src_k, in_range=(0.0, 1.0, -0.5, 0.5), out_range=(0.0, 1.0))\n",
    "    xyz = colour.RGB_to_XYZ(src_rgb, src_colourspace, apply_cctf_decoding=True)\n",
    "    dst_rgb = colour.XYZ_to_RGB(xyz, dst_colourspace, apply_cctf_encoding=True)\n",
    "    dst_yuv = colour.RGB_to_YCbCr(dst_rgb, dst_k, in_range=(0.0, 1.0), out_range=(0.0, 1.0, -0.5, 0.5))\n",
    "    return dst_yuv\n",
    "\n",
    "dst_yuv_ = colourscience_convert(src_yuv.numpy(force=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "423fab4c-4f8e-4e0d-99f7-e72254bee40d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_7812/2375015244.py:1: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)\n",
      "  print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).mean())\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3.3807e-05, dtype=torch.float64)\n",
      "tensor(0.0001, dtype=torch.float64)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_7812/2375015244.py:2: DeprecationWarning: __array_wrap__ must accept context and return_scalar arguments (positionally) in the future. (Deprecated NumPy 2.0)\n",
      "  print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).max())\n"
     ]
    }
   ],
   "source": [
    "print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).mean())\n",
    "print(abs(dst_yuv[..., 2] - dst_yuv_[..., 2]).max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b2f3627a-65b2-4cea-864b-8d724ce88254",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cutcutcodec convesion take 1894.23 ms\n",
      "colourscience convesion take 17219.00 ms\n",
      "cutcutcodec is 9.09022393862781 times faster than colourscience\n"
     ]
    }
   ],
   "source": [
    "number = 5\n",
    "time_cutcutcodec = timeit.repeat(lambda: func(y=src_yuv[..., 0], u=src_yuv[..., 1], v=src_yuv[..., 2]), number=number, repeat=7)\n",
    "print(f\"cutcutcodec convesion take {1000*np.median(time_cutcutcodec)/number:.2f} ms\")\n",
    "time_colourscience = timeit.repeat(lambda: colourscience_convert(src_yuv.numpy(force=True)), number=number, repeat=5)\n",
    "print(f\"colourscience convesion take {1000*np.median(time_colourscience)/number:.2f} ms\")\n",
    "print(f\"cutcutcodec is {np.median(time_colourscience)/np.median(time_cutcutcodec)} times faster than colourscience\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}