experiment_process_data

  1"""
  2This module defines the ExperimentProcessData class, which serves as the central
  3data management hub for the Photologic-Experiment-Rig application.
  4
  5It aggregates references to other data-holding classes (`StimuliData`, `EventData`,
  6`ArduinoData`) and manages core experimental parameters like current trial,
  7state time durations, and the overall experiment schedule DataFrame. It provides
  8methods for generating the experimental schedule, updating parameters from the GUI,
  9calculating runtime, and saving collected data.
 10"""
 11
 12import logging
 13from tkinter import filedialog
 14import numpy as np
 15import pandas as pd
 16from typing import Tuple, List
 17import datetime
 18from pathlib import Path
 19
 20from models.stimuli_data import StimuliData
 21from models.event_data import EventData
 22from models.arduino_data import ArduinoData
 23from views.gui_common import GUIUtils
 24
 25logger = logging.getLogger(__name__)
 26
 27
 28class ExperimentProcessData:
 29    """
 30    Central data management class for the behavioral experiment.
 31
 32    This class acts as the primary container and manager for all data related
 33    to an ongoing experiment. It holds references to specialized data classes
 34    (`EventData`, `StimuliData`, `ArduinoData`) and maintains experiment-wide
 35    state variables, timing intervals, parameter entries (num_stimuli for example), and the generated
 36    program schedule (`program_schedule_df`). It performs the generation
 37    of the trial schedule and state intervals based on user inputs.
 38
 39    Attributes
 40    ----------
 41    - **`start_time`** (*float*): Timestamp marking the absolute start of the program run.
 42    - **`trial_start_time`** (*float*): Timestamp marking the start of the current trial.
 43    - **`state_start_time`** (*float*): Timestamp marking the entry into the current FSM state.
 44    - **`current_trial_number`** (*int*): The 1-indexed number of the trial currently in progress or about to start.
 45    - **`event_data`** (*EventData*): Instance managing the DataFrame of recorded lick/motor events.
 46    - **`stimuli_data`** (*StimuliData*): Instance managing stimuli names and related information.
 47    - **`arduino_data`** (*ArduinoData*): Instance managing Arduino communication data (durations, schedule indices) and processing incoming Arduino messages.
 48    - **`ITI_intervals_final`** (*npt.NDArray[np.int64] | None*): Numpy array holding the calculated Inter-Trial Interval duration (ms) for each trial.
 49    None until schedule generated.
 50    - **`TTC_intervals_final`** (*npt.NDArray[np.int64] | None*): Numpy array holding the calculated Time-To-Contact duration (ms) for each trial.
 51    None until schedule generated.
 52    - **`sample_intervals_final`** (*npt.NDArray[np.int64] | None*): Numpy array holding the calculated Sample period duration (ms) for each trial.
 53    None until schedule generated.
 54    - **`TTC_LICK_THRESHOLD`** (*int*): The number of licks required during the TTC state to trigger an early transition to the SAMPLE state.
 55    - **`interval_vars`** (*dict[str, int]*): Dictionary storing base and random variation values (ms) for timing intervals (ITI, TTC, Sample), sourced from GUI entries.
 56    - **`exp_var_entries`** (*dict[str, int]*): Dictionary storing core experiment parameters (Num Trial Blocks, Num Stimuli), sourced from GUI entries.
 57    `Num Trials` is calculated based on these other values.
 58    - **`program_schedule_df`** (*pd.DataFrame*): Pandas DataFrame holding the generated trial-by-trial schedule, including stimuli presentation, calculated intervals,
 59    and placeholders for results. Initialized empty.
 60
 61    Methods
 62    -------
 63    - `update_model`(...)
 64        Updates internal parameter dictionaries (`interval_vars`, `exp_var_entries`) based on changes from GUI input fields (tkinter vars).
 65    - `get_default_value`(...)
 66        Retrieves the default or current value for a given parameter name from internal dictionaries, used to populate GUI fields initially.
 67    - `generate_schedule`()
 68        Performs the creation of the entire experimental schedule, including intervals and stimuli pairings. Returns status.
 69    - `create_random_intervals`()
 70        Calculates the final ITI, TTC, and Sample intervals for each trial based on base values and randomization ranges.
 71    - `calculate_max_runtime`()
 72        Estimates the maximum possible runtime based on the sum of all generated intervals.
 73    - `create_trial_blocks`()
 74        Determines the unique pairs of stimuli presented within a single block based on the number of stimuli.
 75    - `generate_pairs`(...)
 76        Generates the pseudo-randomized sequence of stimuli pairs across all trials and blocks.
 77    - `build_frame`(...)
 78        Constructs the `program_schedule_df` DataFrame using the generated stimuli sequences and intervals.
 79    - `save_all_data`()
 80        Initiates the process of saving the schedule and event log DataFrames to Excel files.
 81    - `get_paired_index`(...)
 82        Static method to determine the 0-indexed valve number on the opposite side corresponding to a given valve index on side one.
 83    - `save_df_to_xlsx`(...)
 84        Static method to handle saving a pandas DataFrame to an .xlsx file using a file dialog.
 85    - `convert_seconds_to_minutes_seconds`(...)
 86        Static method to convert a total number of seconds into minutes and remaining seconds.
 87    """
 88
 89    def __init__(self):
 90        """
 91        Initializes the ExperimentProcessData central data hub.
 92
 93        Sets initial timestamps to 0.0, starts `current_trial_number` at 1.
 94        Instantiates `EventData`, `StimuliData`, and `ArduinoData` (passing self reference).
 95        Initializes interval arrays (`ITI_intervals_final`, etc.) to None.
 96        Sets the `TTC_LICK_THRESHOLD`.
 97        Defines default values for `interval_vars` and `exp_var_entries`.
 98        Initializes `program_schedule_df` as an empty DataFrame.
 99        """
