diff --git a/code/useful_functions.py b/code/useful_functions.py index c88d006..796ccb5 100644 --- a/code/useful_functions.py +++ b/code/useful_functions.py @@ -6,6 +6,67 @@ import rlxnix as rlx from IPython import embed from scipy.signal import welch +def all_coming_together(freq_array, power_array, points_list, categories, num_harmonics_list, colors, delta=2.5, threshold=0.5): + """ + Process a list of points, calculating integrals, checking validity, and preparing harmonics for valid points. + + Parameters + ---------- + freq_array : np.array + Array of frequencies corresponding to the power values. + power_array : np.array + Array of power spectral density values. + points_list : list + List of harmonic frequency points to process. + categories : list + List of corresponding categories for each point. + num_harmonics_list : list + List of the number of harmonics for each point. + colors : list + List of colors corresponding to each point's category. + delta : float, optional + Radius of the range for integration around each point (default is 2.5). + threshold : float, optional + Threshold value to compare integrals with local mean (default is 0.5). + + Returns + ------- + valid_points : list + A list of valid points with their harmonics. + color_mapping : dict + A dictionary mapping categories to corresponding colors. + category_harmonics : dict + A mapping of categories to their harmonic frequencies. + messages : list + A list of messages for each point, stating whether it was valid or not. + """ + valid_points = [] + color_mapping = {} + category_harmonics = {} + messages = [] + + for i, point in enumerate(points_list): + category = categories[i] + num_harmonics = num_harmonics_list[i] + color = colors[i] + + # Step 1: Calculate the integral for the point + integral, local_mean, _ = calculate_integral(freq_array, power_array, point, delta) + + # Step 2: Check if the point is valid + valid = valid_integrals(integral, local_mean, point, threshold) + if valid: + # Step 3: Prepare harmonics if the point is valid + harmonics, color_map, category_harm = prepare_harmonic(point, category, num_harmonics, color) + valid_points.append((point, harmonics)) + color_mapping.update(color_map) + category_harmonics.update(category_harm) + messages.append(f"The point {point} is valid.") + else: + messages.append(f"The point {point} is not valid.") + + return valid_points, color_mapping, category_harmonics, messages + def AM(EODf, stimulus): """ Calculates the Amplitude Modulation and Nyquist frequency @@ -215,38 +276,36 @@ def power_spectrum(stimulus): freq, power = welch(rate, fs = 1/dt, nperseg = 2**16, noverlap = 2**15) return freq, power -def prepare_harmonics(frequencies, categories, num_harmonics, colors): +def prepare_harmonic(frequency, category, num_harmonics, color): """ - Prepare harmonic frequencies and assign colors based on categories. - + Prepare harmonic frequencies and assign color based on category for a single point. + Parameters ---------- - frequencies : list - Base frequencies to generate harmonics. - categories : list - Corresponding categories for the base frequencies. - num_harmonics : list - Number of harmonics for each base frequency. - colors : list - List of colors corresponding to the categories. + frequency : float + Base frequency to generate harmonics. + category : str + Corresponding category for the base frequency. + num_harmonics : int + Number of harmonics for the base frequency. + color : str + Color corresponding to the category. Returns ------- - points : list - A flat list of harmonic frequencies. + harmonics : list + A list of harmonic frequencies. color_mapping : dict - A dictionary mapping each category to its corresponding color. - points_categories : dict - A mapping of categories to their harmonic frequencies. + A dictionary mapping the category to its corresponding color. + category_harmonics : dict + A mapping of the category to its harmonic frequencies. """ - points_categories = {} - for idx, (freq, category) in enumerate(zip(frequencies, categories)): - points_categories[category] = [freq * (i + 1) for i in range(num_harmonics[idx])] - - points = [p for harmonics in points_categories.values() for p in harmonics] - color_mapping = {category: colors[idx] for idx, category in enumerate(categories)} + harmonics = [frequency * (i + 1) for i in range(num_harmonics)] + + color_mapping = {category: color} + category_harmonics = {category: harmonics} - return points, color_mapping, points_categories + return harmonics, color_mapping, category_harmonics def remove_poor(files): """