fixing recursion depth issues, save best intermediate solutions, show scipy optimize progress, fix mistake in target terms that occurs for some heralding cases

Author: soerenarlt

pytheus/fancy_classes.py

@@ -301,8 +301,7 @@ def getStateCatalog(self, order=None, full=False):
             else:
                 self._state_catalog = th.stateCatalog(th.findEdgeCovers(self.edges, 
                                                         order=order, loops=self.loops))
-        
-        if self._state_catalog_tensor is not None:
+        if self._state_catalog_tensor is None:
             self._state_catalog_tensor = {} 
             for ket, pm_list in self._state_catalog.items():
                 self._state_catalog_tensor[ket] = [[self.complete_graph_edges.index(edge) for edge in pm] for pm in pm_list]

pytheus/leiwand.py

@@ -128,8 +128,8 @@ def leiwand(data, name='graph'):
             poly = reversed(list(dict(sorted(poly.items(), key=lambda x: x[0])).values()))
         for i, coord in enumerate(poly):
             print(r"\node[vertex] ({name}) at ({x},{y}) {xname};".format(name=vertices[i],
-                                                                         xname=r"{\color{fontcolor}" + vertices[
-                                                                             i] + "}", x=coord[0], y=coord[1]),
+                                                                         xname=r"{\color{fontcolor}" + str(vertices[
+                                                                             i]) + "}", x=coord[0], y=coord[1]),
                   file=outf)
 
         # edge_string = r"\path ({v1}) edge[{options}, opacity={opacity}] ({v2});"

pytheus/lossfunctions.py

@@ -180,13 +180,13 @@ def count_rate(graph, target_state, cnfg):
     new_loss = cnfg.get("new_loss", False)
     if not new_loss:
         # set up target equation
-        target = target_state.targetEquation(state_catalog=graph.state_catalog, imaginary=cnfg["imaginary"])
         # get variable names
         variables = th.stringEdges(graph.edges, imaginary=cnfg["imaginary"])
 
         # non-heralded, post-selection case
         if not cnfg["heralding_out"]:
             # only looking at perfect matchings
+            target = target_state.targetEquation(state_catalog=graph.state_catalog, imaginary=cnfg["imaginary"])
             graph.getNorm()
             norm = graph.norm
         # heralded case, more complicated selection rules
@@ -199,6 +199,7 @@ def count_rate(graph, target_state, cnfg):
                 edgecovers = brutal_covers(cnfg, graph)
             cat = th.stateCatalog(edgecovers)
             norm = th.writeNorm(cat, imaginary=cnfg["imaginary"])
+            target = target_state.targetEquation(state_catalog=cat, imaginary=cnfg["imaginary"])
         lambdaloss = "".join(["1-", target, "/(1+", norm, ")"])
         func, lossstring = th.buildLossString(lambdaloss, variables)
     else:
@@ -209,17 +210,32 @@ def count_rate(graph, target_state, cnfg):
         state_catalog_tensor = np.array(graph._state_catalog_tensor)
         target_normed = np.array(target_normed)
 
-        graph_state = lambda edges: edges[state_catalog_tensor].prod(axis=-1).sum(axis=-1)
-        normed_state = lambda state: state / (1 + np.linalg.norm(state, axis=-1))
-        count_rate = lambda state: abs(state @ target_normed)**2
-        func0 = lambda x: 1 - count_rate(normed_state(graph_state(x)))
-
+        # print(state_catalog_tensor)
+        def graph_state(edges):
+            result = edges[state_catalog_tensor].prod(axis=-1).sum(axis=-1)
+            return result
+
+        def normed_state(state):
+            print(f"[normed_state] Input shape: {state.shape}")
+            norm_val = np.linalg.norm(state, axis=-1)
+            result = state / (1 + norm_val)
+            return result
+
+        def count_rate(state):
+            result = abs(state @ target_normed) ** 2
+            return result
+
+        def func0(x):
+            state = graph_state(x)
+            normalized_state = normed_state(state)
+            return 1 - count_rate(normalized_state)
         #matrix that transforms weights of length len(graph.edges) to weights of length len(graph.complete_graph_edges)
         mat = np.zeros((len(graph.complete_graph_edges), len(graph.edges)))
         for i, edge in enumerate(graph.edges):
             mat[graph.complete_graph_edges.index(edge),i] = 1
 
-        func = lambda x: func0(mat @ x)
+        def func(x):
+            return func0(mat @ x)
 
     print('count rate done', flush=True)
     return func