100
101        self.start_time: float = 0.0
102        self.trial_start_time: float = 0.0
103        self.state_start_time: float = 0.0
104
105        # this number is centralized here so that all state classes can access and update it easily without
106        # passing it through to each state every time a state change occurs
107        self.current_trial_number: int = 1
108
109        self.event_data = EventData()
110        self.stimuli_data = StimuliData()
111        self.arduino_data = ArduinoData(self)
112
113        self.ITI_intervals_final = None
114        self.TTC_intervals_final = None
115        self.sample_intervals_final = None
116
117        # this constant should be used in arduino_data to say how many licks moves to sample
118        self.TTC_LICK_THRESHOLD: int = 3
119
120        self.interval_vars: dict[str, int] = {
121            "ITI_var": 30000,
122            "TTC_var": 20000,
123            "sample_var": 15000,
124            "ITI_random_entry": 0,
125            "TTC_random_entry": 5000,
126            "sample_random_entry": 5000,
127        }
128        self.exp_var_entries: dict[str, int] = {
129            "Num Trial Blocks": 10,
130            "Num Stimuli": 4,
131            "Num Trials": 0,
132        }
133
134        # include an initial definition of program_schedule_df here as a blank df to avoid type errors and to
135        # make it clear that this is a class attriute
136        self.program_schedule_df = pd.DataFrame()
137
138    def update_model(self, variable_name: str, value: int | None) -> None:
139        """
140        Updates internal parameter dictionaries from GUI inputs.
141
142        Checks if the `variable_name` corresponds to a key in `interval_vars` or
143        `exp_var_entries` and updates the dictionary value if the provided `value`
144        is not None.
145
146        Parameters
147        ----------
148        - **variable_name** (*str*): The name identifier of the variable being updated (should match a dictionary key).
149        - **value** (*int | None*): The new integer value for the variable, or None if the input was invalid/empty.
150        """
151
152        # if the variable name for the tkinter entry item that we are updating is
153        # in the exp_data interval variables dictionary, update that entry
154        # with the value in the tkinter variable
155        if value is not None:
156            if variable_name in self.interval_vars.keys():
157                self.interval_vars[variable_name] = value
158            elif variable_name in self.exp_var_entries.keys():
159                # update other var types here
160                self.exp_var_entries[variable_name] = value
161
162    def get_default_value(self, variable_name: str) -> int:
163        """
164        Retrieves the current value associated with a parameter name.
165
166        Searches `interval_vars` and `exp_var_entries` for the `variable_name`.
167        Used primarily to populate GUI fields with their initial/default values.
168
169        Parameters
170        ----------
171        - **variable_name** (*str*): The name identifier of the parameter whose value is requested.
172
173        Returns
174        -------
175        - *int*: The integer value associated with `variable_name` in the dictionaries, or -1 if the name is not found.
176        """
177        if variable_name in self.interval_vars.keys():
178            return self.interval_vars[variable_name]
179        elif variable_name in self.exp_var_entries.keys():
180            # update other var types here
181            return self.exp_var_entries[variable_name]
182
183        # couldn't find the value, return -1 as error code
184        return -1
185
186    def generate_schedule(self) -> bool:
187        """
188        Generates the complete pseudo-randomized experimental schedule.
189
190        Calculates the total number of trials based on stimuli count and block count.
191        Validates the number of stimuli. Calls helper methods:
192        - `create_random_intervals`() to determine ITI, TTC, Sample times per trial.
193        - `create_trial_blocks`() to define unique stimuli pairs that must occur per block.
194        - `generate_pairs`() to create the trial-by-trial stimuli sequence.
195        - `build_frame`() to assemble the final `program_schedule_df`.
196
197        Returns
198        -------
199        - *bool*: `True` if schedule generation was successful, `False` if an error occurred (e.g., invalid number of stimuli).
200        """
201        try:
202            # set number of trials based on number of stimuli, and number of trial blocks set by user
203            num_stimuli = self.exp_var_entries["Num Stimuli"]
204
205            if num_stimuli > 8 or num_stimuli < 2 or num_stimuli % 2 != 0:
206                GUIUtils.display_error(
207                    "NUMBER OF STIMULI EXCEEDS CURRENT MAXIMUM",
208                    "Program is currently configured for a minumim of 2 and maximum of 8 TOTAL vavles. You must have an even number of valves. If more are desired, program configuration must be modified.",
209                )
210                return False
211
212            num_trial_blocks = self.exp_var_entries["Num Trial Blocks"]
213
214            self.exp_var_entries["Num Trials"] = (num_stimuli // 2) * num_trial_blocks
215
216            self.create_random_intervals()
217
218            pairs = self.create_trial_blocks()
219
220            # stimuli_1 & stimuli_2 are lists that hold the stimuli to be introduced for each trial on their respective side
221            stimuli_side_one, stimuli_side_two = self.generate_pairs(pairs)
222
223            # args needed -> stimuli_1, stimuli_2
224            # these are lists of stimuli for each trial, for each side respectivelyc:w
225            self.build_frame(
226                stimuli_side_one,
227                stimuli_side_two,
228            )
229            return True
230
231        except Exception as e:
232            logger.error(f"Error generating program schedule {e}.")
233            return False
234
235    def create_random_intervals(self) -> None:
236        """
237        Calculates randomized ITI, TTC, and Sample intervals for all trials.
238
239        For each interval type (ITI, TTC, Sample):
240        - Retrieves the base duration and random variation range from `interval_vars`.
241        - Creates a base array repeating the base duration for the total number of trials.
242        - If random variation is specified, generates an array of random integers within the +/- range.
243        - Adds the random variation array to the base array.
244        - Stores the resulting final interval array in the corresponding class attribute (`ITI_intervals_final`, etc.).
245
246        Raises
247        ------
248        - Propagates NumPy errors (e.g., during random number generation or array addition).
249        - *KeyError*: If expected keys are missing from `interval_vars`.
250        """
251        try:
252            num_trials = self.exp_var_entries["Num Trials"]
253
254            final_intervals = [None, None, None]
255
256            interval_keys = ["ITI_var", "TTC_var", "sample_var"]
257
258            random_interval_keys = [
259                "ITI_random_entry",
260                "TTC_random_entry",
261                "sample_random_entry",
262            ]
263
264            # for each type of interval (iti, ttc, sample), generate final interval times
265            for i in range(len(interval_keys)):
266                base_val = self.interval_vars[interval_keys[i]]
267                # this makes an array that is n=num_trials length of base_val
268                base_intervals = np.repeat(base_val, num_trials)
269
270                # get plus minus value for this interval type
271                random_var = self.interval_vars[random_interval_keys[i]]
272
273                if random_var == 0:
274                    # if there is no value to add/subtract randomly, continue on
275                    final_intervals[i] = base_intervals
276
277                else:
278                    # otherwise make num_trials amount of random integers and add them to the final_intervals
279                    random_interval_arr = np.random.randint(
280                        -random_var, random_var, num_trials
281                    )
282
283                    final_intervals[i] = np.add(base_intervals, random_interval_arr)
284
285            # assign results for respective types to the model for storage and later use
286            self.ITI_intervals_final = final_intervals[0]
287            self.TTC_intervals_final = final_intervals[1]
288            self.sample_intervals_final = final_intervals[2]
289
290            logger.info("Random intervals created.")
291        except Exception as e:
292            logger.error(f"Error creating random intervals: {e}")
293            raise
294
295    def calculate_max_runtime(self) -> tuple[int, int]:
296        """
297        Estimates the maximum possible experiment runtime in minutes and seconds.
298
299        Sums all generated interval durations (ITI, TTC, Sample) across all trials.
300        Converts the total milliseconds to seconds, then uses `convert_seconds_to_minutes_seconds`.
301
302        Returns
303        -------
304        - *tuple[int, int]*: A tuple containing (estimated_max_minutes, estimated_max_seconds).
305        """
306        max_time = (
307            sum(self.ITI_intervals_final)
308            + sum(self.TTC_intervals_final)
309            + sum(self.sample_intervals_final)
310        )
311        minutes, seconds = self.convert_seconds_to_minutes_seconds(max_time / 1000)
312        return (minutes, seconds)
313
314    def create_trial_blocks(self) -> List[tuple[str, str]]:
315        """
316        Determines the unique stimuli pairings within a single trial block.
317
318        Calculates the number of pairs per block (`block_sz = num_stimuli / 2`).
319        Iterates from `i = 0` to `block_sz - 1`. For each `i`, finds the
320        corresponding paired valve index using `get_paired_index`. Appends the
321        tuple of (stimulus_name_at_i, stimulus_name_at_paired_index) to a list.
322
323        Returns
324        -------
325        - *List[tuple]*: A list of tuples, where each tuple represents a unique stimuli pairing (e.g., `[('Odor A', 'Odor B'), ('Odor C', 'Odor D')]`).
326
327        Raises
328        ------
329        - Propagates errors from `get_paired_index`.
330        - *KeyError*: If expected keys are missing from `exp_var_entries`.
331        - *IndexError*: If calculated indices are out of bounds for `stimuli_data.stimuli_vars`.
332        """
333        try:
334            # creating pairs list which stores all possible pairs
335            pairs = []
336            num_stimuli = self.exp_var_entries["Num Stimuli"]
337            block_sz = num_stimuli // 2
338
339            stimuli_names = list(self.stimuli_data.stimuli_vars.values())
340
341            for i in range(block_sz):
342                pair_index = self.get_paired_index(i, num_stimuli)
343                pairs.append((stimuli_names[i], stimuli_names[pair_index]))
344
345            return pairs
346        except Exception as e:
347            logger.error(f"Error creating trial blocks: {e}")
348            raise
349
350    def generate_pairs(self, pairs: list[tuple[str, str]]) -> Tuple[list, list]:
351        """
352        Generates the full, pseudo-randomized sequence of stimuli pairs for all trials.
353
354        Repeats the following for the number of trial blocks specified:
355        - Takes the list of unique `pairs` generated by `create_trial_blocks` that must be included in each block.
356        - Shuffles this list randomly (`np.random.shuffle`).
357        - Extends a master list `pseudo_random_lineup` (adds new list to that list) with the shuffled block.
358        After creating the full sequence, separates it into two lists: one for the
359        stimulus presented on side one each trial, and one for side two.
360
361        Parameters
362        ----------
363        - **pairs** (*List[tuple[str, str]]*): The list of unique stimuli pairings defining one block, generated by `create_trial_blocks`.
364
365        Returns
366        -------
367        - *Tuple[list, list]*: A tuple containing two lists:
368            - `stimulus_1`: List of stimuli names for Port 1 for each trial.
369            - `stimulus_2`: List of stimuli names for Port 2 for each trial.
370
371        Raises
372        ------
373        - Propagates errors from `np.random.shuffle`.
374        """
375        try:
376            pseudo_random_lineup: List[tuple] = []
377
378            stimulus_1: List[str] = []
379            stimulus_2: List[str] = []
380
381            for _ in range(self.exp_var_entries["Num Trial Blocks"]):
382                pairs_copy = [tuple(pair) for pair in pairs]
383                np.random.shuffle(pairs_copy)
384                pseudo_random_lineup.extend(pairs_copy)
385
386            for pair in pseudo_random_lineup:
387                stimulus_1.append(pair[0])
388                stimulus_2.append(pair[1])
389
390            return stimulus_1, stimulus_2
391        except Exception as e:
392            logger.error(f"Error generating pairs: {e}")
393            raise
394
395    def build_frame(
396        self,
397        stimuli_1: list[str],
398        stimuli_2: list[str],
399    ) -> None:
400        """
401        Constructs the main `program_schedule_df` pandas DataFrame.
402
403        Uses the generated stimuli sequences (`stimuli_1`, `stimuli_2`) and the
404        calculated interval arrays (`ITI_intervals_final`, etc.) to build the
405        DataFrame. Includes columns for trial block number, trial number, stimuli
406        per port, calculated intervals, and placeholders (NaN) for results columns
407        like lick counts and actual TTC duration.
408
409        Parameters
410        ----------
411        - **stimuli_1** (*list*): List of stimuli names for Port 1, one per trial.
412        - **stimuli_2** (*list*): List of stimuli names for Port 2, one per trial.
413
414        Raises
415        ------
416        - Propagates pandas errors during DataFrame creation.
417        - *ValueError*: If the lengths of input lists/arrays do not match the expected number of trials.
418        """
419        num_stimuli = self.exp_var_entries["Num Stimuli"]
420        num_trials = self.exp_var_entries["Num Trials"]
421        num_trial_blocks = self.exp_var_entries["Num Trial Blocks"]
422
423        # each block must contain each PAIR -> num pairs = num_stim / 2
424        block_size = int(num_stimuli / 2)
425        try:
426            data = {
427                "Trial Block": np.repeat(
428                    range(1, num_trial_blocks + 1),
429                    block_size,
430                ),
431                "Trial Number": np.repeat(range(1, num_trials + 1), 1),
432                "Port 1": stimuli_1,
433                "Port 2": stimuli_2,
434                "Port 1 Licks": np.full(num_trials, np.nan),
435                "Port 2 Licks": np.full(num_trials, np.nan),
436                "ITI": self.ITI_intervals_final,
437                "TTC": self.TTC_intervals_final,
438                "SAMPLE": self.sample_intervals_final,
439                "TTC Actual": np.full(num_trials, np.nan),
440            }
441
442            self.program_schedule_df = pd.DataFrame(data)
443
444            logger.info("Initialized stimuli dataframe.")
445        except Exception as e:
446            logger.debug(f"Error Building Stimuli Frame: {e}.")
447            raise
448
449    def save_all_data(self):
450        """
451        Saves the experiment schedule and detailed event log DataFrames to Excel files.
452
453        Iterates through a dictionary mapping descriptive names to the DataFrame
454        references (`program_schedule_df`, `event_data.event_dataframe`) and calls
455        the static `save_df_to_xlsx` method for each.
456        """
457        dataframes = {
458            "Experiment Schedule": self.program_schedule_df,
459            "Detailed Event Log Data": self.event_data.event_dataframe,
460        }
461
462        for name, df_reference in dataframes.items():
463            self.save_df_to_xlsx(name, df_reference)
464
465    @staticmethod
466    def get_paired_index(i: int, num_stimuli: int) -> int | None:
467        """
468        Calculates the 0-indexed valve number for the stimulus for side two based on give index (i) for side one.
469
470        Based on the total number of stimuli (`num_stimuli`), determines the index
471        of the valve paired with the valve at index `i`. Handles cases for 2, 4, or 8 total stimuli.
472
473        Parameters
474        ----------
475        - **i** (*int*): The 0-indexed number of the valve on one side.
476        - **num_stimuli** (*int*): The total number of stimuli/valves being used (must be 2, 4, or 8).
477
478        Returns
479        -------
480        - *int | None*: The 0-indexed number of the paired valve, or None if `num_stimuli` is not 2, 4, or 8.
481        """
482        match num_stimuli:
483            case 2:
484                # if num stim is 2, we are only running the same stimulus on both sides,
485                # we will only see i=0 here, and its paired with its sibling i=4
486                return i + 4
487            case 4:
488                # if num stim is 4, we are only running 2 stimuli per side,
489                # we will see i=0 and i =1 here, paired with siblings i=5 & i=4 respectively
490                match i:
491                    case 0:  # valve 1
492                        return i + 5  # (5) / valve 6
493                    case 1:  # valve 2
494                        return i + 3  # (4) / valve 5
495            case 8:
496                match i:
497                    case 0:
498                        return i + 5
499                    case 1:
500                        return i + 3
501                    case 2:
502                        return i + 5
503                    case 3:
504                        return i + 3
505            case _:
506                return None
507
508    @staticmethod
509    def save_df_to_xlsx(name: str, dataframe: pd.DataFrame) -> None:
510        """
511        Saves a pandas DataFrame to an Excel (.xlsx) file using a file save dialog.
512
513        Prompts the user with a standard 'Save As' dialog window. Suggests a default
514        filename including the provided `name` and the current date. Sets the default
515        directory to '../data_outputs'. If the user confirms a filename, saves the
516        DataFrame; otherwise, logs that the save was cancelled.
517
518        Parameters
519        ----------
520        - **name** (*str*): A descriptive name used in the default filename (e.g., "Experiment Schedule").
521        - **dataframe** (*pd.DataFrame*): The pandas DataFrame to be saved.
522
523        Raises
524        ------
525        - Propagates errors from `filedialog.asksaveasfilename` or `dataframe.to_excel`.
526        """
527        try:
528            file_name = filedialog.asksaveasfilename(
529                defaultextension=".xlsx",
530                filetypes=[("Excel Files", "*.xlsx")],
531                initialfile=f"{name}, {datetime.date.today()}",
532                initialdir=Path(__file__).parent.parent.parent.resolve()
533                / "data_outputs",
534                title="Save Excel file",
535            )
536
537            if file_name:
538                dataframe.to_excel(file_name, index=False)
539            else:
540                logger.info("User cancelled saving the stimuli dataframe.")
541
542            logger.info("Data saved to xlsx files.")
543        except Exception as e:
544            logger.error(f"Error saving data to xlsx: {e}")
545            raise
546
547    @staticmethod
548    def convert_seconds_to_minutes_seconds(total_seconds: int) -> tuple[int, int]:
549        """
550        Converts a total number of seconds into whole minutes and remaining seconds.
551
552        Parameters
553        ----------
554        - **total_seconds** (*int*): The total duration in seconds.
555
556        Returns
557        -------
558        - *tuple[int, int]*: A tuple containing (minutes, seconds).
559
560        Raises
561        ------
562        - Propagates `TypeError` if `total_seconds` is not a number.
563        """
564        try:
565            # Integer division to get whole minutes
566            minutes = total_seconds // 60
567            # Modulo to get remaining seconds
568            seconds = total_seconds % 60
569            return minutes, seconds
570        except Exception as e:
571            logger.error(f"Error converting seconds to minutes and seconds: {e}")
572            raise