diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 78d3f8839..000000000 --- a/Dockerfile +++ /dev/null @@ -1,21 +0,0 @@ -FROM ubuntu - -# Install basic tools. -RUN DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ - python3-pip python3-tk git emacs vim locales - -# Configure UTF-8 encoding. -RUN sed -i -e 's/# en_US.UTF-8 UTF-8/en_US.UTF-8 UTF-8/' /etc/locale.gen && locale-gen -ENV LANG en_US.UTF-8 -ENV LANGUAGE en_US:en -ENV LC_ALL en_US.UTF-8 - -# Make python3 default -RUN rm -f /usr/bin/python && ln -s /usr/bin/python3 /usr/bin/python - -# Install Tensor Network with the needed Python libraries. -RUN pip3 install tensornetwork - -WORKDIR /TensorNetwork/examples - -EXPOSE 8888 diff --git a/examples/sat/SATTutorial.ipynb b/examples/sat/SATTutorial.ipynb new file mode 100644 index 000000000..2b6dc9bd9 --- /dev/null +++ b/examples/sat/SATTutorial.ipynb @@ -0,0 +1,465 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "SATTutorial.ipynb", + "provenance": [], + "collapsed_sections": [], + "include_colab_link": true + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "view-in-github", + "colab_type": "text" + }, + "source": [ + "\"Open" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "iJCTaVM1JRVw", + "colab_type": "text" + }, + "source": [ + "# SAT Problem with TensorNetwork\n", + "by Volha Okrut\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "9tNIRZfiKn7Z", + "colab_type": "text" + }, + "source": [ + "## Boolean Logic\n", + "\n", + "Suppose we have a simple [CFN expression](https://en.wikipedia.org/wiki/Conjunctive_normal_form), a logical expression based on logical AND (called conjunction) and logical OR (disjunctions). Strictly defined, CFN expression is a conjunction (AND) of several disjunctions (OR) of logical literals (*Xi*).\n", + "\n", + "Let me come up with the following example of CFN expression:\n", + "\n", + "(True AND False) OR (NOT True AND True)\n", + "\n", + "Now let's simplify it:\n", + "\n", + "False OR False\n", + "\n", + "And at the end we get:\n", + "\n", + "False\n", + "\n", + "That's simple!\n", + "\n", + "Now, instead of logical AND I will use ∨ notation, instead of logical OR - ∧. Additionally, if I want to say NOT True I use ¬True. This is just a formality, and yet it allows us to write these expressions in a more clearer and readable form.\n", + "\n", + "Of course, the concept of boolean expressions would be pretty useless if always you had to start with the same positions of True and False. So let me introduce variables (known as literals) into the formula - this allows me to have the same expression evaluating to different end-results.\n", + "Now instead of our initial formula we have:\n", + "\n", + "(X1 ∨ X2) ∧ (¬X1 ∨ X3)\n", + "\n", + "You can notice that if we place *X1* to be True, *X2* to be False, and *X3* to be True, we would have the same expression as in the begining of the article. If we assign the variables different values, then we of course will get different result." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Q8CjHN1bWo05", + "colab_type": "text" + }, + "source": [ + "## SAT Problem and real-life example\n", + "\n", + "So what is a SAT Problem? SAT Problem - short from *SATISFABILITY* problem - concerns with the number of ways in which you can arrange the given literals in order for the whole expression to be evaluated to True. \n", + "\n", + "Let's start by jumping in with an example of a SAT problem. Suppose that you need to go grocery shopping, and need to visit three stores: Costco, Home Depot, and Walmart. Costco is open in the morning and evening, Home Depot is open in the evening only, and Walmart is open in the morning only. You can only be in one place at a time, and shopping at a given store takes up the entire morning or evening. Can you go to all three stores in a day?\n", + "To a human, it is intuitively obvious that the answer is no. Since Home Depot and Walmart offer us only one time option (evening and morning, respectively), then we have to go there at those times. However, this leaves no time for a Costco trip, so it's evident that this \"puzzle\" has no solution.\n", + "Now suppose instead of three stores, you were given three thousand (each with its own schedule), and instead of two times, you were given all the hours of a day? At this point, the problem becomes intractable for a human. Luckily, though, cruching numbers and analyzing thousands of different options are what computers excel at.\n", + "\n", + "So, how could we encode the above problem in a way that a computer could understand?\n", + "\n", + "One solution would be to re-write the problem involving boolean variables, which can either be true or false. For example, using the example of three stores and two times, let's make six variables:\n", + "\n", + "* *Ce*: Whether we go to Costco in the evening.\n", + "* *Cm*: Whether we go to Costco in the morning.\n", + "* *He*: Whether we go to Home Depot in the evening.\n", + "* *Hm*: Whether we go to Home Depot in the morning.\n", + "* *We*: Whether we go to Walmart in the evening.\n", + "* *Wm*: Whether we go to Walmart in the morning.\n", + "\n", + "Each of these variables if true (or 1) if we visit the store at the corresponding time, and false (or 0) otherwise. Next, we form some constraints on these variables, and express them in a unified form we could feed to a computer.\n", + "\n", + "First, we know that we can only be in one place at a given time. For example, if we are at Costco in the morning (that is, Cm=1\n", + "), then we cannot be at Home Depot or Walmart in the morning (and thus Hm=Wm=0). Using notation introduced above we can express that constrains as:\n", + "\n", + "*Cm ∨ ¬Hm ∨ ¬Wm*\n", + "\n", + "Of course, we know that at a given time, we could go to Costco, Home Depot, or Walmart, so Cm\n", + "doesn't have to be true. Thus, the constraint that we only go to one place in the evening can be represented as:\n", + "\n", + "*( Ce ∧ ¬He ∧ ¬We ) ∨ ( ¬Ce ∧ He ∧ ¬We ) ∨ ( ¬Ce ∧ ¬He ∧ We )*\n", + "\n", + "Similarly, the constraint that we only go to one place in the morning is:\n", + "\n", + "*( Cm ∧ ¬Hm ∧ ¬Wm ) ∨ ( ¬Cm ∧ Hm ∧ ¬Wm ) ∨ ( ¬Cm ∧ ¬Hm ∧ Wm )*\n", + "\n", + "Next, we need a constraint that we go to Costco in either the morning or evening, which we can represent as *Cm ∨ Ce*: either we go to Costco in the morning, or in the evening. We have similar constraints for Walmart and Home Depot, yielding the following constraint to represent that we must go to each store:\n", + "\n", + "*( Cm ∨ Ce) ∧ ( Hm ∨ He ) ∧ ( Wm ∨ We )*\n", + "\n", + "Thus, the full set of constraints for our problem is\n", + "\n", + "( Cm ∨ Ce) ∧ ( Hm ∨ He ) ∧ ( Wm ∨ We ) ∧ ( Cm ∧ ¬Hm ∧ ¬Wm ) ∨ ( ¬Cm ∧ Hm ∧ ¬Wm ) ∨ ( ¬Cm ∧ ¬Hm ∧ Wm ) ∧ ( Ce ∧ ¬He ∧ ¬We ) ∨ ( ¬Ce ∧ He ∧ ¬We ) ∨ ( ¬Ce ∧ ¬He ∧ We )\n", + "\n", + "\n", + "To find out whether we can complete our shopping trip, we must find a set of true or false values for all our boolean variables such that the constraints are satisfied. This type of problem is known as the boolean satisfiability problem, often abbreviated to just \"SAT\". A program that finds solutions to these problems is known as a SAT solver.\n", + "\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "4Itv5A30Stxv", + "colab_type": "text" + }, + "source": [ + "## SAT Problem\n", + "\n", + "SAT problem has been viewed from many different ways, in this tutorial we will learn how to solve this problem using tensors and TensorNetwork library. To be comfartable with tensors you have to know some basics about *Penrose’s Graphical Notation*. Check this nice article on [Medium](https://medium.com/analytics-vidhya/penroses-graphical-notation-fe4c2f24cf3b) that covers this topic extensively.\n", + "\n", + "Suppose we are given four variables: X1, X2, X3, X4. We want to find truth values to all four Xi literals so that the CNF expression is true:\n", + "\n", + "( ¬X1 ∨ ¬X3 ∨ ¬X4 ) ∧ ( X2 ∨ X3 ∨ ¬X4 ) ∧ ( X1 ∨ ¬X2 ∨ X4 ) ∧ ( X1 ∨ X3 ∨ X4 ) ∧ ( ¬X1 ∨ X2 ∨ ¬X3 )\n", + "\n", + "First, we need to define how we encode our input CNF expressions that we want to satisfy:\n", + "\n", + "* Each logical literal is represented as either a positive or negative integer, where i and -i correpond to the logical literals xi and ¬xi, respectively.\n", + "* Each clause in the expression, i.e., disjunction of literals, is represented as a tuple of such encoding of literals, e.g., (-1, 2, -3) represents the disjunction ( ¬x1 ∨ x2 ∨ ¬x3 ).\n", + "* The entire conjunctive expression is a list of such tuples, e.g., the expression above would have encoding:\n", + "[(-1, -3, -4), (2, 3, -4), (1, -2, 4), (1, 3, 4), (-1, 2, -3)]\n", + "\n", + "It is worth to say that we can solve two problems here:\n", + "\n", + "\n", + "1. Find the exact number of all possible solutions to the given SAT problem if these solutions exist.\n", + "2. Find all possible solutions to a given SAT problem if these solutions exist.\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "choohffKZatl", + "colab_type": "text" + }, + "source": [ + "## SAT solver using tensors and TensorNetwork\n", + "\n", + "### Finding all possible solutions to the given SAT \n", + "\n", + "First, we create a function\n", + "```\n", + "# sat_tn(clauses)\n", + "```\n", + "which solves the given 3SAT problem. \n", + "\n", + "We find the maximum indexed logical variable we have, and use that as our count of the number of logical variables. We iterate through each disjunction and calculate absolute value for each variable in the conjunction. The number of logical variables is the maximum element in *var_set* which is then stored in *num_vars*:\n", + "```\n", + "var_set = set()\n", + "for clause in clauses:\n", + " var_set |= {abs(x) for x in clause}\n", + "num_vars = max(var_set)\n", + "```\n", + "After iterating expression [(-1, -3, -4), (2, 3, -4), (1, -2, 4), (1, 3, 4), (-1, 2, -3)], \n", + "I should get the following result:\n", + "```\n", + "var_set final = {1, 3, 4, 2}\n", + "num_vars = 4\n", + "```\n", + "\n", + "Now, we will build the tensor network. Variable nodes (literals) will be represented as *num_vars* tensors with the shape(1,2) filled with ones:\n", + "```\n", + "node [1 1]\n", + "```\n", + "\n", + "This particular shape of nodes is needed for matrix multiplication. I will explain why we need it in just a moment. Since each of the variable nodes is a vector, each of them will have only one edge, which I will store as unconnected edges (dangling edges) in *var_edges*:\n", + "```\n", + "var_edges.append(new_node[0])\n", + "```\n", + "The second step is to create nodes for all clauses. For each clause we will create a tensor of third rank (a 3D matrix) with two fields in each dimenshion as we want as many fields as there are possible solutions (variation of initial literals) to this clause. Each logical variable *Xi* has two posible literals: itself (*Xi*), and its negation (*¬Xi*). Thus for each clause we have *2^3 = 8* solutions and each soltution can be accessed using the coditions of the variables (solution to clause (X1, X2, X3) with X1 = 1, X2 = 0 and X3 = 1 will be found under clause_tensor[1, 0, 1] field and will be 1 (True)).\n", + "The formula (-np.sign(x) + 1) // 2 gives us 0 or 1 depending on the sign of the variable (its negation).\n", + "\n", + "```\n", + "for clause in clauses:\n", + " a, b, c, = clause\n", + " clause_tensor = np.ones((2, 2, 2), dtype=np.int32)\n", + " clause_tensor[(-np.sign(a) + 1) // 2, (-np.sign(b) + 1) // 2,\n", + " (-np.sign(c) + 1) // 2] = 0\n", + " clause_node = tn.Node(clause_tensor)\n", + "```\n", + "\n", + "Now, with everything prepare, I can explain you why tensors are such an elegant solution to this problen. As initially we have several expressions that contain only OR operators unified under AND operator, it might be useful to view those operators as logical summation and multiplication respectively. In other words, for logical operator OR, it doesn't matter how many Falses you have - it takes only one True to bring the expression to True (same as summation). On contrary, while you might have all but one Trues in your expression, operator AND will evaluate it to False if there was at least one constituent set to False (same as multiplication).\n", + "\n", + "The same idea will be applied to tensores in the problem. We have constructed 3D matrices to clauses in such a way that they are filled in with 1 for all possible entries except one (think about it: when the clause consists only of logical OR (summation) it is false only with all of its constituents being evaluated to False). Now, if we are able to correctly multiply all the matricies with each other, all of the configuration that have at least one 0 in it will end up being 0 and the only ones left have all of the clauses being True - exactly what we need! The initial vectoes for literals had to be set in such a way in order for matrices to be reduced after being multiplied with them.\n", + "\n", + "Bear with me, the last step is to connect variable to the clause. Operator (^) is used as a shortcut for tn.connect(clause_node, tensor_node) - function for dot product between matrices introduced in Tensor Network library. The result is stored into the first variable.\n", + "\n", + "For now we just have our clause nodes and literal vectores in place. In order to connect them, for every edge of each clause (they all have three edges - by the numbers of literals in the clause) we will create a copy vectore with the same dimension (3D matrix, extending to two fields to each side). The zero numbered edge of the *copy_tensor_node* we will connect to the one of the edges of the clause matrix. The one numbered edge will be connected to one of the *var_edges* - so called dangling edges - edges that are not yet connected to any other edge. Finally, the last edge (numbered with two) will take the place of the edge from *var_edges* that was just paired up with the one numbered edge.\n", + "\n", + "```\n", + "for i, var in enumerate(clause):\n", + " copy_tensor_node = tn.CopyNode(3, 2)\n", + " clause_node[i] ^ copy_tensor_node[0]\n", + " var_edges[abs(var) - 1] ^ copy_tensor_node[1]\n", + " var_edges[abs(var) - 1] = copy_tensor_node[2]\n", + "```\n", + "\n", + "This process will be repeated until all the edges coming from clause tensores have been paired up with a copy tensor node. By the end of this, in *var_edges* we will have stored all the unconnected edges of this system of multiplied tensors. And that is exactly what will be returened from the function:\n", + "\n", + "```\n", + "return var_edges\n", + "```\n", + "\n", + "\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "0ZWRMj1EkT5k", + "colab_type": "text" + }, + "source": [ + "Let gather all the said above in to one program and run with the given set of variables." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "OdbqdgkB3wXZ", + "colab_type": "code", + "outputId": "813405c4-21da-4781-e7e9-888a0b8fa6c5", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 119 + } + }, + "source": [ + "!pip3 install tensornetwork" + ], + "execution_count": 0, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Requirement already satisfied: tensornetwork in /usr/local/lib/python3.6/dist-packages (0.2.0)\n", + "Requirement already satisfied: numpy>=1.16 in /usr/local/lib/python3.6/dist-packages (from tensornetwork) (1.17.5)\n", + "Requirement already satisfied: graphviz>=0.11.1 in /usr/local/lib/python3.6/dist-packages (from tensornetwork) (0.13.2)\n", + "Requirement already satisfied: opt-einsum>=2.3.0 in /usr/local/lib/python3.6/dist-packages (from tensornetwork) (3.1.0)\n", + "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.6/dist-packages (from tensornetwork) (2.10.0)\n", + "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from h5py>=2.9.0->tensornetwork) (1.12.0)\n" + ], + "name": "stdout" + } + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "qbUls8WGQM6x", + "colab_type": "code", + "colab": {} + }, + "source": [ + "import numpy as np\n", + "from typing import List, Tuple, Set\n", + "import tensornetwork as tn" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "L8IARMaFXwH0", + "colab_type": "code", + "colab": {} + }, + "source": [ + "def sat_tn(clauses: List[Tuple[int, int, int]]\n", + " ) -> List[tn.Edge]:\n", + " \"\"\"Create a 3SAT TensorNetwork of the given 3SAT clauses.\n", + " After full contraction, this network will be a tensor of size (2, 2, ..., 2)\n", + " with the rank being the same as the number of variables. Each element of the\n", + " final tensor represents whether the given assignment satisfies all of the\n", + " clauses. For example, if final_node.get_tensor()[0][1][1] == 1, then the\n", + " assiment (False, True, True) satisfies all clauses.\n", + " Args:\n", + " clauses: A list of 3 int tuples. Each element in the tuple corresponds to a\n", + " variable in the clause. If that int is negative, that variable is negated\n", + " in the clause.\n", + " Returns:\n", + " net: The 3SAT TensorNetwork.\n", + " var_edges: The edges for the given variables.\n", + " Raises:\n", + " ValueError: If any of the clauses have a 0 in them.\n", + " \"\"\"\n", + " for clause in clauses:\n", + " if 0 in clause:\n", + " raise ValueError(\"0's are not allowed in the clauses.\")\n", + " var_set = set()\n", + " for clause in clauses:\n", + " var_set |= {abs(x) for x in clause}\n", + " num_vars = max(var_set)\n", + " var_nodes = []\n", + " var_edges = []\n", + "\n", + " # Prepare the variable nodes.\n", + " for _ in range(num_vars):\n", + " new_node = tn.Node(np.ones(2, dtype=np.int32))\n", + " var_nodes.append(new_node)\n", + " var_edges.append(new_node[0])\n", + "\n", + " # Create the nodes for each clause\n", + " for clause in clauses:\n", + " a, b, c, = clause\n", + " clause_tensor = np.ones((2, 2, 2), dtype=np.int32)\n", + " clause_tensor[(-np.sign(a) + 1) // 2, (-np.sign(b) + 1) // 2,\n", + " (-np.sign(c) + 1) // 2] = 0\n", + " clause_node = tn.Node(clause_tensor)\n", + "\n", + " # Connect the variable to the clause through a copy tensor.\n", + " for i, var in enumerate(clause):\n", + " copy_tensor_node = tn.CopyNode(3, 2)\n", + " clause_node[i] ^ copy_tensor_node[0]\n", + " var_edges[abs(var) - 1] ^ copy_tensor_node[1]\n", + " var_edges[abs(var) - 1] = copy_tensor_node[2]\n", + "\n", + " return var_edges" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "AXpBEmktMfw6", + "colab_type": "text" + }, + "source": [ + "### Find the exact number of all possible solutions to the given SAT\n", + "\n", + "In order to find exact number of all possible solutions to the given SAT problem, we can do full contractions of the adges of the clauses. In other words, we have to calculate a trace of the tensor network we have build in the first part of the tutorial. \n", + "This is done by essentially creating the same tensor net and then connecting all the dangling edges of the first net to the dangling edges of the second." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "v6NvHpWd1AdI", + "colab_type": "code", + "colab": {} + }, + "source": [ + "def sat_count_tn(clauses: List[Tuple[int, int, int]]):\n", + " \"\"\"Create a 3SAT Count TensorNetwork.\n", + " After full contraction, the final node will be the count of all possible\n", + " solutions to the given 3SAT problem.\n", + " Args:\n", + " clauses: A list of 3 int tuples. Each element in the tuple corresponds to a\n", + " variable in the clause. If that int is negative, that variable is negated\n", + " in the clause.\n", + " Returns:\n", + " nodes: The set of nodes\n", + " \"\"\"\n", + " var_edges1 = sat_tn(clauses)\n", + " var_edges2 = sat_tn(clauses)\n", + " for edge1, edge2 in zip(var_edges1, var_edges2):\n", + " edge1 ^ edge2\n", + " return tn.reachable(var_edges1[0].node1)" + ], + "execution_count": 0, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W434kXjqTW1j", + "colab_type": "text" + }, + "source": [ + "Congratulations! You have now learned how to write SAT Solver program with TensorNetwork! Down below you can play with choosing different clauses as your starting points and then seeing with how many ways it can be solved." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "kjxLkM_uOkfV", + "colab_type": "code", + "outputId": "477cf629-d4d6-4831-81a8-f73fea95f03c", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 34 + } + }, + "source": [ + "import numpy as np\n", + "from typing import List, Tuple, Set\n", + "import tensornetwork as tn\n", + "\n", + "my_clause = [(-1, -3, -4), (2, 3, -4), (1, -2, 4), (1, 3, 4), (-1, 2, -3)]\n", + "nodes = sat_count_tn(my_clause)\n", + "count = tn.contractors.greedy(nodes).tensor\n", + "print(\"Number of solutions = \", count)" + ], + "execution_count": 0, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Number of solutions = 7.0\n" + ], + "name": "stdout" + } + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d0OaWmpdJDze", + "colab_type": "text" + }, + "source": [ + "# References:\n", + "\n", + "\n", + "1. [SAT solver example ](http://www.tfinley.net/software/pyglpk/ex_sat.html)\n", + "2. [An exact tensor network for the 3SAT problem](https://arxiv.org/abs/1105.3201) \n", + "1. [Penrose’s Graphical Notation](https://medium.com/analytics-vidhya/penroses-graphical-notation-fe4c2f24cf3b)\n", + "1. https://github.com/google/TensorNetwork\n", + "1. [Writing a SAT Solver](http://andrew.gibiansky.com/blog/verification/writing-a-sat-solver/)\n", + "2. \n", + "\n", + "\n", + "\n", + "\n" + ] + } + ] +} \ No newline at end of file diff --git a/requirements_travis.txt b/requirements_travis.txt index 2b281ebb7..7d226c1a3 100644 --- a/requirements_travis.txt +++ b/requirements_travis.txt @@ -1,6 +1,6 @@ tensorflow>=2.0.0 pytype==2019.06.21 pytest -torch>=1.1 +torch==1.3.1 jax>=0.1.0 jaxlib>=0.1.27 diff --git a/tensornetwork/__init__.py b/tensornetwork/__init__.py index fd6269031..96a490fe4 100644 --- a/tensornetwork/__init__.py +++ b/tensornetwork/__init__.py @@ -10,12 +10,8 @@ from tensornetwork.version import __version__ from tensornetwork.visualization.graphviz import to_graphviz from tensornetwork import contractors -from tensornetwork import config from typing import Text, Optional, Type, Union from tensornetwork.utils import load_nodes, save_nodes from tensornetwork.matrixproductstates.finite_mps import FiniteMPS from tensornetwork.matrixproductstates.infinite_mps import InfiniteMPS - - -def set_default_backend(backend: Union[Text, BaseBackend]) -> None: - config.default_backend = backend +from tensornetwork.backend_contextmanager import DefaultBackend, set_default_backend diff --git a/tensornetwork/backend_contextmanager.py b/tensornetwork/backend_contextmanager.py new file mode 100644 index 000000000..814d6d7bf --- /dev/null +++ b/tensornetwork/backend_contextmanager.py @@ -0,0 +1,41 @@ +from typing import Text, Union +from tensornetwork.backends.base_backend import BaseBackend + +class DefaultBackend(): + """Context manager for setting up backend for nodes""" + + def __init__(self, backend: Union[Text, BaseBackend]) -> None: + if not isinstance(backend, (Text, BaseBackend)): + raise ValueError("Item passed to DefaultBackend " + "must be Text or BaseBackend") + self.backend = backend + + def __enter__(self): + _default_backend_stack.stack.append(self) + + def __exit__(self, exc_type, exc_val, exc_tb): + _default_backend_stack.stack.pop() + +class _DefaultBackendStack(): + """A stack to keep track default backends context manager""" + + def __init__(self): + self.stack = [] + self.default_backend = "numpy" + + def get_current_backend(self): + return self.stack[-1].backend if self.stack else self.default_backend + +_default_backend_stack = _DefaultBackendStack() + +def get_default_backend(): + return _default_backend_stack.get_current_backend() + +def set_default_backend(backend: Union[Text, BaseBackend]) -> None: + if _default_backend_stack.stack: + raise AssertionError("The default backend should not be changed " + "inside the backend context manager") + if not isinstance(backend, (Text, BaseBackend)): + raise ValueError("Item passed to set_default_backend " + "must be Text or BaseBackend") + _default_backend_stack.default_backend = backend diff --git a/tensornetwork/backends/backend_factory.py b/tensornetwork/backends/backend_factory.py index 52d0bfb2e..859d829a4 100644 --- a/tensornetwork/backends/backend_factory.py +++ b/tensornetwork/backends/backend_factory.py @@ -19,7 +19,6 @@ from tensornetwork.backends.shell import shell_backend from tensornetwork.backends.pytorch import pytorch_backend from tensornetwork.backends import base_backend -import tensornetwork.config as config_file _BACKENDS = { "tensorflow": tensorflow_backend.TensorFlowBackend, diff --git a/tensornetwork/backends/backend_test.py b/tensornetwork/backends/backend_test.py index e1a96071a..3e30dcffe 100644 --- a/tensornetwork/backends/backend_test.py +++ b/tensornetwork/backends/backend_test.py @@ -4,6 +4,7 @@ import pytest import numpy as np from tensornetwork import connect, contract, Node +from tensornetwork.backends.base_backend import BaseBackend def clean_tensornetwork_modules(): @@ -146,3 +147,200 @@ def test_basic_network_without_backends_raises_error(): Node(np.ones((2, 2)), backend="tensorflow") with pytest.raises(ImportError): Node(np.ones((2, 2)), backend="pytorch") +[] + +def test_base_backend_name(): + backend = BaseBackend() + assert backend.name == "base backend" + + +def test_base_backend_tensordot_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.tensordot(np.ones((2, 2)), np.ones((2, 2)), axes=[[0], [0]]) + + +def test_base_backend_reshape_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.reshape(np.ones((2, 2)), (4, 1)) + + +def test_base_backend_transpose_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.transpose(np.ones((2, 2)), [0, 1]) + + +def test_base_backend_svd_decompositon_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.svd_decomposition(np.ones((2, 2)), 0) + + +def test_base_backend_qr_decompositon_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.qr_decomposition(np.ones((2, 2)), 0) + + +def test_base_backend_rq_decompositon_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.rq_decomposition(np.ones((2, 2)), 0) + + +def test_base_backend_shape_concat_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.shape_concat([np.ones((2, 2)), np.ones((2, 2))], 0) + + +def test_base_backend_shape_tensor_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.shape_tensor(np.ones((2, 2))) + + +def test_base_backend_shape_tuple_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.shape_tuple(np.ones((2, 2))) + + +def test_base_backend_shape_prod_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.shape_prod(np.ones((2, 2))) + + +def test_base_backend_sqrt_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.sqrt(np.ones((2, 2))) + + +def test_base_backend_diag_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.diag(np.ones((2, 2))) + + +def test_base_backend_convert_to_tensor_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.convert_to_tensor(np.ones((2, 2))) + + +def test_base_backend_trace_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.trace(np.ones((2, 2))) + + +def test_base_backend_outer_product_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.outer_product(np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_einsul_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.einsum("ii", np.ones((2, 2))) + + +def test_base_backend_norm_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.norm(np.ones((2, 2))) + + +def test_base_backend_eye_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.eye(2, dtype=np.float64) + + +def test_base_backend_ones_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.ones((2, 2), dtype=np.float64) + + +def test_base_backend_zeros_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.zeros((2, 2), dtype=np.float64) + + +def test_base_backend_randn_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.randn((2, 2)) + + +def test_base_backend_random_uniforl_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.random_uniform((2, 2)) + + +def test_base_backend_conj_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.conj(np.ones((2, 2))) + + +def test_base_backend_eigh_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.eigh(np.ones((2, 2))) + + +def test_base_backend_eigs_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.eigs(np.ones((2, 2))) + + +def test_base_backend_eigs_lanczos_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.eigsh_lanczos(np.ones((2, 2))) + + +def test_base_backend_addition_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.addition(np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_subtraction_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.subtraction(np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_multiply_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.multiply(np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_divide_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.divide(np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_index_update_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.index_update(np.ones((2, 2)), np.ones((2, 2)), np.ones((2, 2))) + + +def test_base_backend_inv_not_implemented(): + backend = BaseBackend() + with pytest.raises(NotImplementedError): + backend.inv(np.ones((2, 2))) diff --git a/tensornetwork/backends/base_backend.py b/tensornetwork/backends/base_backend.py index 55b9f40ea..9e718998f 100644 --- a/tensornetwork/backends/base_backend.py +++ b/tensornetwork/backends/base_backend.py @@ -135,12 +135,12 @@ def rq_decomposition( raise NotImplementedError( "Backend '{}' has not implemented rq_decomposition.".format(self.name)) - def concat(self, values: Sequence[Tensor], axis) -> Tensor: + def shape_concat(self, values: Sequence[Tensor], axis) -> Tensor: """Concatenate a sequence of tensors together about the given axis.""" raise NotImplementedError("Backend '{}' has not implemented concat.".format( self.name)) - def shape(self, tensor: Tensor) -> Tensor: + def shape_tensor(self, tensor: Tensor) -> Tensor: """Get the shape of a tensor. Args: @@ -163,7 +163,7 @@ def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]: raise NotImplementedError( "Backend '{}' has not implemented shape_tuple.".format(self.name)) - def prod(self, values: Tensor) -> Tensor: + def shape_prod(self, values: Tensor) -> Tensor: """Take the product of all of the elements in values""" raise NotImplementedError("Backend '{}' has not implemented prod.".format( self.name)) @@ -265,6 +265,27 @@ def randn(self, raise NotImplementedError("Backend '{}' has not implemented randn.".format( self.name)) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[Type[np.number]] = None, + seed: Optional[int] = None) -> Tensor: + """Return a random uniform matrix of dimension `dim`. + Depending on specific backends, `dim` has to be either an int + (numpy, torch, tensorflow) or a `ShapeType` object + (for block-sparse backends). Block-sparse + behavior is currently not supported + Args: + shape (int): The dimension of the returned matrix. + boundaries (tuple): The boundaries of the uniform distribution. + dtype: The dtype of the returned matrix. + seed: The seed for the random number generator + Returns: + Tensor : random uniform initialized tensor. + """ + raise NotImplementedError(("Backend '{}' has not implemented " + "random_uniform.").format(self.name)) + def conj(self, tensor: Tensor) -> Tensor: """ Return the complex conjugate of `tensor` @@ -370,6 +391,32 @@ def eigsh_lanczos(self, raise NotImplementedError( "Backend '{}' has not implemented eighs_lanczos.".format(self.name)) + def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + """ + Return the default multiplication of `tensor`. + A backend can override such implementation. + Args: + tensor1: A tensor. + tensor2: A tensor. + Returns: + Tensor + """ + raise NotImplementedError( + "Backend '{}' has not implemented addition.".format(self.name)) + + def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + """ + Return the default multiplication of `tensor`. + A backend can override such implementation. + Args: + tensor1: A tensor. + tensor2: A tensor. + Returns: + Tensor + """ + raise NotImplementedError( + "Backend '{}' has not implemented subtraction.".format(self.name)) + def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: """ Return the default multiplication of `tensor`. @@ -383,6 +430,19 @@ def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: raise NotImplementedError( "Backend '{}' has not implemented multiply.".format(self.name)) + def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + """ + Return the default divide of `tensor`. + A backend can override such implementation. + Args: + tensor1: A tensor. + tensor2: A tensor. + Returns: + Tensor + """ + raise NotImplementedError( + "Backend '{}' has not implemented divide.".format(self.name)) + def index_update(self, tensor: Tensor, mask: Tensor, assignee: Tensor) -> Tensor: """ diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py index f9235391a..a3064912f 100644 --- a/tensornetwork/backends/jax/jax_backend.py +++ b/tensornetwork/backends/jax/jax_backend.py @@ -39,7 +39,7 @@ def convert_to_tensor(self, tensor: Tensor) -> Tensor: result = self.jax.jit(lambda x: x)(tensor) return result - def concat(self, values: Tensor, axis: int) -> Tensor: + def shape_concat(self, values: Tensor, axis: int) -> Tensor: return np.concatenate(values, axis) def randn(self, @@ -72,6 +72,42 @@ def cmplx_randn(complex_dtype, real_dtype): return self.jax.random.normal(key, shape).astype(dtype) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[np.dtype] = None, + seed: Optional[int] = None) -> Tensor: + if not seed: + seed = np.random.randint(0, 2**63) + key = self.jax.random.PRNGKey(seed) + + dtype = dtype if dtype is not None else np.dtype(np.float64) + + def cmplx_random_uniform(complex_dtype, real_dtype): + real_dtype = np.dtype(real_dtype) + complex_dtype = np.dtype(complex_dtype) + + key_2 = self.jax.random.PRNGKey(seed + 1) + + real_part = self.jax.random.uniform(key, shape, dtype=real_dtype, + minval=boundaries[0], + maxval=boundaries[1]) + complex_part = self.jax.random.uniform(key_2, shape, dtype=real_dtype, + minval=boundaries[0], + maxval=boundaries[1]) + unit = ( + np.complex64(1j) + if complex_dtype == np.dtype(np.complex64) else np.complex128(1j)) + return real_part + unit * complex_part + + if np.dtype(dtype) is np.dtype(self.np.complex128): + return cmplx_random_uniform(dtype, self.np.float64) + if np.dtype(dtype) is np.dtype(self.np.complex64): + return cmplx_random_uniform(dtype, self.np.float32) + + return self.jax.random.uniform(key, shape, minval=boundaries[0], + maxval=boundaries[1]).astype(dtype) + def eigs(self, A: Callable, initial_state: Optional[Tensor] = None, diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py index fbed4eebe..4c2868734 100644 --- a/tensornetwork/backends/jax/jax_backend_test.py +++ b/tensornetwork/backends/jax/jax_backend_test.py @@ -35,20 +35,20 @@ def test_transpose(): np.testing.assert_allclose(expected, actual) -def test_concat(): +def test_shape_concat(): backend = jax_backend.JaxBackend() a = backend.convert_to_tensor(2 * np.ones((1, 3, 1))) b = backend.convert_to_tensor(np.ones((1, 2, 1))) - expected = backend.concat((a, b), axis=1) + expected = backend.shape_concat((a, b), axis=1) actual = np.array([[[2.0], [2.0], [2.0], [1.0], [1.0]]]) np.testing.assert_allclose(expected, actual) -def test_shape(): +def test_shape_tensor(): backend = jax_backend.JaxBackend() a = backend.convert_to_tensor(np.ones([2, 3, 4])) - assert isinstance(backend.shape(a), tuple) - actual = backend.shape(a) + assert isinstance(backend.shape_tensor(a), tuple) + actual = backend.shape_tensor(a) expected = np.array([2, 3, 4]) np.testing.assert_allclose(expected, actual) @@ -60,10 +60,10 @@ def test_shape_tuple(): assert actual == (2, 3, 4) -def test_prod(): +def test_shape_prod(): backend = jax_backend.JaxBackend() a = backend.convert_to_tensor(2 * np.ones([1, 2, 3, 4])) - actual = np.array(backend.prod(a)) + actual = np.array(backend.shape_prod(a)) assert actual == 2**24 @@ -161,6 +161,13 @@ def test_randn(dtype): assert a.shape == (4, 4) +@pytest.mark.parametrize("dtype", np_randn_dtypes) +def test_random_uniform(dtype): + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), dtype=dtype) + assert a.shape == (4, 4) + + @pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) def test_randn_non_zero_imag(dtype): backend = jax_backend.JaxBackend() @@ -168,6 +175,13 @@ def test_randn_non_zero_imag(dtype): assert np.linalg.norm(np.imag(a)) != 0.0 +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_random_uniform_non_zero_imag(dtype): + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), dtype=dtype) + assert np.linalg.norm(np.imag(a)) != 0.0 + + @pytest.mark.parametrize("dtype", np_dtypes) def test_eye_dtype(dtype): backend = jax_backend.JaxBackend() @@ -196,6 +210,13 @@ def test_randn_dtype(dtype): assert a.dtype == dtype +@pytest.mark.parametrize("dtype", np_randn_dtypes) +def test_random_uniform_dtype(dtype): + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), dtype=dtype) + assert a.dtype == dtype + + @pytest.mark.parametrize("dtype", np_randn_dtypes) def test_randn_seed(dtype): backend = jax_backend.JaxBackend() @@ -204,6 +225,34 @@ def test_randn_seed(dtype): np.testing.assert_allclose(a, b) +@pytest.mark.parametrize("dtype", np_randn_dtypes) +def test_random_uniform_seed(dtype): + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), seed=10, dtype=dtype) + np.testing.assert_allclose(a, b) + + +@pytest.mark.parametrize("dtype", np_randn_dtypes) +def test_random_uniform_boundaries(dtype): + lb = 1.2 + ub = 4.8 + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype) + assert((a >= 0).all() and (a <= 1).all() and + (b >= lb).all() and (b <= ub).all()) + + +def test_random_uniform_behavior(): + seed = 10 + key = jax.random.PRNGKey(seed) + backend = jax_backend.JaxBackend() + a = backend.random_uniform((4, 4), seed=seed) + b = jax.random.uniform(key, (4, 4)) + np.testing.assert_allclose(a, b) + + def test_conj(): backend = jax_backend.JaxBackend() real = np.random.rand(2, 2, 2) @@ -222,3 +271,27 @@ def index_update(dtype): tensor = np.array(tensor) tensor[tensor > 0.1] = 0.0 np.testing.assert_allclose(tensor, out) + + +def test_base_backend_eigs_not_implemented(): + backend = jax_backend.JaxBackend() + tensor = backend.randn((4, 2, 3), dtype=np.float64) + with pytest.raises(NotImplementedError): + backend.eigs(tensor) + + +def test_base_backend_eigsh_lanczos_not_implemented(): + backend = jax_backend.JaxBackend() + tensor = backend.randn((4, 2, 3), dtype=np.float64) + with pytest.raises(NotImplementedError): + backend.eigsh_lanczos(tensor) + + +@pytest.mark.parametrize("dtype", np_dtypes) +def test_index_update(dtype): + backend = jax_backend.JaxBackend() + tensor = backend.randn((4, 2, 3), dtype=dtype, seed=10) + out = backend.index_update(tensor, tensor > 0.1, 0.0) + np_tensor = np.array(tensor) + np_tensor[np_tensor > 0.1] = 0.0 + np.testing.assert_allclose(out, np_tensor) diff --git a/tensornetwork/backends/numpy/decompositions.py b/tensornetwork/backends/numpy/decompositions.py index 862689f61..0e72f2ac5 100644 --- a/tensornetwork/backends/numpy/decompositions.py +++ b/tensornetwork/backends/numpy/decompositions.py @@ -33,7 +33,7 @@ def svd_decomposition( right_dims = tensor.shape[split_axis:] tensor = np.reshape(tensor, [numpy.prod(left_dims), numpy.prod(right_dims)]) - u, s, vh = np.linalg.svd(tensor) + u, s, vh = np.linalg.svd(tensor, full_matrices=False) if max_singular_values is None: max_singular_values = np.size(s) diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py index c16517391..cea0b712e 100644 --- a/tensornetwork/backends/numpy/decompositions_test.py +++ b/tensornetwork/backends/numpy/decompositions_test.py @@ -63,6 +63,18 @@ def test_max_singular_values(self): self.assertEqual(vh.shape, (7, 10)) self.assertAllClose(trun, np.arange(2, -1, -1)) + def test_max_singular_values_larger_than_bond_dimension(self): + random_matrix = np.random.rand(10, 6) + unitary1, _, unitary2 = np.linalg.svd(random_matrix, full_matrices=False) + singular_values = np.array(range(6)) + val = unitary1.dot(np.diag(singular_values).dot(unitary2.T)) + u, s, vh, _ = decompositions.svd_decomposition( + np, val, 1, max_singular_values=30) + self.assertEqual(u.shape, (10, 6)) + self.assertEqual(s.shape, (6,)) + self.assertEqual(vh.shape, (6, 6)) + + def test_max_truncation_error(self): random_matrix = np.random.rand(10, 10) unitary1, _, unitary2 = np.linalg.svd(random_matrix) diff --git a/tensornetwork/backends/numpy/numpy_backend.py b/tensornetwork/backends/numpy/numpy_backend.py index 0246d32eb..d6a698b67 100644 --- a/tensornetwork/backends/numpy/numpy_backend.py +++ b/tensornetwork/backends/numpy/numpy_backend.py @@ -60,16 +60,16 @@ def rq_decomposition( ) -> Tuple[Tensor, Tensor]: return decompositions.rq_decomposition(self.np, tensor, split_axis) - def concat(self, values: Tensor, axis: int) -> Tensor: + def shape_concat(self, values: Tensor, axis: int) -> Tensor: return self.np.concatenate(values, axis) - def shape(self, tensor: Tensor) -> Tensor: + def shape_tensor(self, tensor: Tensor) -> Tensor: return tensor.shape def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]: return tensor.shape - def prod(self, values: Tensor) -> Tensor: + def shape_prod(self, values: Tensor) -> Tensor: return self.np.prod(values) def sqrt(self, tensor: Tensor) -> Tensor: @@ -131,6 +131,24 @@ def randn(self, dtype) + 1j * self.np.random.randn(*shape).astype(dtype) return self.np.random.randn(*shape).astype(dtype) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[numpy.dtype] = None, + seed: Optional[int] = None) -> Tensor: + + if seed: + self.np.random.seed(seed) + dtype = dtype if dtype is not None else self.np.float64 + if ((self.np.dtype(dtype) is self.np.dtype(self.np.complex128)) or + (self.np.dtype(dtype) is self.np.dtype(self.np.complex64))): + return self.np.random.uniform(boundaries[0], boundaries[1], shape).astype( + dtype) + 1j * self.np.random.uniform(boundaries[0], + boundaries[1], + shape).astype(dtype) + return self.np.random.uniform(boundaries[0], + boundaries[1], shape).astype(dtype) + def conj(self, tensor: Tensor) -> Tensor: return self.np.conj(tensor) @@ -352,9 +370,18 @@ def eigsh_lanczos( eigenvectors.append(state / self.np.linalg.norm(state)) return eigvals[0:numeig], eigenvectors + def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 + tensor2 + + def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 - tensor2 + def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: return tensor1 * tensor2 + def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 / tensor2 + def index_update(self, tensor: Tensor, mask: Tensor, assignee: Tensor) -> Tensor: t = self.np.copy(tensor) diff --git a/tensornetwork/backends/numpy/numpy_backend_test.py b/tensornetwork/backends/numpy/numpy_backend_test.py index 55fe9edb6..6a8068c41 100644 --- a/tensornetwork/backends/numpy/numpy_backend_test.py +++ b/tensornetwork/backends/numpy/numpy_backend_test.py @@ -3,6 +3,7 @@ import numpy as np import pytest from tensornetwork.backends.numpy import numpy_backend +from unittest.mock import Mock np_randn_dtypes = [np.float32, np.float16, np.float64] np_dtypes = np_randn_dtypes + [np.complex64, np.complex128] @@ -33,20 +34,20 @@ def test_transpose(): np.testing.assert_allclose(expected, actual) -def test_concat(): +def test_shape_concat(): backend = numpy_backend.NumPyBackend() a = backend.convert_to_tensor(2 * np.ones((1, 3, 1))) b = backend.convert_to_tensor(np.ones((1, 2, 1))) - expected = backend.concat((a, b), axis=1) + expected = backend.shape_concat((a, b), axis=1) actual = np.array([[[2.0], [2.0], [2.0], [1.0], [1.0]]]) np.testing.assert_allclose(expected, actual) -def test_shape(): +def test_shape_tensor(): backend = numpy_backend.NumPyBackend() a = backend.convert_to_tensor(np.ones([2, 3, 4])) - assert isinstance(backend.shape(a), tuple) - actual = backend.shape(a) + assert isinstance(backend.shape_tensor(a), tuple) + actual = backend.shape_tensor(a) expected = np.array([2, 3, 4]) np.testing.assert_allclose(expected, actual) @@ -58,10 +59,10 @@ def test_shape_tuple(): assert actual == (2, 3, 4) -def test_prod(): +def test_shape_prod(): backend = numpy_backend.NumPyBackend() a = backend.convert_to_tensor(2 * np.ones([1, 2, 3, 4])) - actual = np.array(backend.prod(a)) + actual = np.array(backend.shape_prod(a)) assert actual == 2**24 @@ -159,6 +160,13 @@ def test_randn(dtype): assert a.shape == (4, 4) +@pytest.mark.parametrize("dtype", np_dtypes) +def test_random_uniform(dtype): + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert a.shape == (4, 4) + + @pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) def test_randn_non_zero_imag(dtype): backend = numpy_backend.NumPyBackend() @@ -166,6 +174,13 @@ def test_randn_non_zero_imag(dtype): assert np.linalg.norm(np.imag(a)) != 0.0 +@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) +def test_random_uniform_non_zero_imag(dtype): + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert np.linalg.norm(np.imag(a)) != 0.0 + + @pytest.mark.parametrize("dtype", np_dtypes) def test_eye_dtype(dtype): backend = numpy_backend.NumPyBackend() @@ -194,6 +209,13 @@ def test_randn_dtype(dtype): assert a.dtype == dtype +@pytest.mark.parametrize("dtype", np_dtypes) +def test_random_uniform_dtype(dtype): + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert a.dtype == dtype + + @pytest.mark.parametrize("dtype", np_randn_dtypes) def test_randn_seed(dtype): backend = numpy_backend.NumPyBackend() @@ -202,6 +224,33 @@ def test_randn_seed(dtype): np.testing.assert_allclose(a, b) +@pytest.mark.parametrize("dtype", np_dtypes) +def test_random_uniform_seed(dtype): + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), seed=10, dtype=dtype) + np.testing.assert_allclose(a, b) + + +@pytest.mark.parametrize("dtype", np_randn_dtypes) +def test_random_uniform_boundaries(dtype): + lb = 1.2 + ub = 4.8 + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype) + assert((a >= 0).all() and (a <= 1).all() and + (b >= lb).all() and (b <= ub).all()) + + +def test_random_uniform_behavior(): + backend = numpy_backend.NumPyBackend() + a = backend.random_uniform((4, 4), seed=10) + np.random.seed(10) + b = np.random.uniform(size=(4, 4)) + np.testing.assert_allclose(a, b) + + def test_conj(): backend = numpy_backend.NumPyBackend() real = np.random.rand(2, 2, 2) @@ -262,6 +311,35 @@ def __call__(self, x): np.testing.assert_allclose(v1, v2) +@pytest.mark.parametrize("dtype", [np.float64, np.complex128]) +def test_eigsh_lanczos_reorthogonalize(dtype): + backend = numpy_backend.NumPyBackend() + D = 24 + np.random.seed(10) + tmp = backend.randn((D, D), dtype=dtype, seed=10) + H = tmp + backend.transpose(backend.conj(tmp), (1, 0)) + + class LinearOperator: + + def __init__(self, shape, dtype): + self.shape = shape + self.dtype = dtype + + def __call__(self, x): + return np.dot(H, x) + + mv = LinearOperator(shape=((D,), (D,)), dtype=dtype) + eta1, U1 = backend.eigsh_lanczos(mv, reorthogonalize=True, ndiag=1, + tol=10**(-12), delta=10**(-12)) + eta2, U2 = np.linalg.eigh(H) + v2 = U2[:, 0] + v2 = v2 / sum(v2) + v1 = np.reshape(U1[0], (D)) + v1 = v1 / sum(v1) + np.testing.assert_allclose(eta1[0], min(eta2)) + np.testing.assert_allclose(v1, v2, rtol=10**(-5), atol=10**(-5)) + + def test_eigsh_lanczos_raises(): backend = numpy_backend.NumPyBackend() with pytest.raises(AttributeError): @@ -272,16 +350,99 @@ def test_eigsh_lanczos_raises(): backend.eigsh_lanczos(lambda x: x, numeig=2, reorthogonalize=False) +def test_eigsh_lanczos_raises_error_for_incompatible_shapes(): + backend = numpy_backend.NumPyBackend() + A = backend.randn((4, 4), dtype=np.float64) + init = backend.randn((3, ), dtype=np.float64) + with pytest.raises(ValueError): + backend.eigsh_lanczos(A, initial_state=init) + + +def test_eigsh_lanczos_raises_error_for_untyped_A(): + backend = numpy_backend.NumPyBackend() + A = Mock(spec=[]) + A.shape = Mock(return_value=(2, 2)) + err_msg = "`A` has no attribute `dtype`. Cannot initialize lanczos. " \ + "Please provide a valid `initial_state` with a `dtype` attribute" + with pytest.raises(AttributeError, match=err_msg): + backend.eigsh_lanczos(A) + + +def test_eigsh_lanczos_raises_error_for_bad_initial_state(): + backend = numpy_backend.NumPyBackend() + D = 16 + init = [1]*D + M = backend.randn((D, D), dtype=np.float64) + + def mv(x): + return np.dot(M, x) + + with pytest.raises(TypeError): + backend.eigsh_lanczos(mv, initial_state=init) + + @pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 2), + pytest.param(1., np.ones((1, 2, 3)), 2*np.ones((1, 2, 3))), + pytest.param(2.*np.ones(()), 1., 3.*np.ones((1, 2, 3))), + pytest.param(2.*np.ones(()), 1.*np.ones((1, 2, 3)), 3.*np.ones((1, 2, 3))), +]) +def test_addition(a, b, expected): + backend = numpy_backend.NumPyBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.addition(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 0), + pytest.param(2., 1.*np.ones((1, 2, 3)), 1.*np.ones((1, 2, 3))), + pytest.param(np.ones((1, 2, 3)), 1., np.zeros((1, 2, 3))), + pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.zeros((1, 2, 3))), +]) +def test_subtraction(a, b, expected): + backend = numpy_backend.NumPyBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.subtraction(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 1), + pytest.param(2., 1.*np.ones((1, 2, 3)), 2.*np.ones((1, 2, 3))), + pytest.param(np.ones((1, 2, 3)), 1., np.ones((1, 2, 3))), pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.ones((1, 2, 3))), - pytest.param(2. * np.ones(()), np.ones((1, 2, 3)), 2. * np.ones((1, 2, 3))), ]) def test_multiply(a, b, expected): backend = numpy_backend.NumPyBackend() tensor1 = backend.convert_to_tensor(a) tensor2 = backend.convert_to_tensor(b) + result = backend.multiply(tensor1, tensor2) - np.testing.assert_allclose(backend.multiply(tensor1, tensor2), expected) + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(2., 2., 1.), + pytest.param(2., 0.5*np.ones((1, 2, 3)), 4.*np.ones((1, 2, 3))), + pytest.param(np.ones(()), 2., 0.5*np.ones((1, 2, 3))), + pytest.param(np.ones(()), 2.*np.ones((1, 2, 3)), 0.5*np.ones((1, 2, 3))), +]) +def test_divide(a, b, expected): + backend = numpy_backend.NumPyBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.divide(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype def find(which, vector): @@ -330,6 +491,84 @@ def mv(x): np.testing.assert_allclose(v1, v2) +@pytest.mark.parametrize("dtype", [np.float64, np.complex128]) +@pytest.mark.parametrize("which", ['LM', 'LR', 'SM', 'SR']) +def test_eigs_no_init(dtype, which): + backend = numpy_backend.NumPyBackend() + D = 16 + np.random.seed(10) + H = backend.randn((D, D), dtype=dtype, seed=10) + + class LinearOperator: + + def __init__(self, shape, dtype): + self.shape = shape + self.dtype = dtype + + def __call__(self, x): + return np.dot(H, x) + + mv = LinearOperator(shape=((D,), (D,)), dtype=dtype) + eta1, U1 = backend.eigs(mv, numeig=1, which=which) + eta2, U2 = np.linalg.eig(H) + val, index = find(which, eta2) + v2 = U2[:, index] + v2 = v2 / sum(v2) + v1 = np.reshape(U1[0], (D)) + v1 = v1 / sum(v1) + np.testing.assert_allclose(find(which, eta1)[0], val) + np.testing.assert_allclose(v1, v2) + + +@pytest.mark.parametrize("which", ['SI', 'LI']) +def test_eigs_raises_error_for_unsupported_which(which): + backend = numpy_backend.NumPyBackend() + A = backend.randn((4, 4), dtype=np.float64) + with pytest.raises(ValueError): + backend.eigs(A=A, which=which) + + +def test_eigs_raises_error_for_incompatible_shapes(): + backend = numpy_backend.NumPyBackend() + A = backend.randn((4, 4), dtype=np.float64) + init = backend.randn((3, ), dtype=np.float64) + with pytest.raises(ValueError): + backend.eigs(A, initial_state=init) + + +def test_eigs_raises_error_for_unshaped_A(): + backend = numpy_backend.NumPyBackend() + A = Mock(spec=[]) + print(hasattr(A, "shape")) + err_msg = "`A` has no attribute `shape`. Cannot initialize lanczos. " \ + "Please provide a valid `initial_state`" + with pytest.raises(AttributeError, match=err_msg): + backend.eigs(A) + + +def test_eigs_raises_error_for_untyped_A(): + backend = numpy_backend.NumPyBackend() + A = Mock(spec=[]) + A.shape = Mock(return_value=(2, 2)) + err_msg = "`A` has no attribute `dtype`. Cannot initialize lanczos. " \ + "Please provide a valid `initial_state` with a `dtype` attribute" + with pytest.raises(AttributeError, match=err_msg): + backend.eigs(A) + + +def test_eigs_raises_error_for_bad_initial_state(): + backend = numpy_backend.NumPyBackend() + D = 16 + init = [1]*D + M = backend.randn((D, D), dtype=np.float64) + + def mv(x): + return np.dot(M, x) + + with pytest.raises(TypeError): + backend.eigs(mv, initial_state=init) + + @pytest.mark.parametrize("dtype", [np.float64, np.complex128]) def test_eigh(dtype): backend = numpy_backend.NumPyBackend() diff --git a/tensornetwork/backends/pytorch/pytorch_backend.py b/tensornetwork/backends/pytorch/pytorch_backend.py index 9a979b1c8..788d4258b 100644 --- a/tensornetwork/backends/pytorch/pytorch_backend.py +++ b/tensornetwork/backends/pytorch/pytorch_backend.py @@ -69,16 +69,16 @@ def rq_decomposition( ) -> Tuple[Tensor, Tensor]: return decompositions.rq_decomposition(self.torch, tensor, split_axis) - def concat(self, values: Tensor, axis: int) -> Tensor: + def shape_concat(self, values: Tensor, axis: int) -> Tensor: return np.concatenate(values, axis) - def shape(self, tensor: Tensor) -> Tensor: + def shape_tensor(self, tensor: Tensor) -> Tensor: return self.torch.tensor(list(tensor.shape)) def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]: return tuple(tensor.shape) - def prod(self, values: Tensor) -> int: + def shape_prod(self, values: Tensor) -> int: return np.prod(np.array(values)) def sqrt(self, tensor: Tensor) -> Tensor: @@ -128,6 +128,16 @@ def randn(self, dtype = dtype if dtype is not None else self.torch.float64 return self.torch.randn(shape, dtype=dtype) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[Any] = None, + seed: Optional[int] = None) -> Tensor: + if seed: + self.torch.manual_seed(seed) + dtype = dtype if dtype is not None else self.torch.float64 + return self.torch.empty(shape, dtype=dtype).uniform_(*boundaries) + def conj(self, tensor: Tensor) -> Tensor: return tensor #pytorch does not support complex dtypes @@ -265,9 +275,18 @@ def eigsh_lanczos(self, eigenvectors.append(state / self.torch.norm(state)) return eigvals[0:numeig], eigenvectors + def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 + tensor2 + + def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 - tensor2 + def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: return tensor1 * tensor2 + def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 / tensor2 + def index_update(self, tensor: Tensor, mask: Tensor, assignee: Tensor) -> Tensor: #make a copy diff --git a/tensornetwork/backends/pytorch/pytorch_backend_test.py b/tensornetwork/backends/pytorch/pytorch_backend_test.py index 5e3ead3f9..f03dc26a3 100644 --- a/tensornetwork/backends/pytorch/pytorch_backend_test.py +++ b/tensornetwork/backends/pytorch/pytorch_backend_test.py @@ -3,6 +3,7 @@ from tensornetwork.backends.pytorch import pytorch_backend import torch import pytest +from unittest.mock import Mock torch_dtypes = [torch.float32, torch.float64, torch.int32] torch_eye_dtypes = [torch.float32, torch.float64, torch.int32, torch.int64] @@ -34,20 +35,20 @@ def test_transpose(): np.testing.assert_allclose(expected, actual) -def test_concat(): +def test_shape_concat(): backend = pytorch_backend.PyTorchBackend() a = backend.convert_to_tensor(2 * np.ones((1, 3, 1))) b = backend.convert_to_tensor(np.ones((1, 2, 1))) - expected = backend.concat((a, b), axis=1) + expected = backend.shape_concat((a, b), axis=1) actual = np.array([[[2.0], [2.0], [2.0], [1.0], [1.0]]]) np.testing.assert_allclose(expected, actual) -def test_shape(): +def test_shape_tensor(): backend = pytorch_backend.PyTorchBackend() a = backend.convert_to_tensor(np.ones([2, 3, 4])) - assert isinstance(backend.shape(a), torch.Tensor) - actual = backend.shape(a) + assert isinstance(backend.shape_tensor(a), torch.Tensor) + actual = backend.shape_tensor(a) expected = np.array([2, 3, 4]) np.testing.assert_allclose(expected, actual) @@ -59,10 +60,10 @@ def test_shape_tuple(): assert actual == (2, 3, 4) -def test_prod(): +def test_shape_prod(): backend = pytorch_backend.PyTorchBackend() a = backend.convert_to_tensor(2 * np.ones([1, 2, 3, 4])) - actual = np.array(backend.prod(a)) + actual = np.array(backend.shape_prod(a)) assert actual == 2**24 @@ -148,6 +149,13 @@ def test_randn(dtype): assert a.shape == (4, 4) +@pytest.mark.parametrize("dtype", torch_randn_dtypes) +def test_random_uniform(dtype): + backend = pytorch_backend.PyTorchBackend() + a = backend.random_uniform((4, 4), dtype=dtype) + assert a.shape == (4, 4) + + @pytest.mark.parametrize("dtype", torch_eye_dtypes) def test_eye_dtype(dtype): backend = pytorch_backend.PyTorchBackend() @@ -176,6 +184,13 @@ def test_randn_dtype(dtype): assert a.dtype == dtype +@pytest.mark.parametrize("dtype", torch_randn_dtypes) +def test_random_uniform_dtype(dtype): + backend = pytorch_backend.PyTorchBackend() + a = backend.random_uniform((4, 4), dtype=dtype) + assert a.dtype == dtype + + @pytest.mark.parametrize("dtype", torch_randn_dtypes) def test_randn_seed(dtype): backend = pytorch_backend.PyTorchBackend() @@ -184,6 +199,33 @@ def test_randn_seed(dtype): np.testing.assert_allclose(a, b) +@pytest.mark.parametrize("dtype", torch_randn_dtypes) +def test_random_uniform_seed(dtype): + backend = pytorch_backend.PyTorchBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), seed=10, dtype=dtype) + torch.allclose(a, b) + + +@pytest.mark.parametrize("dtype", torch_randn_dtypes) +def test_random_uniform_boundaries(dtype): + lb = 1.2 + ub = 4.8 + backend = pytorch_backend.PyTorchBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype) + assert(torch.ge(a, 0).byte().all() and torch.le(a, 1).byte().all() and + torch.ge(b, lb).byte().all() and torch.le(b, ub).byte().all()) + + +def test_random_uniform_behavior(): + backend = pytorch_backend.PyTorchBackend() + a = backend.random_uniform((4, 4), seed=10) + torch.manual_seed(10) + b = torch.empty((4, 4), dtype=torch.float64).uniform_() + torch.allclose(a, b) + + def test_conj(): backend = pytorch_backend.PyTorchBackend() real = np.random.rand(2, 2, 2) @@ -196,7 +238,7 @@ def test_conj(): def test_eigsh_lanczos_1(): dtype = torch.float64 backend = pytorch_backend.PyTorchBackend() - D = 16 + D = 24 init = backend.randn((D,), dtype=dtype) tmp = backend.randn((D, D), dtype=dtype) H = tmp + backend.transpose(backend.conj(tmp), (1, 0)) @@ -214,10 +256,11 @@ def mv(x): np.testing.assert_allclose(v1, v2) -def test_eigsh_lanczos_2(): +def test_eigsh_lanczos_reorthogonalize(): dtype = torch.float64 backend = pytorch_backend.PyTorchBackend() D = 16 + init = backend.randn((D,), dtype=dtype) tmp = backend.randn((D, D), dtype=dtype) H = tmp + backend.transpose(backend.conj(tmp), (1, 0)) @@ -231,7 +274,7 @@ def __call__(self, x): return H.mv(x) mv = LinearOperator(shape=((D,), (D,)), dtype=dtype) - eta1, U1 = backend.eigsh_lanczos(mv) + eta1, U1 = backend.eigsh_lanczos(mv, init) eta2, U2 = H.symeig() v2 = U2[:, 0] v2 = v2 / sum(v2) @@ -241,6 +284,34 @@ def __call__(self, x): np.testing.assert_allclose(v1, v2) +def test_eigsh_lanczos_2(): + dtype = torch.float64 + backend = pytorch_backend.PyTorchBackend() + D = 16 + tmp = backend.randn((D, D), dtype=dtype) + H = tmp + backend.transpose(backend.conj(tmp), (1, 0)) + + class LinearOperator: + + def __init__(self, shape, dtype): + self.shape = shape + self.dtype = dtype + + def __call__(self, x): + return H.mv(x) + + mv = LinearOperator(shape=((D,), (D,)), dtype=dtype) + eta1, U1 = backend.eigsh_lanczos(mv, reorthogonalize=True, ndiag=1, + tol=10**(-12), delta=10**(-12)) + eta2, U2 = H.symeig() + v2 = U2[:, 0] + v2 = v2 / sum(v2) + v1 = np.reshape(U1[0], (D)) + v1 = v1 / sum(v1) + np.testing.assert_allclose(eta1[0], min(eta2)) + np.testing.assert_allclose(v1, v2, rtol=10**(-5), atol=10**(-5)) + + def test_eigsh_lanczos_raises(): backend = pytorch_backend.PyTorchBackend() with pytest.raises(AttributeError): @@ -252,15 +323,59 @@ def test_eigsh_lanczos_raises(): @pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 2), + pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), 2.*np.ones((1, 2, 3))), +]) +def test_addition(a, b, expected): + backend = pytorch_backend.PyTorchBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.addition(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 0), + pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.zeros((1, 2, 3))), +]) +def test_subtraction(a, b, expected): + backend = pytorch_backend.PyTorchBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.subtraction(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 1), pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.ones((1, 2, 3))), - pytest.param(2. * np.ones(()), np.ones((1, 2, 3)), 2. * np.ones((1, 2, 3))), ]) def test_multiply(a, b, expected): backend = pytorch_backend.PyTorchBackend() tensor1 = backend.convert_to_tensor(a) tensor2 = backend.convert_to_tensor(b) + result = backend.multiply(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(2., 2., 1.), + pytest.param(np.ones(()), 2.*np.ones((1, 2, 3)), 0.5*np.ones((1, 2, 3))), +]) +def test_divide(a, b, expected): + backend = pytorch_backend.PyTorchBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.divide(tensor1, tensor2) - np.testing.assert_allclose(backend.multiply(tensor1, tensor2), expected) + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype def test_eigh(): @@ -303,3 +418,27 @@ def test_matrix_inv_raises(dtype): matrix = backend.randn((4, 4, 4), dtype=dtype, seed=10) with pytest.raises(ValueError): backend.inv(matrix) + + +def test_eigs_not_implemented(): + backend = pytorch_backend.PyTorchBackend() + with pytest.raises(NotImplementedError): + backend.eigs(np.ones((2, 2))) + + +def test_eigsh_lanczos_raises_error_for_incompatible_shapes(): + backend = pytorch_backend.PyTorchBackend() + A = backend.randn((4, 4), dtype=torch.float64) + init = backend.randn((3, ), dtype=torch.float64) + with pytest.raises(ValueError): + backend.eigsh_lanczos(A, initial_state=init) + + +def test_eigsh_lanczos_raises_error_for_untyped_A(): + backend = pytorch_backend.PyTorchBackend() + A = Mock(spec=[]) + A.shape = Mock(return_value=(2, 2)) + err_msg = "`A` has no attribute `dtype`. Cannot initialize lanczos. " \ + "Please provide a valid `initial_state` with a `dtype` attribute" + with pytest.raises(AttributeError, match=err_msg): + backend.eigsh_lanczos(A) diff --git a/tensornetwork/backends/shell/shell_backend.py b/tensornetwork/backends/shell/shell_backend.py index 4e73f638d..6cc6d9ae2 100644 --- a/tensornetwork/backends/shell/shell_backend.py +++ b/tensornetwork/backends/shell/shell_backend.py @@ -107,7 +107,7 @@ def rq_decomposition(self, tensor: Tensor, r = ShellTensor((center_dim,) + right_dims) return q, r - def concat(self, values: Sequence[Tensor], axis: int) -> Tensor: + def shape_concat(self, values: Sequence[Tensor], axis: int) -> Tensor: shape = values[0].shape if axis < 0: axis += len(shape) @@ -119,20 +119,20 @@ def concat_shape(self, values) -> Sequence: tuple_values = (tuple(v) for v in values) return functools.reduce(operator.concat, tuple_values) - def shape(self, tensor: Tensor) -> Tuple: + def shape_tensor(self, tensor: Tensor) -> Tuple: return tensor.shape def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]: return tensor.shape - def prod(self, values: Tensor) -> int: + def shape_prod(self, values: Tensor) -> int: # This is different from the BaseBackend prod! # prod calculates the product of tensor elements and cannot implemented # for shell tensors # This returns the product of sizes instead - return self.shape_prod(values.shape) + return self.shape_product(values.shape) - def shape_prod(self, shape: Sequence[int]) -> int: + def shape_product(self, shape: Sequence[int]) -> int: return functools.reduce(operator.mul, shape) def sqrt(self, tensor: Tensor) -> Tensor: @@ -207,6 +207,13 @@ def randn(self, seed: Optional[int] = None) -> Tensor: return ShellTensor(shape) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[Type[np.number]] = None, + seed: Optional[int] = None) -> Tensor: + return ShellTensor(shape) + def conj(self, tensor: Tensor) -> Tensor: return tensor @@ -283,11 +290,20 @@ def eigsh_lanczos(self, raise ValueError( '`A` has no attribut shape adn no `initial_state` is given.') + def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + raise NotImplementedError("Shell tensor has not implemented addition( + )") + + def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + raise NotImplementedError("Shell tensor has not implemented subtraction( - )") + def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: a = np.ones(tensor1.shape) b = np.ones(tensor2.shape) return ShellTensor((a * b).shape) + def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + raise NotImplementedError("Shell tensor has not implemented add( / )") + def index_update(self, tensor: Tensor, mask: Tensor, assignee: Tensor) -> Tensor: return ShellTensor(tensor.shape) diff --git a/tensornetwork/backends/shell/shell_backend_test.py b/tensornetwork/backends/shell/shell_backend_test.py index 45a713965..af17c3354 100644 --- a/tensornetwork/backends/shell/shell_backend_test.py +++ b/tensornetwork/backends/shell/shell_backend_test.py @@ -62,16 +62,16 @@ def test_svd_decomposition_with_max_values(): assert x.shape == y.shape -def test_concat(): +def test_shape_concat(): args = { "values": [np.ones([3, 2, 5]), np.zeros([3, 2, 5]), np.ones([3, 3, 5])] } args["axis"] = 1 - assertBackendsAgree("concat", args) + assertBackendsAgree("shape_concat", args) args["axis"] = -2 - assertBackendsAgree("concat", args) + assertBackendsAgree("shape_concat", args) def test_concat_shape(): @@ -80,10 +80,10 @@ def test_concat_shape(): assert result == (5, 2, 3, 4, 6) -def test_shape(): +def test_shape_tensor(): tensor = np.ones([3, 5, 2]) - np_result = numpy_backend.NumPyBackend().shape(tensor) - sh_result = shell_backend.ShellBackend().shape(tensor) + np_result = numpy_backend.NumPyBackend().shape_tensor(tensor) + sh_result = shell_backend.ShellBackend().shape_tensor(tensor) assert np_result == sh_result @@ -94,8 +94,8 @@ def test_shape_tuple(): assert np_result == sh_result -def test_prod(): - result = shell_backend.ShellBackend().prod(np.ones([3, 5, 2])) +def test_shape_prod(): + result = shell_backend.ShellBackend().shape_prod(np.ones([3, 5, 2])) assert result == 30 @@ -157,6 +157,11 @@ def test_randn(): assertBackendsAgree("randn", args) +def test_random_uniform(): + args = {"shape": (10, 4)} + assertBackendsAgree("random_uniform", args) + + def test_eigsh_lanczos_1(): backend = shell_backend.ShellBackend() D = 16 diff --git a/tensornetwork/backends/tensorflow/tensordot2.py b/tensornetwork/backends/tensorflow/tensordot2.py index 00de66c49..94e980f4e 100644 --- a/tensornetwork/backends/tensorflow/tensordot2.py +++ b/tensornetwork/backends/tensorflow/tensordot2.py @@ -19,11 +19,7 @@ Tensor = Any -def tensordot(tf, - a, - b, - axes, - name: Optional[Text] = None) -> Tensor: +def tensordot(tf, a, b, axes, name: Optional[Text] = None) -> Tensor: r"""Tensor contraction of a and b along specified axes. Tensordot (also known as tensor contraction) sums the product of elements from `a` and `b` over the indices specified by `a_axes` and `b_axes`. @@ -80,7 +76,7 @@ def _tensordot_should_flip(contraction_axes: List[int], return bool(np.mean(contraction_axes) < np.mean(free_axes)) return False - def _tranpose_if_necessary(tensor: Tensor, perm: List[int]) -> Tensor: + def _transpose_if_necessary(tensor: Tensor, perm: List[int]) -> Tensor: """Like transpose(), but avoids creating a new tensor if possible. Although the graph optimizer should kill trivial transposes, it is best not to add them in the first place! @@ -89,8 +85,7 @@ def _tranpose_if_necessary(tensor: Tensor, perm: List[int]) -> Tensor: return tensor return tf.transpose(tensor, perm) - def _reshape_if_necessary(tensor: Tensor, - new_shape: List[int]) -> Tensor: + def _reshape_if_necessary(tensor: Tensor, new_shape: List[int]) -> Tensor: """Like reshape(), but avoids creating a new tensor if possible. Assumes shapes are both fully specified.""" cur_shape = tensor.get_shape().as_list() @@ -135,7 +130,7 @@ def _tensordot_reshape( prod_axes = int(np.prod([shape_a[i] for i in axes])) perm = axes + free if flipped else free + axes new_shape = [prod_axes, prod_free] if flipped else [prod_free, prod_axes] - transposed_a = _tranpose_if_necessary(a, perm) + transposed_a = _transpose_if_necessary(a, perm) reshaped_a = _reshape_if_necessary(transposed_a, new_shape) transpose_needed = (not flipped) if is_right_term else flipped return reshaped_a, free_dims, free_dims, transpose_needed @@ -152,7 +147,7 @@ def _tensordot_reshape( axes = tf.convert_to_tensor(axes, dtype=tf.dtypes.int32, name="axes") free = tf.convert_to_tensor(free, dtype=tf.dtypes.int32, name="free") shape_a = tf.shape(a) - transposed_a = _tranpose_if_necessary(a, perm) + transposed_a = _transpose_if_necessary(a, perm) else: free_dims_static = None shape_a = tf.shape(a) @@ -184,8 +179,7 @@ def _tensordot_reshape( transpose_needed = (not flipped) if is_right_term else flipped return reshaped_a, free_dims, free_dims_static, transpose_needed - def _tensordot_axes(a: Tensor, axes - ) -> Tuple[Any, Any]: + def _tensordot_axes(a: Tensor, axes) -> Tuple[Any, Any]: """Generates two sets of contraction axes for the two tensor arguments.""" a_shape = a.get_shape() if isinstance(axes, tf.compat.integral_types): @@ -195,11 +189,11 @@ def _tensordot_axes(a: Tensor, axes if axes > a_shape.ndims: raise ValueError("'axes' must not be larger than the number of " "dimensions of tensor %s." % a) - return (list(range(a_shape.ndims - axes, - a_shape.ndims)), list(range(axes))) + return (list(range(a_shape.ndims - axes, a_shape.ndims)), + list(range(axes))) rank = tf.rank(a) - return (tf.range(rank - axes, rank, - dtype=tf.int32), tf.range(axes, dtype=tf.int32)) + return (tf.range(rank - axes, rank, dtype=tf.int32), + tf.range(axes, dtype=tf.int32)) if isinstance(axes, (list, tuple)): if len(axes) != 2: raise ValueError("'axes' must be an integer or have length 2.") diff --git a/tensornetwork/backends/tensorflow/tensorflow_backend.py b/tensornetwork/backends/tensorflow/tensorflow_backend.py index 2602984ed..2d9c0ed4d 100644 --- a/tensornetwork/backends/tensorflow/tensorflow_backend.py +++ b/tensornetwork/backends/tensorflow/tensorflow_backend.py @@ -64,16 +64,16 @@ def rq_decomposition(self, tensor: Tensor, split_axis: int) -> Tuple[Tensor, Tensor]: return decompositions.rq_decomposition(self.tf, tensor, split_axis) - def concat(self, values: Tensor, axis: int) -> Tensor: + def shape_concat(self, values: Tensor, axis: int) -> Tensor: return self.tf.concat(values, axis) - def shape(self, tensor: Tensor) -> Tensor: + def shape_tensor(self, tensor: Tensor) -> Tensor: return self.tf.shape(tensor) def shape_tuple(self, tensor: Tensor) -> Tuple[Optional[int], ...]: return tuple(tensor.shape.as_list()) - def prod(self, values: Tensor) -> Tensor: + def shape_prod(self, values: Tensor) -> Tensor: return self.tf.reduce_prod(values) def sqrt(self, tensor: Tensor) -> Tensor: @@ -131,6 +131,26 @@ def randn(self, self.tf.random.normal(shape=shape, dtype=dtype.real_dtype)) return self.tf.random.normal(shape=shape, dtype=dtype) + def random_uniform(self, + shape: Tuple[int, ...], + boundaries: Optional[Tuple[float, float]] = (0.0, 1.0), + dtype: Optional[Type[np.number]] = None, + seed: Optional[int] = None) -> Tensor: + if seed: + self.tf.random.set_seed(seed) + + dtype = dtype if dtype is not None else self.tf.float64 + if (dtype is self.tf.complex128) or (dtype is self.tf.complex64): + return self.tf.complex( + self.tf.random.uniform(shape=shape, minval=boundaries[0], + maxval=boundaries[1], dtype=dtype.real_dtype), + self.tf.random.uniform(shape=shape, minval=boundaries[0], + maxval=boundaries[1], dtype=dtype.real_dtype)) + self.tf.random.set_seed(10) + a = self.tf.random.uniform(shape=shape, minval=boundaries[0], + maxval=boundaries[1], dtype=dtype) + return a + def conj(self, tensor: Tensor) -> Tensor: return self.tf.math.conj(tensor) @@ -162,9 +182,18 @@ def eigsh_lanczos(self, raise NotImplementedError( "Backend '{}' has not implemented eighs_lanczos.".format(self.name)) + def addition(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 + tensor2 + + def subtraction(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 - tensor2 + def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: return tensor1 * tensor2 + def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor: + return tensor1 / tensor2 + def index_update(self, tensor: Tensor, mask: Tensor, assignee: Tensor) -> Tensor: #returns a copy (unfortunately) diff --git a/tensornetwork/backends/tensorflow/tensorflow_backend_test.py b/tensornetwork/backends/tensorflow/tensorflow_backend_test.py index 8df3fcce0..4916301ce 100644 --- a/tensornetwork/backends/tensorflow/tensorflow_backend_test.py +++ b/tensornetwork/backends/tensorflow/tensorflow_backend_test.py @@ -34,20 +34,20 @@ def test_transpose(): np.testing.assert_allclose(expected, actual) -def test_concat(): +def test_shape_concat(): backend = tensorflow_backend.TensorFlowBackend() a = backend.convert_to_tensor(2 * np.ones((1, 3, 1))) b = backend.convert_to_tensor(np.ones((1, 2, 1))) - expected = backend.concat((a, b), axis=1) + expected = backend.shape_concat((a, b), axis=1) actual = np.array([[[2.0], [2.0], [2.0], [1.0], [1.0]]]) np.testing.assert_allclose(expected, actual) -def test_shape(): +def test_shape_tensor(): backend = tensorflow_backend.TensorFlowBackend() a = backend.convert_to_tensor(np.ones([2, 3, 4])) - assert isinstance(backend.shape(a), type(a)) - actual = backend.shape(a) + assert isinstance(backend.shape_tensor(a), type(a)) + actual = backend.shape_tensor(a) expected = np.array([2, 3, 4]) np.testing.assert_allclose(expected, actual) @@ -59,10 +59,10 @@ def test_shape_tuple(): assert actual == (2, 3, 4) -def test_prod(): +def test_shape_prod(): backend = tensorflow_backend.TensorFlowBackend() a = backend.convert_to_tensor(2 * np.ones([1, 2, 3, 4])) - actual = np.array(backend.prod(a)) + actual = np.array(backend.shape_prod(a)) assert actual == 2**24 @@ -151,6 +151,13 @@ def test_randn(dtype): assert a.shape == (4, 4) +@pytest.mark.parametrize("dtype", tf_dtypes) +def test_random_uniform(dtype): + backend = tensorflow_backend.TensorFlowBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert a.shape == (4, 4) + + @pytest.mark.parametrize("dtype", [tf.complex64, tf.complex128]) def test_randn_non_zero_imag(dtype): backend = tensorflow_backend.TensorFlowBackend() @@ -158,6 +165,13 @@ def test_randn_non_zero_imag(dtype): assert tf.math.greater(tf.linalg.norm(tf.math.imag(a)), 0.0) +@pytest.mark.parametrize("dtype", [tf.complex64, tf.complex128]) +def test_random_uniform_non_zero_imag(dtype): + backend = tensorflow_backend.TensorFlowBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert tf.math.greater(tf.linalg.norm(tf.math.imag(a)), 0.0) + + @pytest.mark.parametrize("dtype", tf_dtypes) def test_eye_dtype(dtype): backend = tensorflow_backend.TensorFlowBackend() @@ -186,6 +200,13 @@ def test_randn_dtype(dtype): assert a.dtype == dtype +@pytest.mark.parametrize("dtype", tf_dtypes) +def test_random_uniform_dtype(dtype): + backend = tensorflow_backend.TensorFlowBackend() + a = backend.random_uniform((4, 4), dtype=dtype, seed=10) + assert a.dtype == dtype + + @pytest.mark.parametrize("dtype", tf_randn_dtypes) def test_randn_seed(dtype): backend = tensorflow_backend.TensorFlowBackend() @@ -194,6 +215,27 @@ def test_randn_seed(dtype): np.testing.assert_allclose(a, b) +@pytest.mark.parametrize("dtype", tf_dtypes) +def test_random_uniform_seed(dtype): + test = tf.test.TestCase() + backend = tensorflow_backend.TensorFlowBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), seed=10, dtype=dtype) + test.assertAllCloseAccordingToType(a, b) + + +@pytest.mark.parametrize("dtype", tf_randn_dtypes) +def test_random_uniform_boundaries(dtype): + test = tf.test.TestCase() + lb = 1.2 + ub = 4.8 + backend = tensorflow_backend.TensorFlowBackend() + a = backend.random_uniform((4, 4), seed=10, dtype=dtype) + b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype) + test.assertAllInRange(a, 0, 1) + test.assertAllInRange(b, lb, ub) + + def test_conj(): backend = tensorflow_backend.TensorFlowBackend() real = np.random.rand(2, 2, 2) @@ -205,15 +247,59 @@ def test_conj(): @pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 2), + pytest.param(2.*np.ones(()), 1.*np.ones((1, 2, 3)), 3.*np.ones((1, 2, 3))), +]) +def test_addition(a, b, expected): + backend = tensorflow_backend.TensorFlowBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.addition(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 0), + pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.zeros((1, 2, 3))), +]) +def test_subtraction(a, b, expected): + backend = tensorflow_backend.TensorFlowBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.subtraction(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(1, 1, 1), pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.ones((1, 2, 3))), - pytest.param(2. * np.ones(()), np.ones((1, 2, 3)), 2. * np.ones((1, 2, 3))), ]) def test_multiply(a, b, expected): backend = tensorflow_backend.TensorFlowBackend() tensor1 = backend.convert_to_tensor(a) tensor2 = backend.convert_to_tensor(b) + result = backend.multiply(tensor1, tensor2) - np.testing.assert_allclose(backend.multiply(tensor1, tensor2), expected) + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype + + +@pytest.mark.parametrize("a, b, expected", [ + pytest.param(2., 2., 1.), + pytest.param(np.ones(()), 2.*np.ones((1, 2, 3)), 0.5*np.ones((1, 2, 3))), +]) +def test_divide(a, b, expected): + backend = tensorflow_backend.TensorFlowBackend() + tensor1 = backend.convert_to_tensor(a) + tensor2 = backend.convert_to_tensor(b) + result = backend.divide(tensor1, tensor2) + + np.testing.assert_allclose(result, expected) + assert tensor1.dtype == tensor2.dtype == result.dtype @pytest.mark.parametrize("dtype", [tf.float64, tf.complex128]) @@ -256,3 +342,14 @@ def test_matrix_inv_raises(dtype): matrix = backend.randn((4, 4, 4), dtype=dtype, seed=10) with pytest.raises(ValueError): backend.inv(matrix) + +def test_eigs_not_implemented(): + backend = tensorflow_backend.TensorFlowBackend() + with pytest.raises(NotImplementedError): + backend.eigs(np.ones((2, 2))) + + +def test_eigsh_lanczos_not_implemented(): + backend = tensorflow_backend.TensorFlowBackend() + with pytest.raises(NotImplementedError): + backend.eigsh_lanczos(np.ones((2, 2))) diff --git a/tensornetwork/block_tensor/benchmarks.py b/tensornetwork/block_tensor/benchmarks.py new file mode 100644 index 000000000..710c75f7b --- /dev/null +++ b/tensornetwork/block_tensor/benchmarks.py @@ -0,0 +1,171 @@ +import tensornetwork as tn +import numpy as np +import time +from tensornetwork.block_tensor.block_tensor import BlockSparseTensor, tensordot +from tensornetwork.block_tensor.index import Index +from tensornetwork.block_tensor.charge import U1Charge + + +def benchmark_1_U1(): + R = 6 + charges = [ + U1Charge( + np.asarray([-2, -1, -1, 0, -1, 0, 0, 1, -1, 0, 0, 1, 0, 1, 1, 2], + dtype=np.int16)) for n in range(R) + ] + + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 1- U1: {}s'.format(time.time() - t1)) + + +def benchmark_1_U1xU1(): + + R = 6 + charges = [ + U1Charge( + np.asarray([-2, -1, -1, 0, -1, 0, 0, 1, -1, 0, 0, 1, 0, 1, 1, 2], + dtype=np.int16)) + @ U1Charge( + np.asarray([0, -1, 1, 0, -1, -2, 0, -1, 1, 0, 2, 0, 0, -1, 1, 0], + dtype=np.int16)) for n in range(R) + ] + + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 1- U1xU1: {}s'.format(time.time() - t1)) + + +def benchmark_2_U1(): + R = 12 + charges = [ + U1Charge(np.asarray([-1, 0, 0, 1], dtype=np.int16)) for n in range(R) + ] + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 2- U1: {}s'.format(time.time() - t1)) + + +def benchmark_2_U1xU1(): + R = 12 + charges = [ + U1Charge(np.asarray([-1, 0, 0, 1], dtype=np.int16)) @ U1Charge( + np.asarray([0, -1, 1, 0], dtype=np.int16)) for n in range(R) + ] + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 2- U1xU1: {}s'.format(time.time() - t1)) + + +def benchmark_3_U1(): + R = 14 + charges = [ + U1Charge(np.asarray([-1, 0, 0, 1], dtype=np.int16)) for n in range(R) + ] + + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 3- U1: {}s'.format(time.time() - t1)) + + +def benchmark_3_U1xU1(): + R = 14 + charges = [ + U1Charge(np.asarray([-1, 0, 0, 1], dtype=np.int16)) @ U1Charge( + np.asarray([0, -1, 1, 0], dtype=np.int16)) for n in range(R) + ] + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + flowsA = np.asarray([False] * R) + flowsB = np.asarray([False] * R) + flowsB[indsB] = True + A = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsA[n], name='a{}'.format(n)) for n in range(R) + ]) + B = BlockSparseTensor.random(indices=[ + Index(charges[n], flowsB[n], name='b{}'.format(n)) for n in range(R) + ]) + final_order = np.arange(R) + np.random.shuffle(final_order) + t1 = time.time() + res = tensordot(A, B, (indsA, indsB), final_order) + print('BM 3- U1xU1: {}s'.format(time.time() - t1)) + + +benchmark_1_U1() +benchmark_1_U1xU1() +benchmark_2_U1() +benchmark_2_U1xU1() +benchmark_3_U1() +benchmark_3_U1xU1() diff --git a/tensornetwork/block_tensor/block_tensor.py b/tensornetwork/block_tensor/block_tensor.py index ba11e2965..786058c72 100644 --- a/tensornetwork/block_tensor/block_tensor.py +++ b/tensornetwork/block_tensor/block_tensor.py @@ -16,72 +16,104 @@ from __future__ import division from __future__ import print_function import numpy as np -# pylint: disable=line-too-long -from tensornetwork.network_components import Node, contract, contract_between +#from tensornetwork.block_tensor.lookup import lookup from tensornetwork.backends import backend_factory # pylint: disable=line-too-long -from tensornetwork.block_tensor.index import Index, fuse_index_pair, split_index, fuse_charge_pair, fuse_degeneracies, fuse_charges, unfuse +from tensornetwork.block_tensor.index import Index, fuse_index_pair, split_index +# pylint: disable=line-too-long +from tensornetwork.block_tensor.charge import fuse_degeneracies, fuse_charges, fuse_degeneracies, BaseCharge, BaseCharge, fuse_ndarray_charges, intersect import numpy as np +import scipy as sp import itertools import time -from typing import List, Union, Any, Tuple, Type, Optional, Dict, Iterable +from typing import List, Union, Any, Tuple, Type, Optional, Dict, Iterable, Sequence Tensor = Any -def _check_flows(flows) -> None: - if (set(flows) != {1}) and (set(flows) != {-1}) and (set(flows) != {-1, 1}): - raise ValueError( - "flows = {} contains values different from 1 and -1".format(flows)) +def fuse_stride_arrays(dims: np.ndarray, strides: np.ndarray) -> np.ndarray: + return fuse_ndarrays([ + np.arange(0, strides[n] * dims[n], strides[n], dtype=np.uint32) + for n in range(len(dims)) + ]) -def _find_best_partition(charges, flows): - if len(charges) == 1: - raise ValueError( - '_expecting `charges` with a length of at least 2, got `len(charges)={}`' - .format(len(charges))) - dims = np.asarray([len(c) for c in charges]) - min_ind = np.argmin([ - np.abs(np.prod(dims[0:n]) - np.prod(dims[n::])) - for n in range(1, len(charges)) - ]) - fused_left_charges = fuse_charges(charges[0:min_ind + 1], - flows[0:min_ind + 1]) - fused_right_charges = fuse_charges(charges[min_ind + 1::], - flows[min_ind + 1::]) +def compute_sparse_lookup(charges: List[BaseCharge], flows: Iterable[bool], + target_charges: BaseCharge) -> np.ndarray: + """ + Compute lookup table for looking up how dense index positions map + to sparse index positions for the diagonal blocks a symmetric matrix. + Args: + charges: + flows: + target_charges: + + """ + fused_charges = fuse_charges(charges, flows) + unique_charges, inverse, degens = fused_charges.unique( + return_inverse=True, return_counts=True) + common_charges, label_to_unique, label_to_target = unique_charges.intersect( + target_charges, return_indices=True) + + tmp = np.full(len(unique_charges), fill_value=-1, dtype=np.int16) + tmp[label_to_unique] = label_to_unique + lookup = tmp[inverse] + vec = np.empty(len(fused_charges), dtype=np.uint32) + for n in label_to_unique: + vec[lookup == n] = np.arange(degens[n]) + return vec - return fused_left_charges, fused_right_charges +def _get_strides(dims): + return np.flip(np.append(1, np.cumprod(np.flip(dims[1::])))) -def map_to_integer(dims: Union[List, np.ndarray], - table: np.ndarray, - dtype: Optional[Type[np.number]] = np.int64): + +def fuse_ndarrays(arrays: List[Union[List, np.ndarray]]) -> np.ndarray: """ - Map a `table` of integers of shape (N, r) bijectively into - an np.ndarray `integers` of length N of unique numbers. - The mapping is done using - ``` - `integers[n] = table[n,0] * np.prod(dims[1::]) + table[n,1] * np.prod(dims[2::]) + ... + table[n,r-1] * 1` - + Fuse all `arrays` by simple kronecker addition. + Arrays are fused from "right to left", Args: - dims: An iterable of integers. - table: An array of shape (N,r) of integers. - dtype: An optional dtype used for the conversion. - Care should be taken when choosing this to avoid overflow issues. + arrays: A list of arrays to be fused. Returns: - np.ndarray: An array of integers. + np.ndarray: The result of fusing `charges`. + """ + if len(arrays) == 1: + return arrays[0] + fused_arrays = arrays[0] + for n in range(1, len(arrays)): + fused_arrays = np.ravel(np.add.outer(fused_arrays, arrays[n])) + return fused_arrays + + +def _check_flows(flows: List[int]) -> None: + return + + +def _find_best_partition(dims: Iterable[int]) -> int: + """ + """ - converter_table = np.expand_dims( - np.flip(np.append(1, np.cumprod(np.flip(dims[1::])))), 0) - tmp = table * converter_table - integers = np.sum(tmp, axis=1) - return integers + if len(dims) == 1: + raise ValueError( + 'expecting `dims` with a length of at least 2, got `len(dims ) =1`') + diffs = [ + np.abs(np.prod(dims[0:n]) - np.prod(dims[n::])) + for n in range(1, len(dims)) + ] + min_inds = np.nonzero(diffs == np.min(diffs))[0] + if len(min_inds) > 1: + right_dims = [np.prod(dims[min_ind + 1:]) for min_ind in min_inds] + min_ind = min_inds[np.argmax(right_dims)] + else: + min_ind = min_inds[0] + return min_ind + 1 -def compute_fused_charge_degeneracies(charges: List[np.ndarray], - flows: List[Union[bool, int]]) -> Dict: +def compute_fused_charge_degeneracies(charges: List[BaseCharge], + flows: List[bool] + ) -> Tuple[BaseCharge, np.ndarray]: """ For a list of charges, compute all possible fused charges resulting - from fusing `charges`, together with their respective degeneracyn + from fusing `charges`, together with their respective degeneracies Args: charges: List of np.ndarray of int, one for each leg of the underlying tensor. Each np.ndarray `charges[leg]` @@ -92,50 +124,39 @@ def compute_fused_charge_degeneracies(charges: List[np.ndarray], of the charges on each leg. `1` is inflowing, `-1` is outflowing charge. Returns: - dict: Mapping fused charges (int) to degeneracies (int) + Union[BaseCharge, BaseCharge]: The unique fused charges. + np.ndarray of integers: The degeneracies of each unqiue fused charge. """ if len(charges) == 1: - return np.unique(charges[0], return_counts=True) + return (charges[0] * flows[0]).unique(return_counts=True) # get unique charges and their degeneracies on the first leg. # We are fusing from "left" to "right". - accumulated_charges, accumulated_degeneracies = np.unique( - charges[0], return_counts=True) - #multiply the flow into the charges of first leg - accumulated_charges *= flows[0] + accumulated_charges, accumulated_degeneracies = (charges[0] * + flows[0]).unique( + return_counts=True) for n in range(1, len(charges)): - #list of unique charges and list of their degeneracies - #on the next unfused leg of the tensor - leg_charges, leg_degeneracies = np.unique(charges[n], return_counts=True) - - #fuse the unique charges - #Note: entries in `fused_charges` are not unique anymore. - #flow1 = 1 because the flow of leg 0 has already been - #mulitplied above - fused_charges = fuse_charge_pair( - q1=accumulated_charges, flow1=1, q2=leg_charges, flow2=flows[n]) - #compute the degeneracies of `fused_charges` charges - #`fused_degeneracies` is a list of degeneracies such that - # `fused_degeneracies[n]` is the degeneracy of of - # charge `c = fused_charges[n]`. + leg_charges, leg_degeneracies = charges[n].unique(return_counts=True) + fused_charges = accumulated_charges + leg_charges * flows[n] fused_degeneracies = fuse_degeneracies(accumulated_degeneracies, leg_degeneracies) - #compute the new degeneracies resulting from fusing - #`accumulated_charges` and `leg_charges_2` - accumulated_charges = np.unique(fused_charges) + accumulated_charges = fused_charges.unique() accumulated_degeneracies = np.empty( - len(accumulated_charges), dtype=np.int64) + len(accumulated_charges), dtype=np.uint32) + for n in range(len(accumulated_charges)): - accumulated_degeneracies[n] = np.sum( - fused_degeneracies[fused_charges == accumulated_charges[n]]) + accumulated_degeneracies[n] = np.sum(fused_degeneracies[ + fused_charges.charge_labels == accumulated_charges.charge_labels[n]]) + return accumulated_charges, accumulated_degeneracies -def compute_num_nonzero(charges: List[np.ndarray], - flows: List[Union[bool, int]]) -> int: +def compute_unique_fused_charges(charges: List[BaseCharge], + flows: List[Union[bool, int]] + ) -> Tuple[BaseCharge, np.ndarray]: """ - Compute the number of non-zero elements, given the meta-data of - a symmetric tensor. + For a list of charges, compute all possible fused charges resulting + from fusing `charges`. Args: charges: List of np.ndarray of int, one for each leg of the underlying tensor. Each np.ndarray `charges[leg]` @@ -146,741 +167,266 @@ def compute_num_nonzero(charges: List[np.ndarray], of the charges on each leg. `1` is inflowing, `-1` is outflowing charge. Returns: - int: The number of non-zero elements. + Union[BaseCharge, ChargeCollection]: The unique fused charges. + np.ndarray of integers: The degeneracies of each unqiue fused charge. """ - accumulated_charges, accumulated_degeneracies = compute_fused_charge_degeneracies( - charges, flows) - if len(np.nonzero(accumulated_charges == 0)[0]) == 0: - raise ValueError( - "given leg-charges `charges` and flows `flows` are incompatible " - "with a symmetric tensor") - return accumulated_degeneracies[accumulated_charges == 0][0] + if len(charges) == 1: + return (charges[0] * flows[0]).unique() + + accumulated_charges = (charges[0] * flows[0]).unique() + for n in range(1, len(charges)): + leg_charges = charges[n].unique() + fused_charges = accumulated_charges + leg_charges * flows[n] + accumulated_charges = fused_charges.unique() + return accumulated_charges -def compute_nonzero_block_shapes(charges: List[np.ndarray], - flows: List[Union[bool, int]]) -> Dict: +def compute_num_nonzero(charges: List[BaseCharge], flows: List[bool]) -> int: """ - Compute the blocks and their respective shapes of a symmetric tensor, - given its meta-data. + Compute the number of non-zero elements, given the meta-data of + a symmetric tensor. Args: - charges: List of np.ndarray, one for each leg. - Each np.ndarray `charges[leg]` is of shape `(D[leg],)`. + charges: List of np.ndarray of int, one for each leg of the + underlying tensor. Each np.ndarray `charges[leg]` + is of shape `(D[leg],)`. The bond dimension `D[leg]` can vary on each leg. flows: A list of integers, one for each leg, with values `1` or `-1`, denoting the flow direction of the charges on each leg. `1` is inflowing, `-1` is outflowing charge. Returns: - dict: Dictionary mapping a tuple of charges to a shape tuple. - Each element corresponds to a non-zero valued block of the tensor. - """ - #FIXME: this routine is slow - _check_flows(flows) - degeneracies = [] - unique_charges = [] - rank = len(charges) - #find the unique quantum numbers and their degeneracy on each leg - for leg in range(rank): - c, d = np.unique(charges[leg], return_counts=True) - unique_charges.append(c) - degeneracies.append(dict(zip(c, d))) - - #find all possible combination of leg charges c0, c1, ... - #(with one charge per leg 0, 1, ...) - #such that sum([flows[0] * c0, flows[1] * c1, ...]) = 0 - charge_combinations = list( - itertools.product(*[ - unique_charges[leg] * flows[leg] - for leg in range(len(unique_charges)) - ])) - net_charges = np.array([np.sum(c) for c in charge_combinations]) - zero_idxs = np.nonzero(net_charges == 0)[0] - charge_shape_dict = {} - for idx in zero_idxs: - c = charge_combinations[idx] - shapes = [degeneracies[leg][flows[leg] * c[leg]] for leg in range(rank)] - charge_shape_dict[c] = shapes - return charge_shape_dict - - -def find_diagonal_sparse_blocks_version_1( - data: np.ndarray, - row_charges: List[Union[List, np.ndarray]], - column_charges: List[Union[List, np.ndarray]], - row_flows: List[Union[bool, int]], - column_flows: List[Union[bool, int]], - return_data: Optional[bool] = True) -> Dict: - """ - Deprecated - - This version is slow for matrices with shape[0] >> shape[1], but fast otherwise. - - Given the meta data and underlying data of a symmetric matrix, compute - all diagonal blocks and return them in a dict. - `row_charges` and `column_charges` are lists of np.ndarray. The tensor - is viewed as a matrix with rows given by fusing `row_charges` and - columns given by fusing `column_charges`. Note that `column_charges` - are never explicitly fused (`row_charges` are). - Args: - data: An np.ndarray of the data. The number of elements in `data` - has to match the number of non-zero elements defined by `charges` - and `flows` - row_charges: List of np.ndarray, one for each leg of the row-indices. - Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - column_charges: List of np.ndarray, one for each leg of the column-indices. - Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - row_flows: A list of integers, one for each entry in `row_charges`. - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - column_flows: A list of integers, one for each entry in `column_charges`. - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the sparse locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - - Returns: - dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray - or a python list of locations and shapes, depending on the value of `return_data`. - """ - flows = row_flows.copy() - flows.extend(column_flows) - _check_flows(flows) - if len(flows) != (len(row_charges) + len(column_charges)): - raise ValueError( - "`len(flows)` is different from `len(row_charges) + len(column_charges)`" - ) - - #since we are using row-major we have to fuse the row charges anyway. - fused_row_charges = fuse_charges(row_charges, row_flows) - #get the unique row-charges - unique_row_charges, row_dims = np.unique( - fused_row_charges, return_counts=True) - - #get the unique column-charges - #we only care about their degeneracies, not their order; that's much faster - #to compute since we don't have to fuse all charges explicitly - unique_column_charges, column_dims = compute_fused_charge_degeneracies( - column_charges, column_flows) - #get the charges common to rows and columns (only those matter) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - - #convenience container for storing the degeneracies of each - #row and column charge - row_degeneracies = dict(zip(unique_row_charges, row_dims)) - column_degeneracies = dict(zip(unique_column_charges, column_dims)) - - # we only care about charges common to row and columns - mask = np.isin(fused_row_charges, common_charges) - relevant_row_charges = fused_row_charges[mask] - #some numpy magic to get the index locations of the blocks - #we generate a vector of `len(relevant_row_charges) which, - #for each charge `c` in `relevant_row_charges` holds the - #column-degeneracy of charge `c` - degeneracy_vector = np.empty(len(relevant_row_charges), dtype=np.int64) - #for each charge `c` in `common_charges` we generate a boolean mask - #for indexing the positions where `relevant_column_charges` has a value of `c`. - masks = {} - for c in common_charges: - mask = relevant_row_charges == c - masks[c] = mask - degeneracy_vector[mask] = column_degeneracies[-c] - - # the result of the cumulative sum is a vector containing - # the stop positions of the non-zero values of each row - # within the data vector. - # E.g. for `relevant_row_charges` = [0,1,0,0,3], and - # column_degeneracies[0] = 10 - # column_degeneracies[1] = 20 - # column_degeneracies[3] = 30 - # we have - # `stop_positions` = [10, 10+20, 10+20+10, 10+20+10+10, 10+20+10+10+30] - # The starting positions of consecutive elements (in row-major order) in - # each row with charge `c=0` within the data vector are then simply obtained using - # masks[0] = [True, False, True, True, False] - # and `stop_positions[masks[0]] - column_degeneracies[0]` - stop_positions = np.cumsum(degeneracy_vector) - start_positions = stop_positions - degeneracy_vector - blocks = {} - - for c in common_charges: - #numpy broadcasting is substantially faster than kron! - a = np.expand_dims(start_positions[masks[c]], 1) - b = np.expand_dims(np.arange(column_degeneracies[-c]), 0) - if not return_data: - blocks[c] = [ - np.reshape(a + b, row_degeneracies[c] * column_degeneracies[-c]), - (row_degeneracies[c], column_degeneracies[-c]) - ] - else: - blocks[c] = np.reshape( - data[np.reshape(a + b, - row_degeneracies[c] * column_degeneracies[-c])], - (row_degeneracies[c], column_degeneracies[-c])) - return blocks - - -def find_diagonal_sparse_blocks(data: np.ndarray, - row_charges: List[Union[List, np.ndarray]], - column_charges: List[Union[List, np.ndarray]], - row_flows: List[Union[bool, int]], - column_flows: List[Union[bool, int]], - return_data: Optional[bool] = True) -> Dict: - """ - Given the meta data and underlying data of a symmetric matrix, compute - all diagonal blocks and return them in a dict. - `row_charges` and `column_charges` are lists of np.ndarray. The tensor - is viewed as a matrix with rows given by fusing `row_charges` and - columns given by fusing `column_charges`. Note that `column_charges` - are never explicitly fused (`row_charges` are). - Args: - data: An np.ndarray of the data. The number of elements in `data` - has to match the number of non-zero elements defined by `charges` - and `flows` - row_charges: List of np.ndarray, one for each leg of the row-indices. - Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - column_charges: List of np.ndarray, one for each leg of the column-indices. - Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - row_flows: A list of integers, one for each entry in `row_charges`. - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - column_flows: A list of integers, one for each entry in `column_charges`. - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the sparse locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - - Returns: - dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray - or a python list of locations and shapes, depending on the value of `return_data`. + int: The number of non-zero elements. """ - flows = row_flows.copy() - flows.extend(column_flows) - _check_flows(flows) - if len(flows) != (len(row_charges) + len(column_charges)): + accumulated_charges, accumulated_degeneracies = compute_fused_charge_degeneracies( + charges, flows) + res = accumulated_charges == accumulated_charges.identity_charges + nz_inds = np.nonzero(res)[0] + if len(nz_inds) == 0: raise ValueError( - "`len(flows)` is different from `len(row_charges) + len(column_charges)`" - ) - - #get the unique column-charges - #we only care about their degeneracies, not their order; that's much faster - #to compute since we don't have to fuse all charges explicitly - unique_column_charges, column_dims = compute_fused_charge_degeneracies( - column_charges, column_flows) - #convenience container for storing the degeneracies of each - #column charge - column_degeneracies = dict(zip(unique_column_charges, column_dims)) - - if len(row_charges) > 1: - left_row_charges, right_row_charges = _find_best_partition( - row_charges, row_flows) - unique_left = np.unique(left_row_charges) - unique_right = np.unique(right_row_charges) - unique_row_charges = np.unique( - fuse_charges(charges=[unique_left, unique_right], flows=[1, 1])) - - #get the charges common to rows and columns (only those matter) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - - row_locations = find_sparse_positions( - left_charges=left_row_charges, - left_flow=1, - right_charges=right_row_charges, - right_flow=1, - target_charges=common_charges) - elif len(row_charges) == 1: - fused_row_charges = fuse_charges(row_charges, row_flows) - - #get the unique row-charges - unique_row_charges, row_dims = np.unique( - fused_row_charges, return_counts=True) - #get the charges common to rows and columns (only those matter) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - relevant_fused_row_charges = fused_row_charges[np.isin( - fused_row_charges, common_charges)] - row_locations = {} - for c in common_charges: - row_locations[c] = np.nonzero(relevant_fused_row_charges == c)[0] - else: - raise ValueError('Found an empty sequence for `row_charges`') - #some numpy magic to get the index locations of the blocks - degeneracy_vector = np.empty( - np.sum([len(v) for v in row_locations.values()]), dtype=np.int64) - #for each charge `c` in `common_charges` we generate a boolean mask - #for indexing the positions where `relevant_column_charges` has a value of `c`. - masks = {} - for c in common_charges: - degeneracy_vector[row_locations[c]] = column_degeneracies[-c] - - # the result of the cumulative sum is a vector containing - # the stop positions of the non-zero values of each row - # within the data vector. - # E.g. for `relevant_row_charges` = [0,1,0,0,3], and - # column_degeneracies[0] = 10 - # column_degeneracies[1] = 20 - # column_degeneracies[3] = 30 - # we have - # `stop_positions` = [10, 10+20, 10+20+10, 10+20+10+10, 10+20+10+10+30] - # The starting positions of consecutive elements (in row-major order) in - # each row with charge `c=0` within the data vector are then simply obtained using - # masks[0] = [True, False, True, True, False] - # and `stop_positions[masks[0]] - column_degeneracies[0]` - stop_positions = np.cumsum(degeneracy_vector) - start_positions = stop_positions - degeneracy_vector - blocks = {} - - for c in common_charges: - #numpy broadcasting is substantially faster than kron! - a = np.expand_dims(start_positions[np.sort(row_locations[c])], 1) - b = np.expand_dims(np.arange(column_degeneracies[-c]), 0) - inds = np.reshape(a + b, len(row_locations[c]) * column_degeneracies[-c]) - if not return_data: - blocks[c] = [inds, (len(row_locations[c]), column_degeneracies[-c])] - else: - blocks[c] = np.reshape(data[inds], - (len(row_locations[c]), column_degeneracies[-c])) - return blocks + "given leg-charges `charges` and flows `flows` are incompatible " + "with a symmetric tensor") + return np.squeeze(accumulated_degeneracies[nz_inds][0]) -def find_diagonal_sparse_blocks_version_0( - data: np.ndarray, - charges: List[np.ndarray], - flows: List[Union[bool, int]], - return_data: Optional[bool] = True) -> Dict: +def _find_diagonal_sparse_blocks(charges: List[BaseCharge], flows: np.ndarray, + partition: int + ) -> (np.ndarray, np.ndarray, np.ndarray): """ - Deprecated: this version is about 2 times slower (worst case) than the current used - implementation - Given the meta data and underlying data of a symmetric matrix, compute - all diagonal blocks and return them in a dict. - Args: - data: An np.ndarray of the data. The number of elements in `data` - has to match the number of non-zero elements defined by `charges` - and `flows` - charges: List of np.ndarray, one for each leg. - Each np.ndarray `charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - flows: A list of integers, one for each leg, - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - + Find the location of all non-trivial symmetry blocks from the data vector of + of SymTensor (when viewed as a matrix across some prescribed index + bi-partition). + Args: + charges (List[SymIndex]): list of SymIndex. + flows (np.ndarray): vector of bools describing index orientations. + partition_loc (int): location of tensor partition (i.e. such that the + tensor is viewed as a matrix between first partition_loc charges and + the remaining charges). Returns: - dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray - or a python list of locations and shapes, depending on the value of `return_data`. + block_maps (List[np.ndarray]): list of integer arrays, which each + containing the location of a symmetry block in the data vector. + block_qnums (np.ndarray): n-by-m array describing qauntum numbers of each + block, with 'n' the number of symmetries and 'm' the number of blocks. + block_dims (np.ndarray): 2-by-m array describing the dims each block, + with 'm' the number of blocks). """ - if len(charges) != 2: - raise ValueError("input has to be a two-dimensional symmetric matrix") - _check_flows(flows) - if len(flows) != len(charges): - raise ValueError("`len(flows)` is different from `len(charges)`") + num_inds = len(charges) + num_syms = charges[0].num_symmetries - #we multiply the flows into the charges - row_charges = flows[0] * charges[0] # a list of charges on each row - column_charges = flows[1] * charges[1] # a list of charges on each column + if (partition == 0) or (partition == num_inds): + # special cases (matrix of trivial height or width) + num_nonzero = compute_num_nonzero(charges, flows) + block_maps = [np.arange(0, num_nonzero, dtype=np.uint64).ravel()] + block_qnums = np.zeros([num_syms, 1], dtype=np.int16) + block_dims = np.array([[1], [num_nonzero]]) - #get the unique charges - unique_row_charges, row_dims = np.unique(row_charges, return_counts=True) - unique_column_charges, column_dims = np.unique( - column_charges, return_counts=True) - #get the charges common to rows and columns (only those matter) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - - #convenience container for storing the degeneracies of each - #row and column charge - row_degeneracies = dict(zip(unique_row_charges, row_dims)) - column_degeneracies = dict(zip(unique_column_charges, column_dims)) - - # we only care about charges common to row and columns - mask = np.isin(row_charges, common_charges) - relevant_row_charges = row_charges[mask] - - #some numpy magic to get the index locations of the blocks - #we generate a vector of `len(relevant_row_charges) which, - #for each charge `c` in `relevant_row_charges` holds the - #column-degeneracy of charge `c` - degeneracy_vector = np.empty(len(relevant_row_charges), dtype=np.int64) - #for each charge `c` in `common_charges` we generate a boolean mask - #for indexing the positions where `relevant_column_charges` has a value of `c`. - masks = {} - for c in common_charges: - mask = relevant_row_charges == c - masks[c] = mask - degeneracy_vector[mask] = column_degeneracies[-c] - - # the result of the cumulative sum is a vector containing - # the stop positions of the non-zero values of each row - # within the data vector. - # E.g. for `relevant_row_charges` = [0,1,0,0,3], and - # column_degeneracies[0] = 10 - # column_degeneracies[1] = 20 - # column_degeneracies[3] = 30 - # we have - # `stop_positions` = [10, 10+20, 10+20+10, 10+20+10+10, 10+20+10+10+30] - # The starting positions of consecutive elements (in row-major order) in - # each row with charge `c=0` within the data vector are then simply obtained using - # masks[0] = [True, False, True, True, False] - # and `stop_positions[masks[0]] - column_degeneracies[0]` - stop_positions = np.cumsum(degeneracy_vector) - blocks = {} - - for c in common_charges: - #numpy broadcasting is substantially faster than kron! - a = np.expand_dims(stop_positions[masks[c]] - column_degeneracies[-c], 0) - b = np.expand_dims(np.arange(column_degeneracies[-c]), 1) - if not return_data: - blocks[c] = [ - np.reshape(a + b, row_degeneracies[c] * column_degeneracies[-c]), - (row_degeneracies[c], column_degeneracies[-c]) - ] - else: - blocks[c] = np.reshape( - data[np.reshape(a + b, - row_degeneracies[c] * column_degeneracies[-c])], - (row_degeneracies[c], column_degeneracies[-c])) - return blocks - - -def find_diagonal_sparse_blocks_column_major( - data: np.ndarray, - charges: List[np.ndarray], - flows: List[Union[bool, int]], - return_data: Optional[bool] = True) -> Dict: - """ - Deprecated + if partition == len(flows): + block_dims = np.flipud(block_dims) - Given the meta data and underlying data of a symmetric matrix, compute - all diagonal blocks and return them in a dict, assuming column-major - ordering. - Args: - data: An np.ndarray of the data. The number of elements in `data` - has to match the number of non-zero elements defined by `charges` - and `flows` - charges: List of np.ndarray, one for each leg. - Each np.ndarray `charges[leg]` is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - flows: A list of integers, one for each leg, - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - - Returns: - dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray - or a python list of locations and shapes, depending on the value of `return_data`. - """ - if len(charges) != 2: - raise ValueError("input has to be a two-dimensional symmetric matrix") - _check_flows(flows) - if len(flows) != len(charges): - raise ValueError("`len(flows)` is different from `len(charges)`") + return block_maps, block_qnums, block_dims - #we multiply the flows into the charges - row_charges = flows[0] * charges[0] # a list of charges on each row - column_charges = flows[1] * charges[1] # a list of charges on each column + else: + unique_row_qnums, row_degen = compute_fused_charge_degeneracies( + charges[:partition], flows[:partition]) + + unique_col_qnums, col_degen = compute_fused_charge_degeneracies( + charges[partition:], np.logical_not(flows[partition:])) + + block_qnums, row_to_block, col_to_block = intersect( + unique_row_qnums.unique_charges, + unique_col_qnums.unique_charges, + axis=1, + return_indices=True) + num_blocks = block_qnums.shape[1] + if num_blocks == 0: + obj = charges[0].__new__(type(charges[0])) + obj.__init__( + np.zeros(0, dtype=np.int16), np.arange(0, dtype=np.int16), + charges[0].charge_types) + + return [], obj, [] - #get the unique charges - unique_row_charges, row_dims = np.unique(row_charges, return_counts=True) - unique_column_charges, column_dims = np.unique( - column_charges, return_counts=True) - #get the charges common to rows and columns (only those matter) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - - #convenience container for storing the degeneracies of each - #row and column charge - row_degeneracies = dict(zip(unique_row_charges, row_dims)) - column_degeneracies = dict(zip(unique_column_charges, column_dims)) - - # we only care about charges common to row and columns - mask = np.isin(column_charges, -common_charges) - relevant_column_charges = column_charges[mask] - - #some numpy magic to get the index locations of the blocks - #we generate a vector of `len(relevant_column_charges) which, - #for each charge `c` in `relevant_column_charges` holds the - #row-degeneracy of charge `c` - degeneracy_vector = np.empty(len(relevant_column_charges), dtype=np.int64) - #for each charge `c` in `common_charges` we generate a boolean mask - #for indexing the positions where `relevant_column_charges` has a value of `c`. - masks = {} - for c in common_charges: - mask = relevant_column_charges == -c - masks[c] = mask - degeneracy_vector[mask] = row_degeneracies[c] - - # the result of the cumulative sum is a vector containing - # the stop positions of the non-zero values of each column - # within the data vector. - # E.g. for `relevant_column_charges` = [0,1,0,0,3], and - # row_degeneracies[0] = 10 - # row_degeneracies[1] = 20 - # row_degeneracies[3] = 30 - # we have - # `stop_positions` = [10, 10+20, 10+20+10, 10+20+10+10, 10+20+10+10+30] - # The starting positions of consecutive elements (in column-major order) in - # each column with charge `c=0` within the data vector are then simply obtained using - # masks[0] = [True, False, True, True, False] - # and `stop_positions[masks[0]] - row_degeneracies[0]` - stop_positions = np.cumsum(degeneracy_vector) - blocks = {} - - for c in common_charges: - #numpy broadcasting is substantially faster than kron! - a = np.expand_dims(stop_positions[masks[c]] - row_degeneracies[c], 0) - b = np.expand_dims(np.arange(row_degeneracies[c]), 1) - if not return_data: - blocks[c] = [ - np.reshape(a + b, row_degeneracies[c] * column_degeneracies[-c]), - (row_degeneracies[c], column_degeneracies[-c]) - ] else: - blocks[c] = np.reshape( - data[np.reshape(a + b, - row_degeneracies[c] * column_degeneracies[-c])], - (row_degeneracies[c], column_degeneracies[-c])) - return blocks - - -def find_dense_positions(left_charges: np.ndarray, left_flow: int, - right_charges: np.ndarray, right_flow: int, - target_charge: int) -> Dict: - """ - Find the dense locations of elements (i.e. the index-values within the DENSE tensor) - in the vector `fused_charges` (resulting from fusing np.ndarrays - `left_charges` and `right_charges`) that have a value of `target_charge`. - For example, given - ``` - left_charges = [-2,0,1,0,0] - right_charges = [-1,0,2,1] - target_charge = 0 - fused_charges = fuse_charges([left_charges, right_charges],[1,1]) - print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] - ``` - we want to find the all different blocks - that fuse to `target_charge=0`, i.e. where `fused_charges==0`, - together with their corresponding index-values of the data in the dense array. - `find_dense_blocks` returns a dict mapping tuples `(left_charge, right_charge)` - to an array of integers. - For the above example, we get: - * for `left_charge` = -2 and `right_charge` = 2 we get an array [2]. Thus, `fused_charges[2]` - was obtained from fusing -2 and 2. - * for `left_charge` = 0 and `right_charge` = 0 we get an array [5, 13, 17]. Thus, - `fused_charges[5,13,17]` were obtained from fusing 0 and 0. - * for `left_charge` = 1 and `right_charge` = -1 we get an array [8]. Thus, `fused_charges[8]` - was obtained from fusing 1 and -1. - Args: - left_charges: An np.ndarray of integer charges. - left_flow: The flow direction of the left charges. - right_charges: An np.ndarray of integer charges. - right_flow: The flow direction of the right charges. - target_charge: The target charge. - Returns: - dict: Mapping tuples of integers to np.ndarray of integers. - """ - _check_flows([left_flow, right_flow]) - unique_left = np.unique(left_charges) - unique_right = np.unique(right_charges) - fused = fuse_charges([unique_left, unique_right], [left_flow, right_flow]) - left_inds, right_inds = unfuse( - np.nonzero(fused == target_charge)[0], len(unique_left), - len(unique_right)) - left_c = unique_left[left_inds] - right_c = unique_right[right_inds] - len_right_charges = len(right_charges) - linear_positions = {} - for left_charge, right_charge in zip(left_c, right_c): - left_positions = np.nonzero(left_charges == left_charge)[0] - left_offsets = np.expand_dims(left_positions * len_right_charges, 1) - right_offsets = np.expand_dims( - np.nonzero(right_charges == right_charge)[0], 0) - linear_positions[(left_charge, right_charge)] = np.reshape( - left_offsets + right_offsets, - left_offsets.shape[0] * right_offsets.shape[1]) - return linear_positions - - -def find_sparse_positions(left_charges: np.ndarray, left_flow: int, - right_charges: np.ndarray, right_flow: int, - target_charges: Union[List[int], np.ndarray]) -> Dict: - """ - Find the sparse locations of elements (i.e. the index-values within the SPARSE tensor) - in the vector `fused_charges` (resulting from fusing np.ndarrays - `left_charges` and `right_charges`) that have a value of `target_charge`, - assuming that all elements different from `target_charges` are `0`. - For example, given - ``` - left_charges = [-2,0,1,0,0] - right_charges = [-1,0,2,1] - target_charges = [0,1] - fused_charges = fuse_charges([left_charges, right_charges],[1,1]) - print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] - ``` 0 1 2 3 4 5 6 7 8 - we want to find the all different blocks - that fuse to `target_charges=[0,1]`, i.e. where `fused_charges==0` or `1`, - together with their corresponding sparse index-values of the data in the sparse array, - assuming that all elements in `fused_charges` different from `target_charges` are 0. - - `find_sparse_blocks` returns a dict mapping integers `target_charge` - to an array of integers denoting the sparse locations of elements within - `fused_charges`. - For the above example, we get: - * `target_charge=0`: [0,1,3,5,7] - * `target_charge=1`: [2,4,6,8] - Args: - left_charges: An np.ndarray of integer charges. - left_flow: The flow direction of the left charges. - right_charges: An np.ndarray of integer charges. - right_flow: The flow direction of the right charges. - target_charge: The target charge. - Returns: - dict: Mapping integers to np.ndarray of integers. - """ - #FIXME: this is probably still not optimal - - _check_flows([left_flow, right_flow]) - target_charges = np.unique(target_charges) - unique_left = np.unique(left_charges) - unique_right = np.unique(right_charges) - fused = fuse_charges([unique_left, unique_right], [left_flow, right_flow]) - - #compute all unique charges that can add up to - #target_charges - left_inds, right_inds = [], [] - for target_charge in target_charges: - li, ri = unfuse( - np.nonzero(fused == target_charge)[0], len(unique_left), - len(unique_right)) - left_inds.append(li) - right_inds.append(ri) - - #compute the relevant unique left and right charges - unique_left_charges = unique_left[np.unique(np.concatenate(left_inds))] - unique_right_charges = unique_right[np.unique(np.concatenate(right_inds))] - - relevant_left_charges = left_charges[np.isin(left_charges, - unique_left_charges)] - relevant_right_charges = right_charges[np.isin(right_charges, - unique_right_charges)] - unique_right_charges, right_dims = np.unique( - relevant_right_charges, return_counts=True) - right_degeneracies = dict(zip(unique_right_charges, right_dims)) - degeneracy_vector = np.empty(len(relevant_left_charges), dtype=np.int64) - total_row_degeneracies = {} - right_indices = {} - for left_charge in unique_left_charges: - total_degeneracy = np.sum(right_dims[np.isin( - left_flow * left_charge + right_flow * unique_right_charges, - target_charges)]) - tmp_relevant_right_charges = relevant_right_charges[np.isin( - relevant_right_charges, - (target_charges - left_flow * left_charge) * right_flow)] - - for target_charge in target_charges: - right_indices[(left_charge, target_charge)] = np.nonzero( - tmp_relevant_right_charges == - (target_charge - left_flow * left_charge) * right_flow)[0] - - degeneracy_vector[relevant_left_charges == left_charge] = total_degeneracy - - stop_positions = np.cumsum(degeneracy_vector) - start_positions = stop_positions - degeneracy_vector - blocks = {t: [] for t in target_charges} - for left_charge in unique_left_charges: - a = np.expand_dims(start_positions[relevant_left_charges == left_charge], 0) - for target_charge in target_charges: - ri = right_indices[(left_charge, target_charge)] - if len(ri) != 0: - b = np.expand_dims(ri, 1) - tmp = a + b - blocks[target_charge].append(np.reshape(tmp, np.prod(tmp.shape))) - out = {} - for target_charge in target_charges: - out[target_charge] = np.concatenate(blocks[target_charge]) - return out - - -def compute_dense_to_sparse_table(charges: List[np.ndarray], - flows: List[Union[bool, int]], - target_charge: int) -> int: + # calculate number of non-zero elements in each row of the matrix + row_ind = reduce_charges(charges[:partition], flows[:partition], + block_qnums) + row_num_nz = col_degen[col_to_block[row_ind.charge_labels]] + cumulate_num_nz = np.insert(np.cumsum(row_num_nz[0:-1]), 0, + 0).astype(np.uint32) + + # calculate mappings for the position in datavector of each block + if num_blocks < 15: + # faster method for small number of blocks + row_locs = np.concatenate([ + (row_ind.charge_labels == n) for n in range(num_blocks) + ]).reshape(num_blocks, row_ind.dim) + else: + # faster method for large number of blocks + row_locs = np.zeros([num_blocks, row_ind.dim], dtype=bool) + row_locs[row_ind + .charge_labels, np.arange(row_ind.dim)] = np.ones( + row_ind.dim, dtype=bool) + + # block_dims = np.array([row_degen[row_to_block],col_degen[col_to_block]], dtype=np.uint32) + block_dims = np.array( + [[row_degen[row_to_block[n]], col_degen[col_to_block[n]]] + for n in range(num_blocks)], + dtype=np.uint32).T + block_maps = [(cumulate_num_nz[row_locs[n, :]][:, None] + + np.arange(block_dims[1, n])[None, :]).ravel() + for n in range(num_blocks)] + obj = charges[0].__new__(type(charges[0])) + obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16), + charges[0].charge_types) + + return block_maps, obj, block_dims + + +def _find_transposed_diagonal_sparse_blocks( + charges: List[BaseCharge], + flows: np.ndarray, + tr_partition: int, + order: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Compute a table mapping multi-index positions to the linear positions - within the sparse data container. + Find the location of all non-trivial symmetry blocks from the data vector of + of SymTensor after transposition (and then viewed as a matrix across some + prescribed index bi-tr_partition). Produces and equivalent result to + retrieve_blocks acting on a transposed SymTensor, but is much faster. Args: - charges: List of np.ndarray of int, one for each leg of the - underlying tensor. Each np.ndarray `charges[leg]` - is of shape `(D[leg],)`. - The bond dimension `D[leg]` can vary on each leg. - flows: A list of integers, one for each leg, - with values `1` or `-1`, denoting the flow direction - of the charges on each leg. `1` is inflowing, `-1` is outflowing - charge. - target_charge: The total target charge of the blocks to be calculated. + charges (List[SymIndex]): list of SymIndex. + flows (np.ndarray): vector of bools describing index orientations. + tr_partition (int): location of tensor partition (i.e. such that the + tensor is viewed as a matrix between first partition charges and + the remaining charges). + order (np.ndarray): order with which to permute the tensor axes. Returns: - np.ndarray: An (N, r) np.ndarray of dtype np.int16, - with `N` the number of non-zero elements, and `r` - the rank of the tensor. + block_maps (List[np.ndarray]): list of integer arrays, which each + containing the location of a symmetry block in the data vector. + block_qnums (np.ndarray): n-by-m array describing qauntum numbers of each + block, with 'n' the number of symmetries and 'm' the number of blocks. + block_dims (np.ndarray): 2-by-m array describing the dims each block, + with 'm' the number of blocks). """ - #find the best partition (the one where left and right dimensions are - #closest - dims = np.asarray([len(c) for c in charges]) - - # #all legs smaller or equal to `min_ind` are on the left side - # #of the partition. All others are on the right side. - # min_ind = np.argmin([ - # np.abs(np.prod(dims[0:n]) - np.prod(dims[n::])) - # for n in range(1, len(charges)) - # ]) - # fused_left_charges = fuse_charges(charges[0:min_ind + 1], - # flows[0:min_ind + 1]) - # fused_right_charges = fuse_charges(charges[min_ind + 1::], - # flows[min_ind + 1::]) - - fused_charges = fuse_charges(charges, flows) - nz_indices = np.nonzero(fused_charges == target_charge)[0] - - if len(nz_indices) == 0: - raise ValueError( - "`charges` do not add up to a total charge {}".format(target_charge)) + flows = np.asarray(flows) + if np.array_equal(order, None) or (np.array_equal( + np.array(order), np.arange(len(charges)))): + # no transpose order + return _find_diagonal_sparse_blocks(charges, flows, tr_partition) - index_locations = [] - for n in reversed(range(len(charges))): - nz_indices, right_indices = unfuse(nz_indices, np.prod(dims[0:n]), dims[n]) - index_locations.insert(0, right_indices) - return index_locations + else: + # non-trivial transposition is required + num_inds = len(charges) + tensor_dims = np.array([charges[n].dim for n in range(num_inds)], dtype=int) + strides = np.append(np.flip(np.cumprod(np.flip(tensor_dims[1:]))), 1) + + # define properties of new tensor resulting from transposition + new_strides = strides[order] + new_row_charges = [charges[n] for n in order[:tr_partition]] + new_col_charges = [charges[n] for n in order[tr_partition:]] + new_row_flows = flows[order[:tr_partition]] + new_col_flows = flows[order[tr_partition:]] + + # compute qnums of row/cols in transposed tensor + unique_row_qnums, new_row_degen = compute_fused_charge_degeneracies( + new_row_charges, new_row_flows) + + # unique_row_qnums, new_row_degen = compute_qnum_degen( + # new_row_charges, new_row_flows) + unique_col_qnums, new_col_degen = compute_fused_charge_degeneracies( + new_col_charges, np.logical_not(new_col_flows)) + # unique_col_qnums, new_col_degen = compute_qnum_degen( + # new_col_charges, np.logical_not(new_col_flows)) + block_qnums, new_row_map, new_col_map = intersect( + unique_row_qnums.unique_charges, + unique_col_qnums.unique_charges, + axis=1, + return_indices=True) + block_dims = np.array( + [new_row_degen[new_row_map], new_col_degen[new_col_map]], + dtype=np.uint32) + num_blocks = len(new_row_map) + row_ind, row_locs = reduce_charges( + new_row_charges, + new_row_flows, + block_qnums, + return_locations=True, + strides=new_strides[:tr_partition]) + + col_ind, col_locs = reduce_charges( + new_col_charges, + np.logical_not(new_col_flows), + block_qnums, + return_locations=True, + strides=new_strides[tr_partition:]) + # compute qnums of row/cols in original tensor + orig_partition = _find_best_partition(tensor_dims) + orig_width = np.prod(tensor_dims[orig_partition:]) + + orig_unique_row_qnums = compute_unique_fused_charges( + charges[:orig_partition], flows[:orig_partition]) + orig_unique_col_qnums, orig_col_degen = compute_fused_charge_degeneracies( + charges[orig_partition:], np.logical_not(flows[orig_partition:])) + orig_block_qnums, row_map, col_map = intersect( + orig_unique_row_qnums.unique_charges, + orig_unique_col_qnums.unique_charges, + axis=1, + return_indices=True) + orig_num_blocks = orig_block_qnums.shape[1] + + orig_row_ind = fuse_charges(charges[:orig_partition], + flows[:orig_partition]) + orig_col_ind = fuse_charges(charges[orig_partition:], + np.logical_not(flows[orig_partition:])) + + # compute row degens (i.e. number of non-zero elements per row) + inv_row_map = -np.ones( + orig_unique_row_qnums.unique_charges.shape[1], dtype=np.int16) + for n in range(len(row_map)): + inv_row_map[row_map[n]] = n + + all_degens = np.append(orig_col_degen[col_map], + 0)[inv_row_map[orig_row_ind.charge_labels]] + all_cumul_degens = np.cumsum(np.insert(all_degens[:-1], 0, + 0)).astype(np.uint32) + # generate vector which translates from dense row position to sparse row position + dense_to_sparse = np.zeros(orig_width, dtype=np.uint32) + for n in range(orig_num_blocks): + dense_to_sparse[orig_col_ind.charge_labels == col_map[n]] = np.arange( + orig_col_degen[col_map[n]], dtype=np.uint32) + + block_maps = [0] * num_blocks + for n in range(num_blocks): + orig_row_posL, orig_col_posL = np.divmod( + row_locs[row_ind.charge_labels == n], orig_width) + orig_row_posR, orig_col_posR = np.divmod( + col_locs[col_ind.charge_labels == n], orig_width) + block_maps[n] = ( + all_cumul_degens[np.add.outer(orig_row_posL, orig_row_posR)] + + dense_to_sparse[np.add.outer(orig_col_posL, orig_col_posR)]).ravel() + obj = charges[0].__new__(type(charges[0])) + obj.__init__(block_qnums, np.arange(block_qnums.shape[1], dtype=np.int16), + charges[0].charge_types) + + return block_maps, obj, block_dims class BlockSparseTensor: @@ -906,6 +452,9 @@ class BlockSparseTensor: The tensor data is stored in self.data, a 1d np.ndarray. """ + def copy(self): + return BlockSparseTensor(self.data.copy(), [i.copy() for i in self.indices]) + def __init__(self, data: np.ndarray, indices: List[Index]) -> None: """ Args: @@ -919,12 +468,25 @@ def __init__(self, data: np.ndarray, indices: List[Index]) -> None: num_non_zero_elements = compute_num_nonzero(self.charges, self.flows) if num_non_zero_elements != len(data.flat): - raise ValueError("number of tensor elements defined " + raise ValueError("number of tensor elements {} defined " "by `charges` is different from" - " len(data)={}".format(len(data.flat))) + " len(data)={}".format(num_non_zero_elements, + len(data.flat))) self.data = np.asarray(data.flat) #do not copy data + def todense(self) -> np.ndarray: + """ + Map the sparse tensor to dense storage. + + """ + out = np.asarray(np.zeros(self.dense_shape, dtype=self.dtype).flat) + charges = self.charges + out[np.nonzero( + fuse_charges(charges, self.flows) == charges[0].identity_charges) + [0]] = self.data + return np.reshape(out, self.dense_shape) + @classmethod def randn(cls, indices: List[Index], dtype: Optional[Type[np.number]] = None) -> "BlockSparseTensor": @@ -943,6 +505,42 @@ def randn(cls, indices: List[Index], data = backend.randn((num_non_zero_elements,), dtype=dtype) return cls(data=data, indices=indices) + @classmethod + def ones(cls, indices: List[Index], + dtype: Optional[Type[np.number]] = None) -> "BlockSparseTensor": + """ + Initialize a symmetric tensor with ones. + Args: + indices: List of `Index` objecst, one for each leg. + dtype: An optional numpy dtype. The dtype of the tensor + Returns: + BlockSparseTensor + """ + charges = [i.charges for i in indices] + flows = [i.flow for i in indices] + num_non_zero_elements = compute_num_nonzero(charges, flows) + backend = backend_factory.get_backend('numpy') + data = backend.ones((num_non_zero_elements,), dtype=dtype) + return cls(data=data, indices=indices) + + @classmethod + def zeros(cls, indices: List[Index], + dtype: Optional[Type[np.number]] = None) -> "BlockSparseTensor": + """ + Initialize a symmetric tensor with zeros. + Args: + indices: List of `Index` objecst, one for each leg. + dtype: An optional numpy dtype. The dtype of the tensor + Returns: + BlockSparseTensor + """ + charges = [i.charges for i in indices] + flows = [i.flow for i in indices] + num_non_zero_elements = compute_num_nonzero(charges, flows) + backend = backend_factory.get_backend('numpy') + data = backend.zeros((num_non_zero_elements,), dtype=dtype) + return cls(data=data, indices=indices) + @classmethod def random(cls, indices: List[Index], dtype: Optional[Type[np.number]] = None) -> "BlockSparseTensor": @@ -956,8 +554,10 @@ def random(cls, indices: List[Index], """ charges = [i.charges for i in indices] flows = [i.flow for i in indices] + num_non_zero_elements = compute_num_nonzero(charges, flows) - dtype = dtype if dtype is not None else self.np.float64 + + dtype = dtype if dtype is not None else np.float64 def init_random(): if ((np.dtype(dtype) is np.dtype(np.complex128)) or @@ -1003,25 +603,65 @@ def flows(self): def charges(self): return [i.charges for i in self.indices] - def transpose(self, order): + def transpose( + self, + order: Union[List[int], np.ndarray], + ) -> "BlockSparseTensor": """ - Transpose the tensor into the new order `order` + Transpose the tensor in place into the new order `order`. + Args: + order: The new order of indices. + Returns: + BlockSparseTensor: The transposed tensor. """ - raise NotImplementedError('transpose is not implemented!!') + if len(order) != self.rank: + raise ValueError( + "`len(order)={}` is different form `self.rank={}`".format( + len(order), self.rank)) + + #check for trivial permutation + if np.all(order == np.arange(len(order))): + return self + flat_indices, flat_charges, flat_flows, _, flat_order, _ = flatten_meta_data( + self.indices, order, 0) + tr_partition = _find_best_partition( + [len(flat_charges[n]) for n in flat_order]) + + tr_sparse_blocks, tr_charges, tr_shapes = _find_transposed_diagonal_sparse_blocks( + flat_charges, flat_flows, tr_partition, flat_order) + + sparse_blocks, charges, shapes = _find_diagonal_sparse_blocks( + [flat_charges[n] for n in flat_order], + [flat_flows[n] for n in flat_order], tr_partition) + data = np.empty(len(self.data), dtype=self.dtype) + for n in range(len(sparse_blocks)): + sparse_block = sparse_blocks[n] + ind = np.nonzero(tr_charges == charges[n])[0][0] + permutation = tr_sparse_blocks[ind] + data[sparse_block] = self.data[permutation] + self.indices = [self.indices[o] for o in order] + self.data = data + return self def reset_shape(self) -> None: """ Bring the tensor back into its elementary shape. """ + self.indices = self.get_elementary_indices() + + def get_elementary_indices(self) -> List: + """ + Compute the elementary indices of the array. + """ elementary_indices = [] for i in self.indices: elementary_indices.extend(i.get_elementary_indices()) - self.indices = elementary_indices + return elementary_indices def reshape(self, shape: Union[Iterable[Index], Iterable[int]]) -> None: """ - Reshape `tensor` into `shape` in place. + Reshape `tensor` into `shape. `BlockSparseTensor.reshape` works essentially the same as the dense version, with the notable exception that the tensor can only be reshaped into a form compatible with its elementary indices. @@ -1055,155 +695,56 @@ def reshape(self, shape: Union[Iterable[Index], Iterable[int]]) -> None: Returns: BlockSparseTensor: A new tensor reshaped into `shape` """ - dense_shape = [] + new_shape = [] for s in shape: if isinstance(s, Index): - dense_shape.append(s.dimension) + new_shape.append(s.dimension) else: - dense_shape.append(s) + new_shape.append(s) # a few simple checks - if np.prod(dense_shape) != np.prod(self.dense_shape): + if np.prod(new_shape) != np.prod(self.dense_shape): raise ValueError("A tensor with {} elements cannot be " "reshaped into a tensor with {} elements".format( - np.prod(self.shape), np.prod(dense_shape))) + np.prod(self.shape), np.prod(self.dense_shape))) #keep a copy of the old indices for the case where reshaping fails #FIXME: this is pretty hacky! - index_copy = [i.copy() for i in self.indices] + indices = [i.copy() for i in self.indices] + flat_indices = [] + for i in indices: + flat_indices.extend(i.get_elementary_indices()) def raise_error(): #if this error is raised then `shape` is incompatible #with the elementary indices. We then reset the shape #to what is was before the call to `reshape`. - self.indices = index_copy - elementary_indices = [] - for i in self.indices: - elementary_indices.extend(i.get_elementary_indices()) + # self.indices = index_copy raise ValueError("The shape {} is incompatible with the " "elementary shape {} of the tensor.".format( - dense_shape, - tuple([e.dimension for e in elementary_indices]))) - - self.reset_shape() #bring tensor back into its elementary shape - for n in range(len(dense_shape)): - if dense_shape[n] > self.dense_shape[n]: - while dense_shape[n] > self.dense_shape[n]: - #fuse indices - i1, i2 = self.indices.pop(n), self.indices.pop(n) + new_shape, + tuple([e.dimension for e in flat_indices]))) + + for n in range(len(new_shape)): + if new_shape[n] > flat_indices[n].dimension: + while new_shape[n] > flat_indices[n].dimension: + #fuse flat_indices + i1, i2 = flat_indices.pop(n), flat_indices.pop(n) #note: the resulting flow is set to one since the flow #is multiplied into the charges. As a result the tensor #will then be invariant in any case. - self.indices.insert(n, fuse_index_pair(i1, i2)) - if self.dense_shape[n] > dense_shape[n]: + flat_indices.insert(n, fuse_index_pair(i1, i2)) + if flat_indices[n].dimension > new_shape[n]: raise_error() - elif dense_shape[n] < self.dense_shape[n]: + elif new_shape[n] < flat_indices[n].dimension: raise_error() - #at this point the first len(dense_shape) indices of the tensor - #match the `dense_shape`. - while len(dense_shape) < len(self.indices): - i2, i1 = self.indices.pop(), self.indices.pop() - self.indices.append(fuse_index_pair(i1, i2)) + #at this point the first len(new_shape) flat_indices of the tensor + #match the `new_shape`. + while len(new_shape) < len(flat_indices): + i2, i1 = flat_indices.pop(), flat_indices.pop() + flat_indices.append(fuse_index_pair(i1, i2)) - def get_diagonal_blocks(self, return_data: Optional[bool] = True) -> Dict: - """ - Obtain the diagonal blocks of symmetric matrix. - BlockSparseTensor has to be a matrix. - For matrices with shape[0] << shape[1], this routine avoids explicit fusion - of column charges. - - Args: - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - Returns: - dict: Dictionary mapping charge to np.ndarray of rank 2 (a matrix) - - """ - if self.rank != 2: - raise ValueError( - "`get_diagonal_blocks` can only be called on a matrix, but found rank={}" - .format(self.rank)) - - row_indices = self.indices[0].get_elementary_indices() - column_indices = self.indices[1].get_elementary_indices() - - return find_diagonal_sparse_blocks( - data=self.data, - row_charges=[i.charges for i in row_indices], - column_charges=[i.charges for i in column_indices], - row_flows=[i.flow for i in row_indices], - column_flows=[i.flow for i in column_indices], - return_data=return_data) - - def get_diagonal_blocks_version_1(self, - return_data: Optional[bool] = True) -> Dict: - """ - Obtain the diagonal blocks of symmetric matrix. - BlockSparseTensor has to be a matrix. - For matrices with shape[0] << shape[1], this routine avoids explicit fusion - of column charges. - - Args: - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - Returns: - dict: Dictionary mapping charge to np.ndarray of rank 2 (a matrix) - - """ - if self.rank != 2: - raise ValueError( - "`get_diagonal_blocks` can only be called on a matrix, but found rank={}" - .format(self.rank)) - - row_indices = self.indices[0].get_elementary_indices() - column_indices = self.indices[1].get_elementary_indices() - - return find_diagonal_sparse_blocks_version_1( - data=self.data, - row_charges=[i.charges for i in row_indices], - column_charges=[i.charges for i in column_indices], - row_flows=[i.flow for i in row_indices], - column_flows=[i.flow for i in column_indices], - return_data=return_data) - - def get_diagonal_blocks_version_0(self, - return_data: Optional[bool] = True) -> Dict: - """ - Deprecated - - Obtain the diagonal blocks of symmetric matrix. - BlockSparseTensor has to be a matrix. - Args: - return_data: If `True`, the return dictionary maps quantum numbers `q` to - actual `np.ndarray` with the data. This involves a copy of data. - If `False`, the returned dict maps quantum numbers of a list - [locations, shape], where `locations` is an np.ndarray of type np.int64 - containing the locations of the tensor elements within A.data, i.e. - `A.data[locations]` contains the elements belonging to the tensor with - quantum numbers `(q,q). `shape` is the shape of the corresponding array. - Returns: - dict: Dictionary mapping charge to np.ndarray of rank 2 (a matrix) - - """ - if self.rank != 2: - raise ValueError( - "`get_diagonal_blocks` can only be called on a matrix, but found rank={}" - .format(self.rank)) - - return find_diagonal_sparse_blocks_version_0( - data=self.data, - charges=self.charges, - flows=self.flows, - return_data=return_data) + result = BlockSparseTensor(data=self.data, indices=flat_indices) + return result def reshape(tensor: BlockSparseTensor, @@ -1229,20 +770,1425 @@ def reshape(tensor: BlockSparseTensor, i2 = Index(charges=q2,flow=-1) i3 = Index(charges=q3,flow=1) A=BlockSparseTensor.randn(indices=[i1,i2,i3]) - print(A.shape) #prints (6,6,6) + print(nA.shape) #prints (6,6,6) reshape(A, (2,3,6,6)) #raises ValueError ``` raises a `ValueError` since (2,3,6,6) is incompatible with the elementary shape (6,6,6) of the tensor. Args: - tensor: A symmetric tensor. + tensopr: A symmetric tensor. shape: The new shape. Can either be a list of `Index` or a list of `int`. Returns: BlockSparseTensor: A new tensor reshaped into `shape` """ - result = BlockSparseTensor( - data=tensor.data.copy(), indices=[i.copy() for i in tensor.indices]) - result.reshape(shape) + + return tensor.reshape(shape) + + +def transpose(tensor: BlockSparseTensor, + order: Union[List[int], np.ndarray]) -> "BlockSparseTensor": + """ + Transpose `tensor` into the new order `order`. This routine currently shuffles + data. + Args: + tensor: The tensor to be transposed. + order: The new order of indices. + Returns: + BlockSparseTensor: The transposed tensor. + """ + result = tensor.copy() + result.transpose(order) return result + + +def _compute_transposed_sparse_blocks(indices: BlockSparseTensor, + order: Union[List[int], np.ndarray], + transposed_partition: Optional[int] = None + ) -> Tuple[BaseCharge, Dict, int]: + """ + Args: + indices: A symmetric tensor. + order: The new order of indices. + permutation: An np.ndarray of int for reshuffling the data, + typically the output of a prior call to `transpose`. Passing `permutation` + can greatly speed up the transposition. + return_permutation: If `True`, return the the permutation data. + Returns: + + """ + if len(order) != len(indices): + raise ValueError( + "`len(order)={}` is different form `len(indices)={}`".format( + len(order), len(indices))) + flat_indices, flat_charges, flat_flows, flat_strides, flat_order, transposed_partition = flatten_meta_data( + indices, order, transposed_partition) + if transposed_partition is None: + transposed_partition = _find_best_partition( + [len(flat_charges[n]) for n in flat_order]) + + return _find_transposed_diagonal_sparse_blocks( + flat_charges, + flat_flows, + tr_partition=transposed_partition, + order=flat_order) + + +def tensordot(tensor1: BlockSparseTensor, + tensor2: BlockSparseTensor, + axes: Sequence[Sequence[int]], + final_order: Optional[Union[List, np.ndarray]] = None + ) -> BlockSparseTensor: + """ + Contract two `BlockSparseTensor`s along `axes`. + Args: + tensor1: First tensor. + tensor2: Second tensor. + axes: The axes to contract. + final_order: An optional final order for the result + Returns: + BlockSparseTensor: The result of the tensor contraction. + + """ + axes1 = axes[0] + axes2 = axes[1] + if not np.all(np.unique(axes1) == np.sort(axes1)): + raise ValueError( + "Some values in axes[0] = {} appear more than once!".format(axes1)) + if not np.all(np.unique(axes2) == np.sort(axes2)): + raise ValueError( + "Some values in axes[1] = {} appear more than once!".format(axes2n)) + + if max(axes1) >= len(tensor1.shape): + raise ValueError( + "rank of `tensor1` is smaller than `max(axes1) = {}.`".format( + max(axes1))) + + if max(axes2) >= len(tensor2.shape): + raise ValueError( + "rank of `tensor2` is smaller than `max(axes2) = {}`".format( + max(axes1))) + elementary_1, elementary_2 = [], [] + for a in axes1: + elementary_1.extend(tensor1.indices[a].get_elementary_indices()) + for a in axes2: + elementary_2.extend(tensor2.indices[a].get_elementary_indices()) + + if len(elementary_2) != len(elementary_1): + raise ValueError("axes1 and axes2 have incompatible elementary" + " shapes {} and {}".format(elementary_1, elementary_2)) + if not np.all( + np.array([i.flow for i in elementary_1]) == np.array( + [not i.flow for i in elementary_2])): + raise ValueError("axes1 and axes2 have incompatible elementary" + " flows {} and {}".format( + np.array([i.flow for i in elementary_1]), + np.array([i.flow for i in elementary_2]))) + + free_axes1 = sorted(set(np.arange(len(tensor1.shape))) - set(axes1)) + free_axes2 = sorted(set(np.arange(len(tensor2.shape))) - set(axes2)) + if (final_order is not None) and (len(final_order) != + len(free_axes1) + len(free_axes2)): + raise ValueError("`final_order = {}` is not a valid order for " + "a final tensor of rank {}".format( + final_order, + len(free_axes1) + len(free_axes2))) + + if (final_order is not None) and not np.all( + np.sort(final_order) == np.arange(len(final_order))): + raise ValueError( + "`final_order = {}` is not a valid permutation of {} ".format( + final_order, np.arange(len(final_order)))) + + new_order1 = free_axes1 + list(axes1) + new_order2 = list(axes2) + free_axes2 + + #get the flattened indices for the output tensor + left_indices = [] + right_indices = [] + for n in free_axes1: + left_indices.extend(tensor1.indices[n].get_elementary_indices()) + for n in free_axes2: + right_indices.extend(tensor2.indices[n].get_elementary_indices()) + indices = left_indices + right_indices + + _, flat_charges1, flat_flows1, flat_strides1, flat_order1, tr_partition1 = flatten_meta_data( + tensor1.indices, new_order1, len(free_axes1)) + tr_sparse_blocks_1, charges1, shapes_1 = _find_transposed_diagonal_sparse_blocks( + flat_charges1, flat_flows1, tr_partition1, flat_order1) + _, flat_charges2, flat_flows2, flat_strides2, flat_order2, tr_partition2 = flatten_meta_data( + tensor2.indices, new_order2, len(axes2)) + tr_sparse_blocks_2, charges2, shapes_2 = _find_transposed_diagonal_sparse_blocks( + flat_charges2, flat_flows2, tr_partition2, flat_order2) + #common_charges = charges1.intersect(charges2) + common_charges, label_to_common_1, label_to_common_2 = intersect( + charges1.unique_charges, + charges2.unique_charges, + axis=1, + return_indices=True) + + #initialize the data-vector of the output with zeros; + if final_order is not None: + #in this case we view the result of the diagonal multiplication + #as a transposition of the final tensor + final_indices = [indices[n] for n in final_order] + _, reverse_order = np.unique(final_order, return_index=True) + + flat_final_indices, flat_final_charges, flat_final_flows, flat_final_strides, flat_final_order, tr_partition = flatten_meta_data( + final_indices, reverse_order, len(free_axes1)) + + sparse_blocks_final, charges_final, shapes_final = _find_transposed_diagonal_sparse_blocks( + flat_final_charges, flat_final_flows, tr_partition, flat_final_order) + + num_nonzero_elements = np.sum([len(v) for v in sparse_blocks_final]) + data = np.zeros( + num_nonzero_elements, + dtype=np.result_type(tensor1.dtype, tensor2.dtype)) + label_to_common_final = intersect( + charges_final.unique_charges, + common_charges, + axis=1, + return_indices=True)[1] + + for n in range(common_charges.shape[1]): + n1 = label_to_common_1[n] + n2 = label_to_common_2[n] + nf = label_to_common_final[n] + + data[sparse_blocks_final[nf].ravel()] = ( + tensor1.data[tr_sparse_blocks_1[n1].reshape( + shapes_1[:, n1])] @ tensor2.data[tr_sparse_blocks_2[n2].reshape( + shapes_2[:, n2])]).ravel() + + return BlockSparseTensor(data=data, indices=final_indices) + else: + #Note: `cs` may contain charges that are not present in `common_charges` + charges = [i.charges for i in indices] + flows = [i.flow for i in indices] + sparse_blocks, cs, shapes = _find_diagonal_sparse_blocks( + charges, flows, len(left_indices)) + num_nonzero_elements = np.sum([len(v) for v in sparse_blocks]) + #Note that empty is not a viable choice here. + data = np.zeros( + num_nonzero_elements, + dtype=np.result_type(tensor1.dtype, tensor2.dtype)) + + label_to_common_final = intersect( + cs.unique_charges, common_charges, axis=1, return_indices=True)[1] + + for n in range(common_charges.shape[1]): + n1 = label_to_common_1[n] + n2 = label_to_common_2[n] + nf = label_to_common_final[n] + + data[sparse_blocks[nf].ravel()] = ( + tensor1.data[tr_sparse_blocks_1[n1].reshape( + shapes_1[:, n1])] @ tensor2.data[tr_sparse_blocks_2[n2].reshape( + shapes_2[:, n2])]).ravel() + return BlockSparseTensor(data=data, indices=indices) + + +def flatten_meta_data(indices, order, partition): + elementary_indices = {} + flat_elementary_indices = [] + new_partition = 0 + for n in range(len(indices)): + elementary_indices[n] = indices[n].get_elementary_indices() + if n < partition: + new_partition += len(elementary_indices[n]) + flat_elementary_indices.extend(elementary_indices[n]) + flat_index_list = np.arange(len(flat_elementary_indices)) + cum_num_legs = np.append( + 0, np.cumsum([len(elementary_indices[n]) for n in range(len(indices))])) + + flat_charges = [i.charges for i in flat_elementary_indices] + flat_flows = [i.flow for i in flat_elementary_indices] + flat_dims = [len(c) for c in flat_charges] + flat_strides = _get_strides(flat_dims) + flat_order = np.concatenate( + [flat_index_list[cum_num_legs[n]:cum_num_legs[n + 1]] for n in order]) + + return flat_elementary_indices, flat_charges, flat_flows, flat_strides, flat_order, new_partition + + +##################################################### DEPRECATED ROUTINES ############################ + + +def _find_transposed_diagonal_sparse_blocks_old( + charges: List[BaseCharge], flows: List[Union[bool, int]], order: np.ndarray, + tr_partition: int) -> Tuple[BaseCharge, List[np.ndarray]]: + """ + Given the meta data and underlying data of a symmetric matrix, compute the + dense positions of all diagonal blocks and return them in a dict. + `row_charges` and `column_charges` are lists of np.ndarray. The tensor + is viewed as a matrix with rows given by fusing `row_charges` and + columns given by fusing `column_charges`. + + Args: + data: An np.ndarray of the data. The number of elements in `data` + has to match the number of non-zero elements defined by `charges` + and `flows` + row_charges: List of np.ndarray, one for each leg of the row-indices. + Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. + The bond dimension `D[leg]` can vary on each leg. + column_charges: List of np.ndarray, one for each leg of the column-indices. + Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. + The bond dimension `D[leg]` can vary on each leg. + row_flows: A list of integers, one for each entry in `row_charges`. + with values `1` or `-1`, denoting the flow direction + of the charges on each leg. `1` is inflowing, `-1` is outflowing + charge. + column_flows: A list of integers, one for each entry in `column_charges`. + with values `1` or `-1`, denoting the flow direction + of the charges on each leg. `1` is inflowing, `-1` is outflowing + charge. + row_strides: An optional np.ndarray denoting the strides of `row_charges`. + If `None`, natural stride ordering is assumed. + column_strides: An optional np.ndarray denoting the strides of + `column_charges`. If `None`, natural stride ordering is assumed. + + Returns: + List[Union[BaseCharge, ChargeCollection]]: A list of unique charges, one per block. + List[List]: A list containing the blocks information. + For each element `e` in the list `e[0]` is an `np.ndarray` of ints + denoting the dense positions of the non-zero elements and `e[1]` + is a tuple corresponding to the blocks' matrix shape + """ + _check_flows(flows) + if len(flows) != len(charges): + raise ValueError("`len(flows)` is different from `len(charges) ") + if np.all(order == np.arange(len(order))): + return _find_diagonal_sparse_blocks(charges, flows, tr_partition) + + strides = _get_strides([len(c) for c in charges]) + + tr_row_charges = [charges[n] for n in order[:tr_partition]] + tr_row_flows = [flows[n] for n in order[:tr_partition]] + tr_row_strides = [strides[n] for n in order[:tr_partition]] + + tr_column_charges = [charges[n] for n in order[tr_partition:]] + tr_column_flows = [flows[n] for n in order[tr_partition:]] + tr_column_strides = [strides[n] for n in order[tr_partition:]] + + unique_tr_column_charges, tr_column_dims = compute_fused_charge_degeneracies( + tr_column_charges, tr_column_flows) + unique_tr_row_charges = compute_unique_fused_charges(tr_row_charges, + tr_row_flows) + + fused = unique_tr_row_charges + unique_tr_column_charges + tr_li, tr_ri = np.divmod( + np.nonzero(fused == unique_tr_column_charges.identity_charges)[0], + len(unique_tr_column_charges)) + + row_ind, row_locations = reduce_charges( + charges=tr_row_charges, + flows=tr_row_flows, + target_charges=unique_tr_row_charges.charges[:, tr_li], + return_locations=True, + strides=tr_row_strides) + + col_ind, column_locations = reduce_charges( + charges=tr_column_charges, + flows=tr_column_flows, + target_charges=unique_tr_column_charges.charges[:, tr_ri], + return_locations=True, + strides=tr_column_strides) + + partition = _find_best_partition([len(c) for c in charges]) + fused_row_charges = fuse_charges(charges[:partition], flows[:partition]) + fused_column_charges = fuse_charges(charges[partition:], flows[partition:]) + + unique_fused_row, row_inverse = fused_row_charges.unique(return_inverse=True) + unique_fused_column, column_inverse = fused_column_charges.unique( + return_inverse=True) + + unique_column_charges, column_dims = compute_fused_charge_degeneracies( + charges[partition:], flows[partition:]) + unique_row_charges = compute_unique_fused_charges(charges[:partition], + flows[:partition]) + fused = unique_row_charges + unique_column_charges + li, ri = np.divmod( + np.nonzero(fused == unique_column_charges.identity_charges)[0], + len(unique_column_charges)) + + common_charges, label_to_row, label_to_column = unique_row_charges.intersect( + unique_column_charges * True, return_indices=True) + num_blocks = len(label_to_row) + tmp = -np.ones(len(unique_row_charges), dtype=np.int16) + for n in range(len(label_to_row)): + tmp[label_to_row[n]] = n + + degeneracy_vector = np.append(column_dims[label_to_column], + 0)[tmp[row_inverse]] + start_positions = np.cumsum(np.insert(degeneracy_vector[:-1], 0, + 0)).astype(np.uint32) + + column_dimension = np.prod([len(c) for c in charges[partition:]]) + + column_lookup = compute_sparse_lookup(charges[partition:], flows[partition:], + common_charges) + + blocks = [] + for n in range(num_blocks): + rlocs = row_locations[row_ind.charge_labels == n] + clocs = column_locations[col_ind.charge_labels == n] + orig_row_posL, orig_col_posL = np.divmod(rlocs, np.uint32(column_dimension)) + orig_row_posR, orig_col_posR = np.divmod(clocs, np.uint32(column_dimension)) + inds = (start_positions[np.add.outer(orig_row_posL, orig_row_posR)] + + column_lookup[np.add.outer(orig_col_posL, orig_col_posR)]).ravel() + + blocks.append([inds, (len(rlocs), len(clocs))]) + charges_out = unique_tr_row_charges[tr_li] + return charges_out, blocks + + +def _find_diagonal_dense_blocks( + row_charges: List[BaseCharge], + column_charges: List[BaseCharge], + row_flows: List[Union[bool, int]], + column_flows: List[Union[bool, int]], + row_strides: Optional[np.ndarray] = None, + column_strides: Optional[np.ndarray] = None, +) -> Tuple[BaseCharge, List[np.ndarray]]: + """ + + Deprecated + Given the meta data and underlying data of a symmetric matrix, compute the + dense positions of all diagonal blocks and return them in a dict. + `row_charges` and `column_charges` are lists of np.ndarray. The tensor + is viewed as a matrix with rows given by fusing `row_charges` and + columns given by fusing `column_charges`. + + Args: + data: An np.ndarray of the data. The number of elements in `data` + has to match the number of non-zero elements defined by `charges` + and `flows` + row_charges: List of np.ndarray, one for each leg of the row-indices. + Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. + The bond dimension `D[leg]` can vary on each leg. + column_charges: List of np.ndarray, one for each leg of the column-indices. + Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`. + The bond dimension `D[leg]` can vary on each leg. + row_flows: A list of integers, one for each entry in `row_charges`. + with values `1` or `-1`, denoting the flow direction + of the charges on each leg. `1` is inflowing, `-1` is outflowing + charge. + column_flows: A list of integers, one for each entry in `column_charges`. + with values `1` or `-1`, denoting the flow direction + of the charges on each leg. `1` is inflowing, `-1` is outflowing + charge. + row_strides: An optional np.ndarray denoting the strides of `row_charges`. + If `None`, natural stride ordering is assumed. + column_strides: An optional np.ndarray denoting the strides of + `column_charges`. If `None`, natural stride ordering is assumed. + + Returns: + List[Union[BaseCharge, ChargeCollection]]: A list of unique charges, one per block. + List[List]: A list containing the blocks information. + For each element `e` in the list `e[0]` is an `np.ndarray` of ints + denoting the dense positions of the non-zero elements and `e[1]` + is a tuple corresponding to the blocks' matrix shape + + """ + flows = list(row_flows).copy() + flows.extend(column_flows) + _check_flows(flows) + if len(flows) != (len(row_charges) + len(column_charges)): + raise ValueError( + "`len(flows)` is different from `len(row_charges) + len(column_charges)`" + ) + #get the unique column-charges + #we only care about their degeneracies, not their order; that's much faster + #to compute since we don't have to fuse all charges explicitly + #`compute_fused_charge_degeneracies` multiplies flows into the column_charges + unique_column_charges = compute_unique_fused_charges(column_charges, + column_flows) + + unique_row_charges = compute_unique_fused_charges(row_charges, row_flows) + #get the charges common to rows and columns (only those matter) + fused = unique_row_charges + unique_column_charges + li, ri = np.divmod( + np.nonzero(fused == unique_column_charges.identity_charges)[0], + len(unique_column_charges)) + if ((row_strides is None) and + (column_strides is not None)) or ((row_strides is not None) and + (column_strides is None)): + raise ValueError("`row_strides` and `column_strides` " + "have to be passed simultaneously." + " Found `row_strides={}` and " + "`column_strides={}`".format(row_strides, column_strides)) + if row_strides is not None: + row_locations = find_dense_positions( + charges=row_charges, + flows=row_flows, + target_charges=unique_row_charges[li], + strides=row_strides) + + else: + column_dim = np.prod([len(c) for c in column_charges]) + row_locations = find_dense_positions( + charges=row_charges, + flows=row_flows, + target_charges=unique_row_charges[li]) + for v in row_locations.values(): + v *= column_dim + if column_strides is not None: + column_locations = find_dense_positions( + charges=column_charges, + flows=column_flows, + target_charges=unique_column_charges[ri], + strides=column_strides, + store_dual=True) + + else: + column_locations = find_dense_positions( + charges=column_charges, + flows=column_flows, + target_charges=unique_column_charges[ri], + store_dual=True) + blocks = [] + for c in unique_row_charges[li]: + #numpy broadcasting is substantially faster than kron! + rlocs = np.expand_dims(row_locations[c], 1) + clocs = np.expand_dims(column_locations[c], 0) + inds = np.reshape(rlocs + clocs, rlocs.shape[0] * clocs.shape[1]) + blocks.append([inds, (rlocs.shape[0], clocs.shape[1])]) + return unique_row_charges[li], blocks + + +# def find_sparse_positions_2( +# charges: List[Union[BaseCharge, ChargeCollection]], +# flows: List[Union[int, bool]], +# target_charges: Union[BaseCharge, ChargeCollection]) -> Dict: +# """ +# Find the sparse locations of elements (i.e. the index-values within +# the SPARSE tensor) in the vector `fused_charges` (resulting from +# fusing `left_charges` and `right_charges`) +# that have a value of `target_charges`, assuming that all elements +# different from `target_charges` are `0`. +# For example, given +# ``` +# left_charges = [-2,0,1,0,0] +# right_charges = [-1,0,2,1] +# target_charges = [0,1] +# fused_charges = fuse_charges([left_charges, right_charges],[1,1]) +# print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] +# ``` 0 1 2 3 4 5 6 7 8 +# we want to find the all different blocks +# that fuse to `target_charges=[0,1]`, i.e. where `fused_charges==0` or `1`, +# together with their corresponding sparse index-values of the data in the sparse array, +# assuming that all elements in `fused_charges` different from `target_charges` are 0. + +# `find_sparse_blocks` returns a dict mapping integers `target_charge` +# to an array of integers denoting the sparse locations of elements within +# `fused_charges`. +# For the above example, we get: +# * `target_charge=0`: [0,1,3,5,7] +# * `target_charge=1`: [2,4,6,8] +# Args: +# left_charges: An np.ndarray of integer charges. +# left_flow: The flow direction of the left charges. +# right_charges: An np.ndarray of integer charges. +# right_flow: The flow direction of the right charges. +# target_charge: The target charge. +# Returns: +# dict: Mapping integers to np.ndarray of integers. +# """ +# #FIXME: this is probably still not optimal + +# _check_flows(flows) +# if len(charges) == 1: +# fused_charges = charges[0] * flows[0] +# unique_charges = fused_charges.unique() +# target_charges = target_charges.unique() +# relevant_target_charges = unique_charges.intersect(target_charges) +# relevant_fused_charges = fused_charges[fused_charges.isin( +# relevant_target_charges)] +# return { +# c: np.nonzero(relevant_fused_charges == c)[0] +# for c in relevant_target_charges +# } + +# left_charges, right_charges, partition = _find_best_partition(charges, flows) + +# unique_target_charges, inds = target_charges.unique(return_index=True) +# target_charges = target_charges[np.sort(inds)] + +# unique_left = left_charges.unique() +# unique_right = right_charges.unique() +# fused = unique_left + unique_right + +# #compute all unique charges that can add up to +# #target_charges +# left_inds, right_inds = [], [] +# for target_charge in target_charges: +# li, ri = np.divmod(np.nonzero(fused == target_charge)[0], len(unique_right)) +# left_inds.append(li) +# right_inds.append(ri) + +# #now compute the relevant unique left and right charges +# unique_left_charges = unique_left[np.unique(np.concatenate(left_inds))] +# unique_right_charges = unique_right[np.unique(np.concatenate(right_inds))] + +# #only keep those charges that are relevant +# relevant_left_charges = left_charges[left_charges.isin(unique_left_charges)] +# relevant_right_charges = right_charges[right_charges.isin( +# unique_right_charges)] + +# unique_right_charges, right_dims = relevant_right_charges.unique( +# return_counts=True) +# right_degeneracies = dict(zip(unique_right_charges, right_dims)) +# #generate a degeneracy vector which for each value r in relevant_right_charges +# #holds the corresponding number of non-zero elements `relevant_right_charges` +# #that can add up to `target_charges`. +# degeneracy_vector = np.empty(len(relevant_left_charges), dtype=np.int64) +# right_indices = {} + +# for n in range(len(unique_left_charges)): +# left_charge = unique_left_charges[n] +# total_charge = left_charge + unique_right_charges +# total_degeneracy = np.sum(right_dims[total_charge.isin(target_charges)]) +# tmp_relevant_right_charges = relevant_right_charges[ +# relevant_right_charges.isin((target_charges + left_charge * (-1)))] + +# for n in range(len(target_charges)): +# target_charge = target_charges[n] +# right_indices[(left_charge.get_item(0), +# target_charge.get_item(0))] = np.nonzero( +# tmp_relevant_right_charges == ( +# target_charge + left_charge * (-1)))[0] + +# degeneracy_vector[relevant_left_charges == left_charge] = total_degeneracy + +# start_positions = np.cumsum(degeneracy_vector) - degeneracy_vector +# blocks = {t: [] for t in target_charges} +# # iterator returns tuple of `int` for ChargeCollection objects +# # and `int` for Ba seCharge objects (both hashable) +# for left_charge in unique_left_charges: +# a = np.expand_dims(start_positions[relevant_left_charges == left_charge], 0) +# for target_charge in target_charges: +# ri = right_indices[(left_charge, target_charge)] +# if len(ri) != 0: +# b = np.expand_dims(ri, 1) +# tmp = a + b +# blocks[target_charge].append(np.reshape(tmp, np.prod(tmp.shape))) +# out = {} +# for target_charge in target_charges: +# out[target_charge] = np.concatenate(blocks[target_charge]) +# return out + + +def _compute_sparse_lookups(row_charges: BaseCharge, row_flows, column_charges, + column_flows): + """ + Compute lookup tables for looking up how dense index positions map + to sparse index positions for the diagonal blocks a symmetric matrix. + Args: + row_charges: + + """ + column_flows = list(-np.asarray(column_flows)) + fused_column_charges = fuse_charges(column_charges, column_flows) + fused_row_charges = fuse_charges(row_charges, row_flows) + unique_column_charges, column_inverse = fused_column_charges.unique( + return_inverse=True) + unique_row_charges, row_inverse = fused_row_charges.unique( + return_inverse=True) + common_charges, comm_row, comm_col = unique_row_charges.intersect( + unique_column_charges, return_indices=True) + + col_ind_sort = np.argsort(column_inverse, kind='stable') + row_ind_sort = np.argsort(row_inverse, kind='stable') + _, col_charge_degeneracies = compute_fused_charge_degeneracies( + column_charges, column_flows) + _, row_charge_degeneracies = compute_fused_charge_degeneracies( + row_charges, row_flows) + # labelsorted_indices = column_inverse[col_ind_sort] + # tmp = np.nonzero( + # np.append(labelsorted_indices, unique_column_charges.charges.shape[0] + 1) - + # np.append(labelsorted_indices[0], labelsorted_indices))[0] + #charge_degeneracies = tmp - np.append(0, tmp[0:-1]) + + col_start_positions = np.cumsum(np.append(0, col_charge_degeneracies)) + row_start_positions = np.cumsum(np.append(0, row_charge_degeneracies)) + column_lookup = np.empty(len(fused_column_charges), dtype=np.uint32) + row_lookup = np.zeros(len(fused_row_charges), dtype=np.uint32) + for n in range(len(common_charges)): + column_lookup[col_ind_sort[col_start_positions[ + comm_col[n]]:col_start_positions[comm_col[n] + 1]]] = np.arange( + col_charge_degeneracies[comm_col[n]]) + row_lookup[ + row_ind_sort[row_start_positions[comm_row[n]]:row_start_positions[ + comm_row[n] + 1]]] = col_charge_degeneracies[comm_col[n]] + + return np.append(0, np.cumsum(row_lookup[0:-1])), column_lookup + + +def _get_stride_arrays(dims): + strides = np.flip(np.append(1, np.cumprod(np.flip(dims[1::])))) + return [np.arange(dims[n]) * strides[n] for n in range(len(dims))] + + +def reduce_charges(charges: List[BaseCharge], + flows: Iterable[bool], + target_charges: np.ndarray, + return_locations: Optional[bool] = False, + strides: Optional[np.ndarray] = None + ) -> Tuple[BaseCharge, np.ndarray]: + """ + Add quantum numbers arising from combining two or more charges into a + single index, keeping only the quantum numbers that appear in 'target_charges'. + Equilvalent to using "combine_charges" followed by "reduce", but is + generally much more efficient. + Args: + charges (List[SymIndex]): list of SymIndex. + flows (np.ndarray): vector of bools describing index orientations. + target_charges (np.ndarray): n-by-m array describing qauntum numbers of the + qnums which should be kept with 'n' the number of symmetries. + return_locations (bool, optional): if True then return the location of the kept + values of the fused charges + strides (np.ndarray, optional): index strides with which to compute the + return_locations of the kept elements. Defaults to trivial strides (based on + row major order) if ommitted. + Returns: + SymIndex: the fused index after reduction. + np.ndarray: locations of the fused SymIndex qnums that were kept. + """ + + num_inds = len(charges) + tensor_dims = [len(c) for c in charges] + + if len(charges) == 1: + # reduce single index + if strides is None: + strides = np.array([1], dtype=np.uint32) + return charges[0].dual(flows[0]).reduce( + target_charges, return_locations=return_locations, strides=strides[0]) + + else: + # find size-balanced partition of charges + partition = _find_best_partition(tensor_dims) + + # compute quantum numbers for each partition + left_ind = fuse_charges(charges[:partition], flows[:partition]) + right_ind = fuse_charges(charges[partition:], flows[partition:]) + + # compute combined qnums + comb_qnums = fuse_ndarray_charges(left_ind.unique_charges, + right_ind.unique_charges, + charges[0].charge_types) + [unique_comb_qnums, comb_labels] = np.unique( + comb_qnums, return_inverse=True, axis=1) + num_unique = unique_comb_qnums.shape[1] + + # intersect combined qnums and target_charges + reduced_qnums, label_to_unique, label_to_kept = intersect( + unique_comb_qnums, target_charges, axis=1, return_indices=True) + map_to_kept = -np.ones(num_unique, dtype=np.int16) + for n in range(len(label_to_unique)): + map_to_kept[label_to_unique[n]] = n + new_comb_labels = map_to_kept[comb_labels].reshape( + [left_ind.num_unique, right_ind.num_unique]) + if return_locations: + if strides is not None: + # computed locations based on non-trivial strides + row_pos = fuse_stride_arrays(tensor_dims[:partition], strides[:partition]) + col_pos = fuse_stride_arrays(tensor_dims[partition:], strides[partition:]) + + # reduce combined qnums to include only those in target_charges + reduced_rows = [0] * left_ind.num_unique + row_locs = [0] * left_ind.num_unique + for n in range(left_ind.num_unique): + temp_label = new_comb_labels[n, right_ind.charge_labels] + temp_keep = temp_label >= 0 + reduced_rows[n] = temp_label[temp_keep] + row_locs[n] = col_pos[temp_keep] + + reduced_labels = np.concatenate( + [reduced_rows[n] for n in left_ind.charge_labels]) + reduced_locs = np.concatenate([ + row_pos[n] + row_locs[left_ind.charge_labels[n]] + for n in range(left_ind.dim) + ]) + obj = charges[0].__new__(type(charges[0])) + obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types) + return obj, reduced_locs + + else: # trivial strides + # reduce combined qnums to include only those in target_charges + reduced_rows = [0] * left_ind.num_unique + row_locs = [0] * left_ind.num_unique + for n in range(left_ind.num_unique): + temp_label = new_comb_labels[n, right_ind.charge_labels] + temp_keep = temp_label >= 0 + reduced_rows[n] = temp_label[temp_keep] + row_locs[n] = np.where(temp_keep)[0] + + reduced_labels = np.concatenate( + [reduced_rows[n] for n in left_ind.charge_labels]) + reduced_locs = np.concatenate([ + n * right_ind.dim + row_locs[left_ind.charge_labels[n]] + for n in range(left_ind.dim) + ]) + obj = charges[0].__new__(type(charges[0])) + obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types) + + return obj, reduced_locs + + else: + # reduce combined qnums to include only those in target_charges + reduced_rows = [0] * left_ind.num_unique + for n in range(left_ind.num_unique): + temp_label = new_comb_labels[n, right_ind.charge_labels] + reduced_rows[n] = temp_label[temp_label >= 0] + + reduced_labels = np.concatenate( + [reduced_rows[n] for n in left_ind.charge_labels]) + obj = charges[0].__new__(type(charges[0])) + obj.__init__(reduced_qnums, reduced_labels, charges[0].charge_types) + + return obj + + +def reduce_to_target_charges(charges: List[BaseCharge], + flows: List[Union[int, bool]], + target_charges: BaseCharge, + strides: Optional[np.ndarray] = None, + return_positions: Optional[bool] = False + ) -> np.ndarray: + """ + Find the dense locations of elements (i.e. the index-values within the DENSE tensor) + in the vector of `fused_charges` resulting from fusing all elements of `charges` + that have a value of `target_charge`. + For example, given + ``` + charges = [[-2,0,1,0,0],[-1,0,2,1]] + target_charge = 0 + fused_charges = fuse_charges(charges,[1,1]) + print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] + ``` + we want to find the index-positions of charges + that fuse to `target_charge=0`, i.e. where `fused_charges==0`, + within the dense array. As one additional wrinkle, `charges` + is a subset of the permuted charges of a tensor with rank R > len(charges), + and `stride_arrays` are their corresponding range of strides, i.e. + + ``` + R=5 + D = [2,3,4,5,6] + tensor_flows = np.random.randint(-1,2,R) + tensor_charges = [np.random.randing(-5,5,D[n]) for n in range(R)] + order = np.arange(R) + np.random.shuffle(order) + tensor_strides = [360, 120, 30, 6, 1] + + charges = [tensor_charges[order[n]] for n in range(3)] + flows = [tensor_flows[order[n]] for n in range(len(3))] + strides = [tensor_stride[order[n]] for n in range(3)] + _ = _find_transposed_dense_positions(charges, flows, 0, strides) + + ``` + `_find_transposed_dense_blocks` returns an np.ndarray containing the + index-positions of these elements calculated using `stride_arrays`. + The result only makes sense in conjuction with the complementary + data computed from the complementary + elements in`tensor_charges`, + `tensor_strides` and `tensor_flows`. + This routine is mainly used in `_find_diagonal_dense_blocks`. + + Args: + charges: A list of BaseCharge or ChargeCollection. + flows: The flow directions of the `charges`. + target_charge: The target charge. + strides: The strides for the `charges` subset. + if `None`, natural stride ordering is assumed. + + Returns: + np.ndarray: The index-positions within the dense data array + of the elements fusing to `target_charge`. + """ + + _check_flows(flows) + if len(charges) == 1: + fused_charges = charges[0] * flows[0] + unique, inverse = fused_charges.unique(return_inverse=True) + common, label_to_unique, label_to_target = unique.intersect( + target_charges, return_indices=True) + inds = np.nonzero(np.isin(inverse, label_to_unique))[0] + if strides is not None: + permuted_inds = strides[0] * np.arange(len(charges[0])) + if return_positions: + return fused_charges[permuted_inds[inds]], inds + return fused_charges[permuted_inds[inds]] + + if return_positions: + return fused_charges[inds], inds + return fused_charges[inds] + + partition = _find_best_partition([len(c) for c in charges]) + left_charges = fuse_charges(charges[:partition], flows[:partition]) + right_charges = fuse_charges(charges[partition:], flows[partition:]) + + # unique_target_charges, inds = target_charges.unique(return_index=True) + # target_charges = target_charges[np.sort(inds)] + unique_left, left_inverse = left_charges.unique(return_inverse=True) + unique_right, right_inverse = right_charges.unique(return_inverse=True) + + fused = unique_left + unique_right + unique_fused, unique_fused_labels = fused.unique(return_inverse=True) + + relevant_charges, relevant_labels, _ = unique_fused.intersect( + target_charges, return_indices=True) + + tmp = np.full(len(unique_fused), fill_value=-1, dtype=np.int16) + tmp[relevant_labels] = np.arange(len(relevant_labels), dtype=np.int16) + lookup_target = tmp[unique_fused_labels].reshape( + [len(unique_left), len(unique_right)]) + + if return_positions: + if strides is not None: + stride_arrays = [ + np.arange(len(charges[n])) * strides[n] for n in range(len(charges)) + ] + permuted_left_inds = fuse_ndarrays(stride_arrays[0:partition]) + permuted_right_inds = fuse_ndarrays(stride_arrays[partition:]) + + row_locations = [None] * len(unique_left) + final_relevant_labels = [None] * len(unique_left) + for n in range(len(unique_left)): + labels = lookup_target[n, right_inverse] + lookup = labels >= 0 + row_locations[n] = permuted_right_inds[lookup] + final_relevant_labels[n] = labels[lookup] + + charge_labels = np.concatenate( + [final_relevant_labels[n] for n in left_inverse]) + tmp_inds = [ + permuted_left_inds[n] + row_locations[left_inverse[n]] + for n in range(len(left_charges)) + ] + try: + inds = np.concatenate(tmp_inds) + except ValueError: + inds = np.asarray(tmp_inds) + + else: + row_locations = [None] * len(unique_left) + final_relevant_labels = [None] * len(unique_left) + for n in range(len(unique_left)): + labels = lookup_target[n, right_inverse] + lookup = labels >= 0 + row_locations[n] = np.nonzero(lookup)[0] + final_relevant_labels[n] = labels[lookup] + charge_labels = np.concatenate( + [final_relevant_labels[n] for n in left_inverse]) + + inds = np.concatenate([ + n * len(right_charges) + row_locations[left_inverse[n]] + for n in range(len(left_charges)) + ]) + obj = charges[0].__new__(type(charges[0])) + obj.__init__(relevant_charges.unique_charges, charge_labels, + charges[0].charge_types) + return obj, inds + + else: + final_relevant_labels = [None] * len(unique_left) + for n in range(len(unique_left)): + labels = lookup_target[n, right_inverse] + lookup = labels >= 0 + final_relevant_labels[n] = labels[lookup] + charge_labels = np.concatenate( + [final_relevant_labels[n] for n in left_inverse]) + obj = charges[0].__new__(type(charges[0])) + obj.__init__(relevant_charges.unique_charges, charge_labels, + charges[0].charge_types) + return obj + + +def find_sparse_positions_new(charges: List[BaseCharge], + flows: List[Union[int, bool]], + target_charges: BaseCharge, + strides: Optional[np.ndarray] = None, + store_dual: Optional[bool] = False) -> np.ndarray: + """ + Find the dense locations of elements (i.e. the index-values within the DENSE tensor) + in the vector of `fused_charges` resulting from fusing all elements of `charges` + that have a value of `target_charge`. + For example, given + ``` + charges = [[-2,0,1,0,0],[-1,0,2,1]] + target_charge = 0 + fused_charges = fuse_charges(charges,[1,1]) + print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] + ``` + we want to find the index-positions of charges + that fuse to `target_charge=0`, i.e. where `fused_charges==0`, + within the dense array. As one additional wrinkle, `charges` + is a subset of the permuted charges of a tensor with rank R > len(charges), + and `stride_arrays` are their corresponding range of strides, i.e. + + ``` + R=5 + D = [2,3,4,5,6] + tensor_flows = np.random.randint(-1,2,R) + tensor_charges = [np.random.randing(-5,5,D[n]) for n in range(R)] + order = np.arange(R) + np.random.shuffle(order) + tensor_strides = [360, 120, 30, 6, 1] + + charges = [tensor_charges[order[n]] for n in range(3)] + flows = [tensor_flows[order[n]] for n in range(len(3))] + strides = [tensor_stride[order[n]] for n in range(3)] + _ = _find_transposed_dense_positions(charges, flows, 0, strides) + + ``` + `_find_transposed_dense_blocks` returns an np.ndarray containing the + index-positions of these elements calculated using `stride_arrays`. + The result only makes sense in conjuction with the complementary + data computed from the complementary + elements in`tensor_charges`, + `tensor_strides` and `tensor_flows`. + This routine is mainly used in `_find_diagonal_dense_blocks`. + + Args: + charges: A list of BaseCharge or ChargeCollection. + flows: The flow directions of the `charges`. + target_charge: The target charge. + strides: The strides for the `charges` subset. + if `None`, natural stride ordering is assumed. + + Returns: + np.ndarray: The index-positions within the dense data array + of the elements fusing to `target_charge`. + """ + + _check_flows(flows) + if len(charges) == 1: + fused_charges = charges[0] * flows[0] + unique, inverse = fused_charges.unique(return_inverse=True) + common, label_to_unique, label_to_target = unique.intersect( + target_charges, return_indices=True) + inds = np.nonzero(np.isin(inverse, label_to_unique))[0] + if strides is not None: + permuted_inds = strides[0] * np.arange(len(charges[0])) + return fused_charges[permuted_inds[inds]], inds + + return fused_charges[inds], inds + + partition = _find_best_partition([len(c) for c in charges]) + left_charges = fuse_charges(charges[:partition], flows[:partition]) + right_charges = fuse_charges(charges[partition:], flows[partition:]) + + # unique_target_charges, inds = target_charges.unique(return_index=True) + # target_charges = target_charges[np.sort(inds)] + unique_left, left_inverse = left_charges.unique(return_inverse=True) + unique_right, right_inverse, right_degens = right_charges.unique( + return_inverse=True, return_counts=True) + + fused = unique_left + unique_right + + unique_fused, labels_fused = fused.unique(return_inverse=True) + + relevant_charges, label_to_unique_fused, label_to_target = unique_fused.intersect( + target_charges, return_indices=True) + + relevant_fused_positions = np.nonzero( + np.isin(labels_fused, label_to_unique_fused))[0] + relevant_left_labels, relevant_right_labels = np.divmod( + relevant_fused_positions, len(unique_right)) + rel_l_labels = np.unique(relevant_left_labels) + total_degen = { + t: np.sum(right_degens[relevant_right_labels[relevant_left_labels == t]]) + for t in rel_l_labels + } + + relevant_left_inverse = left_inverse[np.isin(left_inverse, rel_l_labels)] + degeneracy_vector = np.empty(len(relevant_left_inverse), dtype=np.uint32) + row_locations = [None] * len(unique_left) + final_relevant_labels = [None] * len(unique_left) + for n in range(len(relevant_left_labels)): + degeneracy_vector[relevant_left_inverse == + relevant_left_labels[n]] = total_degen[ + relevant_left_labels[n]] + start_positions = np.cumsum(degeneracy_vector) - degeneracy_vector + tmp = np.full(len(unique_fused), fill_value=-1, dtype=np.int16) + tmp[label_to_unique_fused] = np.arange( + len(label_to_unique_fused), dtype=np.int16) + lookup_target = tmp[labels_fused].reshape( + [len(unique_left), len(unique_right)]) + + final_relevant_labels = [None] * len(unique_left) + for n in range(len(rel_l_labels)): + labels = lookup_target[rel_l_labels[n], right_inverse] + lookup = labels >= 0 + final_relevant_labels[rel_l_labels[n]] = labels[lookup] + charge_labels = np.concatenate( + [final_relevant_labels[n] for n in relevant_left_inverse]) + inds = np.concatenate([ + start_positions[n] + + np.arange(total_degen[relevant_left_inverse[n]], dtype=np.uint32) + for n in range(len(relevant_left_inverse)) + ]) + + return relevant_charges[charge_labels], inds + + +def find_sparse_positions(charges: List[BaseCharge], + flows: List[Union[int, bool]], + target_charges: BaseCharge) -> Dict: + """ + Find the sparse locations of elements (i.e. the index-values within + the SPARSE tensor) in the vector `fused_charges` (resulting from + fusing `left_charges` and `right_charges`) + that have a value of `target_charges`, assuming that all elements + different from `target_charges` are `0`. + For example, given + ``` + left_charges = [-2,0,1,0,0] + right_charges = [-1,0,2,1] + target_charges = [0,1] + fused_charges = fuse_charges([left_charges, right_charges],[1,1]) + print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] + ``` 0 1 2 3 4 5 6 7 8 + we want to find the all different blocks + that fuse to `target_charges=[0,1]`, i.e. where `fused_charges==0` or `1`, + together with their corresponding sparse index-values of the data in the sparse array, + assuming that all elements in `fused_charges` different from `target_charges` are 0. + + `find_sparse_blocks` returns a dict mapping integers `target_charge` + to an array of integers denoting the sparse locations of elements within + `fused_charges`. + For the above example, we get: + * `target_charge=0`: [0,1,3,5,7] + * `target_charge=1`: [2,4,6,8] + Args: + charges: An np.ndarray of integer charges. + flows: The flow direction of the left charges. + target_charges: The target charges. + Returns: + dict: Mapping integers to np.ndarray of integers. + """ + _check_flows(flows) + if len(charges) == 1: + fused_charges = charges[0] * flows[0] + unique_charges = fused_charges.unique() + target_charges = target_charges.unique() + relevant_target_charges = unique_charges.intersect(target_charges) + relevant_fused_charges = fused_charges[fused_charges.isin( + relevant_target_charges)] + return { + c: np.nonzero(relevant_fused_charges == c)[0] + for c in relevant_target_charges + } + partition = _find_best_partition([len(c) for c in charges]) + left_charges = fuse_charges(charges[:partition], flows[:partition]) + right_charges = fuse_charges(charges[partition:], flows[partition:]) + + # unique_target_charges, inds = target_charges.unique(return_index=True) + # target_charges = target_charges[np.sort(inds)] + unique_left, left_inverse = left_charges.unique(return_inverse=True) + unique_right, right_inverse, right_dims = right_charges.unique( + return_inverse=True, return_counts=True) + + fused_unique = unique_left + unique_right + unique_inds = np.nonzero(fused_unique == target_charges) + relevant_positions = unique_inds[0].astype(np.uint32) + tmp_inds_left, tmp_inds_right = np.divmod(relevant_positions, + len(unique_right)) + + relevant_unique_left_inds = np.unique(tmp_inds_left) + left_lookup = np.empty(np.max(relevant_unique_left_inds) + 1, dtype=np.uint32) + left_lookup[relevant_unique_left_inds] = np.arange( + len(relevant_unique_left_inds)) + relevant_unique_right_inds = np.unique(tmp_inds_right) + right_lookup = np.empty( + np.max(relevant_unique_right_inds) + 1, dtype=np.uint32) + right_lookup[relevant_unique_right_inds] = np.arange( + len(relevant_unique_right_inds)) + + left_charge_labels = np.nonzero( + np.expand_dims(left_inverse, 1) == np.expand_dims( + relevant_unique_left_inds, 0)) + relevant_left_inverse = np.arange(len(left_charge_labels[0])) + + right_charge_labels = np.expand_dims(right_inverse, 1) == np.expand_dims( + relevant_unique_right_inds, 0) + right_block_information = {} + for n in relevant_unique_left_inds: + ri = np.nonzero((unique_left[n] + unique_right).isin(target_charges))[0] + tmp_inds = np.nonzero(right_charge_labels[:, right_lookup[ri]]) + right_block_information[n] = [ri, np.arange(len(tmp_inds[0])), tmp_inds[1]] + + relevant_right_inverse = np.arange(len(right_charge_labels[0])) + + #generate a degeneracy vector which for each value r in relevant_right_charges + #holds the corresponding number of non-zero elements `relevant_right_charges` + #that can add up to `target_charges`. + degeneracy_vector = np.empty(len(left_charge_labels[0]), dtype=np.uint32) + for n in range(len(relevant_unique_left_inds)): + degeneracy_vector[relevant_left_inverse[ + left_charge_labels[1] == n]] = np.sum(right_dims[tmp_inds_right[ + tmp_inds_left == relevant_unique_left_inds[n]]]) + + start_positions = (np.cumsum(degeneracy_vector) - degeneracy_vector).astype( + np.uint32) + out = {} + for n in range(len(target_charges)): + block = [] + if len(unique_inds) > 1: + lis, ris = np.divmod(unique_inds[0][unique_inds[1] == n], + len(unique_right)) + else: + lis, ris = np.divmod(unique_inds[0], len(unique_right)) + + for m in range(len(lis)): + ri_tmp, arange, tmp_inds = right_block_information[lis[m]] + block.append( + np.add.outer( + start_positions[relevant_left_inverse[left_charge_labels[1] == + left_lookup[lis[m]]]], + arange[tmp_inds == np.nonzero( + ri_tmp == ris[m])[0]]).ravel().astype(np.uint32)) + out[target_charges[n]] = np.concatenate(block) + return out + + +# def find_dense_positions(charges: List[Union[BaseCharge, ChargeCollection]], +# flows: List[Union[int, bool]], +# target_charges: Union[BaseCharge, ChargeCollection], +# strides: Optional[np.ndarray] = None, +# store_dual: Optional[bool] = False) -> Dict: +# """ +# Find the dense locations of elements (i.e. the index-values within the DENSE tensor) +# in the vector of `fused_charges` resulting from fusing all elements of `charges` +# that have a value of `target_charge`. +# For example, given +# ``` +# charges = [[-2,0,1,0,0],[-1,0,2,1]] +# target_charge = 0 +# fused_charges = fuse_charges(charges,[1,1]) +# print(fused_charges) # [-3,-2,0,-1,-1,0,2,1,0,1,3,2,-1,0,2,1,-1,0,2,1] +# ``` +# we want to find the index-positions of charges +# that fuse to `target_charge=0`, i.e. where `fused_charges==0`, +# within the dense array. As one additional wrinkle, `charges` +# is a subset of the permuted charges of a tensor with rank R > len(charges), +# and `stride_arrays` are their corresponding range of strides, i.e. + +# ``` +# R=5 +# D = [2,3,4,5,6] +# tensor_flows = np.random.randint(-1,2,R) +# tensor_charges = [np.random.randing(-5,5,D[n]) for n in range(R)] +# order = np.arange(R) +# np.random.shuffle(order) +# tensor_strides = [360, 120, 30, 6, 1] + +# charges = [tensor_charges[order[n]] for n in range(3)] +# flows = [tensor_flows[order[n]] for n in range(len(3))] +# strides = [tensor_stride[order[n]] for n in range(3)] +# _ = _find_transposed_dense_positions(charges, flows, 0, strides) + +# ``` +# `_find_transposed_dense_blocks` returns an np.ndarray containing the +# index-positions of these elements calculated using `stride_arrays`. +# The result only makes sense in conjuction with the complementary +# data computed from the complementary +# elements in`tensor_charges`, +# `tensor_strides` and `tensor_flows`. +# This routine is mainly used in `_find_diagonal_dense_blocks`. + +# Args: +# charges: A list of BaseCharge or ChargeCollection. +# flows: The flow directions of the `charges`. +# target_charge: The target charge. +# strides: The strides for the `charges` subset. +# if `None`, natural stride ordering is assumed. + +# Returns: +# dict +# """ + +# _check_flows(flows) +# out = {} +# if store_dual: +# store_charges = target_charges * (-1) +# else: +# store_charges = target_charges + +# if len(charges) == 1: +# fused_charges = charges[0] * flows[0] +# inds = np.nonzero(fused_charges == target_charges) +# if len(target_charges) > 1: +# for n in range(len(target_charges)): +# i = inds[0][inds[1] == n] +# if len(i) == 0: +# continue +# if strides is not None: +# permuted_inds = strides[0] * np.arange(len(charges[0])) +# out[store_charges.get_item(n)] = permuted_inds[i] +# else: +# out[store_charges.get_item(n)] = i +# return out +# else: +# if strides is not None: +# permuted_inds = strides[0] * np.arange(len(charges[0])) +# out[store_charges.get_item(n)] = permuted_inds[inds[0]] +# else: +# out[store_charges.get_item(n)] = inds[0] +# return out + +# partition = _find_best_partition([len(c) for c in charges]) +# left_charges = fuse_charges(charges[:partition], flows[:partition]) +# right_charges = fuse_charges(charges[partition:], flows[partition:]) +# if strides is not None: +# stride_arrays = [ +# np.arange(len(charges[n])) * strides[n] for n in range(len(charges)) +# ] +# permuted_left_inds = fuse_ndarrays(stride_arrays[0:partition]) +# permuted_right_inds = fuse_ndarrays(stride_arrays[partition:]) + +# # unique_target_charges, inds = target_charges.unique(return_index=True) +# # target_charges = target_charges[np.sort(inds)] +# unique_left, left_inverse = left_charges.unique(return_inverse=True) +# unique_right, right_inverse = right_charges.unique(return_inverse=True) + +# fused_unique = unique_left + unique_right +# unique_inds = np.nonzero(fused_unique == target_charges) + +# relevant_positions = unique_inds[0] +# tmp_inds_left, tmp_inds_right = np.divmod(relevant_positions, +# len(unique_right)) + +# relevant_unique_left_inds = np.unique(tmp_inds_left) +# left_lookup = np.empty(np.max(relevant_unique_left_inds) + 1, dtype=np.uint32) +# left_lookup[relevant_unique_left_inds] = np.arange( +# len(relevant_unique_left_inds)) +# relevant_unique_right_inds = np.unique(tmp_inds_right) +# right_lookup = np.empty( +# np.max(relevant_unique_right_inds) + 1, dtype=np.uint32) +# right_lookup[relevant_unique_right_inds] = np.arange( +# len(relevant_unique_right_inds)) + +# left_charge_labels = np.nonzero( +# np.expand_dims(left_inverse, 1) == np.expand_dims( +# relevant_unique_left_inds, 0)) +# right_charge_labels = np.nonzero( +# np.expand_dims(right_inverse, 1) == np.expand_dims( +# relevant_unique_right_inds, 0)) + +# len_right = len(right_charges) + +# for n in range(len(target_charges)): +# if len(unique_inds) > 1: +# lis, ris = np.divmod(unique_inds[0][unique_inds[1] == n], +# len(unique_right)) +# else: +# lis, ris = np.divmod(unique_inds[0], len(unique_right)) +# dense_positions = [] +# left_positions = [] +# lookup = [] +# for m in range(len(lis)): +# li = lis[m] +# ri = ris[m] +# dense_left_positions = (left_charge_labels[0][ +# left_charge_labels[1] == left_lookup[li]]).astype(np.uint32) +# dense_right_positions = (right_charge_labels[0][ +# right_charge_labels[1] == right_lookup[ri]]).astype(np.uint32) +# if strides is None: +# positions = np.expand_dims(dense_left_positions * len_right, +# 1) + np.expand_dims(dense_right_positions, 0) +# else: +# positions = np.expand_dims( +# permuted_left_inds[dense_left_positions], 1) + np.expand_dims( +# permuted_right_inds[dense_right_positions], 0) + +# dense_positions.append(positions) +# left_positions.append(dense_left_positions) +# lookup.append( +# np.stack([ +# np.arange(len(dense_left_positions), dtype=np.uint32), +# np.full(len(dense_left_positions), fill_value=m, dtype=np.uint32) +# ], +# axis=1)) + +# if len(lookup) > 0: +# ind_sort = np.argsort(np.concatenate(left_positions)) +# it = np.concatenate(lookup, axis=0) +# table = it[ind_sort, :] +# out[store_charges.get_item(n)] = np.concatenate([ +# dense_positions[table[n, 1]][table[n, 0], :].astype(np.uint32) +# for n in range(table.shape[0]) +# ]) +# else: +# out[store_charges.get_item(n)] = np.array([]) + +# return out + +# def _find_diagonal_sparse_blocks_old(charges: List[BaseCharge], +# flows: List[Union[bool, int]], +# partition: int) -> Tuple[BaseCharge, List]: +# """ +# Given the meta data and underlying data of a symmetric matrix, compute +# all diagonal blocks and return them in a dict. +# `row_charges` and `column_charges` are lists of np.ndarray. The tensor +# is viewed as a matrix with rows given by fusing `row_charges` and +# columns given by fusing `column_charges`. + +# Args: +# charges: A list of charges. +# flows: A list of flows. +# partition: The location of the partition of `charges` into rows and colums. +# Returns: +# return common_charges, blocks, start_positions, row_locations, column_degeneracies +# List[Union[BaseCharge, ChargeCollection]]: A list of unique charges, one per block. +# List[np.ndarray]: A list containing the blocks. +# """ +# _check_flows(flows) +# if len(flows) != len(charges): +# raise ValueError("`len(flows)` is different from `len(charges)`") +# row_charges = charges[:partition] +# row_flows = flows[:partition] +# column_charges = charges[partition:] +# column_flows = flows[partition:] + +# #get the unique column-charges +# #we only care about their degeneracies, not their order; that's much faster +# #to compute since we don't have to fuse all charges explicitly +# #`compute_fused_charge_degeneracies` multiplies flows into the column_charges +# unique_column_charges, column_dims = compute_fused_charge_degeneracies( +# column_charges, column_flows) +# unique_row_charges = compute_unique_fused_charges(row_charges, row_flows) +# #get the charges common to rows and columns (only those matter) +# common_charges, label_to_row, label_to_column = unique_row_charges.intersect( +# unique_column_charges * True, return_indices=True) + +# #convenience container for storing the degeneracies of each +# #column charge +# #column_degeneracies = dict(zip(unique_column_charges, column_dims)) +# #column_degeneracies = dict(zip(unique_column_charges * True, column_dims)) +# print(common_charges) +# row_locations = find_sparse_positions( +# charges=row_charges, flows=row_flows, target_charges=common_charges) + +# degeneracy_vector = np.empty( +# np.sum([len(v) for v in row_locations.values()]), dtype=np.uint32) +# #for each charge `c` in `common_charges` we generate a boolean mask +# #for indexing the positions where `relevant_column_charges` has a value of `c`. +# for c in common_charges: +# degeneracy_vector[row_locations[c]] = column_degeneracies[c] + +# start_positions = (np.cumsum(degeneracy_vector) - degeneracy_vector).astype( +# np.uint32) +# blocks = [] + +# for c in common_charges: +# #numpy broadcasting is substantially faster than kron! +# rlocs = row_locations[c] +# rlocs.sort() #sort in place (we need it again later) +# cdegs = column_degeneracies[c] +# inds = np.ravel(np.add.outer(start_positions[rlocs], np.arange(cdegs))) +# blocks.append([inds, (len(rlocs), cdegs)]) +# return common_charges, blocks diff --git a/tensornetwork/block_tensor/block_tensor_test.py b/tensornetwork/block_tensor/block_tensor_test.py index e862d811e..2c8ada9c0 100644 --- a/tensornetwork/block_tensor/block_tensor_test.py +++ b/tensornetwork/block_tensor/block_tensor_test.py @@ -1,8 +1,9 @@ import numpy as np import pytest -# pylint: disable=line-too-long -from tensornetwork.block_tensor.block_tensor import BlockSparseTensor, compute_num_nonzero, compute_dense_to_sparse_table, find_sparse_positions, find_dense_positions, map_to_integer -from index import Index, fuse_charges + +from tensornetwork.block_tensor.charge import U1Charge, fuse_charges +from tensornetwork.block_tensor.index import Index +from tensornetwork.block_tensor.block_tensor import compute_num_nonzero, reduce_charges, BlockSparseTensor, fuse_ndarrays, tensordot np_dtypes = [np.float32, np.float16, np.float64, np.complex64, np.complex128] @@ -12,10 +13,10 @@ def test_block_sparse_init(dtype): D = 10 #bond dimension B = 10 #number of blocks rank = 4 - flows = np.asarray([1 for _ in range(rank)]) - flows[-2::] = -1 + flows = np.asarray([False for _ in range(rank)]) + flows[-2::] = True charges = [ - np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + U1Charge(np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16)) for _ in range(rank) ] indices = [ @@ -32,126 +33,278 @@ def test_block_sparse_init(dtype): assert len(A.data) == num_elements -def test_block_table(): - D = 30 #bond dimension - B = 4 #number of blocks - dtype = np.int16 #the dtype of the quantum numbers - rank = 4 - flows = np.asarray([1 for _ in range(rank)]) - flows[-2::] = -1 - charges = [ - np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - for _ in range(rank) - ] - indices = [ - Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) - for n in range(rank) - ] - num_non_zero = compute_num_nonzero([i.charges for i in indices], - [i.flow for i in indices]) +def test_find_dense_positions(): + left_charges = np.asarray([-2, 0, 1, 0, 0]).astype(np.int16) + right_charges = np.asarray([-1, 0, 2, 1]).astype(np.int16) + target_charge = np.zeros((1, 1), dtype=np.int16) + fused_charges = fuse_ndarrays([left_charges, right_charges]) + dense_positions = reduce_charges( + [U1Charge(left_charges), U1Charge(right_charges)], [False, False], + target_charge, + return_locations=True) + np.testing.assert_allclose( + dense_positions[1], + np.nonzero(fused_charges == target_charge[0, 0])[0]) - inds = compute_dense_to_sparse_table( - charges=charges, flows=flows, target_charge=0) - total = flows[0] * charges[0][inds[0]] + flows[1] * charges[1][ - inds[1]] + flows[2] * charges[2][inds[2]] + flows[3] * charges[3][inds[3]] - assert len(total) == len(np.nonzero(total == 0)[0]) - assert len(total) == num_non_zero +def test_transpose(): + R = 4 + Ds = [20, 3, 4, 5] + final_order = np.arange(R) + np.random.shuffle(final_order) + charges = [U1Charge(np.random.randint(-5, 5, Ds[n])) for n in range(R)] + flows = np.full(R, fill_value=False, dtype=np.bool) + indices = [Index(charges[n], flows[n]) for n in range(R)] + A = BlockSparseTensor.random(indices=indices) + Adense = A.todense() + dense_res = np.transpose(Adense, final_order) + A.transpose(final_order) + np.testing.assert_allclose(dense_res, A.todense()) -def test_find_dense_positions(): - left_charges = [-2, 0, 1, 0, 0] - right_charges = [-1, 0, 2, 1] - target_charge = 0 - fused_charges = fuse_charges([left_charges, right_charges], [1, 1]) - blocks = find_dense_positions(left_charges, 1, right_charges, 1, - target_charge) - np.testing.assert_allclose(blocks[(-2, 2)], [2]) - np.testing.assert_allclose(blocks[(0, 0)], [5, 13, 17]) - np.testing.assert_allclose(blocks[(1, -1)], [8]) - - -def test_find_dense_positions_2(): - D = 40 #bond dimension - B = 4 #number of blocks - dtype = np.int16 #the dtype of the quantum numbers - rank = 4 - flows = np.asarray([1 for _ in range(rank)]) - flows[-2::] = -1 - charges = [ - np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - for _ in range(rank) - ] - indices = [ - Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) - for n in range(rank) - ] - n1 = compute_num_nonzero([i.charges for i in indices], - [i.flow for i in indices]) - row_charges = fuse_charges([indices[n].charges for n in range(rank // 2)], - [1 for _ in range(rank // 2)]) - column_charges = fuse_charges( - [indices[n].charges for n in range(rank // 2, rank)], - [1 for _ in range(rank // 2, rank)]) - - i01 = indices[0] * indices[1] - i23 = indices[2] * indices[3] - blocks = find_dense_positions(i01.charges, 1, i23.charges, 1, 0) - assert sum([len(v) for v in blocks.values()]) == n1 - - tensor = BlockSparseTensor.random(indices=indices, dtype=np.float64) - tensor.reshape((D * D, D * D)) - blocks_2 = tensor.get_diagonal_blocks(return_data=False) - np.testing.assert_allclose([k[0] for k in blocks.keys()], - list(blocks_2.keys())) - for c in blocks.keys(): - assert np.prod(blocks_2[c[0]][1]) == len(blocks[c]) - - -def test_find_sparse_positions(): - D = 40 #bond dimension - B = 4 #number of blocks - dtype = np.int16 #the dtype of the quantum numbers - rank = 4 - flows = np.asarray([1 for _ in range(rank)]) - flows[-2::] = -1 - charges = [ - np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - for _ in range(rank) + +def test_tensordot(): + R = 4 + DsA = [10, 12, 14, 16] + DsB = [14, 16, 18, 20] + chargesA = [U1Charge(np.random.randint(-5, 5, DsA[n])) for n in range(R // 2)] + commoncharges = [ + U1Charge(np.random.randint(-5, 5, DsA[n + R // 2])) for n in range(R // 2) ] - indices = [ - Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) - for n in range(rank) + chargesB = [ + U1Charge(np.random.randint(-5, 5, DsB[n + R // 2])) for n in range(R // 2) ] - n1 = compute_num_nonzero([i.charges for i in indices], - [i.flow for i in indices]) - row_charges = fuse_charges([indices[n].charges for n in range(rank // 2)], - [1 for _ in range(rank // 2)]) - column_charges = fuse_charges( - [indices[n].charges for n in range(rank // 2, rank)], - [1 for _ in range(rank // 2, rank)]) - - i01 = indices[0] * indices[1] - i23 = indices[2] * indices[3] - unique_row_charges = np.unique(i01.charges) - unique_column_charges = np.unique(i23.charges) - common_charges = np.intersect1d( - unique_row_charges, -unique_column_charges, assume_unique=True) - blocks = find_sparse_positions( - i01.charges, 1, i23.charges, 1, target_charges=[0]) - assert sum([len(v) for v in blocks.values()]) == n1 - np.testing.assert_allclose(np.sort(blocks[0]), np.arange(n1)) - - -def test_map_to_integer(): - dims = [4, 3, 2] - dim_prod = [6, 2, 1] - N = 10 - table = np.stack([np.random.randint(0, d, N) for d in dims], axis=1) - integers = map_to_integer(dims, table) - ints = [] - for n in range(N): - i = 0 - for d in range(len(dims)): - i += dim_prod[d] * table[n, d] - ints.append(i) - np.testing.assert_allclose(ints, integers) + indsA = np.random.choice(np.arange(R), R // 2, replace=False) + indsB = np.random.choice(np.arange(R), R // 2, replace=False) + flowsA = np.full(R, False, dtype=np.bool) + flowsB = np.full(R, False, dtype=np.bool) + + flowsB[indsB] = True + indicesA = [None for _ in range(R)] + indicesB = [None for _ in range(R)] + for n in range(len(indsA)): + indicesA[indsA[n]] = Index(commoncharges[n], flowsA[indsA[n]]) + indicesB[indsB[n]] = Index(commoncharges[n], flowsB[indsB[n]]) + compA = list(set(np.arange(R)) - set(indsA)) + compB = list(set(np.arange(R)) - set(indsB)) + + for n in range(len(compA)): + indicesA[compA[n]] = Index(chargesA[n], flowsA[compA[n]]) + indicesB[compB[n]] = Index(chargesB[n], flowsB[compB[n]]) + indices_final = [] + for n in sorted(compA): + indices_final.append(indicesA[n]) + for n in sorted(compB): + indices_final.append(indicesB[n]) + shapes = tuple([i.dimension for i in indices_final]) + A = BlockSparseTensor.random(indices=indicesA) + B = BlockSparseTensor.random(indices=indicesB) + + final_order = np.arange(R) + np.random.shuffle(final_order) + Adense = A.todense() + Bdense = B.todense() + dense_res = np.transpose( + np.tensordot(Adense, Bdense, (indsA, indsB)), final_order) + + res = tensordot(A, B, (indsA, indsB), final_order=final_order) + np.testing.assert_allclose(dense_res, res.todense()) + + +# def test_find_dense_positions_2(): +# D = 40 #bond dimension +# B = 4 #number of blocks +# dtype = np.int16 #the dtype of the quantum numbers +# rank = 4 +# flows = np.asarray([1 for _ in range(rank)]) +# flows[-2::] = -1 +# charges = [ +# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) +# for _ in range(rank) +# ] +# indices = [ +# Index( +# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) +# for n in range(rank) +# ] +# n1 = compute_num_nonzero([i.charges for i in indices], +# [i.flow for i in indices]) + +# i01 = indices[0] * indices[1] +# i23 = indices[2] * indices[3] +# positions = find_dense_positions([i01.charges, i23.charges], [1, 1], +# U1Charge(np.asarray([0]))) +# assert len(positions[0]) == n1 + +# def test_find_sparse_positions(): +# D = 40 #bond dimension +# B = 4 #number of blocks +# dtype = np.int16 #the dtype of the quantum numbers +# rank = 4 +# flows = np.asarray([1 for _ in range(rank)]) +# flows[-2::] = -1 +# charges = [ +# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) +# for _ in range(rank) +# ] +# indices = [ +# Index( +# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) +# for n in range(rank) +# ] +# n1 = compute_num_nonzero([i.charges for i in indices], +# [i.flow for i in indices]) +# i01 = indices[0] * indices[1] +# i23 = indices[2] * indices[3] +# unique_row_charges = np.unique(i01.charges.charges) +# unique_column_charges = np.unique(i23.charges.charges) +# common_charges = np.intersect1d( +# unique_row_charges, -unique_column_charges, assume_unique=True) +# blocks = find_sparse_positions([i01.charges, i23.charges], [1, 1], +# target_charges=U1Charge(np.asarray([0]))) +# assert sum([len(v) for v in blocks.values()]) == n1 +# np.testing.assert_allclose(np.sort(blocks[0]), np.arange(n1)) + +# def test_find_sparse_positions_2(): +# D = 1000 #bond dimension +# B = 4 #number of blocks +# dtype = np.int16 #the dtype of the quantum numbers +# charges = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) +# index = Index(charges=U1Charge(charges), flow=1, name='index0') +# targets = np.asarray([-1, 0, 1]) +# blocks = find_sparse_positions([index.charges], [index.flow], +# target_charges=U1Charge(targets)) + +# inds = np.isin(charges, targets) +# relevant_charges = charges[inds] +# blocks_ = {t: np.nonzero(relevant_charges == t)[0] for t in targets} +# assert np.all( +# np.asarray(list(blocks.keys())) == np.asarray(list(blocks_.keys()))) +# for k in blocks.keys(): +# assert np.all(blocks[k] == blocks_[k]) + +# def test_find_sparse_positions_3(): +# D = 40 #bond dimension +# B = 4 #number of blocks +# dtype = np.int16 #the dtype of the quantum numbers +# flows = [1, -1] + +# rank = len(flows) +# charges = [ +# np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) +# for _ in range(rank) +# ] +# indices = [ +# Index( +# charges=U1Charge(charges[n]), flow=flows[n], name='index{}'.format(n)) +# for n in range(rank) +# ] +# i1, i2 = indices +# common_charges = np.intersect1d(i1.charges.charges, i2.charges.charges) +# row_locations = find_sparse_positions( +# charges=[i1.charges, i2.charges], +# flows=flows, +# target_charges=U1Charge(common_charges)) +# fused = (i1 * i2).charges +# relevant = fused.charges[np.isin(fused.charges, common_charges)] +# for k, v in row_locations.items(): +# np.testing.assert_allclose(np.nonzero(relevant == k)[0], np.sort(v)) + +# # def test_dense_transpose(): +# # Ds = [10, 11, 12] #bond dimension +# # rank = len(Ds) +# # flows = np.asarray([1 for _ in range(rank)]) +# # flows[-2::] = -1 +# # charges = [U1Charge(np.zeros(Ds[n], dtype=np.int16)) for n in range(rank)] +# # indices = [ +# # Index(charges=charges[n], flow=flows[n], name='index{}'.format(n)) +# # for n in range(rank) +# # ] +# # A = BlockSparseTensor.random(indices=indices, dtype=np.float64) +# # B = np.transpose(np.reshape(A.data.copy(), Ds), (1, 0, 2)) +# # A.transpose((1, 0, 2)) +# # np.testing.assert_allclose(A.data, B.flat) + +# # B = np.transpose(np.reshape(A.data.copy(), [11, 10, 12]), (1, 0, 2)) +# # A.transpose((1, 0, 2)) + +# # np.testing.assert_allclose(A.data, B.flat) + +# @pytest.mark.parametrize("R", [1, 2]) +# def test_find_diagonal_dense_blocks(R): +# rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] +# cs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] +# charges = rs + cs + +# left_fused = fuse_charges(charges[0:R], [1] * R) +# right_fused = fuse_charges(charges[R:], [1] * R) +# left_unique = left_fused.unique() +# right_unique = right_fused.unique() +# zero = left_unique.zero_charge +# blocks = {} +# rdim = len(right_fused) +# for lu in left_unique: +# linds = np.nonzero(left_fused == lu)[0] +# rinds = np.nonzero(right_fused == lu * (-1))[0] +# if (len(linds) > 0) and (len(rinds) > 0): +# blocks[lu] = fuse_ndarrays([linds * rdim, rinds]) +# comm, blocks_ = _find_diagonal_dense_blocks(rs, cs, [1] * R, [1] * R) +# for n in range(len(comm)): +# assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) + +# # #@pytest.mark.parametrize("dtype", np_dtypes) +# # def test_find_diagonal_dense_blocks_2(): +# # R = 1 +# # rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] +# # cs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] +# # charges = rs + cs + +# # left_fused = fuse_charges(charges[0:R], [1] * R) +# # right_fused = fuse_charges(charges[R:], [1] * R) +# # left_unique = left_fused.unique() +# # right_unique = right_fused.unique() +# # zero = left_unique.zero_charge +# # blocks = {} +# # rdim = len(right_fused) +# # for lu in left_unique: +# # linds = np.nonzero(left_fused == lu)[0] +# # rinds = np.nonzero(right_fused == lu * (-1))[0] +# # if (len(linds) > 0) and (len(rinds) > 0): +# # blocks[lu] = fuse_ndarrays([linds * rdim, rinds]) +# # comm, blocks_ = _find_diagonal_dense_blocks(rs, cs, [1] * R, [1] * R) +# # for n in range(len(comm)): +# # assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) + +# @pytest.mark.parametrize("R", [1, 2]) +# def test_find_diagonal_dense_blocks_transposed(R): +# order = np.arange(2 * R) +# np.random.shuffle(order) +# rs = [U1Charge(np.random.randint(-4, 4, 50)) for _ in range(R)] +# cs = [U1Charge(np.random.randint(-4, 4, 40)) for _ in range(R)] +# charges = rs + cs +# dims = np.asarray([len(c) for c in charges]) +# strides = np.flip(np.append(1, np.cumprod(np.flip(dims[1::])))) +# stride_arrays = [np.arange(dims[n]) * strides[n] for n in range(2 * R)] + +# left_fused = fuse_charges([charges[n] for n in order[0:R]], [1] * R) +# right_fused = fuse_charges([charges[n] for n in order[R:]], [1] * R) +# lstrides = fuse_ndarrays([stride_arrays[n] for n in order[0:R]]) +# rstrides = fuse_ndarrays([stride_arrays[n] for n in order[R:]]) + +# left_unique = left_fused.unique() +# right_unique = right_fused.unique() +# blocks = {} +# rdim = len(right_fused) +# for lu in left_unique: +# linds = np.nonzero(left_fused == lu)[0] +# rinds = np.nonzero(right_fused == lu * (-1))[0] +# if (len(linds) > 0) and (len(rinds) > 0): +# tmp = fuse_ndarrays([linds * rdim, rinds]) +# blocks[lu] = _find_values_in_fused(tmp, lstrides, rstrides) + +# comm, blocks_ = _find_diagonal_dense_blocks([charges[n] for n in order[0:R]], +# [charges[n] for n in order[R:]], +# [1] * R, [1] * R, +# row_strides=strides[order[0:R]], +# column_strides=strides[order[R:]]) +# for n in range(len(comm)): +# assert np.all(blocks[comm.charges[n]] == blocks_[n][0]) diff --git a/tensornetwork/block_tensor/charge.py b/tensornetwork/block_tensor/charge.py new file mode 100644 index 000000000..760453b46 --- /dev/null +++ b/tensornetwork/block_tensor/charge.py @@ -0,0 +1,887 @@ +# Copyright 2019 The TensorNetwork Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +import numpy as np +from tensornetwork.network_components import Node, contract, contract_between +# pylint: disable=line-too-long +from tensornetwork.backends import backend_factory +import copy +import warnings +from typing import List, Union, Any, Optional, Tuple, Text, Iterable, Type + + +class BaseCharge: + + def __init__(self, + charges: np.ndarray, + charge_labels: Optional[np.ndarray] = None, + charge_types: Optional[List[Type["BaseCharge"]]] = None) -> None: + self.charge_types = charge_types + if charges.ndim == 1: + charges = np.expand_dims(charges, 0) + if charge_labels is None: + self.unique_charges, self.charge_labels = np.unique( + charges.astype(np.int16), return_inverse=True, axis=1) + self.charge_labels = self.charge_labels.astype(np.int16) + else: + if charge_labels.dtype not in (np.int16, np.int16): + raise TypeError("`charge_labels` have to be of dtype `np.int16`") + + self.unique_charges = charges.astype(np.int16) + self.charge_labels = charge_labels.astype(np.int16) + + @property + def dim(self): + return len(self.charge_labels) + + @property + def num_symmetries(self) -> int: + """ + Return the number of different charges in `ChargeCollection`. + """ + return self.unique_charges.shape[0] + + @property + def num_unique(self) -> int: + """ + Return the number of different charges in `ChargeCollection`. + """ + return self.unique_charges.shape[1] + + @property + def identity_charges(self) -> np.ndarray: + """ + Give the identity charge associated to a symmetries type in `charge_types`. + Args: + charge_types: A list of charge-types. + Returns: + nd.array: vector of identity charges for each symmetry in self + """ + unique_charges = np.expand_dims( + np.asarray([ct.identity_charge() for ct in self.charge_types], + dtype=np.int16), 1) + charge_labels = np.zeros(1, dtype=np.int16) + obj = self.__new__(type(self)) + obj.__init__(unique_charges, charge_labels, self.charge_types) + return obj + + def __add__(self, other: "BaseCharge") -> "BaseCharge": + """ + Fuse `self` with `other`. + Args: + other: A `BaseCharge` object. + Returns: + Charge: The result of fusing `self` with `other`. + """ + + # fuse the unique charges from each index, then compute new unique charges + comb_charges = fuse_ndarray_charges(self.unique_charges, + other.unique_charges, self.charge_types) + [unique_charges, charge_labels] = np.unique( + comb_charges, return_inverse=True, axis=1) + charge_labels = charge_labels.reshape(self.unique_charges.shape[1], + other.unique_charges.shape[1]).astype( + np.int16) + + # find new labels using broadcasting + charge_labels = charge_labels[( + self.charge_labels[:, None] + np.zeros([1, len(other)], dtype=np.int16) + ).ravel(), (other.charge_labels[None, :] + + np.zeros([len(self), 1], dtype=np.int16)).ravel()] + + obj = self.__new__(type(self)) + obj.__init__(unique_charges, charge_labels, self.charge_types) + + return obj + + def dual(self, take_dual: Optional[bool] = False) -> np.ndarray: + if take_dual: + unique_dual_charges = np.stack([ + self.charge_types[n].dual_charges(self.unique_charges[n, :]) + for n in range(len(self.charge_types)) + ], + axis=0) + + obj = self.__new__(type(self)) + obj.__init__(unique_dual_charges, self.charge_labels, self.charge_types) + return obj + return self + + @property + def charges(self): + return self.unique_charges[:, self.charge_labels] + + def __repr__(self): + return str( + type(self)) + '\n' + 'charges: \n' + self.charges.__repr__() + '\n' + + def __len__(self): + return len(self.charge_labels) + + def __mul__(self, number: bool) -> "BaseCharge": + if not isinstance(number, (bool, np.bool_)): + print(type(number)) + raise ValueError( + "can only multiply by `True` or `False`, found {}".format(number)) + return self.dual(number) + + def intersect(self, other, assume_unique=False, + return_indices=False) -> (np.ndarray, np.ndarray, np.ndarray): + if isinstance(other, type(self)): + out = intersect( + self.unique_charges, + other.unique_charges, + assume_unique=True, + axis=1, + return_indices=return_indices) + obj = self.__new__(type(self)) + else: + out = intersect( + self.unique_charges, + np.asarray(other), + axis=1, + assume_unique=assume_unique, + return_indices=return_indices) + obj = self.__new__(type(self)) + if return_indices == True: + obj.__init__( + charges=out[0], + charge_labels=np.arange(len(out[0]), dtype=np.int16), + charge_types=self.charge_types, + ) + return obj, out[1], out[2] + return obj + + def unique(self, + return_index=False, + return_inverse=False, + return_counts=False + ) -> Tuple["BaseCharge", np.ndarray, np.ndarray, np.ndarray]: + """ + Compute the unique charges in `BaseCharge`. + See np.unique for a more detailed explanation. This function + does the same but instead of a np.ndarray, it returns the unique + elements in a `BaseCharge` object. + Args: + return_index: If `True`, also return the indices of `self.charges` (along the specified axis, + if provided, or in the flattened array) that result in the unique array. + return_inverse: If `True`, also return the indices of the unique array (for the specified + axis, if provided) that can be used to reconstruct `self.charges`. + return_counts: If `True`, also return the number of times each unique item appears + in `self.charges`. + Returns: + BaseCharge: The sorted unique values. + np.ndarray: The indices of the first occurrences of the unique values in the + original array. Only provided if `return_index` is True. + np.ndarray: The indices to reconstruct the original array from the + unique array. Only provided if `return_inverse` is True. + np.ndarray: The number of times each of the unique values comes up in the + original array. Only provided if `return_counts` is True. + """ + obj = self.__new__(type(self)) + obj.__init__( + charges=self.unique_charges, + charge_labels=np.arange(self.unique_charges.shape[1], dtype=np.int16), + charge_types=self.charge_types) + out = [obj] + if return_index: + _, index = np.unique(self.charge_labels, return_index=True) + out.append(index) + if return_inverse: + out.append(self.charge_labels) + if return_counts: + _, cnts = np.unique(self.charge_labels, return_counts=True) + out.append(cnts) + if len(out) == 1: + return out[0] + if len(out) == 2: + return out[0], out[1] + if len(out) == 3: + return out[0], out[1], out[2] + if len(out) == 4: + return out[0], out[1], out[2], out[3] + + @property + def dtype(self): + return self.unique_charges.dtype + + @property + def degeneracies(self): + return np.sum( + np.expand_dims(self.charge_labels, 1) == np.expand_dims( + np.arange(self.unique_charges.shape[1], dtype=np.int16), 0), + axis=0) + + def reduce(self, + target_charges: np.ndarray, + return_locations: bool = False, + strides: int = 1) -> ("SymIndex", np.ndarray): + """ + Reduce the dim of a SymIndex to keep only the index values that intersect target_charges + Args: + target_charges (np.ndarray): array of unique quantum numbers to keep. + return_locations (bool, optional): if True, also return the output index + locations of target values. + Returns: + SymIndex: index of reduced dimension. + np.ndarray: output index locations of target values. + """ + if isinstance(target_charges, (np.integer, int)): + target_charges = np.asarray([target_charges], dtype=np.int16) + if target_charges.ndim == 1: + target_charges = np.expand_dims(target_charges, 0) + target_charges = np.asarray(target_charges, dtype=np.int16) + # find intersection of index charges and target charges + reduced_charges, label_to_unique, label_to_target = intersect( + self.unique_charges, target_charges, axis=1, return_indices=True) + num_unique = len(label_to_unique) + + # construct the map to the reduced charges + map_to_reduced = np.full(self.dim, fill_value=-1, dtype=np.int16) + map_to_reduced[label_to_unique] = np.arange(num_unique, dtype=np.int16) + + # construct the map to the reduced charges + reduced_ind_labels = map_to_reduced[self.charge_labels] + reduced_locs = reduced_ind_labels >= 0 + new_ind_labels = reduced_ind_labels[reduced_locs].astype(np.int16) + obj = self.__new__(type(self)) + obj.__init__(reduced_charges, new_ind_labels, self.charge_types) + + if return_locations: + return obj, strides * np.flatnonzero(reduced_locs).astype(np.uint32) + return obj + + def __matmul__(self, other): + #some checks + if len(self) != len(other): + raise ValueError( + '__matmul__ requires charges to have the same number of elements') + charges = np.concatenate([self.charges, other.charges], axis=0) + charge_types = self.charge_types + other.charge_types + return BaseCharge( + charges=charges, charge_labels=None, charge_types=charge_types) + + def __getitem__(self, n: Union[np.ndarray, int]) -> "BaseCharge": + """ + Return the charge-element at position `n`, wrapped into a `BaseCharge` + object. + Args: + n: An integer or `np.ndarray`. + Returns: + BaseCharge: The charges at `n`. + """ + + if isinstance(n, (np.integer, int)): + n = np.asarray([n]) + obj = self.__new__(type(self)) + obj.__init__(self.unique_charges, self.charge_labels[n], self.charge_types) + return obj + + def __eq__(self, + target_charges: Union[np.ndarray, "BaseCharge"]) -> np.ndarray: + + if isinstance(target_charges, type(self)): + targets = np.unique( + target_charges.unique_charges[:, target_charges.charge_labels], + axis=1) + else: + print(isinstance(target_charges, type(self))) + print(type(target_charges), type(self)) + targets = np.unique(target_charges, axis=1) + inds = np.nonzero( + np.logical_and.reduce( + np.expand_dims(self.unique_charges, 2) == np.expand_dims( + targets, 1), + axis=0))[0] + return np.expand_dims(self.charge_labels, 1) == np.expand_dims(inds, 0) + + def isin(self, target_charges: Union[np.ndarray, "BaseCharge"]) -> np.ndarray: + + if isinstance(target_charges, type(self)): + targets = target_charges.unique_charges + else: + targets = np.unique(target_charges, axis=1) + tmp = np.expand_dims(self.unique_charges, 2) == np.expand_dims(targets, 1) + inds = np.nonzero( + np.logical_or.reduce(np.logical_and.reduce(tmp, axis=0), axis=1))[0] + + return np.isin(self.charge_labels, inds) + + +class U1Charge(BaseCharge): + + def __init__(self, + charges: np.ndarray, + charge_labels: Optional[np.ndarray] = None, + charge_types: Optional[List[Type["BaseCharge"]]] = None) -> None: + super().__init__(charges, charge_labels, charge_types=[type(self)]) + + @staticmethod + def fuse(charge1, charge2): + return np.add.outer(charge1, charge2).ravel() + + @staticmethod + def dual_charges(charges): + return charges * charges.dtype.type(-1) + + @staticmethod + def identity_charge(): + return np.int16(0) + + @classmethod + def random(cls, minval: int, maxval: int, dimension: tuple): + charges = np.random.randint(minval, maxval, dimension, dtype=np.int16) + return cls(charges=charges) + + +def fuse_ndarray_charges(charges_A: np.ndarray, charges_B: np.ndarray, + charge_types: List[Type[BaseCharge]]) -> np.ndarray: + """ + Fuse the quantum numbers of two indices under their kronecker addition. + Args: + charges_A (np.ndarray): n-by-D1 dimensional array integers encoding charges, + with n the number of symmetries and D1 the index dimension. + charges__B (np.ndarray): n-by-D2 dimensional array of charges. + charge_types: A list of types of the charges. + Returns: + np.ndarray: n-by-(D1 * D2) dimensional array of the fused charges. + """ + comb_charges = [0] * len(charge_types) + for n in range(len(charge_types)): + comb_charges[n] = charge_types[n].fuse(charges_A[n, :], charges_B[n, :]) + + return np.concatenate( + comb_charges, axis=0).reshape(len(charge_types), len(comb_charges[0])) + + +def intersect(A: np.ndarray, + B: np.ndarray, + axis=0, + assume_unique=False, + return_indices=False) -> (np.ndarray, np.ndarray, np.ndarray): + """ + Extends numpy's intersect1d to find the row or column-wise intersection of + two 2d arrays. Takes identical input to numpy intersect1d. + Args: + A, B (np.ndarray): arrays of matching widths and datatypes + Returns: + ndarray: sorted 1D array of common rows/cols between the input arrays + ndarray: the indices of the first occurrences of the common values in A. + Only provided if return_indices is True. + ndarray: the indices of the first occurrences of the common values in B. + Only provided if return_indices is True. + """ + #see https://stackoverflow.com/questions/8317022/get-intersecting-rows-across-two-2d-numpy-arrays + if A.ndim == 1: + return np.intersect1d( + A, B, assume_unique=assume_unique, return_indices=return_indices) + + elif A.ndim == 2: + if axis == 0: + ncols = A.shape[1] + if A.shape[1] != B.shape[1]: + raise ValueError("array widths must match to intersect") + + dtype = { + 'names': ['f{}'.format(i) for i in range(ncols)], + 'formats': ncols * [A.dtype] + } + if return_indices: + C, A_locs, B_locs = np.intersect1d( + A.view(dtype), + B.view(dtype), + assume_unique=assume_unique, + return_indices=return_indices) + return C.view(A.dtype).reshape(-1, ncols), A_locs, B_locs + C = np.intersect1d( + A.view(dtype), B.view(dtype), assume_unique=assume_unique) + return C.view(A.dtype).reshape(-1, ncols) + + elif axis == 1: + #@Glen: why the copy here? + out = intersect( + A.T.copy(), + B.T.copy(), + axis=0, + assume_unique=assume_unique, + return_indices=return_indices) + if return_indices: + return out[0].T, out[1], out[2] + return out.T + + else: + raise NotImplementedError( + "intersection can only be performed on first or second axis") + + else: + raise NotImplementedError( + "intersect is only implemented for 1d or 2d arrays") + + +def fuse_charges(charges: List[BaseCharge], flows: List[bool]) -> BaseCharge: + """ + Fuse all `charges` into a new charge. + Charges are fused from "right to left", + in accordance with row-major order. + + Args: + charges: A list of charges to be fused. + flows: A list of flows, one for each element in `charges`. + Returns: + BaseCharge: The result of fusing `charges`. + """ + if len(charges) != len(flows): + raise ValueError( + "`charges` and `flows` are of unequal lengths {} != {}".format( + len(charges), len(flows))) + fused_charges = charges[0] * flows[0] + for n in range(1, len(charges)): + fused_charges = fused_charges + charges[n] * flows[n] + return fused_charges + + +def fuse_degeneracies(degen1: Union[List, np.ndarray], + degen2: Union[List, np.ndarray]) -> np.ndarray: + """ + Fuse degeneracies `degen1` and `degen2` of two leg-charges + by simple kronecker product. `degen1` and `degen2` typically belong to two + consecutive legs of `BlockSparseTensor`. + Given `degen1 = [1, 2, 3]` and `degen2 = [10, 100]`, this returns + `[10, 100, 20, 200, 30, 300]`. + When using row-major ordering of indices in `BlockSparseTensor`, + the position of `degen1` should be "to the left" of the position of `degen2`. + Args: + degen1: Iterable of integers + degen2: Iterable of integers + Returns: + np.ndarray: The result of fusing `dege1` with `degen2`. + """ + return np.reshape(degen1[:, None] * degen2[None, :], + len(degen1) * len(degen2)) + + +# class BaseCharge: + +# def __init__(self, +# charges: np.ndarray, +# charge_labels: Optional[np.ndarray] = None) -> None: +# if charges.dtype != np.int16: +# raise TypeError("`charges` have to be of dtype `np.int16`") + +# if charge_labels is None: + +# self.unique_charges, charge_labels = np.unique( +# charges, return_inverse=True) +# self.charge_labels = charge_labels.astype(np.int16) + +# else: +# if charge_labels.dtype not in (np.int16, np.int16): +# raise TypeError("`charge_labels` have to be of dtype `np.int16`") + +# self.unique_charges = charges +# self.charge_labels = charge_labels.astype(np.int16) + +# def __add__(self, other: "BaseCharge") -> "BaseCharge": +# # fuse the unique charges from each index, then compute new unique charges +# comb_qnums = self.fuse(self.unique_charges, other.unique_charges) +# [unique_charges, charge_labels] = np.unique(comb_qnums, return_inverse=True) +# charge_labels = charge_labels.reshape( +# len(self.unique_charges), len(other.unique_charges)).astype(np.int16) + +# # find new labels using broadcasting (could use np.tile but less efficient) +# charge_labels = charge_labels[( +# self.charge_labels[:, None] + np.zeros([1, len(other)], dtype=np.int16) +# ).ravel(), (other.charge_labels[None, :] + +# np.zeros([len(self), 1], dtype=np.int16)).ravel()] +# obj = self.__new__(type(self)) +# obj.__init__(unique_charges, charge_labels) +# return obj + +# def __len__(self): +# return len(self.charge_labels) + +# def dual(self, take_dual: Optional[bool] = False) -> "BaseCharge": +# if take_dual: +# obj = self.__new__(type(self)) +# obj.__init__(self.dual_charges(self.unique_charges), self.charge_labels) +# return obj +# return self + +# @property +# def num_symmetries(self) -> int: +# """ +# Return the number of different charges in `ChargeCollection`. +# """ +# return self.unique_charges.shape[0] + +# def charges(self) -> np.ndarray: +# return self.unique_charges[self.charge_labels] + +# @property +# def dim(self): +# return len(self.charge_labels) + +# @property +# def dtype(self): +# return self.unique_charges.dtype + +# def __repr__(self): +# return str( +# type(self)) + '\n' + 'charges: ' + self.charges().__repr__() + '\n' + +# @property +# def degeneracies(self): +# return np.sum( +# np.expand_dims(self.charge_labels, 1) == np.expand_dims( +# np.arange(len(self.unique_charges), dtype=np.int16), 0), +# axis=0) + +# def unique(self, +# return_index=False, +# return_inverse=False, +# return_counts=False +# ) -> Tuple["BaseCharge", np.ndarray, np.ndarray, np.ndarray]: +# """ +# Compute the unique charges in `BaseCharge`. +# See np.unique for a more detailed explanation. This function +# does the same but instead of a np.ndarray, it returns the unique +# elements in a `BaseCharge` object. +# Args: +# return_index: If `True`, also return the indices of `self.charges` (along the specified axis, +# if provided, or in the flattened array) that result in the unique array. +# return_inverse: If `True`, also return the indices of the unique array (for the specified +# axis, if provided) that can be used to reconstruct `self.charges`. +# return_counts: If `True`, also return the number of times each unique item appears +# in `self.charges`. +# Returns: +# BaseCharge: The sorted unique values. +# np.ndarray: The indices of the first occurrences of the unique values in the +# original array. Only provided if `return_index` is True. +# np.ndarray: The indices to reconstruct the original array from the +# unique array. Only provided if `return_inverse` is True. +# np.ndarray: The number of times each of the unique values comes up in the +# original array. Only provided if `return_counts` is True. +# """ +# obj = self.__new__(type(self)) +# obj.__init__( +# self.unique_charges, +# charge_labels=np.arange(len(self.unique_charges), dtype=np.int16)) + +# out = [obj] +# if return_index: +# _, index = np.unique(self.charge_labels, return_index=True) +# out.append(index) +# if return_inverse: +# out.append(self.charge_labels) +# if return_counts: +# out.append(self.degeneracies) +# if len(out) == 1: +# return out[0] +# if len(out) == 2: +# return out[0], out[1] +# if len(out) == 3: +# return out[0], out[1], out[2] + +# def isin(self, targets: Union[int, Iterable, "BaseCharge"]) -> np.ndarray: +# """ +# Test each element of `BaseCharge` if it is in `targets`. Returns +# an `np.ndarray` of `dtype=bool`. +# Args: +# targets: The test elements +# Returns: +# np.ndarray: An array of `bool` type holding the result of the comparison. +# """ +# if isinstance(targets, type(self)): +# targets = targets.unique_charges +# targets = np.asarray(targets) +# common, label_to_unique, label_to_targets = np.intersect1d( +# self.unique_charges, targets, return_indices=True) +# if len(common) == 0: +# return np.full(len(self.charge_labels), fill_value=False, dtype=np.bool) +# return np.isin(self.charge_labels, label_to_unique) + +# def __iter__(self): +# return iter(self.charges()) + +# def __getitem__(self, n: Union[np.ndarray, int]) -> "BaseCharge": +# """ +# Return the charge-element at position `n`, wrapped into a `BaseCharge` +# object. +# Args: +# n: An integer or `np.ndarray`. +# Returns: +# BaseCharge: The charges at `n`. +# """ + +# if isinstance(n, (np.integer, int)): +# n = np.asarray([n]) +# obj = self.__new__(type(self)) +# obj.__init__(self.unique_charges, self.charge_labels[n]) +# return obj + +# def __mul__(self, number: bool) -> "BaseCharge": +# if not isinstance(number, bool): +# raise ValueError( +# "can only multiply by `True` or `False`, found {}".format(number)) + +# return self.dual(number) + +# def intersect(self, other, return_indices=False): +# out = np.intersect1d( +# self.unique_charges, +# other.unique_charges, +# assume_unique=True, +# return_indices=return_indices) +# obj = self.__new__(type(self)) +# if not return_indices: +# obj.__init__(out, np.arange(len(out), dtype=np.int16)) +# return obj +# obj.__init__(out[0], np.arange(len(out), dtype=np.int16)) +# return obj, out[1], out[2] + +# def reduce(self, +# target_charges: np.ndarray, +# return_locs: bool = False, +# strides: int = 1) -> ("SymIndex", np.ndarray): +# """ +# Reduce the dim of a SymIndex to keep only the index values that intersect target_charges +# Args: +# target_charges (np.ndarray): array of unique quantum numbers to keep. +# return_locs (bool, optional): if True, also return the output index +# locations of target values. +# Returns: +# SymIndex: index of reduced dimension. +# np.ndarray: output index locations of target values. +# """ +# if isinstance(target_charges, (int, np.integer)): +# target_charges = np.asarray([target_charges]) +# target_charges = np.asarray(target_charges, dtype=np.int16) +# # find intersection of index charges and target charges +# reduced_charges, label_to_unique, label_to_target = intersect( +# self.unique_charges, target_charges, axis=1, return_indices=True) + +# num_unique = len(label_to_unique) + +# # construct the map to the reduced charges +# map_to_reduced = np.full(self.dim, fill_value=-1, dtype=np.int16) +# map_to_reduced[label_to_unique] = np.arange(num_unique, dtype=np.int16) + +# # construct the map to the reduced charges +# reduced_ind_labels = map_to_reduced[self.charge_labels] +# reduced_locs = reduced_ind_labels >= 0 +# new_ind_labels = reduced_ind_labels[reduced_locs].astype(np.int16) +# obj = self.__new__(type(self)) + +# obj.__init__(reduced_charges, new_ind_labels) +# if return_locs: +# return obj, strides * np.flatnonzero(reduced_locs).astype(np.uint32) +# return obj + +# class ChargeCollection: + +# def __init__(self, +# charge_types: List[Type[BaseCharge]], +# charges: np.ndarray, +# charge_labels: Optional[np.ndarray] = None) -> None: +# self.charge_types = charge_types +# if charges.ndim == 1: +# charges = np.expand_dims(charges, 0) +# if charge_labels is None: +# self.unique_charges, self.charge_labels = np.unique( +# charges.astype(np.int16), return_inverse=True, axis=1) +# self.charge_labels = self.charge_labels.astype(np.int16) +# else: +# if charge_labels.dtype not in (np.int16, np.int16): +# raise TypeError("`charge_labels` have to be of dtype `np.int16`") + +# self.unique_charges = charges.astype(np.int16) +# self.charge_labels = charge_labels.astype(np.int16) + +# @property +# def dim(self): +# return len(self.charge_labels) + +# @property +# def num_symmetries(self) -> int: +# """ +# Return the number of different charges in `ChargeCollection`. +# """ +# return self.unique_charges.shape[0] + +# @property +# def identity_charges(self) -> np.ndarray: +# """ +# Give the identity charge associated to a symmetries type in `charge_types`. +# Args: +# charge_types: A list of charge-types. +# Returns: +# nd.array: vector of identity charges for each symmetry in self +# """ +# unique_charges = np.expand_dims( +# np.asarray([ct.identity for ct in self.charge_types], dtype=np.int16), +# 1) +# charge_labels = np.zeros(1, dtype=np.int16) +# obj = self.__new__(type(self)) +# obj.__init__(unique_charges, charge_labels) +# return obj + +# def __add__(self, other: "ChargeCollection") -> "ChargeCollection": +# """ +# Fuse `self` with `other`. +# Args: +# other: A `ChargeCollection` object. +# Returns: +# Charge: The result of fusing `self` with `other`. +# """ + +# # fuse the unique charges from each index, then compute new unique charges +# comb_charges = fuse_ndarray_charges(self.unique_charges, +# other.unique_charges, self.charge_types) +# [unique_charges, charge_labels] = np.unique( +# comb_charges, return_inverse=True, axis=1) +# charge_labels = charge_labels.reshape(self.unique_charges.shape[1], +# other.unique_charges.shape[1]).astype( +# np.int16) + +# # find new labels using broadcasting +# charge_labels = charge_labels[( +# self.charge_labels[:, None] + np.zeros([1, len(other)], dtype=np.int16) +# ).ravel(), (other.charge_labels[None, :] + +# np.zeros([len(self), 1], dtype=np.int16)).ravel()] +# return ChargeCollection(self.charge_types, unique_charges, charge_labels) + +# def dual(self, take_dual: Optional[bool] = False) -> np.ndarray: +# if take_dual: +# unique_dual_charges = np.stack([ +# self.charge_types[n].dual_charges(self.unique_charges[n, :]) +# for n in range(len(self.charge_types)) +# ], +# axis=0) +# return ChargeCollection(self.charge_types, unique_dual_charges, +# self.charge_labels) +# return self + +# def charges(self): +# return self.unique_charges[:, self.charge_labels] + +# def __repr__(self): +# return str( +# type(self)) + '\n' + 'charges: \n' + self.charges().__repr__() + '\n' + +# def __len__(self): +# return len(self.charge_labels) + +# def __mul__(self, number: bool) -> "ChargeCollection": +# if not isinstance(number, bool): +# raise ValueError( +# "can only multiply by `True` or `False`, found {}".format(number)) +# return self.dual(number) + +# def unique(self, +# return_index=False, +# return_inverse=False, +# return_counts=False +# ) -> Tuple["BaseCharge", np.ndarray, np.ndarray, np.ndarray]: +# """ +# Compute the unique charges in `BaseCharge`. +# See np.unique for a more detailed explanation. This function +# does the same but instead of a np.ndarray, it returns the unique +# elements in a `BaseCharge` object. +# Args: +# return_index: If `True`, also return the indices of `self.charges` (along the specified axis, +# if provided, or in the flattened array) that result in the unique array. +# return_inverse: If `True`, also return the indices of the unique array (for the specified +# axis, if provided) that can be used to reconstruct `self.charges`. +# return_counts: If `True`, also return the number of times each unique item appears +# in `self.charges`. +# Returns: +# BaseCharge: The sorted unique values. +# np.ndarray: The indices of the first occurrences of the unique values in the +# original array. Only provided if `return_index` is True. +# np.ndarray: The indices to reconstruct the original array from the +# unique array. Only provided if `return_inverse` is True. +# np.ndarray: The number of times each of the unique values comes up in the +# original array. Only provided if `return_counts` is True. +# """ + +# obj = ChargeCollection( +# self.charge_types, +# self.unique_charges, +# charge_labels=np.arange(self.unique_charges.shape[1], dtype=np.int16)) + +# out = [obj] +# if return_index: +# _, index = np.unique(self.charge_labels, return_index=True) +# out.append(index) +# if return_inverse: +# out.append(self.charge_labels) +# if return_counts: +# _, cnts = np.unique(self.charge_labels, return_counts=True) +# out.append(cnts) +# if len(out) == 1: +# return out[0] +# if len(out) == 2: +# return out[0], out[1] +# if len(out) == 3: +# return out[0], out[1], out[2] + +# @property +# def dtype(self): +# return self.unique_charges.dtype + +# @property +# def degeneracies(self): +# return np.sum( +# np.expand_dims(self.charge_labels, 1) == np.expand_dims( +# np.arange(self.unique_charges.shape[1], dtype=np.int16), 0), +# axis=0) + +# def reduce(self, +# target_charges: np.ndarray, +# return_locs: bool = False, +# strides: int = 1) -> ("SymIndex", np.ndarray): +# """ +# Reduce the dim of a SymIndex to keep only the index values that intersect target_charges +# Args: +# target_charges (np.ndarray): array of unique quantum numbers to keep. +# return_locs (bool, optional): if True, also return the output index +# locations of target values. +# Returns: +# SymIndex: index of reduced dimension. +# np.ndarray: output index locations of target values. +# """ +# target_charges = np.asarray(target_charges, dtype=np.int16) +# # find intersection of index charges and target charges +# reduced_charges, label_to_unique, label_to_target = intersect( +# self.unique_charges, target_charges, axis=1, return_indices=True) +# num_unique = len(label_to_unique) + +# # construct the map to the reduced charges +# map_to_reduced = np.full(self.dim, fill_value=-1, dtype=np.int16) +# map_to_reduced[label_to_unique] = np.arange(num_unique, dtype=np.int16) + +# # construct the map to the reduced charges +# reduced_ind_labels = map_to_reduced[self.charge_labels] +# reduced_locs = reduced_ind_labels >= 0 +# new_ind_labels = reduced_ind_labels[reduced_locs].astype(np.int16) + +# if return_locs: +# return ChargeCollection( +# self.charge_types, reduced_charges, +# new_ind_labels), strides * np.flatnonzero(reduced_locs).astype( +# np.uint32) +# return ChargeCollection(self.charge_types, reduced_charges, new_ind_labels) diff --git a/tensornetwork/block_tensor/charge_test.py b/tensornetwork/block_tensor/charge_test.py new file mode 100644 index 000000000..3e975a9cd --- /dev/null +++ b/tensornetwork/block_tensor/charge_test.py @@ -0,0 +1,256 @@ +import numpy as np +import pytest +# pylint: disable=line-too-long +from tensornetwork.block_tensor.charge import BaseCharge, U1Charge, fuse_degeneracies +from tensornetwork.block_tensor.block_tensor import fuse_ndarrays + + +def test_fuse_degeneracies(): + d1 = np.asarray([0, 1]) + d2 = np.asarray([2, 3, 4]) + fused_degeneracies = fuse_degeneracies(d1, d2) + np.testing.assert_allclose(fused_degeneracies, np.kron(d1, d2)) + + +def test_U1Charge_charges(): + D = 100 + B = 6 + charges = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + + q1 = U1Charge(charges) + assert np.all(q1.charges == charges) + + +def test_U1Charge_dual(): + D = 100 + B = 6 + charges = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + + q1 = U1Charge(charges) + assert np.all(q1.dual(True).charges == -charges) + + +def test_U1Charge_fusion(): + + def run_test(): + D = 2000 + B = 6 + O1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + O2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + P1 = np.random.randint(0, B + 1, D).astype(np.int16) + P2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + Q1 = np.random.randint(1, B + 1, D).astype(np.int16) + Q2 = np.random.randint(1, B + 1, D).astype(np.int16) + + charges_1 = [O1, O2] + charges_2 = [P1, P2] + charges_3 = [Q1, Q2] + + fused_1 = fuse_ndarrays(charges_1) + fused_2 = fuse_ndarrays(charges_2) + fused_3 = fuse_ndarrays(charges_3) + q1 = U1Charge(O1) @ U1Charge(P1) @ U1Charge(Q1) + q2 = U1Charge(O2) @ U1Charge(P2) @ U1Charge(Q2) + + target = BaseCharge( + charges=np.random.randint(-B, B, (3, 1), dtype=np.int16), + charge_labels=None, + charge_types=[U1Charge, U1Charge, U1Charge]) + q12 = q1 + q2 + + nz_1 = np.nonzero(q12 == target)[0] + i1 = fused_1 == target.charges[0, 0] + i2 = fused_2 == target.charges[1, 0] + i3 = fused_3 == target.charges[2, 0] + nz_2 = np.nonzero(np.logical_and.reduce([i1, i2, i3]))[0] + return nz_1, nz_2 + + nz_1, nz_2 = run_test() + while len(nz_1) == 0: + nz_1, nz_2 = run_test() + + assert np.all(nz_1 == nz_2) + + +def test_U1Charge_multiple_fusion(): + + def run_test(): + D = 300 + B = 4 + O1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + O2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + O3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + P1 = np.random.randint(0, B + 1, D).astype(np.int16) + P2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + P3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + Q1 = np.random.randint(1, B + 1, D).astype(np.int16) + Q2 = np.random.randint(0, B + 1, D).astype(np.int16) + Q3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + + charges_1 = [O1, O2, O3] + charges_2 = [P1, P2, P3] + charges_3 = [Q1, Q2, Q3] + + fused_1 = fuse_ndarrays(charges_1) + fused_2 = fuse_ndarrays(charges_2) + fused_3 = fuse_ndarrays(charges_3) + q1 = U1Charge(O1) @ U1Charge(P1) @ U1Charge(Q1) + q2 = U1Charge(O2) @ U1Charge(P2) @ U1Charge(Q2) + q3 = U1Charge(O3) @ U1Charge(P3) @ U1Charge(Q3) + + target = BaseCharge( + charges=np.random.randint(-B, B, (3, 1), dtype=np.int16), + charge_labels=None, + charge_types=[U1Charge, U1Charge, U1Charge]) + + q123 = q1 + q2 + q3 + + nz_1 = np.nonzero(q123 == target)[0] + i1 = fused_1 == target.charges[0, 0] + i2 = fused_2 == target.charges[1, 0] + i3 = fused_3 == target.charges[2, 0] + nz_2 = np.nonzero(np.logical_and.reduce([i1, i2, i3]))[0] + return nz_1, nz_2 + + nz_1, nz_2 = run_test() + while len(nz_1) == 0: + nz_1, nz_2 = run_test() + assert np.all(nz_1 == nz_2) + + +def test_U1Charge_multiple_fusion_with_flow(): + + def run_test(): + D = 300 + B = 4 + O1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + O2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + O3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + P1 = np.random.randint(0, B + 1, D).astype(np.int16) + P2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + P3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + Q1 = np.random.randint(1, B + 1, D).astype(np.int8) + Q2 = np.random.randint(0, B + 1, D).astype(np.int8) + Q3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + + charges_1 = [O1, -O2, O3] + charges_2 = [P1, -P2, P3] + charges_3 = [Q1, -Q2, Q3] + + fused_1 = fuse_ndarrays(charges_1) + fused_2 = fuse_ndarrays(charges_2) + fused_3 = fuse_ndarrays(charges_3) + q1 = U1Charge(O1) @ U1Charge(P1) @ U1Charge(Q1) + q2 = U1Charge(O2) @ U1Charge(P2) @ U1Charge(Q2) + q3 = U1Charge(O3) @ U1Charge(P3) @ U1Charge(Q3) + + target = BaseCharge( + charges=np.random.randint(-B, B, (3, 1), dtype=np.int16), + charge_labels=None, + charge_types=[U1Charge, U1Charge, U1Charge]) + q123 = q1 + q2 * True + q3 + nz_1 = np.nonzero(q123 == target)[0] + i1 = fused_1 == target.charges[0, 0] + i2 = fused_2 == target.charges[1, 0] + i3 = fused_3 == target.charges[2, 0] + nz_2 = np.nonzero(np.logical_and.reduce([i1, i2, i3]))[0] + return nz_1, nz_2 + + nz_1, nz_2 = run_test() + while len(nz_1) == 0: + nz_1, nz_2 = run_test() + assert np.all(nz_1 == nz_2) + + +def test_U1Charge_fusion_with_flow(): + + def run_test(): + D = 2000 + B = 6 + O1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + O2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int8) + P1 = np.random.randint(0, B + 1, D).astype(np.int16) + P2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + Q1 = np.random.randint(1, B + 1, D).astype(np.int8) + Q2 = np.random.randint(1, B + 1, D).astype(np.int8) + + charges_1 = [O1, -O2] + charges_2 = [P1, -P2] + charges_3 = [Q1, -Q2] + + fused_1 = fuse_ndarrays(charges_1) + fused_2 = fuse_ndarrays(charges_2) + fused_3 = fuse_ndarrays(charges_3) + + q1 = U1Charge(O1) @ U1Charge(P1) @ U1Charge(Q1) + q2 = U1Charge(O2) @ U1Charge(P2) @ U1Charge(Q2) + + target = BaseCharge( + charges=np.random.randint(-B, B, (3, 1), dtype=np.int16), + charge_labels=None, + charge_types=[U1Charge, U1Charge, U1Charge]) + q12 = q1 + q2 * True + + nz_1 = np.nonzero(q12 == target)[0] + i1 = fused_1 == target.charges[0, 0] + i2 = fused_2 == target.charges[1, 0] + i3 = fused_3 == target.charges[2, 0] + nz_2 = np.nonzero(np.logical_and.reduce([i1, i2, i3]))[0] + return nz_1, nz_2 + + nz_1, nz_2 = run_test() + while len(nz_1) == 0: + nz_1, nz_2 = run_test() + assert np.all(nz_1 == nz_2) + + +def test_U1Charge_matmul(): + D = 1000 + B = 5 + C1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + C2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + C3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + + q1 = U1Charge(C1) + q2 = U1Charge(C2) + q3 = U1Charge(C3) + + Q = q1 @ q2 @ q3 + Q_ = BaseCharge( + np.stack([C1, C2, C3], axis=0), + charge_labels=None, + charge_types=[U1Charge, U1Charge, U1Charge]) + assert np.all(Q.charges == Q_.charges) + + +def test_BaseCharge_eq(): + D = 3000 + B = 5 + c1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16) + c2 = np.random.randint(-B // 2 - 1, B // 2 + 2, D).astype(np.int16) + q1 = U1Charge(c1) + q2 = U1Charge(c2) + Q = q1 @ q2 + target_charge = np.asarray([ + np.random.randint(-B // 2, B // 2 + 1, dtype=np.int16), + np.random.randint(-B // 2 - 1, B // 2 + 2, dtype=np.int16) + ]) + T = U1Charge(np.asarray([target_charge[0]])) @ U1Charge( + np.asarray([target_charge[1]])) + assert np.all( + (np.squeeze(Q == T) + ) == np.logical_and(c1 == target_charge[0], c2 == target_charge[1])) + + +def test_BaseCharge_unique(): + D = 3000 + B = 5 + q = np.random.randint(-B // 2, B // 2 + 1, (2, D)).astype(np.int16) + Q = BaseCharge(charges=q, charge_types=[U1Charge, U1Charge]) + expected = np.unique( + q, return_index=True, return_inverse=True, return_counts=True, axis=1) + actual = Q.unique(return_index=True, return_inverse=True, return_counts=True) + assert np.all(actual[0].charges == expected[0]) + assert np.all(actual[1] == expected[1]) + assert np.all(actual[2] == expected[2]) + assert np.all(actual[3] == expected[3]) diff --git a/tensornetwork/block_tensor/index.py b/tensornetwork/block_tensor/index.py index b4bba3cee..fc2d033d9 100644 --- a/tensornetwork/block_tensor/index.py +++ b/tensornetwork/block_tensor/index.py @@ -16,9 +16,7 @@ from __future__ import division from __future__ import print_function import numpy as np -from tensornetwork.network_components import Node, contract, contract_between -# pylint: disable=line-too-long -from tensornetwork.backends import backend_factory +from tensornetwork.block_tensor.charge import BaseCharge import copy from typing import List, Union, Any, Optional, Tuple, Text @@ -31,16 +29,16 @@ class Index: """ def __init__(self, - charges: Union[List, np.ndarray], + charges: BaseCharge, flow: int, name: Optional[Text] = None, left_child: Optional["Index"] = None, right_child: Optional["Index"] = None): - self._charges = np.asarray(charges) + self._charges = charges #ChargeCollection([charges]) self.flow = flow self.left_child = left_child self.right_child = right_child - self.name = name if name else 'index' + self.name = name def __repr__(self): return str(self.dimension) @@ -59,17 +57,17 @@ def _copy_helper(self, index: "Index", copied_index: "Index") -> None: """ if index.left_child != None: left_copy = Index( - charges=copy.copy(index.left_child.charges), - flow=copy.copy(index.left_child.flow), - name=copy.copy(index.left_child.name)) + charges=copy.deepcopy(index.left_child.charges), + flow=copy.deepcopy(index.left_child.flow), + name=copy.deepcopy(index.left_child.name)) copied_index.left_child = left_copy self._copy_helper(index.left_child, left_copy) if index.right_child != None: right_copy = Index( - charges=copy.copy(index.right_child.charges), - flow=copy.copy(index.right_child.flow), - name=copy.copy(index.right_child.name)) + charges=copy.deepcopy(index.right_child.charges), + flow=copy.deepcopy(index.right_child.flow), + name=copy.deepcopy(index.right_child.name)) copied_index.right_child = right_copy self._copy_helper(index.right_child, right_copy) @@ -80,7 +78,9 @@ def copy(self): `Index` are copied as well. """ index_copy = Index( - charges=self._charges.copy(), flow=copy.copy(self.flow), name=self.name) + charges=copy.deepcopy(self._charges), + flow=copy.deepcopy(self.flow), + name=self.name) self._copy_helper(self, index_copy) return index_copy @@ -116,83 +116,12 @@ def __mul__(self, index: "Index") -> "Index": def charges(self): if self.is_leave: return self._charges - fused_charges = fuse_charge_pair( - self.left_child.charges, self.left_child.flow, self.right_child.charges, - self.right_child.flow) - - return fused_charges - - -def fuse_charge_pair(q1: Union[List, np.ndarray], flow1: int, - q2: Union[List, np.ndarray], flow2: int) -> np.ndarray: - """ - Fuse charges `q1` with charges `q2` by simple addition (valid - for U(1) charges). `q1` and `q2` typically belong to two consecutive - legs of `BlockSparseTensor`. - Given `q1 = [0,1,2]` and `q2 = [10,100]`, this returns - `[10, 100, 11, 101, 12, 102]`. - When using row-major ordering of indices in `BlockSparseTensor`, - the position of q1 should be "to the left" of the position of q2. - - Args: - q1: Iterable of integers - flow1: Flow direction of charge `q1`. - q2: Iterable of integers - flow2: Flow direction of charge `q2`. - Returns: - np.ndarray: The result of fusing `q1` with `q2`. - """ - return np.reshape( - flow1 * np.asarray(q1)[:, None] + flow2 * np.asarray(q2)[None, :], - len(q1) * len(q2)) - - -def fuse_charges(charges: List[Union[List, np.ndarray]], - flows: List[int]) -> np.ndarray: - """ - Fuse all `charges` by simple addition (valid - for U(1) charges). Charges are fused from "right to left", - in accordance with row-major order (see `fuse_charges_pair`). - - Args: - chargs: A list of charges to be fused. - flows: A list of flows, one for each element in `charges`. - Returns: - np.ndarray: The result of fusing `charges`. - """ - if len(charges) == 1: - #nothing to do - return charges[0] - fused_charges = charges[0] * flows[0] - for n in range(1, len(charges)): - fused_charges = fuse_charge_pair( - q1=fused_charges, flow1=1, q2=charges[n], flow2=flows[n]) - return fused_charges - - -def fuse_degeneracies(degen1: Union[List, np.ndarray], - degen2: Union[List, np.ndarray]) -> np.ndarray: - """ - Fuse degeneracies `degen1` and `degen2` of two leg-charges - by simple kronecker product. `degen1` and `degen2` typically belong to two - consecutive legs of `BlockSparseTensor`. - Given `degen1 = [1, 2, 3]` and `degen2 = [10, 100]`, this returns - `[10, 100, 20, 200, 30, 300]`. - When using row-major ordering of indices in `BlockSparseTensor`, - the position of `degen1` should be "to the left" of the position of `degen2`. - Args: - degen1: Iterable of integers - degen2: Iterable of integers - Returns: - np.ndarray: The result of fusing `dege1` with `degen2`. - """ - return np.reshape(degen1[:, None] * degen2[None, :], - len(degen1) * len(degen2)) + return self.left_child.charges * self.left_child.flow + self.right_child.charges * self.right_child.flow def fuse_index_pair(left_index: Index, right_index: Index, - flow: Optional[int] = 1) -> Index: + flow: Optional[int] = False) -> Index: """ Fuse two consecutive indices (legs) of a symmetric tensor. Args: @@ -211,7 +140,7 @@ def fuse_index_pair(left_index: Index, charges=None, flow=flow, left_child=left_index, right_child=right_index) -def fuse_indices(indices: List[Index], flow: Optional[int] = 1) -> Index: +def fuse_indices(indices: List[Index], flow: Optional[int] = False) -> Index: """ Fuse a list of indices (legs) of a symmetric tensor. Args: @@ -239,48 +168,3 @@ def split_index(index: Index) -> Tuple[Index, Index]: raise ValueError("cannot split an elementary index") return index.left_child, index.right_child - - -def unfuse(fused_indices: np.ndarray, len_left: int, - len_right: int) -> Tuple[np.ndarray, np.ndarray]: - """ - Given an np.ndarray `fused_indices` of integers denoting - index-positions of elements within a 1d array, `unfuse` - obtains the index-positions of the elements in the left and - right np.ndarrays `left`, `right` which, upon fusion, - are placed at the index-positions given by - `fused_indices` in the fused np.ndarray. - An example will help to illuminate this: - Given np.ndarrays `left`, `right` and the result - of their fusion (`fused`): - - ``` - left = [0,1,0,2] - right = [-1,3,-2] - fused = fuse_charges([left, right], flows=[1,1]) - print(fused) #[-1 3 -2 0 4 -1 -1 3 -2 1 5 0] - ``` - - we want to find which elements in `left` and `right` - fuse to a value of 0. In the above case, there are two - 0 in `fused`: one is obtained from fusing `left[1]` and - `right[0]`, the second one from fusing `left[3]` and `right[2]` - `unfuse` returns the index-positions of these values within - `left` and `right`, that is - - ``` - left_index_values, right_index_values = unfuse(np.nonzero(fused==0)[0], len(left), len(right)) - print(left_index_values) # [1,3] - print(right_index_values) # [0,2] - ``` - - Args: - fused_indices: A 1d np.ndarray of integers. - len_left: The length of the left np.ndarray. - len_right: The length of the right np.ndarray. - Returns: - (np.ndarry, np.ndarray) - """ - right = fused_indices % len_right - left = (fused_indices - right) // len_right - return left, right diff --git a/tensornetwork/block_tensor/index_test.py b/tensornetwork/block_tensor/index_test.py index 0feb2eb15..438984952 100644 --- a/tensornetwork/block_tensor/index_test.py +++ b/tensornetwork/block_tensor/index_test.py @@ -1,103 +1,66 @@ import numpy as np # pylint: disable=line-too-long -from tensornetwork.block_tensor.index import Index, fuse_index_pair, split_index, fuse_charges, fuse_degeneracies, fuse_charge_pair, fuse_indices, unfuse - - -def test_fuse_charge_pair(): - q1 = np.asarray([0, 1]) - q2 = np.asarray([2, 3, 4]) - fused_charges = fuse_charge_pair(q1, 1, q2, 1) - assert np.all(fused_charges == np.asarray([2, 3, 4, 3, 4, 5])) - fused_charges = fuse_charge_pair(q1, 1, q2, -1) - assert np.all(fused_charges == np.asarray([-2, -3, -4, -1, -2, -3])) - - -def test_fuse_charges(): - q1 = np.asarray([0, 1]) - q2 = np.asarray([2, 3, 4]) - fused_charges = fuse_charges([q1, q2], flows=[1, 1]) - assert np.all(fused_charges == np.asarray([2, 3, 4, 3, 4, 5])) - fused_charges = fuse_charges([q1, q2], flows=[1, -1]) - assert np.all(fused_charges == np.asarray([-2, -3, -4, -1, -2, -3])) - - -def test_fuse_degeneracies(): - d1 = np.asarray([0, 1]) - d2 = np.asarray([2, 3, 4]) - fused_degeneracies = fuse_degeneracies(d1, d2) - np.testing.assert_allclose(fused_degeneracies, np.kron(d1, d2)) +from tensornetwork.block_tensor.index import Index, fuse_index_pair, split_index, fuse_indices +from tensornetwork.block_tensor.charge import U1Charge, BaseCharge def test_index_fusion_mul(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 1 - q2 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 2 - i1 = Index(charges=q1, flow=1, name='index1') #index on leg 1 - i2 = Index(charges=q2, flow=1, name='index2') #index on leg 2 - - i12 = i1 * i2 - assert i12.left_child is i1 - assert i12.right_child is i2 - assert np.all(i12.charges == fuse_charge_pair(q1, 1, q2, 1)) + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + i1 = Index(charges=q1, flow=False, name='index1') #index on leg 1 + i2 = Index(charges=q2, flow=False, name='index2') #index on leg 2 -def test_fuse_index_pair(): - D = 10 - B = 4 - dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 1 - q2 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 2 - i1 = Index(charges=q1, flow=1, name='index1') #index on leg 1 - i2 = Index(charges=q2, flow=1, name='index2') #index on leg 2 - - i12 = fuse_index_pair(i1, i2) + i12 = i1 * i2 assert i12.left_child is i1 assert i12.right_child is i2 - assert np.all(i12.charges == fuse_charge_pair(q1, 1, q2, 1)) + for n in range(len(i12.charges.charges)): + assert np.all(i12.charges.charges == (q1 + q2).charges) def test_fuse_indices(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 1 - q2 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 2 - i1 = Index(charges=q1, flow=1, name='index1') #index on leg 1 - i2 = Index(charges=q2, flow=1, name='index2') #index on leg 2 + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + i1 = Index(charges=q1, flow=False, name='index1') #index on leg 1 + i2 = Index(charges=q2, flow=False, name='index2') #index on leg 2 i12 = fuse_indices([i1, i2]) assert i12.left_child is i1 assert i12.right_child is i2 - assert np.all(i12.charges == fuse_charge_pair(q1, 1, q2, 1)) + for n in range(len(i12.charges.charges)): + assert np.all(i12.charges.charges == (q1 + q2).charges) def test_split_index(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 1 - q2 = np.random.randint(-B // 2, B // 2 + 1, - D).astype(dtype) #quantum numbers on leg 2 - i1 = Index(charges=q1, flow=1, name='index1') #index on leg 1 - i2 = Index(charges=q2, flow=1, name='index2') #index on leg 2 + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + + i1 = Index(charges=q1, flow=False, name='index1') #index on leg 1 + i2 = Index(charges=q2, flow=False, name='index2') #index on leg 2 i12 = i1 * i2 i1_, i2_ = split_index(i12) assert i1 is i1_ assert i2 is i2_ - np.testing.assert_allclose(q1, i1.charges) - np.testing.assert_allclose(q2, i2.charges) - np.testing.assert_allclose(q1, i1_.charges) - np.testing.assert_allclose(q2, i2_.charges) + np.testing.assert_allclose(q1.charges, i1.charges.charges) + np.testing.assert_allclose(q2.charges, i2.charges.charges) + np.testing.assert_allclose(q1.charges, i1_.charges.charges) + np.testing.assert_allclose(q2.charges, i2_.charges.charges) assert i1_.name == 'index1' assert i2_.name == 'index2' assert i1_.flow == i1.flow @@ -108,14 +71,14 @@ def test_elementary_indices(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - q2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - q3 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - q4 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - i1 = Index(charges=q1, flow=1, name='index1') - i2 = Index(charges=q2, flow=1, name='index2') - i3 = Index(charges=q3, flow=1, name='index3') - i4 = Index(charges=q4, flow=1, name='index4') + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype)) + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype)) + q3 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype)) + q4 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype)) + i1 = Index(charges=q1, flow=False, name='index1') + i2 = Index(charges=q2, flow=False, name='index2') + i3 = Index(charges=q3, flow=False, name='index3') + i4 = Index(charges=q4, flow=False, name='index4') i12 = i1 * i2 i34 = i3 * i4 @@ -138,20 +101,23 @@ def test_elementary_indices(): assert elmt1234[2].flow == i3.flow assert elmt1234[3].flow == i4.flow - np.testing.assert_allclose(q1, i1.charges) - np.testing.assert_allclose(q2, i2.charges) - np.testing.assert_allclose(q3, i3.charges) - np.testing.assert_allclose(q4, i4.charges) + np.testing.assert_allclose(q1.charges, i1.charges.charges) + np.testing.assert_allclose(q2.charges, i2.charges.charges) + np.testing.assert_allclose(q3.charges, i3.charges.charges) + np.testing.assert_allclose(q4.charges, i4.charges.charges) def test_leave(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - q2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - i1 = Index(charges=q1, flow=1, name='index1') - i2 = Index(charges=q2, flow=1, name='index2') + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + + i1 = Index(charges=q1, flow=False, name='index1') + i2 = Index(charges=q2, flow=False, name='index2') assert i1.is_leave assert i2.is_leave @@ -163,12 +129,15 @@ def test_copy(): D = 10 B = 4 dtype = np.int16 - q1 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - q2 = np.random.randint(-B // 2, B // 2 + 1, D).astype(dtype) - i1 = Index(charges=q1, flow=1, name='index1') - i2 = Index(charges=q2, flow=1, name='index2') - i3 = Index(charges=q1, flow=-1, name='index3') - i4 = Index(charges=q2, flow=-1, name='index4') + q1 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + q2 = U1Charge(np.random.randint(-B // 2, B // 2 + 1, + D).astype(dtype)) #quantum numbers on leg 1 + + i1 = Index(charges=q1, flow=False, name='index1') + i2 = Index(charges=q2, flow=False, name='index2') + i3 = Index(charges=q1, flow=True, name='index3') + i4 = Index(charges=q2, flow=True, name='index4') i12 = i1 * i2 i34 = i3 * i4 @@ -180,17 +149,3 @@ def test_copy(): assert elmt1234[1] is not i2 assert elmt1234[2] is not i3 assert elmt1234[3] is not i4 - - -def test_unfuse(): - q1 = np.random.randint(-4, 5, 10).astype(np.int16) - q2 = np.random.randint(-4, 5, 4).astype(np.int16) - q3 = np.random.randint(-4, 5, 4).astype(np.int16) - q12 = fuse_charges([q1, q2], [1, 1]) - q123 = fuse_charges([q12, q3], [1, 1]) - nz = np.nonzero(q123 == 0)[0] - q12_inds, q3_inds = unfuse(nz, len(q12), len(q3)) - - q1_inds, q2_inds = unfuse(q12_inds, len(q1), len(q2)) - np.testing.assert_allclose(q1[q1_inds] + q2[q2_inds] + q3[q3_inds], - np.zeros(len(q1_inds), dtype=np.int16)) diff --git a/tensornetwork/block_tensor/tutorial.py b/tensornetwork/block_tensor/tutorial.py index 01e5eabf0..81d87bbe3 100644 --- a/tensornetwork/block_tensor/tutorial.py +++ b/tensornetwork/block_tensor/tutorial.py @@ -17,9 +17,9 @@ from __future__ import print_function import tensornetwork as tn import numpy as np -import tensornetwork.block_tensor.block_tensor as BT -import tensornetwork.block_tensor.index as IDX - +from tensornetwork.block_tensor.block_tensor import BlockSparseTensor, reshape +from tensornetwork.block_tensor.index import Index +from tensornetwork.block_tensor.charge import U1Charge B = 4 # possible charges on each leg can be between [0,B) ########################################################## ##### Generate a rank 4 symmetrix tensor ####### @@ -27,31 +27,31 @@ # generate random charges on each leg of the tensor D1, D2, D3, D4 = 4, 6, 8, 10 #bond dimensions on each leg -q1 = np.random.randint(0, B, D1) -q2 = np.random.randint(0, B, D2) -q3 = np.random.randint(0, B, D3) -q4 = np.random.randint(0, B, D4) +q1 = U1Charge(np.random.randint(-B, B + 1, D1)) +q2 = U1Charge(np.random.randint(-B, B + 1, D2)) +q3 = U1Charge(np.random.randint(-B, B + 1, D3)) +q4 = U1Charge(np.random.randint(-B, B + 1, D4)) # generate Index objects for each leg. neccessary for initialization of # BlockSparseTensor -i1 = IDX.Index(charges=q1, flow=1) -i2 = IDX.Index(charges=q2, flow=-1) -i3 = IDX.Index(charges=q3, flow=1) -i4 = IDX.Index(charges=q4, flow=-1) +i1 = Index(charges=q1, flow=1) +i2 = Index(charges=q2, flow=-1) +i3 = Index(charges=q3, flow=1) +i4 = Index(charges=q4, flow=-1) # initialize a random symmetric tensor -A = BT.BlockSparseTensor.randn(indices=[i1, i2, i3, i4], dtype=np.complex128) -B = BT.reshape(A, (4, 48, 10)) #creates a new tensor (copy) -shape_A = A.shape #returns the dense shape of A +A = BlockSparseTensor.randn(indices=[i1, i2, i3, i4], dtype=np.complex128) +B = reshape(A, (4, 48, 10)) #creates a new tensor (copy) +shape_A = A.dense_shape #returns the dense shape of A A.reshape([shape_A[0] * shape_A[1], shape_A[2], shape_A[3]]) #in place reshaping A.reshape(shape_A) #reshape back into original shape -sparse_shape = A.sparse_shape #returns a deep copy of `A.indices`. +sparse_shape = A.shape #returns a deep copy of `A.indices`. new_sparse_shape = [ sparse_shape[0] * sparse_shape[1], sparse_shape[2], sparse_shape[3] ] -B = BT.reshape(A, new_sparse_shape) #return a copy +B = reshape(A, new_sparse_shape) #return a copy A.reshape(new_sparse_shape) #in place reshaping A.reshape(sparse_shape) #bring A back into original shape diff --git a/tensornetwork/config.py b/tensornetwork/config.py index 8c347b4fa..0a54b0cb6 100644 --- a/tensornetwork/config.py +++ b/tensornetwork/config.py @@ -11,5 +11,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - -default_backend = "numpy" diff --git a/tensornetwork/matrixproductstates/finite_mps.py b/tensornetwork/matrixproductstates/finite_mps.py index 0fc308ee0..0e32d3e2e 100644 --- a/tensornetwork/matrixproductstates/finite_mps.py +++ b/tensornetwork/matrixproductstates/finite_mps.py @@ -29,25 +29,25 @@ class FiniteMPS(BaseMPS): """ - An MPS class for finite systems. + An MPS class for finite systems. MPS tensors are stored as a list of `Node` objects in the `FiniteMPS.nodes` attribute. - `FiniteMPS` has a central site, also called orthogonality center. - The position of this central site is stored in `FiniteMPS.center_position`, - and it can be be shifted using the `FiniteMPS.position` method. + `FiniteMPS` has a central site, also called orthogonality center. + The position of this central site is stored in `FiniteMPS.center_position`, + and it can be be shifted using the `FiniteMPS.position` method. `FiniteMPS.position` uses QR and RQ methods to shift `center_position`. - + `FiniteMPS` can be initialized either from a `list` of tensors, or by calling the classmethod `FiniteMPS.random`. - + By default, `FiniteMPS` is initialized in *canonical* form, i.e. - the state is normalized, and all tensors to the left of - `center_position` are left orthogonal, and all tensors + the state is normalized, and all tensors to the left of + `center_position` are left orthogonal, and all tensors to the right of `center_position` are right orthogonal. The tensor at `FiniteMPS.center_position` is neither left nor right orthogonal. - Note that canonicalization can be computationally relatively + Note that canonicalization can be computationally relatively costly and scales :math:`\\propto ND^3`. """ @@ -62,7 +62,7 @@ def __init__(self, tensors: A list of `Tensor` or `BaseNode` objects. center_position: The initial position of the center site. canonicalize: If `True` the mps is canonicalized at initialization. - backend: The name of the backend that should be used to perform + backend: The name of the backend that should be used to perform contractions. Available backends are currently 'numpy', 'tensorflow', 'pytorch', 'jax' """ @@ -160,10 +160,13 @@ def left_envs(self, sites: Sequence[int]) -> Dict: Args: sites (list of int): A list of sites of the MPS. Returns: - `dict` mapping `int` to `Tensor`: The left-reduced density matrices + `dict` mapping `int` to `Tensor`: The left-reduced density matrices at each site in `sites`. """ + if not sites: + return {} + n2 = max(sites) sites = np.array(sites) #enable logical indexing @@ -227,9 +230,11 @@ def right_envs(self, sites: Sequence[int]) -> Dict: Args: sites (list of int): A list of sites of the MPS. Returns: - `dict` mapping `int` to `Tensor`: The right-reduced density matrices + `dict` mapping `int` to `Tensor`: The right-reduced density matrices at each site in `sites`. """ + if not sites: + return {} n1 = min(sites) sites = np.array(sites) diff --git a/tensornetwork/matrixproductstates/finite_mps_test.py b/tensornetwork/matrixproductstates/finite_mps_test.py index 0afe0f4f6..9b779ad8c 100644 --- a/tensornetwork/matrixproductstates/finite_mps_test.py +++ b/tensornetwork/matrixproductstates/finite_mps_test.py @@ -121,3 +121,33 @@ def test_correlation_measurement_finite_mps(backend_dtype_values): actual[N // 2] = 0.25 np.testing.assert_almost_equal(result_1, actual) np.testing.assert_allclose(result_2, np.ones(N) * 0.25) + + +def test_left_envs_empty_seq(backend_dtype_values): + backend = backend_dtype_values[0] + dtype = backend_dtype_values[1] + + D, d, N = 1, 2, 10 + tensors = [np.ones((1, d, D), dtype=dtype)] + [ + np.ones((D, d, D), dtype=dtype) for _ in range(N - 2) + ] + [np.ones((D, d, 1), dtype=dtype)] + mps = FiniteMPS(tensors, center_position=0, backend=backend) + + assert mps.left_envs(()) == {} + assert mps.left_envs([]) == {} + assert mps.left_envs(range(0)) == {} + + +def test_right_envs_empty_seq(backend_dtype_values): + backend = backend_dtype_values[0] + dtype = backend_dtype_values[1] + + D, d, N = 1, 2, 10 + tensors = [np.ones((1, d, D), dtype=dtype)] + [ + np.ones((D, d, D), dtype=dtype) for _ in range(N - 2) + ] + [np.ones((D, d, 1), dtype=dtype)] + mps = FiniteMPS(tensors, center_position=0, backend=backend) + + assert mps.right_envs(()) == {} + assert mps.right_envs([]) == {} + assert mps.right_envs(range(0)) == {} diff --git a/tensornetwork/ncon_interface.py b/tensornetwork/ncon_interface.py index f4dde1593..330bdc044 100644 --- a/tensornetwork/ncon_interface.py +++ b/tensornetwork/ncon_interface.py @@ -16,7 +16,7 @@ import warnings from typing import Any, Sequence, List, Optional, Union, Text, Tuple, Dict from tensornetwork import network_components -from tensornetwork import config +from tensornetwork.backend_contextmanager import get_default_backend from tensornetwork.backends import backend_factory Tensor = Any @@ -67,8 +67,8 @@ def ncon(tensors: Sequence[Union[network_components.BaseNode, Tensor]], structure. con_order: List of edge labels specifying the contraction order. out_order: List of edge labels specifying the output order. - backend: String specifying the backend to use. Defaults to - `tensornetwork.config.default_backend`. + backend: String specifying the backend to use. Defaults to + `tensornetwork.backend_contextmanager.get_default_backend`. Returns: The result of the contraction. The result is returned as a `Node` @@ -78,7 +78,7 @@ def ncon(tensors: Sequence[Union[network_components.BaseNode, Tensor]], if backend and (backend not in backend_factory._BACKENDS): raise ValueError("Backend '{}' does not exist".format(backend)) if backend is None: - backend = config.default_backend + backend = get_default_backend() are_nodes = [isinstance(t, network_components.BaseNode) for t in tensors] nodes = {t for t in tensors if isinstance(t, network_components.BaseNode)} diff --git a/tensornetwork/network_components.py b/tensornetwork/network_components.py index 1c9fa5c45..17d4549ed 100644 --- a/tensornetwork/network_components.py +++ b/tensornetwork/network_components.py @@ -21,10 +21,10 @@ import h5py #pylint: disable=useless-import-alias -import tensornetwork.config as config from tensornetwork import ops from tensornetwork.backends import backend_factory from tensornetwork.backends.base_backend import BaseBackend +from tensornetwork.backend_contextmanager import get_default_backend string_type = h5py.special_dtype(vlen=str) Tensor = Any @@ -69,10 +69,18 @@ def __init__(self, """ self.is_disabled = False - self.name = name if name is not None else '__unnamed_node__' + if name is None: + name = '__unnamed_node__' + else: + if not isinstance(name, str): + raise TypeError("Node name should be str type") + self.name = name self.backend = backend self._shape = shape if axis_names is not None: + for axis_name in axis_names: + if not isinstance(axis_name, str): + raise TypeError("axis_names should be str type") self._edges = [ Edge(node1=self, axis1=i, name=edge_name) for i, edge_name in enumerate(axis_names) @@ -97,6 +105,18 @@ def __init__(self, super().__init__() + def __add__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented addition ( + )") + + def __sub__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented subtraction ( - )") + + def __mul__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented multiply ( * )") + + def __truediv__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented divide ( / )") + @property def dtype(self): #any derived instance of BaseNode always has to have a tensor @@ -125,6 +145,9 @@ def add_axis_names(self, axis_names: List[Text]) -> None: raise ValueError("axis_names is not the same length as the tensor shape." "axis_names length: {}, tensor.shape length: {}".format( len(axis_names), len(self.shape))) + for axis_name in axis_names: + if not isinstance(axis_name, str): + raise TypeError("axis_names should be str type") self.axis_names = axis_names[:] def add_edge(self, @@ -312,6 +335,8 @@ def get_all_dangling(self) -> Set["Edge"]: return {edge for edge in self.edges if edge.is_dangling()} def set_name(self, name) -> None: + if not isinstance(name, str): + raise TypeError("Node name should be str type") self.name = name def has_nondangling_edge(self) -> bool: @@ -373,6 +398,16 @@ def edges(self, edges: List) -> None: self.name)) self._edges = edges + @property + def name(self) -> Text: + return self._name + + @name.setter + def name(self, name) -> None: + if not isinstance(name, str): + raise TypeError("Node name should be str type") + self._name = name + @property def axis_names(self) -> List[Text]: return self._axis_names @@ -382,8 +417,12 @@ def axis_names(self, axis_names: List[Text]) -> None: if len(axis_names) != len(self.shape): raise ValueError("Expected {} names, only got {}.".format( len(self.shape), len(axis_names))) + for axis_name in axis_names: + if not isinstance(axis_name, str): + raise TypeError("axis_names should be str type") self._axis_names = axis_names + @property def signature(self) -> Optional[int]: if self.is_disabled: @@ -498,8 +537,8 @@ def __init__(self, """Create a node. Args: - tensor: The concrete that is represented by this node, or a `BaseNode` - object. If a tensor is passed, it can be + tensor: The concrete that is represented by this node, or a `BaseNode` + object. If a tensor is passed, it can be be either a numpy array or the tensor-type of the used backend. If a `BaseNode` is passed, the passed node has to have the same \ backend as given by `backend`. @@ -515,8 +554,8 @@ def __init__(self, #always use the `Node`'s backend backend = tensor.backend tensor = tensor.tensor - if not backend: - backend = config.default_backend + if backend is None: + backend = get_default_backend() if isinstance(backend, BaseBackend): backend_obj = backend else: @@ -528,6 +567,71 @@ def __init__(self, backend=backend_obj, shape=backend_obj.shape_tuple(self._tensor)) + def op_protection(self, other: Union[int, float, "Node"]) -> "Node": + if not isinstance(other, (int, float, Node)): + raise TypeError("Operand should be one of int, float, Node type") + if not hasattr(self, '_tensor'): + raise AttributeError("Please provide a valid tensor for this Node.") + if isinstance(other, Node): + if not self.backend.name == other.backend.name: + raise TypeError("Operands backend must match.\noperand 1 backend: {}\ + \noperand 2 backend: {}".format(self.backend.name, + other.backend.name)) + if not hasattr(other, '_tensor'): + raise AttributeError("Please provide a valid tensor for this Node.") + else: + other_tensor = self.backend.convert_to_tensor(other) + other = Node(tensor=other_tensor, backend=self.backend.name) + return other + + def __add__(self, other: Union[int, float, "Node"]) -> "Node": + other = self.op_protection(other) + new_tensor = self.backend.addition(self.tensor, other.tensor) + if len(self.axis_names) > len(other.axis_names): + axis_names = self.axis_names + else: + axis_names = other.axis_names + return Node(tensor=new_tensor, + name=self.name, + axis_names=axis_names, + backend=self.backend.name) + + def __sub__(self, other: Union[int, float, "Node"]) -> "Node": + other = self.op_protection(other) + new_tensor = self.backend.subtraction(self.tensor, other.tensor) + if len(self.axis_names) > len(other.axis_names): + axis_names = self.axis_names + else: + axis_names = other.axis_names + return Node(tensor=new_tensor, + name=self.name, + axis_names=axis_names, + backend=self.backend.name) + + def __mul__(self, other: Union[int, float, "Node"]) -> "Node": + other = self.op_protection(other) + new_tensor = self.backend.multiply(self.tensor, other.tensor) + if len(self.axis_names) > len(other.axis_names): + axis_names = self.axis_names + else: + axis_names = other.axis_names + return Node(tensor=new_tensor, + name=self.name, + axis_names=axis_names, + backend=self.backend.name) + + def __truediv__(self, other: Union[int, float, "Node"]) -> "Node": + other = self.op_protection(other) + new_tensor = self.backend.divide(self.tensor, other.tensor) + if len(self.axis_names) > len(other.axis_names): + axis_names = self.axis_names + else: + axis_names = other.axis_names + return Node(tensor=new_tensor, + name=self.name, + axis_names=axis_names, + backend=self.backend.name) + def get_tensor(self) -> Tensor: return self.tensor @@ -606,13 +710,13 @@ def __init__(self, backend: An optional backend for the node. If `None`, a default backend is used dtype: The dtype used to initialize a numpy-copy node. - Note that this dtype has to be a numpy dtype, and it has to be + Note that this dtype has to be a numpy dtype, and it has to be compatible with the dtype of the backend, e.g. for a tensorflow backend with a tf.Dtype=tf.floa32, `dtype` has to be `np.float32`. """ - if not backend: - backend = config.default_backend + if backend is None: + backend = get_default_backend() backend_obj = backend_factory.get_backend(backend) self.rank = rank @@ -626,6 +730,23 @@ def __init__(self, backend=backend_obj, shape=(dimension,) * rank) + def __add__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented addition ( + )") + + def __sub__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented subtraction ( - )") + + def __mul__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented multiply ( * )") + + def __truediv__(self, other: Union[int, float, "BaseNode"]) -> "BaseNode": + raise NotImplementedError("BaseNode has not implemented divide ( / )") + + @property + def dtype(self): + # Override so we don't construct the dense tensor when asked for the dtype! + return self.copy_node_dtype + def get_tensor(self) -> Tensor: return self.tensor @@ -803,8 +924,11 @@ def __init__(self, raise ValueError( "node2 and axis2 must either be both None or both not be None") self.is_disabled = False - if not name: + if name is None: name = '__unnamed_edge__' + else: + if not isinstance(name, str): + raise TypeError("Edge name should be str type") self._name = name self.node1 = node1 self._axis1 = axis1 @@ -839,6 +963,8 @@ def name(self, name) -> None: if self.is_disabled: raise ValueError( 'Edge has been disabled, setting its name is no longer possible') + if not isinstance(name, str): + raise TypeError("Edge name should be str type") self._name = name @property @@ -983,6 +1109,8 @@ def is_being_used(self) -> bool: return result def set_name(self, name: Text) -> None: + if not isinstance(name, str): + raise TypeError("Edge name should be str type") self.name = name def _save_edge(self, edge_group: h5py.Group) -> None: @@ -1053,14 +1181,14 @@ def disconnect(self, edge2_name: Optional[Text] = None) -> Tuple["Edge", "Edge"]: """ Break an existing non-dangling edge. - This updates both Edge.node1 and Edge.node2 by removing the + This updates both Edge.node1 and Edge.node2 by removing the connecting edge from `Edge.node1.edges` and `Edge.node2.edges` and adding new dangling edges instead Args: edge1_name: A name for the new dangling edge at `self.node1` edge2_name: A name for the new dangling edge at `self.node2` Returns: - (new_edge1, new_edge2): The new `Edge` objects of + (new_edge1, new_edge2): The new `Edge` objects of `self.node1` and `self.node2` """ if self.is_dangling(): @@ -1116,7 +1244,7 @@ def get_parallel_edges(edge: Edge) -> Set[Edge]: edge: The given edge. Returns: - A `set` of all of the edges parallel to the given edge + A `set` of all of the edges parallel to the given edge (including the given edge). """ return get_shared_edges(edge.node1, edge.node2) @@ -1157,10 +1285,12 @@ def _flatten_trace_edges(edges: List[Edge], perm_front = set(range(len(node.edges))) - set(perm_back) perm_front = sorted(perm_front) perm = perm_front + perm_back - new_dim = backend.prod([backend.shape(node.tensor)[e.axis1] for e in edges]) + new_dim = backend.shape_prod( + [backend.shape_tensor(node.tensor)[e.axis1] for e in edges]) node.reorder_axes(perm) - unaffected_shape = backend.shape(node.tensor)[:len(perm_front)] - new_shape = backend.concat([unaffected_shape, [new_dim, new_dim]], axis=-1) + unaffected_shape = backend.shape_tensor(node.tensor)[:len(perm_front)] + new_shape = backend.shape_concat( + [unaffected_shape, [new_dim, new_dim]], axis=-1) node.tensor = backend.reshape(node.tensor, new_shape) edge1 = Edge(node1=node, axis1=len(perm_front), name="TraceFront") edge2 = Edge(node1=node, axis1=len(perm_front) + 1, name="TraceBack") @@ -1232,11 +1362,11 @@ def flatten_edges(edges: List[Edge], perm_back.append(node.edges.index(edge)) perm_front = sorted(set(range(len(node.edges))) - set(perm_back)) node.reorder_axes(perm_front + perm_back) - old_tensor_shape = backend.shape(node.tensor) + old_tensor_shape = backend.shape_tensor(node.tensor) # Calculate the new axis dimension as a product of the other # axes dimensions. - flattened_axis_dim = backend.prod(old_tensor_shape[len(perm_front):]) - new_tensor_shape = backend.concat( + flattened_axis_dim = backend.shape_prod(old_tensor_shape[len(perm_front):]) + new_tensor_shape = backend.shape_concat( [old_tensor_shape[:len(perm_front)], [flattened_axis_dim]], axis=-1) new_tensor = backend.reshape(node.tensor, new_tensor_shape) # Modify the node in place. Currently, this is they only method that @@ -1324,8 +1454,8 @@ def _split_trace_edge( perm_front = set(range(len(node.edges))) - set(perm_back) perm_front = sorted(perm_front) node.reorder_axes(perm_front + perm_back) - unaffected_shape = backend.shape(node.tensor)[:len(perm_front)] - new_shape = backend.concat([unaffected_shape, shape, shape], axis=-1) + unaffected_shape = backend.shape_tensor(node.tensor)[:len(perm_front)] + new_shape = backend.shape_concat([unaffected_shape, shape, shape], axis=-1) node.tensor = backend.reshape(node.tensor, new_shape) # Trim edges and add placeholder edges for new axes. node.edges = node.edges[:len(perm_front)] + 2 * len(shape) * [None] @@ -1348,8 +1478,8 @@ def split_edge(edge: Edge, shape: Tuple[int, ...], new_edge_names: Optional[List[Text]] = None) -> List[Edge]: """Split an `Edge` into multiple edges according to `shape`. Reshapes - the underlying tensors connected to the edge accordingly. - + the underlying tensors connected to the edge accordingly. + This method acts as the inverse operation of flattening edges and distinguishes between the following edge cases when adding new edges: 1) standard edge connecting two different nodes: reshape node dimensions @@ -1399,8 +1529,8 @@ def split_edge(edge: Edge, perm_front = set(range(len(node.edges))) - set(perm_back) perm_front = sorted(perm_front) node.reorder_axes(perm_front + perm_back) - unaffected_shape = backend.shape(node.tensor)[:len(perm_front)] - new_shape = backend.concat([unaffected_shape, shape], axis=-1) + unaffected_shape = backend.shape_tensor(node.tensor)[:len(perm_front)] + new_shape = backend.shape_concat([unaffected_shape, shape], axis=-1) node.tensor = backend.reshape(node.tensor, new_shape) # in-place update # Trim edges. node.edges = node.edges[:len(perm_front)] @@ -1731,7 +1861,7 @@ def disconnect(edge, edge2_name: Optional[Text] = None) -> Tuple[Edge, Edge]: """ Break an existing non-dangling edge. - This updates both Edge.node1 and Edge.node2 by removing the + This updates both Edge.node1 and Edge.node2 by removing the connecting edge from `Edge.node1.edges` and `Edge.node2.edges` and adding new dangling edges instead """ @@ -1748,6 +1878,10 @@ def contract_between( ) -> BaseNode: """Contract all of the edges between the two given nodes. + If `output_edge_order` is not set, the output axes will be ordered as: + [...free axes of `node1`..., ...free axes of `node2`...]. Within the axes + of each node, the input order is preserved. + Args: node1: The first node. node2: The second node. @@ -1759,7 +1893,8 @@ def contract_between( contain all edges belonging to, but not shared by `node1` and `node2`. The axes of the new node will be permuted (if necessary) to match this ordering of Edges. - axis_names: An optional list of names for the axis of the new node + axis_names: An optional list of names for the axis of the new node in order + of the output axes. Returns: The new node created. @@ -1779,64 +1914,68 @@ def contract_between( node2.backend.name)) backend = node1.backend + shared_edges = get_shared_edges(node1, node2) # Trace edges cannot be contracted using tensordot. if node1 is node2: flat_edge = flatten_edges_between(node1, node2) if not flat_edge: raise ValueError("No trace edges found on contraction of edges between " "node '{}' and itself.".format(node1)) - return contract(flat_edge, name) - - shared_edges = get_shared_edges(node1, node2) - if not shared_edges: - if allow_outer_product: - return outer_product(node1, node2, name=name, axis_names=axis_names) - raise ValueError("No edges found between nodes '{}' and '{}' " - "and allow_outer_product=False.".format(node1, node2)) - - # Collect the axis of each node corresponding to each edge, in order. - # This specifies the contraction for tensordot. - # NOTE: The ordering of node references in each contraction edge is ignored. - axes1 = [] - axes2 = [] - for edge in shared_edges: - if edge.node1 is node1: - axes1.append(edge.axis1) - axes2.append(edge.axis2) - else: - axes1.append(edge.axis2) - axes2.append(edge.axis1) - - if output_edge_order: - # Determine heuristically if output transposition can be minimized by - # flipping the arguments to tensordot. - node1_output_axes = [] - node2_output_axes = [] - for (i, edge) in enumerate(output_edge_order): - if edge in shared_edges: - raise ValueError( - "Edge '{}' in output_edge_order is shared by the nodes to be " - "contracted: '{}' and '{}'.".format(edge, node1, node2)) - edge_nodes = set(edge.get_nodes()) - if node1 in edge_nodes: - node1_output_axes.append(i) - elif node2 in edge_nodes: - node2_output_axes.append(i) + new_node = contract(flat_edge, name) + elif not shared_edges: + if not allow_outer_product: + raise ValueError("No edges found between nodes '{}' and '{}' " + "and allow_outer_product=False.".format(node1, node2)) + new_node = outer_product(node1, node2, name=name) + else: + # Collect the axis of each node corresponding to each edge, in order. + # This specifies the contraction for tensordot. + # NOTE: The ordering of node references in each contraction edge is ignored. + axes1 = [] + axes2 = [] + for edge in shared_edges: + if edge.node1 is node1: + axes1.append(edge.axis1) + axes2.append(edge.axis2) else: - raise ValueError( - "Edge '{}' in output_edge_order is not connected to node '{}' or " - "node '{}'".format(edge, node1, node2)) - if np.mean(node1_output_axes) > np.mean(node2_output_axes): - node1, node2 = node2, node1 - axes1, axes2 = axes2, axes1 - - new_tensor = backend.tensordot(node1.tensor, node2.tensor, [axes1, axes2]) - new_node = Node( - tensor=new_tensor, name=name, axis_names=axis_names, backend=backend) - # node1 and node2 get new edges in _remove_edges - _remove_edges(shared_edges, node1, node2, new_node) + axes1.append(edge.axis2) + axes2.append(edge.axis1) + + if output_edge_order: + # Determine heuristically if output transposition can be minimized by + # flipping the arguments to tensordot. + node1_output_axes = [] + node2_output_axes = [] + for (i, edge) in enumerate(output_edge_order): + if edge in shared_edges: + raise ValueError( + "Edge '{}' in output_edge_order is shared by the nodes to be " + "contracted: '{}' and '{}'.".format(edge, node1, node2)) + edge_nodes = set(edge.get_nodes()) + if node1 in edge_nodes: + node1_output_axes.append(i) + elif node2 in edge_nodes: + node2_output_axes.append(i) + else: + raise ValueError( + "Edge '{}' in output_edge_order is not connected to node '{}' or " + "node '{}'".format(edge, node1, node2)) + if node1_output_axes and node2_output_axes and ( + np.mean(node1_output_axes) > np.mean(node2_output_axes)): + node1, node2 = node2, node1 + axes1, axes2 = axes2, axes1 + + new_tensor = backend.tensordot(node1.tensor, node2.tensor, [axes1, axes2]) + new_node = Node( + tensor=new_tensor, name=name, backend=backend) + # node1 and node2 get new edges in _remove_edges + _remove_edges(shared_edges, node1, node2, new_node) + if output_edge_order: new_node = new_node.reorder_edges(list(output_edge_order)) + if axis_names: + new_node.add_axis_names(axis_names) + return new_node @@ -1844,9 +1983,9 @@ def outer_product_final_nodes(nodes: Iterable[BaseNode], edge_order: List[Edge]) -> BaseNode: """Get the outer product of `nodes` - For example, if there are 3 nodes remaining in `nodes` with + For example, if there are 3 nodes remaining in `nodes` with shapes :math:`(2, 3)`, :math:`(4, 5, 6)`, and :math:`(7)` - respectively, the newly returned node will have shape + respectively, the newly returned node will have shape :math:`(2, 3, 4, 5, 6, 7)`. Args: diff --git a/tensornetwork/network_operations.py b/tensornetwork/network_operations.py index 0d93f45f4..c2e228792 100644 --- a/tensornetwork/network_operations.py +++ b/tensornetwork/network_operations.py @@ -19,7 +19,6 @@ import numpy as np #pylint: disable=useless-import-alias -import tensornetwork.config as config #pylint: disable=line-too-long from tensornetwork.network_components import BaseNode, Node, CopyNode, Edge, disconnect from tensornetwork.backends import backend_factory @@ -125,23 +124,29 @@ def copy(nodes: Iterable[BaseNode], node_dict: A dictionary mapping the nodes to their copies. edge_dict: A dictionary mapping the edges to their copies. """ - #TODO: add support for copying CopyTensor - if conjugate: - node_dict = { - node: Node( + node_dict = {} + for node in nodes: + if isinstance(node, CopyNode): + node_dict[node] = CopyNode( + node.rank, + node.dimension, + name=node.name, + axis_names=node.axis_names, + backend=node.backend, + dtype=node.dtype) + else: + if conjugate: + node_dict[node] = Node( node.backend.conj(node.tensor), name=node.name, axis_names=node.axis_names, - backend=node.backend) for node in nodes - } - else: - node_dict = { - node: Node( + backend=node.backend) + else: + node_dict[node] = Node( node.tensor, name=node.name, axis_names=node.axis_names, - backend=node.backend) for node in nodes - } + backend=node.backend) edge_dict = {} for edge in get_all_edges(nodes): node1 = edge.node1 @@ -184,9 +189,6 @@ def remove_node(node: BaseNode) -> Tuple[Dict[Text, Edge], Dict[int, Edge]]: the newly broken edges. disconnected_edges_by_axis: A Dictionary mapping `node`'s axis numbers to the newly broken edges. - - Raises: - ValueError: If the node isn't in the network. """ disconnected_edges_by_name = {} disconnected_edges_by_axis = {} @@ -292,8 +294,8 @@ def split_node( # the first axis of vh. If we don't, it's possible one of the other axes of # vh will be the same size as sqrt_s and would multiply across that axis # instead, which is bad. - sqrt_s_broadcast_shape = backend.concat( - [backend.shape(sqrt_s), [1] * (len(vh.shape) - 1)], axis=-1) + sqrt_s_broadcast_shape = backend.shape_concat( + [backend.shape_tensor(sqrt_s), [1] * (len(vh.shape) - 1)], axis=-1) vh_s = vh * backend.reshape(sqrt_s, sqrt_s_broadcast_shape) left_node = Node( u_s, name=left_name, axis_names=left_axis_names, backend=backend) @@ -607,7 +609,7 @@ def reachable(inputs: Union[BaseNode, Iterable[BaseNode], Edge, Iterable[Edge]] Args: inputs: A `BaseNode`/`Edge` or collection of `BaseNodes`/`Edges` Returns: - A list of `BaseNode` objects that can be reached from `node` + A set of `BaseNode` objects that can be reached from `node` via connected edges. Raises: ValueError: If an unknown value for `strategy` is passed. diff --git a/tensornetwork/quantum/quantum.py b/tensornetwork/quantum/quantum.py new file mode 100644 index 000000000..33b1313d5 --- /dev/null +++ b/tensornetwork/quantum/quantum.py @@ -0,0 +1,641 @@ +# Copyright 2019 The TensorNetwork Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Abstractions for quantum vectors and operators. + +Quantum mechanics involves a lot of linear algebra on vector spaces that often +have a preferred tensor-product factorization. Tensor networks are a natural +way to represent vectors and operators (matrices) involving these spaces. Hence +we provide some simple abstractions to ease linear algebra operations in which +the vectors and operators are represented by tensor networks. +""" +from typing import Any, Union, Callable, Optional, Sequence, Collection, Text +from typing import Tuple, Set, List, Type +import numpy as np +from tensornetwork.network_components import BaseNode, Node, Edge, connect +from tensornetwork.network_components import CopyNode +from tensornetwork.network_operations import get_all_nodes, copy, reachable +from tensornetwork.network_operations import get_subgraph_dangling, remove_node +from tensornetwork.contractors import greedy +Tensor = Any + + +def quantum_constructor(out_edges: Sequence[Edge], in_edges: Sequence[Edge], + ref_nodes: Optional[Collection[BaseNode]] = None, + ignore_edges: Optional[Collection[Edge]] = None +) -> "QuOperator": + """Constructs an appropriately specialized QuOperator. + + If there are no edges, creates a QuScalar. If the are only output (input) + edges, creates a QuVector (QuAdjointVector). Otherwise creates a + QuOperator. + + Args: + out_edges: output edges. + in_edges: in edges. + ref_nodes: reference nodes for the tensor network (needed if there is a + scalar component). + ignore_edges: edges to ignore when checking the dimensionality of the + tensor network. + Returns: + The object. + """ + if len(out_edges) == 0 and len(in_edges) == 0: + return QuScalar(ref_nodes, ignore_edges) + if len(out_edges) == 0: + return QuAdjointVector(in_edges, ref_nodes, ignore_edges) + if len(in_edges) == 0: + return QuVector(out_edges, ref_nodes, ignore_edges) + return QuOperator(out_edges, in_edges, ref_nodes, ignore_edges) + + +def identity(space: Sequence[int], backend: Optional[Text] = None, + dtype: Type[np.number] = np.float64) -> "QuOperator": + """Construct a `QuOperator` representing the identity on a given space. + + Internally, this is done by constructing `CopyNode`s for each edge, with + dimension according to `space`. + + Args: + space: A sequence of integers for the dimensions of the tensor product + factors of the space (the edges in the tensor network). + backend: Optionally specify the backend to use for computations. + dtype: The data type (for conversion to dense). + Returns: + The desired identity operator. + """ + nodes = [CopyNode(2, d, backend=backend, dtype=dtype) for d in space] + out_edges = [n[0] for n in nodes] + in_edges = [n[1] for n in nodes] + return quantum_constructor(out_edges, in_edges) + + +def check_spaces(edges_1: Sequence[Edge], edges_2: Sequence[Edge]) -> None: + """Check the vector spaces represented by two lists of edges are compatible. + + The number of edges must be the same and the dimensions of each pair of edges + must match. Otherwise, an exception is raised. + + Args: + edges_1: List of edges representing a many-body Hilbert space. + edges_2: List of edges representing a many-body Hilbert space. + """ + if len(edges_1) != len(edges_2): + raise ValueError("Hilbert-space mismatch: Cannot connect {} subsystems " + "with {} subsystems.".format(len(edges_1), len(edges_2))) + + for (i, (e1, e2)) in enumerate(zip(edges_1, edges_2)): + if e1.dimension != e2.dimension: + raise ValueError("Hilbert-space mismatch on subsystems {}: Input " + "dimension {} != output dimension {}.".format( + i, e1.dimension, e2.dimension)) + + +def eliminate_identities(nodes: Collection[BaseNode]) -> Tuple[dict, dict]: + """Eliminates any connected CopyNodes that are identity matrices. + + This will modify the network represented by `nodes`. + Only identities that are connected to other nodes are eliminated. + + Args: + nodes: Collection of nodes to search. + Returns: + nodes_dict: Dictionary mapping remaining Nodes to any replacements. + dangling_edges_dict: Dictionary specifying all dangling-edge replacements. + """ + nodes_dict = {} + dangling_edges_dict = {} + for n in nodes: + if isinstance(n, CopyNode) and n.get_rank() == 2 and not ( + n[0].is_dangling() and n[1].is_dangling()): + old_edges = [n[0], n[1]] + _, new_edges = remove_node(n) + if 0 in new_edges and 1 in new_edges: + e = connect(new_edges[0], new_edges[1]) + elif 0 in new_edges: # 1 was dangling + dangling_edges_dict[old_edges[1]] = new_edges[0] + elif 1 in new_edges: # 0 was dangling + dangling_edges_dict[old_edges[0]] = new_edges[1] + else: + # Trace of identity, so replace with a scalar node! + d = n.get_dimension(0) + # NOTE: Assume CopyNodes have numpy dtypes. + nodes_dict[n] = Node(np.array(d, dtype=n.dtype), backend=n.backend) + else: + for e in n.get_all_dangling(): + dangling_edges_dict[e] = e + nodes_dict[n] = n + + return nodes_dict, dangling_edges_dict + + +class QuOperator(): + """Represents a linear operator via a tensor network. + + To interpret a tensor network as a linear operator, some of the dangling + edges must be designated as `out_edges` (output edges) and the rest as + `in_edges` (input edges). + + Considered as a matrix, the `out_edges` represent the row index and the + `in_edges` represent the column index. + + The (right) action of the operator on another then consists of connecting + the `in_edges` of the first operator to the `out_edges` of the second. + + Can be used to do simple linear algebra with tensor networks. + """ + __array_priority__ = 100.0 # for correct __rmul__ with scalar ndarrays + + def __init__(self, out_edges: Sequence[Edge], in_edges: Sequence[Edge], + ref_nodes: Optional[Collection[BaseNode]] = None, + ignore_edges: Optional[Collection[Edge]] = None) -> None: + """Creates a new `QuOperator` from a tensor network. + + This encapsulates an existing tensor network, interpreting it as a linear + operator. + + The network is checked for consistency: All dangling edges must either be + in `out_edges`, `in_edges`, or `ignore_edges`. + + Args: + out_edges: The edges of the network to be used as the output edges. + in_edges: The edges of the network to be used as the input edges. + ref_nodes: Nodes used to refer to parts of the tensor network that are + not connected to any input or output edges (for example: a scalar + factor). + ignore_edges: Optional collection of dangling edges to ignore when + performing consistency checks. + """ + # TODO: Decide whether the user must also supply all nodes involved. + # This would enable extra error checking and is probably clearer + # than `ref_nodes`. + if len(in_edges) == 0 and len(out_edges) == 0 and not ref_nodes: + raise ValueError("At least one reference node is required to specify a " + "scalar. None provided!") + self.out_edges = list(out_edges) + self.in_edges = list(in_edges) + self.ignore_edges = set(ignore_edges) if ignore_edges else set() + self.ref_nodes = set(ref_nodes) if ref_nodes else set() + self.check_network() + + @classmethod + def from_tensor(cls, tensor: Tensor, out_axes: Sequence[int], + in_axes: Sequence[int], backend: Optional[Text] = None + ) -> "QuOperator": + """Construct a `QuOperator` directly from a single tensor. + + This first wraps the tensor in a `Node`, then constructs the `QuOperator` + from that `Node`. + + Args: + tensor: The tensor. + out_axes: The axis indices of `tensor` to use as `out_edges`. + in_axes: The axis indices of `tensor` to use as `in_edges`. + backend: Optionally specify the backend to use for computations. + Returns: + The new operator. + """ + n = Node(tensor, backend=backend) + out_edges = [n[i] for i in out_axes] + in_edges = [n[i] for i in in_axes] + return cls(out_edges, in_edges, set([n])) + + @property + def nodes(self) -> Set[BaseNode]: + """All tensor-network nodes involved in the operator. + """ + return reachable( + get_all_nodes(self.out_edges + self.in_edges) | self.ref_nodes) + + @property + def in_space(self) -> List[int]: + return [e.dimension for e in self.in_edges] + + @property + def out_space(self) -> List[int]: + return [e.dimension for e in self.out_edges] + + def is_scalar(self) -> bool: + return len(self.out_edges) == 0 and len(self.in_edges) == 0 + + def is_vector(self) -> bool: + return len(self.out_edges) > 0 and len(self.in_edges) == 0 + + def is_adjoint_vector(self) -> bool: + return len(self.out_edges) == 0 and len(self.in_edges) > 0 + + def check_network(self) -> None: + """Check that the network has the expected dimensionality. + + This checks that all input and output edges are dangling and that there + are no other dangling edges (except any specified in `ignore_edges`). + If not, an exception is raised. + """ + for (i, e) in enumerate(self.out_edges): + if not e.is_dangling(): + raise ValueError("Output edge {} is not dangling!".format(i)) + for (i, e) in enumerate(self.in_edges): + if not e.is_dangling(): + raise ValueError("Input edge {} is not dangling!".format(i)) + for e in self.ignore_edges: + if not e.is_dangling(): + raise ValueError("ignore_edges contains non-dangling edge: {}".format( + str(e))) + + known_edges = set(self.in_edges) | set(self.out_edges) | self.ignore_edges + all_dangling_edges = get_subgraph_dangling(self.nodes) + if known_edges != all_dangling_edges: + raise ValueError("The network includes unexpected dangling edges (that " + "are not members of ignore_edges).") + + def adjoint(self) -> "QuOperator": + """The adjoint of the operator. + + This creates a new `QuOperator` with complex-conjugate copies of all + tensors in the network and with the input and output edges switched. + """ + nodes_dict, edge_dict = copy(self.nodes, True) + out_edges = [edge_dict[e] for e in self.in_edges] + in_edges = [edge_dict[e] for e in self.out_edges] + ref_nodes = [nodes_dict[n] for n in self.ref_nodes] + ignore_edges = [edge_dict[e] for e in self.ignore_edges] + return quantum_constructor( + out_edges, in_edges, ref_nodes, ignore_edges) + + def trace(self) -> "QuOperator": + """The trace of the operator. + """ + return self.partial_trace(range(len(self.in_edges))) + + def norm(self) -> "QuOperator": + """The norm of the operator. + This is the 2-norm (also known as the Frobenius or Hilbert-Schmidt norm). + """ + return (self.adjoint() @ self).trace() + + def partial_trace(self, subsystems_to_trace_out: Collection[int] + ) -> "QuOperator": + """The partial trace of the operator. + + Subsystems to trace out are supplied as indices, so that dangling edges + are connected to eachother as: + `out_edges[i] ^ in_edges[i] for i in subsystems_to_trace_out` + + This does not modify the original network. The original ordering of the + remaining subsystems is maintained. + + Args: + subsystems_to_trace_out: Indices of subsystems to trace out. + Returns: + A new QuOperator or QuScalar representing the result. + """ + out_edges_trace = [self.out_edges[i] for i in subsystems_to_trace_out] + in_edges_trace = [self.in_edges[i] for i in subsystems_to_trace_out] + + check_spaces(in_edges_trace, out_edges_trace) + + nodes_dict, edge_dict = copy(self.nodes, False) + for (e1, e2) in zip(out_edges_trace, in_edges_trace): + edge_dict[e1] = edge_dict[e1] ^ edge_dict[e2] + + # get leftover edges in the original order + out_edges_trace = set(out_edges_trace) + in_edges_trace = set(in_edges_trace) + out_edges = [edge_dict[e] for e in self.out_edges + if e not in out_edges_trace] + in_edges = [edge_dict[e] for e in self.in_edges + if e not in in_edges_trace] + ref_nodes = [n for _, n in nodes_dict.items()] + ignore_edges = [edge_dict[e] for e in self.ignore_edges] + + return quantum_constructor(out_edges, in_edges, ref_nodes, ignore_edges) + + def __matmul__(self, other: "QuOperator") -> "QuOperator": + """The action of this operator on another. + + Given `QuOperator`s `A` and `B`, produces a new `QuOperator` for `A @ B`, + where `A @ B` means: "the action of A, as a linear operator, on B". + + Under the hood, this produces copies of the tensor networks defining `A` + and `B` and then connects the copies by hooking up the `in_edges` of + `A.copy()` to the `out_edges` of `B.copy()`. + """ + check_spaces(self.in_edges, other.out_edges) + + # Copy all nodes involved in the two operators. + # We must do this separately for self and other, in case self and other + # are defined via the same network components (e.g. if self === other). + nodes_dict1, edges_dict1 = copy(self.nodes, False) + nodes_dict2, edges_dict2 = copy(other.nodes, False) + + # connect edges to create network for the result + for (e1, e2) in zip(self.in_edges, other.out_edges): + _ = edges_dict1[e1] ^ edges_dict2[e2] + + in_edges = [edges_dict2[e] for e in other.in_edges] + out_edges = [edges_dict1[e] for e in self.out_edges] + ref_nodes = ([n for _, n in nodes_dict1.items()] + + [n for _, n in nodes_dict2.items()]) + ignore_edges = ([edges_dict1[e] for e in self.ignore_edges] + + [edges_dict2[e] for e in other.ignore_edges]) + + return quantum_constructor(out_edges, in_edges, ref_nodes, ignore_edges) + + def __mul__(self, other: Union["QuOperator", BaseNode, Tensor] + ) -> "QuOperator": + """Scalar multiplication of operators. + + Given two operators `A` and `B`, one of the which is a scalar (it has no + input or output edges), `A * B` produces a new operator representing the + scalar multiplication of `A` and `B`. + + For convenience, one of `A` or `B` may be a number or scalar-valued tensor + or `Node` (it will automatically be wrapped in a `QuScalar`). + + Note: This is a special case of `tensor_product()`. + """ + if not isinstance(other, QuOperator): + if isinstance(other, BaseNode): + node = other + else: + node = Node(other, backend=self.nodes.pop().backend) + if node.shape: + raise ValueError("Cannot perform elementwise multiplication by a " + "non-scalar tensor.") + other = QuScalar([node]) + + if self.is_scalar() or other.is_scalar(): + return self.tensor_product(other) + + raise ValueError("Elementwise multiplication is only supported if at " + "least one of the arguments is a scalar.") + + def __rmul__(self, other: Union["QuOperator", BaseNode, Tensor]) -> "QuOperator": + """Scalar multiplication of operators. See `.__mul__()`. + """ + return self.__mul__(other) + + def tensor_product(self, other: "QuOperator") -> "QuOperator": + """Tensor product with another operator. + + Given two operators `A` and `B`, produces a new operator `AB` representing + `A` ⊗ `B`. The `out_edges` (`in_edges`) of `AB` is simply the + concatenation of the `out_edges` (`in_edges`) of `A.copy()` with that of + `B.copy()`: + + `new_out_edges = [*out_edges_A_copy, *out_edges_B_copy]` + `new_in_edges = [*in_edges_A_copy, *in_edges_B_copy]` + + Args: + other: The other operator (`B`). + Returns: + The result (`AB`). + """ + nodes_dict1, edges_dict1 = copy(self.nodes, False) + nodes_dict2, edges_dict2 = copy(other.nodes, False) + + in_edges = ([edges_dict1[e] for e in self.in_edges] + + [edges_dict2[e] for e in other.in_edges]) + out_edges = ([edges_dict1[e] for e in self.out_edges] + + [edges_dict2[e] for e in other.out_edges]) + ref_nodes = ([n for _, n in nodes_dict1.items()] + + [n for _, n in nodes_dict2.items()]) + ignore_edges = ([edges_dict1[e] for e in self.ignore_edges] + + [edges_dict2[e] for e in other.ignore_edges]) + + return quantum_constructor(out_edges, in_edges, ref_nodes, ignore_edges) + + def contract(self, contractor: Callable = greedy, + final_edge_order: Optional[Sequence[Edge]] = None + ) -> "QuOperator": + """Contract the tensor network in place. + + This modifies the tensor network representation of the operator (or vector, + or scalar), reducing it to a single tensor, without changing the value. + + Args: + contractor: A function that performs the contraction. Defaults to + `greedy`, which uses the greedy algorithm from `opt_einsum` to + determine a contraction order. + final_edge_order: Manually specify the axis ordering of the final tensor. + Returns: + The present object. + """ + nodes_dict, dangling_edges_dict = eliminate_identities(self.nodes) + self.in_edges = [dangling_edges_dict[e] for e in self.in_edges] + self.out_edges = [dangling_edges_dict[e] for e in self.out_edges] + self.ignore_edges = set(dangling_edges_dict[e] for e in self.ignore_edges) + self.ref_nodes = set( + nodes_dict[n] for n in self.ref_nodes if n in nodes_dict) + self.check_network() + + if final_edge_order: + final_edge_order = [dangling_edges_dict[e] for e in final_edge_order] + self.ref_nodes = set( + [contractor(self.nodes, output_edge_order=final_edge_order)]) + else: + self.ref_nodes = set([contractor(self.nodes, ignore_edge_order=True)]) + return self + + def eval(self, contractor: Callable = greedy, + final_edge_order: Optional[Sequence[Edge]] = None) -> Tensor: + """Contracts the tensor network in place and returns the final tensor. + + Note that this modifies the tensor network representing the operator. + + The default ordering for the axes of the final tensor is: + `*out_edges, *in_edges`. + + If there are any "ignored" edges, their axes come first: + `*ignored_edges, *out_edges, *in_edges`. + + Args: + contractor: A function that performs the contraction. Defaults to + `greedy`, which uses the greedy algorithm from `opt_einsum` to + determine a contraction order. + final_edge_order: Manually specify the axis ordering of the final tensor. + The default ordering is determined by `out_edges` and `in_edges` (see + above). + Returns: + The final tensor representing the operator. + """ + if not final_edge_order: + final_edge_order = (list(self.ignore_edges) + self.out_edges + + self.in_edges) + self.contract(contractor, final_edge_order) + nodes = self.nodes + if len(nodes) != 1: + raise ValueError("Node count '{}' > 1 after contraction!".format( + len(nodes))) + return list(nodes)[0].tensor + + +class QuVector(QuOperator): + """Represents a (column) vector via a tensor network. + """ + def __init__(self, subsystem_edges: Sequence[Edge], + ref_nodes: Optional[Collection[BaseNode]] = None, + ignore_edges: Optional[Collection[Edge]] = None) -> None: + """Constructs a new `QuVector` from a tensor network. + + This encapsulates an existing tensor network, interpreting it as a (column) + vector. + + Args: + subsystem_edges: The edges of the network to be used as the output edges. + ref_nodes: Nodes used to refer to parts of the tensor network that are + not connected to any input or output edges (for example: a scalar + factor). + ignore_edges: Optional collection of edges to ignore when performing + consistency checks. + """ + super().__init__(subsystem_edges, [], ref_nodes, ignore_edges) + + @classmethod + def from_tensor(cls, tensor: Tensor, + subsystem_axes: Optional[Sequence[int]] = None, + backend: Optional[Text] = None) -> "QuVector": + """Construct a `QuVector` directly from a single tensor. + + This first wraps the tensor in a `Node`, then constructs the `QuVector` + from that `Node`. + + Args: + tensor: The tensor. + subsystem_axes: Sequence of integer indices specifying the order in which + to interpret the axes as subsystems (output edges). If not specified, + the axes are taken in ascending order. + backend: Optionally specify the backend to use for computations. + Returns: + The new operator. + """ + n = Node(tensor, backend=backend) + if subsystem_axes is not None: + subsystem_edges = [n[i] for i in subsystem_axes] + else: + subsystem_edges = n.get_all_edges() + return cls(subsystem_edges) + + @property + def subsystem_edges(self) -> List[Edge]: + return self.out_edges + + @property + def space(self) -> List[int]: + return self.out_space + + def projector(self) -> "QuOperator": + return self @ self.adjoint() + + def reduced_density(self, subsystems_to_trace_out: Collection[int] + ) -> "QuOperator": + rho = self.projector() + return rho.partial_trace(subsystems_to_trace_out) + + +class QuAdjointVector(QuOperator): + """Represents an adjoint (row) vector via a tensor network. + """ + def __init__(self, subsystem_edges: Sequence[Edge], + ref_nodes: Optional[Collection[BaseNode]] = None, + ignore_edges: Optional[Collection[Edge]] = None) -> None: + """Constructs a new `QuAdjointVector` from a tensor network. + + This encapsulates an existing tensor network, interpreting it as an adjoint + vector (row vector). + + Args: + subsystem_edges: The edges of the network to be used as the input edges. + ref_nodes: Nodes used to refer to parts of the tensor network that are + not connected to any input or output edges (for example: a scalar + factor). + ignore_edges: Optional collection of edges to ignore when performing + consistency checks. + """ + super().__init__([], subsystem_edges, ref_nodes, ignore_edges) + + @classmethod + def from_tensor(cls, tensor: Tensor, + subsystem_axes: Optional[Sequence[int]] = None, + backend: Optional[Text] = None) -> "QuAdjointVector": + """Construct a `QuAdjointVector` directly from a single tensor. + + This first wraps the tensor in a `Node`, then constructs the + `QuAdjointVector` from that `Node`. + + Args: + tensor: The tensor. + subsystem_axes: Sequence of integer indices specifying the order in which + to interpret the axes as subsystems (input edges). If not specified, + the axes are taken in ascending order. + backend: Optionally specify the backend to use for computations. + Returns: + The new operator. + """ + n = Node(tensor, backend=backend) + if subsystem_axes is not None: + subsystem_edges = [n[i] for i in subsystem_axes] + else: + subsystem_edges = n.get_all_edges() + return cls(subsystem_edges) + + @property + def subsystem_edges(self) -> List[Edge]: + return self.in_edges + + @property + def space(self) -> List[int]: + return self.in_space + + def projector(self) -> "QuOperator": + return self.adjoint() @ self + + def reduced_density(self, subsystems_to_trace_out: Collection[int] + ) -> "QuOperator": + rho = self.projector() + return rho.partial_trace(subsystems_to_trace_out) + + +class QuScalar(QuOperator): + """Represents a scalar via a tensor network. + """ + def __init__(self, ref_nodes: Collection[BaseNode], + ignore_edges: Optional[Collection[Edge]] = None) -> None: + """Constructs a new `QuScalar` from a tensor network. + + This encapsulates an existing tensor network, interpreting it as a scalar. + + Args: + ref_nodes: Nodes used to refer to the tensor network (need not be + exhaustive - one node from each disconnected subnetwork is sufficient). + ignore_edges: Optional collection of edges to ignore when performing + consistency checks. + """ + super().__init__([], [], ref_nodes, ignore_edges) + + @classmethod + def from_tensor(cls, tensor: Tensor, backend: Optional[Text] = None + ) -> "QuScalar": + """Construct a `QuScalar` directly from a single tensor. + + This first wraps the tensor in a `Node`, then constructs the + `QuScalar` from that `Node`. + + Args: + tensor: The tensor. + backend: Optionally specify the backend to use for computations. + Returns: + The new operator. + """ + n = Node(tensor, backend=backend) + return cls(set([n])) diff --git a/tensornetwork/quantum/quantum_test.py b/tensornetwork/quantum/quantum_test.py new file mode 100644 index 000000000..4a9ae113d --- /dev/null +++ b/tensornetwork/quantum/quantum_test.py @@ -0,0 +1,201 @@ +# Copyright 2019 The TensorNetwork Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import numpy as np +import tensornetwork as tn +import quantum as qu + + +def test_constructor(backend): + psi_tensor = np.random.rand(2, 2) + psi_node = tn.Node(psi_tensor, backend=backend) + + op = qu.quantum_constructor([psi_node[0]], [psi_node[1]]) + assert not op.is_scalar() + assert not op.is_vector() + assert not op.is_adjoint_vector() + assert len(op.out_edges) == 1 + assert len(op.in_edges) == 1 + assert op.out_edges[0] is psi_node[0] + assert op.in_edges[0] is psi_node[1] + + op = qu.quantum_constructor([psi_node[0], psi_node[1]], []) + assert not op.is_scalar() + assert op.is_vector() + assert not op.is_adjoint_vector() + assert len(op.out_edges) == 2 + assert len(op.in_edges) == 0 + assert op.out_edges[0] is psi_node[0] + assert op.out_edges[1] is psi_node[1] + + op = qu.quantum_constructor([], [psi_node[0], psi_node[1]]) + assert not op.is_scalar() + assert not op.is_vector() + assert op.is_adjoint_vector() + assert len(op.out_edges) == 0 + assert len(op.in_edges) == 2 + assert op.in_edges[0] is psi_node[0] + assert op.in_edges[1] is psi_node[1] + + with pytest.raises(ValueError): + op = qu.quantum_constructor([], [], [psi_node]) + + _ = psi_node[0] ^ psi_node[1] + op = qu.quantum_constructor([], [], [psi_node]) + assert op.is_scalar() + assert not op.is_vector() + assert not op.is_adjoint_vector() + assert len(op.out_edges) == 0 + assert len(op.in_edges) == 0 + + +def test_checks(backend): + node1 = tn.Node(np.random.rand(2, 2), backend=backend) + node2 = tn.Node(np.random.rand(2, 2), backend=backend) + _ = node1[1] ^ node2[0] + + # extra dangling edges must be explicitly ignored + with pytest.raises(ValueError): + _ = qu.QuVector([node1[0]]) + + # correctly ignore the extra edge + _ = qu.QuVector([node1[0]], ignore_edges=[node2[1]]) + + # in/out edges must be dangling + with pytest.raises(ValueError): + _ = qu.QuVector([node1[0], node1[1], node2[1]]) + + +def test_from_tensor(backend): + psi_tensor = np.random.rand(2, 2) + + op = qu.QuOperator.from_tensor(psi_tensor, [0], [1], backend=backend) + assert not op.is_scalar() + assert not op.is_vector() + assert not op.is_adjoint_vector() + np.testing.assert_almost_equal(op.eval(), psi_tensor) + + op = qu.QuVector.from_tensor(psi_tensor, [0, 1], backend=backend) + assert not op.is_scalar() + assert op.is_vector() + assert not op.is_adjoint_vector() + np.testing.assert_almost_equal(op.eval(), psi_tensor) + + op = qu.QuAdjointVector.from_tensor(psi_tensor, [0, 1], backend=backend) + assert not op.is_scalar() + assert not op.is_vector() + assert op.is_adjoint_vector() + np.testing.assert_almost_equal(op.eval(), psi_tensor) + + op = qu.QuScalar.from_tensor(1.0, backend=backend) + assert op.is_scalar() + assert not op.is_vector() + assert not op.is_adjoint_vector() + assert op.eval() == 1.0 + + +def test_identity(backend): + E = qu.identity((2, 3, 4), backend=backend) + for n in E.nodes: + assert isinstance(n, tn.CopyNode) + twentyfour = E.trace() + for n in twentyfour.nodes: + assert isinstance(n, tn.CopyNode) + assert twentyfour.eval() == 24 + + tensor = np.random.rand(2, 2) + psi = qu.QuVector.from_tensor(tensor, backend=backend) + E = qu.identity((2, 2), backend=backend) + np.testing.assert_allclose((E @ psi).eval(), psi.eval()) + + np.testing.assert_allclose( + (psi.adjoint() @ E @ psi).eval(), psi.norm().eval()) + + op = qu.QuOperator.from_tensor(tensor, [0], [1], backend=backend) + op_I = op.tensor_product(E) + op_times_4 = op_I.partial_trace([1, 2]) + np.testing.assert_allclose(op_times_4.eval(), 4 * op.eval()) + + +def test_tensor_product(backend): + psi = qu.QuVector.from_tensor(np.random.rand(2, 2), backend=backend) + psi_psi = psi.tensor_product(psi) + assert len(psi_psi.subsystem_edges) == 4 + np.testing.assert_almost_equal(psi_psi.norm().eval(), psi.norm().eval()**2) + + +def test_matmul(backend): + mat = np.random.rand(2, 2) + op = qu.QuOperator.from_tensor(mat, [0], [1], backend=backend) + res = (op @ op).eval() + np.testing.assert_allclose(res, mat @ mat) + + +def test_mul(backend): + mat = np.eye(2) + scal = np.float64(0.5) + op = qu.QuOperator.from_tensor(mat, [0], [1], backend=backend) + scal_op = qu.QuScalar.from_tensor(scal, backend=backend) + + res = (op * scal_op).eval() + np.testing.assert_allclose(res, mat * 0.5) + + res = (scal_op * op).eval() + np.testing.assert_allclose(res, mat * 0.5) + + res = (scal_op * scal_op).eval() + np.testing.assert_almost_equal(res, 0.25) + + res = (op * np.float64(0.5)).eval() + np.testing.assert_allclose(res, mat * 0.5) + + res = (np.float64(0.5) * op).eval() + np.testing.assert_allclose(res, mat * 0.5) + + with pytest.raises(ValueError): + _ = (op * op) + + with pytest.raises(ValueError): + _ = (op * mat) + +def test_expectations(backend): + if backend == 'pytorch': + psi_tensor = np.random.rand(2, 2, 2) + op_tensor = np.random.rand(2, 2) + else: + psi_tensor = np.random.rand(2, 2, 2) + 1.j * np.random.rand(2, 2, 2) + op_tensor = np.random.rand(2, 2) + 1.j * np.random.rand(2, 2) + + psi = qu.QuVector.from_tensor(psi_tensor, backend=backend) + op = qu.QuOperator.from_tensor(op_tensor, [0], [1], backend=backend) + + op_3 = op.tensor_product( + qu.identity((2, 2), backend=backend, dtype=psi_tensor.dtype)) + res1 = (psi.adjoint() @ op_3 @ psi).eval() + + rho_1 = psi.reduced_density([1, 2]) # trace out sites 2 and 3 + res2 = (op @ rho_1).trace().eval() + + np.testing.assert_almost_equal(res1, res2) + + +def test_projector(backend): + psi_tensor = np.random.rand(2, 2) + psi_tensor /= np.linalg.norm(psi_tensor) + psi = qu.QuVector.from_tensor(psi_tensor, backend=backend) + P = psi.projector() + np.testing.assert_allclose((P @ psi).eval(), psi_tensor) + + np.testing.assert_allclose((P @ P).eval(), P.eval()) diff --git a/tensornetwork/tests/backend_contextmanager_test.py b/tensornetwork/tests/backend_contextmanager_test.py new file mode 100644 index 000000000..644c055d0 --- /dev/null +++ b/tensornetwork/tests/backend_contextmanager_test.py @@ -0,0 +1,59 @@ +import tensornetwork as tn +from tensornetwork.backend_contextmanager import _default_backend_stack +import pytest +import numpy as np + + +def test_contextmanager_simple(): + with tn.DefaultBackend("tensorflow"): + a = tn.Node(np.ones((10,))) + b = tn.Node(np.ones((10,))) + assert a.backend.name == b.backend.name + + +def test_contextmanager_default_backend(): + tn.set_default_backend("pytorch") + with tn.DefaultBackend("numpy"): + assert _default_backend_stack.default_backend == "pytorch" + + +def test_contextmanager_interruption(): + tn.set_default_backend("pytorch") + with pytest.raises(AssertionError): + with tn.DefaultBackend("numpy"): + tn.set_default_backend("tensorflow") + + +def test_contextmanager_nested(): + with tn.DefaultBackend("tensorflow"): + a = tn.Node(np.ones((10,))) + assert a.backend.name == "tensorflow" + with tn.DefaultBackend("numpy"): + b = tn.Node(np.ones((10,))) + assert b.backend.name == "numpy" + c = tn.Node(np.ones((10,))) + assert c.backend.name == "tensorflow" + d = tn.Node(np.ones((10,))) + assert d.backend.name == "numpy" + + +def test_contextmanager_wrong_item(): + a = tn.Node(np.ones((10,))) + with pytest.raises(ValueError): + with tn.DefaultBackend(a): # pytype: disable=wrong-arg-types + pass + + +def test_contextmanager_BaseBackend(): + tn.set_default_backend("pytorch") + a = tn.Node(np.ones((10,))) + with tn.DefaultBackend(a.backend): + b = tn.Node(np.ones((10,))) + assert b.backend.name == "pytorch" + + +def test_set_default_backend_value_error(): + tn.set_default_backend("pytorch") + with pytest.raises(ValueError, match="Item passed to set_default_backend " + "must be Text or BaseBackend"): + tn.set_default_backend(-1) # pytype: disable=wrong-arg-types diff --git a/tensornetwork/tests/network_components_free_test.py b/tensornetwork/tests/network_components_free_test.py index 47632debe..5a1e27ba0 100644 --- a/tensornetwork/tests/network_components_free_test.py +++ b/tensornetwork/tests/network_components_free_test.py @@ -1,12 +1,15 @@ import numpy as np import tensorflow as tf +import torch import pytest +from unittest.mock import patch from collections import namedtuple import h5py import re #pylint: disable=line-too-long -from tensornetwork.network_components import Node, CopyNode, Edge, NodeCollection +from tensornetwork.network_components import Node, CopyNode, Edge, NodeCollection, BaseNode, _remove_trace_edge, _remove_edges import tensornetwork as tn +from tensornetwork.backends.base_backend import BaseBackend string_type = h5py.special_dtype(vlen=str) @@ -15,6 +18,46 @@ 'node1 node2 edge1 edge12 tensor') +class TestNode(BaseNode): + + def get_tensor(self): #pylint: disable=useless-super-delegation + return super().get_tensor() + + def set_tensor(self, tensor): #pylint: disable=useless-super-delegation + return super().set_tensor(tensor) + + def __add__(self, other): #pylint: disable=useless-super-delegation + return super().__add__(other) + + def __sub__(self, other): #pylint: disable=useless-super-delegation + return super().__sub__(other) + + def __mul__(self, other): #pylint: disable=useless-super-delegation + return super().__mul__(other) + + def __truediv__(self, other): #pylint: disable=useless-super-delegation + return super().__truediv__(other) + + @property + def shape(self): + return super().shape + + @property + def tensor(self): + return super().tensor + + @tensor.setter + def tensor(self, tensor): + return super(TestNode, type(self)).tensor.fset(self, tensor) + + def _load_node(self, node_data):# pylint: disable=useless-super-delegation + return super()._load_node(node_data) + + def _save_node(self, node_group): #pylint: disable=useless-super-delegation + return super()._save_node(node_group) + + + @pytest.fixture(name='single_node_edge') def fixture_single_node_edge(backend): tensor = np.ones((1, 2, 2)) @@ -249,6 +292,16 @@ def test_node_reorder_edges_raise_error_trace_edge(single_node_edge): assert "Edge reordering does not support trace edges." in str(e.value) +def test_node_reorder_edges_raise_error_no_tensor(single_node_edge): + node = single_node_edge.node + e2 = tn.connect(node[1], node[2]) + e3 = node[0] + del node._tensor + with pytest.raises(AttributeError) as e: + node.reorder_edges([e2, e3]) + assert "Please provide a valid tensor for this Node." in str(e.value) + + def test_node_magic_getitem(single_node_edge): node = single_node_edge.node edge = single_node_edge.edge @@ -279,13 +332,40 @@ def test_node_magic_lt(double_node_edge): def test_node_magic_lt_raises_error_not_node(single_node_edge): node = single_node_edge.node with pytest.raises(ValueError): - assert node < 0 + node < 0 def test_node_magic_matmul_raises_error_not_node(single_node_edge): node = single_node_edge.node with pytest.raises(TypeError): - assert node @ 0 + node @ 0 + + +def test_node_magic_matmul_raises_error_no_tensor(single_node_edge): + node = single_node_edge.node + del node._tensor + with pytest.raises(AttributeError): + node @ node + + +def test_node_magic_matmul_raises_error_disabled_node(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node @ node + + +def test_node_edges_getter_raises_error_disabled_node(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.edges + +def test_node_edges_setter_raises_error_disabled_node(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.edges = [] def test_node_magic_matmul_raises_error_different_network(single_node_edge): @@ -316,6 +396,320 @@ def test_node_magic_matmul(backend): np.testing.assert_allclose(actual.tensor, expected) +def test_between_node_add_op(backend): + node1 = Node(tensor=np.array([[1, 2], [3, 4]]), backend=backend) + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend=backend) + node3 = Node(tensor=np.array([[1., 2.], [3., 4.]]), backend=backend) + int_node = Node(tensor=np.array(2, dtype=np.int64), backend=backend) + float_node = Node(tensor=np.array(2.5, dtype=np.float64), backend=backend) + + expected = np.array([[11, 12], [13, 14]]) + result = (node1 + node2).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == node2.tensor.dtype == result.dtype + + expected = np.array([[3, 4], [5, 6]]) + result = (node1 + int_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == int_node.tensor.dtype == result.dtype + expected = np.array([[3, 4], [5, 6]]) + result = (int_node + node1).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == int_node.tensor.dtype == result.dtype + + expected = np.array([[3.5, 4.5], [5.5, 6.5]]) + result = (node3 + float_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + expected = np.array([[3.5, 4.5], [5.5, 6.5]]) + result = (float_node + node3).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + + +def test_node_and_scalar_add_op(backend): + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.int32), backend=backend) + expected = np.array([[3, 4], [5, 6]]) + result = (node + 2).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'int64' + else: + assert node.tensor.dtype == result.dtype + + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend) + expected = np.array([[3.5, 4.5], [5.5, 6.5]]) + result = (node + 2.5).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'float64' + else: + assert node.tensor.dtype == result.dtype + + +def test_between_node_sub_op(backend): + node1 = Node(tensor=np.array([[1, 2], [3, 4]]), backend=backend) + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend=backend) + node3 = Node(tensor=np.array([[1., 2.], [3., 4.]]), backend=backend) + int_node = Node(tensor=np.array(2, dtype=np.int64), backend=backend) + float_node = Node(tensor=np.array(2.5, dtype=np.float64), backend=backend) + + expected = np.array([[-9, -8], [-7, -6]]) + result = (node1 - node2).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == node2.tensor.dtype == result.dtype + + expected = np.array([[-1, 0], [1, 2]]) + result = (node1 - int_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == int_node.tensor.dtype == result.dtype + expected = np.array([[1, 0], [-1, -2]]) + result = (int_node - node1).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == int_node.tensor.dtype == result.dtype + + expected = np.array([[-1.5, -0.5], [0.5, 1.5]]) + result = (node3 - float_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + expected = np.array([[1.5, 0.5], [-0.5, -1.5]]) + result = (float_node - node3).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + + +def test_node_and_scalar_sub_op(backend): + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.int32), backend=backend) + expected = np.array([[-1, 0], [1, 2]]) + result = (node - 2).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'int64' + else: + assert node.tensor.dtype == result.dtype + + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend) + expected = np.array([[-1.5, -0.5], [0.5, 1.5]]) + result = (node - 2.5).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'float64' + else: + assert node.tensor.dtype == result.dtype + + +def test_between_node_mul_op(backend): + node1 = Node(tensor=np.array([[1, 2], [3, 4]]), backend=backend) + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend=backend) + node3 = Node(tensor=np.array([[1., 2.], [3., 4.]]), backend=backend) + int_node = Node(tensor=np.array(2, dtype=np.int64), backend=backend) + float_node = Node(tensor=np.array(2.5, dtype=np.float64), backend=backend) + + expected = np.array([[10, 20], [30, 40]]) + result = (node1 * node2).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == node2.tensor.dtype == result.dtype + + expected = np.array([[2, 4], [6, 8]]) + result = (node1 * int_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == int_node.tensor.dtype == result.dtype + result = (int_node * node1).tensor + np.testing.assert_almost_equal(result, expected) + + expected = np.array([[2.5, 5], [7.5, 10]]) + result = (node3 * float_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + result = (float_node * node3).tensor + np.testing.assert_almost_equal(result, expected) + assert node3.dtype == float_node.dtype == result.dtype + + +def test_node_and_scalar_mul_op(backend): + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.int32), backend=backend) + expected = np.array([[2, 4], [6, 8]]) + result = (node * 2).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'int64' + else: + assert node.tensor.dtype == result.dtype + + node = Node(tensor=np.array([[1, 2], [3, 4]], dtype=np.float32), backend=backend) + expected = np.array([[2.5, 5], [7.5, 10]]) + result = (node * 2.5).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'float64' + else: + assert node.tensor.dtype == result.dtype + + +def test_between_node_div_op(backend): + node1 = Node(tensor=np.array([[1., 2.], [3., 4.]]), backend=backend) + node2 = Node(tensor=np.array([[10., 10.], [10., 10.]]), backend=backend) + node3 = Node(tensor=np.array([[1, 2], [3, 4]]), backend=backend) + int_node = Node(tensor=np.array(2, dtype=np.int64), backend=backend) + float_node = Node(tensor=np.array(2.5, dtype=np.float64), backend=backend) + + expected = np.array([[0.1, 0.2], [0.3, 0.4]]) + result = (node1 / node2).tensor + np.testing.assert_almost_equal(result, expected) + assert node1.tensor.dtype == node2.tensor.dtype == result.dtype + + expected = np.array([[0.5, 1.], [1.5, 2.]]) + expected_pytorch = np.array([[0, 1], [1, 2]]) + result = (node3 / int_node).tensor + if backend == 'pytorch': + np.testing.assert_almost_equal(result, expected_pytorch) + assert node3.tensor.dtype == result.dtype == torch.int64 + else: + np.testing.assert_almost_equal(result, expected) + assert node3.tensor.dtype == 'int64' + assert result.dtype == 'float64' + + expected = np.array([[2., 1.], [2/3, 0.5]]) + expected_pytorch = np.array([[2, 1], [0, 0]]) + result = (int_node / node3).tensor + if backend == 'pytorch': + np.testing.assert_almost_equal(result, expected_pytorch) + assert node3.tensor.dtype == result.dtype == torch.int64 + else: + np.testing.assert_almost_equal(result, expected) + assert node3.tensor.dtype == 'int64' + assert result.dtype == 'float64' + + expected = np.array([[4., 4.], [4., 4.]]) + result = (node2 / float_node).tensor + np.testing.assert_almost_equal(result, expected) + assert node2.dtype == float_node.dtype == result.dtype + expected = np.array([[0.25, 0.25], [0.25, 0.25]]) + result = (float_node / node2).tensor + np.testing.assert_almost_equal(result, expected) + assert node2.dtype == float_node.dtype == result.dtype + + +def test_node_and_scalar_div_op(backend): + node = Node(tensor=np.array([[5, 10], [15, 20]], dtype=np.int32), backend=backend) + expected = np.array([[0.5, 1.], [1.5, 2.]]) + expected_pytorch = np.array([[0, 1], [1, 2]]) + result = (node / 10).tensor + if backend == 'pytorch': + np.testing.assert_almost_equal(result, expected_pytorch) + assert node.tensor.dtype == result.dtype == torch.int32 + else: + np.testing.assert_almost_equal(result, expected) + assert result.dtype == 'float64' + assert node.tensor.dtype == 'int32' + + node = Node(tensor=np.array([[5., 10.], [15., 20.]], dtype=np.float32), backend=backend) + expected = np.array([[2., 4.], [6., 8.]]) + result = (node / 2.5).tensor + np.testing.assert_almost_equal(result, expected) + if backend == 'jax': + assert result.dtype == 'float64' + else: + assert node.tensor.dtype == result.dtype + + +def test_node_add_input_error(): + #pylint: disable=unused-variable + #pytype: disable=unsupported-operands + node1 = Node(tensor=2, backend='numpy') + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='numpy') + + del node1._tensor + with pytest.raises(AttributeError): + result = node1 + node2 + result = node2 + node1 + + node1.tensor = 1 + node2 = 'str' + copynode = tn.CopyNode(rank=4, dimension=3) + with pytest.raises(TypeError): + result = node1 + node2 + result = node1 + copynode + + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch') + with pytest.raises(TypeError): + result = node1 + node2 + #pytype: enable=unsupported-operands + + +def test_node_sub_input_error(): + #pylint: disable=unused-variable + #pytype: disable=unsupported-operands + node1 = Node(tensor=2, backend='numpy') + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='numpy') + + del node1._tensor + with pytest.raises(AttributeError): + result = node1 - node2 + result = node2 - node1 + + node1.tensor = 1 + node2 = 'str' + copynode = tn.CopyNode(rank=4, dimension=3) + with pytest.raises(TypeError): + result = node1 - node2 + result = node1 - copynode + + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch') + with pytest.raises(TypeError): + result = node1 - node2 + #pytype: enable=unsupported-operands + + +def test_node_mul_input_error(): + #pylint: disable=unused-variable + #pytype: disable=unsupported-operands + node1 = Node(tensor=2, backend='numpy') + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='numpy') + + del node1._tensor + with pytest.raises(AttributeError): + result = node1 * node2 + result = node2 * node1 + + node1.tensor = 1 + node2 = 'str' + copynode = tn.CopyNode(rank=4, dimension=3) + with pytest.raises(TypeError): + result = node1 * node2 + result = node1 * copynode + + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch') + with pytest.raises(TypeError): + result = node1 * node2 + #pytype: enable=unsupported-operands + + +def test_node_div_input_error(): + #pylint: disable=unused-variable + #pytype: disable=unsupported-operands + node1 = Node(tensor=2, backend='numpy') + node1 = Node(tensor=2, backend='numpy') + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='numpy') + + del node1._tensor + with pytest.raises(AttributeError): + result = node1 / node2 + result = node2 / node1 + + node1.tensor = 1 + node2 = 'str' + copynode = tn.CopyNode(rank=4, dimension=3) + with pytest.raises(TypeError): + result = node1 / node2 + result = node1 / copynode + + node2 = Node(tensor=np.array([[10, 10], [10, 10]]), backend='pytorch') + with pytest.raises(TypeError): + result = node1 / node2 + #pytype: enable=unsupported-operands + + def test_node_save_structure(tmp_path, single_node_edge): node = single_node_edge.node with h5py.File(tmp_path / 'nodes', 'w') as node_file: @@ -918,4 +1312,284 @@ def test_repr_for_Nodes_and_Edges(double_node_edge): assert "[[[1.,1.],[1.,1.]]]" in str(node1) and str(node2) assert "Edge(DanglingEdge)[0]" in str(node1) and str(node2) assert "Edge('test_node1'[1]->'test_node2'[1])" in str(node1) and str(node2) - assert "Edge(DanglingEdge)[2]" in str(node1) and str(node2) \ No newline at end of file + assert "Edge(DanglingEdge)[2]" in str(node1) and str(node2) + + +def test_base_node_name_list_throws_error(): + with pytest.raises(TypeError,): + TestNode(name=["A"], axis_names=['a', 'b']) # pytype: disable=wrong-arg-types + + +def test_base_node_name_int_throws_error(): + with pytest.raises(TypeError): + TestNode(name=1, axis_names=['a', 'b']) # pytype: disable=wrong-arg-types + + +def test_base_node_axis_names_int_throws_error(): + with pytest.raises(TypeError): + TestNode(axis_names=[0, 1]) # pytype: disable=wrong-arg-types + + +def test_base_node_no_axis_names_no_shapes_throws_error(): + with pytest.raises(ValueError): + TestNode(name='a') + + +def test_node_add_axis_names_int_throws_error(): + n1 = Node(np.eye(2), axis_names=['a', 'b']) + with pytest.raises(TypeError): + n1.add_axis_names([0, 1]) # pytype: disable=wrong-arg-types + + +def test_node_axis_names_setter_throws_shape_large_mismatch_error(): + n1 = Node(np.eye(2), axis_names=['a', 'b']) + with pytest.raises(ValueError): + n1.axis_names = ['a', 'b', 'c'] + + +def test_node_axis_names_setter_throws_shape_small_mismatch_error(): + n1 = Node(np.eye(2), axis_names=['a', 'b']) + with pytest.raises(ValueError): + n1.axis_names = ['a'] + + +def test_node_axis_names_setter_throws_value_error(): + n1 = Node(np.eye(2), axis_names=['a', 'b']) + with pytest.raises(TypeError): + n1.axis_names = [0, 1] + + +def test_node_dtype(backend): + n1 = Node(np.random.rand(2), backend=backend) + assert n1.dtype == n1.tensor.dtype + + +@pytest.mark.parametrize("name", [1, ['1']]) +def test_node_set_name_raises_type_error(backend, name): + n1 = Node(np.random.rand(2), backend=backend) + with pytest.raises(TypeError): + n1.set_name(name) + + +@pytest.mark.parametrize("name", [1, ['1']]) +def test_node_name_setter_raises_type_error(backend, name): + n1 = Node(np.random.rand(2), backend=backend) + with pytest.raises(TypeError): + n1.name = name + + +def test_base_node_get_tensor(): + n1 = TestNode(name="n1", axis_names=['a'], shape=(1,)) + assert n1.get_tensor() is None + + +def test_base_node_set_tensor(): + n1 = TestNode(name="n1", axis_names=['a'], shape=(1,)) + assert n1.set_tensor(np.random.rand(2)) is None + assert n1.tensor is None + + +def test_base_node_shape(): + n1 = TestNode(name="n1", axis_names=['a'], shape=(1,)) + n1._shape = None + with pytest.raises(ValueError): + n1.shape + + +def test_base_node_tensor_getter(): + n1 = TestNode(name="n1", axis_names=['a'], shape=(1,)) + assert n1.tensor is None + + +def test_base_node_tensor_setter(): + n1 = TestNode(name="n1", axis_names=['a'], shape=(1,)) + n1.tensor = np.random.rand(2) + assert n1.tensor is None + + +def test_node_has_dangling_edge_false(double_node_edge): + node1 = double_node_edge.node1 + node2 = double_node_edge.node2 + tn.connect(node1["a"], node2["a"]) + tn.connect(node1["c"], node2["c"]) + assert not node1.has_dangling_edge() + + +def test_node_has_dangling_edge_true(single_node_edge): + assert single_node_edge.node.has_dangling_edge() + + +def test_node_get_item(single_node_edge): + node = single_node_edge.node + edge = single_node_edge.edge + node.add_edge(edge, axis=0) + assert node[0] == edge + assert edge in node[0:2] + + +def test_node_signature_getter_disabled_throws_error(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.signature + + +def test_node_signature_setter_disabled_throws_error(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.signature = "signature" + + +def test_node_disabled_disabled_throws_error(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.disable() + + +def test_node_disabled_shape_throws_error(single_node_edge): + node = single_node_edge.node + node.is_disabled = True + with pytest.raises(ValueError): + node.shape + + +def test_copy_node_get_partners_with_trace(backend): + node1 = CopyNode(4, 2, backend=backend) + node2 = Node(np.random.rand(2, 2), backend=backend, name="node2") + tn.connect(node1[0], node1[1]) + tn.connect(node1[2], node2[0]) + tn.connect(node1[3], node2[1]) + assert node1.get_partners() == {node2: {0, 1}} + + +@pytest.mark.parametrize("name", [1, ['1']]) +def test_edge_name_throws_type_error(single_node_edge, name): + with pytest.raises(TypeError): + Edge(node1=single_node_edge.node, axis1=0, name=name) + + +def test_edge_name_setter_disabled_throws_error(single_node_edge): + edge = Edge(node1=single_node_edge.node, axis1=0) + edge.is_disabled = True + with pytest.raises(ValueError): + edge.name = 'edge' + + +def test_edge_name_getter_disabled_throws_error(single_node_edge): + edge = Edge(node1=single_node_edge.node, axis1=0) + edge.is_disabled = True + with pytest.raises(ValueError): + edge.name + + +@pytest.mark.parametrize("name", [1, ['1']]) +def test_edge_name_setter_throws_type_error(single_node_edge, name): + edge = Edge(node1=single_node_edge.node, axis1=0) + with pytest.raises(TypeError): + edge.name = name + + +def test_edge_signature_getter_disabled_throws_error(single_node_edge): + edge = Edge(node1=single_node_edge.node, axis1=0) + edge.is_disabled = True + with pytest.raises(ValueError): + edge.signature + + +def test_edge_signature_setter_disabled_throws_error(single_node_edge): + edge = Edge(node1=single_node_edge.node, axis1=0) + edge.is_disabled = True + with pytest.raises(ValueError): + edge.signature = "signature" + + +def test_edge_node1_throws_value_error(single_node_edge): + edge = Edge(node1=single_node_edge.node, axis1=0, name="edge") + edge._node1 = None + err_msg = "node1 for edge 'edge' no longer exists." + with pytest.raises(ValueError, match=err_msg): + edge.node1 + + + +def test_edge_node2_throws_value_error(single_node_edge): + edge = tn.connect(single_node_edge.node[1], single_node_edge.node[2]) + edge.name = 'edge' + edge._node2 = None + err_msg = "node2 for edge 'edge' no longer exists." + with pytest.raises(ValueError, match=err_msg): + edge.node2 + + +@pytest.mark.parametrize("name", [1, ['1']]) +def test_edge_set_name_throws_type_error(single_node_edge, name): + edge = Edge(node1=single_node_edge.node, axis1=0) + with pytest.raises(TypeError): + edge.set_name(name) + + +@patch.object(Edge, "name", None) +def test_edge_str(single_node_edge): + single_node_edge.edge.name = None + assert str(single_node_edge.edge) == "__unnamed_edge__" + + +def test_get_all_dangling_single_node(single_node_edge): + node = single_node_edge.node + assert set(tn.get_all_dangling({node})) == set(node.edges) + + +def test_get_all_dangling_double_node(double_node_edge): + node1 = double_node_edge.node1 + node2 = double_node_edge.node2 + assert set(tn.get_all_dangling({node1, node2})) == {node1[0], node1[2], + node2[0], node2[2]} + + +def test_flatten_edges_different_backend_raises_value_error(single_node_edge): + node1 = single_node_edge.node + node2 = tn.Node(np.random.rand(2, 2, 2)) + node2.backend = BaseBackend() + with pytest.raises(ValueError): + tn.flatten_edges(node1.get_all_edges()+node2.get_all_edges()) + + +def test_split_edge_trivial(single_node_edge): + edge = single_node_edge.edge + assert tn.split_edge(edge, (1,)) == [edge] + + +def test_split_edge_different_backend_raises_value_error(single_node_edge): + if single_node_edge.node.backend.name == "numpy": + pytest.skip("numpy comparing to all the others") + node1 = single_node_edge.node + node2 = tn.Node(np.random.rand(2, 2, 2), backend="numpy") + edge = tn.connect(node1[1], node2[1]) + with pytest.raises(ValueError, match="Not all backends are the same."): + tn.split_edge(edge, (2, 1)) + + +def test_remove_trace_edge_dangling_edge_raises_value_error(single_node_edge): + node = single_node_edge.node + edge = node[0] + edge.name = "e" + with pytest.raises(ValueError, match="Attempted to remove dangling edge 'e"): + _remove_trace_edge(edge, node) + + +def test_remove_trace_edge_non_trace_raises_value_error(double_node_edge): + node1 = double_node_edge.node1 + node2 = double_node_edge.node2 + edge = tn.connect(node1[0], node2[0]) + edge.name = "e" + with pytest.raises(ValueError, match="Edge 'e' is not a trace edge."): + _remove_trace_edge(edge, node1) + + +def test_remove_edges_trace_raises_value_error(single_node_edge): + node = single_node_edge.node + edge = tn.connect(node[1], node[2]) + with pytest.raises(ValueError): + _remove_edges(edge, node, node, node) # pytype: disable=wrong-arg-types diff --git a/tensornetwork/tests/network_operations_test.py b/tensornetwork/tests/network_operations_test.py index 189c08594..4b13d5656 100644 --- a/tensornetwork/tests/network_operations_test.py +++ b/tensornetwork/tests/network_operations_test.py @@ -15,6 +15,7 @@ import tensornetwork as tn import pytest import numpy as np +from tensornetwork.backends.base_backend import BaseBackend def test_split_node_full_svd_names(backend): @@ -334,3 +335,89 @@ def test_switch_backend(backend): nodes = [a, b, c] tn.switch_backend(nodes, backend) assert nodes[0].backend.name == backend + + +def test_norm_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.norm(node) + + +def test_conj_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.conj(node) + + +def test_transpose_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.transpose(node, permutation=[]) + + +def test_split_node_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.split_node(node, left_edges=[], right_edges=[]) + + +def test_split_node_qr_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.split_node_qr(node, left_edges=[], right_edges=[]) + + +def test_split_node_rq_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.split_node_rq(node, left_edges=[], right_edges=[]) + + +def test_split_node_full_svd_of_node_without_backend_raises_error(): + node = np.random.rand(3, 3, 3) + with pytest.raises(AttributeError): + tn.split_node_full_svd(node, left_edges=[], right_edges=[]) + + +def test_reachable_raises_value_error(): + with pytest.raises(ValueError): + tn.reachable({}) + + +def test_check_correct_raises_value_error_1(backend): + a = tn.Node(np.random.rand(3, 3, 3), backend=backend) + b = tn.Node(np.random.rand(3, 3, 3), backend=backend) + edge = a.edges[0] + edge.node1 = b + edge.node2 = b + with pytest.raises(ValueError): + tn.check_correct({a, b}) + + +def test_check_correct_raises_value_error_2(backend): + a = tn.Node(np.random.rand(3, 3, 3), backend=backend) + b = tn.Node(np.random.rand(3, 3, 3), backend=backend) + edge = a.edges[0] + edge.axis1 = -1 + with pytest.raises(ValueError): + tn.check_correct({a, b}) + + +def test_get_all_nodes(backend): + a = tn.Node(np.random.rand(3, 3, 3), backend=backend) + b = tn.Node(np.random.rand(3, 3, 3), backend=backend) + edge = tn.connect(a[0], b[0]) + assert tn.get_all_nodes({edge}) == {a, b} + + +def test_contract_trace_edges(backend): + a = tn.Node(np.random.rand(3, 3, 3), backend=backend) + with pytest.raises(ValueError): + tn.contract_trace_edges(a) + + +def test_switch_backend_raises_error(backend): + a = tn.Node(np.random.rand(3, 3, 3)) + a.backend = BaseBackend() + with pytest.raises(NotImplementedError): + tn.switch_backend({a}, backend) diff --git a/tensornetwork/tests/network_test.py b/tensornetwork/tests/network_test.py index 8efcd9d4d..f4e83a892 100644 --- a/tensornetwork/tests/network_test.py +++ b/tensornetwork/tests/network_test.py @@ -483,26 +483,45 @@ def test_flatten_all_edges(backend): def test_contract_between(backend): - a_val = np.ones((2, 3, 4, 5)) - b_val = np.ones((3, 5, 4, 2)) + a_val = np.random.rand(2, 3, 4, 5) + b_val = np.random.rand(3, 5, 6, 2) a = tn.Node(a_val, backend=backend) b = tn.Node(b_val, backend=backend) tn.connect(a[0], b[3]) tn.connect(b[1], a[3]) tn.connect(a[1], b[0]) - edge_a = a[2] - edge_b = b[2] - c = tn.contract_between(a, b, name="New Node") - c.reorder_edges([edge_a, edge_b]) + output_axis_names = ["a2", "b2"] + c = tn.contract_between(a, b, name="New Node", axis_names=output_axis_names) tn.check_correct({c}) # Check expected values. a_flat = np.reshape(np.transpose(a_val, (2, 1, 0, 3)), (4, 30)) - b_flat = np.reshape(np.transpose(b_val, (2, 0, 3, 1)), (4, 30)) + b_flat = np.reshape(np.transpose(b_val, (2, 0, 3, 1)), (6, 30)) final_val = np.matmul(a_flat, b_flat.T) assert c.name == "New Node" + assert c.axis_names == output_axis_names np.testing.assert_allclose(c.tensor, final_val) +def test_contract_between_output_edge_order(backend): + a_val = np.random.rand(2, 3, 4, 5) + b_val = np.random.rand(3, 5, 6, 2) + a = tn.Node(a_val, backend=backend) + b = tn.Node(b_val, backend=backend) + tn.connect(a[0], b[3]) + tn.connect(b[1], a[3]) + tn.connect(a[1], b[0]) + output_axis_names = ["b2", "a2"] + c = tn.contract_between(a, b, name="New Node", axis_names=output_axis_names, + output_edge_order=[b[2], a[2]]) + # Check expected values. + a_flat = np.reshape(np.transpose(a_val, (2, 1, 0, 3)), (4, 30)) + b_flat = np.reshape(np.transpose(b_val, (2, 0, 3, 1)), (6, 30)) + final_val = np.matmul(a_flat, b_flat.T) + assert c.name == "New Node" + assert c.axis_names == output_axis_names + np.testing.assert_allclose(c.tensor, final_val.T) + + def test_contract_between_no_outer_product_value_error(backend): a_val = np.ones((2, 3, 4)) b_val = np.ones((5, 6, 7)) @@ -517,8 +536,45 @@ def test_contract_between_outer_product_no_value_error(backend): b_val = np.ones((5, 6, 7)) a = tn.Node(a_val, backend=backend) b = tn.Node(b_val, backend=backend) - c = tn.contract_between(a, b, allow_outer_product=True) + output_axis_names = ["a0", "a1", "a2", "b0", "b1", "b2"] + c = tn.contract_between(a, b, allow_outer_product=True, + axis_names=output_axis_names) assert c.shape == (2, 3, 4, 5, 6, 7) + assert c.axis_names == output_axis_names + + +def test_contract_between_outer_product_output_edge_order(backend): + a_val = np.ones((2, 3, 4)) + b_val = np.ones((5, 6, 7)) + a = tn.Node(a_val, backend=backend) + b = tn.Node(b_val, backend=backend) + output_axis_names = ["b0", "b1", "a0", "b2", "a1", "a2"] + c = tn.contract_between( + a, b, + allow_outer_product=True, + output_edge_order=[b[0], b[1], a[0], b[2], a[1], a[2]], + axis_names=output_axis_names) + assert c.shape == (5, 6, 2, 7, 3, 4) + assert c.axis_names == output_axis_names + + +def test_contract_between_trace(backend): + a_val = np.ones((2, 3, 2, 4)) + a = tn.Node(a_val, backend=backend) + tn.connect(a[0], a[2]) + c = tn.contract_between(a, a, axis_names=["1", "3"]) + assert c.shape == (3, 4) + assert c.axis_names == ["1", "3"] + + +def test_contract_between_trace_output_edge_order(backend): + a_val = np.ones((2, 3, 2, 4)) + a = tn.Node(a_val, backend=backend) + tn.connect(a[0], a[2]) + c = tn.contract_between(a, a, output_edge_order=[a[3], a[1]], + axis_names=["3", "1"]) + assert c.shape == (4, 3) + assert c.axis_names == ["3", "1"] def test_contract_parallel(backend): diff --git a/tensornetwork/tests/tensornetwork_test.py b/tensornetwork/tests/tensornetwork_test.py index 318d2053c..10cb4ce38 100644 --- a/tensornetwork/tests/tensornetwork_test.py +++ b/tensornetwork/tests/tensornetwork_test.py @@ -13,6 +13,7 @@ # limitations under the License. import tensornetwork as tn +from tensornetwork.backend_contextmanager import _default_backend_stack import pytest import numpy as np import tensorflow as tf @@ -345,6 +346,21 @@ def test_reorder_axes(backend): assert a.shape == (4, 2, 3) +def test_reorder_axes_raises_error_no_tensor(backend): + a = tn.Node(np.zeros((2, 3, 4)), backend=backend) + del a._tensor + with pytest.raises(AttributeError) as e: + a.reorder_axes([2, 0, 1]) + assert "Please provide a valid tensor for this Node." in str(e.value) + + +def test_reorder_axes_raises_error_bad_permutation(backend): + a = tn.Node(np.zeros((2, 3, 4)), backend=backend) + with pytest.raises(ValueError) as e: + a.reorder_axes([2, 0]) + assert "A full permutation was not passed." in str(e.value) + + def test_flatten_consistent_result(backend): a_val = np.ones((3, 5, 5, 6)) b_val = np.ones((5, 6, 4, 5)) @@ -507,7 +523,7 @@ def test_set_node2(backend): def test_set_default(backend): tn.set_default_backend(backend) - assert tn.config.default_backend == backend + assert _default_backend_stack.default_backend == backend a = tn.Node(np.eye(2)) assert a.backend.name == backend diff --git a/tensornetwork/version.py b/tensornetwork/version.py index fb3fd385e..45642f983 100644 --- a/tensornetwork/version.py +++ b/tensornetwork/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = '0.2.0' +__version__ = '0.2.1'