Skip to content

Instantly share code, notes, and snippets.

@jessegrabowski

jessegrabowski/Jacobian Stuff.ipynb

Created October 19, 2023 17:39
Show Gist options
  • Save jessegrabowski/eea15d82d5b9436dcc76b6dc93d27834 to your computer and use it in GitHub Desktop.
Save jessegrabowski/eea15d82d5b9436dcc76b6dc93d27834 to your computer and use it in GitHub Desktop.
Download ZIP
Jacobians of high-dimension (but independent) systems
Display the source blob
Display the rendered blob
Raw
Jacobian Stuff.ipynb
{
"cells": [
{
"cell_type": "markdown",
"id": "92208cd0",
"metadata": {},
"source": [
"# Symbolic Walk-through"
]
},
{
"cell_type": "code",
"execution_count": 195,
"id": "56ae3339",
"metadata": {},
"outputs": [],
"source": [
"import sympy as sp"
]
},
{
"cell_type": "code",
"execution_count": 199,
"id": "a426bcc0",
"metadata": {},
"outputs": [],
"source": [
"i, j = sp.Idx('i'), sp.Idx('j')\n",
"y, x = sp.IndexedBase('y'), sp.IndexedBase('x')\n",
"y_ij, x_ij = y[i, j], x[i, j]\n"
]
},
{
"cell_type": "code",
"execution_count": 225,
"id": "19d7e479",
"metadata": {},
"outputs": [],
"source": [
"inputs = sp.Matrix([[x[p, q] for p in range(10)] for q in range(2)]).T\n",
"eqs = sp.Matrix([[sp.Function('f')(x[p, q]) for p in range(10)] for q in range(2)]).T"
]
},
{
"cell_type": "markdown",
"id": "ec4179f5",
"metadata": {},
"source": [
"Inputs are a 2d vector of variables"
]
},
{
"cell_type": "code",
"execution_count": 226,
"id": "b7d9213a",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}{x}_{0,0} & {x}_{0,1}\\\\{x}_{1,0} & {x}_{1,1}\\\\{x}_{2,0} & {x}_{2,1}\\\\{x}_{3,0} & {x}_{3,1}\\\\{x}_{4,0} & {x}_{4,1}\\\\{x}_{5,0} & {x}_{5,1}\\\\{x}_{6,0} & {x}_{6,1}\\\\{x}_{7,0} & {x}_{7,1}\\\\{x}_{8,0} & {x}_{8,1}\\\\{x}_{9,0} & {x}_{9,1}\\end{matrix}\\right]$"
],
"text/plain": [
"Matrix([\n",
"[x[0, 0], x[0, 1]],\n",
"[x[1, 0], x[1, 1]],\n",
"[x[2, 0], x[2, 1]],\n",
"[x[3, 0], x[3, 1]],\n",
"[x[4, 0], x[4, 1]],\n",
"[x[5, 0], x[5, 1]],\n",
"[x[6, 0], x[6, 1]],\n",
"[x[7, 0], x[7, 1]],\n",
"[x[8, 0], x[8, 1]],\n",
"[x[9, 0], x[9, 1]]])"
]
},
"execution_count": 226,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"inputs"
]
},
{
"cell_type": "markdown",
"id": "4e8add71",
"metadata": {},
"source": [
"Each row is a system of 2 equations, each one is a function of **only** the variable in that position, so there's no cross-terms anywhere. This is because all the operations you're doing (cdf, pdf, exp) are elemwise."
]
},
{
"cell_type": "code",
"execution_count": 227,
"id": "2fcd92c4",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}f{\\left({x}_{0,0} \\right)} & f{\\left({x}_{0,1} \\right)}\\\\f{\\left({x}_{1,0} \\right)} & f{\\left({x}_{1,1} \\right)}\\\\f{\\left({x}_{2,0} \\right)} & f{\\left({x}_{2,1} \\right)}\\\\f{\\left({x}_{3,0} \\right)} & f{\\left({x}_{3,1} \\right)}\\\\f{\\left({x}_{4,0} \\right)} & f{\\left({x}_{4,1} \\right)}\\\\f{\\left({x}_{5,0} \\right)} & f{\\left({x}_{5,1} \\right)}\\\\f{\\left({x}_{6,0} \\right)} & f{\\left({x}_{6,1} \\right)}\\\\f{\\left({x}_{7,0} \\right)} & f{\\left({x}_{7,1} \\right)}\\\\f{\\left({x}_{8,0} \\right)} & f{\\left({x}_{8,1} \\right)}\\\\f{\\left({x}_{9,0} \\right)} & f{\\left({x}_{9,1} \\right)}\\end{matrix}\\right]$"
],
"text/plain": [
"Matrix([\n",
"[f(x[0, 0]), f(x[0, 1])],\n",
"[f(x[1, 0]), f(x[1, 1])],\n",
"[f(x[2, 0]), f(x[2, 1])],\n",
"[f(x[3, 0]), f(x[3, 1])],\n",
"[f(x[4, 0]), f(x[4, 1])],\n",
"[f(x[5, 0]), f(x[5, 1])],\n",
"[f(x[6, 0]), f(x[6, 1])],\n",
"[f(x[7, 0]), f(x[7, 1])],\n",
"[f(x[8, 0]), f(x[8, 1])],\n",
"[f(x[9, 0]), f(x[9, 1])]])"
]
},
"execution_count": 227,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eqs"
]
},
{
"cell_type": "markdown",
"id": "815dd87f",
"metadata": {},
"source": [
"**Base case**: If you had only two functions, each of which is the result of elemwise on one of two variables, you'd get a diagonal 2x2 jacobian like this:"
]
},
{
"cell_type": "code",
"execution_count": 228,
"id": "3f275a43",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{0,0}} f{\\left({x}_{0,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{0,1}} f{\\left({x}_{0,1} \\right)}\\end{matrix}\\right]$"
],
"text/plain": [
"Matrix([\n",
"[Derivative(f(x[0, 0]), x[0, 0]), 0],\n",
"[ 0, Derivative(f(x[0, 1]), x[0, 1])]])"
]
},
"execution_count": 228,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eqs[0, :].jacobian(inputs[0, :])"
]
},
{
"cell_type": "markdown",
"id": "37c6b739",
"metadata": {},
"source": [
"So what you want is a (10, 2, 2) tensor of diagonal jacobians, like this:"
]
},
{
"cell_type": "code",
"execution_count": 245,
"id": "a3896d44",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{0,0}} f{\\left({x}_{0,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{0,1}} f{\\left({x}_{0,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{1,0}} f{\\left({x}_{1,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{1,1}} f{\\left({x}_{1,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{2,0}} f{\\left({x}_{2,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{2,1}} f{\\left({x}_{2,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{3,0}} f{\\left({x}_{3,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{3,1}} f{\\left({x}_{3,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{4,0}} f{\\left({x}_{4,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{4,1}} f{\\left({x}_{4,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{5,0}} f{\\left({x}_{5,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{5,1}} f{\\left({x}_{5,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{6,0}} f{\\left({x}_{6,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{6,1}} f{\\left({x}_{6,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{7,0}} f{\\left({x}_{7,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{7,1}} f{\\left({x}_{7,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{8,0}} f{\\left({x}_{8,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{8,1}} f{\\left({x}_{8,1} \\right)}\\end{matrix}\\right]\\\\\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{9,0}} f{\\left({x}_{9,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{9,1}} f{\\left({x}_{9,1} \\right)}\\end{matrix}\\right]\\end{matrix}\\right]$"
],
"text/plain": [
"[[[[Derivative(f(x[0, 0]), x[0, 0]), 0], [0, Derivative(f(x[0, 1]), x[0, 1])]]], [[[Derivative(f(x[1, 0]), x[1, 0]), 0], [0, Derivative(f(x[1, 1]), x[1, 1])]]], [[[Derivative(f(x[2, 0]), x[2, 0]), 0], [0, Derivative(f(x[2, 1]), x[2, 1])]]], [[[Derivative(f(x[3, 0]), x[3, 0]), 0], [0, Derivative(f(x[3, 1]), x[3, 1])]]], [[[Derivative(f(x[4, 0]), x[4, 0]), 0], [0, Derivative(f(x[4, 1]), x[4, 1])]]], [[[Derivative(f(x[5, 0]), x[5, 0]), 0], [0, Derivative(f(x[5, 1]), x[5, 1])]]], [[[Derivative(f(x[6, 0]), x[6, 0]), 0], [0, Derivative(f(x[6, 1]), x[6, 1])]]], [[[Derivative(f(x[7, 0]), x[7, 0]), 0], [0, Derivative(f(x[7, 1]), x[7, 1])]]], [[[Derivative(f(x[8, 0]), x[8, 0]), 0], [0, Derivative(f(x[8, 1]), x[8, 1])]]], [[[Derivative(f(x[9, 0]), x[9, 0]), 0], [0, Derivative(f(x[9, 1]), x[9, 1])]]]]"
]
},
"execution_count": 245,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sp.tensor.tensorcontraction([[sp.Array(eqs[i, :].jacobian(inputs[i, :]))] for i in range(10)])"
]
},
{
"cell_type": "markdown",
"id": "423ad67c",
"metadata": {},
"source": [
"But getting this is non-trivial. Instead, you can ravel and reshape. Ravel the equations:"
]
},
{
"cell_type": "code",
"execution_count": 248,
"id": "0233b363",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}f{\\left({x}_{0,0} \\right)}\\\\f{\\left({x}_{0,1} \\right)}\\\\f{\\left({x}_{1,0} \\right)}\\\\f{\\left({x}_{1,1} \\right)}\\\\f{\\left({x}_{2,0} \\right)}\\\\f{\\left({x}_{2,1} \\right)}\\\\f{\\left({x}_{3,0} \\right)}\\\\f{\\left({x}_{3,1} \\right)}\\\\f{\\left({x}_{4,0} \\right)}\\\\f{\\left({x}_{4,1} \\right)}\\\\f{\\left({x}_{5,0} \\right)}\\\\f{\\left({x}_{5,1} \\right)}\\\\f{\\left({x}_{6,0} \\right)}\\\\f{\\left({x}_{6,1} \\right)}\\\\f{\\left({x}_{7,0} \\right)}\\\\f{\\left({x}_{7,1} \\right)}\\\\f{\\left({x}_{8,0} \\right)}\\\\f{\\left({x}_{8,1} \\right)}\\\\f{\\left({x}_{9,0} \\right)}\\\\f{\\left({x}_{9,1} \\right)}\\end{matrix}\\right]$"
],
"text/plain": [
"Matrix([\n",
"[f(x[0, 0])],\n",
"[f(x[0, 1])],\n",
"[f(x[1, 0])],\n",
"[f(x[1, 1])],\n",
"[f(x[2, 0])],\n",
"[f(x[2, 1])],\n",
"[f(x[3, 0])],\n",
"[f(x[3, 1])],\n",
"[f(x[4, 0])],\n",
"[f(x[4, 1])],\n",
"[f(x[5, 0])],\n",
"[f(x[5, 1])],\n",
"[f(x[6, 0])],\n",
"[f(x[6, 1])],\n",
"[f(x[7, 0])],\n",
"[f(x[7, 1])],\n",
"[f(x[8, 0])],\n",
"[f(x[8, 1])],\n",
"[f(x[9, 0])],\n",
"[f(x[9, 1])]])"
]
},
"execution_count": 248,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eqs.reshape(20, 1)"
]
},
{
"cell_type": "markdown",
"id": "8421f69c",
"metadata": {},
"source": [
"Notice how the organization is [(system 1), (system 2), (system 3), ...] where a \"system\" is a row of the original equation matrix. We can get the jacobian of this directly:"
]
},
{
"cell_type": "code",
"execution_count": 251,
"id": "edb53b48",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{array}{cccccccccccccccccccc}\\frac{\\partial}{\\partial {x}_{0,0}} f{\\left({x}_{0,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{0,1}} f{\\left({x}_{0,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & \\frac{\\partial}{\\partial {x}_{1,0}} f{\\left({x}_{1,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{1,1}} f{\\left({x}_{1,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{2,0}} f{\\left({x}_{2,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{2,1}} f{\\left({x}_{2,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{3,0}} f{\\left({x}_{3,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{3,1}} f{\\left({x}_{3,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{4,0}} f{\\left({x}_{4,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{4,1}} f{\\left({x}_{4,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{5,0}} f{\\left({x}_{5,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{5,1}} f{\\left({x}_{5,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{6,0}} f{\\left({x}_{6,0} \\right)} & 0 & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{6,1}} f{\\left({x}_{6,1} \\right)} & 0 & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{7,0}} f{\\left({x}_{7,0} \\right)} & 0 & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{7,1}} f{\\left({x}_{7,1} \\right)} & 0 & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{8,0}} f{\\left({x}_{8,0} \\right)} & 0 & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{8,1}} f{\\left({x}_{8,1} \\right)} & 0 & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{9,0}} f{\\left({x}_{9,0} \\right)} & 0\\\\0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & \\frac{\\partial}{\\partial {x}_{9,1}} f{\\left({x}_{9,1} \\right)}\\end{array}\\right]$"
],
"text/plain": [
"Matrix([\n",
"[Derivative(f(x[0, 0]), x[0, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, Derivative(f(x[0, 1]), x[0, 1]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, Derivative(f(x[1, 0]), x[1, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, Derivative(f(x[1, 1]), x[1, 1]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, Derivative(f(x[2, 0]), x[2, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, Derivative(f(x[2, 1]), x[2, 1]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, Derivative(f(x[3, 0]), x[3, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[3, 1]), x[3, 1]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[4, 0]), x[4, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[4, 1]), x[4, 1]), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[5, 0]), x[5, 0]), 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[5, 1]), x[5, 1]), 0, 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[6, 0]), x[6, 0]), 0, 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[6, 1]), x[6, 1]), 0, 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[7, 0]), x[7, 0]), 0, 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[7, 1]), x[7, 1]), 0, 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[8, 0]), x[8, 0]), 0, 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[8, 1]), x[8, 1]), 0, 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[9, 0]), x[9, 0]), 0],\n",
"[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Derivative(f(x[9, 1]), x[9, 1])]])"
]
},
"execution_count": 251,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"eqs.reshape(20, 1).jacobian(inputs.reshape(20, 1))"
]
},
{
"cell_type": "markdown",
"id": "a24d1cee",
"metadata": {},
"source": [
"This new thing is a block-diagonal matrix with all the 2x2 jacobians. We can now get the thing we want by either clever slicing, or by clever reshaping.\n",
"\n",
"Note that in pytensor, we don't have to ravel the inputs, so what we'll get back is actually this:"
]
},
{
"cell_type": "code",
"execution_count": 254,
"id": "53b11c96",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{0,0}} f{\\left({x}_{0,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & \\frac{\\partial}{\\partial {x}_{0,1}} f{\\left({x}_{0,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\\\frac{\\partial}{\\partial {x}_{1,0}} f{\\left({x}_{1,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{1,1}} f{\\left({x}_{1,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{2,0}} f{\\left({x}_{2,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{2,1}} f{\\left({x}_{2,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{3,0}} f{\\left({x}_{3,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{3,1}} f{\\left({x}_{3,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{4,0}} f{\\left({x}_{4,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{4,1}} f{\\left({x}_{4,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{5,0}} f{\\left({x}_{5,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{5,1}} f{\\left({x}_{5,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{6,0}} f{\\left({x}_{6,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{6,1}} f{\\left({x}_{6,1} \\right)}\\\\0 & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{7,0}} f{\\left({x}_{7,0} \\right)} & 0\\\\0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{7,1}} f{\\left({x}_{7,1} \\right)}\\\\0 & 0\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{8,0}} f{\\left({x}_{8,0} \\right)} & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{8,1}} f{\\left({x}_{8,1} \\right)}\\\\0 & 0\\end{matrix}\\right]\\\\\\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\\\frac{\\partial}{\\partial {x}_{9,0}} f{\\left({x}_{9,0} \\right)} & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{9,1}} f{\\left({x}_{9,1} \\right)}\\end{matrix}\\right]\\end{matrix}\\right]$"
],
"text/plain": [
"[[[[Derivative(f(x[0, 0]), x[0, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, Derivative(f(x[0, 1]), x[0, 1])], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [Derivative(f(x[1, 0]), x[1, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, Derivative(f(x[1, 1]), x[1, 1])], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [Derivative(f(x[2, 0]), x[2, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, Derivative(f(x[2, 1]), x[2, 1])], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [Derivative(f(x[3, 0]), x[3, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, Derivative(f(x[3, 1]), x[3, 1])], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[4, 0]), x[4, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[4, 1]), x[4, 1])], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[5, 0]), x[5, 0]), 0], [0, 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[5, 1]), x[5, 1])], [0, 0], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[6, 0]), x[6, 0]), 0], [0, 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[6, 1]), x[6, 1])], [0, 0], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[7, 0]), x[7, 0]), 0], [0, 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[7, 1]), x[7, 1])], [0, 0], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[8, 0]), x[8, 0]), 0], [0, 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[8, 1]), x[8, 1])], [0, 0]]], [[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [Derivative(f(x[9, 0]), x[9, 0]), 0]], [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, Derivative(f(x[9, 1]), x[9, 1])]]]]"
]
},
"execution_count": 254,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"block_diag_jac_pt = sp.Array(eqs.reshape(20, 1).jacobian(inputs.reshape(20, 1))).reshape(10, 2, 10, 2)\n",
"block_diag_jac_pt"
]
},
{
"cell_type": "markdown",
"id": "7ed988c3",
"metadata": {},
"source": [
"We're not allowed to do dimension reduction operations in Sympy, but if you take a slice from the 2nd dimension, you can see why the sum works: "
]
},
{
"cell_type": "code",
"execution_count": 256,
"id": "8a695e28",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{0,0}} f{\\left({x}_{0,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{0,1}} f{\\left({x}_{0,1} \\right)}\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right]\\end{matrix}\\right]$"
],
"text/plain": [
"[[[Derivative(f(x[0, 0]), x[0, 0]), 0], [0, Derivative(f(x[0, 1]), x[0, 1])]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]]"
]
},
"execution_count": 256,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"block_diag_jac_pt[:, :, 0, :]"
]
},
{
"cell_type": "code",
"execution_count": 257,
"id": "187c7ee9",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{1,0}} f{\\left({x}_{1,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{1,1}} f{\\left({x}_{1,1} \\right)}\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right]\\end{matrix}\\right]$"
],
"text/plain": [
"[[[0, 0], [0, 0]], [[Derivative(f(x[1, 0]), x[1, 0]), 0], [0, Derivative(f(x[1, 1]), x[1, 1])]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]]"
]
},
"execution_count": 257,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"block_diag_jac_pt[:, :, 1, :]"
]
},
{
"cell_type": "code",
"execution_count": 258,
"id": "d30adb57",
"metadata": {},
"outputs": [
{
"data": {
"text/latex": [
"$\\displaystyle \\left[\\begin{matrix}\\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}\\frac{\\partial}{\\partial {x}_{2,0}} f{\\left({x}_{2,0} \\right)} & 0\\\\0 & \\frac{\\partial}{\\partial {x}_{2,1}} f{\\left({x}_{2,1} \\right)}\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right] & \\left[\\begin{matrix}0 & 0\\\\0 & 0\\end{matrix}\\right]\\end{matrix}\\right]$"
],
"text/plain": [
"[[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[Derivative(f(x[2, 0]), x[2, 0]), 0], [0, Derivative(f(x[2, 1]), x[2, 1])]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]]"
]
},
"execution_count": 258,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"block_diag_jac_pt[:, :, 2, :]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "26aa3cfa",
"metadata": {},
"outputs": [],
"source": []
}
],
"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.11.6"
},
"toc": {
"base_numbering": 1,
"nav_menu": {},
"number_sections": true,
"sideBar": true,
"skip_h1_title": false,
"title_cell": "Table of Contents",
"title_sidebar": "Contents",
"toc_cell": false,
"toc_position": {},
"toc_section_display": true,
"toc_window_display": false
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment