Add colorbar to the multiple xyz overlay plot

xyzoverlay.py

@@ -14,6 +14,7 @@
 from itertools import cycle                                     #for color cycling
 import io                                                       #IO for (easy) saving multi xyz
 import re                                                       #regex (get atom x y z from multi xyz)
+import matplotlib.colors as mcolors                             #to add the colorbar
 
 #var for check if xyz is multi xyz
 is_trj = 0
@@ -151,6 +152,11 @@ def align_xyz(vec1,vec2,coord):
     type= str,
     help='use Matplotlib colormap in view mode')
 
+#add color bar
+parser.add_argument('-cb', '--cbar',
+    type= bool, default=False,
+    help='in case of using Matplotlib colormap, add a colorbar as well')
+
 #add +x% to radius (for view mode)
 parser.add_argument('-r','--radius',
     default=8,
@@ -179,6 +185,8 @@ def align_xyz(vec1,vec2,coord):
     default=0, action='store_true',
     help='save superimposed / aligned data as (multi) xyz or xyz trajectory file')
 
+
+
 #parse arguments
 args = parser.parse_args()
 
@@ -401,6 +409,7 @@ def align_xyz(vec1,vec2,coord):
         try:
             new_colors = [plt.get_cmap(args.cmap)(1. * i/num_of_xyz) for i in range(num_of_xyz)]
             color_cycle = cycler('color',new_colors)
+            num_xyz = len(xyz_df_list)
         except ValueError:
             print('Warning! Not a valid Matplotlib colormap. Applying default colormap instead.')
 
@@ -432,7 +441,7 @@ def align_xyz(vec1,vec2,coord):
     #generate list of data frames of bonding atoms 
     for bond_mat in bond_mat_list:
         bond_mat.values[np.triu_indices_from(bond_mat, k=1)] = np.nan
-        pd.set_option("future.no_silent_downcasting", True)
+        #pd.set_option("future.no_silent_downcasting", True)
         bond_mat = bond_mat.replace(0, np.nan)
         bond_mat = bond_mat.unstack().dropna()
         bond_mat = bond_mat.reset_index(level=1)
@@ -451,14 +460,28 @@ def align_xyz(vec1,vec2,coord):
 
         #scatter (atom) plot
         ax.scatter(*xyz_df[['x','y','z']].to_numpy().T,s=np.log10(atom_scaler/num_atom_xyz),alpha=alpha_atoms,**style)
-
         #bonds
         for bonds in atom1_2_coord:
             ax.plot(*bonds.T,linewidth=np.log10(bond_scaler/num_atom_xyz),alpha=alpha_bonds,**style)
+
     #no axes
     ax.set_axis_off()
-    #tight layout 
+
+    #add color bar to account for the number of cartesian coordinates
+    if args.cbar:
+      norm = mcolors.Normalize(vmin=0, vmax=num_xyz)
+      cmap = plt.get_cmap(args.cmap)
+      # Create a ScalarMappable object to generate the colorbar
+      sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+      sm.set_array([])  # N
+      cbar = fig.colorbar(sm, ax=ax, shrink=0.6)
+      cbar.set_label('Number of Cartesian Coordinates', fontsize=14, labelpad=20)
+      cbar.ax.tick_params(labelsize=12)
+
+    #tight layout
     fig.tight_layout()
+
+    # Generate data
     #adjust 3d drawing behavior, otherwise molecules are not correctly displayes
     set_axes_equal(ax)
     #show the plot
@@ -622,7 +645,6 @@ def align_xyz(vec1,vec2,coord):
         #bonds
         for bonds in atom1_2_coord:
             ax.plot(*bonds.T,linewidth=np.log10(bond_scaler/num_atom_xyz),alpha=alpha_bonds,color='gray')
-
     #no axes
     ax.set_axis_off()
     #tight layout