{
"cells": [
{
"cell_type": "markdown",
"id": "69e58b47-c160-40c4-a12e-e1299918b5af",
"metadata": {},
"source": [
"# Contraction\n",
"\n",
"```{hint}\n",
"This page describes the low level contraction functionality of `cotengra`, see\n",
"the [high level interface functions](high-level-interface) for standard use.\n",
"```\n",
"\n",
"You can pass the `cotengra` optimizers, or \n",
"[paths](cotengra.ContractionTree.get_path) \n",
"generated with them, to \n",
"[`quimb`](https://quimb.readthedocs.io/en/latest/), \n",
"[`opt_einsum`](https://dgasmith.github.io/opt_einsum/), \n",
"and other libraries (generally via the `optimize=` kwarg), but `cotengra` also \n",
"provides its own contraction functionality encapsulated in the \n",
"[`ContractionTree.contract`](cotengra.ContractionTree.contract) method.\n",
"\n",
"For an example let's generate a square lattice contraction:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "a079cf1f",
"metadata": {},
"outputs": [],
"source": [
"%config InlineBackend.figure_formats = ['svg']\n",
"import numpy as np\n",
"import cotengra as ctg\n",
"\n",
"inputs, output, shapes, size_dict = ctg.utils.lattice_equation([10, 10])\n",
"arrays = [np.random.uniform(size=shape) for shape in shapes]"
]
},
{
"cell_type": "markdown",
"id": "ffc62ec9",
"metadata": {},
"source": [
"Then find a high quality tree for it (using `minimize='combo'` is generally \n",
"good for real world contraction performance):"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "02cf1f61",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"log2[SIZE]: 10.00 log10[FLOPs]: 5.28: 100%|██████████| 128/128 [00:04<00:00, 28.46it/s]\n"
]
}
],
"source": [
"opt = ctg.HyperOptimizer(\n",
" minimize='combo',\n",
" reconf_opts={},\n",
" progbar=True,\n",
")\n",
"tree = opt.search(inputs, output, size_dict)"
]
},
{
"cell_type": "markdown",
"id": "b509f66e",
"metadata": {},
"source": [
"We can have a quick look at the tree:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "7427229f",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"(, )"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.plot_tent(order=True)"
]
},
{
"cell_type": "markdown",
"id": "fd488036",
"metadata": {},
"source": [
"## Basic contraction\n",
"\n",
"Contracting is then as simple as:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "20ec1c0b",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array(8.02217303e+22)"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.contract(arrays)"
]
},
{
"cell_type": "markdown",
"id": "4630a4eb",
"metadata": {},
"source": [
"We can also show live progress for feedback with `progbar=True`.\n",
"And if the value of contraction is going to be very large or small (e.g. some \n",
"weighted counting problems), then setting `strip_exponents=True` will actively \n",
"rescale intermediate tensors and return the result as a scaled output array \n",
"(the 'mantissa') and the exponent separately:\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "d7842ae3",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 99/99 [00:00<00:00, 16029.81it/s]\n"
]
},
{
"data": {
"text/plain": [
"(1.0, 22.904292025082466)"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.contract(arrays, progbar=True, strip_exponent=True)"
]
},
{
"cell_type": "markdown",
"id": "aedd0ed4",
"metadata": {},
"source": [
"such that the full output would be `mantissa * 10**exponent`."
]
},
{
"cell_type": "markdown",
"id": "743835eb",
"metadata": {},
"source": [
"## Sliced contraction\n",
"\n",
"One of the main reasons to use the contraction abilities of `cotengra` is that\n",
"it handles sliced contractions. Let's generate a sliced tree with a minimum \n",
"of 32 slices:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "47d99a5c",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"log2[SIZE]: 8.00 log10[FLOPs]: 5.91: 100%|██████████| 128/128 [00:04<00:00, 25.76it/s]\n"
]
}
],
"source": [
"opt = ctg.HyperOptimizer(\n",
" minimize='combo',\n",
" slicing_opts={'target_slices': 32}, \n",
" reconf_opts={},\n",
" progbar=True,\n",
")\n",
"tree = opt.search(inputs, output, size_dict)"
]
},
{
"cell_type": "markdown",
"id": "3a072b2c",
"metadata": {},
"source": [
"The sliced indices appear as dashed lines if we plot the tree:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "97920dbd",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {
"needs_background": "light"
},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"(, )"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.plot_tent(order=True)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "4820e980",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"({'ä': SliceInfo(inner=True, ind='ä', size=2, project=None),\n",
" 'æ': SliceInfo(inner=True, ind='æ', size=2, project=None),\n",
" 'è': SliceInfo(inner=True, ind='è', size=2, project=None),\n",
" 'ê': SliceInfo(inner=True, ind='ê', size=2, project=None),\n",
" 'þ': SliceInfo(inner=True, ind='þ', size=2, project=None)},\n",
" 32)"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.sliced_inds, tree.nslices"
]
},
{
"cell_type": "markdown",
"id": "4edd971a",
"metadata": {},
"source": [
"### Automated\n",
"\n",
"`cotengra` can perform all of the necessary slicing and summing in the \n",
"background for you just by calling the `ContractionTree.contract` method as \n",
"usual:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "0c1e9c66",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 32/32 [00:00<00:00, 489.62it/s]\n"
]
},
{
"data": {
"text/plain": [
"8.022173030954557e+22"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tree.contract(arrays, progbar=True)"
]
},