QROAMClean musical score diagram can fail to render

[fdmalone]

Taking the prep_thc bloq from resource_estimation.ipynb results in:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/projects/qualtran/qualtran/drawing/musical_score.py:226, in _update_assign_from_vals(regs, binst, vals, soq_assign, seq_x, topo_gen, manager)
    225 try:
--> 226     arr: Union[RegPosition, NDArray[RegPosition]] = vals[reg.name]
    227 except KeyError:

KeyError: 'junk_target1_'

During handling of the above exception, another exception occurred:

IndexError                                Traceback (most recent call last)
Cell In[16], line 2
      1 from qualtran.drawing import get_musical_score_data, draw_musical_score
----> 2 msd = get_musical_score_data(prep_thc.decompose_bloq())
      3 fig, ax = draw_musical_score(msd)
      4 fig.set_size_inches(12, 8)

File ~/projects/qualtran/qualtran/drawing/musical_score.py:582, in get_musical_score_data(bloq, manager)
    569 """Get the musical score data for a (composite) bloq.
    570 
    571 This will first walk through the compute graph to assign each soquet
   (...)
    578         to control the allocation of horizontal (register) lines.
    579 """
    580 cbloq = bloq.as_composite_bloq()
--> 582 _, soq_assign, manager = _cbloq_musical_score(
    583     cbloq.signature, binst_graph=cbloq._binst_graph, manager=manager
    584 )
    585 msd = MusicalScoreData(
    586     max_x=max(v.seq_x for v in soq_assign.values()),
    587     max_y=max(v.y for v in soq_assign.values()),
   (...)
    590     hlines=sorted(manager.hlines),
    591 )
    593 for hline in manager.hlines:

File ~/projects/qualtran/qualtran/drawing/musical_score.py:332, in _cbloq_musical_score(signature, binst_graph, manager)
    330             continue
    331         pred_cxns, succ_cxns = _binst_to_cxns(binst, binst_graph=binst_graph)
--> 332         _binst_assign_line(
    333             binst, pred_cxns, soq_assign, seq_x=seq_x, topo_gen=topo_gen, manager=manager
    334         )
    335         seq_x += 1
    337 # Track bloq-to-dangle name changes

File ~/projects/qualtran/qualtran/drawing/musical_score.py:284, in _binst_assign_line(binst, pred_cxns, soq_assign, seq_x, topo_gen, manager)
    278 partial_out_vals = {
    279     reg.name: in_vals[reg.name] for reg in bloq.signature if reg.side is Side.THRU
    280 }
    282 # The following will use `partial_out_vals` to re-use existing THRU lines; otherwise
    283 # the following will allocate new lines.
--> 284 _update_assign_from_vals(
    285     bloq.signature.rights(),
    286     binst,
    287     partial_out_vals,
    288     soq_assign,
    289     seq_x=seq_x,
    290     topo_gen=topo_gen,
    291     manager=manager,
    292 )
    294 # Free any purely-left registers.
    295 for reg in bloq.signature:

File ~/projects/qualtran/qualtran/drawing/musical_score.py:228, in _update_assign_from_vals(regs, binst, vals, soq_assign, seq_x, topo_gen, manager)
    226     arr: Union[RegPosition, NDArray[RegPosition]] = vals[reg.name]
    227 except KeyError:
--> 228     arr = manager.new(
    229         binst=cast(BloqInstance, binst), reg=reg, seq_x=seq_x, topo_gen=topo_gen
    230     )
    232 if reg.shape:
    233     arr = np.asarray(arr)

File ~/projects/qualtran/qualtran/drawing/musical_score.py:156, in LineManager.new(self, binst, reg, seq_x, topo_gen)
    154 arg = np.zeros(reg.shape, dtype=object)
    155 for idx in reg.all_idxs():
--> 156     y = self.new_y(binst, reg, idx)
    157     self.hlines.add(HLine(y=y, seq_x_start=seq_x))
    158     arg[idx] = RegPosition(y=y, seq_x=seq_x, topo_gen=topo_gen)

File ~/projects/qualtran/qualtran/drawing/musical_score.py:98, in LineManager.new_y(self, binst, reg, idx)
     96 def new_y(self, binst: BloqInstance, reg: Register, idx=None):
     97     """Allocate a new y position (i.e. a new qubit or register)."""
---> 98     return heapq.heappop(self.available)

IndexError: index out of range

My guess is there are so many junk registers that the default max_n_lines=100 isn't enough.

The issue is to see if there's a way of drawing the diagram without the extra junk registers which can be quite large (in this case the block size is 64 so each junk register of of shape (63,) to begin with.

[mpharrigan]

Yes, by design the musical score diagram will fail if it requires more than that number of lines. The real solution to this issue would be to make sure you're drawing what you want to. Other things to track/consider

• Better error message when this happens
• We could increase the default

[fdmalone]

So QROAMClean now generates an extra junk register of shape (K-1), for each target register for the QROM, I wonder if we can skip drawing these in the diagrams as 1) K could be large and 2) they don't really add any information to the diagram. Maybe the awkwardness of switching the block size momentarily for the purpose of drawing in the notebook outweighs the effort and specializing things.