From 89f13f2a872af29b6f4f1165b7ef401ec1136e85 Mon Sep 17 00:00:00 2001
From: Johnnie Gray <johnniemcgray@gmail.com>
Date: Fri, 20 Sep 2024 15:19:18 -0700
Subject: [PATCH] fixes for arbgeom tebd ordering

---
 quimb/tensor/__init__.py            |  3 +--
 quimb/tensor/circuit.py             |  2 +-
 quimb/tensor/tensor_arbgeom_tebd.py | 12 ++++++++----
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/quimb/tensor/__init__.py b/quimb/tensor/__init__.py
index 0ba8f222..da946f8c 100644
--- a/quimb/tensor/__init__.py
+++ b/quimb/tensor/__init__.py
@@ -1,5 +1,4 @@
-"""Tensor and tensor network functionality.
-"""
+"""Tensor and tensor network functionality."""
 
 from .circuit import (
     Circuit,
diff --git a/quimb/tensor/circuit.py b/quimb/tensor/circuit.py
index 39157e35..113fe104 100644
--- a/quimb/tensor/circuit.py
+++ b/quimb/tensor/circuit.py
@@ -1373,7 +1373,7 @@ def sample_bitstring_from_prob_ndarray(p):
         >>> sample_bitstring_from_prob_ndarray(p)
         '01011'
     """
-    b = np.random.choice(np.arange(p.size), p=p.flat)
+    b = np.random.choice(p.size, p=p.flat)
     return f"{b:0>{p.ndim}b}"
 
 
diff --git a/quimb/tensor/tensor_arbgeom_tebd.py b/quimb/tensor/tensor_arbgeom_tebd.py
index d7b0e530..f6092d50 100644
--- a/quimb/tensor/tensor_arbgeom_tebd.py
+++ b/quimb/tensor/tensor_arbgeom_tebd.py
@@ -4,7 +4,7 @@
 import itertools
 import collections
 
-from autoray import do, to_numpy, dag
+from autoray import do, to_numpy
 
 from ..core import eye, kron, qarray
 from ..utils import ensure_dict
@@ -54,11 +54,16 @@ def edge_coloring(
         coloring.setdefault(color, []).append(edge)
 
     if group:
-        return tuple(tuple(coloring[color]) for color in sorted(coloring))
+        return tuple(
+            tuple(tuple(tuple(sorted(edge)) for edge in coloring[color]))
+            for color in sorted(coloring)
+        )
     else:
         # flatten sorted groups
         return tuple(
-            edge for color in sorted(coloring) for edge in coloring[color]
+            tuple(sorted(edge))
+            for color in sorted(coloring)
+            for edge in coloring[color]
         )
 
 
@@ -657,7 +662,6 @@ def _check_energy(self):
         if self.keep_best and en < self.best["energy"]:
             self.best["energy"] = en
             self.best["state"] = self.state
-            self.best["gauges"] = self.gauges.copy()
             self.best["it"] = self._n
 
         return self.energies[-1]