Autoshape and fock improvements

[ziofil]

Context:
A recommended Fock shape of the components of a circuit can often be determined by the nature of the components and their connectivity. Also, for certain components like States we can determine a good choice of shape by stopping e.g. at 99.9% of the norm.

Description of the Change:

• CircuitComponents now have the property .manual_shape which can be set manually. It starts as [None,...] and it can be modified like e.g. my_state.manual_shape[1]=9 and now it will be [None, 9, None, ...].
• CircuitComponents now have the method .auto_shape() which computes and returns a tuple with a guesstimate of the recommended shape to the best of our knowledge. It will respect the values in .manual_shape that have already been set.
• Circuit now has the method propagate_shapes() which exploits the context of each component to produce a more compact shape than auto_shape() would give in isolation. This gives the Circuit the power to contract an arrangement of components in the Fock representation in a close to optimal way (excluding throwing SVDs around), when combined with the optimal contraction path (upcoming).
• Faster autocutoff method (now autoshape_numba). It doesn't any longer travel back and forth between phase space and Bargmann for each mode, and it's the same as Robbe's diagonal method, but specialized for this task, so it's faster.
• General "fock-related" functionality has been refactored/improved
  • removed the converter module
  • localized the fock conversion into the CircuitComponent so that wire information is available
  • components keep the original bargmann data when they are converted to fock
  • can no longer matmul across representations (this is good)

[codecov]

Codecov Report

Attention: Patch coverage is 94.77912% with 13 lines in your changes missing coverage. Please review.

Project coverage is 88.03%. Comparing base (8946ea5) to head (ba98716).

Additional details and impacted files

@@             Coverage Diff             @@
##           develop     #397      +/-   ##
===========================================
+ Coverage    87.95%   88.03%   +0.07%     
===========================================
  Files           81       80       -1     
  Lines         6310     6443     +133     
===========================================
+ Hits          5550     5672     +122     
- Misses         760      771      +11     

Files	Coverage Δ
mrmustard/lab/abstract/state.py	92.66% <ø> (ø)
mrmustard/lab/circuit_drawer.py	92.15% <ø> (ø)
mrmustard/lab_dev/circuit_components_utils.py	100.00% <100.00%> (ø)
...mustard/lab_dev/transformations/transformations.py	100.00% <100.00%> (ø)
mrmustard/lab_dev/wires.py	98.27% <100.00%> (-0.82%)	⬇️
mrmustard/math/backend_numpy.py	100.00% <ø> (ø)
mrmustard/math/backend_tensorflow.py	100.00% <ø> (ø)
mrmustard/math/lattice/strategies/vanilla.py	100.00% <100.00%> (ø)
mrmustard/math/parameters.py	95.45% <ø> (ø)
mrmustard/physics/ansatze.py	96.71% <100.00%> (+0.03%)	⬆️
... and 8 more

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8946ea5...ba98716. Read the comment docs.

[elib20]

Awesome work!

I found a few spots with unexpected behaviour/edge case bugs, but nothing that can't be easily squashed.

[aplund]

Again, I may be "doing it wrong(TM)", but I was expecting this to change the shape as the squeezing increases. Am I missing a step somewhere:

import numpy as np
import mrmustard as mm
from mrmustard.lab_dev import *

for r in np.linspace(0.0, 3.0, 31):
    cc = Number([0], 0) >> Sgate([0], r)
    cc.auto_shape()
    f = cc.fock()
    tail = np.sqrt(np.sum(np.abs(f[..., -10:])**2))
    print(f"r={r:0.2f}, shape={cc.fock_shape}, tail={tail}")

Here's my output:

r=0.00, shape=[50], tail=0.0
r=0.10, shape=[50], tail=3.320813549783774e-21
r=0.20, shape=[50], tail=2.877796850159443e-15
r=0.30, shape=[50], tail=7.005423842389163e-12
r=0.40, shape=[50], tail=1.4442895820803e-09
r=0.50, shape=[50], tail=7.391845079536454e-08
r=0.60, shape=[50], tail=1.527182338092643e-06
r=0.70, shape=[50], tail=1.6608089924466752e-05
r=0.80, shape=[50], tail=0.0001119152559711487
r=0.90, shape=[50], tail=0.0005210655981549683
r=1.00, shape=[50], tail=0.0018108286711155058
r=1.10, shape=[50], tail=0.004974852404081372
r=1.20, shape=[50], tail=0.011292920255741735
r=1.30, shape=[50], tail=0.02193479266151498
r=1.40, shape=[50], tail=0.03749458889213709
r=1.50, shape=[50], tail=0.057708136029986395
r=1.60, shape=[50], tail=0.08148204007111474
r=1.70, shape=[50], tail=0.10718002773885009
r=1.80, shape=[50], tail=0.1330086283344037
r=1.90, shape=[50], tail=0.15735305275690456
r=2.00, shape=[50], tail=0.178984595152066
r=2.10, shape=[50], tail=0.1971322790005048
r=2.20, shape=[50], tail=0.2114532602089405
r=2.30, shape=[50], tail=0.22194764197883732
r=2.40, shape=[50], tail=0.22885591269080502
r=2.50, shape=[50], tail=0.23256342153828174
r=2.60, shape=[50], tail=0.23352358559759126
r=2.70, shape=[50], tail=0.23220268304523672
r=2.80, shape=[50], tail=0.22904418030721152
r=2.90, shape=[50], tail=0.2244485284290604
r=3.00, shape=[50], tail=0.21876408762339897

Should the shape not increase as r increases?

[ziofil]

Again, I may be "doing it wrong(TM)", but I was expecting this to change the shape as the squeezing increases. Am I missing a step somewhere:
[...]
Should the shape not increase as r increases?

@aplund I think you wanted this:

for r in np.linspace(0.0, 3.0, 31):
    cc = Vacuum([0]) >> Sgate([0], r)
    cc.auto_shape()
    print(f"r={r:0.2f}, shape={cc.fock_shape}")  # works

The issue with your code was that Number has a Fock representation, so the Sgate needs to go to Fock itself and the best it can do is to match the shape (of 1) on the vacuum side and on the output side it has no way to determine how far to go and so it defaults to the value in the settings. On the other hand, Vacuum has a Bargmann representation so we can play all the tricks in our book.

[aplund]

The issue with your code was that Number has a Fock representation, so the Sgate needs to go to Fock itself and the best it can do is to match the shape (of 1) on the vacuum side and on the output side it has no way to determine how far to go and so it defaults to the value in the settings. On the other hand, Vacuum has a Bargmann representation so we can play all the tricks in our book.

So there's an implicit difference between Number([0], 0) and Vacuum([0]) that makes them not interchangeable? If so, this should probably be bolded and italicised somewhere, as I think most people would consider them the same.

When I run the code you suggested, I get the output:

r=0.00, shape=[None]
r=0.10, shape=[None]
r=0.20, shape=[None]
r=0.30, shape=[None]
r=0.40, shape=[None]
r=0.50, shape=[None]
r=0.60, shape=[None]
r=0.70, shape=[None]
r=0.80, shape=[None]
r=0.90, shape=[None]
r=1.00, shape=[None]
r=1.10, shape=[None]
r=1.20, shape=[None]
r=1.30, shape=[None]
r=1.40, shape=[None]
r=1.50, shape=[None]
r=1.60, shape=[None]
r=1.70, shape=[None]
r=1.80, shape=[None]
r=1.90, shape=[None]
r=2.00, shape=[None]
r=2.10, shape=[None]
r=2.20, shape=[None]
r=2.30, shape=[None]
r=2.40, shape=[None]
r=2.50, shape=[None]
r=2.60, shape=[None]
r=2.70, shape=[None]
r=2.80, shape=[None]
r=2.90, shape=[None]
r=3.00, shape=[None]

Is this the expected behaviour?

[ziofil]

So there's an implicit difference between Number([0], 0) and Vacuum([0]) that makes them not interchangeable? If so, this should probably be bolded and italicised somewhere, as I think most people would consider them the same.

Sure, we can make it clearer in the docstrings.

When I run the code you suggested, I get the output:

r=0.00, shape=[None]
...

Is this the expected behaviour?

Yes, now it is. I have removed the lines where fock_shape is updated by auto_shape, which takes about 100us per mode to run, so it's okay if we run it multiple times.

[elib20]

My only remaining comment is to please include NotImplementedError's for autoshape on batched Bargmann DMs and Kets, since if the method was used as currently written, it would produce incorrect results.

Since that is an easy change to make, and since my only other comments are just catching a couple typos in docstrings, I'm providing an approval so that you don't need to wait on a follow-up review from me after you make those changes.

Not sure if @timmysilv or @aplund wanted to have a second look either?

[timmysilv]

looks good! just a few little things