@@ -230,8 +246,6 @@ def fidelity(graph, target_state, cnfg):
     new_loss = cnfg.get("new_loss", False)
 
     if not new_loss:
-        # set up target equation
-        target = target_state.targetEquation(state_catalog=graph.state_catalog, imaginary=cnfg["imaginary"])
         # get variable names
         variables = th.stringEdges(graph.edges, imaginary=cnfg["imaginary"])
 
@@ -240,6 +254,7 @@ def fidelity(graph, target_state, cnfg):
             # only looking at perfect matchings
             graph.getNorm()
             norm = graph.norm
+            target = target_state.targetEquation(state_catalog=graph.state_catalog, imaginary=cnfg["imaginary"])
         # heralded case, more complicated selection rules
         else:
             if not cnfg["brutal_covers"]:
@@ -250,6 +265,7 @@ def fidelity(graph, target_state, cnfg):
                 edgecovers = brutal_covers(cnfg, graph)
             cat = th.stateCatalog(edgecovers)
             norm = th.writeNorm(cat, imaginary=cnfg["imaginary"])
+            target = target_state.targetEquation(state_catalog=cat, imaginary=cnfg["imaginary"])
         lambdaloss = "".join(["1-", target, "/(0+", norm, ")"])
         func, lossstring = th.buildLossString(lambdaloss, variables)
 

pytheus/main.py

@@ -421,7 +421,7 @@ def setup_for_target(cnfg, state_cat=True):
             if edge[0] == connection[1] and edge[1] == connection[0]:
                 graph.remove(edge)
 
-    print(f'start graph has {len(edge_list)} edges.')
+    print(f'start graph has {len(graph.edges)} edges.')
     return target_state, graph, cnfg
 
 

pytheus/optimizer.py

@@ -13,7 +13,7 @@
 import json
 
 import logging
-
+import time
 log = logging.getLogger(__name__)
 
 
