Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ processing \ aurora \ frequency_bands.py: 93%
71 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-10 00:11 -0800
1"""
2Module containing FrequencyBands class representing a collection of Frequency Band objects.
3"""
5from typing import Generator, List, Optional, Union
7import numpy as np
8import pandas as pd
9from loguru import logger
11from mt_metadata.common.band import Band
14class FrequencyBands:
15 """
16 Collection of Band objects, typically used at a single decimation level.
18 Attributes
19 ----------
20 _band_edges : pd.DataFrame
21 DataFrame with columns ['lower_bound', 'upper_bound'] containing
22 frequency band boundaries
23 """
25 def __init__(
26 self,
27 band_edges: Optional[Union[np.ndarray, pd.DataFrame]] = None,
28 ):
29 """
30 Parameters
31 ----------
32 band_edges : np.ndarray or pd.DataFrame, optional
33 If numpy array: 2D array with columns [lower_bound, upper_bound]
34 If DataFrame: Must have columns ['lower_bound', 'upper_bound']
35 """
36 if band_edges is not None:
37 self.band_edges = band_edges
38 else:
39 self._band_edges = pd.DataFrame(columns=["lower_bound", "upper_bound"])
41 def __str__(self) -> str:
42 """Returns a Description of frequency bands"""
43 intro = "Frequency Bands:"
44 return f"{intro} \n{self._band_edges}"
46 def __repr__(self):
47 return self.__str__()
49 @property
50 def band_edges(self) -> pd.DataFrame:
51 """Get band edges as a DataFrame"""
52 return self._band_edges
54 @band_edges.setter
55 def band_edges(self, value: Union[np.ndarray, pd.DataFrame]) -> None:
56 """
57 Set band edges from either numpy array or DataFrame
59 Parameters
60 ----------
61 value : np.ndarray or pd.DataFrame
62 Band edge definitions
63 """
64 if isinstance(value, np.ndarray):
65 if value.ndim != 2 or value.shape[1] != 2:
66 raise ValueError("band_edges array must be 2D with shape (n_bands, 2)")
67 self._band_edges = pd.DataFrame(
68 value, columns=["lower_bound", "upper_bound"]
69 )
70 elif isinstance(value, pd.DataFrame):
71 required_cols = ["lower_bound", "upper_bound"]
72 if not all(col in value.columns for col in required_cols):
73 raise ValueError(f"DataFrame must contain columns {required_cols}")
74 self._band_edges = value[required_cols].copy()
75 else:
76 raise TypeError("band_edges must be numpy array or DataFrame")
78 # Reset index to ensure 0-based integer indexing
79 self._band_edges.reset_index(drop=True, inplace=True)
81 @property
82 def number_of_bands(self) -> int:
83 """Number of frequency bands"""
84 return len(self._band_edges)
86 @property
87 def array(self) -> np.ndarray:
88 """Get band edges as numpy array"""
89 return self._band_edges.values
91 def sort(self, by: str = "center_frequency", ascending: bool = True) -> None:
92 """
93 Sort bands by specified criterion.
95 Parameters
96 ----------
97 by : str
98 Criterion to sort by:
99 - "lower_bound": Sort by lower frequency bound
100 - "upper_bound": Sort by upper frequency bound
101 - "center_frequency": Sort by geometric center frequency (default)
102 ascending : bool
103 If True, sort in ascending order, else descending
104 """
105 if by in ["lower_bound", "upper_bound"]:
106 self._band_edges.sort_values(by=by, ascending=ascending, inplace=True)
107 elif by == "center_frequency":
108 centers = self.band_centers()
109 self._band_edges = self._band_edges.iloc[
110 np.argsort(centers)[:: (-1 if not ascending else 1)]
111 ].reset_index(drop=True)
112 else:
113 raise ValueError(
114 f"Invalid sort criterion: {by}. Must be one of: "
115 "'lower_bound', 'upper_bound', 'center_frequency'"
116 )
118 def bands(
119 self,
120 direction: str = "increasing_frequency",
121 sortby: Optional[str] = None,
122 rtype: str = "list",
123 ) -> Union[List[Band], Generator[Band, None, None]]:
124 """
125 Generate Band objects in specified order.
127 Parameters
128 ----------
129 direction : str
130 Order of iteration: "increasing_frequency" or "increasing_period"
131 sortby : str, optional
132 Sort bands before iteration:
133 - "lower_bound": Sort by lower frequency bound
134 - "upper_bound": Sort by upper frequency bound
135 - "center_frequency": Sort by geometric center frequency
136 If None, uses existing order
137 rtype : str
138 Return type: "list" or "generator". Default is "list" for easier reuse.
139 Use "generator" for memory efficiency when bands are only iterated once.
141 Returns
142 -------
143 Union[List[Band], Generator[Band, None, None]]
144 Band objects for each frequency band, either as a list or generator
145 depending on rtype parameter.
146 """
147 if sortby is not None or direction == "increasing_period":
148 # Create a copy to avoid modifying original
149 temp_bands = FrequencyBands(self._band_edges.copy())
150 temp_bands.sort(
151 by=sortby or "center_frequency",
152 ascending=(direction == "increasing_frequency"),
153 )
154 bands_to_iterate = temp_bands
155 else:
156 bands_to_iterate = self
158 # Create generator
159 def band_generator():
160 for idx in range(bands_to_iterate.number_of_bands):
161 yield bands_to_iterate.band(idx)
163 # Return as requested type
164 if rtype == "generator":
165 return band_generator()
166 elif rtype == "list":
167 return list(band_generator())
168 else:
169 raise ValueError("rtype must be either 'list' or 'generator'")
171 def band(self, i_band: int) -> Band:
172 """
173 Get specific frequency band.
175 Parameters
176 ----------
177 i_band : int
178 Index of band to return (zero-based)
180 Returns
181 -------
182 Band
183 Frequency band object
184 """
185 row = self._band_edges.iloc[i_band]
186 return Band(frequency_min=row["lower_bound"], frequency_max=row["upper_bound"])
188 def band_centers(self, frequency_or_period: str = "frequency") -> np.ndarray:
189 """
190 Calculate center frequencies/periods for all bands.
192 Parameters
193 ----------
194 frequency_or_period : str
195 Return values in "frequency" (Hz) or "period" (s)
197 Returns
198 -------
199 np.ndarray
200 Center frequencies/periods for each band
201 """
202 band_centers = np.array(
203 [self.band(i).center_frequency for i in range(self.number_of_bands)]
204 )
206 if frequency_or_period == "period":
207 band_centers = 1.0 / band_centers
209 return band_centers
211 def validate(self) -> None:
212 """
213 Validate and potentially reorder bands based on center frequencies.
214 """
215 band_centers = self.band_centers()
217 # Check if band centers are monotonically increasing
218 if not np.all(band_centers[1:] > band_centers[:-1]):
219 logger.warning(
220 "Band centers are not monotonic. Attempting to reorganize bands."
221 )
222 self.sort(by="center_frequency")