Skip to content

Commit

Permalink
expose qtn.edge_coloring
Browse files Browse the repository at this point in the history
  • Loading branch information
@jcmgray
jcmgray committed Sep 20, 2024
1 parent fd78337 commit de46163
Show file tree
Hide file tree
Showing 3 changed files with 98 additions and 82 deletions.
9 changes: 9 additions & 0 deletions docs/changelog.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,15 @@

Release notes for `quimb`.

(whats-new-1-8-5)=
## v1.8.5 (unreleased)

**Enhancements:**

- expose [`qtn.edge_coloring`](quimb.tensor.tensor_arbgeom_tebd.edge_coloring)
as top level function and allow layers to be returned grouped.


(whats-new-1-8-4)=
## v1.8.4 (2024-07-20)

Expand Down
2 changes: 2 additions & 0 deletions quimb/tensor/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@
LocalHamGen,
SimpleUpdateGen,
TEBDGen,
edge_coloring,
)
from .tensor_builder import (
MPS_COPY,
Expand Down Expand Up @@ -242,6 +243,7 @@
"DMRG1",
"DMRG2",
"DMRGX",
"edge_coloring",
"edges_1d_chain",
"edges_2d_hexagonal",
"edges_2d_kagome",
Expand Down
169 changes: 87 additions & 82 deletions quimb/tensor/tensor_arbgeom_tebd.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
"""Tools for performing TEBD like algorithms on arbitrary lattices.
"""
"""Tools for performing TEBD like algorithms on arbitrary lattices."""

import random
import itertools
Expand All @@ -14,6 +13,55 @@
from .drawing import get_colors, get_positions


def edge_coloring(
edges,
strategy="smallest_last",
interchange=True,
group=True,
):
"""Generate an edge coloring for the graph given by ``edges``, using
``networkx.coloring.greedy_color``.
Parameters
----------
edges : sequence[tuple[hashable, hashable]]
The edges of the graph.
strategy : str or callable, optional
The strategy to use for coloring the edges. Can be:
- 'largest_first'
- 'smallest_last'
- 'random_sequential'
...
interchange : bool, optional
Whether to use the interchange heuristic. Usually generates better
colorings but can be slower.
group : bool, optional
Whether to group the edges by color or return a flat list.
"""
import networkx as nx

# find vertex coloring of line graph
G = nx.Graph(tuple(edges))
edge_colors = nx.coloring.greedy_color(
nx.line_graph(G), strategy, interchange=interchange
)

# group the edges by color
coloring = {}
for edge, color in edge_colors.items():
coloring.setdefault(color, []).append(edge)

if group:
return tuple(tuple(coloring[color]) for color in sorted(coloring))
else:
# flatten sorted groups
return tuple(
edge for color in sorted(coloring) for edge in coloring[color]
)


class LocalHamGen:
"""Representation of a local hamiltonian defined on a general graph. This
combines all two site and one site terms into a single interaction per
Expand Down Expand Up @@ -104,7 +152,6 @@ def __init__(self, H2, H1=None):

# now absorb the single site terms evenly into the two site terms
for site, H in H1s.items():

# get interacting terms which cover the site
pairs = self._sites_to_covering_terms[site]
num_pairs = len(pairs)
Expand All @@ -124,8 +171,7 @@ def __init__(self, H2, H1=None):

@property
def nsites(self):
"""The number of sites in the system.
"""
"""The number of sites in the system."""
return len(self.sites)

def items(self):
Expand Down Expand Up @@ -171,7 +217,7 @@ def _op_id_cached(self, x):
key = id(x)
if key not in cache:
xn = to_numpy(x)
d = int(xn.size ** 0.5)
d = int(xn.size**0.5)
Id = eye(d, dtype=xn.dtype)
XI = do("array", kron(xn, Id), like=x)
cache[key] = XI
Expand All @@ -182,7 +228,7 @@ def _id_op_cached(self, x):
key = id(x)
if key not in cache:
xn = to_numpy(x)
d = int(xn.size ** 0.5)
d = int(xn.size**0.5)
Id = eye(d, dtype=xn.dtype)
IX = do("array", kron(Id, xn), like=x)
cache[key] = IX
Expand All @@ -198,8 +244,7 @@ def _expm_cached(self, x, y):
return cache[key]

def get_gate(self, where):
"""Get the local term for pair ``where``, cached.
"""
"""Get the local term for pair ``where``, cached."""
return self.terms[tuple(sorted(where))]

def get_gate_expm(self, where, x):
Expand All @@ -209,33 +254,10 @@ def get_gate_expm(self, where, x):
return self._expm_cached(self.get_gate(where), x)

def apply_to_arrays(self, fn):
"""Apply the function ``fn`` to all the arrays representing terms.
"""
"""Apply the function ``fn`` to all the arrays representing terms."""
for k, x in self.terms.items():
self.terms[k] = fn(x)

def _nx_color_ordering(self, strategy="smallest_first", interchange=True):
"""Generate a term ordering based on a coloring on the line graph.
"""
import networkx as nx

G = nx.Graph(tuple(self.terms))

coloring = list(
nx.coloring.greedy_color(
nx.line_graph(G), strategy, interchange=interchange
).items()
)

# sort into color groups
coloring.sort(key=lambda coo_color: coo_color[1])

return [
# networkx doesn't preserve node order of edge spec
tuple(sorted(coo)) for
coo, _ in coloring
]

def get_auto_ordering(self, order="sort", **kwargs):
"""Get an ordering of the terms to use with TEBD, for example. The
default is to sort the coordinates then greedily group them into
Expand Down Expand Up @@ -278,7 +300,7 @@ def get_auto_ordering(self, order="sort", **kwargs):
random.shuffle(pairs)
return pairs
else:
return self._nx_color_ordering(order, **kwargs)
return edge_coloring(self.terms, order, group=False, **kwargs)

pairs = {x: None for x in pairs}

Expand Down Expand Up @@ -334,7 +356,7 @@ def draw(
import matplotlib.pyplot as plt

if figsize is None:
L = self.nsites ** 0.5 + 1
L = self.nsites**0.5 + 1
figsize = (L, L)

ax_supplied = ax is not None
Expand Down Expand Up @@ -472,9 +494,9 @@ def __init__(
if D is None:
D = self._psi.max_bond()
self.gate_opts = ensure_dict(gate_opts)
self.gate_opts['max_bond'] = D
self.gate_opts.setdefault('cutoff', 0.0)
self.gate_opts.setdefault('contract', 'reduce-split')
self.gate_opts["max_bond"] = D
self.gate_opts.setdefault("cutoff", 0.0)
self.gate_opts.setdefault("contract", "reduce-split")

# parse energy computation options
self.compute_energy_opts = ensure_dict(compute_energy_opts)
Expand All @@ -487,7 +509,7 @@ def __init__(
if ordering is None:

def dynamic_random():
return self.ham.get_auto_ordering('random_sequential')
return self.ham.get_auto_ordering("random_sequential")

self.ordering = dynamic_random
elif isinstance(ordering, str):
Expand All @@ -506,7 +528,7 @@ def dynamic_random():
self.energies = []

self.keep_best = bool(keep_best)
self.best = dict(energy=float('inf'), state=None, it=None)
self.best = dict(energy=float("inf"), state=None, it=None)

def sweep(self, tau):
r"""Perform a full sweep of gates at every pair.
Expand Down Expand Up @@ -555,17 +577,17 @@ def evolve(self, steps, tau=None, progbar=None):
if progbar is None:
progbar = self.progbar

pbar = Progbar(total=steps, disable=self.progbar is not True)
pbar = Progbar(total=steps, disable=progbar is not True)

try:
for i in range(steps):
# anything required by both energy and sweep
self.presweep(i)

# possibly compute the energy
should_compute_energy = (
bool(self.compute_energy_every) and
(i % self.compute_energy_every == 0))
should_compute_energy = bool(self.compute_energy_every) and (
i % self.compute_energy_every == 0
)
if should_compute_energy:
self._check_energy()
self._update_progbar(pbar)
Expand All @@ -592,8 +614,7 @@ def evolve(self, steps, tau=None, progbar=None):

@property
def state(self):
"""Return a copy of the current state.
"""
"""Return a copy of the current state."""
return self.get_state()

@state.setter
Expand All @@ -602,25 +623,21 @@ def state(self, psi):

@property
def n(self):
"""The number of sweeps performed.
"""
"""The number of sweeps performed."""
return self._n

@property
def D(self):
"""The maximum bond dimension.
"""
return self.gate_opts['max_bond']
"""The maximum bond dimension."""
return self.gate_opts["max_bond"]

@D.setter
def D(self, value):
"""The maximum bond dimension.
"""
self.gate_opts['max_bond'] = round(value)
"""The maximum bond dimension."""
self.gate_opts["max_bond"] = round(value)

def _check_energy(self):
"""Logic for maybe computing the energy if needed.
"""
"""Logic for maybe computing the energy if needed."""
if self.its and (self._n == self.its[-1]):
# only compute if haven't already
return self.energies[-1]
Expand All @@ -637,17 +654,17 @@ def _check_energy(self):
self.taus.append(float(self.last_tau))
self.its.append(self._n)

if self.keep_best and en < self.best['energy']:
self.best['energy'] = en
self.best['state'] = self.state
self.best['it'] = self._n
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]

@property
def energy(self):
"""Return the energy of current state, computing it only if necessary.
"""
"""Return the energy of current state, computing it only if necessary."""
return self._check_energy()

# ------- abstract methods that subclasses might want to override ------- #
Expand Down Expand Up @@ -681,17 +698,12 @@ def compute_energy(self):
override this with a custom method to compute the energy.
"""
return self._psi.compute_local_expectation_cluster(
terms=self.ham.terms,
**self.compute_energy_opts
terms=self.ham.terms, **self.compute_energy_opts
)

@default_to_neutral_style
def plot(
self,
zoom="auto",
xscale="symlog",
xscale_linthresh=20,
hlines=()
self, zoom="auto", xscale="symlog", xscale_linthresh=20, hlines=()
):
import numpy as np
import matplotlib.pyplot as plt
Expand All @@ -702,13 +714,13 @@ def plot(
xs = np.array(self.its)
ys = np.array(self.energies)

ax.plot(xs, ys, '.-')
ax.plot(xs, ys, ".-")
ax.set_xlabel("Iteration")
ax.set_ylabel("Energy")

if xscale == "symlog":
ax.set_xscale(xscale, linthresh=xscale_linthresh)
ax.axvline(xscale_linthresh, color=(.5, .5, .5), ls="-", lw=0.5)
ax.axvline(xscale_linthresh, color=(0.5, 0.5, 0.5), ls="-", lw=0.5)
else:
ax.set_xscale(xscale)

Expand All @@ -730,22 +742,17 @@ def plot(

def __repr__(self):
s = "<{}(n={}, tau={}, D={})>"
return s.format(
self.__class__.__name__, self.n, self.tau, self.D)
return s.format(self.__class__.__name__, self.n, self.tau, self.D)


class SimpleUpdateGen(TEBDGen):
"""Simple update for arbitrary geometry hamiltonians.
"""
"""Simple update for arbitrary geometry hamiltonians."""

def gate(self, U, where):
self._psi.gate_simple_(
U, where, gauges=self.gauges, **self.gate_opts
)
self._psi.gate_simple_(U, where, gauges=self.gauges, **self.gate_opts)

def normalize(self):
"""Normalize the state and simple gauges.
"""
"""Normalize the state and simple gauges."""
self._psi.normalize_simple(self.gauges)

def compute_energy(self):
Expand Down Expand Up @@ -773,5 +780,3 @@ def set_state(self, psi, gauges=None):
self._psi.gauge_all_simple_(gauges=self.gauges)
else:
self.gauges = dict(gauges)


0 comments on commit de46163

Please sign in to comment.