{ "cells": [ { "cell_type": "markdown", "id": "ea8df034-75fc-4835-8cc6-c656fc884de9", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "# Basic MERA Manipulations & Optimization\n", "\n", "`quimb` has preliminary support for the [MERA](https://arxiv.org/abs/quant-ph/0610099) class of states. Specifically, there is the {class}`~quimb.tensor.tensor_mera.MERA` class which constructs a {class}`~quimb.tensor.tensor_core.TensorNetwork` if supplied with isometries and unitaries.\n", "\n", "{class}`~quimb.tensor.tensor_mera.MERA` inherits from {class}`~quimb.tensor.tensor_1d.TensorNetwork1DVector` and {class}`~quimb.tensor.tensor_1d.TensorNetwork1D` and so in fact shares many methods with {class}`~quimb.tensor.tensor_1d.MatrixProductState`." ] }, { "cell_type": "code", "execution_count": 1, "id": "3b08ebeb-91fe-44c1-95cf-64aadfdb2e54", "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "import quimb as qu\n", "import quimb.tensor as qtn" ] }, { "cell_type": "markdown", "id": "81b10e14-73ac-43a3-a4e1-e1488d3fcd8e", "metadata": {}, "source": [ "First we create the random MERA state (currently the number of sites ``n`` must be a power of 2):" ] }, { "cell_type": "code", "execution_count": 2, "id": "4251b3ce-1c5b-4560-b98b-36f59e75bffa", "metadata": {}, "outputs": [], "source": [ "n = 128\n", "mera = qtn.MERA.rand_invar(n)" ] }, { "cell_type": "markdown", "id": "43f8e1dc-7871-4018-8433-03b449fdaf76", "metadata": {}, "source": [ "We also can set up some default drawing options, namely, to pin the physical (outer) indices in circle:" ] }, { "cell_type": "code", "execution_count": 3, "id": "cd821d5d-46d7-4519-9c06-8c6f35bcdc4d", "metadata": {}, "outputs": [], "source": [ "from math import cos, sin, pi\n", "\n", "fix = {\n", " f'k{i}': (sin(2 * pi * i / n), cos(2 * pi * i / n))\n", " for i in range(n)\n", "}\n", "\n", "# reduce the 'spring constant' k as well\n", "draw_opts = dict(fix=fix, k=0.01)" ] }, { "cell_type": "markdown", "id": "ce42cf98-48f3-4a49-a640-ce36c101ef90", "metadata": {}, "source": [ "By default, the MERA constructor adds the ``'_ISO'`` and ``'_UNI'`` tensor tags to\n", "demark the isometries and unitaries respectively:" ] }, { "cell_type": "code", "execution_count": 4, "id": "3e31ba83-4457-4041-950a-41e3d8d85fbb", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "mera.draw(color=['_UNI', '_ISO'], **draw_opts)" ] }, { "cell_type": "markdown", "id": "38f49da9-26e2-481c-9524-fda4bc98c3ef", "metadata": {}, "source": [ "It also tags each layer with ``'_LAYER2'`` etc.:" ] }, { "cell_type": "code", "execution_count": 5, "id": "6870a42b-2b9f-42e6-9470-862f4fd70a64", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "mera.draw(color=[f'_LAYER{i}' for i in range(7)], **draw_opts)" ] }, { "cell_type": "markdown", "id": "a6ccdf87-ee55-43ed-a0f9-e7ac1534f187", "metadata": {}, "source": [ "Finally, the site-tags of each initial tensor (``'I0'``, ``'I1'``, ``'I3'``, etc.) are propagated up through the isometries and unitaries, forming effective 'lightcones' for each site. Here, for example, we plot the lightcone of site 0, 40, and 80:" ] }, { "cell_type": "code", "execution_count": 6, "id": "a7520ad1-965b-441c-a9af-26b80aa9a6d1", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "mera.draw(color=['I0', 'I40', 'I80'], **draw_opts)" ] }, { "cell_type": "markdown", "id": "0f18ec2f-27fb-46d7-8940-3b099021f65c", "metadata": {}, "source": [ "Computing Local Quantities\n", "----------------------------------------\n", "\n", "In a MERA state, local quantities depend only on this lightcone. The way that ``quimb.tensor`` works supports this very naturally. Firstly, you can easily select all the tensors with site tag ``i``, i.e. the causal cone, with ``MERA.select(i)``:" ] }, { "cell_type": "code", "execution_count": 7, "id": "5856c953-1b14-4e01-a405-1ac0ec14c82e", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "# select all tensors relevant for site-0\n", "mera.select(0).draw(color=[f'_LAYER{i}' for i in range(7)])" ] }, { "cell_type": "markdown", "id": "4cea8df9-ffdb-4ffe-95a7-76cf48a35524", "metadata": {}, "source": [ "Secondly, when combined with its conjugate network, all the dangling indices automatically match up. As an example, consider the state norm, but calculated for site 80 only:" ] }, { "cell_type": "code", "execution_count": 8, "id": "a8cc8211-6dd8-4237-b00a-1703e7df4061", "metadata": {}, "outputs": [], "source": [ "nrm80 = mera.select(80).H & mera.select(80)" ] }, { "cell_type": "code", "execution_count": 9, "id": "55ce920a-508b-4731-a495-accf0d334aac", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "nrm80.draw(color=[f'_LAYER{i}' for i in range(7)])" ] }, { "cell_type": "markdown", "id": "f4c2238f-710c-4ebb-aa33-6c0c9e92ad77", "metadata": {}, "source": [ "We can contract this whole subnetwork efficiently to compute the actual value:" ] }, { "cell_type": "code", "execution_count": 10, "id": "e61104cd-cda2-4d47-a392-097251508a76", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.9999999999999938" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "nrm80 ^ all" ] }, { "cell_type": "markdown", "id": "5823f321-44f3-4f89-b350-f1ee12485f93", "metadata": {}, "source": [ "As expected. Or consider we want to measure $\\langle \\psi | X_i Z_j | \\psi \\rangle$:" ] }, { "cell_type": "code", "execution_count": 11, "id": "19b8df98-ff71-4441-82a8-3b842483670d", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "('I50', 'I100')" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "i, j = 50, 100\n", "ij_tags = mera.site_tag(i), mera.site_tag(j)\n", "ij_tags" ] }, { "cell_type": "markdown", "id": "207e2d85-7353-43dc-94f8-285a8c473c35", "metadata": {}, "source": [ "Now we can select the subnetwork of tensors with *either* the site 50 or site 100 lightcone (and also conjugate to form $\\langle \\psi |$):" ] }, { "cell_type": "code", "execution_count": 12, "id": "7eeed5dc-e0f9-4412-a3ab-a44d2ccdfada", "metadata": {}, "outputs": [], "source": [ "mera_ij_H = mera.select(ij_tags, which='any').H" ] }, { "cell_type": "markdown", "id": "4ce0de05-560b-4850-8a7e-6c9ad1cc1d34", "metadata": {}, "source": [ "For $X_i Z_j | \\psi \\rangle$ we'll first apply the X and Z operators. By default the gate operation propagates the site tags to the applied operators as well, or we could use ``contract=True`` to actively contract them into the MERA:" ] }, { "cell_type": "code", "execution_count": 13, "id": "6b58d6e3-e5a2-408f-a791-083fdf21dfa3", "metadata": {}, "outputs": [], "source": [ "X = qu.pauli('X')\n", "Z = qu.pauli('X')\n", "\n", "XY_mera_ij = (\n", " mera\n", " .gate(X, i)\n", " .gate(Z, j)\n", " .select(ij_tags, which='any')\n", ")" ] }, { "cell_type": "markdown", "id": "43efc0b1-33fc-4240-ba91-485690e3b5bf", "metadata": {}, "source": [ "Now we can lazily form the tensor network of this expectation value:" ] }, { "cell_type": "code", "execution_count": 14, "id": "bfa3bae9-2168-493e-beaf-aba865e417f8", "metadata": {}, "outputs": [], "source": [ "exp_XZ_ij = (mera_ij_H & XY_mera_ij)" ] }, { "cell_type": "code", "execution_count": 15, "id": "0de19a83-1299-4d14-8ed0-3a3f87c6e117", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "exp_XZ_ij.draw(color=[f'_LAYER{i}' for i in range(7)])" ] }, { "cell_type": "markdown", "id": "f226d441-2490-49c3-a27f-f3c4e2b2870e", "metadata": {}, "source": [ "Which we can efficiently contract:" ] }, { "cell_type": "code", "execution_count": 16, "id": "d5127451-3f7b-4107-b95d-62e35e0fd4f4", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.15883517895692317" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "exp_XZ_ij ^ all" ] }, { "cell_type": "code", "execution_count": 17, "id": "7de364e7-1446-4372-85c6-79c72fbc4daa", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2.14 ms ± 43.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "%%timeit\n", "exp_XZ_ij ^ all" ] }, { "cell_type": "raw", "id": "b5e31a47-c562-43ec-bb73-25819391de3a", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "Forming a density matrix\n", "------------------------\n", "\n", "For some quantities, e.g. entanglement measures, one needs a representation of the density matrix rather than just local operator expectations. We form this by combining a 'bra' and 'ket' of the MERA, then tracing out an environment. Again, :mod:`quimb.tensor` handles this very naturally, all we need to do is reindex *only* the physical indices we want to keep on the conjugated state, any that are left will then naturally be contracted, exactly like a partial trace.\n", "\n", "Let's form a density matrix for the first 20 qubits:" ] }, { "cell_type": "code", "execution_count": 18, "id": "21f0ee3c-095e-4694-8b17-86fd1590b9df", "metadata": {}, "outputs": [], "source": [ "# generate the 'bra' state\n", "mera_H = mera.H.reindex_sites('b{}', range(20))" ] }, { "cell_type": "markdown", "id": "36663f94-ac3a-41c2-bdbc-2e71dfc5e90d", "metadata": {}, "source": [ "We again only need the tensors in the causal cones of these 20 sites:" ] }, { "cell_type": "code", "execution_count": 19, "id": "9562cbbc-b60f-4f37-a046-0e5ddbfb9544", "metadata": {}, "outputs": [], "source": [ "# NB we have to slice *before* combining the subnetworks here.\n", "# this is because paired indices are mangled when joining\n", "# two networks -> only dangling indices are guranteed to\n", "# retain their value\n", "rho = (\n", " mera_H.select(slice(20), which='any') &\n", " mera.select(slice(20), which='any')\n", ")" ] }, { "cell_type": "markdown", "id": "b17c5110-84a7-4703-9f6b-f2c698feac5d", "metadata": {}, "source": [ "We can see what this density operator looks like as a tensor network:" ] }, { "cell_type": "code", "execution_count": 20, "id": "e27a3791-21bd-4ab8-983a-fabb39138bca", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "rho.draw(color=[f'_LAYER{i}' for i in range(7)])" ] }, { "cell_type": "markdown", "id": "59df2a77-49ad-4c0e-9f27-665de63fd93d", "metadata": {}, "source": [ "Or we can plot the sites (note how higher and higher tensors have more and more site tags):" ] }, { "cell_type": "code", "execution_count": 21, "id": "70372f8f-63e3-46a4-9648-7e435933be5a", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "rho.draw(color=[f'I{i}' for i in range(20)])" ] }, { "cell_type": "markdown", "id": "bfa104d2-3201-4da8-88b2-912e8b130863", "metadata": {}, "source": [ "This density matrix is too big to explicitly form (it would need $2^{40}$, about a trillion, elements).\n", "On the other hand we can treat it as linear operator, in which case we only need to compute its action on\n", "a vector of size $2^{20}$. This allows the computation of 'spectral' quantities of the form $\\text{Tr}(f(\\rho))$.\n", "\n", "One such quantity is the entropy $-\\text{Tr} \\left( \\rho \\log_2 \\rho \\right) $:" ] }, { "cell_type": "code", "execution_count": 22, "id": "51d78632-ed6f-4646-808d-2c1a2e3032c0", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "<1048576x1048576 TNLinearOperator with dtype=float64>" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# mark the indices as belonging to either the 'left' or\n", "# 'right' hand side of the effective operator\n", "left_ix = [f'k{i}' for i in range(20)]\n", "rght_ix = [f'b{i}' for i in range(20)]\n", "\n", "# form the linear operator\n", "rho_ab = rho.aslinearoperator(left_ix, rght_ix, backend='torch')\n", "rho_ab" ] }, { "cell_type": "raw", "id": "2eaeba15-fbd4-4c85-951c-767057313ddd", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "Now we can compute using :func:`~quimb.linalg.approx_spectral.approx_spectral_function` (note this can be made radically faster if a GPU is available and `cupy `_ is installed):" ] }, { "cell_type": "code", "execution_count": 23, "id": "641920e8-a590-403c-8f37-1c4f855f03c1", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "rho_entropy ~ 5.454770371829187\n" ] } ], "source": [ "f = qu.xlogx\n", "S = - qu.approx_spectral_function(rho_ab, f, tol=0.02)\n", "print(\"rho_entropy ~\", S)" ] }, { "cell_type": "markdown", "id": "65d128ea-fb8c-440b-8442-a0a62eb90bc1", "metadata": {}, "source": [ "To compute a genuine entanglement measure we need a further small trick. Specifically, if we are computing the negativity between subsystem A and subsystem B, we can perform the [partial transpose](https://en.wikipedia.org/wiki/Peres%E2%80%93Horodecki_criterion) simply by swapping subsystem B's 'left' indices for right indices. This creates a linear operator of $\\rho_{AB}^{T_B}$, which we can compute the logarithmic negativity for,\n", "$\\mathcal{E} = \\log_2 \\text{Tr} |\\rho_{AB}^{T_B}|$:" ] }, { "cell_type": "code", "execution_count": 24, "id": "9c9b551e-7603-4c57-9602-0a3a31690888", "metadata": {}, "outputs": [], "source": [ "# partition 20 spins in two\n", "sysa = range(0, 10)\n", "sysb = range(10, 20)\n", "\n", "# k0, k1, k2, ... b10, b11, b12, ...\n", "left_ix = [f'k{i}' for i in sysa] + [f'b{i}' for i in sysb]\n", "# b0, b1, b2, ... k10, k11, k12, ...\n", "rght_ix = [f'b{i}' for i in sysa] + [f'k{i}' for i in sysb]\n", "\n", "rho_ab_pt = rho.aslinearoperator(left_ix, rght_ix, backend='torch')" ] }, { "cell_type": "markdown", "id": "699c906e-5d0c-4af4-830c-c618a1fc5304", "metadata": {}, "source": [ "Now we just to to take ``abs`` as the function $f$ and scale the result with $\\log_2$:" ] }, { "cell_type": "code", "execution_count": 25, "id": "34fd3231-26de-4458-909e-4e5ddf15a43e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "rho_ab logarithmic negativity ~ 1.5515528003023578\n" ] } ], "source": [ "f = abs\n", "neg = qu.approx_spectral_function(rho_ab_pt, f, tol=0.02)\n", "print(\"rho_ab logarithmic negativity ~\", qu.log2(neg))" ] }, { "cell_type": "raw", "id": "ad2a1e83-b8d7-43e0-8a21-5326ffc5cb74", "metadata": { "raw_mimetype": "text/restructuredtext" }, "source": [ "Globally Optimizing a MERA\n", "--------------------------\n", "\n", "Since we can efficiently contract local quantities we can evaluate the energy\n", "of a local hamiltonian as a simple sum. If we pass such a function to a\n", ":class:`~quimb.tensor.optimize.TNOptimizer` it can differentiate the\n", "calculation and efficiently globally optimize every parameter in the MERA\n", "with respect the energy.\n", "\n", "Here we'll demonstrate that for the Heisenberg model (without assuming\n", "any translational or scale invariance). First the initial MERA:" ] }, { "cell_type": "code", "execution_count": 26, "id": "754a5d20-c0b2-4774-a682-79ca2a43ad08", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# total length (currently must be power of 2)\n", "L = 2**6\n", "\n", "# max bond dimension\n", "D = 8\n", "\n", "# use single precision for quick GPU optimization\n", "dtype = 'float32'\n", "\n", "mera = qtn.MERA.rand(L, max_bond=D, dtype=dtype)\n", "\n", "# this is the function that projects all tensors\n", "# with ``left_inds`` into unitary / isometric form\n", "mera.unitize_()" ] }, { "cell_type": "markdown", "id": "153ce978-c013-423c-9a27-0babf37ed29e", "metadata": {}, "source": [ "Then the dictionary of local terms describing the Hamiltonian:" ] }, { "cell_type": "code", "execution_count": 27, "id": "db3c239b-fb31-4fd9-ad97-f9f5b64767d9", "metadata": {}, "outputs": [], "source": [ "H2 = qu.ham_heis(2).real.astype(dtype)\n", "terms = {(i, (i + 1) % L): H2 for i in range(L)}" ] }, { "cell_type": "markdown", "id": "0eb04f86-a4ae-4d7d-95af-208baa157a99", "metadata": {}, "source": [ "For which we can compute the exact energy:" ] }, { "cell_type": "code", "execution_count": 28, "id": "059293cd-abeb-45fd-9de6-62a4ac9fb004", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "-28.374337597812207" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "if L <= 20:\n", " # numerical result\n", " en = qu.groundenergy(qu.ham_heis(L, cyclic=True, sparse=True))\n", "else:\n", " # analytic result for long (PBC) chains\n", " en = qu.heisenberg_energy(L)\n", "\n", "en" ] }, { "cell_type": "markdown", "id": "05d16eba-6f35-4448-b7d5-33e803613548", "metadata": {}, "source": [ "Then we need two functions to supply to the optimizer:" ] }, { "cell_type": "markdown", "id": "05445508-55a3-4bac-9daf-c1e798fedd9c", "metadata": {}, "source": [ "* ``norm_fn``, which projects the TN into the constrained form (i.e. isometric)" ] }, { "cell_type": "code", "execution_count": 29, "id": "f7153942-5c1a-4b2c-8ea9-019151ed8718", "metadata": {}, "outputs": [], "source": [ "def norm_fn(mera):\n", " # there are a few methods to do the projection\n", " # exp works well for optimization\n", " return mera.unitize(method='exp')" ] }, { "cell_type": "markdown", "id": "77e4cfc5-4f38-4f8c-a24d-8a1f1c95cac3", "metadata": {}, "source": [ "* ``loss_fn``, which takes the projected TN and computes the scalar to minimize" ] }, { "cell_type": "code", "execution_count": 30, "id": "5fba9f11-391e-43a3-9c0a-6c25d4bf5589", "metadata": {}, "outputs": [], "source": [ "def local_expectation(mera, terms, where, optimize='auto-hq'):\n", " \"\"\"Compute the energy for a single local term.\n", " \"\"\"\n", " # get the lightcone for `where`\n", " tags = [mera.site_tag(coo) for coo in where]\n", " mera_ij = mera.select(tags, 'any')\n", "\n", " # apply the local gate\n", " G = terms[where]\n", " mera_ij_G = mera_ij.gate(terms[where], where)\n", "\n", " # compute the overlap - this is where the real computation happens\n", " mera_ij_ex = (mera_ij_G & mera_ij.H)\n", " return mera_ij_ex.contract(all, optimize=optimize)\n", "\n", "\n", "def loss_fn(mera, terms, **kwargs):\n", " \"\"\"Compute the total energy as a sum of all terms.\n", " \"\"\"\n", " return sum(\n", " local_expectation(mera, terms, where, **kwargs)\n", " for where in terms\n", " )" ] }, { "cell_type": "markdown", "id": "22b29cdd-edaa-4cac-b957-9b5fbbbe9598", "metadata": {}, "source": [ "To find a high quality contraction path for each term we'll also instantiate a ``cotengra`` optimizer:" ] }, { "cell_type": "code", "execution_count": 31, "id": "c9112e4e-15a3-43af-a094-c01ea054fa21", "metadata": {}, "outputs": [], "source": [ "import cotengra as ctg\n", "\n", "opt = ctg.ReusableHyperOptimizer(\n", " progbar=True,\n", " reconf_opts={},\n", " max_repeats=16,\n", " # directory= # set this for persistent cache\n", ")" ] }, { "cell_type": "markdown", "id": "2b1106fd-bd62-4dba-922d-a1be85a8b6b5", "metadata": {}, "source": [ "The first time each contraction is encountered it will be optimized, after with\n", "the path will be saved for matching contractions:" ] }, { "cell_type": "code", "execution_count": 32, "id": "02ae13ae-678e-40d7-8640-0ef8c5c70f86", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.95it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.94it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:02<00:00, 5.35it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.27it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.55it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.24it/s]\n", "log2[SIZE]: 18.00 log10[FLOPs]: 8.61: 100%|██████████| 16/16 [00:02<00:00, 5.34it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.86it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 5.01it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.85it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.81it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.31it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.25it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.78it/s]\n", "log2[SIZE]: 12.00 log10[FLOPs]: 6.87: 100%|██████████| 16/16 [00:03<00:00, 5.28it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.53it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.13it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.45it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:02<00:00, 5.65it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.87it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 5.13it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.34it/s]\n", "log2[SIZE]: 18.00 log10[FLOPs]: 8.61: 100%|██████████| 16/16 [00:02<00:00, 5.61it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:02<00:00, 5.36it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.86it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.91it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.62it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.69it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:03<00:00, 4.69it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.49it/s]\n", "log2[SIZE]: 12.00 log10[FLOPs]: 6.81: 100%|██████████| 16/16 [00:02<00:00, 6.24it/s]\n", "log2[SIZE]: 19.00 log10[FLOPs]: 9.01: 100%|██████████| 16/16 [00:04<00:00, 3.68it/s]\n" ] }, { "data": { "text/plain": [ "-0.021592675417196006" ] }, "execution_count": 32, "metadata": {}, "output_type": "execute_result" } ], "source": [ "loss_fn(norm_fn(mera), terms, optimize=opt)" ] }, { "cell_type": "markdown", "id": "45dc524d-c5d8-475d-9aca-24c31807c0b4", "metadata": {}, "source": [ "These sizes and FLOPS are all very manageable. And a second call should be\n", "much faster:" ] }, { "cell_type": "code", "execution_count": 33, "id": "f1a1cafc-a810-43d2-9694-f9ec923b7f48", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 5.06 s, sys: 39.5 ms, total: 5.1 s\n", "Wall time: 662 ms\n" ] }, { "data": { "text/plain": [ "-0.021592675417196006" ] }, "execution_count": 33, "metadata": {}, "output_type": "execute_result" } ], "source": [ "%%time\n", "loss_fn(norm_fn(mera), terms, optimize=opt)" ] }, { "cell_type": "markdown", "id": "35af1d10-a85e-4231-aa21-0670f9aa4dc1", "metadata": {}, "source": [ "Now we are ready to set-up our optimizer object:" ] }, { "cell_type": "code", "execution_count": 34, "id": "10a0647d-d240-4499-91b1-2327dd457e18", "metadata": {}, "outputs": [], "source": [ "tnopt = qtn.TNOptimizer(\n", " mera,\n", " loss_fn=loss_fn,\n", " norm_fn=norm_fn,\n", " loss_constants={'terms': terms},\n", " loss_kwargs={'optimize': opt},\n", " autodiff_backend='torch', device='cuda', jit_fn=True,\n", ")" ] }, { "cell_type": "markdown", "id": "01359c35-6058-439e-8843-4fb8d36151c9", "metadata": {}, "source": [ "Since we set ``jit_fn=True``, the first step involves compiling the computation,\n", "which might take some time and print some (ignorable) warnings:" ] }, { "cell_type": "code", "execution_count": 35, "id": "1adb5e8a-73ab-4bdc-8dda-6a6713dd4861", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/1 [00:00" ] }, "execution_count": 35, "metadata": {}, "output_type": "execute_result" } ], "source": [ "%%time\n", "tnopt.optimize(1)" ] }, { "cell_type": "markdown", "id": "b548ad81-26c3-475f-a9a9-0dd420b85f87", "metadata": {}, "source": [ "At this point every iteration should be very fast:" ] }, { "cell_type": "code", "execution_count": 36, "id": "54ce22b5-0fc5-43de-96f5-09f2bc2fefdf", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "-28.175458908081055: 100%|██████████| 999/999 [05:34<00:00, 2.99it/s]\n" ] } ], "source": [ "tnopt.optimizer = 'l-bfgs-b' # the default\n", "mera_opt = tnopt.optimize(999)" ] }, { "cell_type": "code", "execution_count": 37, "id": "f15e1856-46c3-45c5-b586-d0bbe1379cf5", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "-28.269123077392578: 100%|██████████| 1000/1000 [05:01<00:00, 3.32it/s]\n" ] } ], "source": [ "tnopt.optimizer = 'adam' # useful for final iterations\n", "mera_opt = tnopt.optimize(1000)" ] }, { "cell_type": "code", "execution_count": 38, "id": "5523b6b2-032c-47d0-98e4-cef761c44e74", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 38, "metadata": {}, "output_type": "execute_result" }, { "data": { "image/svg+xml": [ "" ], "text/plain": [ "

" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "rel_err = (np.array(tnopt.losses) - en) / abs(en)\n", "\n", "plt.plot(rel_err)\n", "plt.xlabel('Iteration [ l-bfgs-m -> adam ]')\n", "plt.ylabel('Relative Error')\n", "plt.yscale('log')\n", "plt.axvline(len(tnopt.losses) - 1000, color='black')" ] } ], "metadata": { "celltoolbar": "Raw Cell Format", "kernelspec": { "display_name": "Python 3", "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" } }, "nbformat": 4, "nbformat_minor": 4 }