
{
"cell_type": "markdown",
"id": "3fc2609b",
"metadata": {},
"source": [
"If you are just interested in the memory savings of slicing then this is \n",
"usually sufficient. \n",
"\n",
"```{note}\n",
"* For sliced trees progress is shown across slices, not individual contractions. \n",
"* You can still set `strip_exponents=True` - but be aware the \n",
"overall exponent is fixed by the first slice contracted.\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "660240c4",
"metadata": {},
"source": [
"### Manual\n",
"\n",
"To control for example parallelization you have to take a slightly more \n",
"hands-on approach using the some combination of the following methods:\n",
"\n",
"* [`ContractionTree.slice_arrays`](cotengra.ContractionTree.slice_arrays) -\n",
" slice the arrays for slice number `i`\n",
"* [`ContractionTree.contract_core`](cotengra.ContractionTree.contract_core) -\n",
" contract a single set of already sliced arrays, e.g. from above function\n",
"* [`ContractionTree.contract_slice`](cotengra.ContractionTree.contract_slice) -\n",
" contract slice `i`, this simply combines the above two functions\n",
"* [`ContractionTree.gather_slices`](cotengra.ContractionTree.gather_slices) -\n",
" gather all the separate sliced outputs and combine into the full output\n",
"\n",
"The slices are enumerated by `range(tree.nslices)`. The following demonstrates\n",
"an explicit loop that you could distribute:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "2c99204b",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"8.022173030954557e+22"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# note this can be a lazy generator\n",
"output_slices = (\n",
" tree.contract_slice(arrays, i)\n",
" for i in range(tree.nslices)\n",
")\n",
"\n",
"tree.gather_slices(output_slices)"
]
},
{
"cell_type": "markdown",
"id": "00ae063b",
"metadata": {},
"source": [
"There are more examples in the `examples` folder, demonstrating using:\n",
"\n",
"* an executor pool\n",
"* GPU\n",
"* MPI \n",
"\n",
"If no sliced indices are in the output then the final gather is just a sum,\n",
"however if this isn't the case then the final gather also involves calls to\n",
"`stack`. If you are just interested in iterating over the output slices (for \n",
"example performing a map-reduce like $\\sum f(T)$) then you can do so lazily with \n",
"[ContractionTree.gen_output_chunks](cotengra.ContractionTree.gen_output_chunks)\n",
"which avoids constructing the full output tensor."
]
},
{
"cell_type": "markdown",
"id": "88776e18",
"metadata": {},
"source": [
"## Backends and `autoray`\n",
"\n",
"`cotengra` dispatches the functions \n",
"[`tensordot`](https://numpy.org/doc/stable/reference/generated/numpy.tensordot.html), \n",
"[`einsum`](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html), \n",
"[`transpose`](https://numpy.org/doc/stable/reference/generated/numpy.transpose.html) \n",
"and \n",
"[`stack`](https://numpy.org/doc/stable/reference/generated/numpy.stack.html)\n",
"using [`autoray`](https://github.com/jcmgray/autoray) thereby \n",
"supporting a wide range of backend tensor libraries, as well as functionality\n",
"like **intermediate reuse**, **lazy tracing**, and **compilation**. This \n",
"usually requires no extra input:"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "ec41916e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array(8.0221730e+22)"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import cupy as cp\n",
"\n",
"gpu_arrays = [cp.asarray(array) for array in arrays]\n",
"tree.contract(gpu_arrays)"
]
},
{
"cell_type": "markdown",
"id": "d18c9959",
"metadata": {},
"source": [
"but you can specify the `backend` kwarg to explicitly control which module to \n",
"lookup the functions from, or directly register your own implementations with \n",
"`autoray`."
]
},
{
"cell_type": "markdown",
"id": "9c446c52",
"metadata": {},
"source": [
"### `prefer_einsum`\n",
"\n",
"If you would like to use `einsum` for every contraction, not just those that\n",
"`tensordot` doesn't support, then set `prefer_einsum=True`.\n",
"\n",
"\n",
"### `autojit`\n",
"\n",
"This is an experimental feature of `autoray` that tries to just in time compile\n",
"and then cache the contraction expression based on the backend. For `numpy` and \n",
"`cupy` for example, this just strips away some python scaffold and offers a \n",
"small speedup if a contraction will be called many times. For `jax` and others \n",
"it makes use of their full compilation abilities."
]
},
{
"cell_type": "markdown",
"id": "273bc654",
"metadata": {},
"source": [
"\n",
"## Differences with opt_einsum\n",
"\n",
"The major difference between contraction in `cotengra` and `opt_einsum` is just\n",
"that `cotengra` handles slicing. There are also a few other minor differences \n",
"between the two:\n",
"\n",
"* `opt_einsum` prefers to use `einsum` (if available) for *some* operations that \n",
" `tensordot` could be used for (e.g. outer products), `cotengra` prefers \n",
" `tensordot` for all operations that are possible with it\n",
"* cotengra uses [`autoray`](https://github.com/jcmgray/autoray) to dispatch \n",
" pairwise `tensordot`, `transpose` and `einsum` operations\n",
"* the costs `opt_einsum` reports assumes a real datatype for the tensors, which\n",
" is generally twice as high compared to generic 'ops' (i.e. \n",
" [`tree.contraction_cost()`](cotengra.ContractionTree.contraction_cost)) \n",
" and four times lower than assuming a complex datatype (i.e. \n",
" [`tree.total_flops('complex')`](cotengra.ContractionTree.total_flops))."
]
}
],
"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.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}