Update tracks.py

src/fasttrackpy/tracks.py

@@ -8,6 +8,7 @@
                                            get_big_df
 from sklearn.experimental import enable_iterative_imputer
 from sklearn.impute import IterativeImputer
+import matplotlib.pyplot as mp
 
 import polars as pl
 
@@ -211,8 +212,8 @@ def to_df(self, output = "formants"):
             return self._param_df
 
         raise ValueError("output must be either 'formants' or 'param'")
-
-    def spectrogram(self, maximum_frequency=3500, tracks = True, dynamic_range=60):
+    
+    def spectrogram(self, formants = 3, maximum_frequency=3500, tracks = True, dynamic_range=60):
 
         spctgrm = self.sound.to_spectrogram(maximum_frequency=maximum_frequency)
         Time, Hz = spctgrm.x_grid(), spctgrm.y_grid()
@@ -230,10 +231,10 @@ def spectrogram(self, maximum_frequency=3500, tracks = True, dynamic_range=60):
             mp.scatter (point_times, self.formants[0], c="red")
             mp.scatter (point_times, self.formants[1], c="blue")
             mp.scatter (point_times, self.formants[2], c="green")
-
+            if formants == 4:
+                mp.scatter (point_times, self.formants[3], c="darkturquoise")    
 
-
-
+
 
 
 class CandidateTracks(Track):
@@ -371,10 +372,9 @@ def to_df(self, which = "winner", output = "formants"):
 
         if output == "param":
             return self._param_df
-
-
-
-    def spectrograms(self, maximum_frequency = 3500, dynamic_range=60):
+
+
+    def spectrograms(self, formants = 3, maximum_frequency = 3500, dynamic_range=60):
 
         spectrogram = self.sound.to_spectrogram(maximum_frequency=maximum_frequency,time_step=0.005)
         Time, Hz = spectrogram.x_grid(), spectrogram.y_grid()
@@ -388,6 +388,9 @@ def spectrograms(self, maximum_frequency = 3500, dynamic_range=60):
         gs = fig.add_gridspec(4,5, hspace=0.05, wspace=0.05)
         axs = gs.subplots(sharex='col', sharey='row')
 
+        #gs = fig.add_gridspec(3, hspace=0)
+        #axs = gs.subplots(sharex=True, sharey=True)
+
         for i in range (4):
             for j in range(5):
                 axs[i, j].pcolormesh(Time, Hz, db, vmin=min_shown, cmap='magma')
@@ -396,6 +399,8 @@ def spectrograms(self, maximum_frequency = 3500, dynamic_range=60):
                 axs[i, j].scatter (point_times, self.candidates[analysis].formants[0], c="red", s = 5)
                 axs[i, j].scatter (point_times, self.candidates[analysis].formants[1], c="blue", s = 5)
                 axs[i, j].scatter (point_times, self.candidates[analysis].formants[2], c="green", s = 5)    
+                if formants == 4:
+                    axs[i, j].scatter (point_times, self.candidates[analysis].formants[3], c="darkturquoise", s = 5)    
 
         for ax in fig.get_axes():
             ax.label_outer()