{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "c54f04b8-37a4-44bb-8969-9d8370ab98f2", "metadata": {}, "source": [ "(tensor-network-optimization)=\n", "\n", "# Optimization\n", "\n", "`quimb` supports optimizing an abtirary tensor network with respect to an arbitrary 'loss' function using automatic differentiation. This is encapsulated in the [`TNOptimizer`](quimb.tensor.optimize.TNOptimizer) object. Internally, this makes use of one of various autodiff libraries to compute the neccesary tensor gradients, then maps these tensor parameters into a 'ravelled' single real vector for optimization in an outer loop by {func}`scipy.optimize.minimize`.\n", "\n", "```{note}\n", "You can also use `quimb` simply as a way to orchestrate operations on e.g. `torch` or `jax` arrays, and then use optimiztion libraries from those frameworks directly. This is demonstrated in the examples {ref}`quimb-within-torch` and {ref}`quimb-within-jax`.\n", "```\n", "\n", "The basic steps are:\n", "\n", "1. Define a **target tensor network** (or pytree[^pytree] of `TensorNetwork`, `Tensor`, raw array, or `Circuit` objects).\n", "2. *Optionally* if its a single tensor network, define sets of tags that specify which tensors to optimize, keep constant, or share parameters among.\n", "3. *Optionally* define a `norm_fn` that takes the target and projects or constrains it to some valid form, e.g. normalizing or mapping into unitary form. This is the identity by default.\n", "4. Define a **loss function** that takes `tn` (or `norm_fn(tn)`) and returns a single real scalar to minimize.\n", "\n", "Here we'll demonstrate this with a PBC MPS optimization for the Heisenberg model.\n", "\n", "[^pytree]: Any nested combination of `dict`, `list`, and `tuple` objects." ] }, { "cell_type": "code", "execution_count": 3, "id": "652da1b7-7dba-4a4f-8675-892474d4eb09", "metadata": {}, "outputs": [], "source": [ "%config InlineBackend.figure_formats = ['svg']\n", "import quimb as qu\n", "import quimb.tensor as qtn" ] }, { "cell_type": "code", "execution_count": 4, "id": "61708615-4fed-46b8-b9c8-b286acc3fb4b", "metadata": {}, "outputs": [], "source": [ "L = 64\n", "D = 16\n", "pbc = True\n", "\n", "# create a random MPS as our initial target to optimize\n", "psi = qtn.MPS_rand_state(L, bond_dim=D, cyclic=pbc)\n", "\n", "# create the hamiltonian MPO, this is a constant TN not to be optimized\n", "ham = qtn.MPO_ham_heis(L, cyclic=pbc)" ] }, { "cell_type": "markdown", "id": "1d70582b-a30e-4457-a924-d4bfca4b94c9", "metadata": {}, "source": [ "Next we define our `norm_fn`, which here just normalizes the MPS, and our `loss_fn`, which computes the energy of the Heisenberg model by exactly contracting an MPS-MPO-MPS overlap." ] }, { "cell_type": "code", "execution_count": 5, "id": "c57c1f96-fbad-4c47-b8e4-afec61590e9f", "metadata": {}, "outputs": [], "source": [ "def norm_fn(psi):\n", " # we could always define this within the loss function, but separating it\n", " # out can be clearer - it's also called before returning the optimized TN\n", " nfact = (psi.H @ psi)**0.5\n", " return psi.multiply(1 / nfact, spread_over='all')\n", "\n", "\n", "def loss_fn(psi, ham):\n", " b, h, k = qtn.tensor_network_align(psi.H, ham, psi)\n", " energy_tn = b | h | k\n", " return energy_tn ^ ..." ] }, { "cell_type": "markdown", "id": "90ccfc56-7641-4053-8324-9f9ae791879f", "metadata": {}, "source": [ "We can check the initial loss value with:" ] }, { "cell_type": "code", "execution_count": 6, "id": "9535c227-0f2d-42db-9cf6-e545bf03366d", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "-0.16003514944501307" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "loss_fn(norm_fn(psi), ham)" ] }, { "cell_type": "markdown", "id": "0f09042b-fa87-4ad9-91fc-e09ef2cf696a", "metadata": {}, "source": [ "Next we supply these to a [`TNOptimizer`](quimb.tensor.optimize.TNOptimizer) object. Since we have an extra tensor object `ham` that is needed to compute the loss, but should not be optimized, we pass it in `loss_constants`, that allows it to be converted to the correct backend etc." ] }, { "cell_type": "code", "execution_count": 7, "id": "090eb325-5af4-418f-979b-bfb539893f4e", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt = qtn.TNOptimizer(\n", " # the tensor network we want to optimize\n", " psi,\n", " # the functions specfying the loss and normalization\n", " loss_fn=loss_fn,\n", " norm_fn=norm_fn,\n", " # we specify constants so that the arguments can be converted\n", " # to the desired autodiff backend automatically\n", " loss_constants={\"ham\": ham},\n", " # the underlying algorithm to use for the optimization\n", " # 'l-bfgs-b' is the default and often good for fast initial progress\n", " optimizer=\"adam\",\n", " # which gradient computation backend to use\n", " autodiff_backend=\"jax\",\n", ")\n", "tnopt" ] }, { "cell_type": "markdown", "id": "3668f6cf-d946-4728-b325-792101d977e3", "metadata": {}, "source": [ "```{hint}\n", "You can also pass general non-numeric or tensor options in ``loss_kwargs``.\n", "```\n", "\n", "We can see there are 32,768 parameters to optimize, which would be tricky without gradients. We are ready to start optimizing (note for backens like `jax` which compile the computation by default, there will be some initial overhead):" ] }, { "cell_type": "code", "execution_count": 8, "id": "97f603f1-b5d7-4d38-b639-4fde48bc3e6b", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "-28.295518875122 [best: -28.295518875122] : : 1001it [01:03, 15.81it/s] \n" ] } ], "source": [ "psi_opt = tnopt.optimize(1000)" ] }, { "cell_type": "markdown", "id": "a24d9191-a1aa-40ff-8b2a-96aba4258d7b", "metadata": {}, "source": [ "There is a simple [`tnopt.plot`](quimb.tensor.optimize.TNOptimizer.plot) method to visualize the loss progress (note by default the first 20 points are shown on a linear plot, the rest on a log plot):" ] }, { "cell_type": "code", "execution_count": 9, "id": "6ac927f4-4c1a-46d6-9ab0-a255a95b8068", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "text/plain": [ "(
, )" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt.plot(hlines={'analytic': qu.heisenberg_energy(L)})" ] }, { "cell_type": "markdown", "id": "23879a36-5819-43b5-b7f1-1be6c5ad3b0f", "metadata": {}, "source": [ "We can check the returned `psi_opt` optimized target indeed matches loss:" ] }, { "cell_type": "code", "execution_count": 10, "id": "429c3953-ee06-4ce1-9b89-b79468daf879", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "-28.295518668947235" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "loss_fn(psi_opt, ham)" ] }, { "cell_type": "markdown", "id": "56d72083-c659-41a3-b8d0-554ba2a7fbf8", "metadata": {}, "source": [ "Note this TN (which can be retrieved from [`tnopt.get_tn_opt`](quimb.tensor.optimize.TNOptimizer.get_tn_opt)) is a copy of the original target TN, with the optimized parameters set. It has also been passed through `norm_fn` so is in normalized/projected form, and converted back to `numpy` backed arrays." ] }, { "attachments": {}, "cell_type": "markdown", "id": "f67f1e14-bf72-4424-8906-69409478a5b8", "metadata": {}, "source": [ "## Using tags to opt in, opt out or group tensors\n", "\n", "There are three mutually exclusive options when it comes to specifying exactly which tensors to optimize.\n", "\n", "1. **Opt in**: specify `tags` that tensors must have to be **optimized**, any tensors without these tags are assumed to be **constant**.\n", "2. **Opt out**: specify a set of `constant_tags` that tensors must have to be **constant**, any tensors without these tags are assumed to be **optimized**.\n", "3. **pytree**: supply an arbitrary pytree of objects to use however within the `norm_fn` and `loss_fn`, in this case all tensors are assumed to be **optimized**.\n", "\n", "In all cases you can supply `loss_constants` that are passed to the `loss_fn` but not optimized, this should be a `dict` containing arbitrary pytree values.\n", "\n", "In the first case, you can also specify `shared_tags`, which we demonstrate here. Here every tensor with one of shared tags it assumed to share parameters with every other tensor with the same tag. For example, we can create a unit cell of size 2, and specify that our MPS be compsed of two repeated tensors A and B only:" ] }, { "cell_type": "code", "execution_count": 11, "id": "d988cc62-c9f5-43f9-89a2-e627eb581f80", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
MatrixProductState(tensors=64, indices=128, L=64, max_bond=16)
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAA, _8db77cAAAAB, k0], tags={I0, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAB, _8db77cAAAAC, k1], tags={I1, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAC, _8db77cAAAAD, k2], tags={I2, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAD, _8db77cAAAAE, k3], tags={I3, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAE, _8db77cAAAAF, k4], tags={I4, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAF, _8db77cAAAAG, k5], tags={I5, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAG, _8db77cAAAAH, k6], tags={I6, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAH, _8db77cAAAAI, k7], tags={I7, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAI, _8db77cAAAAJ, k8], tags={I8, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAJ, _8db77cAAAAK, k9], tags={I9, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAK, _8db77cAAAAL, k10], tags={I10, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAL, _8db77cAAAAM, k11], tags={I11, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAM, _8db77cAAAAN, k12], tags={I12, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAN, _8db77cAAAAO, k13], tags={I13, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAO, _8db77cAAAAP, k14], tags={I14, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAP, _8db77cAAAAQ, k15], tags={I15, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAQ, _8db77cAAAAR, k16], tags={I16, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAR, _8db77cAAAAS, k17], tags={I17, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAS, _8db77cAAAAT, k18], tags={I18, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAT, _8db77cAAAAU, k19], tags={I19, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAU, _8db77cAAAAV, k20], tags={I20, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAV, _8db77cAAAAW, k21], tags={I21, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAW, _8db77cAAAAX, k22], tags={I22, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAX, _8db77cAAAAY, k23], tags={I23, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAY, _8db77cAAAAZ, k24], tags={I24, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAZ, _8db77cAAAAa, k25], tags={I25, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAa, _8db77cAAAAb, k26], tags={I26, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAb, _8db77cAAAAc, k27], tags={I27, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAc, _8db77cAAAAd, k28], tags={I28, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAd, _8db77cAAAAe, k29], tags={I29, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAe, _8db77cAAAAf, k30], tags={I30, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAf, _8db77cAAAAg, k31], tags={I31, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAg, _8db77cAAAAh, k32], tags={I32, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAh, _8db77cAAAAi, k33], tags={I33, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAi, _8db77cAAAAj, k34], tags={I34, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAj, _8db77cAAAAk, k35], tags={I35, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAk, _8db77cAAAAl, k36], tags={I36, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAl, _8db77cAAAAm, k37], tags={I37, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAm, _8db77cAAAAn, k38], tags={I38, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAn, _8db77cAAAAo, k39], tags={I39, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAo, _8db77cAAAAp, k40], tags={I40, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAp, _8db77cAAAAq, k41], tags={I41, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAq, _8db77cAAAAr, k42], tags={I42, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAr, _8db77cAAAAs, k43], tags={I43, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAs, _8db77cAAAAt, k44], tags={I44, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAt, _8db77cAAAAu, k45], tags={I45, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAu, _8db77cAAAAv, k46], tags={I46, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAv, _8db77cAAAAw, k47], tags={I47, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAw, _8db77cAAAAx, k48], tags={I48, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAx, _8db77cAAAAy, k49], tags={I49, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAy, _8db77cAAAAz, k50], tags={I50, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAz, _8db77cAAABA, k51], tags={I51, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABA, _8db77cAAABB, k52], tags={I52, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABB, _8db77cAAABC, k53], tags={I53, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABC, _8db77cAAABD, k54], tags={I54, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABD, _8db77cAAABE, k55], tags={I55, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABE, _8db77cAAABF, k56], tags={I56, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABF, _8db77cAAABG, k57], tags={I57, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABG, _8db77cAAABH, k58], tags={I58, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABH, _8db77cAAABI, k59], tags={I59, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABI, _8db77cAAABJ, k60], tags={I60, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABJ, _8db77cAAABK, k61], tags={I61, B}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABK, _8db77cAAABL, k62], tags={I62, A}),backend=numpy, dtype=float64, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABL, _8db77cAAAAA, k63], tags={I63, B}),backend=numpy, dtype=float64, data=...
" ], "text/plain": [ "MatrixProductState(tensors=64, indices=128, L=64, max_bond=16)" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "for site in psi.sites:\n", " t = psi[site]\n", " if site % 2 == 0:\n", " t.add_tag(\"A\")\n", " t.modify(data=psi[0].data)\n", " else:\n", " t.add_tag(\"B\")\n", " t.modify(data=psi[1].data)\n", "\n", "psi.normalize()\n", "psi.equalize_norms_()" ] }, { "cell_type": "code", "execution_count": 12, "id": "0eba2cc5-fad6-486c-b247-af717f8b6696", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "psi.draw([\"A\", \"B\"], figsize=(4, 4))" ] }, { "cell_type": "code", "execution_count": 13, "id": "83427e03-ba0c-494a-98d3-169c70ae692f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt = qtn.TNOptimizer(\n", " psi,\n", " loss_fn=loss_fn,\n", " norm_fn=norm_fn,\n", " loss_constants={\"ham\": ham},\n", " optimizer=\"adam\",\n", " autodiff_backend=\"jax\",\n", " # only optimize the tensors with these tags (in this case all)\n", " tags=[\"A\", \"B\"],\n", " # within those, group all with each of these tags together\n", " shared_tags=[\"A\", \"B\"],\n", ")\n", "tnopt" ] }, { "cell_type": "markdown", "id": "b96c6021-7ce6-4724-8770-f36f5779da8c", "metadata": {}, "source": [ "You can see the dramatic reduction in the number of parameters to optimize, from 32,768 to 1,024. The optimization proceeds in the same way as before, but now the tensors A and B are constrained to be the same." ] }, { "cell_type": "code", "execution_count": 14, "id": "2ef1de72-8919-419d-bec4-0d80eaa86b3d", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/1000 [00:00
MatrixProductState(tensors=64, indices=128, L=64, max_bond=16)
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAA, _8db77cAAAAB, k0], tags={I0, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAB, _8db77cAAAAC, k1], tags={I1, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAC, _8db77cAAAAD, k2], tags={I2, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAD, _8db77cAAAAE, k3], tags={I3, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAE, _8db77cAAAAF, k4], tags={I4, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAF, _8db77cAAAAG, k5], tags={I5, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAG, _8db77cAAAAH, k6], tags={I6, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAH, _8db77cAAAAI, k7], tags={I7, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAI, _8db77cAAAAJ, k8], tags={I8, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAJ, _8db77cAAAAK, k9], tags={I9, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAK, _8db77cAAAAL, k10], tags={I10, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAL, _8db77cAAAAM, k11], tags={I11, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAM, _8db77cAAAAN, k12], tags={I12, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAN, _8db77cAAAAO, k13], tags={I13, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAO, _8db77cAAAAP, k14], tags={I14, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAP, _8db77cAAAAQ, k15], tags={I15, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAQ, _8db77cAAAAR, k16], tags={I16, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAR, _8db77cAAAAS, k17], tags={I17, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAS, _8db77cAAAAT, k18], tags={I18, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAT, _8db77cAAAAU, k19], tags={I19, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAU, _8db77cAAAAV, k20], tags={I20, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAV, _8db77cAAAAW, k21], tags={I21, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAW, _8db77cAAAAX, k22], tags={I22, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAX, _8db77cAAAAY, k23], tags={I23, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAY, _8db77cAAAAZ, k24], tags={I24, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAZ, _8db77cAAAAa, k25], tags={I25, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAa, _8db77cAAAAb, k26], tags={I26, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAb, _8db77cAAAAc, k27], tags={I27, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAc, _8db77cAAAAd, k28], tags={I28, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAd, _8db77cAAAAe, k29], tags={I29, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAe, _8db77cAAAAf, k30], tags={I30, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAf, _8db77cAAAAg, k31], tags={I31, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAg, _8db77cAAAAh, k32], tags={I32, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAh, _8db77cAAAAi, k33], tags={I33, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAi, _8db77cAAAAj, k34], tags={I34, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAj, _8db77cAAAAk, k35], tags={I35, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAk, _8db77cAAAAl, k36], tags={I36, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAl, _8db77cAAAAm, k37], tags={I37, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAm, _8db77cAAAAn, k38], tags={I38, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAn, _8db77cAAAAo, k39], tags={I39, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAo, _8db77cAAAAp, k40], tags={I40, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAp, _8db77cAAAAq, k41], tags={I41, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAq, _8db77cAAAAr, k42], tags={I42, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAr, _8db77cAAAAs, k43], tags={I43, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAs, _8db77cAAAAt, k44], tags={I44, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAt, _8db77cAAAAu, k45], tags={I45, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAu, _8db77cAAAAv, k46], tags={I46, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAv, _8db77cAAAAw, k47], tags={I47, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAw, _8db77cAAAAx, k48], tags={I48, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAx, _8db77cAAAAy, k49], tags={I49, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAy, _8db77cAAAAz, k50], tags={I50, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAAAz, _8db77cAAABA, k51], tags={I51, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABA, _8db77cAAABB, k52], tags={I52, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABB, _8db77cAAABC, k53], tags={I53, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABC, _8db77cAAABD, k54], tags={I54, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABD, _8db77cAAABE, k55], tags={I55, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABE, _8db77cAAABF, k56], tags={I56, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABF, _8db77cAAABG, k57], tags={I57, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABG, _8db77cAAABH, k58], tags={I58, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABH, _8db77cAAABI, k59], tags={I59, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABI, _8db77cAAABJ, k60], tags={I60, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABJ, _8db77cAAABK, k61], tags={I61, B}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABK, _8db77cAAABL, k62], tags={I62, A}),backend=numpy, dtype=float32, data=...
Tensor(shape=(16, 16, 2), inds=[_8db77cAAABL, _8db77cAAAAA, k63], tags={I63, B}),backend=numpy, dtype=float32, data=...
" ], "text/plain": [ "MatrixProductState(tensors=64, indices=128, L=64, max_bond=16)" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt.optimize(1000)" ] }, { "cell_type": "code", "execution_count": 15, "id": "3575733a-5ba8-4d52-b635-84c034012191", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "text/plain": [ "(
, )" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt.plot(hlines={'analytic': qu.heisenberg_energy(L)})" ] }, { "cell_type": "markdown", "id": "d22f796f-5933-4b28-a379-d96e8e16048f", "metadata": {}, "source": [ "The reduction in parameters also helps the optimization converge faster.\n", "\n", "As a sanity check, we can explicitly visualize the first four tensors in the optimized MPS, and see that they do indeed repeat:" ] }, { "cell_type": "code", "execution_count": 16, "id": "911fd7d0-b69a-4b52-8932-9fa721ac3891", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" } ], "source": [ "psi_opt = tnopt.get_tn_opt()\n", "psi_opt[:4].visualize_tensors(\"row\")" ] }, { "cell_type": "markdown", "id": "b1ad5a90-1fe3-4f50-b20b-b30d80ae62fb", "metadata": {}, "source": [ "(optimizing-circuits)=\n", "## Optimizing `Circuit` objects\n", "\n", "One special case of optimizing pytrees, is when [`Circuit`](quimb.tensor.circuit.Circuit) objects are encountered. These contain a tensor network representation of the quantum circuit, but only the gates/tensors which have been specified as `parametrized` are optimized." ] }, { "cell_type": "code", "execution_count": 17, "id": "cb6a603e-ba03-4857-8af5-4f5cdd67ca15", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "-0.749999999208 [best: -0.749999999208] : : 13it [00:00, 39.17it/s] \n" ] } ], "source": [ "import autoray as ar\n", "\n", "rng = ar.do('random.default_rng', 42, like=\"numpy\")\n", "\n", "circ = qtn.Circuit(2)\n", "circ.u3(*rng.uniform(size=3, high=2 * qu.pi), 0, parametrize=True)\n", "circ.u3(*rng.uniform(size=3, high=2 * qu.pi), 1, parametrize=True)\n", "circ.cnot(0, 1)\n", "circ.u3(*rng.uniform(size=3, high=2 * qu.pi), 0, parametrize=True)\n", "circ.u3(*rng.uniform(size=3, high=2 * qu.pi), 1, parametrize=True)\n", "\n", "H = qu.ham_heis(2).astype(\"complex128\")\n", "\n", "def loss(circ, H):\n", " en = circ.local_expectation(H, (0, 1), simplify_sequence=\"ADCRS\")\n", " # we use `autoray.do` to allow arbitrary autodiff backends\n", " return ar.do(\"real\", en)\n", "\n", "tnopt = qtn.TNOptimizer(\n", " circ,\n", " loss,\n", " loss_constants=dict(H=H),\n", " # because we are using dynamic (entry dependent) simplification\n", " autodiff_backend=\"autograd\",\n", ")\n", "circ_opt = tnopt.optimize(10)" ] }, { "cell_type": "markdown", "id": "82c25e32-7692-4aeb-9835-c1c18a5f1268", "metadata": {}, "source": [ "The returned circuit now has the optimized parameters set." ] }, { "cell_type": "code", "execution_count": 18, "id": "cd68221f-f40e-4149-be06-eecb788c358f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(,\n", " ,\n", " ,\n", " ,\n", " )" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "circ_opt.gates" ] }, { "cell_type": "code", "execution_count": 19, "id": "509e0fe8-8264-4c3f-8bec-32a01ec0fe52", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "-0.7499999992075422" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "loss(circ_opt, H)" ] }, { "cell_type": "markdown", "id": "6ddfef05-04fa-46e3-85c5-080c6a5f7095", "metadata": {}, "source": [ "But the initial state and constant gates have not been changed, and becuase the parametrized tensors are manifestly unitary, it is always normalized.\n", "\n", "```{note}\n", "Note also, that because we used `simplify_sequence=\"ADCRS\"`, which performs TN simplifications that depend on the tensor entries, we cannot use a statically compiled autodiff backend like `jax`. Instead we use an 'eager' library `autograd` here. Another option would be to instead use `simplify_sequence=\"R\"`.\n", "```\n", "\n", "If you want fine control over which gates to optimize and share parameters among in the circuit, it is best to extract the TN representation first, and optimize it directly. You can then call [`circ.update_params_from(tn)`](quimb.tensor.circuit.Circuit.update_params_from) to update the circuit parameters from an optimized tensor network." ] }, { "cell_type": "markdown", "id": "df7cfd7a-33d6-4e0a-ad05-9a541ed8fbee", "metadata": {}, "source": [ "## Optimizing `PTensor` objects\n", "\n", "The circuit object paramterized gates behind the scenes use a 'paramterized' tensor object, [`PTensor`](quimb.tensor.PTensor), which holds a [`PArray`](quimb.tensor.array_ops.PArray). This is a generalization of `Tensor` whose data is defined by a function and some parameters (kind of like a local `norm_fn` that is always applied). You can use these directly for even finer control.\n", "\n", "Here we show a very roundabout way of trying to diagonalize a non-symmetric matrix using two orthogonal matrices `U` and `V`." ] }, { "cell_type": "code", "execution_count": 20, "id": "bc0fbd42-63d3-4832-8848-c80861ee36db", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "text/plain": [ "(
, )" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# each parametrized tensor has implicit `data = fn(params)`\n", "\n", "Ua = qtn.PTensor(\n", " # project into isometric / unitary / isometric form\n", " fn=qtn.decomp.isometrize_cayley,\n", " # parameters here are some arbitrary array\n", " params=qu.identity(10, dtype=\"float64\"),\n", " inds=[\"w\", \"x\"],\n", " tags=\"Ua\",\n", ")\n", "\n", "G = qtn.Tensor(\n", " data=qu.randn((10, 10)),\n", " inds=['x', 'y'],\n", " tags=\"G\",\n", ")\n", "\n", "Ub = qtn.PTensor(\n", " fn=qtn.decomp.isometrize_cayley,\n", " params=qu.identity(10, dtype=\"float64\"),\n", " inds=[\"y\", \"z\"],\n", " tags=\"Ub\",\n", ")\n", "\n", "(Ua | G | Ub).draw([\"Ua\", \"G\", \"Ub\"])\n", "(Ua | G | Ub).contract().visualize()" ] }, { "cell_type": "code", "execution_count": 21, "id": "33201a4e-4d04-4c9a-975f-02c42ffbc05b", "metadata": {}, "outputs": [], "source": [ "def loss(target, G):\n", " Ua, Ub = target[\"Ua\"], target[\"Ub\"]\n", " UGU = (Ua | G | Ub).contract(output_inds=['w', 'z']).data\n", " # minimize off-diagonal elements of UGU\n", " return ar.do(\"linalg.norm\", UGU - ar.do('diag', ar.do('diag', UGU)))" ] }, { "cell_type": "markdown", "id": "469826ed-a3cc-4f6b-9665-9bc82a665b5f", "metadata": {}, "source": [ "We'll also here make use of automatic hessian-vector product computation by some backends (e.g. `jax`), which is can be used with second order optimization methods like `Newton-CG`." ] }, { "cell_type": "code", "execution_count": 22, "id": "7bfcc1df-0fcf-4f2f-9456-38e0617528ec", "metadata": {}, "outputs": [], "source": [ "tnopt = qtn.TNOptimizer(\n", " {\"Ua\": Ua, \"Ub\": Ub},\n", " loss,\n", " loss_constants=dict(G=G),\n", " optimizer=\"newton-cg\",\n", " autodiff_backend=\"jax\",\n", ")" ] }, { "cell_type": "code", "execution_count": 23, "id": "c639c11f-f59c-467d-b40d-8559d6378ade", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "+0.000001728355 [best: +0.000001728355] : 60%|██████ | 601/1000 [00:01<00:01, 307.78it/s]\n" ] } ], "source": [ "to = tnopt.optimize(1000, hessp=True)" ] }, { "cell_type": "code", "execution_count": 24, "id": "7f732bc6-e98d-42f0-b7f9-1ec208d8824e", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "text/plain": [ "(
, )" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tnopt.plot()" ] }, { "cell_type": "code", "execution_count": 25, "id": "c9569835-803f-4800-9a96-1e11b698041c", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
PTensor(shape=(10, 10), inds=[w, x], tags={Ua}),backend=numpy, dtype=None, data=array([[ 0.32642075, -0.6051291 , 0.2978906 , 0.39854315, 0.02511937,\n", " 0.0838748 , -0.29238603, 0.0739949 , -0.21712632, -0.36594826],\n", " [ 0.0908875 , 0.33829898, 0.42404962, -0.4448276 , 0.18260145,\n", " 0.42017654, 0.08901135, 0.26795745, -0.29496846, -0.35068703],\n", " [-0.19292217, -0.27753505, 0.0356112 , 0.15946661, 0.47263956,\n", " 0.2622383 , 0.54696834, -0.28224143, -0.32863796, 0.282954 ],\n", " [ 0.07516661, 0.08561171, -0.151499 , 0.3938459 , -0.45643058,\n", " 0.49859524, 0.24238792, 0.48964697, -0.01888227, 0.23057617],\n", " [ 0.09385245, -0.0597567 , -0.7957913 , -0.08307468, 0.39757654,\n", " 0.09958041, -0.1210546 , 0.2673096 , -0.14329785, -0.26983204],\n", " [-0.34093434, -0.35484943, -0.00559698, -0.38959286, -0.08984588,\n", " 0.3742761 , -0.5319735 , 0.02843514, -0.0771872 , 0.41001788],\n", " [-0.35343474, -0.46307853, -0.07277285, -0.31267455, -0.3303163 ,\n", " -0.0156951 , 0.4368167 , 0.0749092 , 0.23010063, -0.44593632],\n", " [ 0.18945608, -0.24993688, 0.17366748, -0.21469343, 0.22132577,\n", " -0.4433413 , 0.16982655, 0.63924885, 0.03606365, 0.3755888 ],\n", " [ 0.46199533, -0.12276153, 0.02312034, -0.13623571, 0.23001385,\n", " 0.3858519 , 0.06653047, -0.15374139, 0.71851975, 0.07920738],\n", " [ 0.5856843 , -0.11639551, -0.19472294, -0.37635726, -0.39644542,\n", " -0.05807046, 0.16411862, -0.30522007, -0.40060037, 0.15082498]],\n", " dtype=float32)
" ], "text/plain": [ "PTensor(shape=(10, 10), inds=('w', 'x'), tags=oset(['Ua']))" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "Uao, Ubo = to[\"Ua\"], to[\"Ub\"]\n", "Uao" ] }, { "cell_type": "code", "execution_count": 26, "id": "30945e98-807a-4b7f-89d8-f86be0b8fff1", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(10.000000000000002, 10.000000000000002)" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# still orthogonal\n", "Uao.norm()**2, Ubo.norm()**2" ] }, { "cell_type": "code", "execution_count": 27, "id": "d2494938-0019-4bf9-ba68-a98aaa7d05c9", "metadata": {}, "outputs": [ { "data": { "image/svg+xml": [ "" ], "text/plain": [ "
" ] }, "metadata": { "needs_background": "light" }, "output_type": "display_data" }, { "data": { "text/plain": [ "(
, )" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# check gauged G\n", "(Uao | G | Ubo).contract().visualize()" ] } ], "metadata": { "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" }, "vscode": { "interpreter": { "hash": "6132c5c0a7d26b7c311caf7f55df83b87474b489906668e67e2d71a3b39ab16a" } } }, "nbformat": 4, "nbformat_minor": 2 }