improving basis states plotting

EQuS · May 2, 2024 · 292cb4f · 292cb4f
1 parent 4adf928
commit 292cb4f
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 31 deletions.
diff --git a/bosonic_jax/codes/base.py b/bosonic_jax/codes/base.py
@@ -264,15 +264,21 @@ def plot_code_states(self, qp_type: str = jqt.WIGNER, **kwargs):
             pos = (i // 3, i % 3)
             ax = axs[pos]
             w_plt = self._plot_single(state, ax=ax, qp_type=qp_type, **kwargs)
-            ax.set_title(label)
+            ax.set_title(f"|{label}" + r"$\rangle$")
+            ax.set_xlabel(r"Re[$\alpha$]")
+            ax.set_ylabel(r"Im[$\alpha$]")
+
         fig.suptitle(self.name)
         fig.tight_layout()
         fig.subplots_adjust(right=0.8, hspace=0.2, wspace=0.2)
-        cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7])
+        fig.align_xlabels(axs)
+        fig.align_ylabels(axs)
+        cbar_ax = fig.add_axes([0.85 + 0.2, 0.15, 0.05, 0.7])
 
         ticks = np.linspace(vmin, vmax, 5)
         fig.colorbar(w_plt, cax=cbar_ax, ticks=ticks)
 
         cbar_ax.set_title(cbar_title, pad=20)
+        fig.tight_layout()
         plt.show()
 
diff --git a/tutorials/binomial/1-codes.ipynb b/tutorials/binomial/1-codes.ipynb
diff --git a/tutorials/cat/1-codes.ipynb b/tutorials/cat/1-codes.ipynb
diff --git a/tutorials/gkp/1-codes.ipynb b/tutorials/gkp/1-codes.ipynb