@@ -12,6 +12,7 @@ def plot_waterfall(
1212 exclude_labels : Optional [list ] = None ,
1313 suppress_deviations : bool = True ,
1414 suppress_deviations_threshold : float = 40.0 ,
15+ use_interpolated : bool = True ,
1516) -> None :
1617 """
1718 Recreate the waterfall similar to Fig 1d. from "Neural signal propagation atlas of Caenorhabditis elegans".
@@ -45,10 +46,12 @@ def plot_waterfall(
4546 Delta = 5.0
4647
4748 # Fetch data from NWB source
48- green_signal = segmentation_nwbfile .processing ["ophys" ]["GreenSignals" ].microscopy_response_series ["GreenSignal" ]
49+ signal_name = "InterpolatedGreenSignal" if use_interpolated else "BaseGreenSignal"
50+ green_signal = segmentation_nwbfile .processing ["ophys" ]["GreenSignals" ].microscopy_response_series [signal_name ]
4951 time = green_signal .timestamps [:]
5052
5153 neuropal_rois = segmentation_nwbfile .processing ["ophys" ]["NeuroPALSegmentations" ].microscopy_plane_segmentations ["NeuroPALPlaneSegmentation" ]
54+ neuropal_ids = neuropal_rois .id [:]
5255 green_rois = segmentation_nwbfile .processing ["ophys" ]["PumpProbeGreenSegmentations" ]
5356
5457 coregistered_neuropal_id_to_green_ids = {
@@ -59,6 +62,7 @@ def plot_waterfall(
5962 )
6063 if coregistered_neuropal_id != ""
6164 }
65+ coregistered_green_ids_to_neuropal_ids = {int (value ): key for key , value in coregistered_neuropal_id_to_green_ids .items ()}
6266
6367 neuropal_labels = neuropal_rois .labels .data [:]
6468 alphabetized_valid_neuropal_labels_with_ids = [
@@ -106,9 +110,9 @@ def plot_waterfall(
106110 matplotlib .pyplot .rc ("xtick" , labelsize = 8 )
107111
108112 plotted_neuropal_ids = []
109- colors = []
113+ neuropal_label_to_colors = dict ()
110114 baseline = []
111- for plot_index , (label , neuropal_id ) in enumerate (alphabetized_valid_neuropal_labels_with_ids ):
115+ for plot_index , (neuropal_label , neuropal_id ) in enumerate (alphabetized_valid_neuropal_labels_with_ids ):
112116 green_id = coregistered_neuropal_id_to_green_ids [neuropal_id ]
113117
114118 roi_response = green_signal .data [:, green_id ]
@@ -144,12 +148,13 @@ def plot_waterfall(
144148 # In particular, the "M5", "AWAR", "RIMR", and "I3" traces all had massive deviations
145149 # from the figure in the paper that were not smoothed by any of the previous steps
146150 # (spike removal, photobleaching correction, or Savitzky-Golay filtering).
147- smoothed_deviations = max (smoothed [13 :]) - min (smoothed [13 :])
151+ smoothed_deviations = max (smoothed [savgol_filter_size :]) - min (smoothed [savgol_filter_size :])
148152 if suppress_deviations is True and smoothed_deviations > suppress_deviations_threshold :
149153 line , = ax .plot ([], [], lw = 0.8 ) # Still plotting an empty line to increment coloration to match
150154
151155 color = line .get_color ()
152- colors .append (color )
156+ neuropal_label_to_colors .update ({neuropal_label : color })
157+
153158 continue
154159
155160 # There are a few other lines in the plot that don't precisely match the figure from the paper
@@ -166,19 +171,55 @@ def plot_waterfall(
166171 line , = ax .plot (time [13 :], plot [13 :], lw = 0.8 )
167172
168173 color = line .get_color ()
169- colors .append (color )
174+ neuropal_label_to_colors .update ({neuropal_label : color })
175+
176+ ax .annotate (text = neuropal_label , xy = (- 150 - 120 * (plot_index % 2 ), plot_index * Delta ), c = color , fontsize = 8 )
177+
178+ # Optogenetic stimulation table
179+ stimulation_times = segmentation_nwbfile .intervals ["OptogeneticStimulusTable" ]["start_time" ][:]
180+
181+ stimulation_labels = []
182+ for green_id in segmentation_nwbfile .intervals ["OptogeneticStimulusTable" ]["target_pumpprobe_id" ][:]:
183+ if numpy .isnan (green_id ):
184+ stimulation_labels .append ("" )
185+ continue
186+
187+ neuropal_id = coregistered_green_ids_to_neuropal_ids .get (int (green_id ), None )
188+ if neuropal_id is None or neuropal_id == "" :
189+ stimulation_labels .append ("" )
190+ continue
191+
192+ stimulation_label = neuropal_labels [neuropal_id ]
193+ if (
194+ stimulation_label in excluded_labels
195+ or stimulation_label not in neuropal_label_to_colors
196+ ):
197+ stimulation_labels .append ("" )
198+ continue
199+
200+ stimulation_labels .append (stimulation_label )
201+
202+ for stimulation_time , stimulation_label in zip (stimulation_times , stimulation_labels ):
203+ ax .axvline (x = stimulation_time , c = "k" , alpha = 0.5 , lw = 1 , ymax = 0.98 )
170204
171- ax .annotate (label , (- 150 - 120 * (plot_index % 2 ), plot_index * Delta ), c = color , fontsize = 8 )
205+ if stimulation_label != "" :
206+ ax .annotate (
207+ text = stimulation_label ,
208+ xy = (stimulation_time , (plot_index + 6 ) * Delta ),
209+ c = neuropal_label_to_colors [stimulation_label ],
210+ fontsize = 8 ,
211+ rotation = 90 ,
212+ )
172213
173214 ax .set_xlim (- 270 , time [- 1 ])
174215 ax .set_ylim (- Delta , Delta * (plot_index + 6 ))
175216 ax .set_yticks ([])
176217 ax .set_xlabel ("Time (s)" , fontsize = 10 )
177218 ax .set_ylabel ("Responding neuron" , fontsize = 10 )
178219
179- ax .spines [' right' ].set_visible (False )
180- ax .spines [' top' ].set_visible (False )
181- ax .spines [' left' ].set_visible (False )
220+ ax .spines [" right" ].set_visible (False )
221+ ax .spines [" top" ].set_visible (False )
222+ ax .spines [" left" ].set_visible (False )
182223
183224 return None
184225
0 commit comments