Problem computing tags in reverse light cone

steve-jeffrey commented 2 years ago

What happened?

The tags in the reverse lightcone are computed when calling quimb.tensor.circuit.Circuit.local_expectation().

The calculation does not include gates that have more than 2 indices (e.g. raw gates defined by the user), consequently:

the lightcone is incomplete, leading to incorrect results, and/or;
code failure if there is only one gate in the circuit with an index matching the where index, and the gate has more than 2 indices, and the index matching the where index is not one of the last two indices (the code fails because the light cone is empty).

I have provided draft code which appears to fix the problem (it it is indeed a bug).

What did you expect to happen?

The lightcone should include all gates with indices matching the where index.

Minimal Complete Verifiable Example

import numpy as np
import quimb.tensor as qtn
import quimb as qu

# Construct a custom gate
gate = np.array(
    [[1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j],
    [0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
    [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 1.+0.j]]
).reshape([2] * 6)

circ = qtn.Circuit(N=4)
circ.apply_gate_raw(gate, [3, 2, 0], contract=False, tags='custom_gate')
circ.apply_gate('H', 1)
ev = circ.local_expectation(qu.pauli('Z'), 3)
print(f"Local expectation value: {ev}")

Relevant log output

Traceback (most recent call last):
  File "my_code.py", line 19, in <module>
    ev = circ.local_expectation(qu.pauli('Z'), 3)
  File "site-packages/quimb/tensor/circuit.py", line 1670, in local_expectation
    rho = self.get_rdm_lightcone_simplified(where=where, **fs_opts)
  File "site-packages/quimb/tensor/circuit.py", line 1324, in get_rdm_lightcone_simplified
    rho_lc.full_simplify_(seq=seq, atol=atol, output_inds=output_inds,
  File "site-packages/quimb/tensor/tensor_core.py", line 7926, in full_simplify
    tn.rank_simplify_(output_inds=ix_o, cache=cache,
  File "site-packages/quimb/tensor/tensor_core.py", line 7039, in rank_simplify
    queue = oset(sorted(count, key=rank_weight))
  File "site-packages/quimb/tensor/tensor_core.py", line 7036, in rank_weight
    return (tn.ind_size(ind), -sum(tn.tensor_map[tid].ndim
  File "site-packages/quimb/tensor/tensor_core.py", line 6689, in ind_size
    tid = next(iter(self.ind_map[ind]))
KeyError: 'b3'

Anything else we need to know?

I think the problem is due to this code block:

if isinstance(gate[-2], numbers.Integral):
     regs = set(gate[-2:])
else:
     regs = set(gate[-1:])

which appears to only take the last one or two indices in the gate array, which is insufficient when the gate has more than 2 indices.

If the above block of code is replaced with this:

# Extract the regs from the gate array. The maximum number of regs that could be in the gate array
# will be limited by: 
# - the length of the gate, less 1 (the first entry is the name or ID of the gate),
# - the dimension N, less 1 (if the indices are zero based, they will range from 0 to N-1)
max_regs = min(self.N-1, len(gate)-1)
gate[-max_regs:]
regs = set()
for reg in gate[-max_regs:]:
    # Check the reg is less than N 
    # - raw gates are assigned an integer ID, where the value is a large integer (e.g. 139776274826224).
    #   If the value is integral and below N, we assume it is a reg and not the ID of a raw gate
    if isinstance(reg, numbers.Integral) and reg < self.N:
        regs.add(reg)

the tags appear to be computed correctly (at least on the example I provided above).

Environment

Quimb commit cf1898f (from 17 May 2022) Python 3.8.12 Fedora 35

jcmgray commented 2 years ago

Hi @steve-jeffrey, Thanks very much for the detailed issue. Does the code snippet fix you posted account for parameters supplied to gates?

Keeping each gate as a tuple of (label, *params, *regs), is pretty unsatisfactory I have to say, this might be a good moment to think about using a simple Gate dataclass, with gate.label, gate.params, gate.regs, gate.size etc. I will try and look into making this change.

steve-jeffrey commented 2 years ago

Hi Johnnie,

No, it does not account for parameters.

The code snippet attempts to work backwards through the array, so it only picks up regs. I took this approach as: (i) I didn't know what the array contained (I thought the first element was the gate name/ID and the last elements were the regs, but I didn't know what else it could contain) and (ii): this approach seemed similar to the approach you used (I think you were checking only the last element in the array, or the last two elements). However, I think my code snippet will fail if the array contains parameters which are integers, as the code will assume the integer parameters are regs. Hopefully you can implement something better :-)

jcmgray commented 2 years ago

OK, I've added a much more robust Gate class that should fix the issue (see commit above). Let me know if things aren't working.

It should be much easier to add custom gates and multi-qubit gates as well if you have particular things in mind.

steve-jeffrey commented 2 years ago

It seems to be working fine. Thanks Johnnie.

jcmgray / quimb