From bcd95aa86da140587716561a8687454b334ea13e Mon Sep 17 00:00:00 2001
From: Danil Lykov <lkv97dn@gmail.com>
Date: Wed, 3 Jun 2020 00:12:52 -0500
Subject: [PATCH] add option to choose algorithm; update readme

---
 scratchpad/README.md      | 42 ++++++++++++++++++++++++++++++++++++++-
 scratchpad/contr_sheme.py | 28 +++++++++++++++++++++-----
 scratchpad/expr.py        |  1 -
 scratchpad/expr_optim.py  |  9 +++++++--
 4 files changed, 71 insertions(+), 9 deletions(-)

diff --git a/scratchpad/README.md b/scratchpad/README.md
index 1386b5a1..ee6caacf 100644
--- a/scratchpad/README.md
+++ b/scratchpad/README.md
@@ -1,6 +1,46 @@
 # Usage
 
-`python problem_graph_json.py 30 |mongocat -W -d tensim problem_graphs`
+`python expr.py 30 |mongocat -W -d tensim problem_graphs`
+
+## Search for a tensor expression
+
+```bash
+⟩ mongocat -F -d tensim expr_graphs --query '{"n_qubits":100}'
+[...bson....]
+```
+
+## Optimize a tensor expression
+
+```bash
+ptyhon contr_sheme.py --ordering <ordering>
+```
+
+where ordering is one of 'qbb' or 'nghs'
+
+
+## Compositions
+
+The `contr_sheme.py` script accepts input in json, which means you can combine the 
+commands. Feel the power of UNIX philosophy!
+
+### Create tensor expression and pipe it to ordering
+
+```bash
+⟩ python expr.py 47 --qaoa-layers 1 | python contr_sheme.py --ordering nghs \
+  | mongocat -W -d tensim contr_schemes
+```
+
+
+### Optimize straight from db and store the result
+
+for example, take 100 qubits and p=1 tensor expression, and run QuickBB on it
+
+```bash
+⟩ mongocat -F -d tensim expr_graphs --query '{"n_qubits":100, "extra.p":1}' |\
+    python contr_sheme.py --ordering qbb \|
+    mongocat -W -d tensim contr_schemes
+```
+
 
 ## Count objects in DB
 
diff --git a/scratchpad/contr_sheme.py b/scratchpad/contr_sheme.py
index 153fb685..9d1aeaa4 100644
--- a/scratchpad/contr_sheme.py
+++ b/scratchpad/contr_sheme.py
@@ -1,6 +1,7 @@
 import json
 import fire
 import sys
+import time
 import networkx as nx
 
 import expr_optim
@@ -24,11 +25,15 @@ def read_expressions():
 
 
 
-def as_json(read=True):
+def as_json(read=True
+            , ordering='qbb'
+           ):
     """
     Return a JSON representation of optimized expression
 
     """
+    if ordering not in ['qbb', 'nghs']:
+        raise Exception("ordering should be one of 'nghs', 'qbb'")
     if read:
         exprs = read_expressions()
     else:
@@ -41,11 +46,14 @@ def as_json(read=True):
     for G, expr_id in exprs:
         print(G.nodes())
 
-        peo, tw, G = expr_optim.optimize_order(G)
+        
+        start_t = time.time()
+        peo, tw, G = expr_optim.optimize_order(G, ordering_algo=ordering)
+        end_t = time.time()
 
         to_db =  {}
         # TODO: generator of id should be separate
-        to_db['_id'] = 'ordered_qbb.'+ expr_id
+        to_db['_id'] = f'ordered_{ordering}.'+ expr_id
         to_db['expr_id'] = expr_id
         # Note: mongodb will try to use all the nested dicts,
         #       so store the graph as string
@@ -56,9 +64,19 @@ def as_json(read=True):
             except KeyError:
                 pass
 
+        def get_machine_id():
+            import socket
+            import getpass
+            host = socket.gethostname()
+            user = getpass.getuser()
+            return user + '@' + host
+
         to_db['transforms'] =[{
-            'type':'order',
-            'params': {'order':peo, 'tw':tw, 'algo':'QuickBB'}
+            'type':'order'
+            ,'contract': {'order':peo}
+            ,'extra': {'tw':tw, 'algo': ordering
+                       ,'perf':{'time':end_t-start_t, 'run_by':get_machine_id()}
+                      }
         }]
 
         str = json.dumps(to_db)
diff --git a/scratchpad/expr.py b/scratchpad/expr.py
index 171f0e64..81e88dab 100644
--- a/scratchpad/expr.py
+++ b/scratchpad/expr.py
@@ -19,7 +19,6 @@ def get_tensors_from_graph(graph):
     return tensors
 
 def get_id(**args):
-    print(args)
     id = "p{p}_expr.S{S}_{type}_d{degree}_seed{seed}".format(**args)
     return id
 
diff --git a/scratchpad/expr_optim.py b/scratchpad/expr_optim.py
index 1986fe15..2a53f89a 100644
--- a/scratchpad/expr_optim.py
+++ b/scratchpad/expr_optim.py
@@ -5,6 +5,7 @@
 sys.path.append('..')
 
 import utils_qaoa
+import utils
 import expr
 import qtree
 
@@ -89,8 +90,12 @@ def test_reorder_tensors_small():
     assert perm_tensors[1]['indices'] == (0,0)
     assert perm_tensors[2]['indices'] == (1,2,3)
 
-def optimize_order(graph):
-    peo, tw = qtree.graph_model.get_peo(graph, int_vars=True)
+def optimize_order(graph, ordering_algo='qbb'):
+    if ordering_algo=='qbb':
+        peo, tw = qtree.graph_model.get_peo(graph, int_vars=True)
+    if ordering_algo=='nghs':
+        peo, _nghs = utils.get_locale_peo(graph, utils.n_neighbors)
+        tw = max(_nghs) - 1
     mapping = {v:i for i, v in enumerate(peo)}
     tensors = expr.get_tensors_from_graph(graph)
     reorder_tensors(tensors, peo)