@@ -29,11 +29,16 @@ def __init__(self, start_graph: Graph, saver: saver, ent_dic=None, target_state=
         else:
             self.target = target_state  # object of State class
 
+        self.best_loss = np.inf
         # do preoptimization on complete starting graph, this might already take some time
         self.graph = self.pre_optimize_start_graph(start_graph)
-        self.saver = saver
+
         self.save_hist = safe_history
         self.history = []
+        
+        self.saver = saver
+        #save preoptimized graph to file
+        self.saver.save_graph(self)
 
     def check(self, result: object, lossfunctions: object):
         """
@@ -57,13 +62,9 @@ def check(self, result: object, lossfunctions: object):
             if abs(result.fun) - abs(self.loss_val[0]) > self.config['thresholds'][0]:
                 return False
         else:
-            # uncomment to see where checks fail
-            # print(result.fun, self.config['thresholds'][0])
-            if result.fun > self.config['thresholds'][0]:
-                # if check fails return false
-                return False
             # check if all loss functions are under the corresponding threshold
-            for ii in range(1, len(lossfunctions)):
+            for ii in range(len(lossfunctions)):
+                print(f'loss {ii}: {lossfunctions[ii](result.x):.4f}, threshold: {self.config["thresholds"][ii]:.4f}')
                 if lossfunctions[ii](result.x) > self.config['thresholds'][ii]:
                     # if check fails return false
                     return False
@@ -166,7 +167,8 @@ def get_loss_functions(self, current_graph: Graph):
         for loss in callable_loss:
             try:
                 loss(testinit)
-            except Exception:
+            except Exception as e:
+                print(f'Error in loss function {loss}: {e}')
                 raise RuntimeError('Loss function gives error for a test input, so it is not properly defined. This could be due to configuration parameters given for the optimization. Trying to compute perfect matchings for an odd number of total particles (main+ancilla) will lead to a meaningless loss function (0/0 --> division by zero).')
         return callable_loss
 
@@ -204,94 +206,115 @@ def pre_optimize_start_graph(self, graph) -> Graph:
         preopt_graph
 
         """
-        # losses is a list of callable lossfunctions, e.g. [countrate(x), fidelity(x)], where x is a vector of edge weights
-        # that can be given to scipy.optimize
-        log.info('loading losses')
-        losses = self.get_loss_functions(graph)
-        log.info('losses done')
-        valid = False
-        counter = 0
-        # repeat optimization of complete graph until a good solution is found (which satifies self.check())
-        while not valid:
-            # prepare optimizer
-            initial_values, bounds = self.prepOptimizer(len(graph))
-            # optimization with scipy
-            log.info('begin preopt')
-            best_result = optimize.minimize(losses[0], x0=initial_values,
-                                            bounds=bounds,
-                                            method=self.config['optimizer'],
-                                            options={'ftol': self.config['ftol']})
-            log.info('end preopt')
-            self.loss_val = self.update_losses(best_result, losses)
-            # check if solution is valid
-            valid = self.check(best_result, losses)
-            counter += 1
-            # print a warning if preoptimization is stuck in a loop
-            if counter % 10 == 0:
-                print('10 invalid preoptimization, consider changing parameters.')
-                log.info('10 invalid preoptimization, consider changing parameters.')
-            if counter % 100 == 0:
-                print('100 invalid preoptimization, state cannot be found.')
-                log.info('100 invalid preoptimization, state cannot be found.')
-                raise ValueError('100 invalid preoptimization steps. Conclusion: State cannot be created with provides parameters. Consider adding more ancillas or using less restrictions if possible (e.g. removed_connections).')
-                
-        # if num_pre is set to larger than 1 in config, do num_pre preoptimization and choose the best one.
-        # for optimizations with concrete target state, num_pre = 1 is enough
-        for __ in range(self.config['num_pre'] - 1):
-            initial_values, bounds = self.prepOptimizer(len(graph))
-            result = optimize.minimize(losses[0], x0=initial_values,
-                                       bounds=bounds,
-                                       method=self.config['optimizer'],
-                                       options={'ftol': self.config['ftol']})
-
-            if result.fun < best_result.fun:
-                best_result = result
-        self.loss_val = self.update_losses(best_result, losses)
-        print(f'best result from pre-opt: {abs(best_result.fun)}')
-        log.info(f'best result from pre-opt: {abs(best_result.fun)}')
-
-        for ii, edge in enumerate(graph.edges):
-            graph[edge] = best_result.x[ii]
-        preopt_graph = graph.copy()
-
-        try:
-            bulk_thr = self.config['bulk_thr']
-        except:
-            bulk_thr = 0
-        if bulk_thr > 0:
-            # cut all edges smaller than bulk_thr and optimize again
-            # this can save a lot of time
-            cont = True
-            num_deleted = 0
-            while cont:
-                # delete smallest edges one by one
-                min_edge = preopt_graph.minimum()
-                amplitude = preopt_graph[min_edge]
-                if self.imaginary == 'polar':
-                    amplitude = amplitude[0]
-                if abs(amplitude) < bulk_thr:
-                    preopt_graph.remove(min_edge, update=True)
-                    num_deleted += 1
+        if 'init_graph' in self.config:
+            print('SKIPPING PREOPTIMIZATION, USING INIT_GRAPH FROM CONFIG', flush=True)
+            init_edges = self.config['init_graph']
+            init_edges = {eval(k): v for k, v in init_edges.items()}
+            print(init_edges)
+            print(graph.edges)
+            for edge in list(graph.edges):
+                if edge in init_edges:
+                    graph[edge] = init_edges[edge]
                 else:
-                    cont = False
-            print(f'{num_deleted} edges deleted')
-            log.info(f'{num_deleted} edges deleted')
+                    graph.remove(edge)
+            print(graph)
+            preopt_graph = graph.copy()
+        else:
+            # losses is a list of callable lossfunctions, e.g. [countrate(x), fidelity(x)], where x is a vector of edge weights
+            # that can be given to scipy.optimize
+            log.info('loading losses')
+            losses = self.get_loss_functions(graph)
+            log.info('losses done')
             valid = False
+            counter = 0
+            # repeat optimization of complete graph until a good solution is found (which satifies self.check())
             while not valid:
-                # it is necessary that the truncated graph passes the checks
-                initial_values, bounds = self.prepOptimizer(len(preopt_graph))
-                losses = self.get_loss_functions(preopt_graph)
-                trunc_result = optimize.minimize(losses[0], x0=initial_values,
-                                                 bounds=bounds,
-                                                 method=self.config['optimizer'],
-                                                 options={'ftol': self.config['ftol']})
-                self.loss_val = self.update_losses(trunc_result, losses)
-                print(f'result after truncation: {abs(trunc_result.fun)}')
-                log.info(f'result after truncation: {abs(trunc_result.fun)}')
-                valid = self.check(trunc_result, losses)
-        
-            for ii, edge in enumerate(preopt_graph.edges):
-                preopt_graph[edge] = trunc_result.x[ii]
+                # prepare optimizer
+                initial_values, bounds = self.prepOptimizer(len(graph))
+                # optimization with scipy
+                log.info('begin preopt')
+                self.tt = time.time()
+                def callback(xk):
+                    print(f'preopt step with loss {losses[0](xk):.3f}', flush=True)
+                    print(f'time for last step: {time.time()-self.tt:.1f}s', flush=True)
+                    self.tt = time.time()
+                best_result = optimize.minimize(losses[0], x0=initial_values,
+                                                bounds=bounds,
+                                                method=self.config['optimizer'],
+                                                options={'ftol': self.config['ftol']},
+                                                callback=callback)
+                log.info('end preopt')
+                self.loss_val = self.update_losses(best_result, losses)
+                # check if solution is valid
+                valid = self.check(best_result, losses)
+                counter += 1
+                # print a warning if preoptimization is stuck in a loop
+                if counter % 10 == 0:
+                    print('10 invalid preoptimization, consider changing parameters.')
+                    log.info('10 invalid preoptimization, consider changing parameters.')
+                if counter % 100 == 0:
+                    print('100 invalid preoptimization, state cannot be found.')
+                    log.info('100 invalid preoptimization, state cannot be found.')
+                    raise ValueError('100 invalid preoptimization steps. Conclusion: State cannot be created with provides parameters. Consider adding more ancillas or using less restrictions if possible (e.g. removed_connections).')
+                    
+            # if num_pre is set to larger than 1 in config, do num_pre preoptimization and choose the best one.
+            # for optimizations with concrete target state, num_pre = 1 is enough
+            for __ in range(self.config['num_pre'] - 1):
+                initial_values, bounds = self.prepOptimizer(len(graph))
+                result = optimize.minimize(losses[0], x0=initial_values,
+                                        bounds=bounds,
+                                        method=self.config['optimizer'],
+                                        options={'ftol': self.config['ftol']})
+
+                if result.fun < best_result.fun:
+                    best_result = result
+            self.loss_val = self.update_losses(best_result, losses)
+            print(f'best result from pre-opt: {abs(best_result.fun)}')
+            log.info(f'best result from pre-opt: {abs(best_result.fun)}')
+            self.best_loss = best_result.fun
+
+            for ii, edge in enumerate(graph.edges):
+                graph[edge] = best_result.x[ii]
+            preopt_graph = graph.copy()
+
+            try:
+                bulk_thr = self.config['bulk_thr']
+            except:
+                bulk_thr = 0
+            if bulk_thr > 0:
+                # cut all edges smaller than bulk_thr and optimize again
+                # this can save a lot of time
+                cont = True
+                num_deleted = 0
+                while cont:
+                    # delete smallest edges one by one
+                    min_edge = preopt_graph.minimum()
+                    amplitude = preopt_graph[min_edge]
+                    if self.imaginary == 'polar':
+                        amplitude = amplitude[0]
+                    if abs(amplitude) < bulk_thr:
+                        preopt_graph.remove(min_edge, update=True)
+                        num_deleted += 1
+                    else:
+                        cont = False
+                print(f'{num_deleted} edges deleted')
+                log.info(f'{num_deleted} edges deleted')
+                valid = False
+                while not valid:
+                    # it is necessary that the truncated graph passes the checks
+                    initial_values, bounds = self.prepOptimizer(len(preopt_graph))
+                    losses = self.get_loss_functions(preopt_graph)
+                    trunc_result = optimize.minimize(losses[0], x0=initial_values,
+                                                    bounds=bounds,
+                                                    method=self.config['optimizer'],
+                                                    options={'ftol': self.config['ftol']})
+                    self.loss_val = self.update_losses(trunc_result, losses)
+                    print(f'result after truncation: {abs(trunc_result.fun)}')
+                    log.info(f'result after truncation: {abs(trunc_result.fun)}')
+                    valid = self.check(trunc_result, losses)
+            
+                for ii, edge in enumerate(preopt_graph.edges):
+                    preopt_graph[edge] = trunc_result.x[ii]
 
         return preopt_graph
 
@@ -419,6 +442,9 @@ def optimize_one_edge(self, num_edge: int,
                 # to a file even if topological optimization can be continued
                 if all(np.array(abs(self.graph)) > 0.95):
                     self.saver.save_graph(self)
+                elif result.fun < self.best_loss:
+                    self.best_loss = result.fun
+                    self.saver.save_graph(self)
                 if self.save_hist:
                     self.history.append([str(self.graph),self.loss_val])
                 # return updated result graph