Change .flatten() to .ravel() to avoid copies

[stavros11]

While working on a qibo callback I realized that qibojit creates copies, instead of using in-place updates, when simulating density matrices. This is true on CPU and GPU.

Here is a script to reproduce this issue:

import numpy as np
from qibo import gates, K
from qibo.models import Circuit

state = np.zeros((4, 4), dtype=np.complex128)
state[0, 0] = 1

state = K.cast(state)
c = Circuit(2, density_matrix=True)
c.add(gates.H(0))
c.add(gates.H(1))
final_state = c(initial_state=state)

print(state)
print(final_state.state())

will print

[Qibo 0.1.8.dev0|INFO|2022-03-10 14:20:02]: Using qibojit (numba) backend on /CPU:0
[[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]
 [0.+0.j 0.+0.j 0.+0.j 0.+0.j]]
[[0.25+0.j 0.25+0.j 0.25+0.j 0.25+0.j]
 [0.25+0.j 0.25+0.j 0.25+0.j 0.25+0.j]
 [0.25+0.j 0.25+0.j 0.25+0.j 0.25+0.j]
 [0.25+0.j 0.25+0.j 0.25+0.j 0.25+0.j]]

which means that the matrix was copied somewhere, as the initial state was not modified. The equivalent for state vectors

import numpy as np
from qibo import gates, K
from qibo.models import Circuit

state = np.zeros(4, dtype=np.complex128)
state[0] = 1

state = K.cast(state)
c = Circuit(2)
c.add(gates.H(0))
c.add(gates.H(1))
final_state = c(initial_state=state)

print(state)
print(final_state.state())

will print

[Qibo 0.1.8.dev0|INFO|2022-03-10 14:21:06]: Using qibojit (numba) backend on /CPU:0
[0.5+0.j 0.5+0.j 0.5+0.j 0.5+0.j]
[0.5+0.j 0.5+0.j 0.5+0.j 0.5+0.j]

so it modifies the initial state too.

Apparently using state.flatten() creates a copy of the state. Here I am replacing it with .ravel() which returns a flattened version without copying. I tested with cupy and .flatten() indeed duplicates the used GPU memory, while .ravel() does not. The above example is also fixed with this branch. I have not checked how the overal performance is affected for density matrix simulation, but it would be good to do so before merging.

[codecov]

Codecov Report

Merging #74 (c02e38e) into main (38fb97c) will not change coverage.
The diff coverage is 100.00%.

@@            Coverage Diff            @@
##              main       #74   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            9         9           
  Lines         1080      1080           
=========================================
  Hits          1080      1080           

Flag	Coverage Δ
unittests	100.00% <100.00%> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
src/qibojit/custom_operators/platforms.py	100.00% <ø> (ø)
src/qibojit/custom_operators/__init__.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

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

[andrea-pasquale]

Thanks for spotting this @stavros11.
Looks good to me.
Basically with the current main we have memory issues similar to those one from a month ago on GPU, but we didn't notice it because we were running just state vector simulation. Interesting.