Wave function compact visualization (#376)

* wfn compact visualization * minor edit * Update src/tequila/wavefunction/qubit_wavefunction.py Co-authored-by: Oliver Hüttenhofer <[email protected]> * Update src/tequila/wavefunction/qubit_wavefunction.py Co-authored-by: Oliver Hüttenhofer <[email protected]> * Update test_chemistry.py for h2 hcb is a bit slower (overhead due to changes in orbopt function, not dramatic) --------- Co-authored-by: Oliver Hüttenhofer <[email protected]>
tequilahub · Dec 15, 2024 · aab6c33 · aab6c33
1 parent a12e33f
commit aab6c33
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/tequila/wavefunction/qubit_wavefunction.py b/src/tequila/wavefunction/qubit_wavefunction.py
@@ -373,7 +373,10 @@ def __repr__(self):
             index = index.integer
             if self.numbering == BitNumbering.LSB:
                 index = reverse_int_bits(index, self._n_qubits)
-            result += f"{coeff} |{index:0{self._n_qubits}b}> "
+            if np.isclose(coeff.imag, 0.0):
+                result += f"{coeff.real:+2.4f} |{index:0{self._n_qubits}b}> "
+            else:
+                result += f"({coeff.real:+2.4f} + {coeff.imag:+2.4f}i) |{index:0{self._n_qubits}b}> "
         # If the wavefunction contains no states
         if not result:
             result = "empty wavefunction"

diff --git a/tests/test_chemistry.py b/tests/test_chemistry.py
@@ -704,7 +704,7 @@ def test_hcb_rdms(geometry, optimize):
 
 
 @pytest.mark.skipif(condition=not HAS_PYSCF, reason="you don't have pyscf")
-@pytest.mark.parametrize("geometry", ["H 0.0 0.0 0.0\nH 0.0 0.0 4.5", "Li 0.0 0.0 0.0\nH 0.0 0.0 3.0", "Be 0.0 0.0 0.0\nH 0.0 0.0 3.0\nH 0.0 0.0 -3.0"])
+@pytest.mark.parametrize("geometry", ["Li 0.0 0.0 0.0\nH 0.0 0.0 3.0", "Be 0.0 0.0 0.0\nH 0.0 0.0 3.0\nH 0.0 0.0 -3.0"])
 def test_orbital_optimization_hcb(geometry):
 
     mol = tq.Molecule(geometry=geometry, basis_set="sto-3g")